<commit_before>// Copyright (c) 2015, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

#![allow(bad_style)]
#![unstable]

extern crate netmap_sys;
extern crate libc;

use self::netmap_sys::netmap_user::{NETMAP_BUF, NETMAP_FD, NETMAP_TXRING, nm_close, nm_desc,
                                    nm_nextpkt, nm_open, nm_pkthdr, nm_ring_next};
use self::netmap_sys::netmap::{netmap_slot, nm_ring_empty};

use std::ffi::CString;
use std::path::Path;
use std::fs::File;
use std::io;
use std::io::Read;
use std::mem;
use std::num;
use std::ptr;
use std::raw;
use std::sync::Arc;

use datalink::{DataLinkChannelIterator, DataLinkChannelType, DataLinkReceiver, DataLinkSender};
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::NetworkInterface;

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
#[repr(C)]
struct pollfd {
    fd: libc::c_int,
    events: libc::c_short,
    revents: libc::c_short,
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
const POLLIN: libc::c_short = 0x0001;
#[cfg(any(target_os = "linux", target_os = "freebsd"))]
const POLLOUT: libc::c_short = 0x0004;

#[cfg(target_os = "freebsd")]
type nfds_t = libc::c_uint;
#[cfg(target_os = "linux")]
type nfds_t = libc::c_ulong;

extern {
    fn poll(fds: *mut pollfd, nfds: nfds_t, timeout: libc::c_int) -> libc::c_int;
}

struct NmDesc {
    desc: *mut nm_desc,
    buf_size: libc::c_uint,
}

impl NmDesc {
    fn new(iface: &NetworkInterface) -> io::Result<NmDesc> {
        let ifname = CString::new(("netmap:".to_string() + &iface.name[..]).as_bytes());
        let desc = unsafe { nm_open(ifname.unwrap().as_ptr(), ptr::null(), 0, ptr::null()) };

        if desc.is_null() {
            Err(io::Error::last_os_error())
        } else {
            let mut f = try!(File::open(&Path::new("/sys/module/netmap/parameters/buf_size")));
            let mut num_str = String::new();
            try!(f.read_to_string(&mut num_str));
            let buf_size = num_str.trim_right().parse().unwrap();

            Ok(NmDesc {
                desc: desc,
                buf_size: buf_size,
            })
        }
    }
}

impl Drop for NmDesc {
    fn drop(&mut self) {
        unsafe {
            nm_close(self.desc);
        }
    }
}

#[inline]
pub fn datalink_channel(network_interface: &NetworkInterface,
                        _write_buffer_size: usize,
                        _read_buffer_size: usize,
                        _channel_type: DataLinkChannelType)
    -> io::Result<(Box<DataLinkSenderImpl>, Box<DataLinkReceiverImpl>)> {
    // FIXME probably want one for each of send/recv
    let desc = NmDesc::new(network_interface);
    match desc {
        Ok(desc) => {
            let arc = Arc::new(desc);

            Ok((Box::new(DataLinkSenderImpl { desc: arc.clone() }),
                Box::new(DataLinkReceiverImpl { desc: arc })))
        }
        Err(e) => Err(e),
    }
}

pub struct DataLinkSenderImpl {
    desc: Arc<NmDesc>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        assert!(num::cast::<usize, u16>(packet_size).unwrap() as libc::c_uint <=
                self.desc.buf_size);
        let desc = self.desc.desc;
        let mut fds = pollfd {
            fd: unsafe { NETMAP_FD(desc) },
            events: POLLOUT,
            revents: 0,
        };
        let mut packet_idx = 0usize;
        while packet_idx < num_packets {
            unsafe {
                if poll(&mut fds, 1, -1) < 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
                let ring = NETMAP_TXRING((*desc).nifp, 0);
                while !nm_ring_empty(ring) && packet_idx < num_packets {
                    let i = (*ring).cur;
                    let slot_ptr: *mut netmap_slot = mem::transmute(&mut (*ring).slot);
                    let buf = NETMAP_BUF(ring, (*slot_ptr.offset(i as isize)).buf_idx as isize);
                    let slice = raw::Slice {
                        data: buf,
                        len: packet_size,
                    };
                    let meh = MutableEthernetPacket::new(mem::transmute(slice));
                    (*slot_ptr.offset(i as isize)).len = packet_size as u16;
                    func(meh);
                    let next = nm_ring_next(ring, i);
                    (*ring).head = next;
                    (*ring).cur = next;
                    packet_idx += 1;
                }
            }
        }

        Some(Ok(()))
    }

    #[inline]
    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        use packet::MutablePacket;
        self.build_and_send(1,
                            packet.packet().len(),
                            &mut |mut eh: MutableEthernetPacket| {
                                eh.clone_from(packet);
                            })
    }
}

pub struct DataLinkReceiverImpl {
    desc: Arc<NmDesc>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    // FIXME Layer 3
    fn iter<'a>(&'a mut self) -> Box<DataLinkChannelIterator + 'a> {
        Box::new(DataLinkChannelIteratorImpl { pc: self })
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
}

impl<'a> DataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a> {
    fn next<'c>(&'c mut self) -> io::Result<EthernetPacket<'c>> {
        let desc = self.pc.desc.desc;
        let mut h: nm_pkthdr = unsafe { mem::uninitialized() };
        let mut buf = unsafe { nm_nextpkt(desc, &mut h) };
        if buf.is_null() {
            let mut fds = pollfd {
                fd: unsafe { NETMAP_FD(desc) },
                events: POLLIN,
                revents: 0,
            };
            if unsafe { poll(&mut fds, 1, -1) } < 0 {
                return Err(io::Error::last_os_error());
            }
            buf = unsafe { nm_nextpkt(desc, &mut h) };
        }
        Ok(EthernetPacket::new(unsafe {
            mem::transmute(raw::Slice {
                data: buf,
                len: h.len as usize,
            })
        }))
    }
}
<commit_msg>Fix netmap build.<commit_after>// Copyright (c) 2015, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

#![allow(bad_style)]

extern crate netmap_sys;
extern crate libc;

use self::netmap_sys::netmap_user::{NETMAP_BUF, NETMAP_FD, NETMAP_TXRING, nm_close, nm_desc,
                                    nm_nextpkt, nm_open, nm_pkthdr, nm_ring_next};
use self::netmap_sys::netmap::{netmap_slot, nm_ring_empty};

use std::ffi::CString;
use std::path::Path;
use std::fs::File;
use std::io;
use std::io::Read;
use std::mem;
use std::num;
use std::ptr;
use std::raw;
use std::sync::Arc;

use datalink::{DataLinkChannelIterator, DataLinkChannelType, DataLinkReceiver, DataLinkSender};
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::NetworkInterface;

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
#[repr(C)]
struct pollfd {
    fd: libc::c_int,
    events: libc::c_short,
    revents: libc::c_short,
}

#[cfg(any(target_os = "linux", target_os = "freebsd"))]
const POLLIN: libc::c_short = 0x0001;
#[cfg(any(target_os = "linux", target_os = "freebsd"))]
const POLLOUT: libc::c_short = 0x0004;

#[cfg(target_os = "freebsd")]
type nfds_t = libc::c_uint;
#[cfg(target_os = "linux")]
type nfds_t = libc::c_ulong;

extern {
    fn poll(fds: *mut pollfd, nfds: nfds_t, timeout: libc::c_int) -> libc::c_int;
}

struct NmDesc {
    desc: *mut nm_desc,
    buf_size: libc::c_uint,
}

unsafe impl Send for NmDesc {}
unsafe impl Sync for NmDesc {}

impl NmDesc {
    fn new(iface: &NetworkInterface) -> io::Result<NmDesc> {
        let ifname = CString::new(("netmap:".to_string() + &iface.name[..]).as_bytes());
        let desc = unsafe { nm_open(ifname.unwrap().as_ptr(), ptr::null(), 0, ptr::null()) };

        if desc.is_null() {
            Err(io::Error::last_os_error())
        } else {
            let mut f = try!(File::open(&Path::new("/sys/module/netmap/parameters/buf_size")));
            let mut num_str = String::new();
            try!(f.read_to_string(&mut num_str));
            let buf_size = num_str.trim_right().parse().unwrap();

            Ok(NmDesc {
                desc: desc,
                buf_size: buf_size,
            })
        }
    }
}

impl Drop for NmDesc {
    fn drop(&mut self) {
        unsafe {
            nm_close(self.desc);
        }
    }
}

#[inline]
pub fn datalink_channel(network_interface: &NetworkInterface,
                        _write_buffer_size: usize,
                        _read_buffer_size: usize,
                        _channel_type: DataLinkChannelType)
    -> io::Result<(Box<DataLinkSender>, Box<DataLinkReceiver>)> {
    // FIXME probably want one for each of send/recv
    let desc = NmDesc::new(network_interface);
    match desc {
        Ok(desc) => {
            let arc = Arc::new(desc);

            Ok((Box::new(DataLinkSenderImpl { desc: arc.clone() }),
                Box::new(DataLinkReceiverImpl { desc: arc })))
        }
        Err(e) => Err(e),
    }
}

pub struct DataLinkSenderImpl {
    desc: Arc<NmDesc>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        assert!(packet_size <= self.desc.buf_size as usize);
        let desc = self.desc.desc;
        let mut fds = pollfd {
            fd: unsafe { NETMAP_FD(desc) },
            events: POLLOUT,
            revents: 0,
        };
        let mut packet_idx = 0usize;
        while packet_idx < num_packets {
            unsafe {
                if poll(&mut fds, 1, -1) < 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
                let ring = NETMAP_TXRING((*desc).nifp, 0);
                while !nm_ring_empty(ring) && packet_idx < num_packets {
                    let i = (*ring).cur;
                    let slot_ptr: *mut netmap_slot = mem::transmute(&mut (*ring).slot);
                    let buf = NETMAP_BUF(ring, (*slot_ptr.offset(i as isize)).buf_idx as isize);
                    let slice = raw::Slice {
                        data: buf,
                        len: packet_size,
                    };
                    let meh = MutableEthernetPacket::new(mem::transmute(slice)).unwrap();
                    (*slot_ptr.offset(i as isize)).len = packet_size as u16;
                    func(meh);
                    let next = nm_ring_next(ring, i);
                    (*ring).head = next;
                    (*ring).cur = next;
                    packet_idx += 1;
                }
            }
        }

        Some(Ok(()))
    }

    #[inline]
    fn send_to(&mut self,
               packet: &EthernetPacket,
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        use packet::MutablePacket;
        self.build_and_send(1,
                            packet.packet().len(),
                            &mut |mut eh: MutableEthernetPacket| {
                                eh.clone_from(packet);
                            })
    }
}

pub struct DataLinkReceiverImpl {
    desc: Arc<NmDesc>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    // FIXME Layer 3
    fn iter<'a>(&'a mut self) -> Box<DataLinkChannelIterator + 'a> {
        Box::new(DataLinkChannelIteratorImpl { pc: self })
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
}

impl<'a> DataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a> {
    fn next<'c>(&'c mut self) -> io::Result<EthernetPacket<'c>> {
        let desc = self.pc.desc.desc;
        let mut h: nm_pkthdr = unsafe { mem::uninitialized() };
        let mut buf = unsafe { nm_nextpkt(desc, &mut h) };
        if buf.is_null() {
            let mut fds = pollfd {
                fd: unsafe { NETMAP_FD(desc) },
                events: POLLIN,
                revents: 0,
            };
            if unsafe { poll(&mut fds, 1, -1) } < 0 {
                return Err(io::Error::last_os_error());
            }
            buf = unsafe { nm_nextpkt(desc, &mut h) };
        }
        Ok(EthernetPacket::new(unsafe {
            mem::transmute(raw::Slice {
                data: buf,
                len: h.len as usize,
            })
        }).unwrap())
    }
}
<|endoftext|>


<commit_before>extern crate crossbeam;
use std;
use decoders::NUM_CORES;

#[derive(Debug, Copy, Clone)]
pub struct Endian {
  big: bool,
}

impl Endian {
  pub fn ru32(&self, buf: &[u8], pos: usize) -> u32 {
    if self.big {
      BEu32(buf,pos)
    } else {
      LEu32(buf,pos)
    }
  }

  pub fn ru16(&self, buf: &[u8], pos: usize) -> u16 {
    if self.big {
      BEu16(buf,pos)
    } else {
      LEu16(buf,pos)
    }
  }
}

pub static BIG_ENDIAN: Endian = Endian{big: true};
pub static LITTLE_ENDIAN: Endian = Endian{big: false};

#[allow(non_snake_case)] pub fn BEu32(buf: &[u8], pos: usize) -> u32 {
  (buf[pos] as u32) << 24 |
  (buf[pos+1] as u32) << 16 |
  (buf[pos+2] as u32) << 8 |
  (buf[pos+3] as u32)
}

#[allow(non_snake_case)] pub fn LEu32(buf: &[u8], pos: usize) -> u32 {
  (buf[pos] as u32) |
  (buf[pos+1] as u32) << 8 |
  (buf[pos+2] as u32) << 16 |
  (buf[pos+3] as u32) << 24
}

#[allow(non_snake_case)] pub fn BEu16(buf: &[u8], pos: usize) -> u16 {
  (buf[pos] as u16) << 8 | (buf[pos+1] as u16)
}

#[allow(non_snake_case)] pub fn LEu16(buf: &[u8], pos: usize) -> u16 {
  (buf[pos] as u16) | (buf[pos+1] as u16) << 8
}

pub fn decode_12be(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*12/8) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(3)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;
      let g3: u16 = i[2] as u16;

      o[0] = (g1 << 4) | (g2 >> 4);
      o[1] = ((g2 & 0x0f) << 8) | g3;
    }
  }))
}

pub fn decode_12le(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*12/8) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(3)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;
      let g3: u16 = i[2] as u16;

      o[0] = ((g2 & 0x0f) << 8) | g1;
      o[1] = (g3 << 4) | (g2 >> 4);
    }
  }))
}

pub fn decode_12be_unpacked(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*2) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(2)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;

      o[0] = ((g1 & 0x0f) << 8) | g2;
    }
  }))
}

pub fn decode_16le(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*2) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(2)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;

      o[0] = (g2 << 8) | g1;
    }
  }))
}

pub fn decode_threaded<F>(width: usize, height: usize, closure: &F) -> Vec<u16>
  where F : Fn(&mut [u16], usize, usize, usize)+std::marker::Sync {
  let mut out: Vec<u16> = vec![0; (width*height) as usize];

  // Default to the number of available physical cores (no hyperthreading)
  let threads = *NUM_CORES;

  // If we're only using one thread do it all sequentially and be done with it
  if threads < 2 || height < threads {
    closure(&mut out, 0, width, height);
    return out
  }

  let mut split = height/threads;
  if split*threads < height { // Make sure the last split is always the smallest
    split += 1;
  }

  crossbeam::scope(|scope| {
    let mut handles = Vec::new();
    for (i,out_part) in (&mut out[..]).chunks_mut(split*width).enumerate() {
      let start = split*(i);
      let tall = out_part.len()/width;
      let handle = scope.spawn(move || {
        closure(out_part, start, width, tall);
      });
      handles.push(handle);
    }

    for h in handles { h.join() };
  });

  out
}

#[derive(Debug, Clone)]
pub struct LookupTable {
  table: Vec<(u16, u16, u16)>,
}

impl LookupTable {
  pub fn new(table: &[u16]) -> LookupTable {
    let mut tbl = vec![(0,0,0); table.len()];
    for i in 0..table.len() {
      let center = table[i];
      let lower = if i > 0 {table[i-1]} else {center};
      let upper = if i < (table.len()-1) {table[i+1]} else {center};
      let base = center - ((upper - lower + 2) / 4);
      let delta = upper - lower;
      tbl[i] = (center, base, delta);
    }
    LookupTable {
      table: tbl,
    }
  }

//  pub fn lookup(&self, value: u16) -> u16 {
//    let (val, _, _) = self.table[value as usize];
//    val
//  }

  pub fn dither(&self, value: u16, rand: &mut u32) -> u16 {
    let (_, sbase, sdelta) = self.table[value as usize];
    let base = sbase as u32;
    let delta = sdelta as u32;
    let pixel = base + ((delta * (*rand & 2047) + 1024) >> 12);
    *rand = 15700 * (*rand & 65535) + (*rand >> 16);
    pixel as u16
  }
}

pub struct BitPump<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPump<'a> {
  pub fn new(src: &'a [u8]) -> BitPump {
    BitPump {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }

  fn fill_bits(&mut self) {
    let inbits: u64 = LEu32(self.buffer, self.pos) as u64;
    self.pos += 4;
    self.bits = ((inbits << 32) | (self.bits << (32-self.nbits))) >> (32-self.nbits);
    self.nbits += 32;
  }

  pub fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      self.fill_bits();
    }
    (self.bits & (0x0ffffffffu64 >> (32-num))) as u32
  }

  pub fn get_bits(&mut self, num: u32) -> u32 {
    let val = self.peek_bits(num);
    self.bits = self.bits >> num;
    self.nbits -= num;
    val
  }
}
<commit_msg>Fix broken 16bit decoders<commit_after>extern crate crossbeam;
use std;
use decoders::NUM_CORES;

#[derive(Debug, Copy, Clone)]
pub struct Endian {
  big: bool,
}

impl Endian {
  pub fn ru32(&self, buf: &[u8], pos: usize) -> u32 {
    if self.big {
      BEu32(buf,pos)
    } else {
      LEu32(buf,pos)
    }
  }

  pub fn ru16(&self, buf: &[u8], pos: usize) -> u16 {
    if self.big {
      BEu16(buf,pos)
    } else {
      LEu16(buf,pos)
    }
  }
}

pub static BIG_ENDIAN: Endian = Endian{big: true};
pub static LITTLE_ENDIAN: Endian = Endian{big: false};

#[allow(non_snake_case)] pub fn BEu32(buf: &[u8], pos: usize) -> u32 {
  (buf[pos] as u32) << 24 |
  (buf[pos+1] as u32) << 16 |
  (buf[pos+2] as u32) << 8 |
  (buf[pos+3] as u32)
}

#[allow(non_snake_case)] pub fn LEu32(buf: &[u8], pos: usize) -> u32 {
  (buf[pos] as u32) |
  (buf[pos+1] as u32) << 8 |
  (buf[pos+2] as u32) << 16 |
  (buf[pos+3] as u32) << 24
}

#[allow(non_snake_case)] pub fn BEu16(buf: &[u8], pos: usize) -> u16 {
  (buf[pos] as u16) << 8 | (buf[pos+1] as u16)
}

#[allow(non_snake_case)] pub fn LEu16(buf: &[u8], pos: usize) -> u16 {
  (buf[pos] as u16) | (buf[pos+1] as u16) << 8
}

pub fn decode_12be(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*12/8) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(3)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;
      let g3: u16 = i[2] as u16;

      o[0] = (g1 << 4) | (g2 >> 4);
      o[1] = ((g2 & 0x0f) << 8) | g3;
    }
  }))
}

pub fn decode_12le(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*12/8) as usize)..];

    for (o, i) in out.chunks_mut(2).zip(inb.chunks(3)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;
      let g3: u16 = i[2] as u16;

      o[0] = ((g2 & 0x0f) << 8) | g1;
      o[1] = (g3 << 4) | (g2 >> 4);
    }
  }))
}

pub fn decode_12be_unpacked(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*2) as usize)..];

    for (o, i) in out.chunks_mut(1).zip(inb.chunks(2)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;

      o[0] = ((g1 & 0x0f) << 8) | g2;
    }
  }))
}

pub fn decode_16le(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
  decode_threaded(width, height, &(|out: &mut [u16], start, width, _| {
    let inb = &buf[((start*width*2) as usize)..];

    for (o, i) in out.chunks_mut(1).zip(inb.chunks(2)) {
      let g1: u16 = i[0] as u16;
      let g2: u16 = i[1] as u16;

      o[0] = (g2 << 8) | g1;
    }
  }))
}

pub fn decode_threaded<F>(width: usize, height: usize, closure: &F) -> Vec<u16>
  where F : Fn(&mut [u16], usize, usize, usize)+std::marker::Sync {
  let mut out: Vec<u16> = vec![0; (width*height) as usize];

  // Default to the number of available physical cores (no hyperthreading)
  let threads = *NUM_CORES;

  // If we're only using one thread do it all sequentially and be done with it
  if threads < 2 || height < threads {
    closure(&mut out, 0, width, height);
    return out
  }

  let mut split = height/threads;
  if split*threads < height { // Make sure the last split is always the smallest
    split += 1;
  }

  crossbeam::scope(|scope| {
    let mut handles = Vec::new();
    for (i,out_part) in (&mut out[..]).chunks_mut(split*width).enumerate() {
      let start = split*(i);
      let tall = out_part.len()/width;
      let handle = scope.spawn(move || {
        closure(out_part, start, width, tall);
      });
      handles.push(handle);
    }

    for h in handles { h.join() };
  });

  out
}

#[derive(Debug, Clone)]
pub struct LookupTable {
  table: Vec<(u16, u16, u16)>,
}

impl LookupTable {
  pub fn new(table: &[u16]) -> LookupTable {
    let mut tbl = vec![(0,0,0); table.len()];
    for i in 0..table.len() {
      let center = table[i];
      let lower = if i > 0 {table[i-1]} else {center};
      let upper = if i < (table.len()-1) {table[i+1]} else {center};
      let base = center - ((upper - lower + 2) / 4);
      let delta = upper - lower;
      tbl[i] = (center, base, delta);
    }
    LookupTable {
      table: tbl,
    }
  }

//  pub fn lookup(&self, value: u16) -> u16 {
//    let (val, _, _) = self.table[value as usize];
//    val
//  }

  pub fn dither(&self, value: u16, rand: &mut u32) -> u16 {
    let (_, sbase, sdelta) = self.table[value as usize];
    let base = sbase as u32;
    let delta = sdelta as u32;
    let pixel = base + ((delta * (*rand & 2047) + 1024) >> 12);
    *rand = 15700 * (*rand & 65535) + (*rand >> 16);
    pixel as u16
  }
}

pub struct BitPump<'a> {
  buffer: &'a [u8],
  pos: usize,
  bits: u64,
  nbits: u32,
}

impl<'a> BitPump<'a> {
  pub fn new(src: &'a [u8]) -> BitPump {
    BitPump {
      buffer: src,
      pos: 0,
      bits: 0,
      nbits: 0,
    }
  }

  fn fill_bits(&mut self) {
    let inbits: u64 = LEu32(self.buffer, self.pos) as u64;
    self.pos += 4;
    self.bits = ((inbits << 32) | (self.bits << (32-self.nbits))) >> (32-self.nbits);
    self.nbits += 32;
  }

  pub fn peek_bits(&mut self, num: u32) -> u32 {
    if num > self.nbits {
      self.fill_bits();
    }
    (self.bits & (0x0ffffffffu64 >> (32-num))) as u32
  }

  pub fn get_bits(&mut self, num: u32) -> u32 {
    let val = self.peek_bits(num);
    self.bits = self.bits >> num;
    self.nbits -= num;
    val
  }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hash_map::Entry;

pub fn count(nucleotide: char, input: &str) -> usize {
    input.chars().filter(|&c| c == nucleotide).count()
}

pub fn nucleotide_counts(input: &str) -> HashMap<char, usize> {
    let mut map: HashMap<char, usize> = "ACGT".chars().map(|c| (c, 0)).collect();
    for nucleotide in input.chars() {
        match map.entry(nucleotide) {
            Entry::Vacant(view) => view.insert(1),
            Entry::Occupied(mut view) => {
                *view.get_mut() += 1;
                view.into_mut()
            }
        };
    }
    map
}
<commit_msg>nucleotide-count: Use `Entry#or_insert`<commit_after>use std::collections::HashMap;

pub fn count(nucleotide: char, input: &str) -> usize {
    input.chars().filter(|&c| c == nucleotide).count()
}

pub fn nucleotide_counts(input: &str) -> HashMap<char, usize> {
    let mut map: HashMap<char, usize> = "ACGT".chars().map(|c| (c, 0)).collect();
    for nucleotide in input.chars() {
        *map.entry(nucleotide).or_insert(0) += 1;
    }
    map
}
<|endoftext|>


<commit_before>#![feature(async_await)]
#![deny(warnings)]
extern crate hyper;
extern crate pretty_env_logger;
extern crate serde_json;

use hyper::{Body, Chunk, Client, Method, Request, Response, Server, StatusCode, header};
use hyper::client::HttpConnector;
use hyper::service::{service_fn, make_service_fn};
use futures_util::{TryStreamExt};

type GenericError = Box<dyn std::error::Error + Send + Sync>;

static NOTFOUND: &[u8] = b"Not Found";
static URL: &str = "http://127.0.0.1:1337/json_api";
static INDEX: &[u8] = b"test.html";
static POST_DATA: &str = r#"{"original": "data"}"#;

async fn client_request_response(client: &Client<HttpConnector>)
    -> Result<Response, GenericError>
{
     let req = Request::builder()
         .method(Method::POST)
         .uri(URL)
         .header(header::CONTENT_TYPE, "application/json")
         .body(POST_DATA.into())
         .unwrap();

     let web_res = client.request(req).await?;
    // Compare the JSON we sent (before) with what we received (after):
    let body = Body::wrap_stream(web_res.into_body().map_ok(|b| {
        Chunk::from(format!("POST request body: {}
Response: {}",
                            POST_DATA,
                            std::str::from_utf8(&b).unwrap()))
    }));

    Ok(Response::new(body))
}

async fn api_post_response(req: Request)
    -> Result<Response, GenericError>
{
    // A web api to run against
    let entire_body = req.into_body().try_concat().await?;
    // TODO: Replace all unwraps with proper error handling
    let str = String::from_utf8(entire_body.to_vec())?;
    let mut data : serde_json::Value = serde_json::from_str(&str)?;
    data["test"] = serde_json::Value::from("test_value");
    let json = serde_json::to_string(&data)?;
    let response = Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/json")
        .body(Body::from(json))?;
    Ok(response)
}

async fn api_get_response() -> Result<Response, GenericError> {
    let data = vec!["foo", "bar"];
    let res = match serde_json::to_string(&data) {
        Ok(json) => {
            Response::builder()
                .header(header::CONTENT_TYPE, "application/json")
                .body(Body::from(json))
                .unwrap()
        }
        Err(_) => {
            Response::builder()
                .status(StatusCode::INTERNAL_SERVER_ERROR)
                .body(Body::from("Internal Server Error"))
                .unwrap()
        }
    };
    Ok(res)
}

async fn response_examples(req: Request, client: &Client<HttpConnector>)
    -> Result<Response, GenericError>
{
    match (req.method(), req.uri().path()) {
        (&Method::GET, "/") | (&Method::GET, "/index.html") => {
            let body = Body::from(INDEX);
            Ok(Response::new(body))
        },
        (&Method::GET, "/test.html") => {
           client_request_response(client).await
        },
        (&Method::POST, "/json_api") => {
            api_post_response(req).await
        },
        (&Method::GET, "/json_api") => {
            api_get_response().await
        }
        _ => {
            // Return 404 not found response.
            let body = Body::from(NOTFOUND);
            Ok(Response::builder()
                        .status(StatusCode::NOT_FOUND)
                        .body(body)
                        .unwrap())
        }
    }
}

#[hyper::rt::main]
async fn main() -> Result<(), GenericError> {
    pretty_env_logger::init();

    let addr = "127.0.0.1:1337".parse().unwrap();

    // Share a `Client` with all `Service`s
    let client = Client::new();

    let new_service = make_service_fn(move |_| {
        // Move a clone of `client` into the `service_fn`.
        let client = client.clone();
        async {
            Ok::<_, GenericError>(service_fn(move |req| {
                response_examples(req, &client)
            }))
        }
    });

    let server = Server::bind(&addr)
        .serve(new_service);

    println!("Listening on http://{}", addr);

    server.await?;

    Ok(())
}
<commit_msg>docs(examples): fix web_api shared client lifetime issue<commit_after>#![feature(async_await)]
#![deny(warnings)]
extern crate hyper;
extern crate pretty_env_logger;
extern crate serde_json;

use hyper::{Body, Chunk, Client, Method, Request, Response, Server, StatusCode, header};
use hyper::client::HttpConnector;
use hyper::service::{service_fn, make_service_fn};
use futures_util::{TryStreamExt};

type GenericError = Box<dyn std::error::Error + Send + Sync>;

static NOTFOUND: &[u8] = b"Not Found";
static URL: &str = "http://127.0.0.1:1337/json_api";
static INDEX: &[u8] = b"test.html";
static POST_DATA: &str = r#"{"original": "data"}"#;

async fn client_request_response(client: &Client<HttpConnector>)
    -> Result<Response, GenericError>
{
     let req = Request::builder()
         .method(Method::POST)
         .uri(URL)
         .header(header::CONTENT_TYPE, "application/json")
         .body(POST_DATA.into())
         .unwrap();

     let web_res = client.request(req).await?;
    // Compare the JSON we sent (before) with what we received (after):
    let body = Body::wrap_stream(web_res.into_body().map_ok(|b| {
        Chunk::from(format!("POST request body: {}
Response: {}",
                            POST_DATA,
                            std::str::from_utf8(&b).unwrap()))
    }));

    Ok(Response::new(body))
}

async fn api_post_response(req: Request)
    -> Result<Response, GenericError>
{
    // A web api to run against
    let entire_body = req.into_body().try_concat().await?;
    // TODO: Replace all unwraps with proper error handling
    let str = String::from_utf8(entire_body.to_vec())?;
    let mut data : serde_json::Value = serde_json::from_str(&str)?;
    data["test"] = serde_json::Value::from("test_value");
    let json = serde_json::to_string(&data)?;
    let response = Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/json")
        .body(Body::from(json))?;
    Ok(response)
}

async fn api_get_response() -> Result<Response, GenericError> {
    let data = vec!["foo", "bar"];
    let res = match serde_json::to_string(&data) {
        Ok(json) => {
            Response::builder()
                .header(header::CONTENT_TYPE, "application/json")
                .body(Body::from(json))
                .unwrap()
        }
        Err(_) => {
            Response::builder()
                .status(StatusCode::INTERNAL_SERVER_ERROR)
                .body(Body::from("Internal Server Error"))
                .unwrap()
        }
    };
    Ok(res)
}

async fn response_examples(req: Request, client: Client<HttpConnector>)
    -> Result<Response, GenericError>
{
    match (req.method(), req.uri().path()) {
        (&Method::GET, "/") | (&Method::GET, "/index.html") => {
            let body = Body::from(INDEX);
            Ok(Response::new(body))
        },
        (&Method::GET, "/test.html") => {
            client_request_response(&client).await
        },
        (&Method::POST, "/json_api") => {
            api_post_response(req).await
        },
        (&Method::GET, "/json_api") => {
            api_get_response().await
        }
        _ => {
            // Return 404 not found response.
            let body = Body::from(NOTFOUND);
            Ok(Response::builder()
                        .status(StatusCode::NOT_FOUND)
                        .body(body)
                        .unwrap())
        }
    }
}

#[hyper::rt::main]
async fn main() -> Result<(), GenericError> {
    pretty_env_logger::init();

    let addr = "127.0.0.1:1337".parse().unwrap();

    // Share a `Client` with all `Service`s
    let client = Client::new();

    let new_service = make_service_fn(move |_| {
        // Move a clone of `client` into the `service_fn`.
        let client = client.clone();
        async {
            Ok::<_, GenericError>(service_fn(move |req| {
                // Clone again to ensure that client outlives this closure.
                response_examples(req, client.to_owned())
            }))
        }
    });

    let server = Server::bind(&addr)
        .serve(new_service);

    println!("Listening on http://{}", addr);

    server.await?;

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture};
use futures::future::Either;
use futures::sync::mpsc;
use futures::unsync;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Retry, Action};
use tokio_retry::strategy::{FixedInterval, jitter};

use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

macro_rules! try_future_boxed {
    ($e:expr) =>(
        match $e {
            Ok(v) => v,
            Err(e) => {
                return Err(into_other(e)).into_future().boxed();
            }
        }
    )
}

impl NetworkPart {
    pub fn run(self) -> Result<(), ()> {
        let mut core = Core::new().unwrap();
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);
        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();
        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let cancel_sender = cancel_sender.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());
                        // Enable retry feature for outgoing connection.
                        let requests_tx = requests_tx.clone();
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FixedInterval::from_millis(timeout).map(jitter).take(
                                max_tries,
                            ),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            .and_then(move |sock| {
                                try_future_boxed!(sock.set_nodelay(network_config.tcp_nodelay));
                                try_future_boxed!(sock.set_keepalive(
                                    network_config.tcp_keep_alive.map(Duration::from_millis),
                                ));

                                trace!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok).map_err(into_other);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                                    .boxed()
                            })
                            .then(move |res| {
                                trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);
                    handle.spawn(send_handle);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(forget_result)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle()).unwrap();
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                trace!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        trace!("Received handshake message={:?}", connect);
                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        })
                    })
                    .map(forget_result)
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .into_future()
                .map(|_| trace!("Network thread shutdown"))
                .map_err(|_| error!("An error during shutdown occured")),
        )
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}
<commit_msg>Remove unwraps<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use futures::{Future, Stream, Sink, IntoFuture};
use futures::future::Either;
use futures::sync::mpsc;
use futures::unsync;
use tokio_core::net::{TcpListener, TcpStream, TcpStreamNew};
use tokio_core::reactor::{Core, Handle};
use tokio_io::AsyncRead;
use tokio_retry::{Retry, Action};
use tokio_retry::strategy::{FixedInterval, jitter};

use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;

use messages::{Any, Connect, RawMessage};
use node::{ExternalMessage, NodeTimeout};
use helpers::Milliseconds;

use super::EventHandler;
use super::error::{other_error, result_ok, forget_result, into_other, log_error};
use super::codec::MessagesCodec;
use super::EventsAggregator;

#[derive(Debug)]
pub enum NetworkEvent {
    MessageReceived(SocketAddr, RawMessage),
    PeerConnected(SocketAddr, Connect),
    PeerDisconnected(SocketAddr),
}

#[derive(Debug, Clone)]
pub enum NetworkRequest {
    SendMessage(SocketAddr, RawMessage),
    DisconnectWithPeer(SocketAddr),
    Shutdown,
}

#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
pub struct NetworkConfiguration {
    // TODO: think more about config parameters
    pub max_incoming_connections: usize,
    pub max_outgoing_connections: usize,
    pub tcp_nodelay: bool,
    pub tcp_keep_alive: Option<u64>,
    pub tcp_connect_retry_timeout: Milliseconds,
    pub tcp_connect_max_retries: u64,
}

impl Default for NetworkConfiguration {
    fn default() -> NetworkConfiguration {
        NetworkConfiguration {
            max_incoming_connections: 128,
            max_outgoing_connections: 128,
            tcp_keep_alive: None,
            tcp_nodelay: false,
            tcp_connect_retry_timeout: 15_000,
            tcp_connect_max_retries: 10,
        }
    }
}

#[derive(Debug)]
pub struct HandlerPart<H: EventHandler> {
    pub core: Core,
    pub handler: H,
    pub timeout_rx: mpsc::Receiver<NodeTimeout>,
    pub network_rx: mpsc::Receiver<NetworkEvent>,
    pub api_rx: mpsc::Receiver<ExternalMessage>,
}

#[derive(Debug)]
pub struct NetworkPart {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub network_requests: (mpsc::Sender<NetworkRequest>, mpsc::Receiver<NetworkRequest>),
    pub network_tx: mpsc::Sender<NetworkEvent>,
}

impl<H: EventHandler> HandlerPart<H> {
    pub fn run(self) -> Result<(), ()> {
        let mut core = self.core;
        let mut handler = self.handler;

        let events_handle = EventsAggregator::new(self.timeout_rx, self.network_rx, self.api_rx)
            .for_each(move |event| {
                handler.handle_event(event);
                Ok(())
            });
        core.run(events_handle)
    }
}

macro_rules! try_future_boxed {
    ($e:expr) =>(
        match $e {
            Ok(v) => v,
            Err(e) => {
                return Err(into_other(e)).into_future().boxed();
            }
        }
    )
}

impl NetworkPart {
    pub fn run(self) -> Result<(), io::Error> {
        let mut core = Core::new()?;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);
        // Outgoing connections handler
        let mut outgoing_connections: HashMap<SocketAddr, mpsc::Sender<RawMessage>> =
            HashMap::new();
        // Requests handler
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_tx = self.network_requests.0.clone();
        let network_config = self.network_config;
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let cancel_sender = cancel_sender.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(&peer).cloned() {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, conn_tx.clone());
                        // Enable retry feature for outgoing connection.
                        let requests_tx = requests_tx.clone();
                        let timeout = network_config.tcp_connect_retry_timeout;
                        let max_tries = network_config.tcp_connect_max_retries as usize;
                        let connect_handle = Retry::spawn(
                            handle.clone(),
                            FixedInterval::from_millis(timeout).map(jitter).take(
                                max_tries,
                            ),
                            TcpStreamConnectAction {
                                handle: handle.clone(),
                                peer,
                            },
                        ).map_err(into_other)
                            .and_then(move |sock| {
                                try_future_boxed!(sock.set_nodelay(network_config.tcp_nodelay));
                                try_future_boxed!(sock.set_keepalive(
                                    network_config.tcp_keep_alive.map(Duration::from_millis),
                                ));

                                trace!("Established connection with peer={}", peer);

                                let stream = sock.framed(MessagesCodec);
                                let (sink, stream) = stream.split();

                                let writer = conn_rx
                                    .map_err(|_| other_error("Can't send data into socket"))
                                    .forward(sink);
                                let reader = stream.for_each(result_ok).map_err(into_other);

                                reader
                                    .select2(writer)
                                    .map_err(|_| other_error("Socket error"))
                                    .and_then(|res| match res {
                                        Either::A((_, _reader)) => Ok(()).into_future(),
                                        Either::B((_, _writer)) => Ok(()).into_future(),
                                    })
                                    .boxed()
                            })
                            .then(move |res| {
                                trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                let request = NetworkRequest::DisconnectWithPeer(peer);
                                requests_tx
                                    .clone()
                                    .send(request)
                                    .map(forget_result)
                                    .map_err(into_other)
                            })
                            .map_err(log_error);
                        handle.spawn(connect_handle);
                        conn_tx
                    };

                    let send_handle = conn_tx.send(msg).map(forget_result).map_err(log_error);
                    handle.spawn(send_handle);
                }
                NetworkRequest::DisconnectWithPeer(peer) => {
                    outgoing_connections.remove(&peer);

                    let event = NetworkEvent::PeerDisconnected(peer);
                    let event_handle = network_tx.clone().send(event).map(forget_result).map_err(
                        log_error,
                    );
                    handle.spawn(event_handle);
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(forget_result)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections handler
        let listener = TcpListener::bind(&self.listen_address, &core.handle())?;
        let network_tx = self.network_tx.clone();
        let handle = core.handle();
        let server = listener
            .incoming()
            .for_each(move |(sock, addr)| {
                trace!("Accepted incoming connection with peer={}", addr);

                let stream = sock.framed(MessagesCodec);
                let (_, stream) = stream.split();
                let network_tx = network_tx.clone();
                let connection_handler = stream
                    .into_future()
                    .map_err(|e| e.0)
                    .and_then(move |(raw, stream)| match raw.map(Any::from_raw) {
                        Some(Ok(Any::Connect(msg))) => Ok((msg, stream)),
                        Some(Ok(other)) => Err(other_error(&format!(
                            "First message is not Connect, got={:?}",
                            other
                        ))),
                        Some(Err(e)) => Err(into_other(e)),
                        None => Err(other_error("Incoming socket closed")),
                    })
                    .and_then(move |(connect, stream)| {
                        trace!("Received handshake message={:?}", connect);
                        let event = NetworkEvent::PeerConnected(addr, connect);
                        let stream = network_tx
                            .clone()
                            .send(event)
                            .map_err(into_other)
                            .and_then(move |_| Ok(stream))
                            .flatten_stream();

                        stream.for_each(move |raw| {
                            let event = NetworkEvent::MessageReceived(addr, raw);
                            network_tx.clone().send(event).map(forget_result).map_err(
                                into_other,
                            )
                        })
                    })
                    .map(forget_result)
                    .map_err(log_error);
                handle.spawn(connection_handler);
                Ok(())
            })
            .map_err(log_error);

        core.handle().spawn(server);
        core.handle().spawn(requests_handle);
        core.run(
            cancel_handler
                .into_future()
                .map(|_| trace!("Network thread shutdown"))
                .map_err(|_| other_error("An error during `Core` shutdown occured")),
        )
    }
}

// Tcp streams source for retry action.
struct TcpStreamConnectAction {
    handle: Handle,
    peer: SocketAddr,
}

impl Action for TcpStreamConnectAction {
    type Future = TcpStreamNew;
    type Item = <TcpStreamNew as Future>::Item;
    type Error = <TcpStreamNew as Future>::Error;

    fn run(&mut self) -> Self::Future {
        TcpStream::connect(&self.peer, &self.handle)
    }
}
<|endoftext|>


<commit_before>//\TODO don't reload whitelisted_peers if path the same
use std::collections::BTreeSet;

use crypto::PublicKey;

/// `Whitelist` is special set to keep peers that can connect to us.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Whitelist {
    whitelist_on: bool,
    whitelisted_peers: BTreeSet<PublicKey>,

    #[serde(default)]
    #[serde(skip_serializing)]
    #[serde(skip_deserializing)]
    validators_list: BTreeSet<PublicKey>,
}

impl Whitelist {
    /// is this `peer` can connect or not
    pub fn allow(&self, peer: &PublicKey) -> bool {
        !self.whitelist_on || self.validators_list.contains(peer) ||
        self.whitelisted_peers.contains(peer)
    }

    /// append `peer` to whitelist
    pub fn add(&mut self, peer: PublicKey) {
        self.whitelisted_peers.insert(peer);
    }

    /// get list of peers in whitelist
    pub fn collect_allowed(&self) -> Vec<&PublicKey> {
        self.whitelisted_peers
            .iter()
            .chain(self.validators_list.iter())
            .collect()
    }

    pub fn set_validators<I>(&mut self, list: I)
        where I: IntoIterator<Item = PublicKey>
    {
        self.validators_list = list.into_iter().collect();
    }

    /// check if we support whitelist, or keep connection politics open
    /// if it return false, everybody can connect to us
    pub fn is_enabled(&self) -> bool {
        !self.whitelist_on
    }
}

#[cfg(test)]
mod test {
    use super::Whitelist;
    use crypto::PublicKey;

    static VALIDATORS: [[[u8; 32]; 2]; 2] = [[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                                              [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]],
                                             [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1],
                                              [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2]]];
    static REGULAR_PEERS: [[u8; 32]; 4] = [[2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                                           [2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2],
                                           [2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3],
                                           [2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                                            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4]];

    fn make_keys(source: &[[u8; 32]]) -> Vec<PublicKey> {
        source
            .iter()
            .map(|k| PublicKey::from_slice(k).unwrap())
            .collect()
    }

    fn check_in_whitelist(whitelist: &Whitelist,
                          keys: &[PublicKey],
                          in_whitelist: &[usize],
                          not_in_whitelist: &[usize]) {
        for i in in_whitelist {
            assert_eq!(whitelist.allow(&keys[*i]), true);
        }
        for i in not_in_whitelist {
            assert_eq!(whitelist.allow(&keys[*i]), false);
        }
    }

    #[test]
    fn test_whitelist() {
        let regular = make_keys(&REGULAR_PEERS);

        let mut whitelist = Whitelist::default();
        whitelist.whitelist_on = true;
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        whitelist.add(regular[0]);
        check_in_whitelist(&whitelist, &regular, &[0], &[1, 2, 3]);
        whitelist.add(regular[2]);
        check_in_whitelist(&whitelist, &regular, &[0, 2], &[1, 3]);
        assert_eq!(whitelist.collect_allowed().len(), 2);
    }

    #[test]
    fn test_wildcard() {
        let regular = make_keys(&REGULAR_PEERS);

        let whitelist = Whitelist::default();
        check_in_whitelist(&whitelist, &regular, &[0, 1, 2, 3], &[]);
        assert_eq!(whitelist.collect_allowed().len(), 0);
    }

    #[test]
    fn test_validators_in_whitelist() {
        let regular = make_keys(&REGULAR_PEERS);
        let validators = make_keys(&VALIDATORS[0]);
        let mut whitelist = Whitelist::default();
        whitelist.whitelist_on = true;
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        check_in_whitelist(&whitelist, &validators, &[], &[0, 1]);
        assert_eq!(whitelist.collect_allowed().len(), 0);
        whitelist.set_validators(validators.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        check_in_whitelist(&whitelist, &validators, &[0, 1], &[]);
    }

    #[test]
    fn test_update_validators() {
        let validators0 = make_keys(&VALIDATORS[0]);
        let validators1 = make_keys(&VALIDATORS[1]);
        let mut whitelist = Whitelist::default();
        whitelist.whitelist_on = true;
        assert_eq!(whitelist.collect_allowed().len(), 0);
        whitelist.set_validators(validators0.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &validators0, &[0, 1], &[]);
        check_in_whitelist(&whitelist, &validators1, &[], &[0, 1]);
        whitelist.set_validators(validators1.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &validators0, &[], &[0, 1]);
        check_in_whitelist(&whitelist, &validators1, &[0, 1], &[]);
    }
}
<commit_msg>generate testdata from rng<commit_after>//\TODO don't reload whitelisted_peers if path the same
use std::collections::BTreeSet;

use crypto::PublicKey;

/// `Whitelist` is special set to keep peers that can connect to us.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Whitelist {
    whitelist_enabled: bool,
    whitelisted_peers: BTreeSet<PublicKey>,

    #[serde(default)]
    #[serde(skip_serializing)]
    #[serde(skip_deserializing)]
    validators_list: BTreeSet<PublicKey>,
}

impl Whitelist {
    /// is this `peer` can connect or not
    pub fn allow(&self, peer: &PublicKey) -> bool {
        !self.whitelist_enabled || self.validators_list.contains(peer) ||
        self.whitelisted_peers.contains(peer)
    }

    /// append `peer` to whitelist
    pub fn add(&mut self, peer: PublicKey) {
        self.whitelisted_peers.insert(peer);
    }

    /// get list of peers in whitelist
    pub fn collect_allowed(&self) -> Vec<&PublicKey> {
        self.whitelisted_peers
            .iter()
            .chain(self.validators_list.iter())
            .collect()
    }

    pub fn set_validators<I>(&mut self, list: I)
        where I: IntoIterator<Item = PublicKey>
    {
        self.validators_list = list.into_iter().collect();
    }

    /// check if we support whitelist, or keep connection politics open
    /// if it return false, everybody can connect to us
    pub fn is_enabled(&self) -> bool {
        !self.whitelist_enabled
    }
}

#[cfg(test)]
mod test {
    use super::Whitelist;
    use crypto::PublicKey;
    use rand::{Rand, SeedableRng, XorShiftRng};

    static VALIDATORS: [[u32; 4]; 2] = [[123, 45, 67, 89], [223, 45, 67, 98]];
    static REGULAR_PEERS: [u32; 4] = [5, 6, 7, 9];

    fn make_keys(source: [u32; 4], count: usize) -> Vec<PublicKey> {
        let mut rng = XorShiftRng::from_seed(source);
        (0..count)
            .into_iter()
            .map(|_| PublicKey::from_slice(&<[u8;32] as Rand>::rand(&mut rng)).unwrap())
            .collect()
    }

    fn check_in_whitelist(whitelist: &Whitelist,
                          keys: &[PublicKey],
                          in_whitelist: &[usize],
                          not_in_whitelist: &[usize]) {
        for i in in_whitelist {
            assert_eq!(whitelist.allow(&keys[*i]), true);
        }
        for i in not_in_whitelist {
            assert_eq!(whitelist.allow(&keys[*i]), false);
        }
    }

    #[test]
    fn test_whitelist() {
        let regular = make_keys(REGULAR_PEERS, 4);

        let mut whitelist = Whitelist::default();
        whitelist.whitelist_enabled = true;
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        whitelist.add(regular[0]);
        check_in_whitelist(&whitelist, &regular, &[0], &[1, 2, 3]);
        whitelist.add(regular[2]);
        check_in_whitelist(&whitelist, &regular, &[0, 2], &[1, 3]);
        assert_eq!(whitelist.collect_allowed().len(), 2);
    }

    #[test]
    fn test_wildcard() {
        let regular = make_keys(REGULAR_PEERS, 4);

        let whitelist = Whitelist::default();
        check_in_whitelist(&whitelist, &regular, &[0, 1, 2, 3], &[]);
        assert_eq!(whitelist.collect_allowed().len(), 0);
    }

    #[test]
    fn test_validators_in_whitelist() {
        let regular = make_keys(REGULAR_PEERS, 4);
        let validators = make_keys(VALIDATORS[0], 2);
        let mut whitelist = Whitelist::default();
        whitelist.whitelist_enabled = true;
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        check_in_whitelist(&whitelist, &validators, &[], &[0, 1]);
        assert_eq!(whitelist.collect_allowed().len(), 0);
        whitelist.set_validators(validators.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &regular, &[], &[0, 1, 2, 3]);
        check_in_whitelist(&whitelist, &validators, &[0, 1], &[]);
    }

    #[test]
    fn test_update_validators() {
        let validators0 = make_keys(VALIDATORS[0], 2);
        let validators1 = make_keys(VALIDATORS[1], 2);
        let mut whitelist = Whitelist::default();
        whitelist.whitelist_enabled = true;
        assert_eq!(whitelist.collect_allowed().len(), 0);
        whitelist.set_validators(validators0.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &validators0, &[0, 1], &[]);
        check_in_whitelist(&whitelist, &validators1, &[], &[0, 1]);
        whitelist.set_validators(validators1.clone());
        assert_eq!(whitelist.collect_allowed().len(), 2);
        check_in_whitelist(&whitelist, &validators0, &[], &[0, 1]);
        check_in_whitelist(&whitelist, &validators1, &[0, 1], &[]);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{auth::with_auth, context::Context, types::auth::Claims};
use aptos_config::config::PeerRole;
use aptos_logger::{debug, error};
use warp::{filters::BoxedFilter, hyper::body::Bytes, reply, Filter, Rejection, Reply};

pub fn metrics_ingest(context: Context) -> BoxedFilter<(impl Reply,)> {
    warp::path!("push-metrics")
        .and(warp::post())
        .and(context.clone().filter())
        .and(with_auth(
            context,
            vec![PeerRole::Validator, PeerRole::Unknown],
        ))
        .and(warp::body::bytes())
        .and_then(handle_metrics_ingest)
        .boxed()
}

pub async fn handle_metrics_ingest(
    context: Context,
    claims: Claims,
    metrics_body: Bytes,
) -> anyhow::Result<impl Reply, Rejection> {
    let extra_labels = vec![
        format!("peer_id={}", claims.peer_id),
        format!("peer_role={}", claims.peer_role as u16),
        format!("chain_name={}", claims.chain_id),
        format!("namespace={}", "telemetry-service"),
        format!(
            "kubernetes_pod_name={}/{}",
            claims.peer_role as u16, claims.peer_id
        ),
    ];

    let res = context
        .victoria_metrics_client
        .unwrap()
        .post_prometheus_metrics(metrics_body, extra_labels)
        .await;

    match res {
        Ok(res) => {
            if res.status().is_success() {
                debug!("remote write to victoria metrics succeeded");
            } else {
                error!(
                    "remote write failed to victoria_metrics: {}",
                    res.error_for_status().err().unwrap()
                );
            }
        }
        Err(err) => {
            error!("error sending remote write request: {}", err);
        }
    }

    Ok(reply::reply())
}
<commit_msg>telemetry-service: optimize label naming for readability and compatibility (#3519)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{auth::with_auth, context::Context, types::auth::Claims};
use aptos_config::config::PeerRole;
use aptos_logger::{debug, error};
use warp::{filters::BoxedFilter, hyper::body::Bytes, reply, Filter, Rejection, Reply};

pub fn metrics_ingest(context: Context) -> BoxedFilter<(impl Reply,)> {
    warp::path!("push-metrics")
        .and(warp::post())
        .and(context.clone().filter())
        .and(with_auth(
            context,
            vec![PeerRole::Validator, PeerRole::Unknown],
        ))
        .and(warp::body::bytes())
        .and_then(handle_metrics_ingest)
        .boxed()
}

pub async fn handle_metrics_ingest(
    context: Context,
    claims: Claims,
    metrics_body: Bytes,
) -> anyhow::Result<impl Reply, Rejection> {
    let extra_labels = claims_to_extra_labels(&claims);

    let res = context
        .victoria_metrics_client
        .unwrap()
        .post_prometheus_metrics(metrics_body, extra_labels)
        .await;

    match res {
        Ok(res) => {
            if res.status().is_success() {
                debug!("remote write to victoria metrics succeeded");
            } else {
                error!(
                    "remote write failed to victoria_metrics: {}",
                    res.error_for_status().err().unwrap()
                );
            }
        }
        Err(err) => {
            error!("error sending remote write request: {}", err);
        }
    }

    Ok(reply::reply())
}

fn claims_to_extra_labels(claims: &Claims) -> Vec<String> {
    vec![
        format!("role={}", claims.peer_role),
        format!("chain_name={}", claims.chain_id),
        format!("namespace={}", "telemetry-service"),
        // for community nodes we cannot determine which pod name they run in (or whether they run in k8s at all), so we use the peer id as an approximation/replacement for pod_name
        // This works well with our existing grafana dashboards
        format!(
            "kubernetes_pod_name=peer_id:{}",
            claims.peer_id.to_hex_literal()
        ),
    ]
}

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use super::*;
    use aptos_types::{chain_id::ChainId, PeerId};
    #[test]
    fn verify_labels() {
        let claims = claims_to_extra_labels(&super::Claims {
            chain_id: ChainId::new(25),
            peer_id: PeerId::from_str("0x1").unwrap(),
            peer_role: PeerRole::Validator,
            epoch: 3,
            exp: 123,
            iat: 123,
        });
        assert_eq!(
            claims,
            vec![
                "role=validator",
                "chain_name=25",
                "namespace=telemetry-service",
                "kubernetes_pod_name=peer_id:0x1",
            ]
        );
    }
}
<|endoftext|>


<commit_before>#![feature(buf_stream)]

extern crate podio;

pub use protocol::Protocol;
pub use protocol::Error;
pub use transport::Transport;

pub mod protocol;
pub mod transport;
pub mod server;
pub mod processor;

#[derive(Debug)]
pub enum ThriftErr {
    TransportError(std::io::Error),
    ProtocolError1(protocol::Error),
    // Received a user-defined exception from the far end
    Exception,
}

impl std::convert::From<std::io::Error> for ThriftErr {
    fn from(err: std::io::Error) -> ThriftErr {
        ThriftErr::TransportError(err)
    }
}

impl std::convert::From<protocol::Error> for ThriftErr {
    fn from(err: protocol::Error) -> ThriftErr {
        ThriftErr::ProtocolError1(err)
    }
}

pub type TResult<T> = Result<T, ThriftErr>;
<commit_msg>Fixed the ProtocolError name<commit_after>#![feature(buf_stream)]

extern crate podio;

pub use protocol::Protocol;
pub use protocol::Error;
pub use transport::Transport;

pub mod protocol;
pub mod transport;
pub mod server;
pub mod processor;

#[derive(Debug)]
pub enum ThriftErr {
    TransportError(std::io::Error),
    ProtocolError(protocol::Error),
    // Received a user-defined exception from the far end
    Exception,
}

impl std::convert::From<std::io::Error> for ThriftErr {
    fn from(err: std::io::Error) -> ThriftErr {
        ThriftErr::TransportError(err)
    }
}

impl std::convert::From<protocol::Error> for ThriftErr {
    fn from(err: protocol::Error) -> ThriftErr {
        ThriftErr::ProtocolError(err)
    }
}

pub type TResult<T> = Result<T, ThriftErr>;
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::DocumentBinding::DocumentMethods;
use dom::bindings::codegen::Bindings::NodeBinding::NodeMethods;
use dom::bindings::inheritance::Castable;
use dom::bindings::js::RootedReference;
use dom::bindings::str::DOMString;
use dom::htmlheadelement::HTMLHeadElement;
use dom::node::Node;
use servo_config::opts;
use servo_config::resource_files::resources_dir_path;
use std::borrow::ToOwned;
use std::fs::read_dir;
use std::path::PathBuf;

pub fn load_script(head: &HTMLHeadElement) {
    if let Some(ref path_str) = opts::get().userscripts {
        let node = head.upcast::<Node>();
        let first_child = node.GetFirstChild();
        let doc = node.owner_doc();

        let path = if &**path_str == "" {
            if let Ok(mut p) = resources_dir_path() {
                p.push("user-agent-js");
                p
            } else {
                return
            }
        } else {
            PathBuf::from(path_str)
        };

        let mut files = read_dir(&path).expect("Bad path passed to --userscripts")
                                       .filter_map(|e| e.ok())
                                       .map(|e| e.path()).collect::<Vec<_>>();

        files.sort();

        for file in files {
            let name = match file.into_os_string().into_string() {
                Ok(ref s) if s.ends_with(".js") => "file://".to_owned() + &s[..],
                _ => continue
            };
            let new_script = doc.CreateElement(DOMString::from("script")).unwrap();
            new_script.set_string_attribute(&local_name!("src"), DOMString::from(name));
            node.InsertBefore(new_script.upcast(), first_child.r()).unwrap();
        }
    }
}
<commit_msg>Auto merge of #15874 - sendilkumarn:user-script, r=nox<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::inheritance::Castable;
use dom::globalscope::GlobalScope;
use dom::htmlheadelement::HTMLHeadElement;
use dom::node::Node;
use js::jsval::UndefinedValue;
use servo_config::opts;
use servo_config::resource_files::resources_dir_path;
use std::fs::{File, read_dir};
use std::io::Read;
use std::path::PathBuf;

pub fn load_script(head: &HTMLHeadElement) {
    if let Some(ref path_str) = opts::get().userscripts {
        let node = head.upcast::<Node>();
        let doc = node.owner_doc();
        let win = doc.window();
        let cx = win.get_cx();
        rooted!(in(cx) let mut rval = UndefinedValue());

        let path = if &**path_str == "" {
            if let Ok(mut p) = resources_dir_path() {
                p.push("user-agent-js");
                p
            } else {
                return
            }
        } else {
            PathBuf::from(path_str)
        };

        let mut files = read_dir(&path).expect("Bad path passed to --userscripts")
                                       .filter_map(|e| e.ok())
                                       .map(|e| e.path()).collect::<Vec<_>>();

        files.sort();

        for file in files {
            let mut f = File::open(&file).unwrap();
            let mut contents = vec![];
            f.read_to_end(&mut contents).unwrap();
            let script_text = String::from_utf8_lossy(&contents);
            win.upcast::<GlobalScope>().evaluate_js_on_global_with_result(&script_text, rval.handle_mut());
        }
    }
}
<|endoftext|>


<commit_before>mod iter;
mod coord;
mod test;

use std::mem;

use traits::Grid;
use traits::Cell;
use traits::EvolutionState;
use traits::Nhood;
use traits::Coord;

use self::iter::Iter;
pub use self::coord::GridCoord;

pub struct TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    cells: Vec<C>,
    old_cells: Vec<C>,
    evolution_state: C::State,
    nhood: N,
    neighbors: Vec<Vec<Option<usize>>>,
    dimensions: GridCoord,
    rows: i32,
    cols: i32,
}

impl<C, N> TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    pub fn new(rows: i32, cols: i32, nhood: N, state: C::State) -> Self {

        let len = (rows * cols) as usize;

        let cells = Vec::with_capacity(len);
        let old_cells = Vec::with_capacity(len);
        let neighbors = Vec::with_capacity(len);

        let mut grid = TwodimGrid {
            cells: cells,
            old_cells: old_cells,
            evolution_state: state,
            nhood: nhood,
            neighbors: neighbors,
            rows: rows,
            cols: cols,
            dimensions: GridCoord::from_2d(cols, rows),
        };

        grid.init();

        grid
    }

    fn init(&mut self) {

        for offset in 0..self.rows * self.cols {

            let coord = GridCoord::from_offset(offset as i32, self.rows, self.cols);

            // init neighbors
            let neighbors = self.get_neighbors(&coord);
            self.neighbors.push(neighbors);

            // init cells
            let cell = C::with_coord(coord);
            self.cells.push(cell.clone());
            self.old_cells.push(cell.clone());
        }
    }

    fn get_neighbors(&self, coord: &GridCoord) -> Vec<Option<usize>> {

        let neighbors_count = self.nhood.neighbors_count();
        let mut neighbors = Vec::with_capacity(neighbors_count);

        for coord in self.nhood.neighbors(coord).iter() {

            if coord.x() >= 0 && coord.x() < self.cols && coord.y() >= 0 && coord.y() < self.rows {
                neighbors.push(Some(self.offset(coord)));
            } else {
                neighbors.push(None);
            }
        }

        neighbors
    }

    fn neighbors_iter<'b>(&self,
                          cells: &'b Vec<C>,
                          neighbors: &'b Vec<Option<usize>>)
                          -> Iter<'b, C> {

        Iter::new(cells, neighbors, self.nhood.neighbors_count())
    }

    #[inline]
    pub fn offset<Crd: Coord>(&self, coord: &Crd) -> usize {
        (coord.y() * self.cols + coord.x()) as usize
    }
}

impl<C, N> Grid for TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    type Cell = C;
    type Coord = GridCoord;

    fn update(&mut self) {
        mem::swap(&mut self.cells, &mut self.old_cells);

        for i in 0..self.cells.len() {
            unsafe {
                let ref neighbors = self.neighbors.get_unchecked(i);
                let neighbors_iter = self.neighbors_iter(&self.old_cells, &neighbors);

                let old = self.old_cells.get_unchecked(i);
                let cell = self.cells.get_unchecked_mut(i);
                cell.update(old, neighbors_iter, &self.evolution_state);
            }
        }

        self.evolution_state.update();
    }

    fn set_cells(&mut self, new_cells: Vec<Self::Cell>) {

        for cell in new_cells.into_iter() {

            let index;

            {
                index = self.offset(cell.coord());
            }

            self.cells[index] = cell;
        }
    }

    fn cells(&self) -> &Vec<Self::Cell> { &self.cells }

    fn state(&self) -> &<<Self as Grid>::Cell as Cell>::State { &self.evolution_state }

    fn dimensions(&self) -> Self::Coord { self.dimensions.clone() }
}
<commit_msg>Note about TwodimGrid. Offset is made private.<commit_after>//! 2D grid with neighbors iter and custom internal coordinate.

mod iter;
mod coord;
mod test;

use std::mem;

use traits::Grid;
use traits::Cell;
use traits::EvolutionState;
use traits::Nhood;
use traits::Coord;

use self::iter::Iter;
pub use self::coord::GridCoord;

/// 2D grid. Implemented with two buffers.
/// They are swapped on every evolution step.
/// Old buffer is used for read-only neighbors data.
/// New buffer is writable and mutated through update
/// process.
pub struct TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    cells: Vec<C>,
    old_cells: Vec<C>,
    evolution_state: C::State,
    nhood: N,
    neighbors: Vec<Vec<Option<usize>>>,
    dimensions: GridCoord,
    rows: i32,
    cols: i32,
}

impl<C, N> TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    /// Constructs TwodimGrid with given ROWSxCOLS, neighborhood
    /// strategy and initial evolution state.
    pub fn new(rows: i32, cols: i32, nhood: N, state: C::State) -> Self {

        let len = (rows * cols) as usize;

        let cells = Vec::with_capacity(len);
        let old_cells = Vec::with_capacity(len);
        let neighbors = Vec::with_capacity(len);

        let mut grid = TwodimGrid {
            cells: cells,
            old_cells: old_cells,
            evolution_state: state,
            nhood: nhood,
            neighbors: neighbors,
            rows: rows,
            cols: cols,
            dimensions: GridCoord::from_2d(cols, rows),
        };

        grid.init();

        grid
    }

    fn init(&mut self) {

        for offset in 0..self.rows * self.cols {

            let coord = GridCoord::from_offset(offset as i32, self.rows, self.cols);

            // init neighbors
            let neighbors = self.get_neighbors(&coord);
            self.neighbors.push(neighbors);

            // init cells
            let cell = C::with_coord(coord);
            self.cells.push(cell.clone());
            self.old_cells.push(cell.clone());
        }
    }

    fn get_neighbors(&self, coord: &GridCoord) -> Vec<Option<usize>> {

        let neighbors_count = self.nhood.neighbors_count();
        let mut neighbors = Vec::with_capacity(neighbors_count);

        for coord in self.nhood.neighbors(coord).iter() {

            if coord.x() >= 0 && coord.x() < self.cols && coord.y() >= 0 && coord.y() < self.rows {
                neighbors.push(Some(self.offset(coord)));
            } else {
                neighbors.push(None);
            }
        }

        neighbors
    }

    fn neighbors_iter<'b>(&self,
                          cells: &'b Vec<C>,
                          neighbors: &'b Vec<Option<usize>>)
                          -> Iter<'b, C> {

        Iter::new(cells, neighbors, self.nhood.neighbors_count())
    }

    #[inline]
    fn offset<Crd: Coord>(&self, coord: &Crd) -> usize {
        (coord.y() * self.cols + coord.x()) as usize
    }
}

impl<C, N> Grid for TwodimGrid<C, N>
    where C: Cell + Clone,
          N: Nhood<Coord = GridCoord>,
{
    type Cell = C;
    type Coord = GridCoord;

    fn update(&mut self) {
        mem::swap(&mut self.cells, &mut self.old_cells);

        for i in 0..self.cells.len() {
            unsafe {
                let ref neighbors = self.neighbors.get_unchecked(i);
                let neighbors_iter = self.neighbors_iter(&self.old_cells, &neighbors);

                let old = self.old_cells.get_unchecked(i);
                let cell = self.cells.get_unchecked_mut(i);
                cell.update(old, neighbors_iter, &self.evolution_state);
            }
        }

        self.evolution_state.update();
    }

    fn set_cells(&mut self, new_cells: Vec<Self::Cell>) {

        for cell in new_cells.into_iter() {

            let index;

            {
                index = self.offset(cell.coord());
            }

            self.cells[index] = cell;
        }
    }

    fn cells(&self) -> &Vec<Self::Cell> { &self.cells }

    fn state(&self) -> &<<Self as Grid>::Cell as Cell>::State { &self.evolution_state }

    fn dimensions(&self) -> Self::Coord { self.dimensions.clone() }
}
<|endoftext|>


<commit_before>use crate::ast::rules;

use super::utils::{interpret_rule, interpret_rule_env};

#[test]
fn test_closure_eval() {
    let (val, mut _env) = interpret_rule("function () break; end", rules::functiondef);
    assert_eq!(val, "Function { id: 1, parameters: [], varargs: false, body: Block { statements: [Break, Terminal(SEMICOLONS)], retstat: None }, env: 0 }");

    let (val, mut _env) = interpret_rule("function (b, c, ...) break; end", rules::functiondef);
    assert_eq!(val, r#"Function { id: 1, parameters: ["b", "c"], varargs: true, body: Block { statements: [Break, Terminal(SEMICOLONS)], retstat: None }, env: 0 }"#);
}

#[test]
fn test_function_eval() {
    let (_val, env) = interpret_rule("function t (...) break end", rules::stat);
    assert_eq!(env, r#"{"t": RefCell { value: Function { id: 1, parameters: [], varargs: true, body: Block { statements: [Break], retstat: None }, env: 0 } }}"#);

    let (_val, env) = interpret_rule("t = {}; function t:f(b, c, ...) break end", rules::block);
    assert_eq!(env, r#"{"t": RefCell { value: Table { id: 1, map: {String("f"): RefCell { value: Function { id: 2, parameters: ["self", "b", "c"], varargs: true, body: Block { statements: [Break], retstat: None }, env: 0 } }}, metatable: {}, border: 0 } }}"#);
}

#[test]
fn test_functioncall_eval() {
    let (_, mut env) = interpret_rule("function sum1(x) return x + 1; end", rules::stat);

    let (val, _) = interpret_rule_env("sum1(5)", rules::functioncall, &mut env);
    assert_eq!(val, "Number(6.0)");
}
<commit_msg>Add method test<commit_after>use crate::ast::rules;

use super::utils::{interpret_rule, interpret_rule_env};

#[test]
fn test_closure_eval() {
    let (val, mut _env) = interpret_rule("function () break; end", rules::functiondef);
    assert_eq!(val, "Function { id: 1, parameters: [], varargs: false, body: Block { statements: [Break, Terminal(SEMICOLONS)], retstat: None }, env: 0 }");

    let (val, mut _env) = interpret_rule("function (b, c, ...) break; end", rules::functiondef);
    assert_eq!(val, r#"Function { id: 1, parameters: ["b", "c"], varargs: true, body: Block { statements: [Break, Terminal(SEMICOLONS)], retstat: None }, env: 0 }"#);
}

#[test]
fn test_function_eval() {
    let (_val, env) = interpret_rule("function t (...) break end", rules::stat);
    assert_eq!(env, r#"{"t": RefCell { value: Function { id: 1, parameters: [], varargs: true, body: Block { statements: [Break], retstat: None }, env: 0 } }}"#);

    let (_val, env) = interpret_rule("t = {}; function t:f(b, c, ...) break end", rules::block);
    assert_eq!(env, r#"{"t": RefCell { value: Table { id: 1, map: {String("f"): RefCell { value: Function { id: 2, parameters: ["self", "b", "c"], varargs: true, body: Block { statements: [Break], retstat: None }, env: 0 } }}, metatable: {}, border: 0 } }}"#);
}

#[test]
fn test_functioncall_eval() {
    let (_, mut env) = interpret_rule("function sum1(x) return x + 1; end", rules::stat);

    let (val, _) = interpret_rule_env("sum1(5)", rules::functioncall, &mut env);
    assert_eq!(val, "Number(6.0)");
}

#[test]
fn test_method_eval() {
    let (_, mut env) = interpret_rule("tab = {}; tab.x = 5;", rules::block);
    let (_, mut env) = interpret_rule_env("function tab:sum(x) return x + self.x; end", rules::stat, &mut env);

    let (val, _) = interpret_rule_env("tab:sum(5)", rules::functioncall, &mut env);
    assert_eq!(val, "Number(10.0)");
}
<|endoftext|>