<commit_before>pub type Register = u8;
pub type Address = u16;
pub type Value = u8;

use std::fmt;

#[derive(Debug)]
#[derive(PartialEq)]
pub enum Instruction {
    Clear,
    Return,
    Jump(Address),
    Call(Address),
    SkipIfEqual(Register, Value),
    SkipIfNotEqual(Register, Value),
    LoadVal(Register, Value),
    AddVal(Register, Value),
    LoadReg(Register, Register),
    Or(Register, Register),
    And(Register, Register),
    Xor(Register, Register),
    AddReg(Register, Register),
    SubReg(Register, Register),
    ShiftRight(Register),
    ShiftLeft(Register),
    SetIndexRegister(Address),
    Random(Register, Value),
    Draw(Register, Register, Value),
    AddIndex(Register),
    LoadDigit(Register),
    StoreIndex(Register),
    ReadIndex(Register),
    InvalidOperation,
}

impl Instruction {
    pub fn parse(val: u16) -> Instruction {
        match val & 0xF000 {
            0x0000 => {
                match val & 0x000F {
                    0x0000 => Instruction::Clear,
                    0x000E => Instruction::Return,
                    _ => Instruction::InvalidOperation,
                }
            }
            0x1000 => Instruction::Jump(get_address(val)),
            0x2000 => Instruction::Call(get_address(val)),
            0x3000 => Instruction::SkipIfEqual(get_first_register(val), get_value(val)),
            0x4000 => Instruction::SkipIfNotEqual(get_first_register(val), get_value(val)),
            0x6000 => Instruction::LoadVal(get_first_register(val), get_value(val)),
            0x7000 => Instruction::AddVal(get_first_register(val), get_value(val)),
            0x8000 => {
                match val & 0x000F {
                    0x0000 => {
                        Instruction::LoadReg(get_first_register(val), get_second_register(val))
                    }
                    0x0001 => Instruction::Or(get_first_register(val), get_second_register(val)),
                    0x0002 => Instruction::And(get_first_register(val), get_second_register(val)),
                    0x0003 => Instruction::Xor(get_first_register(val), get_second_register(val)),
                    0x0004 => {
                        Instruction::AddReg(get_first_register(val), get_second_register(val))
                    }
                    0x0005 => {
                        Instruction::SubReg(get_first_register(val), get_second_register(val))
                    }
                    0x0006 => Instruction::ShiftRight(get_first_register(val)),
                    0x000E => Instruction::ShiftLeft(get_first_register(val)),
                    _ => Instruction::InvalidOperation,
                }
            }
            0xA000 => Instruction::SetIndexRegister(get_address(val)),
            0xC000 => Instruction::Random(get_first_register(val), get_value(val)),
            0xD000 => {
                Instruction::Draw(
                    get_first_register(val),
                    get_second_register(val),
                    get_value(val),
                )
            }
            0xF000 => {
                match val & 0x00FF {
                    0x001E => Instruction::AddIndex(get_first_register(val)),
                    0x0029 => Instruction::LoadDigit(get_first_register(val)),
                    0x0055 => Instruction::StoreIndex(get_first_register(val)),
                    0x0065 => Instruction::ReadIndex(get_first_register(val)),
                    _ => Instruction::InvalidOperation,
                }
            }
            _ => Instruction::InvalidOperation,
        }
    }
}

impl fmt::Display for Instruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pretty_instruction = match *self {
            Instruction::Clear => format!("CLS"),
            Instruction::Return => format!("RET"),
            Instruction::Jump(a) => format!("JP 0x{:X}", a),
            Instruction::Call(a) => format!("CALL 0x{:X}", a),
            Instruction::SkipIfEqual(r, v) => format!("SE V{:X}, {}", r, v),
            Instruction::SkipIfNotEqual(r, v) => format!("SNE V{:X}, {}", r, v),
            Instruction::LoadVal(r, v) => format!("LD V{:X}, {}", r, v),
            Instruction::AddVal(r, v) => format!("ADD V{:X}, {}", r, v),
            Instruction::LoadReg(r1, r2) => format!("LD V{:X}, V{:X}", r1, r2),
            Instruction::Or(r1, r2) => format!("OR V{:X}, V{:X}", r1, r2),
            Instruction::And(r1, r2) => format!("AND V{:X}, V{:X}", r1, r2),
            Instruction::Xor(r1, r2) => format!("XOR V{:X}, V{:X}", r1, r2),
            Instruction::AddReg(r1, r2) => format!("ADD V{:X}, V{:X}", r1, r2),
            Instruction::SubReg(r1, r2) => format!("SUB V{:X}, V{:X}", r1, r2),
            Instruction::ShiftRight(r) => format!("SHR V{:X}", r),
            Instruction::ShiftLeft(r) => format!("SHL V{:X}", r),
            Instruction::SetIndexRegister(a) => format!("LD I, 0x{:X}", a),
            Instruction::Random(r, v) => format!("RND V{:X}, {}", r, v),
            Instruction::Draw(r1, r2, v) => format!("DRW V{:X}, V{:X}, {}", r1, r2, v),
            Instruction::AddIndex(r) => format!("ADD I, V{:X}", r),
            Instruction::LoadDigit(r) => format!("LD F, V{:X}", r),
            Instruction::StoreIndex(r) => format!("LD [I], V{:X}", r),
            Instruction::ReadIndex(r) => format!("LD V{:X} [I]", r),
            Instruction::InvalidOperation => format!("INVALID OPERATION"),
        };
        write!(f, "{}", pretty_instruction)
    }
}

fn get_first_register(val: u16) -> u8 {
    ((val & 0x0F00) >> 8) as u8
}

fn get_second_register(val: u16) -> u8 {
    ((val & 0x00F0) >> 4) as u8
}

fn get_value(val: u16) -> u8 {
    (val & 0x00FF) as u8
}

fn get_address(val: u16) -> u16 {
    (val & 0x0FFF)
}<commit_msg>Renamed opcode parsing methods<commit_after>pub type Register = u8;
pub type Address = u16;
pub type Value = u8;

use std::fmt;

#[derive(Debug)]
#[derive(PartialEq)]
pub enum Instruction {
    Clear,
    Return,
    Jump(Address),
    Call(Address),
    SkipIfEqual(Register, Value),
    SkipIfNotEqual(Register, Value),
    LoadVal(Register, Value),
    AddVal(Register, Value),
    LoadReg(Register, Register),
    Or(Register, Register),
    And(Register, Register),
    Xor(Register, Register),
    AddReg(Register, Register),
    SubReg(Register, Register),
    ShiftRight(Register),
    ShiftLeft(Register),
    SetIndexRegister(Address),
    Random(Register, Value),
    Draw(Register, Register, Value),
    AddIndex(Register),
    LoadDigit(Register),
    StoreIndex(Register),
    ReadIndex(Register),
    InvalidOperation,
}

impl Instruction {
    pub fn parse(val: u16) -> Instruction {
        match val & 0xF000 {
            0x0000 => {
                match val & 0x000F {
                    0x0000 => Instruction::Clear,
                    0x000E => Instruction::Return,
                    _ => Instruction::InvalidOperation,
                }
            }
            0x1000 => Instruction::Jump(addr(val)),
            0x2000 => Instruction::Call(addr(val)),
            0x3000 => Instruction::SkipIfEqual(reg1(val), value(val)),
            0x4000 => Instruction::SkipIfNotEqual(reg1(val), value(val)),
            0x6000 => Instruction::LoadVal(reg1(val), value(val)),
            0x7000 => Instruction::AddVal(reg1(val), value(val)),
            0x8000 => {
                match val & 0x000F {
                    0x0000 => Instruction::LoadReg(reg1(val), reg2(val)),
                    0x0001 => Instruction::Or(reg1(val), reg2(val)),
                    0x0002 => Instruction::And(reg1(val), reg2(val)),
                    0x0003 => Instruction::Xor(reg1(val), reg2(val)),
                    0x0004 => Instruction::AddReg(reg1(val), reg2(val)),
                    0x0005 => Instruction::SubReg(reg1(val), reg2(val)),
                    0x0006 => Instruction::ShiftRight(reg1(val)),
                    0x000E => Instruction::ShiftLeft(reg1(val)),
                    _ => Instruction::InvalidOperation,
                }
            }
            0xA000 => Instruction::SetIndexRegister(addr(val)),
            0xC000 => Instruction::Random(reg1(val), value(val)),
            0xD000 => Instruction::Draw(reg1(val), reg2(val), value(val)),
            0xF000 => {
                match val & 0x00FF {
                    0x001E => Instruction::AddIndex(reg1(val)),
                    0x0029 => Instruction::LoadDigit(reg1(val)),
                    0x0055 => Instruction::StoreIndex(reg1(val)),
                    0x0065 => Instruction::ReadIndex(reg1(val)),
                    _ => Instruction::InvalidOperation,
                }
            }
            _ => Instruction::InvalidOperation,
        }
    }
}

impl fmt::Display for Instruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pretty_instruction = match *self {
            Instruction::Clear => format!("CLS"),
            Instruction::Return => format!("RET"),
            Instruction::Jump(a) => format!("JP 0x{:X}", a),
            Instruction::Call(a) => format!("CALL 0x{:X}", a),
            Instruction::SkipIfEqual(r, v) => format!("SE V{:X}, {}", r, v),
            Instruction::SkipIfNotEqual(r, v) => format!("SNE V{:X}, {}", r, v),
            Instruction::LoadVal(r, v) => format!("LD V{:X}, {}", r, v),
            Instruction::AddVal(r, v) => format!("ADD V{:X}, {}", r, v),
            Instruction::LoadReg(r1, r2) => format!("LD V{:X}, V{:X}", r1, r2),
            Instruction::Or(r1, r2) => format!("OR V{:X}, V{:X}", r1, r2),
            Instruction::And(r1, r2) => format!("AND V{:X}, V{:X}", r1, r2),
            Instruction::Xor(r1, r2) => format!("XOR V{:X}, V{:X}", r1, r2),
            Instruction::AddReg(r1, r2) => format!("ADD V{:X}, V{:X}", r1, r2),
            Instruction::SubReg(r1, r2) => format!("SUB V{:X}, V{:X}", r1, r2),
            Instruction::ShiftRight(r) => format!("SHR V{:X}", r),
            Instruction::ShiftLeft(r) => format!("SHL V{:X}", r),
            Instruction::SetIndexRegister(a) => format!("LD I, 0x{:X}", a),
            Instruction::Random(r, v) => format!("RND V{:X}, {}", r, v),
            Instruction::Draw(r1, r2, v) => format!("DRW V{:X}, V{:X}, {}", r1, r2, v),
            Instruction::AddIndex(r) => format!("ADD I, V{:X}", r),
            Instruction::LoadDigit(r) => format!("LD F, V{:X}", r),
            Instruction::StoreIndex(r) => format!("LD [I], V{:X}", r),
            Instruction::ReadIndex(r) => format!("LD V{:X} [I]", r),
            Instruction::InvalidOperation => format!("INVALID OPERATION"),
        };
        write!(f, "{}", pretty_instruction)
    }
}

fn reg1(val: u16) -> u8 {
    ((val & 0x0F00) >> 8) as u8
}

fn reg2(val: u16) -> u8 {
    ((val & 0x00F0) >> 4) as u8
}

fn value(val: u16) -> u8 {
    (val & 0x00FF) as u8
}

fn addr(val: u16) -> u16 {
    (val & 0x0FFF)
}
<|endoftext|>


<commit_before>//pub use neuralops_kernels::*;
//pub use neuralops_omp_kernels::*;

use libc::*;

#[cfg(not(feature = "iomp"))]
#[link(name = "gomp")]
extern "C" {}

#[cfg(feature = "iomp")]
#[link(name = "iomp5")]
extern "C" {}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_rect_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_rect_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
  pub fn neuralops_logistic_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_logistic_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_batchnorm2d_fwd_mean(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *mut f32);
  pub fn neuralops_batchnorm2d_fwd_var(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      var: *mut f32);
  pub fn neuralops_batchnorm2d_fwd_output(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      out_buf: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_var(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      out_grad: *const f32,
      var_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_mean(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      var_grad: *const f32,
      out_grad: *const f32,
      mean_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_input(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      mean_grad: *const f32,
      var: *const f32,
      //run_var: *const f32,
      var_grad: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_conv2d_bias_fwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      bias: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_conv2d_bias_bwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      out_grad: *const f32,
      bias_grad: *mut f32);
      //in_grad: *mut f32);
  pub fn neuralops_conv2d_scale_bias_fwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      scale: *const f32,
      bias: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_conv2d_scale_bias_bwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      scale: *const f32,
      out_grad: *const f32,
      scale_grad: *mut f32,
      bias_grad: *mut f32,
      in_grad: *mut f32);
  pub fn neuralops_caffe_im2col(
      data_im: *const f32,
      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: *mut f32);
  pub fn neuralops_caffe_col2im(
      data_col: *const f32,
      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: *mut f32);
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_image_crop(
      in_width: size_t,
      in_height: size_t,
      chan: size_t,
      crop_w: size_t,
      crop_h: size_t,
      offset_x: ptrdiff_t,
      offset_y: ptrdiff_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
  pub fn neuralops_image_flip(
      width: size_t,
      height: size_t,
      chan: size_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_interpolate2d_catmullrom(
      in_width: size_t,
      in_height: size_t,
      chan: size_t,
      out_width: size_t,
      out_height: size_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_avgpool2d_2x2_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_avgpool2d_2x2_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32);
  pub fn neuralops_avgpool2d_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_buf: *mut f32,
      pool_w: usize,
      pool_h: usize);
  pub fn neuralops_avgpool2d_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
      pool_w: usize,
      pool_h: usize);
}

#[link(name = "neuralops_omp_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_omp_rect_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_omp_rect_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
  pub fn neuralops_omp_logistic_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_omp_logistic_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );

  pub fn neuralops_omp_caffe_im2col(
      data_im: *const f32,
      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: *mut f32);
  pub fn neuralops_omp_caffe_col2im(
      data_col: *const f32,
      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: *mut f32);

  pub fn neuralops_omp_caffe_avgpool2d_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_width: usize,
      out_height: usize,
      out_buf: *mut f32,
      pool_w: usize,
      pool_h: usize,
      stride_w: usize,
      stride_h: usize,
      pad_w: usize,
      pad_h: usize);
  pub fn neuralops_omp_caffe_avgpool2d_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      //in_buf: *const f32,
      out_width: usize,
      out_height: usize,
      out_grad: *const f32,
      in_grad: *mut f32,
      pool_w: usize,
      pool_h: usize,
      stride_w: usize,
      stride_h: usize,
      pad_w: usize,
      pad_h: usize);
}
<commit_msg>When compiling Intel OMP versions, link to icc runtime libraries.<commit_after>//pub use neuralops_kernels::*;
//pub use neuralops_omp_kernels::*;

use libc::*;

#[cfg(not(feature = "iomp"))]
#[link(name = "gomp")]
extern "C" {}

#[cfg(feature = "iomp")]
#[link(name = "iomp5")]
extern "C" {}

#[cfg(feature = "iomp")]
#[link(name = "imf")]
extern "C" {}

#[cfg(feature = "iomp")]
#[link(name = "irc")]
extern "C" {}

#[cfg(feature = "iomp")]
#[link(name = "irc_s")]
extern "C" {}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_rect_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_rect_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
  pub fn neuralops_logistic_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_logistic_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_batchnorm2d_fwd_mean(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *mut f32);
  pub fn neuralops_batchnorm2d_fwd_var(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      var: *mut f32);
  pub fn neuralops_batchnorm2d_fwd_output(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      out_buf: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_var(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      out_grad: *const f32,
      var_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_mean(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      var: *const f32,
      //run_var: *const f32,
      var_grad: *const f32,
      out_grad: *const f32,
      mean_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
  pub fn neuralops_batchnorm2d_bwd_input(
      batch_sz: size_t,
      width: size_t,
      height: size_t,
      chan: size_t,
      in_buf: *const f32,
      mean: *const f32,
      //run_mean: *const f32,
      mean_grad: *const f32,
      var: *const f32,
      //run_var: *const f32,
      var_grad: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
      //gamma: f32,
      epsilon: f32);
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_conv2d_bias_fwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      bias: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_conv2d_bias_bwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      out_grad: *const f32,
      bias_grad: *mut f32);
      //in_grad: *mut f32);
  pub fn neuralops_conv2d_scale_bias_fwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      scale: *const f32,
      bias: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_conv2d_scale_bias_bwd(
      batch_sz: usize,
      out_width: usize,
      out_height: usize,
      out_chan: usize,
      in_buf: *const f32,
      scale: *const f32,
      out_grad: *const f32,
      scale_grad: *mut f32,
      bias_grad: *mut f32,
      in_grad: *mut f32);
  pub fn neuralops_caffe_im2col(
      data_im: *const f32,
      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: *mut f32);
  pub fn neuralops_caffe_col2im(
      data_col: *const f32,
      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: *mut f32);
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_image_crop(
      in_width: size_t,
      in_height: size_t,
      chan: size_t,
      crop_w: size_t,
      crop_h: size_t,
      offset_x: ptrdiff_t,
      offset_y: ptrdiff_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
  pub fn neuralops_image_flip(
      width: size_t,
      height: size_t,
      chan: size_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_interpolate2d_catmullrom(
      in_width: size_t,
      in_height: size_t,
      chan: size_t,
      out_width: size_t,
      out_height: size_t,
      in_pixels: *const f32,
      out_pixels: *mut f32,
  );
}

#[link(name = "neuralops_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_avgpool2d_2x2_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_buf: *mut f32);
  pub fn neuralops_avgpool2d_2x2_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32);
  pub fn neuralops_avgpool2d_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_buf: *mut f32,
      pool_w: usize,
      pool_h: usize);
  pub fn neuralops_avgpool2d_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
      pool_w: usize,
      pool_h: usize);
}

#[link(name = "neuralops_omp_kernels", kind = "static")]
extern "C" {
  pub fn neuralops_omp_rect_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_omp_rect_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );
  pub fn neuralops_omp_logistic_fwd(
      batch_sz: size_t,
      dim: size_t,
      in_buf: *const f32,
      out_buf: *mut f32,
  );
  pub fn neuralops_omp_logistic_bwd(
      batch_sz: size_t,
      dim: size_t,
      out_buf: *const f32,
      out_grad: *const f32,
      in_grad: *mut f32,
  );

  pub fn neuralops_omp_caffe_im2col(
      data_im: *const f32,
      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: *mut f32);
  pub fn neuralops_omp_caffe_col2im(
      data_col: *const f32,
      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: *mut f32);

  pub fn neuralops_omp_caffe_avgpool2d_fwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      in_buf: *const f32,
      out_width: usize,
      out_height: usize,
      out_buf: *mut f32,
      pool_w: usize,
      pool_h: usize,
      stride_w: usize,
      stride_h: usize,
      pad_w: usize,
      pad_h: usize);
  pub fn neuralops_omp_caffe_avgpool2d_bwd(
      batch_sz: usize,
      in_width: usize,
      in_height: usize,
      chan: usize,
      //in_buf: *const f32,
      out_width: usize,
      out_height: usize,
      out_grad: *const f32,
      in_grad: *mut f32,
      pool_w: usize,
      pool_h: usize,
      stride_w: usize,
      stride_h: usize,
      pad_w: usize,
      pad_h: usize);
}
<|endoftext|>


<commit_before>
#[cfg(test)]
mod test {
    use super::super::*;
    use sdl2::keyboard::Keycode;

    #[test]
    fn test_default()
    {
        let mut keypad = Keypad::default();
        let keys: [Keycode; 16] = [X, Num1, Num2, Num3, Q, W, E, A, S, D, Z, C, Num4, R, F, V];

        for i in 0..keys.len()
        {
            assert_eq!(keypad.pressed(i), false);
        }
    }

    #[test]
    fn test_pressing_keys()
    {
        let mut keypad = Keypad::default();
        let keys: [Keycode; 16] = [X, Num1, Num2, Num3, Q, W, E, A, S, D, Z, C, Num4, R, F, V];

        for i in 0..keys.len()
        {
            keypad.press(keys[i], true);
            assert!(keypad.pressed(i));

            keypad.press(keys[i], false);
            assert_eq!(keypad.pressed(i), false);
        }
    }
}<commit_msg>whitespace<commit_after>#[cfg(test)]
mod test {
    use super::super::*;
    use sdl2::keyboard::Keycode;

    #[test]
    fn test_default()
    {
        let mut keypad = Keypad::default();
        // Ordered based on order in file
        let keys: [Keycode; 16] = [X, Num1, Num2, Num3, Q, W, E, A, S, D, Z, C, Num4, R, F, V];

        for i in 0..keys.len()
        {
            assert_eq!(keypad.pressed(i), false);
        }
    }

    #[test]
    fn test_pressing_keys()
    {
        let mut keypad = Keypad::default();
        // Ordered based on order in file
        let keys: [Keycode; 16] = [X, Num1, Num2, Num3, Q, W, E, A, S, D, Z, C, Num4, R, F, V];

        for i in 0..keys.len()
        {
            keypad.press(keys[i], true);
            assert!(keypad.pressed(i));

            keypad.press(keys[i], false);
            assert_eq!(keypad.pressed(i), false);
        }
    }
}<|endoftext|>


<commit_before>//! It's a tree!

use std::ptr;

use rustwlc::handle::WlcView;

use super::container::*;

#[derive(Debug, PartialEq, Eq)]
pub struct Node {
    // We need a mut pointer so we can modify the parent
    parent: *mut Node,
    val: Container,
    children: Vec<Node>
}

impl Node {
    /// Create a new node with the existing value.
    /// For root-style constructors.
    pub fn new(val: Container) -> Node {
        Node {
            parent: ptr::null_mut(),
            val: val,
            children: Vec::new()
        }
    }

    /// Add a new child node to this node, using a value
    pub fn new_child(&mut self, val: Container) -> &mut Node {
        let self_mut = self as *mut Node;
        self.children.push(Node {
            parent: self_mut,
            val: val,
            children: Vec::new()
        });
        let last_ix = self.children.len() -1;
        &mut self.children[last_ix]
    }

    /// Whether this node has a (currently-reachable) parent
    #[allow(dead_code)]
    pub fn has_parent(&self) -> bool {
        !self.parent.is_null()
    }

    /// Gets the type of container this node holds
    pub fn get_container_type(&self) -> ContainerType {
        self.val.get_type()
    }

    /// Tries to get an ancestor of the requested type
    pub fn get_ancestor_of_type(&self, container_type: ContainerType)
                                -> Option<&Node> {
        let mut maybe_parent = self.get_parent();
        loop {
            if let Some(parent) = maybe_parent {
                if parent.get_container_type() == container_type {
                    return Some(parent);
                }
                maybe_parent = parent.get_parent();
            }
            else {
                return None;
            }
        }
    }

    /// Gets a node by handle
    pub fn find_view_by_handle(&self, view_handle: &WlcView) -> Option<&Node>{
        match self.get_val() {
            &Container::View { ref handle, .. } => {
                if view_handle == handle {
                    Some(self)
                } else {
                    None
                }
            },
            _ => {
                for child in self.get_children() {
                    if let Some(view) = child.find_view_by_handle(view_handle) {
                        return Some(view);
                    }
                }
                return None;
            }
        }
    }

    /// Gets the parent of this node (if it exists)
    pub fn get_parent(&self) -> Option<&mut Node> {
        if self.parent.is_null() {
            return None;
        }
        unsafe {
            return Some(&mut *self.parent);
        }
    }

    /// Borrow the children of this node.
    pub fn get_children(&self) -> &[Node] {
        &self.children
    }

    /// Mutably borrow the children of this mutable node
    pub fn get_children_mut(&mut self) -> &mut[Node] {
        &mut self.children
    }

    /// Remove a child at the given index
    #[allow(dead_code)]
    pub fn remove_child_at(&mut self, index: usize) -> Node {
        let mut child = self.children.remove(index);
        child.parent = ptr::null_mut();
        child
    }

    /// Moves another node to be a sibling of this node.
    #[allow(dead_code)]
    pub fn add_sibling(&self, node: Node) -> Result<(), ()> {
        if let Some(parent) = self.get_parent() {
            node.move_to(parent);
            Ok(())
        }
        else {
            Err(())
        }
    }

    /// Whether this node is a parent of another node
    #[allow(dead_code)]
    pub fn is_parent_of(&self, other: &Node) -> bool {
        // Fun fact, other.parent == self as *const Node won't compile
        self as *const Node == other.parent
    }

    /// Remove a node from its parent.
    /// This method will mutate the parent if it exists.
    pub fn remove_from_parent(&mut self) -> Option<Node> {
        let mut result: Option<Node> = None;
        if let Some(mut parent) = self.get_parent() {
            if let Some(index) = parent.children.iter().position(|c| c == self) {
                result = Some(parent.children.remove(index));
            }
        }
        self.parent = ptr::null_mut();
        result
    }

    /// Removes a child from self
    pub fn remove_child(&mut self, other: &Node) -> Option<Node> {
        if let Some(index) = self.children.iter().position(|c| c == other) {
            let mut child = self.children.remove(index);
            child.parent = ptr::null_mut();
            Some(child)
        }
        else {
            None
        }
    }

    pub fn move_to(mut self, new_parent: &mut Node) {
        self.remove_from_parent();
        self.parent = new_parent as *mut Node;
        new_parent.children.push(self);
    }

    pub fn get_val(&self) -> &Container {
        &self.val
    }
}

impl Drop for Node {
    fn drop(&mut self) {
        println!("Dropping a node.");
        let children: &mut Vec<Node> = &mut self.children;
        for mut child in children {
            child.parent = ptr::null_mut();
        }
    }
}

unsafe impl Sync for Node {}
unsafe impl Send for Node {}

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

    /// Nodes can have children added to them
    #[test]
    fn add_child() {
        let mut root = Node::new(Container::Root);
        root.new_child(Container::Root);
        root.new_child(Container::Root); // This is okay
        {
            let mut third_child = root.new_child(Container::Root);
            third_child.new_child(Container::Root);
            //root.new_child(Root); // Have to wait for 3rd child to drop
        }
        root.new_child(Container::Root); // Now this works
        assert_eq!(root.children.len(), 4);
    }

    #[test]
    fn has_get_parent() {
        let mut root = Node::new(Container::Root);
        //assert!(!root.has_parent(), "Root has a parent");
        assert_eq!(root.get_parent(), None);

        let child = root.new_child(Container::Root);
        //assert!(child.has_parent(), "Child does not have parent");
        assert!(child.get_parent().is_some(), "Child does not have parent");
        let parent = child.get_parent().expect("Asserted child has parent");
        assert_eq!(parent.get_container_type(), ContainerType::Root);

    }
}
<commit_msg>Change remove_child_at to Option, add many Node tests<commit_after>//! It's a tree!

use std::ptr;

use rustwlc::handle::WlcView;

use super::container::*;

#[derive(Debug, PartialEq, Eq)]
pub struct Node {
    // We need a mut pointer so we can modify the parent
    parent: *mut Node,
    val: Container,
    children: Vec<Node>
}

impl Node {
    /// Create a new node with the existing value.
    /// For root-style constructors.
    pub fn new(val: Container) -> Node {
        Node {
            parent: ptr::null_mut(),
            val: val,
            children: Vec::new()
        }
    }

    /// Add a new child node to this node, using a value
    pub fn new_child(&mut self, val: Container) -> &mut Node {
        let self_mut = self as *mut Node;
        self.children.push(Node {
            parent: self_mut,
            val: val,
            children: Vec::new()
        });
        let last_ix = self.children.len() -1;
        &mut self.children[last_ix]
    }

    /// Whether this node has a (currently-reachable) parent
    #[allow(dead_code)]
    pub fn has_parent(&self) -> bool {
        !self.parent.is_null()
    }

    /// Gets the type of container this node holds
    pub fn get_container_type(&self) -> ContainerType {
        self.val.get_type()
    }

    /// Tries to get an ancestor of the requested type
    pub fn get_ancestor_of_type(&self, container_type: ContainerType)
                                -> Option<&Node> {
        let mut maybe_parent = self.get_parent();
        loop {
            if let Some(parent) = maybe_parent {
                if parent.get_container_type() == container_type {
                    return Some(parent);
                }
                maybe_parent = parent.get_parent();
            }
            else {
                return None;
            }
        }
    }

    /// Gets a node by handle
    pub fn find_view_by_handle(&self, view_handle: &WlcView) -> Option<&Node>{
        match self.get_val() {
            &Container::View { ref handle, .. } => {
                if view_handle == handle {
                    Some(self)
                } else {
                    None
                }
            },
            _ => {
                for child in self.get_children() {
                    if let Some(view) = child.find_view_by_handle(view_handle) {
                        return Some(view);
                    }
                }
                return None;
            }
        }
    }

    /// Gets the parent of this node (if it exists)
    pub fn get_parent(&self) -> Option<&mut Node> {
        if self.parent.is_null() {
            return None;
        }
        unsafe {
            return Some(&mut *self.parent);
        }
    }

    /// Borrow the children of this node.
    pub fn get_children(&self) -> &[Node] {
        &self.children
    }

    /// Mutably borrow the children of this mutable node
    pub fn get_children_mut(&mut self) -> &mut[Node] {
        &mut self.children
    }

    /// Remove a child at the given index
    #[allow(dead_code)]
    pub fn remove_child_at(&mut self, index: usize) -> Option<Node> {
        if index > self.children.len() - 1 {
            return None;
        }
        let mut child = self.children.remove(index);
        child.parent = ptr::null_mut();
        return Some(child);
    }

    /// Moves another node to be a sibling of this node.
    #[allow(dead_code)]
    pub fn add_sibling(&self, node: Node) -> Result<(), ()> {
        if let Some(parent) = self.get_parent() {
            node.move_to(parent);
            Ok(())
        }
        else {
            Err(())
        }
    }

    /// Whether this node is a parent of another node
    #[allow(dead_code)]
    pub fn is_parent_of(&self, other: &Node) -> bool {
        // Fun fact, other.parent == self as *const Node won't compile
        self as *const Node == other.parent
    }

    /// Remove a node from its parent.
    /// This method will mutate the parent if it exists.
    pub fn remove_from_parent(&mut self) -> Option<Node> {
        let mut result: Option<Node> = None;
        if let Some(mut parent) = self.get_parent() {
            if let Some(index) = parent.children.iter().position(|c| c == self) {
                result = Some(parent.children.remove(index));
            }
        }
        self.parent = ptr::null_mut();
        result
    }

    /// Removes a child from self
    pub fn remove_child(&mut self, other: &Node) -> Option<Node> {
        if let Some(index) = self.children.iter().position(|c| c == other) {
            let mut child = self.children.remove(index);
            child.parent = ptr::null_mut();
            Some(child)
        }
        else {
            None
        }
    }

    pub fn move_to(mut self, new_parent: &mut Node) {
        self.remove_from_parent();
        self.parent = new_parent as *mut Node;
        new_parent.children.push(self);
    }

    pub fn get_val(&self) -> &Container {
        &self.val
    }
}

impl Drop for Node {
    fn drop(&mut self) {
        println!("Dropping a node.");
        let children: &mut Vec<Node> = &mut self.children;
        for mut child in children {
            child.parent = ptr::null_mut();
        }
    }
}

unsafe impl Sync for Node {}
unsafe impl Send for Node {}

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

    /// Nodes can have children added to them
    #[test]
    fn new_child() {
        let mut root = Node::new(Container::Root);
        root.new_child(Container::Root);
        root.new_child(Container::Root); // This is okay
        {
            let mut third_child = root.new_child(Container::Root);
            third_child.new_child(Container::Root);
            //root.new_child(Root); // Have to wait for 3rd child to drop
        }
        root.new_child(Container::Root); // Now this works
        assert_eq!(root.children.len(), 4);
    }

    #[test]
    fn has_get_parent() {
        let mut root = Node::new(Container::Root);
        assert!(!root.has_parent(), "Root has a parent");
        assert_eq!(root.get_parent(), None);

        let child = root.new_child(Container::Root);
        assert!(child.has_parent(), "Child does not have parent");
        assert!(child.get_parent().is_some(), "Child does not have parent");
        let parent = child.get_parent().expect("Asserted child has parent");
        assert_eq!(parent.get_container_type(), ContainerType::Root);
    }

    #[test]
    fn get_container_type() {
        let mut root = Node::new(Container::Root);
        assert_eq!(root.get_container_type(), ContainerType::Root);
        {
            let wksp = root.new_child(
                Container::Workspace { name: "Foo".to_string(),
                                       focused: false });
            assert_eq!(wksp.get_container_type(), ContainerType::Workspace);
        }
        {
            let container = root.new_child(Container::Container {
                layout: Layout::Horizontal, visible: false,
                floating: false, focused: false
            });
            assert_eq!(container.get_container_type(), ContainerType::Container);
        }
    }

    #[test]
    fn get_children() {
        // Create a root with 3 children. The 3rd child has 2 children.
        let mut root = Node::new(Container::Root);
        root.new_child(Container::Root);
        root.new_child(Container::Root);
        {
            let third_child = root.new_child(Container::Root);
            third_child.new_child(Container::Root);
            third_child.new_child(Container::Root);
        }
        let root_children = root.get_children();
        assert!(!root_children.is_empty(), "Root has children");
        assert_eq!(root_children.len(), 3);
        assert!(root_children.last().is_some(), "Root has final child");
        let third_child = root_children.last().expect("Asserted unwrap!");
        assert_eq!(third_child.get_children().len(), 2);
        let third_child_second_child = third_child.get_children().last()
            .expect("Asserted unwrap!");
        assert!(third_child_second_child.get_children().is_empty(),
                "Grandchild doesn't have children");
    }

    #[test]
    fn get_children_mut() {
        // Start out with one child, use get_mut_children to add grandchild.
        let mut root = Node::new(Container::Root);
        root.new_child(Container::Root);

        assert!(root.get_children().last().expect("Root has child")
                .get_children().is_empty(), "Root has no grandchildren");

        {
            let mut children = root.get_children_mut();
            assert_eq!(children.len(), 1);

            let mut child = &mut children[0];
            child.new_child(Container::Root);
        }

        assert_eq!(root.get_children().last().expect("Asserted unwrap!")
                .get_children().len(), 1);
    }

    #[test]
    fn remove_child_at() {
        let mut root = Node::new(Container::Root);
        root.new_child(Container::Root);
        root.new_child(Container::Workspace {
            name: "Foo".to_string(), focused: false
        });
        root.new_child(Container::Root);

        let worksp = root.remove_child_at(1).expect("Index should be valid");
        assert_eq!(worksp.get_container_type(), ContainerType::Workspace);
        assert_eq!(root.get_children().len(), 2)
    }
}
<|endoftext|>


<commit_before>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::read::MemDb;
use super::super::config::Config;
use super::super::sdk::SdkInfo;
use super::super::s3::S3;
use super::super::utils::{ProgressIndicator, copy_with_progress, HumanDuration,
                          IgnorePatterns};
use super::super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
    ignore_patterns: IgnorePatterns,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
    revision: Option<u64>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
    revision: u64,
    offline: bool,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The local filename the SDK has in the stash folder
    pub fn local_filename(&self) -> &str {
        self.filename.trim_right_matches('z')
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, filename: &str) -> Option<&RemoteSdk> {
        self.sdks.get(filename)
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.local_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }

    pub fn sdk_count(&self) -> usize {
        self.sdks.len()
    }
}

impl SyncStatus {

    /// Indicates that the server is running offline (no S3 access)
    pub fn is_offline(&self) -> bool {
        self.offline
    }

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        if self.offline {
            true
        } else {
            let total = self.remote_total as f32;
            let lag = self.lag() as f32;
            lag / total < 0.10
        }
    }

    /// Returns the revision of the stash
    pub fn revision(&self) -> u64 {
        self.revision
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
            ignore_patterns: config.get_ignore_patterns()?.clone(),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv: SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(io::BufReader::new(f))
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.local_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks, revision: None })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.local_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.local_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Returns the current revision
    pub fn get_revision(&self) -> Result<u64> {
        Ok(self.read_local_state()?.revision.unwrap_or(0))
    }

    /// Returns the number of local SDKs
    pub fn sdk_count(&self) -> Result<usize> {
        Ok(self.get_local_state()?.sdk_count())
    }

    /// Returns a list of all synched SDKs
    pub fn list_sdks(&self) -> Result<Vec<SdkInfo>> {
        let local_state = self.get_local_state()?;
        let mut rv : Vec<_> = local_state.sdks().map(|sdk| {
            sdk.info().clone()
        }).collect();
        rv.sort();
        Ok(rv)
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;
        let mut offline = false;

        match self.fetch_remote_state() {
            Ok(remote_state) => {
                for sdk in remote_state.sdks() {
                    if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                        if local_sdk != sdk {
                            different += 1;
                        }
                    } else {
                        missing += 1;
                    }
                    remote_total += 1;
                }
            }
            Err(err) => {
                if let &ErrorKind::S3Unavailable(_) = err.kind() {
                    offline = true;
                } else {
                    return Err(err);
                }
            }
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
            revision: local_state.revision.unwrap_or(0),
            offline: offline,
        })
    }

    /// Checks if the SDK is ignored by config
    pub fn sdk_is_ignored(&self, info: &SdkInfo) -> bool {
        self.ignore_patterns.is_match(&info.sdk_id())
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let mut remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        let mut changed = false;

        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.local_filename().to_string()));

        for sdk in remote_state.sdks() {
            let mut changed_something = false;
            if !self.sdk_is_ignored(&sdk.info()) {
                if let Some(local_sdk) = local_state.get_sdk(sdk.local_filename()) {
                    if local_sdk != sdk {
                        self.update_sdk(&sdk, &options)?;
                        changed_something = true;
                    } else if options.user_facing {
                        println!("  ⸰ Unchanged {}", sdk.info());
                    } else {
                        debug!("unchanged sdk {}", sdk.info());
                    }
                } else {
                    self.update_sdk(&sdk, &options)?;
                    changed_something = true;
                }
            } else {
                if options.user_facing {
                    println!("  ⸰ Ignored {} by config", sdk.info());
                } else {
                    debug!("ignored sdk {} by config", sdk.info());
                }
            }

            to_delete.remove(sdk.local_filename());
            if changed_something {
                changed = true;
                local_state.update_sdk(&sdk);
                local_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
                self.save_local_state(&local_state)?;
            }
        }

        for local_filename in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(local_filename) {
                self.remove_sdk(sdk, &options)?;
                self.memdbs.write().unwrap().remove(&sdk.info());
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else if changed {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        remote_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        // try to fetch it from the local mapping.  The sync method will
        // remove it from here automatically.
        if let Some(arc) = self.memdbs.read().unwrap().get(info) {
            return Ok(arc.clone());
        }

        let local_state = self.get_local_state()?;
        let filename = info.memdb_filename();

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&filename).is_some() {
            let memdb = MemDb::from_path(self.path.join(&filename))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }

    /// Given an SDK info this returns an array of fuzzy matches for it.
    pub fn fuzzy_match_sdk_id(&self, sdk_id: &str) -> Result<Vec<SdkInfo>> {
        let local_state = self.get_local_state()?;
        let mut rv = vec![];

        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            // find all sdks that have a fuzzy match
            for other in local_state.sdks() {
                if let Some(quality) = other.info().get_fuzzy_match(&sdk_info) {
                    rv.push((quality, other.info().clone()));
                }
            }
        }

        rv.sort_by_key(|&(q, _)| q);

        Ok(rv.into_iter().map(|(_, info)| info).collect())
    }
}
<commit_msg>Sync SDKs in order<commit_after>//! Provide access to locally cached memdb SDKs
//!
//! The `MemDbStash` pulls in remote SDKs from an S3 bucket and provides
//! access to it.  This is used by the symbol server to manage the local
//! cache and also to refer to memdb files that are mmap'ed in.
use std::fs;
use std::io;
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::collections::hash_map::Values as HashMapValuesIter;
use std::sync::{Arc, RwLock};

use serde_json;
use xz2::write::XzDecoder;
use chrono::UTC;

use super::read::MemDb;
use super::super::config::Config;
use super::super::sdk::SdkInfo;
use super::super::s3::S3;
use super::super::utils::{ProgressIndicator, copy_with_progress, HumanDuration,
                          IgnorePatterns};
use super::super::{Result, ResultExt, ErrorKind};

/// Helper for synching
pub struct SyncOptions {
    pub user_facing: bool,
}

/// The main memdb stash type
pub struct MemDbStash {
    path: PathBuf,
    s3: S3,
    local_state: RwLock<Option<Arc<SdkSyncState>>>,
    memdbs: RwLock<HashMap<SdkInfo, Arc<MemDb<'static>>>>,
    ignore_patterns: IgnorePatterns,
}

/// Information about a remotely available SDK
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct RemoteSdk {
    filename: String,
    info: SdkInfo,
    size: u64,
    etag: String,
}

#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct SdkSyncState {
    sdks: HashMap<String, RemoteSdk>,
    revision: Option<u64>,
}

/// Information about the health of the stash sync
#[derive(Debug)]
pub struct SyncStatus {
    remote_total: u32,
    missing: u32,
    different: u32,
    revision: u64,
    offline: bool,
}

impl RemoteSdk {
    /// Creates a remote SDK object from some information
    pub fn new(filename: String, info: SdkInfo, etag: String, size: u64) -> RemoteSdk {
        RemoteSdk {
            filename: filename,
            info: info,
            etag: etag,
            size: size,
        }
    }

    /// The remotely visible filename for the SDK
    pub fn filename(&self) -> &str {
        &self.filename
    }

    /// The size of the SDK in bytes
    pub fn size(&self) -> u64 {
        self.size
    }

    /// Returns the SDK info
    pub fn info(&self) -> &SdkInfo {
        &self.info
    }
}

/// Iterator over the SDKs
pub type RemoteSdkIter<'a> = HashMapValuesIter<'a, String, RemoteSdk>;

impl SdkSyncState {
    pub fn get_sdk(&self, info: &SdkInfo) -> Option<&RemoteSdk> {
        self.sdks.get(&info.memdb_filename())
    }

    pub fn update_sdk(&mut self, sdk: &RemoteSdk) {
        self.sdks.insert(sdk.info().memdb_filename().to_string(), sdk.clone());
    }

    pub fn sdks<'a>(&'a self) -> RemoteSdkIter<'a> {
        self.sdks.values()
    }

    pub fn sdk_count(&self) -> usize {
        self.sdks.len()
    }
}

impl SyncStatus {

    /// Indicates that the server is running offline (no S3 access)
    pub fn is_offline(&self) -> bool {
        self.offline
    }

    /// Returns the lag (number of SDKs behind upstream)
    pub fn lag(&self) -> u32 {
        self.missing + self.different
    }

    /// Returns true if the local sync is still considered healthy
    pub fn is_healthy(&self) -> bool {
        if self.offline {
            true
        } else {
            let total = self.remote_total as f32;
            let lag = self.lag() as f32;
            lag / total < 0.10
        }
    }

    /// Returns the revision of the stash
    pub fn revision(&self) -> u64 {
        self.revision
    }
}

impl Default for SyncOptions {
    fn default() -> SyncOptions {
        SyncOptions {
            user_facing: false,
        }
    }
}

impl MemDbStash {
    /// Opens a stash for a given config.
    pub fn new(config: &Config) -> Result<MemDbStash> {
        Ok(MemDbStash {
            path: config.get_symbol_dir()?.to_path_buf(),
            s3: S3::from_config(config)?,
            local_state: RwLock::new(None),
            memdbs: RwLock::new(HashMap::new()),
            ignore_patterns: config.get_ignore_patterns()?.clone(),
        })
    }

    fn get_local_sync_state_filename(&self) -> PathBuf {
        self.path.join("sync.state")
    }

    fn save_state(&self, new_state: &SdkSyncState, filename: &Path) -> Result<()> {
        let mut tmp_filename = filename.to_path_buf();
        tmp_filename.set_extension("tempstate");
        {
            let mut f = fs::File::create(&tmp_filename)?;
            serde_json::to_writer(&mut f, new_state)
                .chain_err(|| "Could not update sync state")?;
        }
        fs::rename(&tmp_filename, &filename)?;
        Ok(())
    }

    fn read_local_state(&self) -> Result<SdkSyncState> {
        let rv: SdkSyncState = match fs::File::open(&self.get_local_sync_state_filename()) {
            Ok(f) => serde_json::from_reader(io::BufReader::new(f))
                .chain_err(|| "Parsing error on loading sync state")?,
            Err(err) => {
                if err.kind() == io::ErrorKind::NotFound {
                    Default::default()
                } else {
                    return Err(err).chain_err(|| "Error loading sync state");
                }
            }
        };
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(rv.clone()));
        Ok(rv)
    }

    fn get_local_state(&self) -> Result<Arc<SdkSyncState>> {
        if let Some(ref arc) = *self.local_state.read().unwrap() {
            return Ok(arc.clone());
        }
        self.read_local_state()?;
        Ok(self.local_state.read().unwrap().as_ref().unwrap().clone())
    }

    fn save_local_state(&self, new_state: &SdkSyncState) -> Result<()> {
        self.save_state(new_state, &self.get_local_sync_state_filename())?;
        let mut opt = self.local_state.write().unwrap();
        *opt = Some(Arc::new(new_state.clone()));
        Ok(())
    }

    fn fetch_remote_state(&self) -> Result<SdkSyncState> {
        let mut sdks = HashMap::new();
        for remote_sdk in self.s3.list_upstream_sdks()? {
            sdks.insert(remote_sdk.info().memdb_filename().into(), remote_sdk);
        }
        Ok(SdkSyncState { sdks: sdks, revision: None })
    }

    fn update_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        // XXX: the progress bar here can stall out because we currently
        // need to buffer the download into memory in the s3 code :(
        let progress = if options.user_facing {
            ProgressIndicator::new(sdk.size() as usize)
        } else {
            info!("updating {}", sdk.info());
            ProgressIndicator::disabled()
        };
        let started = UTC::now();
        progress.set_message(&format!("Synchronizing {}", sdk.info()));
        let mut src = self.s3.download_sdk(sdk)?;
        let dst = fs::File::create(self.path.join(sdk.info().memdb_filename()))?;
        let mut dst = XzDecoder::new(dst);
        copy_with_progress(&progress, &mut src, &mut dst)?;
        progress.finish(&format!("Synchronized {}", sdk.info()));

        let duration = UTC::now() - started;
        if !options.user_facing {
            info!("updated {} in {}", sdk.info(), HumanDuration(duration));
        }
        Ok(())
    }

    fn remove_sdk(&self, sdk: &RemoteSdk, options: &SyncOptions) -> Result<()> {
        if options.user_facing {
            info!("  Deleting {}", sdk.info());
        } else {
            info!("removing {}", sdk.info());
        }
        if let Err(err) = fs::remove_file(self.path.join(sdk.info().memdb_filename())) {
            if err.kind() != io::ErrorKind::NotFound {
                return Err(err.into());
            }
        }
        Ok(())
    }

    /// Returns the current revision
    pub fn get_revision(&self) -> Result<u64> {
        Ok(self.read_local_state()?.revision.unwrap_or(0))
    }

    /// Returns the number of local SDKs
    pub fn sdk_count(&self) -> Result<usize> {
        Ok(self.get_local_state()?.sdk_count())
    }

    /// Returns a list of all synched SDKs
    pub fn list_sdks(&self) -> Result<Vec<SdkInfo>> {
        let local_state = self.get_local_state()?;
        let mut rv : Vec<_> = local_state.sdks().map(|sdk| {
            sdk.info().clone()
        }).collect();
        rv.sort();
        Ok(rv)
    }

    /// Checks the local stash against the server
    pub fn get_sync_status(&self) -> Result<SyncStatus> {
        let local_state = self.read_local_state()?;

        let mut remote_total = 0;
        let mut missing = 0;
        let mut different = 0;
        let mut offline = false;

        match self.fetch_remote_state() {
            Ok(remote_state) => {
                for sdk in remote_state.sdks() {
                    if let Some(local_sdk) = local_state.get_sdk(sdk.info()) {
                        if local_sdk != sdk {
                            different += 1;
                        }
                    } else {
                        missing += 1;
                    }
                    remote_total += 1;
                }
            }
            Err(err) => {
                if let &ErrorKind::S3Unavailable(_) = err.kind() {
                    offline = true;
                } else {
                    return Err(err);
                }
            }
        }

        Ok(SyncStatus {
            remote_total: remote_total as u32,
            missing: missing as u32,
            different: different as u32,
            revision: local_state.revision.unwrap_or(0),
            offline: offline,
        })
    }

    /// Checks if the SDK is ignored by config
    pub fn sdk_is_ignored(&self, info: &SdkInfo) -> bool {
        self.ignore_patterns.is_match(&info.sdk_id())
    }

    /// Synchronize the local stash with the server
    pub fn sync(&self, options: SyncOptions) -> Result<()> {
        let mut local_state = self.read_local_state()?;
        let mut remote_state = self.fetch_remote_state()?;
        let started = UTC::now();
        let mut changed = false;
        let mut to_delete : HashSet<_> = HashSet::from_iter(
            local_state.sdks().map(|x| x.info().clone()));
        let mut sdks : Vec<_> = remote_state.sdks()
            .map(|x| x.info().clone()).collect();
        sdks.sort();

        for sdk_info in sdks.iter() {
            if !self.sdk_is_ignored(sdk_info) {
                let mut changed_something = false;
                let sdk = remote_state.get_sdk(sdk_info).unwrap();
                if let Some(local_sdk) = local_state.get_sdk(sdk_info) {
                    if local_sdk != sdk {
                        self.update_sdk(&sdk, &options)?;
                        changed_something = true;
                    } else if options.user_facing {
                        println!("  ⸰ Unchanged {}", sdk_info);
                    } else {
                        debug!("unchanged sdk {}", sdk_info);
                    }
                } else {
                    self.update_sdk(&sdk, &options)?;
                    changed_something = true;
                }
                if changed_something {
                    changed = true;
                    local_state.update_sdk(&sdk);
                    local_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
                    self.save_local_state(&local_state)?;
                }
            } else {
                if options.user_facing {
                    println!("  ⸰ Ignored {} by config", sdk_info);
                } else {
                    debug!("ignored sdk {} by config", sdk_info);
                }
            }

            to_delete.remove(sdk_info);
        }

        for sdk_info in to_delete.iter() {
            if let Some(sdk) = local_state.get_sdk(sdk_info) {
                self.remove_sdk(sdk, &options)?;
                self.memdbs.write().unwrap().remove(&sdk.info());
            }
        }

        let duration = UTC::now() - started;
        if options.user_facing {
            println!("Sync done in {}", HumanDuration(duration));
        } else if changed {
            info!("finished sync in {}", HumanDuration(duration));
        }

        // if we get this far, the remote state is indeed the local state.
        remote_state.revision = Some(local_state.revision.unwrap_or(0) + 1);
        self.save_local_state(&remote_state)?;

        Ok(())
    }

    /// Looks up an memdb by an SDK info if it's available.
    ///
    /// This returns a memdb wrapped in an arc as internally the system
    /// might try to unload the memdb if no longer needed.  If the MemDb
    /// does not exist, a `UnknownSdk` error is returned.
    pub fn get_memdb(&self, info: &SdkInfo) -> Result<Arc<MemDb<'static>>> {
        // try to fetch it from the local mapping.  The sync method will
        // remove it from here automatically.
        if let Some(arc) = self.memdbs.read().unwrap().get(info) {
            return Ok(arc.clone());
        }

        let local_state = self.get_local_state()?;

        // make sure we check in the local state first if the SDK exists.
        // if we go directly to the memdbs array or look at the file system
        // we might start to consider things that are not available yet or
        // not available any longer.
        if local_state.get_sdk(&info).is_some() {
            let memdb = MemDb::from_path(self.path.join(&info.memdb_filename()))?;
            self.memdbs.write().unwrap().insert(info.clone(), Arc::new(memdb));
            if let Some(arc) = self.memdbs.read().unwrap().get(info) {
                return Ok(arc.clone());
            }
        }

        Err(ErrorKind::UnknownSdk.into())
    }

    /// Looks up an memdb by an SDK info as string if available.
    pub fn get_memdb_from_sdk_id(&self, sdk_id: &str) -> Result<Arc<MemDb<'static>>> {
        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            self.get_memdb(&sdk_info)
        } else {
            Err(ErrorKind::UnknownSdk.into())
        }
    }

    /// Given an SDK info this returns an array of fuzzy matches for it.
    pub fn fuzzy_match_sdk_id(&self, sdk_id: &str) -> Result<Vec<SdkInfo>> {
        let local_state = self.get_local_state()?;
        let mut rv = vec![];

        if let Some(sdk_info) = SdkInfo::from_filename(sdk_id) {
            // find all sdks that have a fuzzy match
            for other in local_state.sdks() {
                if let Some(quality) = other.info().get_fuzzy_match(&sdk_info) {
                    rv.push((quality, other.info().clone()));
                }
            }
        }

        rv.sort_by_key(|&(q, _)| q);

        Ok(rv.into_iter().map(|(_, info)| info).collect())
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use chrono::{DateTime, UTC};
use log;

use level::Level;
use errors::{Result, ErrorKind};
use util;

pub use self::chunked_message::{ChunkedMessage, ChunkSize};
pub use self::compression::MessageCompression;
pub use self::wire_message::WireMessage;

mod chunked_message;
mod compression;
mod wire_message;

/// Message is thre representation of a GELF message.
///
/// `Message` provides a fluid setter and getter  interface to all of GELF's
/// features. Only the `host`-field is not available. It is managed the
/// `Logger`.
///
/// A `Message` can also be constructed from a `log::LogRecord`. All
/// available metadata is transferred over to the message object.
pub struct Message<'a> {
    short_message: String,
    full_message: Option<String>,
    timestamp: Option<DateTime<UTC>>,
    level: Level,
    metadata: HashMap<&'a str, String>,
}

impl<'a> Message<'a> {
    /// Construct a new log message.
    ///
    /// All fields will use their defaults. This means usually Option::None.
    /// A notable exception is `level`. The GELF spec requires this field to
    /// default to Level::Alert.
    pub fn new(short_message: String) -> Self {
        Self::new_with_level(short_message, Level::Alert)
    }

    /// Construct a new log message with a defined level
    ///
    /// All fields will use their defaults. This means usually Option::None.
    pub fn new_with_level(short_message: String, level: Level) -> Self {
        Message {
            short_message: short_message,
            level: level,
            full_message: None,
            timestamp: None,
            metadata: HashMap::new(),
        }
    }

    /// Return the `short_message`
    pub fn short_message(&self) -> &String {
        &self.short_message
    }

    /// Set the `short_message`
    pub fn set_short_message(&mut self, msg: String) -> &mut Self {
        self.short_message = msg;
        self
    }

    /// Return the `full_message`
    pub fn full_message(&self) -> &Option<String> {
        &self.full_message
    }

    /// Set the `full_message`
    pub fn set_full_message(&mut self, msg: String) -> &mut Self {
        self.full_message = Some(msg);
        self
    }

    // Clear the `full_message`
    pub fn clear_full_message(&mut self) -> &mut Self {
        self.full_message = None;
        self
    }

    /// Return the `timestamp`
    pub fn timestamp(&self) -> &Option<DateTime<UTC>> {
        &self.timestamp
    }

    /// Set the `timestamp`
    pub fn set_timestamp(&mut self, ts: DateTime<UTC>) -> &mut Self {
        self.timestamp = Some(ts);
        self
    }

    /// Clear the `timestamp`
    pub fn clear_timestamp(&mut self) -> &mut Self {
        self.timestamp = None;
        self
    }

    /// Return the `level`
    pub fn level(&self) -> Level {
        self.level
    }

    /// Set the `level`
    pub fn set_level(&mut self, level: Level) -> &mut Self {
        self.level = level;
        self
    }

    /// Return a metadata field with given key
    pub fn metadata(&self, key: &str) -> Option<&String> {
        self.metadata.get(key)
    }

    /// Return all metadata
    pub fn all_metadata(&self) -> &HashMap<&'a str, String> {
        &self.metadata
    }

    /// Set a metadata field with given key to value
    pub fn set_metadata(&mut self, key: &'a str, value: String) -> Result<&mut Self> {
        if key == "id" {
            bail!(ErrorKind::IllegalNameForAdditional(String::from(key)));
        }

        self.metadata.insert(key, value);

        Ok(self)
    }
}

impl<'a> From<&'a log::LogRecord<'a>> for Message<'a> {
    /// Create a `Message` from given `log::LogRecord` including all metadata
    fn from(record: &'a log::LogRecord) -> Message<'a> {
        // Create message with given text and level
        let mut msg = Message::new(format!("{}", record.args()));

        msg.set_timestamp(UTC::now());

        // Add location meta-data
        msg.metadata.insert("file", record.location().file().to_owned());
        msg.metadata.insert("line", record.location().line().to_string());
        msg.metadata.insert("module_path", record.location().module_path().to_owned());

        // Add runtime meta-data
        msg.metadata.insert("process_id", util::pid().to_string());

        msg
    }
}
<commit_msg>support loglevel<commit_after>use std::collections::HashMap;
use chrono::{DateTime, UTC};
use log;

use level::Level;
use errors::{Result, ErrorKind};
use util;

pub use self::chunked_message::{ChunkedMessage, ChunkSize};
pub use self::compression::MessageCompression;
pub use self::wire_message::WireMessage;

mod chunked_message;
mod compression;
mod wire_message;

/// Message is thre representation of a GELF message.
///
/// `Message` provides a fluid setter and getter  interface to all of GELF's
/// features. Only the `host`-field is not available. It is managed the
/// `Logger`.
///
/// A `Message` can also be constructed from a `log::LogRecord`. All
/// available metadata is transferred over to the message object.
pub struct Message<'a> {
    short_message: String,
    full_message: Option<String>,
    timestamp: Option<DateTime<UTC>>,
    level: Level,
    metadata: HashMap<&'a str, String>,
}

impl<'a> Message<'a> {
    /// Construct a new log message.
    ///
    /// All fields will use their defaults. This means usually Option::None.
    /// A notable exception is `level`. The GELF spec requires this field to
    /// default to Level::Alert.
    pub fn new(short_message: String) -> Self {
        Self::new_with_level(short_message, Level::Alert)
    }

    /// Construct a new log message with a defined level
    ///
    /// All fields will use their defaults. This means usually Option::None.
    pub fn new_with_level(short_message: String, level: Level) -> Self {
        Message {
            short_message: short_message,
            level: level,
            full_message: None,
            timestamp: None,
            metadata: HashMap::new(),
        }
    }

    /// Return the `short_message`
    pub fn short_message(&self) -> &String {
        &self.short_message
    }

    /// Set the `short_message`
    pub fn set_short_message(&mut self, msg: String) -> &mut Self {
        self.short_message = msg;
        self
    }

    /// Return the `full_message`
    pub fn full_message(&self) -> &Option<String> {
        &self.full_message
    }

    /// Set the `full_message`
    pub fn set_full_message(&mut self, msg: String) -> &mut Self {
        self.full_message = Some(msg);
        self
    }

    // Clear the `full_message`
    pub fn clear_full_message(&mut self) -> &mut Self {
        self.full_message = None;
        self
    }

    /// Return the `timestamp`
    pub fn timestamp(&self) -> &Option<DateTime<UTC>> {
        &self.timestamp
    }

    /// Set the `timestamp`
    pub fn set_timestamp(&mut self, ts: DateTime<UTC>) -> &mut Self {
        self.timestamp = Some(ts);
        self
    }

    /// Clear the `timestamp`
    pub fn clear_timestamp(&mut self) -> &mut Self {
        self.timestamp = None;
        self
    }

    /// Return the `level`
    pub fn level(&self) -> Level {
        self.level
    }

    /// Set the `level`
    pub fn set_level(&mut self, level: Level) -> &mut Self {
        self.level = level;
        self
    }

    /// Return a metadata field with given key
    pub fn metadata(&self, key: &str) -> Option<&String> {
        self.metadata.get(key)
    }

    /// Return all metadata
    pub fn all_metadata(&self) -> &HashMap<&'a str, String> {
        &self.metadata
    }

    /// Set a metadata field with given key to value
    pub fn set_metadata(&mut self, key: &'a str, value: String) -> Result<&mut Self> {
        if key == "id" {
            bail!(ErrorKind::IllegalNameForAdditional(String::from(key)));
        }

        self.metadata.insert(key, value);

        Ok(self)
    }
}

impl<'a> From<&'a log::LogRecord<'a>> for Message<'a> {
    /// Create a `Message` from given `log::LogRecord` including all metadata
    fn from(record: &'a log::LogRecord) -> Message<'a> {
        // Create message with given text and level
        let mut msg = Message::new_with_level(format!("{}", record.args()), record.level().into());

        msg.set_timestamp(UTC::now());

        // Add location meta-data
        msg.metadata.insert("file", record.location().file().to_owned());
        msg.metadata.insert("line", record.location().line().to_string());
        msg.metadata.insert("module_path", record.location().module_path().to_owned());

        // Add runtime meta-data
        msg.metadata.insert("process_id", util::pid().to_string());

        msg
    }
}
<|endoftext|>