<commit_before>extern crate libloading as libl;

use std::collections::HashMap;
use std::ffi::CString;
use std::io;
use std::os::raw::{c_char, c_int, c_uint, c_void};

fn to_c_str<T: Into<Vec<u8>>>(t: T) -> *const c_char {
    CString::new(t).unwrap().as_ptr() as *const c_char
}

pub struct CoreClr {
    host_handle: *const c_void,
    domain_id: c_uint
}

pub type CoreClrInitializeFn = unsafe extern fn(
    *const c_char,
    *const c_char,
    c_int,
    *const *const c_char,
    *const *const c_char,
    *const *const c_void,
    *const c_uint) -> c_int;

pub type CoreClrShutdownFn = unsafe extern fn(*const c_void, c_uint) -> c_int;
pub type CoreClrShutdown2Fn = unsafe extern fn(*const c_void, c_uint, *const c_int) -> c_int;

impl CoreClr {
    fn library() -> libl::Result<libl::Library> {
        libl::Library::new("/usr/local/share/dotnet/shared/Microsoft.NETCore.App/2.0.0")
    }

    pub fn init(
        exe_path: &str,
        app_domain_friendly_name: &str,
        properties_option: Option<HashMap<&str, &str>>) -> libl::Result<CoreClr>
    {
        // Create the host handle and its ref
        let host_handle = 0 as *const c_void;
        let host_handle_ref = &host_handle as *const *const c_void;

        // Create the domain id and its ref
        let domain_id = 0 as c_uint;
        let domain_id_ref = &domain_id as *const c_uint;

        // Raw C string refs from exe_path and app_domain_friendly_name slices
        let exe_path_raw = to_c_str(exe_path);
        let app_domain_friendly_name_raw = to_c_str(app_domain_friendly_name);

        // Either use the provided option or a blank hashmap
        let properties = properties_option.unwrap_or_else(HashMap::new);

        // Count for upcoming vecs
        let properties_count = properties.len() as c_int;

        // Collect property keys into a vec
        let properties_keys: Vec<*const c_char> = properties.keys()
                .map(|&k| to_c_str(k))
                .collect();

        // Collect equivalent property values into another vec
        let properties_values: Vec<*const c_char> = properties.values()
                .map(|&v| to_c_str(v))
                .collect();

        // Grab refs of the two to pass into the actual function call
        let properties_keys_ref = &properties_keys[0] as *const *const c_char;
        let properties_values_ref = &properties_values[0] as *const *const c_char;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_initialize_fn: libl::Symbol<CoreClrInitializeFn> = coreclr_library.get(b"coreclr_initialize")?;

            assert_eq!(
                coreclr_initialize_fn(
                    exe_path_raw,
                    app_domain_friendly_name_raw,
                    properties_count,
                    properties_keys_ref,
                    properties_values_ref,
                    host_handle_ref,
                    domain_id_ref), 0);
        }

        Ok(CoreClr {
            host_handle: host_handle,
            domain_id: domain_id
        })
    }

    pub fn shutdown(self: Self) -> io::Result<()> {
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_fn: libl::Symbol<CoreClrShutdownFn> = coreclr_library.get(b"coreclr_shutdown")?;

            match coreclr_shutdown_fn(self.host_handle, self.domain_id) {
                0 => Ok(()),
                _ => panic!("Failed to shutdown")
            }
        }
    }

    pub fn shutdown_2(self: Self) -> io::Result<c_int> {
        let latched_exit_code = -1 as c_int;
        let latched_exit_code_ref = &latched_exit_code as *const c_int;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_2_fn: libl::Symbol<CoreClrShutdown2Fn> = coreclr_library.get(b"coreclr_shutdown_2")?;

            match coreclr_shutdown_2_fn(self.host_handle, self.domain_id, latched_exit_code_ref) {
                0 => Ok(latched_exit_code),
                _ => panic!("Failed to shutdown")
            }
        }
    }
}<commit_msg>Updated init function with error handling more inline with the shutdown functions<commit_after>extern crate libloading as libl;

use std::collections::HashMap;
use std::ffi::CString;
use std::io;
use std::os::raw::{c_char, c_int, c_uint, c_void};

fn to_c_str<T: Into<Vec<u8>>>(t: T) -> *const c_char {
    CString::new(t).unwrap().as_ptr() as *const c_char
}

pub struct CoreClr {
    host_handle: *const c_void,
    domain_id: c_uint
}

pub type CoreClrInitializeFn = unsafe extern fn(
    *const c_char,
    *const c_char,
    c_int,
    *const *const c_char,
    *const *const c_char,
    *const *const c_void,
    *const c_uint) -> c_int;

pub type CoreClrShutdownFn = unsafe extern fn(*const c_void, c_uint) -> c_int;
pub type CoreClrShutdown2Fn = unsafe extern fn(*const c_void, c_uint, *const c_int) -> c_int;

impl CoreClr {
    fn library() -> libl::Result<libl::Library> {
        libl::Library::new("/usr/local/share/dotnet/shared/Microsoft.NETCore.App/2.0.0")
    }

    pub fn init(
        exe_path: &str,
        app_domain_friendly_name: &str,
        properties_option: Option<HashMap<&str, &str>>) -> libl::Result<CoreClr>
    {
        // Create the host handle and its ref
        let host_handle = 0 as *const c_void;
        let host_handle_ref = &host_handle as *const *const c_void;

        // Create the domain id and its ref
        let domain_id = 0 as c_uint;
        let domain_id_ref = &domain_id as *const c_uint;

        // Raw C string refs from exe_path and app_domain_friendly_name slices
        let exe_path_raw = to_c_str(exe_path);
        let app_domain_friendly_name_raw = to_c_str(app_domain_friendly_name);

        // Either use the provided option or a blank hashmap
        let properties = properties_option.unwrap_or_else(HashMap::new);

        // Count for upcoming vecs
        let properties_count = properties.len() as c_int;

        // Collect property keys into a vec
        let properties_keys: Vec<*const c_char> = properties.keys()
                .map(|&k| to_c_str(k))
                .collect();

        // Collect equivalent property values into another vec
        let properties_values: Vec<*const c_char> = properties.values()
                .map(|&v| to_c_str(v))
                .collect();

        // Grab refs of the two to pass into the actual function call
        let properties_keys_ref = &properties_keys[0] as *const *const c_char;
        let properties_values_ref = &properties_values[0] as *const *const c_char;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_initialize_fn: libl::Symbol<CoreClrInitializeFn> = coreclr_library.get(b"coreclr_initialize")?;

            match coreclr_initialize_fn(
                exe_path_raw,
                app_domain_friendly_name_raw,
                properties_count,
                properties_keys_ref,
                properties_values_ref,
                host_handle_ref,
                domain_id_ref)
            {
                0 => Ok(CoreClr {
                    host_handle: host_handle,
                    domain_id: domain_id
                }),
                _ => panic!("Failed to initialize")
            }
        }
    }

    pub fn shutdown(self: Self) -> io::Result<()> {
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_fn: libl::Symbol<CoreClrShutdownFn> = coreclr_library.get(b"coreclr_shutdown")?;

            match coreclr_shutdown_fn(self.host_handle, self.domain_id) {
                0 => Ok(()),
                _ => panic!("Failed to shutdown")
            }
        }
    }

    pub fn shutdown_2(self: Self) -> io::Result<c_int> {
        let latched_exit_code = -1 as c_int;
        let latched_exit_code_ref = &latched_exit_code as *const c_int;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_2_fn: libl::Symbol<CoreClrShutdown2Fn> = coreclr_library.get(b"coreclr_shutdown_2")?;

            match coreclr_shutdown_2_fn(self.host_handle, self.domain_id, latched_exit_code_ref) {
                0 => Ok(latched_exit_code),
                _ => panic!("Failed to shutdown")
            }
        }
    }
}<|endoftext|>


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

#[cfg(feature = "heapsize")] use heapsize::HeapSizeOf;
use std::cmp;
use std::fmt::{self, Formatter};
use std::io;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
use std::vec;
use parser::{ParseResult, ParseError};
use percent_encoding::percent_decode;
use idna;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum HostInternal {
    None,
    Domain,
    Ipv4(Ipv4Addr),
    Ipv6(Ipv6Addr),
}

#[cfg(feature = "heapsize")]
known_heap_size!(0, HostInternal);

impl<S> From<Host<S>> for HostInternal {
    fn from(host: Host<S>) -> HostInternal {
        match host {
            Host::Domain(_) => HostInternal::Domain,
            Host::Ipv4(address) => HostInternal::Ipv4(address),
            Host::Ipv6(address) => HostInternal::Ipv6(address),
        }
    }
}

/// The host name of an URL.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum Host<S=String> {
    /// A DNS domain name, as '.' dot-separated labels.
    /// Non-ASCII labels are encoded in punycode per IDNA.
    Domain(S),

    /// An IPv4 address.
    /// `Url::host_str` returns the serialization of this address,
    /// as four decimal integers separated by `.` dots.
    Ipv4(Ipv4Addr),

    /// An IPv6 address.
    /// `Url::host_str` returns the serialization of that address between `[` and `]` brackets,
    /// in the format per [RFC 5952 *A Recommendation
    /// for IPv6 Address Text Representation*](https://tools.ietf.org/html/rfc5952):
    /// lowercase hexadecimal with maximal `::` compression.
    Ipv6(Ipv6Addr),
}

#[cfg(feature = "heapsize")]
impl<S: HeapSizeOf> HeapSizeOf for Host<S> {
    fn heap_size_of_children(&self) -> usize {
        match *self {
            Host::Domain(ref s) => s.heap_size_of_children(),
            _ => 0,
        }
    }
}

impl<'a> Host<&'a str> {
    /// Return a copy of `self` that owns an allocated `String` but does not borrow an `&Url`.
    pub fn to_owned(&self) -> Host<String> {
        match *self {
            Host::Domain(domain) => Host::Domain(domain.to_owned()),
            Host::Ipv4(address) => Host::Ipv4(address),
            Host::Ipv6(address) => Host::Ipv6(address),
        }
    }
}

impl Host<String> {
    /// Parse a host: either an IPv6 address in [] square brackets, or a domain.
    ///
    /// https://url.spec.whatwg.org/#host-parsing
    pub fn parse(input: &str) -> Result<Self, ParseError> {
        if input.starts_with("[") {
            if !input.ends_with("]") {
                return Err(ParseError::InvalidIpv6Address)
            }
            return parse_ipv6addr(&input[1..input.len() - 1]).map(Host::Ipv6)
        }
        let domain = percent_decode(input.as_bytes()).decode_utf8_lossy();
        let domain = try!(idna::domain_to_ascii(&domain));
        if domain.find(|c| matches!(c,
            '\0' | '\t' | '\n' | '\r' | ' ' | '#' | '%' | '/' | ':' | '?' | '@' | '[' | '\\' | ']'
        )).is_some() {
            return Err(ParseError::InvalidDomainCharacter)
        }
        if let Some(address) = try!(parse_ipv4addr(&domain)) {
            Ok(Host::Ipv4(address))
        } else {
            Ok(Host::Domain(domain.into()))
        }
    }
}

impl<S: AsRef<str>> fmt::Display for Host<S> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Host::Domain(ref domain) => domain.as_ref().fmt(f),
            Host::Ipv4(ref addr) => addr.fmt(f),
            Host::Ipv6(ref addr) => {
                try!(f.write_str("["));
                try!(write_ipv6(addr, f));
                f.write_str("]")
            }
        }
    }
}

/// This mostly exists because coherence rules don’t allow us to implement
/// `ToSocketAddrs for (Host<S>, u16)`.
pub struct HostAndPort<S=String> {
    pub host: Host<S>,
    pub port: u16,
}

impl<'a> HostAndPort<&'a str> {
    /// Return a copy of `self` that owns an allocated `String` but does not borrow an `&Url`.
    pub fn to_owned(&self) -> HostAndPort<String> {
        HostAndPort {
            host: self.host.to_owned(),
            port: self.port
        }
    }
}

impl<S: AsRef<str>> ToSocketAddrs for HostAndPort<S> {
    type Iter = SocketAddrs;

    fn to_socket_addrs(&self) -> io::Result<Self::Iter> {
        let port = self.port;
        match self.host {
            Host::Domain(ref domain) => Ok(SocketAddrs {
                // FIXME: use std::net::lookup_host when it’s stable.
                state: SocketAddrsState::Domain(try!((domain.as_ref(), port).to_socket_addrs()))
            }),
            Host::Ipv4(address) => Ok(SocketAddrs {
                state: SocketAddrsState::One(SocketAddr::V4(SocketAddrV4::new(address, port)))
            }),
            Host::Ipv6(address) => Ok(SocketAddrs {
                state: SocketAddrsState::One(SocketAddr::V6(SocketAddrV6::new(address, port, 0, 0)))
            }),
        }
    }
}

/// Socket addresses for an URL.
pub struct SocketAddrs {
    state: SocketAddrsState
}

enum SocketAddrsState {
    Domain(vec::IntoIter<SocketAddr>),
    One(SocketAddr),
    Done,
}

impl Iterator for SocketAddrs {
    type Item = SocketAddr;
    fn next(&mut self) -> Option<SocketAddr> {
        match self.state {
            SocketAddrsState::Domain(ref mut iter) => iter.next(),
            SocketAddrsState::One(s) => {
                self.state = SocketAddrsState::Done;
                Some(s)
            }
            SocketAddrsState::Done => None
        }
    }
}

fn write_ipv6(addr: &Ipv6Addr, f: &mut Formatter) -> fmt::Result {
    let segments = addr.segments();
    let (compress_start, compress_end) = longest_zero_sequence(&segments);
    let mut i = 0;
    while i < 8 {
        if i == compress_start {
            try!(f.write_str(":"));
            if i == 0 {
                try!(f.write_str(":"));
            }
            if compress_end < 8 {
                i = compress_end;
            } else {
                break;
            }
        }
        try!(write!(f, "{:x}", segments[i as usize]));
        if i < 7 {
            try!(f.write_str(":"));
        }
        i += 1;
    }
    Ok(())
}

fn longest_zero_sequence(pieces: &[u16; 8]) -> (isize, isize) {
    let mut longest = -1;
    let mut longest_length = -1;
    let mut start = -1;
    macro_rules! finish_sequence(
        ($end: expr) => {
            if start >= 0 {
                let length = $end - start;
                if length > longest_length {
                    longest = start;
                    longest_length = length;
                }
            }
        };
    );
    for i in 0..8 {
        if pieces[i as usize] == 0 {
            if start < 0 {
                start = i;
            }
        } else {
            finish_sequence!(i);
            start = -1;
        }
    }
    finish_sequence!(8);
    (longest, longest + longest_length)
}

fn parse_ipv4number(mut input: &str) -> Result<u32, ()> {
    let mut r = 10;
    if input.starts_with("0x") || input.starts_with("0X") {
        input = &input[2..];
        r = 16;
    } else if input.len() >= 2 && input.starts_with("0") {
        input = &input[1..];
        r = 8;
    }
    if input.is_empty() {
        return Ok(0);
    }
    if input.starts_with("+") {
        return Err(())
    }
    match u32::from_str_radix(&input, r) {
        Ok(number) => Ok(number),
        Err(_) => Err(()),
    }
}

fn parse_ipv4addr(input: &str) -> ParseResult<Option<Ipv4Addr>> {
    if input.is_empty() {
        return Ok(None)
    }
    let mut parts: Vec<&str> = input.split('.').collect();
    if parts.last() == Some(&"") {
        parts.pop();
    }
    if parts.len() > 4 {
        return Ok(None);
    }
    let mut numbers: Vec<u32> = Vec::new();
    for part in parts {
        if part == "" {
            return Ok(None);
        }
        if let Ok(n) = parse_ipv4number(part) {
            numbers.push(n);
        } else {
            return Ok(None);
        }
    }
    let mut ipv4 = numbers.pop().expect("a non-empty list of numbers");
    // Equivalent to: ipv4 >= 256 ** (4 − numbers.len())
    if ipv4 > u32::max_value() >> (8 * numbers.len() as u32)  {
        return Err(ParseError::InvalidIpv4Address);
    }
    if numbers.iter().any(|x| *x > 255) {
        return Err(ParseError::InvalidIpv4Address);
    }
    for (counter, n) in numbers.iter().enumerate() {
        ipv4 += n << (8 * (3 - counter as u32))
    }
    Ok(Some(Ipv4Addr::from(ipv4)))
}

fn parse_ipv6addr(input: &str) -> ParseResult<Ipv6Addr> {
    let input = input.as_bytes();
    let len = input.len();
    let mut is_ip_v4 = false;
    let mut pieces = [0, 0, 0, 0, 0, 0, 0, 0];
    let mut piece_pointer = 0;
    let mut compress_pointer = None;
    let mut i = 0;

    if len < 2 {
        return Err(ParseError::InvalidIpv6Address)
    }

    if input[0] == b':' {
        if input[1] != b':' {
            return Err(ParseError::InvalidIpv6Address)
        }
        i = 2;
        piece_pointer = 1;
        compress_pointer = Some(1);
    }

    while i < len {
        if piece_pointer == 8 {
            return Err(ParseError::InvalidIpv6Address)
        }
        if input[i] == b':' {
            if compress_pointer.is_some() {
                return Err(ParseError::InvalidIpv6Address)
            }
            i += 1;
            piece_pointer += 1;
            compress_pointer = Some(piece_pointer);
            continue
        }
        let start = i;
        let end = cmp::min(len, start + 4);
        let mut value = 0u16;
        while i < end {
            match (input[i] as char).to_digit(16) {
                Some(digit) => {
                    value = value * 0x10 + digit as u16;
                    i += 1;
                },
                None => break
            }
        }
        if i < len {
            match input[i] {
                b'.' => {
                    if i == start {
                        return Err(ParseError::InvalidIpv6Address)
                    }
                    i = start;
                    is_ip_v4 = true;
                },
                b':' => {
                    i += 1;
                    if i == len {
                        return Err(ParseError::InvalidIpv6Address)
                    }
                },
                _ => return Err(ParseError::InvalidIpv6Address)
            }
        }
        if is_ip_v4 {
            break
        }
        pieces[piece_pointer] = value;
        piece_pointer += 1;
    }

    if is_ip_v4 {
        if piece_pointer > 6 {
            return Err(ParseError::InvalidIpv6Address)
        }
        let mut dots_seen = 0;
        while i < len {
            let mut value = None;
            while i < len {
                let digit = match input[i] {
                    c @ b'0' ... b'9' => c - b'0',
                    _ => break
                };
                match value {
                    None => value = Some(digit as u16),
                    Some(0) => return Err(ParseError::InvalidIpv6Address),  // No leading zero
                    Some(ref mut v) => {
                        *v = *v * 10 + digit as u16;
                        if *v > 255 {
                            return Err(ParseError::InvalidIpv6Address)
                        }
                    }
                }
                i += 1;
            }
            if dots_seen < 3 && !(i < len && input[i] == b'.') {
                return Err(ParseError::InvalidIpv6Address)
            }
            pieces[piece_pointer] = if let Some(v) = value {
                pieces[piece_pointer] * 0x100 + v
            } else {
                return Err(ParseError::InvalidIpv6Address)
            };
            if dots_seen == 1 || dots_seen == 3 {
                piece_pointer += 1;
            }
            i += 1;
            if dots_seen == 3 && i < len {
                return Err(ParseError::InvalidIpv6Address)
            }
            dots_seen += 1;
        }
    }

    match compress_pointer {
        Some(compress_pointer) => {
            let mut swaps = piece_pointer - compress_pointer;
            piece_pointer = 7;
            while swaps > 0 {
                pieces[piece_pointer] = pieces[compress_pointer + swaps - 1];
                pieces[compress_pointer + swaps - 1] = 0;
                swaps -= 1;
                piece_pointer -= 1;
            }
        }
        _ => if piece_pointer != 8 {
            return Err(ParseError::InvalidIpv6Address)
        }
    }
    Ok(Ipv6Addr::new(pieces[0], pieces[1], pieces[2], pieces[3],
                     pieces[4], pieces[5], pieces[6], pieces[7]))
}
<commit_msg>Auto merge of #250 - amikhalev:HostAndPort-derive-Clone, r=SimonSapin<commit_after>// Copyright 2013-2016 The rust-url developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[cfg(feature = "heapsize")] use heapsize::HeapSizeOf;
use std::cmp;
use std::fmt::{self, Formatter};
use std::io;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs};
use std::vec;
use parser::{ParseResult, ParseError};
use percent_encoding::percent_decode;
use idna;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum HostInternal {
    None,
    Domain,
    Ipv4(Ipv4Addr),
    Ipv6(Ipv6Addr),
}

#[cfg(feature = "heapsize")]
known_heap_size!(0, HostInternal);

impl<S> From<Host<S>> for HostInternal {
    fn from(host: Host<S>) -> HostInternal {
        match host {
            Host::Domain(_) => HostInternal::Domain,
            Host::Ipv4(address) => HostInternal::Ipv4(address),
            Host::Ipv6(address) => HostInternal::Ipv6(address),
        }
    }
}

/// The host name of an URL.
#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum Host<S=String> {
    /// A DNS domain name, as '.' dot-separated labels.
    /// Non-ASCII labels are encoded in punycode per IDNA.
    Domain(S),

    /// An IPv4 address.
    /// `Url::host_str` returns the serialization of this address,
    /// as four decimal integers separated by `.` dots.
    Ipv4(Ipv4Addr),

    /// An IPv6 address.
    /// `Url::host_str` returns the serialization of that address between `[` and `]` brackets,
    /// in the format per [RFC 5952 *A Recommendation
    /// for IPv6 Address Text Representation*](https://tools.ietf.org/html/rfc5952):
    /// lowercase hexadecimal with maximal `::` compression.
    Ipv6(Ipv6Addr),
}

#[cfg(feature = "heapsize")]
impl<S: HeapSizeOf> HeapSizeOf for Host<S> {
    fn heap_size_of_children(&self) -> usize {
        match *self {
            Host::Domain(ref s) => s.heap_size_of_children(),
            _ => 0,
        }
    }
}

impl<'a> Host<&'a str> {
    /// Return a copy of `self` that owns an allocated `String` but does not borrow an `&Url`.
    pub fn to_owned(&self) -> Host<String> {
        match *self {
            Host::Domain(domain) => Host::Domain(domain.to_owned()),
            Host::Ipv4(address) => Host::Ipv4(address),
            Host::Ipv6(address) => Host::Ipv6(address),
        }
    }
}

impl Host<String> {
    /// Parse a host: either an IPv6 address in [] square brackets, or a domain.
    ///
    /// https://url.spec.whatwg.org/#host-parsing
    pub fn parse(input: &str) -> Result<Self, ParseError> {
        if input.starts_with("[") {
            if !input.ends_with("]") {
                return Err(ParseError::InvalidIpv6Address)
            }
            return parse_ipv6addr(&input[1..input.len() - 1]).map(Host::Ipv6)
        }
        let domain = percent_decode(input.as_bytes()).decode_utf8_lossy();
        let domain = try!(idna::domain_to_ascii(&domain));
        if domain.find(|c| matches!(c,
            '\0' | '\t' | '\n' | '\r' | ' ' | '#' | '%' | '/' | ':' | '?' | '@' | '[' | '\\' | ']'
        )).is_some() {
            return Err(ParseError::InvalidDomainCharacter)
        }
        if let Some(address) = try!(parse_ipv4addr(&domain)) {
            Ok(Host::Ipv4(address))
        } else {
            Ok(Host::Domain(domain.into()))
        }
    }
}

impl<S: AsRef<str>> fmt::Display for Host<S> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Host::Domain(ref domain) => domain.as_ref().fmt(f),
            Host::Ipv4(ref addr) => addr.fmt(f),
            Host::Ipv6(ref addr) => {
                try!(f.write_str("["));
                try!(write_ipv6(addr, f));
                f.write_str("]")
            }
        }
    }
}

/// This mostly exists because coherence rules don’t allow us to implement
/// `ToSocketAddrs for (Host<S>, u16)`.
#[derive(Clone)]
pub struct HostAndPort<S=String> {
    pub host: Host<S>,
    pub port: u16,
}

impl<'a> HostAndPort<&'a str> {
    /// Return a copy of `self` that owns an allocated `String` but does not borrow an `&Url`.
    pub fn to_owned(&self) -> HostAndPort<String> {
        HostAndPort {
            host: self.host.to_owned(),
            port: self.port
        }
    }
}

impl<S: AsRef<str>> ToSocketAddrs for HostAndPort<S> {
    type Iter = SocketAddrs;

    fn to_socket_addrs(&self) -> io::Result<Self::Iter> {
        let port = self.port;
        match self.host {
            Host::Domain(ref domain) => Ok(SocketAddrs {
                // FIXME: use std::net::lookup_host when it’s stable.
                state: SocketAddrsState::Domain(try!((domain.as_ref(), port).to_socket_addrs()))
            }),
            Host::Ipv4(address) => Ok(SocketAddrs {
                state: SocketAddrsState::One(SocketAddr::V4(SocketAddrV4::new(address, port)))
            }),
            Host::Ipv6(address) => Ok(SocketAddrs {
                state: SocketAddrsState::One(SocketAddr::V6(SocketAddrV6::new(address, port, 0, 0)))
            }),
        }
    }
}

/// Socket addresses for an URL.
pub struct SocketAddrs {
    state: SocketAddrsState
}

enum SocketAddrsState {
    Domain(vec::IntoIter<SocketAddr>),
    One(SocketAddr),
    Done,
}

impl Iterator for SocketAddrs {
    type Item = SocketAddr;
    fn next(&mut self) -> Option<SocketAddr> {
        match self.state {
            SocketAddrsState::Domain(ref mut iter) => iter.next(),
            SocketAddrsState::One(s) => {
                self.state = SocketAddrsState::Done;
                Some(s)
            }
            SocketAddrsState::Done => None
        }
    }
}

fn write_ipv6(addr: &Ipv6Addr, f: &mut Formatter) -> fmt::Result {
    let segments = addr.segments();
    let (compress_start, compress_end) = longest_zero_sequence(&segments);
    let mut i = 0;
    while i < 8 {
        if i == compress_start {
            try!(f.write_str(":"));
            if i == 0 {
                try!(f.write_str(":"));
            }
            if compress_end < 8 {
                i = compress_end;
            } else {
                break;
            }
        }
        try!(write!(f, "{:x}", segments[i as usize]));
        if i < 7 {
            try!(f.write_str(":"));
        }
        i += 1;
    }
    Ok(())
}

fn longest_zero_sequence(pieces: &[u16; 8]) -> (isize, isize) {
    let mut longest = -1;
    let mut longest_length = -1;
    let mut start = -1;
    macro_rules! finish_sequence(
        ($end: expr) => {
            if start >= 0 {
                let length = $end - start;
                if length > longest_length {
                    longest = start;
                    longest_length = length;
                }
            }
        };
    );
    for i in 0..8 {
        if pieces[i as usize] == 0 {
            if start < 0 {
                start = i;
            }
        } else {
            finish_sequence!(i);
            start = -1;
        }
    }
    finish_sequence!(8);
    (longest, longest + longest_length)
}

fn parse_ipv4number(mut input: &str) -> Result<u32, ()> {
    let mut r = 10;
    if input.starts_with("0x") || input.starts_with("0X") {
        input = &input[2..];
        r = 16;
    } else if input.len() >= 2 && input.starts_with("0") {
        input = &input[1..];
        r = 8;
    }
    if input.is_empty() {
        return Ok(0);
    }
    if input.starts_with("+") {
        return Err(())
    }
    match u32::from_str_radix(&input, r) {
        Ok(number) => Ok(number),
        Err(_) => Err(()),
    }
}

fn parse_ipv4addr(input: &str) -> ParseResult<Option<Ipv4Addr>> {
    if input.is_empty() {
        return Ok(None)
    }
    let mut parts: Vec<&str> = input.split('.').collect();
    if parts.last() == Some(&"") {
        parts.pop();
    }
    if parts.len() > 4 {
        return Ok(None);
    }
    let mut numbers: Vec<u32> = Vec::new();
    for part in parts {
        if part == "" {
            return Ok(None);
        }
        if let Ok(n) = parse_ipv4number(part) {
            numbers.push(n);
        } else {
            return Ok(None);
        }
    }
    let mut ipv4 = numbers.pop().expect("a non-empty list of numbers");
    // Equivalent to: ipv4 >= 256 ** (4 − numbers.len())
    if ipv4 > u32::max_value() >> (8 * numbers.len() as u32)  {
        return Err(ParseError::InvalidIpv4Address);
    }
    if numbers.iter().any(|x| *x > 255) {
        return Err(ParseError::InvalidIpv4Address);
    }
    for (counter, n) in numbers.iter().enumerate() {
        ipv4 += n << (8 * (3 - counter as u32))
    }
    Ok(Some(Ipv4Addr::from(ipv4)))
}

fn parse_ipv6addr(input: &str) -> ParseResult<Ipv6Addr> {
    let input = input.as_bytes();
    let len = input.len();
    let mut is_ip_v4 = false;
    let mut pieces = [0, 0, 0, 0, 0, 0, 0, 0];
    let mut piece_pointer = 0;
    let mut compress_pointer = None;
    let mut i = 0;

    if len < 2 {
        return Err(ParseError::InvalidIpv6Address)
    }

    if input[0] == b':' {
        if input[1] != b':' {
            return Err(ParseError::InvalidIpv6Address)
        }
        i = 2;
        piece_pointer = 1;
        compress_pointer = Some(1);
    }

    while i < len {
        if piece_pointer == 8 {
            return Err(ParseError::InvalidIpv6Address)
        }
        if input[i] == b':' {
            if compress_pointer.is_some() {
                return Err(ParseError::InvalidIpv6Address)
            }
            i += 1;
            piece_pointer += 1;
            compress_pointer = Some(piece_pointer);
            continue
        }
        let start = i;
        let end = cmp::min(len, start + 4);
        let mut value = 0u16;
        while i < end {
            match (input[i] as char).to_digit(16) {
                Some(digit) => {
                    value = value * 0x10 + digit as u16;
                    i += 1;
                },
                None => break
            }
        }
        if i < len {
            match input[i] {
                b'.' => {
                    if i == start {
                        return Err(ParseError::InvalidIpv6Address)
                    }
                    i = start;
                    is_ip_v4 = true;
                },
                b':' => {
                    i += 1;
                    if i == len {
                        return Err(ParseError::InvalidIpv6Address)
                    }
                },
                _ => return Err(ParseError::InvalidIpv6Address)
            }
        }
        if is_ip_v4 {
            break
        }
        pieces[piece_pointer] = value;
        piece_pointer += 1;
    }

    if is_ip_v4 {
        if piece_pointer > 6 {
            return Err(ParseError::InvalidIpv6Address)
        }
        let mut dots_seen = 0;
        while i < len {
            let mut value = None;
            while i < len {
                let digit = match input[i] {
                    c @ b'0' ... b'9' => c - b'0',
                    _ => break
                };
                match value {
                    None => value = Some(digit as u16),
                    Some(0) => return Err(ParseError::InvalidIpv6Address),  // No leading zero
                    Some(ref mut v) => {
                        *v = *v * 10 + digit as u16;
                        if *v > 255 {
                            return Err(ParseError::InvalidIpv6Address)
                        }
                    }
                }
                i += 1;
            }
            if dots_seen < 3 && !(i < len && input[i] == b'.') {
                return Err(ParseError::InvalidIpv6Address)
            }
            pieces[piece_pointer] = if let Some(v) = value {
                pieces[piece_pointer] * 0x100 + v
            } else {
                return Err(ParseError::InvalidIpv6Address)
            };
            if dots_seen == 1 || dots_seen == 3 {
                piece_pointer += 1;
            }
            i += 1;
            if dots_seen == 3 && i < len {
                return Err(ParseError::InvalidIpv6Address)
            }
            dots_seen += 1;
        }
    }

    match compress_pointer {
        Some(compress_pointer) => {
            let mut swaps = piece_pointer - compress_pointer;
            piece_pointer = 7;
            while swaps > 0 {
                pieces[piece_pointer] = pieces[compress_pointer + swaps - 1];
                pieces[compress_pointer + swaps - 1] = 0;
                swaps -= 1;
                piece_pointer -= 1;
            }
        }
        _ => if piece_pointer != 8 {
            return Err(ParseError::InvalidIpv6Address)
        }
    }
    Ok(Ipv6Addr::new(pieces[0], pieces[1], pieces[2], pieces[3],
                     pieces[4], pieces[5], pieces[6], pieces[7]))
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::ops::{DerefMut, Deref};

use toml::Value;

use libimagrt::runtime::Runtime;
use libimagrt::edit::{Edit, EditResult};
use libimagstore::storeid::IntoStoreId;
use libimagstore::storeid::StoreId;
use libimagstore::storeid::StoreIdIterator;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagentrytag::tag::{Tag, TagSlice};
use libimagentrytag::tagable::Tagable;
use libimagentrytag::result::Result as TagResult;

use module_path::ModuleEntryPath;
use result::Result;
use error::NoteError as NE;
use error::NoteErrorKind as NEK;

#[derive(Debug)]
pub struct Note<'a> {
    entry: FileLockEntry<'a>,
}

impl<'a> Note<'a> {

    pub fn new(store: &Store, name: String, text: String) -> Result<Note> {
        use std::ops::DerefMut;

        debug!("Creating new Note: '{}'", name);
        let fle = {
            let lockentry = store.create(ModuleEntryPath::new(name.clone()).into_storeid());
            if lockentry.is_err() {
                return Err(NE::new(NEK::StoreWriteError, Some(Box::new(lockentry.unwrap_err()))));
            }
            let mut lockentry = lockentry.unwrap();

            {
                let mut entry  = lockentry.deref_mut();

                {
                    let mut header = entry.get_header_mut();
                    let setres = header.set("note", Value::Table(BTreeMap::new()));
                    if setres.is_err() {
                        let kind = NEK::StoreWriteError;
                        return Err(NE::new(kind, Some(Box::new(setres.unwrap_err()))));
                    }

                    let setres = header.set("note.name", Value::String(name));
                    if setres.is_err() {
                        let kind = NEK::StoreWriteError;
                        return Err(NE::new(kind, Some(Box::new(setres.unwrap_err()))));
                    }
                }

                *entry.get_content_mut() = text;
            }

            lockentry
        };

        Ok(Note { entry: fle })
    }

    pub fn set_name(&mut self, n: String) -> Result<()> {
        let mut header = self.entry.get_header_mut();
        header.set("note.name", Value::String(n))
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
            .map(|_| ())
    }

    pub fn get_name(&self) -> Result<String> {
        let header = self.entry.get_header();
        match header.read("note.name") {
            Ok(Some(Value::String(s))) => Ok(String::from(s)),
            Ok(_)                => {
                let e = NE::new(NEK::HeaderTypeError, None);
                Err(NE::new(NEK::StoreReadError, Some(Box::new(e))))
            },
            Err(e) => Err(NE::new(NEK::StoreReadError, Some(Box::new(e))))
        }
    }

    pub fn set_text(&mut self, n: String) {
        *self.entry.get_content_mut() = n
    }

    pub fn get_text(&self) -> &String {
        self.entry.get_content()
    }

    pub fn delete(store: &Store, name: String) -> Result<()> {
        store.delete(ModuleEntryPath::new(name).into_storeid())
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
    }

    pub fn retrieve(store: &Store, name: String) -> Result<Note> {
        store.retrieve(ModuleEntryPath::new(name).into_storeid())
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
            .map(|entry| Note { entry: entry })
    }

    pub fn all_notes(store: &Store) -> Result<NoteIterator> {
        store.retrieve_for_module("notes")
            .map(|iter| NoteIterator::new(store, iter))
            .map_err(|e| NE::new(NEK::StoreReadError, Some(Box::new(e))))
    }

}

impl<'a> Edit for Note<'a> {

    fn edit_content(&mut self, rt: &Runtime) -> EditResult<()> {
        self.entry.edit_content(rt)
    }

}

impl<'a> Tagable for Note<'a> {

    fn get_tags(&self) -> TagResult<Vec<Tag>> {
        self.entry.get_tags()
    }

    fn set_tags(&mut self, ts: &[Tag]) -> TagResult<()> {
        self.entry.set_tags(ts)
    }

    fn add_tag(&mut self, t: Tag) -> TagResult<()> {
        self.entry.add_tag(t)
    }

    fn remove_tag(&mut self, t: Tag) -> TagResult<()> {
        self.entry.remove_tag(t)
    }

    fn has_tag(&self, t: TagSlice) -> TagResult<bool> {
        self.entry.has_tag(t)
    }

    fn has_tags(&self, ts: &[Tag]) -> TagResult<bool> {
        self.entry.has_tags(ts)
    }

}

trait FromStoreId {
    fn from_storeid(&Store, StoreId) -> Result<Note>;
}

impl<'a> FromStoreId for Note<'a> {

    fn from_storeid(store: &Store, id: StoreId) -> Result<Note> {
        debug!("Loading note from storeid: '{:?}'", id);
        match store.retrieve(id) {
            Err(e)    => Err(NE::new(NEK::StoreReadError, Some(Box::new(e)))),
            Ok(entry) => Ok(Note { entry: entry }),
        }
    }

}

impl<'a> Deref for Note<'a> {

    type Target = FileLockEntry<'a>;

    fn deref(&self) -> &FileLockEntry<'a> {
        &self.entry
    }

}

#[derive(Debug)]
pub struct NoteIterator<'a> {
    store: &'a Store,
    iditer: StoreIdIterator,
}

impl<'a> NoteIterator<'a> {

    pub fn new(store: &'a Store, iditer: StoreIdIterator) -> NoteIterator<'a> {
        NoteIterator {
            store: store,
            iditer: iditer,
        }
    }

}

impl<'a> Iterator for NoteIterator<'a> {
    type Item = Result<Note<'a>>;

    fn next(&mut self) -> Option<Result<Note<'a>>> {
        self.iditer
            .next()
            .map(|id| Note::from_storeid(self.store, id))
    }

}

<commit_msg>Remove unused import<commit_after>use std::collections::BTreeMap;
use std::ops::Deref;

use toml::Value;

use libimagrt::runtime::Runtime;
use libimagrt::edit::{Edit, EditResult};
use libimagstore::storeid::IntoStoreId;
use libimagstore::storeid::StoreId;
use libimagstore::storeid::StoreIdIterator;
use libimagstore::store::FileLockEntry;
use libimagstore::store::Store;
use libimagentrytag::tag::{Tag, TagSlice};
use libimagentrytag::tagable::Tagable;
use libimagentrytag::result::Result as TagResult;

use module_path::ModuleEntryPath;
use result::Result;
use error::NoteError as NE;
use error::NoteErrorKind as NEK;

#[derive(Debug)]
pub struct Note<'a> {
    entry: FileLockEntry<'a>,
}

impl<'a> Note<'a> {

    pub fn new(store: &Store, name: String, text: String) -> Result<Note> {
        use std::ops::DerefMut;

        debug!("Creating new Note: '{}'", name);
        let fle = {
            let lockentry = store.create(ModuleEntryPath::new(name.clone()).into_storeid());
            if lockentry.is_err() {
                return Err(NE::new(NEK::StoreWriteError, Some(Box::new(lockentry.unwrap_err()))));
            }
            let mut lockentry = lockentry.unwrap();

            {
                let mut entry  = lockentry.deref_mut();

                {
                    let mut header = entry.get_header_mut();
                    let setres = header.set("note", Value::Table(BTreeMap::new()));
                    if setres.is_err() {
                        let kind = NEK::StoreWriteError;
                        return Err(NE::new(kind, Some(Box::new(setres.unwrap_err()))));
                    }

                    let setres = header.set("note.name", Value::String(name));
                    if setres.is_err() {
                        let kind = NEK::StoreWriteError;
                        return Err(NE::new(kind, Some(Box::new(setres.unwrap_err()))));
                    }
                }

                *entry.get_content_mut() = text;
            }

            lockentry
        };

        Ok(Note { entry: fle })
    }

    pub fn set_name(&mut self, n: String) -> Result<()> {
        let mut header = self.entry.get_header_mut();
        header.set("note.name", Value::String(n))
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
            .map(|_| ())
    }

    pub fn get_name(&self) -> Result<String> {
        let header = self.entry.get_header();
        match header.read("note.name") {
            Ok(Some(Value::String(s))) => Ok(String::from(s)),
            Ok(_)                => {
                let e = NE::new(NEK::HeaderTypeError, None);
                Err(NE::new(NEK::StoreReadError, Some(Box::new(e))))
            },
            Err(e) => Err(NE::new(NEK::StoreReadError, Some(Box::new(e))))
        }
    }

    pub fn set_text(&mut self, n: String) {
        *self.entry.get_content_mut() = n
    }

    pub fn get_text(&self) -> &String {
        self.entry.get_content()
    }

    pub fn delete(store: &Store, name: String) -> Result<()> {
        store.delete(ModuleEntryPath::new(name).into_storeid())
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
    }

    pub fn retrieve(store: &Store, name: String) -> Result<Note> {
        store.retrieve(ModuleEntryPath::new(name).into_storeid())
            .map_err(|e| NE::new(NEK::StoreWriteError, Some(Box::new(e))))
            .map(|entry| Note { entry: entry })
    }

    pub fn all_notes(store: &Store) -> Result<NoteIterator> {
        store.retrieve_for_module("notes")
            .map(|iter| NoteIterator::new(store, iter))
            .map_err(|e| NE::new(NEK::StoreReadError, Some(Box::new(e))))
    }

}

impl<'a> Edit for Note<'a> {

    fn edit_content(&mut self, rt: &Runtime) -> EditResult<()> {
        self.entry.edit_content(rt)
    }

}

impl<'a> Tagable for Note<'a> {

    fn get_tags(&self) -> TagResult<Vec<Tag>> {
        self.entry.get_tags()
    }

    fn set_tags(&mut self, ts: &[Tag]) -> TagResult<()> {
        self.entry.set_tags(ts)
    }

    fn add_tag(&mut self, t: Tag) -> TagResult<()> {
        self.entry.add_tag(t)
    }

    fn remove_tag(&mut self, t: Tag) -> TagResult<()> {
        self.entry.remove_tag(t)
    }

    fn has_tag(&self, t: TagSlice) -> TagResult<bool> {
        self.entry.has_tag(t)
    }

    fn has_tags(&self, ts: &[Tag]) -> TagResult<bool> {
        self.entry.has_tags(ts)
    }

}

trait FromStoreId {
    fn from_storeid(&Store, StoreId) -> Result<Note>;
}

impl<'a> FromStoreId for Note<'a> {

    fn from_storeid(store: &Store, id: StoreId) -> Result<Note> {
        debug!("Loading note from storeid: '{:?}'", id);
        match store.retrieve(id) {
            Err(e)    => Err(NE::new(NEK::StoreReadError, Some(Box::new(e)))),
            Ok(entry) => Ok(Note { entry: entry }),
        }
    }

}

impl<'a> Deref for Note<'a> {

    type Target = FileLockEntry<'a>;

    fn deref(&self) -> &FileLockEntry<'a> {
        &self.entry
    }

}

#[derive(Debug)]
pub struct NoteIterator<'a> {
    store: &'a Store,
    iditer: StoreIdIterator,
}

impl<'a> NoteIterator<'a> {

    pub fn new(store: &'a Store, iditer: StoreIdIterator) -> NoteIterator<'a> {
        NoteIterator {
            store: store,
            iditer: iditer,
        }
    }

}

impl<'a> Iterator for NoteIterator<'a> {
    type Item = Result<Note<'a>>;

    fn next(&mut self) -> Option<Result<Note<'a>>> {
        self.iditer
            .next()
            .map(|id| Note::from_storeid(self.store, id))
    }

}

<|endoftext|>


<commit_before>use std::error::Error;
use std::io::Write;
use std::io::stderr;

/// Print an Error type and its cause recursively
///
/// The error is printed with "Error NNNN :" as prefix, where "NNNN" is a number which increases
/// which each recursion into the errors cause. The error description is used to visualize what
/// failed and if there is a cause "-- caused by:" is appended, and the cause is printed on the next
/// line.
///
/// Example output:
///
/// ```ignore
/// Error    1 : Some error -- caused by:
/// Error    2 : Some other error -- caused by:
/// Error    3 : Yet another Error -- caused by:
/// ...
///
/// Error <NNNN> : <Error description>
/// ```
pub fn trace_error(e: &Error) {
    print_trace_maxdepth(count_error_causes(e), e, ::std::u64::MAX);
    write!(stderr(), "\n");
}

/// Print an Error type and its cause recursively, but only `max` levels
///
/// Output is the same as for `trace_error()`, though there are only `max` levels printed.
pub fn trace_error_maxdepth(e: &Error, max: u64) {
    let n = count_error_causes(e);
    write!(stderr(), "{}/{} Levels of errors will be printed\n", (if max > n { n } else { max }), n);
    print_trace_maxdepth(n, e, max);
    write!(stderr(), "");
}

/// Print an Error type and its cause recursively with the debug!() macro
///
/// Output is the same as for `trace_error()`.
pub fn trace_error_dbg(e: &Error) {
    print_trace_dbg(0, e);
}

/// Helper function for `trace_error()` and `trace_error_maxdepth()`.
///
/// Returns the cause of the last processed error in the recursion, so `None` if all errors where
/// processed.
fn print_trace_maxdepth(idx: u64, e: &Error, max: u64) -> Option<&Error> {
    if e.cause().is_some() && idx > 0 {
        match print_trace_maxdepth(idx - 1, e.cause().unwrap(), max) {
            None => write!(stderr(), "\n"),
            Some(_) => write!(stderr(), " -- caused:\n"),
        };
    } else {
        write!(stderr(), "\n");
    }
    write!(stderr(), "Error {:>4} : {}", idx, e.description());
    e.cause()
}

/// Count errors in Error::cause() recursively
fn count_error_causes(e: &Error) -> u64 {
    1 + if e.cause().is_some() { count_error_causes(e.cause().unwrap()) } else { 0 }
}

fn print_trace_dbg(idx: u64, e: &Error) {
    debug!("Error {:>4} : {}", idx, e.description());
    if e.cause().is_some() {
        debug!(" -- caused by:");
        print_trace_dbg(idx + 1, e.cause().unwrap());
    }
}

<commit_msg>Use unused results by calling write!() with .ok()<commit_after>use std::error::Error;
use std::io::Write;
use std::io::stderr;

/// Print an Error type and its cause recursively
///
/// The error is printed with "Error NNNN :" as prefix, where "NNNN" is a number which increases
/// which each recursion into the errors cause. The error description is used to visualize what
/// failed and if there is a cause "-- caused by:" is appended, and the cause is printed on the next
/// line.
///
/// Example output:
///
/// ```ignore
/// Error    1 : Some error -- caused by:
/// Error    2 : Some other error -- caused by:
/// Error    3 : Yet another Error -- caused by:
/// ...
///
/// Error <NNNN> : <Error description>
/// ```
pub fn trace_error(e: &Error) {
    print_trace_maxdepth(count_error_causes(e), e, ::std::u64::MAX);
    write!(stderr(), "\n").ok();
}

/// Print an Error type and its cause recursively, but only `max` levels
///
/// Output is the same as for `trace_error()`, though there are only `max` levels printed.
pub fn trace_error_maxdepth(e: &Error, max: u64) {
    let n = count_error_causes(e);
    write!(stderr(),
    "{}/{} Levels of errors will be printed\n", (if max > n { n } else { max }), n).ok();
    print_trace_maxdepth(n, e, max);
    write!(stderr(), "").ok();
}

/// Print an Error type and its cause recursively with the debug!() macro
///
/// Output is the same as for `trace_error()`.
pub fn trace_error_dbg(e: &Error) {
    print_trace_dbg(0, e);
}

/// Helper function for `trace_error()` and `trace_error_maxdepth()`.
///
/// Returns the cause of the last processed error in the recursion, so `None` if all errors where
/// processed.
fn print_trace_maxdepth(idx: u64, e: &Error, max: u64) -> Option<&Error> {
    if e.cause().is_some() && idx > 0 {
        match print_trace_maxdepth(idx - 1, e.cause().unwrap(), max) {
            None    => write!(stderr(), "\n").ok(),
            Some(_) => write!(stderr(), " -- caused:\n").ok(),
        };
    } else {
        write!(stderr(), "\n").ok();
    }
    write!(stderr(), "Error {:>4} : {}", idx, e.description()).ok();
    e.cause()
}

/// Count errors in Error::cause() recursively
fn count_error_causes(e: &Error) -> u64 {
    1 + if e.cause().is_some() { count_error_causes(e.cause().unwrap()) } else { 0 }
}

fn print_trace_dbg(idx: u64, e: &Error) {
    debug!("Error {:>4} : {}", idx, e.description());
    if e.cause().is_some() {
        debug!(" -- caused by:");
        print_trace_dbg(idx + 1, e.cause().unwrap());
    }
}

<|endoftext|>


<commit_before>#[cfg(feature="serde-1")]
pub(super) mod rand_size_serde {
    const RAND_SIZE_LEN: usize = 8;
    const RAND_SIZE: usize = 1 << RAND_SIZE_LEN;

    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use serde::de::{Visitor,SeqAccess};
    use serde::de;

    use std::fmt;

    pub fn serialize<T, S>(arr: &[T;RAND_SIZE], ser: S) -> Result<S::Ok, S::Error>
    where
        T: Serialize,
        S: Serializer
    {
        use serde::ser::SerializeTuple;

        let mut seq = ser.serialize_tuple(RAND_SIZE)?;

        for e in arr.iter() {
            seq.serialize_element(&e)?;
        }

        seq.end()
    }

    #[inline]
    pub fn deserialize<'de, T, D>(de: D) -> Result<[T;RAND_SIZE], D::Error>
    where
        T: Deserialize<'de>+Default+Copy,
        D: Deserializer<'de>,
    {
        use std::marker::PhantomData;
        struct ArrayVisitor<T> {
            _pd: PhantomData<T>,
        };
        impl<'de,T> Visitor<'de> for ArrayVisitor<T>
        where
            T: Deserialize<'de>+Default+Copy
        {
            type Value = [T; RAND_SIZE];

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("Isaac state array")
            }

            #[inline]
            fn visit_seq<A>(self, mut seq: A) -> Result<[T; RAND_SIZE], A::Error>
            where
                A: SeqAccess<'de>,
            {
                let mut out = [Default::default();RAND_SIZE];

                for i in 0..RAND_SIZE {
                    match seq.next_element()? {
                        Some(val) => out[i] = val,
                        None => return Err(de::Error::invalid_length(i, &self)),
                    };
                }

                Ok(out)
            }
        }

        de.deserialize_tuple(RAND_SIZE, ArrayVisitor{_pd: PhantomData})
    }
}
<commit_msg>Add a header and remove redundant gate<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! ISAAC serde helper functions.

pub(super) mod rand_size_serde {
    const RAND_SIZE_LEN: usize = 8;
    const RAND_SIZE: usize = 1 << RAND_SIZE_LEN;

    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use serde::de::{Visitor,SeqAccess};
    use serde::de;

    use std::fmt;

    pub fn serialize<T, S>(arr: &[T;RAND_SIZE], ser: S) -> Result<S::Ok, S::Error>
    where
        T: Serialize,
        S: Serializer
    {
        use serde::ser::SerializeTuple;

        let mut seq = ser.serialize_tuple(RAND_SIZE)?;

        for e in arr.iter() {
            seq.serialize_element(&e)?;
        }

        seq.end()
    }

    #[inline]
    pub fn deserialize<'de, T, D>(de: D) -> Result<[T;RAND_SIZE], D::Error>
    where
        T: Deserialize<'de>+Default+Copy,
        D: Deserializer<'de>,
    {
        use std::marker::PhantomData;
        struct ArrayVisitor<T> {
            _pd: PhantomData<T>,
        };
        impl<'de,T> Visitor<'de> for ArrayVisitor<T>
        where
            T: Deserialize<'de>+Default+Copy
        {
            type Value = [T; RAND_SIZE];

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("Isaac state array")
            }

            #[inline]
            fn visit_seq<A>(self, mut seq: A) -> Result<[T; RAND_SIZE], A::Error>
            where
                A: SeqAccess<'de>,
            {
                let mut out = [Default::default();RAND_SIZE];

                for i in 0..RAND_SIZE {
                    match seq.next_element()? {
                        Some(val) => out[i] = val,
                        None => return Err(de::Error::invalid_length(i, &self)),
                    };
                }

                Ok(out)
            }
        }

        de.deserialize_tuple(RAND_SIZE, ArrayVisitor{_pd: PhantomData})
    }
}
<|endoftext|>


<commit_before>//! This module contains functions for kmerizing a longer sequence and various
//! utilities for dealing with these kmers.
//!
//!
use std::borrow::Cow;

#[inline]
fn complement(n: &u8) -> u8 {
    //! Returns the complementary base for a given IUPAC base code.
    //!
    //! Does not work for RNA sequences (maybe we should raise an error or something?)
    match *n as char {
        'a' => 't' as u8,
        'c' => 'g' as u8,
        'g' => 'c' as u8,
        't' => 'a' as u8,
        'A' => 'T' as u8,
        'C' => 'G' as u8,
        'G' => 'C' as u8,
        'T' => 'A' as u8,

        // IUPAC codes
        'r' => 'y' as u8,
        'y' => 'r' as u8,
        's' => 's' as u8,
        'w' => 'w' as u8,
        'k' => 'm' as u8,
        'm' => 'k' as u8,
        'b' => 'v' as u8,
        'd' => 'h' as u8,
        'h' => 'd' as u8,
        'v' => 'b' as u8,
        'R' => 'Y' as u8,
        'Y' => 'R' as u8,
        'S' => 'S' as u8,
        'W' => 'W' as u8,
        'K' => 'M' as u8,
        'M' => 'K' as u8,
        'B' => 'V' as u8,
        'D' => 'H' as u8,
        'H' => 'D' as u8,
        'V' => 'B' as u8,

        // anything else just pass through
        // 'u' | 'U' => panic!("Does not support complements of U"),
        x => x as u8,
    }
}

pub fn normalize<'a>(seq: &'a [u8], iupac: bool) -> Cow<'a, [u8]> {
    //! Transform a FASTX sequence into it's "normalized" form.
    //!
    //! The normalized form is:
    //!  - only AGCTN and possibly . (for gaps)
    //!  - lowercase versions of these are uppercased
    //!  - U is converted to T (make everything a DNA sequence)
    //!  - some other punctuation is converted to gaps
    //!  - IUPAC bases may be converted to N's depending on the parameter passed in
    //!  - everything else is considered a N
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    let mut original_was_bad = false;

    for n in seq.iter() {
        let mut good_char = false;
        buf.push(match (*n as char, iupac) {
            c @ ('A', _) | c @ ('C', _) | c @ ('G', _) | c @ ('T', _) | c @ ('N', _) |
            c @ ('.', _) => {
                good_char = true;
                c.0 as u8
            },
            ('a', _) => 'A' as u8,
            ('c', _) => 'C' as u8,
            ('g', _) => 'G' as u8,
            // normalize uridine to thymine
            ('t', _) | ('u', _) | ('U', _) => 'T' as u8,
            ('-', _) | ('~', _) | (' ', _) => '.' as u8,
            // logic for IUPAC bases (a little messy)
            c @ ('B', true) | c @ ('D', true) | c @ ('H', true) | c @ ('V', true) |
            c @ ('R', true) | c @ ('Y', true) | c @ ('S', true) | c @ ('W', true) |
            c @ ('K', true) | c @ ('M', true) => {
                good_char = true;
                c.0 as u8
            },
            ('b', true) => 'B' as u8,
            ('d', true) => 'D' as u8,
            ('h', true) => 'H' as u8,
            ('v', true) => 'V' as u8,
            ('r', true) => 'R' as u8,
            ('y', true) => 'Y' as u8,
            ('s', true) => 'S' as u8,
            ('w', true) => 'W' as u8,
            ('k', true) => 'K' as u8,
            ('m', true) => 'M' as u8,
            _ => 'N' as u8,
        });
        if !good_char {
            original_was_bad = true;
        }
    }
    match original_was_bad {
        true => Cow::Owned(buf),
        false => Cow::Borrowed(seq),
    }
}

pub fn canonical<'a>(seq: &'a [u8]) -> Cow<'a, [u8]> {
    //! Taking in a sequence string, return the canonical form of the sequence
    //! (e.g. the lexigraphically lowest of either the original sequence or its
    //! reverse complement)
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    // enough just keeps our comparisons from happening after they need to
    let mut enough = false;
    let mut original_was_canonical = false;

    // loop through the kmer and its reverse complement simultaneously
    for (rn, n) in seq.iter().rev().map(|n| complement(n)).zip(seq.iter()) {
        buf.push(rn);
        if !enough && n < &rn {
            original_was_canonical = true;
            break;
        } else if !enough && &rn < n {
            enough = true;
        }
        // unstated if branch: if rn == n, keep comparing
    }
    match (original_was_canonical, enough) {
        (true, true) => panic!("Bug: should never set original_was_canonical if enough == true"),
        (true, false) => Cow::Borrowed(seq),
        (false, true) => Cow::Owned(buf),
        // the sequences were completely equal, return the ref
        (false, false) => Cow::Borrowed(seq),
    }
}

#[test]
fn can_canonicalize() {
    assert!(canonical(b"A") == Cow::Borrowed(b"A"));
    assert!(canonical(b"T") == Cow::Owned::<[u8]>(b"A".to_vec()));
    assert!(canonical(b"AAGT") == Cow::Borrowed(b"AAGT"));
    assert!(canonical(b"ACTT") == Cow::Owned::<[u8]>(b"AAGT".to_vec()));
    assert!(canonical(b"GC") == Cow::Borrowed(b"GC"));
}

/// Find the lexigraphically smallest substring of `seq` of length `length`
///
/// There's probably a faster algorithm for this somewhere...
pub fn minimizer<'a>(seq: &'a [u8], length: usize) -> Cow<'a, [u8]> {
    let reverse_complement: Vec<u8> = seq.iter().rev().map(|n| complement(n)).collect();
    let mut minmer = Cow::Borrowed(&seq[..length]);

    for (kmer, rc_kmer) in seq.windows(length).zip(reverse_complement.windows(length)) {
        if kmer < &minmer[..] {
            minmer = Cow::Borrowed(kmer);
        }
        if rc_kmer < &minmer[..] {
            minmer = Cow::Owned(rc_kmer.to_vec());
        }
    }
    minmer
}

#[test]
fn can_minimize() {
    let minmer = minimizer(&b"ATTTCG"[..], 3);
    assert_eq!(&minmer[..], b"AAA");
}

// TODO
// pub fn skip_n<'a, T>(iter: T) -> T where T: Iterator<Item=&'a [u8]> {
//    iter.filter(|kmer| kmer.contains(&('N' as u8)) || kmer.contains(&('n' as u8)))
// }

fn is_good_base(chr: u8) -> bool {
    match chr as char {
        'a' | 'c' | 'g' | 't' | 'A' | 'C' | 'G' | 'T' => true,
        _ => false,
    }
}

fn has_no_n<'a>(seq: &'a [u8]) -> bool {
    //! Determines if a sequence has any non-primary four bases
    //! characters in it
    seq.iter().all(|n| is_good_base(*n))
}

#[test]
fn can_detect_no_n() {
    assert!(has_no_n(b"AAGT"));
    assert!(!has_no_n(b"NAGT"));
}

pub struct NuclKmer<'a> {
    k: u8,
    start_pos: usize,
    buffer: &'a [u8],
}

fn update_position(start_pos: &mut usize, k: u8, buffer: &[u8], initial: bool) -> bool {
    // check if we have enough "physical" space for one more kmer
    if *start_pos + k as usize > buffer.len() {
        return false;
    }

    let mut kmer_len = (k - 1) as usize;
    let mut stop_len = k as usize;
    if initial {
        kmer_len = 0;
        stop_len = (k - 1) as usize;
    }

    while kmer_len < stop_len {
        if is_good_base(buffer[*start_pos + kmer_len]) {
            kmer_len += 1;
        } else {
            kmer_len = 0;
            *start_pos += kmer_len + 1;
            if *start_pos + k as usize > buffer.len() {
                return false;
            }
        }
    }
    true
}

impl<'a> NuclKmer<'a> {
    //! A kmer-izer for a nucleotide/amino acid sequence; returning slices to the original data
    pub fn new(buffer: &'a [u8], k: u8) -> NuclKmer<'a> {
        let mut start_pos = 0;
        update_position(&mut start_pos, k, buffer, true);
        NuclKmer {
            k: k,
            start_pos: start_pos,
            buffer: buffer,
        }
    }
}

impl<'a> Iterator for NuclKmer<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<&'a [u8]> {
        if !update_position(&mut self.start_pos, self.k, self.buffer, false) {
            return None;
        }

        let result = Some(&self.buffer[self.start_pos..self.start_pos + self.k as usize]);
        self.start_pos += 1;
        result
    }
}

#[test]
fn can_kmerize() {
    // test general function
    let mut i = 0;
    for k in NuclKmer::new(b"AGCT", 1) {
        match i {
            0 => assert_eq!(k, &b"A"[..]),
            1 => assert_eq!(k, &b"G"[..]),
            2 => assert_eq!(k, &b"C"[..]),
            3 => assert_eq!(k, &b"T"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that we skip over N's
    i = 0;
    for k in NuclKmer::new(b"ACNGT", 2) {
        match i {
            0 => assert_eq!(k, &b"AC"[..]),
            1 => assert_eq!(k, &b"GT"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that we skip over N's and handle short kmers
    i = 0;
    for k in NuclKmer::new(b"ACNG", 2) {
        match i {
            0 => assert_eq!(k, &b"AC"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that the minimum length works
    for k in NuclKmer::new(b"AC", 2) {
        assert_eq!(k, &b"AC"[..]);
    }
}
<commit_msg>Make has_no_n public<commit_after>//! This module contains functions for kmerizing a longer sequence and various
//! utilities for dealing with these kmers.
//!
//!
use std::borrow::Cow;

#[inline]
fn complement(n: &u8) -> u8 {
    //! Returns the complementary base for a given IUPAC base code.
    //!
    //! Does not work for RNA sequences (maybe we should raise an error or something?)
    match *n as char {
        'a' => 't' as u8,
        'c' => 'g' as u8,
        'g' => 'c' as u8,
        't' => 'a' as u8,
        'A' => 'T' as u8,
        'C' => 'G' as u8,
        'G' => 'C' as u8,
        'T' => 'A' as u8,

        // IUPAC codes
        'r' => 'y' as u8,
        'y' => 'r' as u8,
        's' => 's' as u8,
        'w' => 'w' as u8,
        'k' => 'm' as u8,
        'm' => 'k' as u8,
        'b' => 'v' as u8,
        'd' => 'h' as u8,
        'h' => 'd' as u8,
        'v' => 'b' as u8,
        'R' => 'Y' as u8,
        'Y' => 'R' as u8,
        'S' => 'S' as u8,
        'W' => 'W' as u8,
        'K' => 'M' as u8,
        'M' => 'K' as u8,
        'B' => 'V' as u8,
        'D' => 'H' as u8,
        'H' => 'D' as u8,
        'V' => 'B' as u8,

        // anything else just pass through
        // 'u' | 'U' => panic!("Does not support complements of U"),
        x => x as u8,
    }
}

pub fn normalize<'a>(seq: &'a [u8], iupac: bool) -> Cow<'a, [u8]> {
    //! Transform a FASTX sequence into it's "normalized" form.
    //!
    //! The normalized form is:
    //!  - only AGCTN and possibly . (for gaps)
    //!  - lowercase versions of these are uppercased
    //!  - U is converted to T (make everything a DNA sequence)
    //!  - some other punctuation is converted to gaps
    //!  - IUPAC bases may be converted to N's depending on the parameter passed in
    //!  - everything else is considered a N
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    let mut original_was_bad = false;

    for n in seq.iter() {
        let mut good_char = false;
        buf.push(match (*n as char, iupac) {
            c @ ('A', _) | c @ ('C', _) | c @ ('G', _) | c @ ('T', _) | c @ ('N', _) |
            c @ ('.', _) => {
                good_char = true;
                c.0 as u8
            },
            ('a', _) => 'A' as u8,
            ('c', _) => 'C' as u8,
            ('g', _) => 'G' as u8,
            // normalize uridine to thymine
            ('t', _) | ('u', _) | ('U', _) => 'T' as u8,
            ('-', _) | ('~', _) | (' ', _) => '.' as u8,
            // logic for IUPAC bases (a little messy)
            c @ ('B', true) | c @ ('D', true) | c @ ('H', true) | c @ ('V', true) |
            c @ ('R', true) | c @ ('Y', true) | c @ ('S', true) | c @ ('W', true) |
            c @ ('K', true) | c @ ('M', true) => {
                good_char = true;
                c.0 as u8
            },
            ('b', true) => 'B' as u8,
            ('d', true) => 'D' as u8,
            ('h', true) => 'H' as u8,
            ('v', true) => 'V' as u8,
            ('r', true) => 'R' as u8,
            ('y', true) => 'Y' as u8,
            ('s', true) => 'S' as u8,
            ('w', true) => 'W' as u8,
            ('k', true) => 'K' as u8,
            ('m', true) => 'M' as u8,
            _ => 'N' as u8,
        });
        if !good_char {
            original_was_bad = true;
        }
    }
    match original_was_bad {
        true => Cow::Owned(buf),
        false => Cow::Borrowed(seq),
    }
}

pub fn canonical<'a>(seq: &'a [u8]) -> Cow<'a, [u8]> {
    //! Taking in a sequence string, return the canonical form of the sequence
    //! (e.g. the lexigraphically lowest of either the original sequence or its
    //! reverse complement)
    let mut buf: Vec<u8> = Vec::with_capacity(seq.len());
    // enough just keeps our comparisons from happening after they need to
    let mut enough = false;
    let mut original_was_canonical = false;

    // loop through the kmer and its reverse complement simultaneously
    for (rn, n) in seq.iter().rev().map(|n| complement(n)).zip(seq.iter()) {
        buf.push(rn);
        if !enough && n < &rn {
            original_was_canonical = true;
            break;
        } else if !enough && &rn < n {
            enough = true;
        }
        // unstated if branch: if rn == n, keep comparing
    }
    match (original_was_canonical, enough) {
        (true, true) => panic!("Bug: should never set original_was_canonical if enough == true"),
        (true, false) => Cow::Borrowed(seq),
        (false, true) => Cow::Owned(buf),
        // the sequences were completely equal, return the ref
        (false, false) => Cow::Borrowed(seq),
    }
}

#[test]
fn can_canonicalize() {
    assert!(canonical(b"A") == Cow::Borrowed(b"A"));
    assert!(canonical(b"T") == Cow::Owned::<[u8]>(b"A".to_vec()));
    assert!(canonical(b"AAGT") == Cow::Borrowed(b"AAGT"));
    assert!(canonical(b"ACTT") == Cow::Owned::<[u8]>(b"AAGT".to_vec()));
    assert!(canonical(b"GC") == Cow::Borrowed(b"GC"));
}

/// Find the lexigraphically smallest substring of `seq` of length `length`
///
/// There's probably a faster algorithm for this somewhere...
pub fn minimizer<'a>(seq: &'a [u8], length: usize) -> Cow<'a, [u8]> {
    let reverse_complement: Vec<u8> = seq.iter().rev().map(|n| complement(n)).collect();
    let mut minmer = Cow::Borrowed(&seq[..length]);

    for (kmer, rc_kmer) in seq.windows(length).zip(reverse_complement.windows(length)) {
        if kmer < &minmer[..] {
            minmer = Cow::Borrowed(kmer);
        }
        if rc_kmer < &minmer[..] {
            minmer = Cow::Owned(rc_kmer.to_vec());
        }
    }
    minmer
}

#[test]
fn can_minimize() {
    let minmer = minimizer(&b"ATTTCG"[..], 3);
    assert_eq!(&minmer[..], b"AAA");
}

// TODO
// pub fn skip_n<'a, T>(iter: T) -> T where T: Iterator<Item=&'a [u8]> {
//    iter.filter(|kmer| kmer.contains(&('N' as u8)) || kmer.contains(&('n' as u8)))
// }

fn is_good_base(chr: u8) -> bool {
    match chr as char {
        'a' | 'c' | 'g' | 't' | 'A' | 'C' | 'G' | 'T' => true,
        _ => false,
    }
}

pub fn has_no_n<'a>(seq: &'a [u8]) -> bool {
    //! Determines if a sequence has any non-primary four bases
    //! characters in it
    seq.iter().all(|n| is_good_base(*n))
}

#[test]
fn can_detect_no_n() {
    assert!(has_no_n(b"AAGT"));
    assert!(!has_no_n(b"NAGT"));
}

pub struct NuclKmer<'a> {
    k: u8,
    start_pos: usize,
    buffer: &'a [u8],
}

fn update_position(start_pos: &mut usize, k: u8, buffer: &[u8], initial: bool) -> bool {
    // check if we have enough "physical" space for one more kmer
    if *start_pos + k as usize > buffer.len() {
        return false;
    }

    let mut kmer_len = (k - 1) as usize;
    let mut stop_len = k as usize;
    if initial {
        kmer_len = 0;
        stop_len = (k - 1) as usize;
    }

    while kmer_len < stop_len {
        if is_good_base(buffer[*start_pos + kmer_len]) {
            kmer_len += 1;
        } else {
            kmer_len = 0;
            *start_pos += kmer_len + 1;
            if *start_pos + k as usize > buffer.len() {
                return false;
            }
        }
    }
    true
}

impl<'a> NuclKmer<'a> {
    //! A kmer-izer for a nucleotide/amino acid sequence; returning slices to the original data
    pub fn new(buffer: &'a [u8], k: u8) -> NuclKmer<'a> {
        let mut start_pos = 0;
        update_position(&mut start_pos, k, buffer, true);
        NuclKmer {
            k: k,
            start_pos: start_pos,
            buffer: buffer,
        }
    }
}

impl<'a> Iterator for NuclKmer<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<&'a [u8]> {
        if !update_position(&mut self.start_pos, self.k, self.buffer, false) {
            return None;
        }

        let result = Some(&self.buffer[self.start_pos..self.start_pos + self.k as usize]);
        self.start_pos += 1;
        result
    }
}

#[test]
fn can_kmerize() {
    // test general function
    let mut i = 0;
    for k in NuclKmer::new(b"AGCT", 1) {
        match i {
            0 => assert_eq!(k, &b"A"[..]),
            1 => assert_eq!(k, &b"G"[..]),
            2 => assert_eq!(k, &b"C"[..]),
            3 => assert_eq!(k, &b"T"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that we skip over N's
    i = 0;
    for k in NuclKmer::new(b"ACNGT", 2) {
        match i {
            0 => assert_eq!(k, &b"AC"[..]),
            1 => assert_eq!(k, &b"GT"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that we skip over N's and handle short kmers
    i = 0;
    for k in NuclKmer::new(b"ACNG", 2) {
        match i {
            0 => assert_eq!(k, &b"AC"[..]),
            _ => assert!(false),
        }
        i += 1;
    }

    // test that the minimum length works
    for k in NuclKmer::new(b"AC", 2) {
        assert_eq!(k, &b"AC"[..]);
    }
}
<|endoftext|>


<commit_before>
use nom::{HexDisplay, IResult, Needed, ErrorKind};

struct Document {
    header: Header,
    body: Vec<Element>,
}

struct Header {}

#[derive(Debug)]
enum ElementData {
    Signed(i64),
    Unsigned(u64),
    Float(f64),
    PlainString(String),
    UTF8String(String),
    Date(u64),
    Master(Vec<Element>),
    Binary(Vec<u8>),
    Unknown(u64),
}

#[derive(Debug)]
pub struct Element {
    id: u64,
    size: u64,
    data: ElementData,
}

pub fn vint(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = (v ^ (1 << (7 - len))) as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

// The ID are represented in the specification as their binary representation
// do not drop the marker bit.
pub fn vid(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = v as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

fn parse_master(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    map!(input,
         many0!(parse_element),
         |elem| ElementData::Master(elem))
}

fn parse_uint(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], ElementData::Unsigned(val))
}

pub fn parse_uint_data(input: &[u8], size: u64) -> IResult<&[u8], u64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], val)
}

pub fn parse_int_data(input: &[u8], size: u64) -> IResult<&[u8], i64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    //FIXME: is that right?
    IResult::Done(&input[size as usize..], val as i64)
}

use std::str;

fn parse_str(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( ElementData::PlainString(String::from_utf8(s.to_owned()).unwrap()) )
    )
}

pub fn parse_str_data(input: &[u8], size: u64) -> IResult<&[u8], String> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( String::from_utf8(s.to_owned()).unwrap() ) //FIXME: maybe do not unwrap here
    )
}

pub fn parse_binary_data(input: &[u8], size: u64) -> IResult<&[u8], Vec<u8>> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( s.to_owned() )
    )
}

//FIXME: handle default values
//FIXME: is that really following IEEE_754-1985 ?
pub fn parse_float_data(input: &[u8], size: u64) -> IResult<&[u8], f64> {
    use nom::{be_f32,be_f64};
    if size == 0 {
      IResult::Done(input, 0f64)
    } else if size == 4 {
      map!(input, flat_map!(take!(4), be_f32), |val| val as f64)
    } else if size == 8 {
      flat_map!(input, take!(8), be_f64)
    } else {
      IResult::Error(ErrorKind::Custom(1))
    }
}

fn parse_element_id(input: &[u8], id: u64, size: u64) -> IResult<&[u8], ElementData> {
    // println!("id: 0x{:X} size: {}", id, size);
    if input.len() < size as usize {
        return IResult::Incomplete(Needed::Size(size as usize));
    }

    match id {
        0x1A45DFA3 => parse_master(input, size),
        0x4286 => parse_uint(input, size),
        0x42F7 => parse_uint(input, size),
        0x42F2 => parse_uint(input, size),
        0x42F3 => parse_uint(input, size),
        0x4282 => parse_str(input, size),
        0x4287 => parse_uint(input, size),
        0x4285 => parse_uint(input, size),
        _ => IResult::Done(&input[size as usize..], ElementData::Unknown(id)),
    }
}

named!(pub parse_element<Element>,
    do_parse!(
        id : vid >>
        size: vint >>
        data: call!(parse_element_id, id, size) >>
        (Element { id, size, data })
    )
);

#[macro_export]
macro_rules! ebml_uint (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_uint_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_uint_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_int (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_int_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_int_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_float (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_float_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_float_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_str (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_str_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_str_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_binary (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_binary_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_binary_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_master (
  ($i: expr, $id:expr, $submac:ident!( $($args:tt)* )) => ({
    use $crate::ebml::{vid, vint};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: flat_map!(take!(size as usize), $submac!($($args)*))
      >> (data)
    )
  })
);

#[derive(Debug,Clone,PartialEq)]
pub struct EBMLHeader {
    pub version: u64,
    pub read_version: u64,
    pub max_id_length: u64,
    pub max_size_length: u64,
    pub doc_type: String,
    pub doc_type_version: u64,
    pub doc_type_read_version: u64,
}

named!(pub ebml_header<EBMLHeader>,
  ebml_master!(0x1A45DFA3,
    do_parse!(
      t: permutation!(
        ebml_uint!(0x4286), // version
        ebml_uint!(0x42F7), // read_version
        ebml_uint!(0x42F2), // max id length
        ebml_uint!(0x42F3), // max size length
        ebml_str!(0x4282),  // doctype
        ebml_uint!(0x4287), // doctype version
        ebml_uint!(0x4285)  // doctype_read version
      ) >>
      ({
        EBMLHeader {
          version:               t.0,
          read_version:          t.1,
          max_id_length:         t.2,
          max_size_length:       t.3,
          doc_type:              t.4,
          doc_type_version:      t.5,
          doc_type_read_version: t.6,

        }
      })
    )
  )
);

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

    const single_stream: &'static [u8] = include_bytes!("../assets/single_stream.mkv");

    #[test]
    fn variable_integer() {
        let val01 = [0b10000000];
        //        let val01 = [0b01000000, 0b1];

        match vint(&val01) {
            IResult::Done(_, v) => assert!(0 == v),
            _ => panic!(),
        }
    }

    /*
    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:#?}", parse_element(single_stream));
        panic!();
    }*/

    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:?}", ebml_header(&single_stream[..100]));
        panic!();
    }
}
<commit_msg>format<commit_after>
use nom::{HexDisplay, IResult, Needed, ErrorKind};

struct Document {
    header: Header,
    body: Vec<Element>,
}

struct Header {}

#[derive(Debug)]
enum ElementData {
    Signed(i64),
    Unsigned(u64),
    Float(f64),
    PlainString(String),
    UTF8String(String),
    Date(u64),
    Master(Vec<Element>),
    Binary(Vec<u8>),
    Unknown(u64),
}

#[derive(Debug)]
pub struct Element {
    id: u64,
    size: u64,
    data: ElementData,
}

pub fn vint(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = (v ^ (1 << (7 - len))) as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

// The ID are represented in the specification as their binary representation
// do not drop the marker bit.
pub fn vid(input: &[u8]) -> IResult<&[u8], u64> {
    if input.len() == 0 {
        return IResult::Incomplete(Needed::Size(1));
    }

    let v = input[0];
    let len = v.leading_zeros();

    if len == 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    if input.len() <= len as usize {
        return IResult::Incomplete(Needed::Size(1));
    }

    let mut val = v as u64;

    // println!("vint {:08b} {:08b} {:08b} {}", val, v, (1 << (8 - len)), len);

    for i in 0..len as usize {
        val = (val << 8) | input[i + 1] as u64;
    }

    IResult::Done(&input[len as usize + 1..], val)
}

fn parse_master(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    map!(input,
         many0!(parse_element),
         |elem| ElementData::Master(elem))
}

fn parse_uint(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], ElementData::Unsigned(val))
}

pub fn parse_uint_data(input: &[u8], size: u64) -> IResult<&[u8], u64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    IResult::Done(&input[size as usize..], val)
}

pub fn parse_int_data(input: &[u8], size: u64) -> IResult<&[u8], i64> {
    let mut val = 0;

    if size > 8 {
        return IResult::Error(ErrorKind::Custom(1));
    }

    for i in 0..size as usize {
        val = (val << 8) | input[i] as u64;
    }

    //FIXME: is that right?
    IResult::Done(&input[size as usize..], val as i64)
}

use std::str;

fn parse_str(input: &[u8], size: u64) -> IResult<&[u8], ElementData> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( ElementData::PlainString(String::from_utf8(s.to_owned()).unwrap()) )
    )
}

pub fn parse_str_data(input: &[u8], size: u64) -> IResult<&[u8], String> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( String::from_utf8(s.to_owned()).unwrap() ) //FIXME: maybe do not unwrap here
    )
}

pub fn parse_binary_data(input: &[u8], size: u64) -> IResult<&[u8], Vec<u8>> {
    do_parse!(input,
        s: take_s!(size as usize) >>
        ( s.to_owned() )
    )
}

//FIXME: handle default values
//FIXME: is that really following IEEE_754-1985 ?
pub fn parse_float_data(input: &[u8], size: u64) -> IResult<&[u8], f64> {
    use nom::{be_f32, be_f64};
    if size == 0 {
        IResult::Done(input, 0f64)
    } else if size == 4 {
        map!(input, flat_map!(take!(4), be_f32), |val| val as f64)
    } else if size == 8 {
        flat_map!(input, take!(8), be_f64)
    } else {
        IResult::Error(ErrorKind::Custom(1))
    }
}

fn parse_element_id(input: &[u8], id: u64, size: u64) -> IResult<&[u8], ElementData> {
    // println!("id: 0x{:X} size: {}", id, size);
    if input.len() < size as usize {
        return IResult::Incomplete(Needed::Size(size as usize));
    }

    match id {
        0x1A45DFA3 => parse_master(input, size),
        0x4286 => parse_uint(input, size),
        0x42F7 => parse_uint(input, size),
        0x42F2 => parse_uint(input, size),
        0x42F3 => parse_uint(input, size),
        0x4282 => parse_str(input, size),
        0x4287 => parse_uint(input, size),
        0x4285 => parse_uint(input, size),
        _ => IResult::Done(&input[size as usize..], ElementData::Unknown(id)),
    }
}

named!(pub parse_element<Element>,
    do_parse!(
        id : vid >>
        size: vint >>
        data: call!(parse_element_id, id, size) >>
        (Element { id, size, data })
    )
);

#[macro_export]
macro_rules! ebml_uint (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_uint_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_uint_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_int (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_int_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_int_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_float (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_float_data};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_float_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_str (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_str_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_str_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_binary (
  ($i: expr, $id:expr) => ({
    use $crate::ebml::{vid, vint, parse_binary_data};

    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: apply!(parse_binary_data, size)
      >> (data)
    )
  })
);

#[macro_export]
macro_rules! ebml_master (
  ($i: expr, $id:expr, $submac:ident!( $($args:tt)* )) => ({
    use $crate::ebml::{vid, vint};
    do_parse!($i,
               verify!(vid, |val:u64| val == $id)
      >> size: vint
      >> data: flat_map!(take!(size as usize), $submac!($($args)*))
      >> (data)
    )
  })
);

#[derive(Debug,Clone,PartialEq)]
pub struct EBMLHeader {
    pub version: u64,
    pub read_version: u64,
    pub max_id_length: u64,
    pub max_size_length: u64,
    pub doc_type: String,
    pub doc_type_version: u64,
    pub doc_type_read_version: u64,
}

named!(pub ebml_header<EBMLHeader>,
  ebml_master!(0x1A45DFA3,
    do_parse!(
      t: permutation!(
        ebml_uint!(0x4286), // version
        ebml_uint!(0x42F7), // read_version
        ebml_uint!(0x42F2), // max id length
        ebml_uint!(0x42F3), // max size length
        ebml_str!(0x4282),  // doctype
        ebml_uint!(0x4287), // doctype version
        ebml_uint!(0x4285)  // doctype_read version
      ) >>
      ({
        EBMLHeader {
          version:               t.0,
          read_version:          t.1,
          max_id_length:         t.2,
          max_size_length:       t.3,
          doc_type:              t.4,
          doc_type_version:      t.5,
          doc_type_read_version: t.6,

        }
      })
    )
  )
);

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

    const single_stream: &'static [u8] = include_bytes!("../assets/single_stream.mkv");

    #[test]
    fn variable_integer() {
        let val01 = [0b10000000];
        //        let val01 = [0b01000000, 0b1];

        match vint(&val01) {
            IResult::Done(_, v) => assert!(0 == v),
            _ => panic!(),
        }
    }

    /*
    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:#?}", parse_element(single_stream));
        panic!();
    }*/

    #[test]
    fn header() {
        println!("{}", single_stream[..8].to_hex(8));
        println!("{:b} {:b}", single_stream[0], single_stream[1]);
        println!("{:?}", ebml_header(&single_stream[..100]));
        panic!();
    }
}
<|endoftext|>