<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, IntoFuture, Sink, Stream};
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::{Action, Retry};
use tokio_retry::strategy::{jitter, FixedInterval};

use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

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

use super::EventHandler;
use super::error::{forget_result, into_other, log_error, other_error, result_ok};
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>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }
}

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 network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        // Requests handler
        let mut core = Core::new()?;
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let network_tx = network_tx.clone();
            let cancel_sender = cancel_sender.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(peer) {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, &conn_tx);
                        // Enable retry feature for outgoing connection.
                        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);
                                outgoing_connections.clone().remove(&peer);
                                network_tx
                                    .clone()
                                    .send(NetworkEvent::PeerDisconnected(peer))
                                    .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);
                    Self::send_event(&handle, &network_tx, NetworkEvent::PeerDisconnected(peer));
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => cancel_sender
                    .clone()
                    .send(())
                    .map(forget_result)
                    .map_err(log_error)
                    .wait()?,
            }

            Ok(())
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // 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)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    return Ok(());
                }
                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(|_| {
                        // Ensure that holder lives until the stream ends.
                        let _holder = holder;
                    })
                    .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")),
        )
    }

    fn send_event(handle: &Handle, network_tx: &mpsc::Sender<NetworkEvent>, event: NetworkEvent) {
        handle.spawn(
            network_tx
                .clone()
                .send(event)
                .map(forget_result)
                .map_err(log_error),
        );
    }
}

// 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>Fix clippy warning<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, IntoFuture, Sink, Stream};
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::{Action, Retry};
use tokio_retry::strategy::{jitter, FixedInterval};

use std::io;
use std::net::SocketAddr;
use std::time::Duration;
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;

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

use super::EventHandler;
use super::error::{forget_result, into_other, log_error, other_error, result_ok};
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>,
}

#[derive(Debug, Default, Clone)]
struct ConnectionsPool {
    inner: Rc<RefCell<HashMap<SocketAddr, mpsc::Sender<RawMessage>>>>,
}

impl ConnectionsPool {
    fn new() -> ConnectionsPool {
        ConnectionsPool::default()
    }

    fn insert(&self, peer: SocketAddr, sender: &mpsc::Sender<RawMessage>) {
        self.inner.borrow_mut().insert(peer, sender.clone());
    }

    fn remove(&self, peer: &SocketAddr) {
        self.inner.borrow_mut().remove(peer);
    }

    fn get(&self, peer: SocketAddr) -> Option<mpsc::Sender<RawMessage>> {
        self.inner.borrow_mut().get(&peer).cloned()
    }
}

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)
    }
}

impl NetworkPart {
    pub fn run(self) -> Result<(), io::Error> {
        let network_config = self.network_config;
        // Cancelation token
        let (cancel_sender, cancel_handler) = unsync::mpsc::channel(1);
        // Outgoing connections
        let outgoing_connections = ConnectionsPool::new();
        // Requests handler
        let mut core = Core::new()?;
        let handle = core.handle();
        let network_tx = self.network_tx.clone();
        let requests_handle = self.network_requests.1.for_each(move |request| {
            let network_tx = network_tx.clone();
            let cancel_sender = cancel_sender.clone();
            let outgoing_connections = outgoing_connections.clone();
            match request {
                NetworkRequest::SendMessage(peer, msg) => {
                    let conn_tx = if let Some(conn_tx) = outgoing_connections.get(peer) {
                        conn_tx
                    } else {
                        // Register outgoing channel.
                        let (conn_tx, conn_rx) = mpsc::channel(10);
                        outgoing_connections.insert(peer, &conn_tx);
                        // Enable retry feature for outgoing connection.
                        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)
                            // Configure socket
                            .and_then(move |sock| {
                                sock.set_nodelay(network_config.tcp_nodelay)?;
                                Ok(sock)
                            })
                            .and_then(move |sock| {
                                let duration =
                                    network_config.tcp_keep_alive.map(Duration::from_millis);
                                sock.set_keepalive(duration)?;
                                Ok(sock)
                            })
                            // Connect socket with the outgoing channel
                            .and_then(move |sock| {
                                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(),
                                    })
                            })
                            .then(move |res| {
                                trace!("Connection with peer={} closed, reason={:?}", peer, res);
                                outgoing_connections.clone().remove(&peer);
                                network_tx
                                    .clone()
                                    .send(NetworkEvent::PeerDisconnected(peer))
                                    .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);
                    Self::send_event(&handle, &network_tx, NetworkEvent::PeerDisconnected(peer));
                }
                // Immediately stop the event loop.
                NetworkRequest::Shutdown => {
                    cancel_sender
                        .clone()
                        .send(())
                        .map(forget_result)
                        .map_err(log_error)
                        .wait()?
                }
            }

            Ok(())
        });

        // Incoming connections limiter
        let incoming_connections_limit = network_config.max_incoming_connections;
        let incoming_connections_counter: Rc<()> = Rc::default();
        // 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)| {
                // Increment reference counter
                let holder = Rc::downgrade(&incoming_connections_counter);
                // Check incoming connections count
                let connections_count = Rc::weak_count(&incoming_connections_counter);
                if connections_count > incoming_connections_limit {
                    warn!(
                        "Rejected incoming connection with peer={}, \
                         connections limit reached.",
                        addr
                    );
                    return Ok(());
                }
                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(|_| {
                        // Ensure that holder lives until the stream ends.
                        let _holder = holder;
                    })
                    .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")),
        )
    }

    fn send_event(handle: &Handle, network_tx: &mpsc::Sender<NetworkEvent>, event: NetworkEvent) {
        handle.spawn(network_tx.clone().send(event).map(forget_result).map_err(
            log_error,
        ));
    }
}

// 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>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }
}
<commit_msg>Add tests for ELF note iterators<commit_after>// Defined note types for GNU systems.

// ABI information.  The descriptor consists of words:
//    word 0: OS descriptor
//    word 1: major version of the ABI
//    word 2: minor version of the ABI
//    word 3: subminor version of the ABI

pub const NT_GNU_ABI_TAG: u32 = 1;
// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
//    NT_GNU_ABI_TAG note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

// Synthetic hwcap information.  The descriptor begins with two words:
//    word 0: number of entries
//    word 1: bitmask of enabled entries
//    Then follow variable-length entries, one byte followed by a
//    '\0'-terminated hwcap name string.  The byte gives the bit
//    number to test if enabled, (1U << bit) & bitmask.
pub const NT_GNU_HWCAP: u32 = 2;

// Build ID bits as generated by ld --build-id.
//    The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "std", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_std! {
    use error;
    use container;
    use scroll::{ctx, Pread};

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e.into()))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        type Size = usize;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, Self::Size), Self::Error> {
            let offset = &mut 0;
            let header: NoteHeader = {
                match alignment {
                    4 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    // this is a guess; i haven't seen gcc/clang compilers emit 64-bit notes, and i don't have any non gcc/clang compilers
                    8 => bytes.gread_with::<Nhdr64>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        const NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        const CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<|endoftext|>


<commit_before>pub struct Server{
    hostname: &'static str,
    port: u16
}

impl Server {
    pub fn hostname(&self) -> &'static str {
        self.hostname
    }

    pub fn port(&self) -> u16 {
        self.port
    }
}

pub struct User {
    nick: &'static str,
    ident: &'static str,
    hostname: &'static str,
    realname: &'static str
}

impl User {
    pub fn nick(&self) -> &'static str {
        self.nick
    }

    pub fn ident(&self) -> &'static str {
        self.ident
    }

    pub fn host(&self) -> &'static str {
        self.hostname
    }

    pub fn realname(&self) -> &'static str {
        self.realname
    }
}

pub struct Channel {
    name: &'static str,
    usercount: u16
}

impl Channel {
    pub fn usercount(&self) -> u16 {
        self.usercount
    }

    pub fn name(&self) -> &'static str {
        self.name
    }
}
<commit_msg>Basic constructors<commit_after>pub struct Server{
    hostname: &'static str,
    port: u16
}

impl Server {

    pub fn new(hostname: &'static str, port: u16) -> Server {
        Server { hostname:hostname, port:port }
    }

    pub fn hostname(&self) -> &'static str {
        self.hostname
    }

    pub fn port(&self) -> u16 {
        self.port
    }
}

pub struct User {
    nick: &'static str,
    ident: &'static str,
    hostname: &'static str,
    realname: &'static str
}

impl User {

    pub fn new(nick: &'static str, ident: &'static str, hostname: &'static str, realname: &'static str) -> User {
        User { nick:nick, ident:ident, hostname:hostname, realname:realname }
    }

    pub fn nick(&self) -> &'static str {
        self.nick
    }

    pub fn ident(&self) -> &'static str {
        self.ident
    }

    pub fn host(&self) -> &'static str {
        self.hostname
    }

    pub fn realname(&self) -> &'static str {
        self.realname
    }
}

pub struct Channel {
    name: &'static str,
    usercount: u16
}

impl Channel {
    pub fn new(name: &'static str) -> Channel {
        Channel { name:name, usercount: 0u16 }
    }

    pub fn usercount(&self) -> u16 {
        self.usercount
    }

    pub fn name(&self) -> &'static str {
        self.name
    }
}
<|endoftext|>


<commit_before>use serde::Deserialize;
use std::fmt::{self, Display, Formatter};
use std::fs::{remove_file};
use std::path::{PathBuf};
use std::process::{self, Command, Output};

const RUSTC_COLOR_ARGS: &[&str] = &["--color", "always"];

fn temp_file() -> String {
    format!("./temp_{}", process::id())
}

#[derive(Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Mode {
    Compile,
    Test,
}

#[derive(Deserialize)]
pub struct ExerciseList {
    pub exercises: Vec<Exercise>,
}

#[derive(Deserialize)]
pub struct Exercise {
    pub path: PathBuf,
    pub mode: Mode,
}

impl Exercise {
    pub fn compile(&self) -> Output {
        match self.mode {
            Mode::Compile => Command::new("rustc")
                .args(&[self.path.to_str().unwrap(), "-o", &temp_file()])
                .args(RUSTC_COLOR_ARGS)
                .output(),
            Mode::Test => Command::new("rustc")
                .args(&["--test", self.path.to_str().unwrap(), "-o", &temp_file()])
                .args(RUSTC_COLOR_ARGS)
                .args(&["--", "--nocapture"])
                .output(),
        }
        .expect("Failed to run 'compile' command.")
    }

    pub fn run(&self) -> Output {
        Command::new(&temp_file())
            .output()
            .expect("Failed to run 'run' command")
    }

    pub fn clean(&self) {
        let _ignored = remove_file(&temp_file());
    }
}

impl Display for Exercise {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}", self.path.to_str().unwrap())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::path::Path;
    use std::fs::File;

    #[test]
    fn test_clean() {
        File::create(&temp_file()).unwrap();
        let exercise = Exercise {
            path: PathBuf::from("example.rs"),
            mode: Mode::Test,
        };
        exercise.clean();
        assert!(!Path::new(&temp_file()).exists());
    }
}
<commit_msg>fix the --nocapture functionality<commit_after>use serde::Deserialize;
use std::fmt::{self, Display, Formatter};
use std::fs::{remove_file};
use std::path::{PathBuf};
use std::process::{self, Command, Output};

const RUSTC_COLOR_ARGS: &[&str] = &["--color", "always"];

fn temp_file() -> String {
    format!("./temp_{}", process::id())
}

#[derive(Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Mode {
    Compile,
    Test,
}

#[derive(Deserialize)]
pub struct ExerciseList {
    pub exercises: Vec<Exercise>,
}

#[derive(Deserialize)]
pub struct Exercise {
    pub path: PathBuf,
    pub mode: Mode,
}

impl Exercise {
    pub fn compile(&self) -> Output {
        match self.mode {
            Mode::Compile => Command::new("rustc")
                .args(&[self.path.to_str().unwrap(), "-o", &temp_file()])
                .args(RUSTC_COLOR_ARGS)
                .output(),
            Mode::Test => Command::new("rustc")
                .args(&["--test", self.path.to_str().unwrap(), "-o", &temp_file()])
                .args(RUSTC_COLOR_ARGS)
                .output(),
        }
        .expect("Failed to run 'compile' command.")
    }

    pub fn run(&self) -> Output {
        Command::new(&temp_file())
            .args(&["--", "--nocapture"])
            .output()
            .expect("Failed to run 'run' command")
    }

    pub fn clean(&self) {
        let _ignored = remove_file(&temp_file());
    }
}

impl Display for Exercise {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}", self.path.to_str().unwrap())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::path::Path;
    use std::fs::File;

    #[test]
    fn test_clean() {
        File::create(&temp_file()).unwrap();
        let exercise = Exercise {
            path: PathBuf::from("example.rs"),
            mode: Mode::Test,
        };
        exercise.clean();
        assert!(!Path::new(&temp_file()).exists());
    }
}
<|endoftext|>


<commit_before>use std::ops::AddAssign;
use std::collections:: { HashMap };
use std::collections::hash_map:: { Iter };
use lcov_parser:: { FunctionData };

type FunctionName = String;
type ExecutionCount = u32;

#[derive(Debug, Clone)]
pub struct Functions {
    functions: HashMap<FunctionName, ExecutionCount>
}

impl Functions {
    pub fn new() -> Self {
        Functions {
            functions: HashMap::new()
        }
    }
    pub fn iter(&self) -> Iter<FunctionName, ExecutionCount> {
        self.functions.iter()
    }
    pub fn get(&self, key: &FunctionName) -> Option<&ExecutionCount> {
        self.functions.get(key)
    }
}

impl<'a> AddAssign<&'a FunctionData> for Functions {
    fn add_assign(&mut self, data: &'a FunctionData) {
        if !self.functions.contains_key(&data.name) {
            self.functions.insert(data.name.clone(), data.count.clone());
            return;
        }
        let mut func_count = self.functions.get_mut(&data.name).unwrap();
        *func_count += data.count;
    }
}

impl<'a> AddAssign<&'a Functions> for Functions {
    fn add_assign(&mut self, other: &'a Functions) {
        for (name, count) in other.iter() {
            if !self.functions.contains_key(name) {
                self.functions.insert(name.clone(), *count);
                continue;
            }
            let current_count = self.functions.get_mut(name).unwrap();
            *current_count += *count;
        }
    }
}

#[cfg(test)]
mod tests {
    use lcov_parser:: { FunctionData };
    use function:: { Functions };

    #[test]
    fn add_function_data() {
        let mut functions = Functions::new();
        functions += &FunctionData { name: "main".to_string(), count: 1 };
        functions += &FunctionData { name: "main".to_string(), count: 1 };

        let result = functions.clone();
        assert_eq!( result.get(&"main".to_string()), Some(&2u32) );
    }

    #[test]
    fn add_lines_data() {
        let mut functions = Functions::new();
        functions += &FunctionData { name: "main".to_string(), count: 1 };

        let ref cloned_functions = functions.clone();
        functions += cloned_functions;

        assert_eq!( functions.get(&"main".to_string()), Some(&2u32) );
    }
}
<commit_msg>rename test function<commit_after>use std::ops::AddAssign;
use std::collections:: { HashMap };
use std::collections::hash_map:: { Iter };
use lcov_parser:: { FunctionData };

type FunctionName = String;
type ExecutionCount = u32;

#[derive(Debug, Clone)]
pub struct Functions {
    functions: HashMap<FunctionName, ExecutionCount>
}

impl Functions {
    pub fn new() -> Self {
        Functions {
            functions: HashMap::new()
        }
    }
    pub fn iter(&self) -> Iter<FunctionName, ExecutionCount> {
        self.functions.iter()
    }
    pub fn get(&self, key: &FunctionName) -> Option<&ExecutionCount> {
        self.functions.get(key)
    }
}

impl<'a> AddAssign<&'a FunctionData> for Functions {
    fn add_assign(&mut self, data: &'a FunctionData) {
        if !self.functions.contains_key(&data.name) {
            self.functions.insert(data.name.clone(), data.count.clone());
            return;
        }
        let mut func_count = self.functions.get_mut(&data.name).unwrap();
        *func_count += data.count;
    }
}

impl<'a> AddAssign<&'a Functions> for Functions {
    fn add_assign(&mut self, other: &'a Functions) {
        for (name, count) in other.iter() {
            if !self.functions.contains_key(name) {
                self.functions.insert(name.clone(), *count);
                continue;
            }
            let current_count = self.functions.get_mut(name).unwrap();
            *current_count += *count;
        }
    }
}

#[cfg(test)]
mod tests {
    use lcov_parser:: { FunctionData };
    use function:: { Functions };

    #[test]
    fn add_function_data() {
        let mut functions = Functions::new();
        functions += &FunctionData { name: "main".to_string(), count: 1 };
        functions += &FunctionData { name: "main".to_string(), count: 1 };

        let result = functions.clone();
        assert_eq!( result.get(&"main".to_string()), Some(&2u32) );
    }

    #[test]
    fn add_function_data() {
        let mut functions = Functions::new();
        functions += &FunctionData { name: "main".to_string(), count: 1 };

        let ref cloned_functions = functions.clone();
        functions += cloned_functions;

        assert_eq!( functions.get(&"main".to_string()), Some(&2u32) );
    }
}
<|endoftext|>


<commit_before>use std::ops::Deref;

use syntax::{ast, abi, codemap};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::parse::token;

use syntax::ext::build::AstBuilder;

use test::Test;
use bench::Bench;

use describe::{DescribeState, SubBlock};

/// Trait meaning something can be turned into an ast::Item with configuration.
pub trait Generate<Cfg> {
    /// Turn Self into an ast::Item with a configuration object.
    fn generate(self, codemap::Span, &mut base::ExtCtxt, Cfg) -> P<ast::Item>;
}

impl<'a> Generate<&'a DescribeState> for Test {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, state: &'a DescribeState) -> P<ast::Item> {
        let Test { description, block, failing } = self;

        // Create the #[test] attribute.
        let test_attribute = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("test")));

        // Create the #[should_panic] attribute.
        let should_panic = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("should_panic")));

        let non_snake_word = cx.meta_word(sp, token::InternedString::new("non_snake_case"));
        let allow_non_snake_case = cx.meta_list(sp, token::InternedString::new("allow"),
                                                vec![non_snake_word]);
        let allow_non_snake_case = cx.attribute(sp, allow_non_snake_case);

        // Create the full test body by splicing in the statements and view items of the before and
        // after blocks if they are present.
        let test_body = match (&state.before_each, &state.after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    stmts: before.stmts.iter().chain(&*block.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    stmts: block.stmts.iter().chain(&*after.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    stmts: before.stmts.iter().chain(&*block.stmts).chain(&*after.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            }
        };

        // Constructing attributes:
        // #[test] - no way without it
        // #[allow(non_snake_case_attr)] as description may contain upper case
        // #[should_panic] if specified
        let mut attrs = vec![test_attribute, allow_non_snake_case];
        if failing {
            attrs.push(should_panic);
        }

        // Create the final Item that represents the test.
        P(ast::Item {
            // Name it with a snake_case version of the description.
            ident: cx.ident_of(&description.replace(" ", "_")),
            attrs: attrs,
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                // Takes no arguments and returns ()
                P(ast::FnDecl {
                    inputs: vec![],
                    output: ast::DefaultReturn(sp),
                    variadic: false
                }),
                // All the usual types.
                ast::Unsafety::Normal,
                ast::Constness::NotConst,
                abi::Rust,
                ast::Generics::default(),

                // Add the body of the function.
                test_body
            ),
            // Inherited visibility (not pub)
            vis: ast::Inherited,
            span: sp
        })
    }
}

impl Generate<()> for Bench {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, _: ()) -> P<ast::Item> {
        let Bench { bench, description, block } = self;

        // Create the #[bench] attribute.
        let bench_attribute = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("bench")));

        // Create the final Item that represents the benchmark.
        P(ast::Item {
            // Name it with a snake_case version of the description.
            ident: cx.ident_of(&description.replace(" ", "_")),

            // Add #[test] and possibly #[should_panic]
            attrs: vec![bench_attribute],
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                // Takes one argument of &mut Bencher
                P(ast::FnDecl {
                    inputs: vec![ast::Arg {
                        ty: quote_ty!(cx, &mut ::test::Bencher),
                        pat: quote_pat!(cx, $bench),
                        id: ast::DUMMY_NODE_ID
                    }],
                    output: ast::DefaultReturn(sp),
                    variadic: false
                }),

                // All the usual types.
                ast::Unsafety::Normal,
                ast::Constness::Const,
                abi::Rust,
                ast::Generics::default(),

                // Add the body of the function.
                block
            ),
            // Inherited visibility (not pub)
            vis: ast::Inherited,
            span: sp
        })
    }
}

impl<'a> Generate<&'a DescribeState> for SubBlock {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, state: &'a DescribeState) -> P<ast::Item> {
        match self {
            SubBlock::Test(test) => test.generate(sp, cx, state),
            SubBlock::Bench(bench) => bench.generate(sp, cx, ()),
            SubBlock::Describe(item) => item.generate(sp, cx, Some(state))
        }
    }
}

impl<'a> Generate<Option<&'a DescribeState>> for DescribeState {
    fn generate(mut self, sp: codemap::Span, cx: &mut base::ExtCtxt,
                state: Option<&'a DescribeState>) -> P<ast::Item> {
        // Get the name of this mod.
        let name = self.name.clone().unwrap();

        if let Some(state) = state {
            if let Some(ref parent) = state.before_each {
                self.before_each = match self.before_each {
                    Some(ref now) => Some(P(ast::Block {
                        stmts: parent.stmts.iter().chain(&*now.stmts).cloned().collect(),
                        ..now.deref().clone()
                    })),
                    None => Some(P(parent.deref().clone()))
                };
            }

            if let Some(ref parent) = state.after_each {
                self.after_each = match self.after_each {
                    Some(ref now) => Some(P(ast::Block {
                        stmts: now.stmts.iter().chain(&*parent.stmts).cloned().collect(),
                        ..now.deref().clone()
                    })),
                    None => Some(P(parent.deref().clone()))
                };
            }
        }

        // Get a glob import of all items in scope to the module that `describe!` is called from.
        //
        // This glob is `pub use super::*` so that nested `describe!` blocks (which will also contain
        // this glob) will be able to see all the symbols.
        let super_glob = cx.item_use_glob(sp, ast::Public, vec![cx.ident_of("super")]);
        let mut items = vec![super_glob];

        // Create subblocks from a full DescribeState
        items.extend(self.subblocks.clone().into_iter().map(|block| {
            block.generate(sp, cx, &self)
        }));

        // Generate the new module.
        cx.item_mod(sp, sp, name, vec![], items)
    }
}
<commit_msg>fix for rustc nightly<commit_after>use std::ops::Deref;

use syntax::{ast, abi, codemap};
use syntax::ptr::P;
use syntax::ext::base;
use syntax::parse::token;

use syntax::ext::build::AstBuilder;

use test::Test;
use bench::Bench;

use describe::{DescribeState, SubBlock};

/// Trait meaning something can be turned into an ast::Item with configuration.
pub trait Generate<Cfg> {
    /// Turn Self into an ast::Item with a configuration object.
    fn generate(self, codemap::Span, &mut base::ExtCtxt, Cfg) -> P<ast::Item>;
}

impl<'a> Generate<&'a DescribeState> for Test {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, state: &'a DescribeState) -> P<ast::Item> {
        let Test { description, block, failing } = self;

        // Create the #[test] attribute.
        let test_attribute = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("test")));

        // Create the #[should_panic] attribute.
        let should_panic = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("should_panic")));

        let non_snake_word = cx.meta_word(sp, token::InternedString::new("non_snake_case"));
        let allow_non_snake_case = cx.meta_list(sp, token::InternedString::new("allow"),
                                                vec![non_snake_word]);
        let allow_non_snake_case = cx.attribute(sp, allow_non_snake_case);

        // Create the full test body by splicing in the statements and view items of the before and
        // after blocks if they are present.
        let test_body = match (&state.before_each, &state.after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    stmts: before.stmts.iter().chain(&*block.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    stmts: block.stmts.iter().chain(&*after.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    stmts: before.stmts.iter().chain(&*block.stmts).chain(&*after.stmts).cloned().collect(),
                    ..block.deref().clone()
                })
            }
        };

        // Constructing attributes:
        // #[test] - no way without it
        // #[allow(non_snake_case_attr)] as description may contain upper case
        // #[should_panic] if specified
        let mut attrs = vec![test_attribute, allow_non_snake_case];
        if failing {
            attrs.push(should_panic);
        }

        // Create the final Item that represents the test.
        P(ast::Item {
            // Name it with a snake_case version of the description.
            ident: cx.ident_of(&description.replace(" ", "_")),
            attrs: attrs,
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemKind::Fn(
                // Takes no arguments and returns ()
                P(ast::FnDecl {
                    inputs: vec![],
                    output: ast::FunctionRetTy::Default(sp),
                    variadic: false
                }),
                // All the usual types.
                ast::Unsafety::Normal,
                ast::Constness::NotConst,
                abi::Abi::Rust,
                ast::Generics::default(),

                // Add the body of the function.
                test_body
            ),
            // Inherited visibility (not pub)
            vis: ast::Visibility::Inherited,
            span: sp
        })
    }
}

impl Generate<()> for Bench {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, _: ()) -> P<ast::Item> {
        let Bench { bench, description, block } = self;

        // Create the #[bench] attribute.
        let bench_attribute = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("bench")));

        // Create the final Item that represents the benchmark.
        P(ast::Item {
            // Name it with a snake_case version of the description.
            ident: cx.ident_of(&description.replace(" ", "_")),

            // Add #[test] and possibly #[should_panic]
            attrs: vec![bench_attribute],
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemKind::Fn(
                // Takes one argument of &mut Bencher
                P(ast::FnDecl {
                    inputs: vec![ast::Arg {
                        ty: quote_ty!(cx, &mut ::test::Bencher),
                        pat: quote_pat!(cx, $bench),
                        id: ast::DUMMY_NODE_ID
                    }],
                    output: ast::FunctionRetTy::Default(sp),
                    variadic: false
                }),

                // All the usual types.
                ast::Unsafety::Normal,
                ast::Constness::Const,
                abi::Abi::Rust,
                ast::Generics::default(),

                // Add the body of the function.
                block
            ),
            // Inherited visibility (not pub)
            vis: ast::Visibility::Inherited,
            span: sp
        })
    }
}

impl<'a> Generate<&'a DescribeState> for SubBlock {
    fn generate(self, sp: codemap::Span, cx: &mut base::ExtCtxt, state: &'a DescribeState) -> P<ast::Item> {
        match self {
            SubBlock::Test(test) => test.generate(sp, cx, state),
            SubBlock::Bench(bench) => bench.generate(sp, cx, ()),
            SubBlock::Describe(item) => item.generate(sp, cx, Some(state))
        }
    }
}

impl<'a> Generate<Option<&'a DescribeState>> for DescribeState {
    fn generate(mut self, sp: codemap::Span, cx: &mut base::ExtCtxt,
                state: Option<&'a DescribeState>) -> P<ast::Item> {
        // Get the name of this mod.
        let name = self.name.clone().unwrap();

        if let Some(state) = state {
            if let Some(ref parent) = state.before_each {
                self.before_each = match self.before_each {
                    Some(ref now) => Some(P(ast::Block {
                        stmts: parent.stmts.iter().chain(&*now.stmts).cloned().collect(),
                        ..now.deref().clone()
                    })),
                    None => Some(P(parent.deref().clone()))
                };
            }

            if let Some(ref parent) = state.after_each {
                self.after_each = match self.after_each {
                    Some(ref now) => Some(P(ast::Block {
                        stmts: now.stmts.iter().chain(&*parent.stmts).cloned().collect(),
                        ..now.deref().clone()
                    })),
                    None => Some(P(parent.deref().clone()))
                };
            }
        }

        // Get a glob import of all items in scope to the module that `describe!` is called from.
        //
        // This glob is `pub use super::*` so that nested `describe!` blocks (which will also contain
        // this glob) will be able to see all the symbols.
        let super_glob = cx.item_use_glob(sp, ast::Visibility::Public, vec![cx.ident_of("super")]);
        let mut items = vec![super_glob];

        // Create subblocks from a full DescribeState
        items.extend(self.subblocks.clone().into_iter().map(|block| {
            block.generate(sp, cx, &self)
        }));

        // Generate the new module.
        cx.item_mod(sp, sp, name, vec![], items)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster, Christopher Schröder.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Fasta reading and writing.
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::fasta;
//! let reader = fasta::Reader::new(io::stdin());
//! ```

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;
use std::collections;
use std::fs;
use std::path::Path;
use std::convert::AsRef;

use csv;

pub struct Reader<R: io::Read> {
    reader: io::BufReader<R>,
    line: String
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(|f| Reader::new(f))
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new FastQ reader.
    pub fn new(reader: R) -> Self {
        Reader { reader: io::BufReader::new(reader), line: String::new() }
    }

    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        if self.line.is_empty() {
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() {
                return Ok(());
            }
        }

        if !self.line.starts_with(">") {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "Expected > at record start."
            ));
        }
        record.header.push_str(&self.line);
        loop {
            self.line.clear();
            try!(self.reader.read_line(&mut self.line));
            record.seq.push_str(&self.line.trim_right());
            if self.line.is_empty() || self.line.starts_with(">") {
                break;
            }
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

pub struct Index {
    inner: collections::HashMap<Vec<u8>, IndexRecord>,
    seqs: Vec<Vec<u8>>,
}

impl Index {
    pub fn new<R: io::Read>(fai: R) -> csv::Result<Self> {
        let mut inner = collections::HashMap::new();
        let mut seqs = vec![];
        let mut fai_reader = csv::Reader::from_reader(fai).delimiter(b'\t').has_headers(false);
        for row in fai_reader.decode() {
            let (name, record): (String, IndexRecord) = try!(row);
            seqs.push(name.clone().into_bytes());
            inner.insert(name.into_bytes(), record);
        }
        Ok(Index { inner: inner, seqs: seqs })
    }

    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        match fs::File::open(path) {
            Ok(fai) => Self::new(fai),
            Err(e)  => Err(csv::Error::Io(e))
        }
    }

    pub fn with_fasta_file<P: AsRef<Path>>(fasta_path: &P) -> csv::Result<Self> {
        let mut ext = fasta_path.as_ref().extension().unwrap().to_str().unwrap().to_string();
        ext.push_str(".fai");
        let fai_path = fasta_path.as_ref().with_extension(ext);

        Self::from_file(&fai_path)
    }

    pub fn sequences(&self) -> Vec<Sequence> {
        self.seqs.iter().map(|name| Sequence { name: name.clone(), len: self.inner.get(name).unwrap().len }).collect()
    }
}

pub struct IndexedReader<R: io::Read + io::Seek> {
    reader: io::BufReader<R>,
    pub index: Index,
}

impl IndexedReader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        let index = try!(Index::with_fasta_file(path));

        match fs::File::open(path) {
            Ok(fasta) => Ok(IndexedReader::with_index(fasta, index)),
            Err(e)    => Err(csv::Error::Io(e))
        }
    }
}

impl<R: io::Read + io::Seek> IndexedReader<R> {
    pub fn new<I: io::Read>(fasta: R, fai: I) -> csv::Result<Self> {
        let index = try!(Index::new(fai));
        Ok(IndexedReader { reader: io::BufReader::new(fasta), index: index })
    }

    pub fn with_index(fasta: R, index: Index) -> Self {
        IndexedReader { reader: io::BufReader::new(fasta), index: index }
    }

    pub fn read_all(&mut self, seqname: &[u8], seq: &mut Vec<u8>) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(&idx) => self.read(seqname, 0, idx.len, seq),
            None      => Err(
                io::Error::new(
                    io::ErrorKind::Other,
                    "Unknown sequence name."
                )
            )
        }
    }

    pub fn read(&mut self, seqname: &[u8], start: u64, stop: u64, seq: &mut Vec<u8>) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(idx) => {
                seq.clear();
                // derived from
                // http://www.allenyu.info/item/24-quickly-fetch-sequence-from-samtools-faidx-indexed-fasta-sequences.html
                let line = start / idx.line_bases * idx.line_bytes;
                let line_offset = start % idx.line_bases;
                let offset = idx.offset + line + line_offset;
                let lines = stop / idx.line_bases * idx.line_bytes - line;
                let line_stop = stop % idx.line_bases - if lines == 0 { line_offset } else { 0 };

                try!(self.reader.seek(io::SeekFrom::Start(offset)));
                let mut buf = vec![0u8; idx.line_bases as usize];
                for _ in 0..lines {
                    // read full lines
                    try!(self.reader.read(&mut buf));
                    seq.push_all(&buf);
                }
                // read last line
                println!("linestop {}", line_stop);
                try!(self.reader.read(&mut buf[..line_stop as usize]));
                seq.push_all(&buf[..line_stop as usize]);
                Ok(())
            },
            None      => Err(
                io::Error::new(
                    io::ErrorKind::Other,
                    "Unknown sequence name."
                )
            )
        }
    }
}

#[derive(RustcDecodable, Debug, Copy, Clone)]
struct IndexRecord {
    len: u64,
    offset: u64,
    line_bases: u64,
    line_bytes: u64,
}

pub struct Sequence {
    pub name: Vec<u8>,
    pub len: u64,
}

/// A Fasta writer.
pub struct Writer<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl Writer<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(|f| Writer::new(f))
    }
}

impl<W: io::Write> Writer<W> {
    /// Create a new Fasta writer.
    pub fn new(writer: W) -> Self {
        Writer { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a Fasta record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), record.desc(), record.seq())
    }

    /// Write a Fasta record with given values.
    ///
    /// # Arguments
    ///
    /// * `id` - the record id
    /// * `desc` - the optional descriptions
    /// * `seq` - the sequence
    pub fn write(&mut self, id: &str, desc: Option<&str>, seq: &[u8]) -> io::Result<()> {
        try!(self.writer.write(b">"));
        try!(self.writer.write(id.as_bytes()));
        if desc.is_some() {
            try!(self.writer.write(b" "));
            try!(self.writer.write(desc.unwrap().as_bytes()));
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

pub struct Record {
    header: String,
    seq: String,
}

impl Record {
    pub fn new() -> Self {
        Record { header: String::new(), seq: String::new() }
    }

    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty()
    }

    /// Check validity of Fasta record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, " ").next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, " ").skip(1).next()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> &[u8] {
        self.seq.as_bytes()
    }

    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
    }
}

/// An iterator over the records of a Fasta file.
pub struct Records<R: io::Read> {
    reader: Reader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

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

    const FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
";
    const FAI_FILE: &'static [u8] = b"id\t12\t9\t60\t61
";

    #[test]
    fn test_reader() {
        let reader = Reader::new(FASTA_FILE);
        let records: Vec<io::Result<Record>> = reader.records().collect();
        assert!(records.len() == 1);
        for res in records {
            let record = res.ok().unwrap();
            assert_eq!(record.check(), Ok(()));
            assert_eq!(record.id(), Some("id"));
            assert_eq!(record.desc(), Some("desc"));
            assert_eq!(record.seq(), b"ACCGTAGGCTGA");
        }
    }

    #[test]
    fn test_indexed_reader() {
        let mut reader = IndexedReader::new(io::Cursor::new(FASTA_FILE), FAI_FILE).ok().expect("Error reading index");
        let mut seq = Vec::new();
        reader.read(b"id", 1, 5, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGT");
    }

    #[test]
    fn test_writer() {
        let mut writer = Writer::new(Vec::new());
        writer.write("id", Some("desc"), b"ACCGTAGGCTGA").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &FASTA_FILE);
    }
}
<commit_msg>Fix FASTA reader<commit_after>// Copyright 2014 Johannes Köster, Christopher Schröder.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Fasta reading and writing.
//!
//! # Example
//!
//! ```
//! use std::io;
//! use bio::io::fasta;
//! let reader = fasta::Reader::new(io::stdin());
//! ```

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;
use std::collections;
use std::fs;
use std::path::Path;
use std::convert::AsRef;

use csv;

pub struct Reader<R: io::Read> {
    reader: io::BufReader<R>,
    line: String
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::open(path).map(|f| Reader::new(f))
    }
}

impl<R: io::Read> Reader<R> {
    /// Create a new FastQ reader.
    pub fn new(reader: R) -> Self {
        Reader { reader: io::BufReader::new(reader), line: String::new() }
    }

    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        if self.line.is_empty() {
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() {
                return Ok(());
            }
        }

        if !self.line.starts_with(">") {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "Expected > at record start."
            ));
        }
        record.header.push_str(&self.line);
        loop {
            self.line.clear();
            try!(self.reader.read_line(&mut self.line));
            if self.line.is_empty() || self.line.starts_with(">") {
                break;
            }
            record.seq.push_str(&self.line.trim_right());
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

pub struct Index {
    inner: collections::HashMap<Vec<u8>, IndexRecord>,
    seqs: Vec<Vec<u8>>,
}

impl Index {
    pub fn new<R: io::Read>(fai: R) -> csv::Result<Self> {
        let mut inner = collections::HashMap::new();
        let mut seqs = vec![];
        let mut fai_reader = csv::Reader::from_reader(fai).delimiter(b'\t').has_headers(false);
        for row in fai_reader.decode() {
            let (name, record): (String, IndexRecord) = try!(row);
            seqs.push(name.clone().into_bytes());
            inner.insert(name.into_bytes(), record);
        }
        Ok(Index { inner: inner, seqs: seqs })
    }

    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        match fs::File::open(path) {
            Ok(fai) => Self::new(fai),
            Err(e)  => Err(csv::Error::Io(e))
        }
    }

    pub fn with_fasta_file<P: AsRef<Path>>(fasta_path: &P) -> csv::Result<Self> {
        let mut ext = fasta_path.as_ref().extension().unwrap().to_str().unwrap().to_string();
        ext.push_str(".fai");
        let fai_path = fasta_path.as_ref().with_extension(ext);

        Self::from_file(&fai_path)
    }

    pub fn sequences(&self) -> Vec<Sequence> {
        self.seqs.iter().map(|name| Sequence { name: name.clone(), len: self.inner.get(name).unwrap().len }).collect()
    }
}

pub struct IndexedReader<R: io::Read + io::Seek> {
    reader: io::BufReader<R>,
    pub index: Index,
}

impl IndexedReader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: &P) -> csv::Result<Self> {
        let index = try!(Index::with_fasta_file(path));

        match fs::File::open(path) {
            Ok(fasta) => Ok(IndexedReader::with_index(fasta, index)),
            Err(e)    => Err(csv::Error::Io(e))
        }
    }
}

impl<R: io::Read + io::Seek> IndexedReader<R> {
    pub fn new<I: io::Read>(fasta: R, fai: I) -> csv::Result<Self> {
        let index = try!(Index::new(fai));
        Ok(IndexedReader { reader: io::BufReader::new(fasta), index: index })
    }

    pub fn with_index(fasta: R, index: Index) -> Self {
        IndexedReader { reader: io::BufReader::new(fasta), index: index }
    }

    pub fn read_all(&mut self, seqname: &[u8], seq: &mut Vec<u8>) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(&idx) => self.read(seqname, 0, idx.len, seq),
            None      => Err(
                io::Error::new(
                    io::ErrorKind::Other,
                    "Unknown sequence name."
                )
            )
        }
    }

    pub fn read(&mut self, seqname: &[u8], start: u64, stop: u64, seq: &mut Vec<u8>) -> io::Result<()> {
        match self.index.inner.get(seqname) {
            Some(idx) => {
                seq.clear();
                // derived from
                // http://www.allenyu.info/item/24-quickly-fetch-sequence-from-samtools-faidx-indexed-fasta-sequences.html
                let line = start / idx.line_bases * idx.line_bytes;
                let line_offset = start % idx.line_bases;
                let offset = idx.offset + line + line_offset;
                let lines = stop / idx.line_bases * idx.line_bytes - line;
                let line_stop = stop % idx.line_bases - if lines == 0 { line_offset } else { 0 };

                try!(self.reader.seek(io::SeekFrom::Start(offset)));
                let mut buf = vec![0u8; idx.line_bases as usize];
                for _ in 0..lines {
                    // read full lines
                    try!(self.reader.read(&mut buf));
                    seq.push_all(&buf);
                }
                // read last line
                println!("linestop {}", line_stop);
                try!(self.reader.read(&mut buf[..line_stop as usize]));
                seq.push_all(&buf[..line_stop as usize]);
                Ok(())
            },
            None      => Err(
                io::Error::new(
                    io::ErrorKind::Other,
                    "Unknown sequence name."
                )
            )
        }
    }
}

#[derive(RustcDecodable, Debug, Copy, Clone)]
struct IndexRecord {
    len: u64,
    offset: u64,
    line_bases: u64,
    line_bytes: u64,
}

pub struct Sequence {
    pub name: Vec<u8>,
    pub len: u64,
}

/// A Fasta writer.
pub struct Writer<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl Writer<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        fs::File::create(path).map(|f| Writer::new(f))
    }
}

impl<W: io::Write> Writer<W> {
    /// Create a new Fasta writer.
    pub fn new(writer: W) -> Self {
        Writer { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a Fasta record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), record.desc(), record.seq())
    }

    /// Write a Fasta record with given values.
    ///
    /// # Arguments
    ///
    /// * `id` - the record id
    /// * `desc` - the optional descriptions
    /// * `seq` - the sequence
    pub fn write(&mut self, id: &str, desc: Option<&str>, seq: &[u8]) -> io::Result<()> {
        try!(self.writer.write(b">"));
        try!(self.writer.write(id.as_bytes()));
        if desc.is_some() {
            try!(self.writer.write(b" "));
            try!(self.writer.write(desc.unwrap().as_bytes()));
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

pub struct Record {
    header: String,
    seq: String,
}

impl Record {
    pub fn new() -> Self {
        Record { header: String::new(), seq: String::new() }
    }

    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty()
    }

    /// Check validity of Fasta record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, " ").next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Option<&str> {
        self.header[1..].trim_right().splitn(2, " ").skip(1).next()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> &[u8] {
        self.seq.as_bytes()
    }

    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
    }
}

/// An iterator over the records of a Fasta file.
pub struct Records<R: io::Read> {
    reader: Reader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

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

    const FASTA_FILE: &'static [u8] = b">id desc
ACCGTAGGCTGA
>id2
ATTGTTGTTTTA
";
    const FAI_FILE: &'static [u8] = b"id\t12\t9\t60\t61
";

    #[test]
    fn test_reader() {
        let reader = Reader::new(FASTA_FILE);
        let ids = [Some("id"), Some("id2")];
        let descs = [Some("desc"), None];
        let seqs : [&[u8]; 2] = [b"ACCGTAGGCTGA", b"ATTGTTGTTTTA"];

        for (i, r) in reader.records().enumerate() {
          let record = r.ok().expect("Error reading record");
          assert_eq!(record.check(), Ok(()));
          assert_eq!(record.id(), ids[i]);
          assert_eq!(record.desc(), descs[i]);
          assert_eq!(record.seq(), seqs[i]);
        }

        // let record = records.ok().nth(1).unwrap();
    }

    #[test]
    fn test_indexed_reader() {
        let mut reader = IndexedReader::new(io::Cursor::new(FASTA_FILE), FAI_FILE).ok().expect("Error reading index");
        let mut seq = Vec::new();
        reader.read(b"id", 1, 5, &mut seq).ok().expect("Error reading sequence.");
        assert_eq!(seq, b"CCGT");
    }

    #[test]
    fn test_writer() {
        let mut writer = Writer::new(Vec::new());
        writer.write("id", Some("desc"), b"ACCGTAGGCTGA").ok().expect("Expected successful write");
        writer.write("id2", None, b"ATTGTTGTTTTA").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &FASTA_FILE);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2018 Erik Nordstrøm <erik@nordstroem.no>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

use arrayvec::ArrayVec;

use std::ops::Deref;

use cards::Card;

macro_rules! impl_cardstack
{
    ($t:ident, $a:ty) =>
    {
        impl_cardstack_ops!($t, $a);

        impl $t
        {
            pub fn new () -> $t
            {
                $t::from(ArrayVec::<$a>::new())
            }
        }
    }
}

impl_cardstack!(TableauSlot, [Card; 19]);

pub struct Table
{
    pub tableau: [TableauSlot; 7],
    //pub foundation: [FoundationSlot; 4],
}
<commit_msg>Table cont.<commit_after>/*
 * Copyright (c) 2018 Erik Nordstrøm <erik@nordstroem.no>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

use arrayvec::ArrayVec;

use std::ops::Deref;

use cards::Card;

macro_rules! impl_cardstack
{
    ($t:ident, $a:ty) =>
    {
        impl_cardstack_ops!($t, $a);

        impl $t
        {
            pub fn new () -> $t
            {
                $t::from(ArrayVec::<$a>::new())
            }
        }
    }
}

impl_cardstack!(StockSlot,      [Card; 21]); // 52 - (1 + 2 + 3 + 4 + 5 + 6 + 7) = 21
impl_cardstack!(WastePileSlot,  [Card; 21]);
impl_cardstack!(FoundationSlot, [Card; 13]);
impl_cardstack!(TableauSlot,    [Card; 19]);

pub struct Table
{
    pub stock:       StockSlot,
    pub waste_pile:  WastePileSlot,
    pub foundations: [FoundationSlot; 4],
    pub tableau:     [TableauSlot; 7],
}
<|endoftext|>