<commit_before>// Copyright (c) 2014-2016 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.

//! Support for sending and receiving data link layer packets using the WinPcap library

extern crate libc;

use std::cmp;
use std::collections::VecDeque;
use std::ffi::CString;
use std::str::from_utf8_unchecked;
use std::io;
use std::mem;
use std::slice;
use std::sync::Arc;
use std::net::IpAddr;

use bindings::{bpf, winpcap};
use datalink::{self, NetworkInterface};
use datalink::Channel::Ethernet;
use datalink::{EthernetDataLinkChannelIterator, EthernetDataLinkReceiver, EthernetDataLinkSender};
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::MacAddr;

struct WinPcapAdapter {
    adapter: winpcap::LPADAPTER,
}

impl Drop for WinPcapAdapter {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketCloseAdapter(self.adapter);
        }
    }
}

struct WinPcapPacket {
    packet: winpcap::LPPACKET,
}

impl Drop for WinPcapPacket {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketFreePacket(self.packet);
        }
    }
}

/// WinPcap specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
        }
    }
}

/// Create a datalink channel using the WinPcap library
#[inline]
pub fn channel(network_interface: &NetworkInterface, config: &Config)
    -> io::Result<datalink::Channel> {
    let mut read_buffer = Vec::new();
    read_buffer.resize(config.read_buffer_size, 0u8);

    let mut write_buffer = Vec::new();
    write_buffer.resize(config.write_buffer_size, 0u8);

    let adapter = unsafe {
        let net_if_str = CString::new(network_interface.name.as_bytes()).unwrap();
        winpcap::PacketOpenAdapter(net_if_str.as_ptr() as *mut libc::c_char)
    };
    if adapter.is_null() {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetHwFilter(adapter, winpcap::NDIS_PACKET_TYPE_PROMISCUOUS) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Set kernel buffer size
    let ret = unsafe { winpcap::PacketSetBuff(adapter, config.read_buffer_size as libc::c_int) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetReadTimeout(adapter, 0) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Immediate mode
    let ret = unsafe { winpcap::PacketSetMinToCopy(adapter, 1) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let read_packet = unsafe { winpcap::PacketAllocatePacket() };
    if read_packet.is_null() {
        unsafe {
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(read_packet,
                                  read_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.read_buffer_size as winpcap::UINT)
    }

    let write_packet = unsafe { winpcap::PacketAllocatePacket() };
    if write_packet.is_null() {
        unsafe {
            winpcap::PacketFreePacket(read_packet);
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(write_packet,
                                  write_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.write_buffer_size as winpcap::UINT)
    }

    let adapter = Arc::new(WinPcapAdapter { adapter: adapter });
    let sender = Box::new(DataLinkSenderImpl {
        adapter: adapter.clone(),
        _write_buffer: write_buffer,
        packet: WinPcapPacket { packet: write_packet },
    });
    let receiver = Box::new(DataLinkReceiverImpl {
        adapter: adapter,
        _read_buffer: read_buffer,
        packet: WinPcapPacket { packet: read_packet },
    });
    Ok(Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    adapter: Arc<WinPcapAdapter>,
    _write_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl EthernetDataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= unsafe { (*self.packet.packet).Length } as usize {
            None
        } else {
            let min = unsafe { cmp::min((*self.packet.packet).Length as usize, len) };
            let slice: &mut [u8] = unsafe {
                slice::from_raw_parts_mut((*self.packet.packet).Buffer as *mut u8, min)
            };
            for chunk in slice.chunks_mut(packet_size) {
                {
                    let eh = MutableEthernetPacket::new(chunk).unwrap();
                    func(eh);
                }

                // Make sure the right length of packet is sent
                let old_len = unsafe { (*self.packet.packet).Length };
                unsafe {
                    (*self.packet.packet).Length = packet_size as u32;
                }

                let ret = unsafe {
                    winpcap::PacketSendPacket(self.adapter.adapter, self.packet.packet, 0)
                };

                unsafe {
                    (*self.packet.packet).Length = old_len;
                }

                if ret == 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
            }
            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| {
                                eh.clone_from(packet);
                            })
    }
}

unsafe impl Send for DataLinkSenderImpl {}
unsafe impl Sync for DataLinkSenderImpl {}

struct DataLinkReceiverImpl {
    adapter: Arc<WinPcapAdapter>,
    _read_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl EthernetDataLinkReceiver for DataLinkReceiverImpl {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        let buflen = unsafe { (*self.packet.packet).Length } as usize;
        Box::new(DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: VecDeque::with_capacity(buflen / 64),
        })
    }
}

unsafe impl Send for DataLinkReceiverImpl {}
unsafe impl Sync for DataLinkReceiverImpl {}

struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: VecDeque<(usize, usize)>,
}

impl<'a> EthernetDataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a> {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        // NOTE Most of the logic here is identical to FreeBSD/OS X
        if self.packets.is_empty() {
            let ret = unsafe {
                winpcap::PacketReceivePacket(self.pc.adapter.adapter, self.pc.packet.packet, 0)
            };
            let buflen = match ret {
                0 => return Err(io::Error::last_os_error()),
                _ => unsafe { (*self.pc.packet.packet).ulBytesReceived },
            };
            let mut ptr = unsafe { (*self.pc.packet.packet).Buffer };
            let end = unsafe { (*self.pc.packet.packet).Buffer.offset(buflen as isize) };
            while ptr < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                (*self.pc.packet.packet).Buffer as isize;
                    self.packets.push_back((start as usize, (*packet).bh_caplen as usize));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        let slice = unsafe {
            let data = (*self.pc.packet.packet).Buffer as usize + start;
            slice::from_raw_parts(data as *const u8, len)
        };
        Ok(EthernetPacket::new(slice).unwrap())
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use
        // GetAdaptersAddresses
        winpcap::GetAdaptersInfo(&mut tmp, &mut adapters_size);
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    let mut adapters = Vec::with_capacity(vec_size as usize);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as usize);
    while !cursor.is_null() {
        let mac = unsafe {
            MacAddr((*cursor).Address[0],
                    (*cursor).Address[1],
                    (*cursor).Address[2],
                    (*cursor).Address[3],
                    (*cursor).Address[4],
                    (*cursor).Address[5])
        };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING };
        let mut ips: Vec<IpAddr> = Vec::new();
        while !ip_cursor.is_null() {
            let ip_str_ptr = unsafe { &(*ip_cursor) }.IpAddress.String.as_ptr() as *const i8;
            let bytes = unsafe { CStr::from_ptr(ip_str_ptr).to_bytes() };
            let ip_str = unsafe { from_utf8_unchecked(bytes).to_owned() };
            ips.push(ip_str.parse().unwrap());
            ip_cursor = unsafe { (*ip_cursor).Next };
        }

        unsafe {
            let name_str_ptr = (*cursor).AdapterName.as_ptr() as *const i8;

            let bytes = CStr::from_ptr(name_str_ptr).to_bytes();
            let name_str = from_utf8_unchecked(bytes).to_owned();

            all_ifaces.push(NetworkInterface {
                name: name_str,
                index: (*cursor).Index,
                mac: Some(mac),
                ips: Some(ips),
                // flags: (*cursor).Type, // FIXME [windows]
                flags: 0,
            });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = [0u8; 4096];
    let mut buflen = buf.len() as u32;

    // Gets list of supported adapters in form:
    // adapter1\0adapter2\0\0desc1\0desc2\0\0
    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        // FIXME [windows] Should allocate a buffer big enough and try again
        //        - size should be buf.len() + buflen (buflen is overwritten)
        panic!("FIXME [windows] unable to get interface list");
    }

    let buf_str = unsafe { from_utf8_unchecked(&buf) };
    let iface_names = buf_str.split("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_owned();
                let next = all_ifaces.iter().filter(|x| name[..].ends_with(&x.name[..])).next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        }
        None => (),
    };

    vec
}
<commit_msg>Add missing CStr import<commit_after>// Copyright (c) 2014-2016 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.

//! Support for sending and receiving data link layer packets using the WinPcap library

extern crate libc;

use std::cmp;
use std::collections::VecDeque;
use std::ffi::{CString, CStr};
use std::str::from_utf8_unchecked;
use std::io;
use std::mem;
use std::slice;
use std::sync::Arc;
use std::net::IpAddr;

use bindings::{bpf, winpcap};
use datalink::{self, NetworkInterface};
use datalink::Channel::Ethernet;
use datalink::{EthernetDataLinkChannelIterator, EthernetDataLinkReceiver, EthernetDataLinkSender};
use packet::Packet;
use packet::ethernet::{EthernetPacket, MutableEthernetPacket};
use util::MacAddr;

struct WinPcapAdapter {
    adapter: winpcap::LPADAPTER,
}

impl Drop for WinPcapAdapter {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketCloseAdapter(self.adapter);
        }
    }
}

struct WinPcapPacket {
    packet: winpcap::LPPACKET,
}

impl Drop for WinPcapPacket {
    fn drop(&mut self) {
        unsafe {
            winpcap::PacketFreePacket(self.packet);
        }
    }
}

/// WinPcap specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
        }
    }
}

/// Create a datalink channel using the WinPcap library
#[inline]
pub fn channel(network_interface: &NetworkInterface, config: &Config)
    -> io::Result<datalink::Channel> {
    let mut read_buffer = Vec::new();
    read_buffer.resize(config.read_buffer_size, 0u8);

    let mut write_buffer = Vec::new();
    write_buffer.resize(config.write_buffer_size, 0u8);

    let adapter = unsafe {
        let net_if_str = CString::new(network_interface.name.as_bytes()).unwrap();
        winpcap::PacketOpenAdapter(net_if_str.as_ptr() as *mut libc::c_char)
    };
    if adapter.is_null() {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetHwFilter(adapter, winpcap::NDIS_PACKET_TYPE_PROMISCUOUS) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Set kernel buffer size
    let ret = unsafe { winpcap::PacketSetBuff(adapter, config.read_buffer_size as libc::c_int) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let ret = unsafe { winpcap::PacketSetReadTimeout(adapter, 0) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    // Immediate mode
    let ret = unsafe { winpcap::PacketSetMinToCopy(adapter, 1) };
    if ret == 0 {
        return Err(io::Error::last_os_error());
    }

    let read_packet = unsafe { winpcap::PacketAllocatePacket() };
    if read_packet.is_null() {
        unsafe {
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(read_packet,
                                  read_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.read_buffer_size as winpcap::UINT)
    }

    let write_packet = unsafe { winpcap::PacketAllocatePacket() };
    if write_packet.is_null() {
        unsafe {
            winpcap::PacketFreePacket(read_packet);
            winpcap::PacketCloseAdapter(adapter);
        }
        return Err(io::Error::last_os_error());
    }

    unsafe {
        winpcap::PacketInitPacket(write_packet,
                                  write_buffer.as_mut_ptr() as winpcap::PVOID,
                                  config.write_buffer_size as winpcap::UINT)
    }

    let adapter = Arc::new(WinPcapAdapter { adapter: adapter });
    let sender = Box::new(DataLinkSenderImpl {
        adapter: adapter.clone(),
        _write_buffer: write_buffer,
        packet: WinPcapPacket { packet: write_packet },
    });
    let receiver = Box::new(DataLinkReceiverImpl {
        adapter: adapter,
        _read_buffer: read_buffer,
        packet: WinPcapPacket { packet: read_packet },
    });
    Ok(Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    adapter: Arc<WinPcapAdapter>,
    _write_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl EthernetDataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(MutableEthernetPacket))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= unsafe { (*self.packet.packet).Length } as usize {
            None
        } else {
            let min = unsafe { cmp::min((*self.packet.packet).Length as usize, len) };
            let slice: &mut [u8] = unsafe {
                slice::from_raw_parts_mut((*self.packet.packet).Buffer as *mut u8, min)
            };
            for chunk in slice.chunks_mut(packet_size) {
                {
                    let eh = MutableEthernetPacket::new(chunk).unwrap();
                    func(eh);
                }

                // Make sure the right length of packet is sent
                let old_len = unsafe { (*self.packet.packet).Length };
                unsafe {
                    (*self.packet.packet).Length = packet_size as u32;
                }

                let ret = unsafe {
                    winpcap::PacketSendPacket(self.adapter.adapter, self.packet.packet, 0)
                };

                unsafe {
                    (*self.packet.packet).Length = old_len;
                }

                if ret == 0 {
                    return Some(Err(io::Error::last_os_error()));
                }
            }
            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| {
                                eh.clone_from(packet);
                            })
    }
}

unsafe impl Send for DataLinkSenderImpl {}
unsafe impl Sync for DataLinkSenderImpl {}

struct DataLinkReceiverImpl {
    adapter: Arc<WinPcapAdapter>,
    _read_buffer: Vec<u8>,
    packet: WinPcapPacket,
}

impl EthernetDataLinkReceiver for DataLinkReceiverImpl {
    fn iter<'a>(&'a mut self) -> Box<EthernetDataLinkChannelIterator + 'a> {
        let buflen = unsafe { (*self.packet.packet).Length } as usize;
        Box::new(DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: VecDeque::with_capacity(buflen / 64),
        })
    }
}

unsafe impl Send for DataLinkReceiverImpl {}
unsafe impl Sync for DataLinkReceiverImpl {}

struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: VecDeque<(usize, usize)>,
}

impl<'a> EthernetDataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a> {
    fn next(&mut self) -> io::Result<EthernetPacket> {
        // NOTE Most of the logic here is identical to FreeBSD/OS X
        if self.packets.is_empty() {
            let ret = unsafe {
                winpcap::PacketReceivePacket(self.pc.adapter.adapter, self.pc.packet.packet, 0)
            };
            let buflen = match ret {
                0 => return Err(io::Error::last_os_error()),
                _ => unsafe { (*self.pc.packet.packet).ulBytesReceived },
            };
            let mut ptr = unsafe { (*self.pc.packet.packet).Buffer };
            let end = unsafe { (*self.pc.packet.packet).Buffer.offset(buflen as isize) };
            while ptr < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                (*self.pc.packet.packet).Buffer as isize;
                    self.packets.push_back((start as usize, (*packet).bh_caplen as usize));
                    let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        let slice = unsafe {
            let data = (*self.pc.packet.packet).Buffer as usize + start;
            slice::from_raw_parts(data as *const u8, len)
        };
        Ok(EthernetPacket::new(slice).unwrap())
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    use bindings::winpcap;

    let mut adapters_size = 0u32;

    unsafe {
        let mut tmp: winpcap::IP_ADAPTER_INFO = mem::zeroed();
        // FIXME [windows] This only gets IPv4 addresses - should use
        // GetAdaptersAddresses
        winpcap::GetAdaptersInfo(&mut tmp, &mut adapters_size);
    }

    let vec_size = adapters_size / mem::size_of::<winpcap::IP_ADAPTER_INFO>() as u32;

    let mut adapters = Vec::with_capacity(vec_size as usize);

    // FIXME [windows] Check return code
    unsafe {
        winpcap::GetAdaptersInfo(adapters.as_mut_ptr(), &mut adapters_size);
    }

    // Create a complete list of NetworkInterfaces for the machine
    let mut cursor = adapters.as_mut_ptr();
    let mut all_ifaces = Vec::with_capacity(vec_size as usize);
    while !cursor.is_null() {
        let mac = unsafe {
            MacAddr((*cursor).Address[0],
                    (*cursor).Address[1],
                    (*cursor).Address[2],
                    (*cursor).Address[3],
                    (*cursor).Address[4],
                    (*cursor).Address[5])
        };
        let mut ip_cursor = unsafe { &mut (*cursor).IpAddressList as winpcap::PIP_ADDR_STRING };
        let mut ips: Vec<IpAddr> = Vec::new();
        while !ip_cursor.is_null() {
            let ip_str_ptr = unsafe { &(*ip_cursor) }.IpAddress.String.as_ptr() as *const i8;
            let bytes = unsafe { CStr::from_ptr(ip_str_ptr).to_bytes() };
            let ip_str = unsafe { from_utf8_unchecked(bytes).to_owned() };
            ips.push(ip_str.parse().unwrap());
            ip_cursor = unsafe { (*ip_cursor).Next };
        }

        unsafe {
            let name_str_ptr = (*cursor).AdapterName.as_ptr() as *const i8;

            let bytes = CStr::from_ptr(name_str_ptr).to_bytes();
            let name_str = from_utf8_unchecked(bytes).to_owned();

            all_ifaces.push(NetworkInterface {
                name: name_str,
                index: (*cursor).Index,
                mac: Some(mac),
                ips: Some(ips),
                // flags: (*cursor).Type, // FIXME [windows]
                flags: 0,
            });

            cursor = (*cursor).Next;
        }
    }

    let mut buf = [0u8; 4096];
    let mut buflen = buf.len() as u32;

    // Gets list of supported adapters in form:
    // adapter1\0adapter2\0\0desc1\0desc2\0\0
    if unsafe { winpcap::PacketGetAdapterNames(buf.as_mut_ptr() as *mut i8, &mut buflen) } == 0 {
        // FIXME [windows] Should allocate a buffer big enough and try again
        //        - size should be buf.len() + buflen (buflen is overwritten)
        panic!("FIXME [windows] unable to get interface list");
    }

    let buf_str = unsafe { from_utf8_unchecked(&buf) };
    let iface_names = buf_str.split("\0\0").next();
    let mut vec = Vec::new();

    // Return only supported adapters
    match iface_names {
        Some(iface_names) => {
            for iface in iface_names.split('\0') {
                let name = iface.to_owned();
                let next = all_ifaces.iter().filter(|x| name[..].ends_with(&x.name[..])).next();
                if next.is_some() {
                    let mut iface = next.unwrap().clone();
                    iface.name = name;
                    vec.push(iface);
                }
            }
        }
        None => (),
    };

    vec
}
<|endoftext|>


<commit_before>//! This module is UNSTABLE, the reason is - recently added.

use std::convert::From;
use std::io::{Cursor, Read};
use std::str::{from_utf8, Utf8Error};

use super::{read_marker};
use super::{
    ReadError,
    MarkerReadError,
};
use super::{BigEndianRead};

use super::super::init::Marker;
use super::super::value::{Float, Integer, ValueRef};

trait USizeCast {
    fn from(v: Self) -> Option<usize> where Self: Sized;
}

impl USizeCast for u8 {
    fn from(v: u8) -> Option<usize> {
        // Impossible to panic, since u8 always fits in usize.
        Some(v as usize)
    }
}

impl USizeCast for u16 {
    fn from(v: u16) -> Option<usize> {
        // TODO: This can overflow on 8-bit systems.
        Some(v as usize)
    }
}

impl USizeCast for u32 {
    fn from(v: u32) -> Option<usize> {
        // TODO: This can overflow on 8- and 16-bit systems.
        Some(v as usize)
    }
}

// TODO: Display trait.
#[derive(Debug)]
pub enum Error<'r> {
    /// Failed to read the type marker value.
    InvalidMarkerRead(ReadError),
    /// Failed to read string/array/map size.
    InvalidLengthRead(ReadError),
    /// Failed to read packed non-marker data.
    InvalidDataRead(ReadError),
    /// Failed to cast the length read to machine size.
    InvalidLengthSize,
    /// Failed to interpret a byte slice as a UTF-8 string.
    ///
    /// Contains untouched bytearray with the underlying decoding error.
    InvalidUtf8(&'r [u8], Utf8Error),
    /// Failed to read ext type.
    InvalidExtTypeRead(ReadError),
    /// Using Reserved type found.
    TypeMismatch,
}

impl<'r> From<MarkerReadError> for Error<'r> {
    fn from(err: MarkerReadError) -> Error<'r> {
        Error::InvalidMarkerRead(From::from(err))
    }
}

fn read_num<'a, R, D>(mut rd: &mut R) -> Result<D, Error<'a>>
    where R: BufRead<'a>,
          D: BigEndianRead
{
    D::read(&mut rd).map_err(|err| Error::InvalidDataRead(From::from(err)))
}

// TODO: Rename to `read_len`.
fn read_length<R, D>(rd: &mut R) -> Result<D, ReadError>
    where R: Read,
          D: BigEndianRead
{
    D::read(rd).map_err(From::from)
}

fn read_str<'a, R>(buf: &mut R, len: usize) -> Result<&'a str, Error<'a>>
    where R: BufRead<'a>
{
    let buf = try!(read_bin(buf, len));

    // Try to decode sliced buffer as UTF-8.
    let res = try!(from_utf8(buf).map_err(|err| Error::InvalidUtf8(buf, err)));

    Ok(res)
}

fn read_bin<'a, R>(rd: &mut R, len: usize) -> Result<&'a [u8], Error<'a>>
    where R: BufRead<'a>
{
    let buf = rd.fill_buf();

    if len > buf.len() {
        return Err(Error::InvalidDataRead(ReadError::UnexpectedEOF));
    }

    // Take a slice.
    let buf = &buf[..len];
    rd.consume(len);

    Ok(buf)
}

// Helper function that reads a single byte from the given `Read` and interpret it as an Ext type.
fn read_ext_type<R>(rd: &mut R) -> Result<i8, ReadError>
    where R: Read
{
    i8::read(rd).map_err(From::from)
}

fn read_ext<'a, R>(mut buf: &mut R, len: usize) -> Result<(i8, &'a [u8]), Error<'a>>
    where R: BufRead<'a>
{
    let ty  = try!(read_ext_type(&mut buf).map_err(|err| Error::InvalidExtTypeRead(err)));
    let buf = try!(read_bin(buf, len));

    Ok((ty, buf))
}

#[inline]
fn read_str_value<'a, R, U>(buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_str(buf, len));

    Ok(ValueRef::String(res))
}

#[inline]
fn read_bin_value<'a, R, U>(buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_bin(buf, len));

    Ok(ValueRef::Binary(res))
}

#[inline]
fn read_ext_value<'a, R, U>(mut buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let (ty, buf) = try!(read_ext(buf, len));

    Ok(ValueRef::Ext(ty, buf))
}

#[inline]
fn read_array_value<'a, R, U>(rd: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let vec = try!(read_array(rd, len));

    Ok(ValueRef::Array(vec))
}

#[inline]
fn read_map_value<'a, R, U>(rd: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let map = try!(read_map(rd, len));

    Ok(ValueRef::Map(map))
}

fn read_array<'a, R>(rd: &mut R, len: usize) -> Result<Vec<ValueRef<'a>>, Error<'a>>
    where R: BufRead<'a>
{
    let mut vec = Vec::with_capacity(len);

    for _ in 0..len {
        let val = try!(read_value_ref(rd));

        vec.push(val);
    }

    Ok(vec)
}

fn read_map<'a, R>(rd: &mut R, len: usize) -> Result<Vec<(ValueRef<'a>, ValueRef<'a>)>, Error<'a>>
    where R: BufRead<'a>
{
    let mut vec = Vec::with_capacity(len);

    for _ in 0..len {
        let key = try!(read_value_ref(rd));
        let val = try!(read_value_ref(rd));

        vec.push((key, val));
    }

    Ok(vec)
}

/// A BufRead is a type of Reader which has an internal buffer.
///
/// This magic trait acts like a standard BufRead but unlike the standard this has an explicit
/// internal buffer lifetime, which allows to borrow from underlying buffer while consuming bytes.
pub trait BufRead<'a>: Read {
    /// Returns the buffer contents.
    ///
    /// This function is a lower-level call. It needs to be paired with the consume method to
    /// function properly. When calling this method, none of the contents will be "read" in the
    /// sense that later calling read may return the same contents. As such, consume must be called
    /// with the number of bytes that are consumed from this buffer to ensure that the bytes are
    /// never returned twice.
    ///
    /// An empty buffer returned indicates that the stream has reached EOF.
    fn fill_buf(&self) -> &'a [u8];

    /// Tells this buffer that len bytes have been consumed from the buffer, so they should no
    /// longer be returned in calls to read.
    fn consume(&mut self, len: usize);
}

impl<'a> BufRead<'a> for &'a [u8] {
    fn fill_buf(&self) -> &'a [u8] {
        self
    }

    fn consume(&mut self, len: usize) {
        *self = &(*self)[len..];
    }
}

/// Useful when you want to know how much bytes has been consumed during ValueRef decoding.
impl<'a> BufRead<'a> for Cursor<&'a [u8]> {
    fn fill_buf(&self) -> &'a [u8] {
        use std::cmp;

        let len = cmp::min(self.position(), self.get_ref().len() as u64);
        &self.get_ref()[len as usize..]
    }

    fn consume(&mut self, len: usize) {
        let pos = self.position();
        self.set_position(pos + len as u64);
    }
}

// TODO: Doc, examples.
pub fn read_value_ref<'a, R>(rd: &mut R) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>
{
    let mut rd = rd;

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `rd`.
    let marker = try!(read_marker(rd));

    let val = match marker {
        Marker::Null => ValueRef::Nil,
        Marker::True => ValueRef::Boolean(true),
        Marker::False => ValueRef::Boolean(false),
        Marker::PositiveFixnum(val) => {
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U8 => {
            let val: u8 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U16 => {
            let val: u16 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U32 => {
            let val: u32 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U64 => {
            let val: u64 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val))
        }
        Marker::NegativeFixnum(val) => {
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I8 => {
            let val: i8 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I16 => {
            let val: i16 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I32 => {
            let val: i32 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I64 => {
            let val: i64 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val))
        }
        Marker::F32 => {
            let val: f32 = try!(read_num(rd));
            ValueRef::Float(Float::F32(val))
        }
        Marker::F64 => {
            let val: f64 = try!(read_num(rd));
            ValueRef::Float(Float::F64(val))
        }
        Marker::FixedString(len) => {
            try!(read_str_value(rd, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::FixedArray(len) => {
            try!(read_array_value(rd, len))
        }
        Marker::Array16 => {
            let len: u16 = try!(read_length(&mut rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_array_value(rd, len))
        }
        Marker::Array32 => {
            let len: u32 = try!(read_length(&mut rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_array_value(rd, len))
        }
        Marker::FixedMap(len) => {
            try!(read_map_value(rd, len))
        }
        Marker::Map16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_map_value(rd, len))
        }
        Marker::Map32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_map_value(rd, len))
        }
        Marker::FixExt1 => {
            try!(read_ext_value(rd, 1u8))
        }
        Marker::FixExt2 => {
            try!(read_ext_value(rd, 2u8))
        }
        Marker::FixExt4 => {
            try!(read_ext_value(rd, 4u8))
        }
        Marker::FixExt8 => {
            try!(read_ext_value(rd, 8u8))
        }
        Marker::FixExt16 => {
            try!(read_ext_value(rd, 16u8))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Reserved => return Err(Error::TypeMismatch),
    };

    Ok(val)
}
<commit_msg>ValueRef: minor renaming.<commit_after>//! This module is UNSTABLE, the reason is - recently added.

use std::convert::From;
use std::io::{Cursor, Read};
use std::str::{from_utf8, Utf8Error};

use super::{read_marker};
use super::{
    ReadError,
    MarkerReadError,
};
use super::{BigEndianRead};

use super::super::init::Marker;
use super::super::value::{Float, Integer, ValueRef};

trait ToUnsigned {
    fn from(v: Self) -> Option<usize> where Self: Sized;
}

impl ToUnsigned for u8 {
    fn from(v: u8) -> Option<usize> {
        // Impossible to panic, since u8 always fits in usize.
        Some(v as usize)
    }
}

impl ToUnsigned for u16 {
    fn from(v: u16) -> Option<usize> {
        // TODO: This can overflow on 8-bit systems.
        Some(v as usize)
    }
}

impl ToUnsigned for u32 {
    fn from(v: u32) -> Option<usize> {
        // TODO: This can overflow on 8- and 16-bit systems.
        Some(v as usize)
    }
}

// TODO: Display trait.
#[derive(Debug)]
pub enum Error<'r> {
    /// Failed to read the type marker value.
    InvalidMarkerRead(ReadError),
    /// Failed to read string/array/map size.
    InvalidLengthRead(ReadError),
    /// Failed to read packed non-marker data.
    InvalidDataRead(ReadError),
    /// Failed to cast the length read to machine size.
    InvalidLengthSize,
    /// Failed to interpret a byte slice as a UTF-8 string.
    ///
    /// Contains untouched bytearray with the underlying decoding error.
    InvalidUtf8(&'r [u8], Utf8Error),
    /// Failed to read ext type.
    InvalidExtTypeRead(ReadError),
    /// Using Reserved type found.
    TypeMismatch,
}

impl<'r> From<MarkerReadError> for Error<'r> {
    fn from(err: MarkerReadError) -> Error<'r> {
        Error::InvalidMarkerRead(From::from(err))
    }
}

fn read_num<'a, R, D>(mut rd: &mut R) -> Result<D, Error<'a>>
    where R: BufRead<'a>,
          D: BigEndianRead
{
    D::read(&mut rd).map_err(|err| Error::InvalidDataRead(From::from(err)))
}

// TODO: Rename to `read_len`.
fn read_length<R, D>(rd: &mut R) -> Result<D, ReadError>
    where R: Read,
          D: BigEndianRead
{
    D::read(rd).map_err(From::from)
}

fn read_str<'a, R>(buf: &mut R, len: usize) -> Result<&'a str, Error<'a>>
    where R: BufRead<'a>
{
    let buf = try!(read_bin(buf, len));

    // Try to decode sliced buffer as UTF-8.
    let res = try!(from_utf8(buf).map_err(|err| Error::InvalidUtf8(buf, err)));

    Ok(res)
}

fn read_bin<'a, R>(rd: &mut R, len: usize) -> Result<&'a [u8], Error<'a>>
    where R: BufRead<'a>
{
    let buf = rd.fill_buf();

    if len > buf.len() {
        return Err(Error::InvalidDataRead(ReadError::UnexpectedEOF));
    }

    // Take a slice.
    let buf = &buf[..len];
    rd.consume(len);

    Ok(buf)
}

// Helper function that reads a single byte from the given `Read` and interpret it as an Ext type.
fn read_ext_type<R>(rd: &mut R) -> Result<i8, ReadError>
    where R: Read
{
    i8::read(rd).map_err(From::from)
}

fn read_ext<'a, R>(mut buf: &mut R, len: usize) -> Result<(i8, &'a [u8]), Error<'a>>
    where R: BufRead<'a>
{
    let ty  = try!(read_ext_type(&mut buf).map_err(|err| Error::InvalidExtTypeRead(err)));
    let buf = try!(read_bin(buf, len));

    Ok((ty, buf))
}

#[inline]
fn read_str_value<'a, R, U>(buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: ToUnsigned
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_str(buf, len));

    Ok(ValueRef::String(res))
}

#[inline]
fn read_bin_value<'a, R, U>(buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: ToUnsigned
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_bin(buf, len));

    Ok(ValueRef::Binary(res))
}

#[inline]
fn read_ext_value<'a, R, U>(mut buf: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: ToUnsigned
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let (ty, buf) = try!(read_ext(buf, len));

    Ok(ValueRef::Ext(ty, buf))
}

#[inline]
fn read_array_value<'a, R, U>(rd: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: ToUnsigned
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let vec = try!(read_array(rd, len));

    Ok(ValueRef::Array(vec))
}

#[inline]
fn read_map_value<'a, R, U>(rd: &mut R, len: U) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>,
          U: ToUnsigned
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let map = try!(read_map(rd, len));

    Ok(ValueRef::Map(map))
}

fn read_array<'a, R>(rd: &mut R, len: usize) -> Result<Vec<ValueRef<'a>>, Error<'a>>
    where R: BufRead<'a>
{
    let mut vec = Vec::with_capacity(len);

    for _ in 0..len {
        let val = try!(read_value_ref(rd));

        vec.push(val);
    }

    Ok(vec)
}

fn read_map<'a, R>(rd: &mut R, len: usize) -> Result<Vec<(ValueRef<'a>, ValueRef<'a>)>, Error<'a>>
    where R: BufRead<'a>
{
    let mut vec = Vec::with_capacity(len);

    for _ in 0..len {
        let key = try!(read_value_ref(rd));
        let val = try!(read_value_ref(rd));

        vec.push((key, val));
    }

    Ok(vec)
}

/// A BufRead is a type of Reader which has an internal buffer.
///
/// This magic trait acts like a standard BufRead but unlike the standard this has an explicit
/// internal buffer lifetime, which allows to borrow from underlying buffer while consuming bytes.
pub trait BufRead<'a>: Read {
    /// Returns the buffer contents.
    ///
    /// This function is a lower-level call. It needs to be paired with the consume method to
    /// function properly. When calling this method, none of the contents will be "read" in the
    /// sense that later calling read may return the same contents. As such, consume must be called
    /// with the number of bytes that are consumed from this buffer to ensure that the bytes are
    /// never returned twice.
    ///
    /// An empty buffer returned indicates that the stream has reached EOF.
    fn fill_buf(&self) -> &'a [u8];

    /// Tells this buffer that len bytes have been consumed from the buffer, so they should no
    /// longer be returned in calls to read.
    fn consume(&mut self, len: usize);
}

impl<'a> BufRead<'a> for &'a [u8] {
    fn fill_buf(&self) -> &'a [u8] {
        self
    }

    fn consume(&mut self, len: usize) {
        *self = &(*self)[len..];
    }
}

/// Useful when you want to know how much bytes has been consumed during ValueRef decoding.
impl<'a> BufRead<'a> for Cursor<&'a [u8]> {
    fn fill_buf(&self) -> &'a [u8] {
        use std::cmp;

        let len = cmp::min(self.position(), self.get_ref().len() as u64);
        &self.get_ref()[len as usize..]
    }

    fn consume(&mut self, len: usize) {
        let pos = self.position();
        self.set_position(pos + len as u64);
    }
}

// TODO: Doc, examples.
pub fn read_value_ref<'a, R>(rd: &mut R) -> Result<ValueRef<'a>, Error<'a>>
    where R: BufRead<'a>
{
    let mut rd = rd;

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `rd`.
    let marker = try!(read_marker(rd));

    let val = match marker {
        Marker::Null => ValueRef::Nil,
        Marker::True => ValueRef::Boolean(true),
        Marker::False => ValueRef::Boolean(false),
        Marker::PositiveFixnum(val) => {
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U8 => {
            let val: u8 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U16 => {
            let val: u16 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U32 => {
            let val: u32 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val as u64))
        }
        Marker::U64 => {
            let val: u64 = try!(read_num(rd));
            ValueRef::Integer(Integer::U64(val))
        }
        Marker::NegativeFixnum(val) => {
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I8 => {
            let val: i8 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I16 => {
            let val: i16 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I32 => {
            let val: i32 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val as i64))
        }
        Marker::I64 => {
            let val: i64 = try!(read_num(rd));
            ValueRef::Integer(Integer::I64(val))
        }
        Marker::F32 => {
            let val: f32 = try!(read_num(rd));
            ValueRef::Float(Float::F32(val))
        }
        Marker::F64 => {
            let val: f64 = try!(read_num(rd));
            ValueRef::Float(Float::F64(val))
        }
        Marker::FixedString(len) => {
            try!(read_str_value(rd, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(rd, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(rd, len))
        }
        Marker::FixedArray(len) => {
            try!(read_array_value(rd, len))
        }
        Marker::Array16 => {
            let len: u16 = try!(read_length(&mut rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_array_value(rd, len))
        }
        Marker::Array32 => {
            let len: u32 = try!(read_length(&mut rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_array_value(rd, len))
        }
        Marker::FixedMap(len) => {
            try!(read_map_value(rd, len))
        }
        Marker::Map16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_map_value(rd, len))
        }
        Marker::Map32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_map_value(rd, len))
        }
        Marker::FixExt1 => {
            try!(read_ext_value(rd, 1u8))
        }
        Marker::FixExt2 => {
            try!(read_ext_value(rd, 2u8))
        }
        Marker::FixExt4 => {
            try!(read_ext_value(rd, 4u8))
        }
        Marker::FixExt8 => {
            try!(read_ext_value(rd, 8u8))
        }
        Marker::FixExt16 => {
            try!(read_ext_value(rd, 16u8))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(rd).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(rd, len))
        }
        Marker::Reserved => return Err(Error::TypeMismatch),
    };

    Ok(val)
}
<|endoftext|>


<commit_before>use crate::atomicmin::AtomicMin;
use crate::{PngError, PngResult};
use libdeflater::{CompressionError, CompressionLvl, Compressor};

pub fn deflate(data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
    let mut compressor = Compressor::new(CompressionLvl::best());
    let capacity = max_size.get().unwrap_or(data.len() / 2);
    let mut dest = Vec::with_capacity(capacity);
    unsafe {
        // This is ok because the Vec contains Copy-able data (u8)
        // and because libdeflater wrapper doesn't try to read
        // the bytes from the target.
        //
        // That said, it should be able to accept MaybeUninit instead,
        // so I raised an upstream issue that should make this safer:
        // https://github.com/adamkewley/libdeflater/issues/1
        dest.set_len(capacity);
    }
    let len = compressor
        .zlib_compress(data, &mut dest)
        .map_err(|err| match err {
            CompressionError::InsufficientSpace => PngError::DeflatedDataTooLong(capacity),
        })?;
    if let Some(max) = max_size.get() {
        if len > max {
            return Err(PngError::DeflatedDataTooLong(max));
        }
    }
    dest.truncate(len);
    Ok(dest)
}
<commit_msg>Fix a hard error in clippy<commit_after>use crate::atomicmin::AtomicMin;
use crate::{PngError, PngResult};
use libdeflater::{CompressionError, CompressionLvl, Compressor};

pub fn deflate(data: &[u8], max_size: &AtomicMin) -> PngResult<Vec<u8>> {
    let mut compressor = Compressor::new(CompressionLvl::best());
    let capacity = max_size.get().unwrap_or(data.len() / 2);
    let mut dest = vec![0; capacity];
    let len = compressor
        .zlib_compress(data, &mut dest)
        .map_err(|err| match err {
            CompressionError::InsufficientSpace => PngError::DeflatedDataTooLong(capacity),
        })?;
    if let Some(max) = max_size.get() {
        if len > max {
            return Err(PngError::DeflatedDataTooLong(max));
        }
    }
    dest.truncate(len);
    Ok(dest)
}
<|endoftext|>


<commit_before>#[legacy_exports];
export FreeTypeNativeFont, with_test_native_font, create;

use util::*;
use vec_as_buf = vec::as_imm_buf;
use ptr::{addr_of, null};
use cast::reinterpret_cast;
use glyph::GlyphIndex;
use font::FontMetrics;
use azure::freetype;
use freetype::{ FT_Error, FT_Library, FT_Face, FT_Long, FT_ULong, FT_UInt, FT_GlyphSlot };
use freetype::bindgen::{
    FT_Init_FreeType,
    FT_Done_FreeType,
    FT_New_Memory_Face,
    FT_Done_Face,
    FT_Get_Char_Index,
    FT_Load_Glyph,
    FT_Set_Char_Size
};

fn float_to_fixed_ft(f: float) -> i32 {
    float_to_fixed(26, f)
}

fn fixed_to_float_ft(f: i32) -> float {
    fixed_to_float(26, f)
}

struct FreeTypeNativeFont {
    /// The font binary. This must stay valid for the lifetime of the font
    buf: @~[u8],
    face: FT_Face,

    drop {
        assert self.face.is_not_null();
        if !FT_Done_Face(self.face).succeeded() {
            fail ~"FT_Done_Face failed";
        }
    }
}

fn FreeTypeNativeFont(face: FT_Face, buf: @~[u8]) -> FreeTypeNativeFont {
    assert face.is_not_null();
    FreeTypeNativeFont { buf: buf, face: face }
}

impl FreeTypeNativeFont {

    fn glyph_index(codepoint: char) -> Option<GlyphIndex> {
        assert self.face.is_not_null();
        let idx = FT_Get_Char_Index(self.face, codepoint as FT_ULong);
        return if idx != 0 as FT_UInt {
            Some(idx as GlyphIndex)
        } else {
            debug!("Invalid codepoint: %?", codepoint);
            None
        };
    }

    // FIXME: What unit is this returning? Let's have a custom type
    fn glyph_h_advance(glyph: GlyphIndex) -> Option<FractionalPixel> {
        assert self.face.is_not_null();
        let res =  FT_Load_Glyph(self.face, glyph as FT_UInt, 0);
        if res.succeeded() {
            unsafe {
                let void_glyph = (*self.face).glyph;
                let slot: FT_GlyphSlot = reinterpret_cast(&void_glyph);
                assert slot.is_not_null();
                let advance = (*slot).metrics.horiAdvance;
                debug!("h_advance for %? is %?", glyph, advance);
                return Some(fixed_to_float_ft(advance) as FractionalPixel);
            }
        } else {
            debug!("Unable to load glyph %?. reason: %?", glyph, res);
            return None;
        }
    }

    fn get_metrics() -> FontMetrics {
        /* TODO: complete me (Issue #76) */
        return FontMetrics {
            underline_size:   0.0,
            underline_offset: 0.0,
            leading:          0.0,
            x_height:         0.0,
            em_ascent:        0.0,
            em_descent:       0.0,
            em_height:        0.0,
            max_advance:      0.0
        }
    }
}

pub fn create(lib: &FT_Library, buf: @~[u8]) -> Result<FreeTypeNativeFont, ()> {
    assert lib.is_not_null();
    let face: FT_Face = null();
    return vec_as_buf(*buf, |cbuf, _len| {
           if FT_New_Memory_Face(*lib, cbuf, (*buf).len() as FT_Long,
                                 0 as FT_Long, addr_of(&face)).succeeded() {
               // FIXME: These values are placeholders
               let res = FT_Set_Char_Size(face, 0, 20*64, 0, 72);
               if !res.succeeded() { fail ~"unable to set font char size" }
               Ok(FreeTypeNativeFont(face, buf))
           } else {
               Err(())
           }
    })
}

trait FTErrorMethods {
    fn succeeded() -> bool;
}

impl FT_Error : FTErrorMethods {
    fn succeeded() -> bool { self == 0 as FT_Error }
}

fn with_test_native_font(f: fn@(nf: &NativeFont)) {
    use font::test_font_bin;
    use unwrap_result = result::unwrap;

    with_lib(|lib| {
        let buf = @test_font_bin();
        let font = unwrap_result(create(lib, move buf));
        f(&font);
    })
}

fn with_lib(f: fn@((&FT_Library))) {
    let lib: FT_Library = null();
    assert FT_Init_FreeType(addr_of(&lib)).succeeded();
    f(&lib);
    FT_Done_FreeType(lib);
}

#[test]
fn create_should_return_err_if_buf_is_bogus() {
    with_lib(|lib| {
        let buf = @~[];
        assert create(lib, buf).is_err();
    })
}
<commit_msg>Make freetype code build<commit_after>#[legacy_exports];
export FreeTypeNativeFont, with_test_native_font, create;

use font::{FontMetrics, FractionalPixel};

use util::*;
use vec_as_buf = vec::as_imm_buf;
use ptr::{addr_of, null};
use cast::reinterpret_cast;
use glyph::GlyphIndex;
use font::FontMetrics;
use azure::freetype;
use freetype::{ FT_Error, FT_Library, FT_Face, FT_Long, FT_ULong, FT_UInt, FT_GlyphSlot };
use freetype::bindgen::{
    FT_Init_FreeType,
    FT_Done_FreeType,
    FT_New_Memory_Face,
    FT_Done_Face,
    FT_Get_Char_Index,
    FT_Load_Glyph,
    FT_Set_Char_Size
};

fn float_to_fixed_ft(f: float) -> i32 {
    float_to_fixed(26, f)
}

fn fixed_to_float_ft(f: i32) -> float {
    fixed_to_float(26, f)
}

struct FreeTypeNativeFont {
    /// The font binary. This must stay valid for the lifetime of the font
    buf: @~[u8],
    face: FT_Face,

    drop {
        assert self.face.is_not_null();
        if !FT_Done_Face(self.face).succeeded() {
            fail ~"FT_Done_Face failed";
        }
    }
}

fn FreeTypeNativeFont(face: FT_Face, buf: @~[u8]) -> FreeTypeNativeFont {
    assert face.is_not_null();
    FreeTypeNativeFont { buf: buf, face: face }
}

impl FreeTypeNativeFont {

    fn glyph_index(codepoint: char) -> Option<GlyphIndex> {
        assert self.face.is_not_null();
        let idx = FT_Get_Char_Index(self.face, codepoint as FT_ULong);
        return if idx != 0 as FT_UInt {
            Some(idx as GlyphIndex)
        } else {
            debug!("Invalid codepoint: %?", codepoint);
            None
        };
    }

    // FIXME: What unit is this returning? Let's have a custom type
    fn glyph_h_advance(glyph: GlyphIndex) -> Option<FractionalPixel> {
        assert self.face.is_not_null();
        let res =  FT_Load_Glyph(self.face, glyph as FT_UInt, 0);
        if res.succeeded() {
            unsafe {
                let void_glyph = (*self.face).glyph;
                let slot: FT_GlyphSlot = reinterpret_cast(&void_glyph);
                assert slot.is_not_null();
                let advance = (*slot).metrics.horiAdvance;
                debug!("h_advance for %? is %?", glyph, advance);
                let advance = advance as i32;
                return Some(fixed_to_float_ft(advance) as FractionalPixel);
            }
        } else {
            debug!("Unable to load glyph %?. reason: %?", glyph, res);
            return None;
        }
    }

    fn get_metrics() -> FontMetrics {
        /* TODO: complete me (Issue #76) */
        return FontMetrics {
            underline_size:   0.0,
            underline_offset: 0.0,
            leading:          0.0,
            x_height:         0.0,
            em_size:          gfx::geometry::au(0),
            em_height:        0.0,
            em_ascent:        0.0,
            em_descent:       0.0,
            max_advance:      0.0
        }
    }
}

pub fn create(lib: &FT_Library, buf: @~[u8]) -> Result<FreeTypeNativeFont, ()> {
    assert lib.is_not_null();
    let face: FT_Face = null();
    return vec_as_buf(*buf, |cbuf, _len| {
           if FT_New_Memory_Face(*lib, cbuf, (*buf).len() as FT_Long,
                                 0 as FT_Long, addr_of(&face)).succeeded() {
               // FIXME: These values are placeholders
               let res = FT_Set_Char_Size(face, 0, 20*64, 0, 72);
               if !res.succeeded() { fail ~"unable to set font char size" }
               Ok(FreeTypeNativeFont(face, buf))
           } else {
               Err(())
           }
    })
}

trait FTErrorMethods {
    fn succeeded() -> bool;
}

impl FT_Error : FTErrorMethods {
    fn succeeded() -> bool { self == 0 as FT_Error }
}

fn with_test_native_font(f: fn@(nf: &NativeFont)) {
    use font::test_font_bin;
    use unwrap_result = result::unwrap;

    with_lib(|lib| {
        let buf = @test_font_bin();
        let font = unwrap_result(create(lib, move buf));
        f(&font);
    })
}

fn with_lib(f: fn@((&FT_Library))) {
    let lib: FT_Library = null();
    assert FT_Init_FreeType(addr_of(&lib)).succeeded();
    f(&lib);
    FT_Done_FreeType(lib);
}

#[test]
fn create_should_return_err_if_buf_is_bogus() {
    with_lib(|lib| {
        let buf = @~[];
        assert create(lib, buf).is_err();
    })
}
<|endoftext|>


<commit_before>/// Immediate IO model: performs IO in the current thread.
///
/// Use for diagnosing bugs or working around any unexpected issues with the
/// threaded code paths.
use super::{perform, Executor, Item};

pub struct ImmediateUnpacker {}
impl ImmediateUnpacker {
    pub fn new() -> ImmediateUnpacker {
        ImmediateUnpacker {}
    }
}

impl Executor for ImmediateUnpacker {
    fn dispatch(&mut self, mut item: Item) -> Box<dyn '_ + Iterator<Item = Item>> {
        perform(&mut item);
        Box::new(Some(item).into_iter())
    }

    fn join(&mut self) -> Box<dyn Iterator<Item = Item>> {
        Box::new(None.into_iter())
    }

    fn completed(&mut self) -> Box<dyn Iterator<Item = Item>> {
        Box::new(None.into_iter())
    }
}
<commit_msg>Fix clippy: new_without_default<commit_after>/// Immediate IO model: performs IO in the current thread.
///
/// Use for diagnosing bugs or working around any unexpected issues with the
/// threaded code paths.
use super::{perform, Executor, Item};

#[derive(Default)]
pub struct ImmediateUnpacker {}
impl ImmediateUnpacker {
    pub fn new() -> ImmediateUnpacker {
        ImmediateUnpacker {}
    }
}

impl Executor for ImmediateUnpacker {
    fn dispatch(&mut self, mut item: Item) -> Box<dyn '_ + Iterator<Item = Item>> {
        perform(&mut item);
        Box::new(Some(item).into_iter())
    }

    fn join(&mut self) -> Box<dyn Iterator<Item = Item>> {
        Box::new(None.into_iter())
    }

    fn completed(&mut self) -> Box<dyn Iterator<Item = Item>> {
        Box::new(None.into_iter())
    }
}
<|endoftext|>


<commit_before>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        // .with_dtypes_slice(Some(&[DataType::Utf8, DataType::Utf8, DataType::Float64, DataType::Float64]))
        .finish()
        .unwrap_or_else(|x| panic!("Failed on {:?} with error {:?}", path, x));

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<Expr>::new();

    for column in input.get_column_names_owned() {
        if column.starts_with('g') {
            input.apply(&column, baseline_get_median_gt_zero).unwrap();
            meas.push(col(&column));
        }
    }
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(
            any_exprs(&meas[0..meas.len() / 2]).or(any_exprs(&meas[meas.len() / 2..meas.len()])),
        ) // If we try to do the whole thing at once, we get a stack overflow
        .collect()
        .unwrap()
}

fn lowercase_fn(ent: &Option<&str>) -> Option<String> {
    match ent {
        None => None,
        Some(ent) => Some(ent.to_lowercase()),
    }
}

fn fix_bad_infinities(str_val: &Series) -> Series {
    let lowercased = str_val
        .utf8()
        .unwrap()
        .into_iter()
        .map(|x| lowercase_fn(&x))
        .collect::<Vec<Option<String>>>()
        .into_iter()
        .map(|x| match x { None => None, Some(x) => Some(x.parse::<f64>().unwrap()) } )
        .collect::<Vec<Option<f64>>>();

    Series::from_iter(lowercased)

}

fn filter_differential(input: &mut DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            // Check for badly parsed infinities
            if !input.column(&column).unwrap().dtype().is_numeric() {
                // input.column(&column).unwrap().utf8().into_iter().map(|y| y.map(|x| x.to_lowercase()) ).collect();
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            if !input.column(&column).unwrap().dtype().is_numeric() {
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            lf = col(&column).neq(lit(0.0));
            exprs.push(lf);
        }
    }
    input
        .clone()
        .lazy()
        .filter(all_exprs(exprs))
        .collect()
        .unwrap_or_else(|x| panic!("Error with dataframe: \n {:?}\nAnd error: {:?}", input, x))
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        println!("{:?}", infile);
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<commit_msg>Ensure unique gene IDs<commit_after>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        // .with_dtypes_slice(Some(&[DataType::Utf8, DataType::Utf8, DataType::Float64, DataType::Float64]))
        .finish()
        .unwrap_or_else(|x| panic!("Failed on {:?} with error {:?}", path, x));

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<Expr>::new();

    for column in input.get_column_names_owned() {
        if column.starts_with('g') {
            input.apply(&column, baseline_get_median_gt_zero).unwrap();
            meas.push(col(&column));
        }
    }
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(
            any_exprs(&meas[0..meas.len() / 2]).or(any_exprs(&meas[meas.len() / 2..meas.len()])),
        ) // If we try to do the whole thing at once, we get a stack overflow
        .collect()
        .unwrap()
}

fn lowercase_fn(ent: &Option<&str>) -> Option<String> {
    match ent {
        None => None,
        Some(ent) => Some(ent.to_lowercase()),
    }
}

fn fix_bad_infinities(str_val: &Series) -> Series {
    let lowercased = str_val
        .utf8()
        .unwrap()
        .into_iter()
        .map(|x| lowercase_fn(&x))
        .collect::<Vec<Option<String>>>()
        .into_iter()
        .map(|x| match x { None => None, Some(x) => Some(x.parse::<f64>().unwrap()) } )
        .collect::<Vec<Option<f64>>>();

    Series::from_iter(lowercased)

}

fn filter_differential(input: &mut DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            // Check for badly parsed infinities
            if !input.column(&column).unwrap().dtype().is_numeric() {
                // input.column(&column).unwrap().utf8().into_iter().map(|y| y.map(|x| x.to_lowercase()) ).collect();
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            if !input.column(&column).unwrap().dtype().is_numeric() {
                input.apply(&column, fix_bad_infinities);
                // input.column(&column).unwrap().cast(&DataType::Float64);
            }
            lf = col(&column).neq(lit(0.0));
            exprs.push(lf);
        }
    }
    input
        .clone()
        .lazy()
        .filter(all_exprs(exprs))
        .collect()
        .unwrap_or_else(|x| panic!("Error with dataframe: \n {:?}\nAnd error: {:?}", input, x))
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        println!("{:?}", infile);
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    output = output.unique(Some(&["GeneID".to_string()]), UniqueKeepStrategy::First).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<|endoftext|>