<commit_before>extern crate c_ares_sys;
extern crate libc;

use std::ffi::{CStr, CString};
use std::mem;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::os::unix::io;
use std::ptr;

#[derive(Debug)]
pub enum AresError {
    ENODATA = 1,
    EFORMERR,
    ESERVFAIL,
    ENOTFOUND,
    ENOTIMP,
    EREFUSED,
    EBADQUERY,
    EBADNAME,
    EBADFAMILY,
    EBADRESP,
    ECONNREFUSED,
    ETIMEOUT,
    EOF,
    EFILE,
    ENOMEM,
    EDESTRUCTION,
    EBADSTR,
    EBADFLAGS,
    ENONAME,
    EBADHINTS,
    ENOTINITIALIZED,
    ELOADIPHLPAPI,
    EADDRGETNETWORKPARAMS,
    ECANCELLED,
    UNKNOWN,
}

pub const SOCKET_BAD: io::RawFd = c_ares_sys::ARES_SOCKET_BAD;

pub struct AResult {
    pub ip_addrs: Vec<Ipv4Addr>,
}

pub struct AAAAResult {
    pub ip_addrs: Vec<Ipv6Addr>,
}

pub struct Channel {
    ares_channel: c_ares_sys::ares_channel,
}

unsafe impl Send for Channel { }

extern "C" fn socket_callback<F>(
    data: *mut libc::c_void,
    socket_fd: c_ares_sys::ares_socket_t,
    readable: libc::c_int,
    writable: libc::c_int)
    where F: FnOnce(io::RawFd, bool, bool) + 'static {
    let handler: Box<F> = unsafe { mem::transmute(data) };
    handler(socket_fd as io::RawFd, readable != 0, writable != 0);
}

extern "C" fn query_a_callback<F>(
    arg: *mut libc::c_void,
    status: libc::c_int,
    _timeouts: libc::c_int,
    abuf: *mut libc::c_uchar,
    alen: libc::c_int)
    where F: FnOnce(Result<AResult, AresError>) + 'static {
    let result = if status != c_ares_sys::ARES_SUCCESS {
        Err(ares_error(status))
    } else {
        let mut hostent: *mut hostent = ptr::null_mut();
        let parse_status = unsafe {
            c_ares_sys::ares_parse_a_reply(
                abuf,
                alen,
                &mut hostent as *mut *mut _ as *mut *mut c_ares_sys::Struct_hostent,
                ptr::null_mut(),
                ptr::null_mut())
        };
        if parse_status != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(parse_status))
        } else {
            let mut answers = Vec::new();
            unsafe {
                let mut ptr = (*hostent).h_addr_list;
                while !(*ptr).is_null() {
                    let h_addr = *ptr;
                    let ip_addr = Ipv4Addr::new(
                        *h_addr as u8,
                        *h_addr.offset(1) as u8,
                        *h_addr.offset(2) as u8,
                        *h_addr.offset(3) as u8);
                    answers.push(ip_addr);
                    ptr = ptr.offset(1);
                }
                c_ares_sys::ares_free_hostent(hostent as *mut c_ares_sys::Struct_hostent);
            }
            Ok(AResult { ip_addrs: answers })
        }
    };

    let handler: Box<F> = unsafe { mem::transmute(arg) };
    handler(result);
}

extern "C" fn query_aaaa_callback<F>(
    arg: *mut libc::c_void,
    status: libc::c_int,
    _timeouts: libc::c_int,
    abuf: *mut libc::c_uchar,
    alen: libc::c_int)
    where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
    let result = if status != c_ares_sys::ARES_SUCCESS {
        Err(ares_error(status))
    } else {
        let mut hostent: *mut hostent = ptr::null_mut();
        let parse_status = unsafe {
            c_ares_sys::ares_parse_aaaa_reply(
                abuf,
                alen,
                &mut hostent as *mut *mut _ as *mut *mut c_ares_sys::Struct_hostent,
                ptr::null_mut(),
                ptr::null_mut())
        };
        if parse_status != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(parse_status))
        } else {
            let mut answers = Vec::new();
            unsafe {
                let mut ptr = (*hostent).h_addr_list;
                while !(*ptr).is_null() {
                    let h_addr = *ptr;
                    let ip_addr = Ipv6Addr::new(
                        ((*h_addr as u16) << 8) + *h_addr.offset(1) as u16,
                        ((*h_addr.offset(2) as u16) << 8) + *h_addr.offset(3) as u16,
                        ((*h_addr.offset(4) as u16) << 8) + *h_addr.offset(5) as u16,
                        ((*h_addr.offset(6) as u16) << 8) + *h_addr.offset(7) as u16,
                        ((*h_addr.offset(8) as u16) << 8) + *h_addr.offset(9) as u16,
                        ((*h_addr.offset(10) as u16) << 8) + *h_addr.offset(11) as u16,
                        ((*h_addr.offset(12) as u16) << 8) + *h_addr.offset(13) as u16,
                        ((*h_addr.offset(14) as u16) << 8) + *h_addr.offset(15) as u16);
                    answers.push(ip_addr);
                    ptr = ptr.offset(1);
                }
                c_ares_sys::ares_free_hostent(hostent as *mut c_ares_sys::Struct_hostent);
            }
            Ok(AAAAResult { ip_addrs: answers })
        }
    };

    let handler: Box<F> = unsafe { mem::transmute(arg) };
    handler(result);
}

impl Channel {
    pub fn new<F>(callback: F) -> Result<Channel, AresError>
        where F: FnOnce(io::RawFd, bool, bool) + 'static {
        let lib_rc = unsafe {
            c_ares_sys::ares_library_init(c_ares_sys::ARES_LIB_INIT_ALL)
        };
        match lib_rc {
            c_ares_sys::ARES_SUCCESS => (),
            _ => return Err(ares_error(lib_rc))
        }

        // TODO suport user-provided options
        let mut ares_channel = ptr::null_mut();
        let mut options = c_ares_sys::Struct_ares_options::default();
        options.flags = c_ares_sys::ARES_FLAG_STAYOPEN;
        options.timeout = 500;
        options.tries = 3;
        options.sock_state_cb = Some(socket_callback::<F>);
        options.sock_state_cb_data = unsafe {
            mem::transmute(Box::new(callback))
        };
        let optmask =
            c_ares_sys::ARES_OPT_FLAGS |
            c_ares_sys::ARES_OPT_TIMEOUT |
            c_ares_sys::ARES_OPT_TRIES |
            c_ares_sys::ARES_OPT_SOCK_STATE_CB;
        let channel_rc = unsafe {
            c_ares_sys::ares_init_options(&mut ares_channel, &mut options, optmask)
        };

        match channel_rc {
            c_ares_sys::ARES_SUCCESS => (),
            _ => {
                unsafe { c_ares_sys::ares_library_cleanup(); }
                return Err(ares_error(channel_rc))
            }
        }
        let channel = Channel {
            ares_channel: ares_channel,
        };

        // TODO ares_set_servers() here too?
        Ok(channel)
    }

    pub fn query_a<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::A as libc::c_int,
                Some(query_a_callback::<F>),
                c_arg);
        }
    }

    pub fn query_aaaa<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::AAAA as libc::c_int,
                Some(query_aaaa_callback::<F>),
                c_arg);
        }
    }

    pub fn process_fd(&mut self, read_fd: io::RawFd, write_fd: io::RawFd) {
        unsafe {
            c_ares_sys::ares_process_fd(self.ares_channel, read_fd, write_fd);
        }
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        unsafe {
            c_ares_sys::ares_destroy(self.ares_channel);
            c_ares_sys::ares_library_cleanup();
        }
    }
}

pub fn str_error<'a>(code: AresError) -> &'a str {
    let buf = unsafe {
        let ptr = c_ares_sys::ares_strerror(code as libc::c_int);
        CStr::from_ptr(ptr).to_bytes()
    };
    std::str::from_utf8(buf).unwrap()
}

// Convert an error code from the library into a more strongly typed AresError.
fn ares_error(code: libc::c_int) -> AresError {
    match code {
        c_ares_sys::ARES_ENODATA => AresError::ENODATA,
        c_ares_sys::ARES_EFORMERR => AresError::EFORMERR,
        c_ares_sys::ARES_ESERVFAIL => AresError::ESERVFAIL,
        c_ares_sys::ARES_ENOTFOUND => AresError::ENOTFOUND,
        c_ares_sys::ARES_ENOTIMP => AresError::ENOTIMP,
        c_ares_sys::ARES_EREFUSED => AresError::EREFUSED,
        c_ares_sys::ARES_EBADQUERY => AresError::EBADQUERY,
        c_ares_sys::ARES_EBADNAME => AresError::EBADNAME,
        c_ares_sys::ARES_EBADFAMILY => AresError::EBADFAMILY,
        c_ares_sys::ARES_EBADRESP => AresError::EBADRESP,
        c_ares_sys::ARES_ECONNREFUSED => AresError::ECONNREFUSED,
        c_ares_sys::ARES_ETIMEOUT => AresError::ETIMEOUT,
        c_ares_sys::ARES_EOF => AresError::EOF,
        c_ares_sys::ARES_EFILE => AresError::EFILE,
        c_ares_sys::ARES_ENOMEM => AresError::ENOMEM,
        c_ares_sys::ARES_EDESTRUCTION => AresError::EDESTRUCTION,
        c_ares_sys::ARES_EBADSTR => AresError::EBADSTR,
        c_ares_sys::ARES_EBADFLAGS => AresError::EBADFLAGS,
        c_ares_sys::ARES_ENONAME => AresError::ENONAME,
        c_ares_sys::ARES_EBADHINTS => AresError::EBADHINTS,
        c_ares_sys::ARES_ENOTINITIALIZED => AresError::ENOTINITIALIZED,
        c_ares_sys::ARES_ELOADIPHLPAPI => AresError::ELOADIPHLPAPI,
        c_ares_sys::ARES_EADDRGETNETWORKPARAMS => AresError::EADDRGETNETWORKPARAMS,
        c_ares_sys::ARES_ECANCELLED => AresError::ECANCELLED,
        _ => AresError::UNKNOWN,
    }
}

#[repr(C)]
struct hostent {
    h_name: *mut libc::c_char,
    h_aliases: *mut *mut libc::c_char,
    h_addrtype: libc::c_int,
    h_length: libc::c_int,
    h_addr_list: *mut *mut libc::c_char,
}

// See arpa/nameser.h
enum QueryType {
    A = 1,
    AAAA = 28,
}

// See arpa/nameser.h
enum DnsClass {
   IN = 1,
}
<commit_msg>Shuffle code around so that related pieces are near one another<commit_after>extern crate c_ares_sys;
extern crate libc;

use std::ffi::{CStr, CString};
use std::mem;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::os::unix::io;
use std::ptr;

#[derive(Debug)]
pub enum AresError {
    ENODATA = 1,
    EFORMERR,
    ESERVFAIL,
    ENOTFOUND,
    ENOTIMP,
    EREFUSED,
    EBADQUERY,
    EBADNAME,
    EBADFAMILY,
    EBADRESP,
    ECONNREFUSED,
    ETIMEOUT,
    EOF,
    EFILE,
    ENOMEM,
    EDESTRUCTION,
    EBADSTR,
    EBADFLAGS,
    ENONAME,
    EBADHINTS,
    ENOTINITIALIZED,
    ELOADIPHLPAPI,
    EADDRGETNETWORKPARAMS,
    ECANCELLED,
    UNKNOWN,
}

pub const SOCKET_BAD: io::RawFd = c_ares_sys::ARES_SOCKET_BAD;

pub struct AResult {
    pub ip_addrs: Vec<Ipv4Addr>,
}

pub struct AAAAResult {
    pub ip_addrs: Vec<Ipv6Addr>,
}

pub struct Channel {
    ares_channel: c_ares_sys::ares_channel,
}

unsafe impl Send for Channel { }

impl Channel {
    pub fn new<F>(callback: F) -> Result<Channel, AresError>
        where F: FnOnce(io::RawFd, bool, bool) + 'static {
        let lib_rc = unsafe {
            c_ares_sys::ares_library_init(c_ares_sys::ARES_LIB_INIT_ALL)
        };
        match lib_rc {
            c_ares_sys::ARES_SUCCESS => (),
            _ => return Err(ares_error(lib_rc))
        }

        // TODO suport user-provided options
        let mut ares_channel = ptr::null_mut();
        let mut options = c_ares_sys::Struct_ares_options::default();
        options.flags = c_ares_sys::ARES_FLAG_STAYOPEN;
        options.timeout = 500;
        options.tries = 3;
        options.sock_state_cb = Some(socket_callback::<F>);
        options.sock_state_cb_data = unsafe {
            mem::transmute(Box::new(callback))
        };
        let optmask =
            c_ares_sys::ARES_OPT_FLAGS |
            c_ares_sys::ARES_OPT_TIMEOUT |
            c_ares_sys::ARES_OPT_TRIES |
            c_ares_sys::ARES_OPT_SOCK_STATE_CB;
        let channel_rc = unsafe {
            c_ares_sys::ares_init_options(&mut ares_channel, &mut options, optmask)
        };

        match channel_rc {
            c_ares_sys::ARES_SUCCESS => (),
            _ => {
                unsafe { c_ares_sys::ares_library_cleanup(); }
                return Err(ares_error(channel_rc))
            }
        }
        let channel = Channel {
            ares_channel: ares_channel,
        };

        // TODO ares_set_servers() here too?
        Ok(channel)
    }

    pub fn query_a<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::A as libc::c_int,
                Some(query_a_callback::<F>),
                c_arg);
        }
    }

    pub fn query_aaaa<F>(&mut self, name: &str, handler: F)
        where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
        let c_name = CString::new(name).unwrap();
        unsafe {
            let c_arg: *mut libc::c_void = mem::transmute(Box::new(handler));
            c_ares_sys::ares_query(
                self.ares_channel,
                c_name.as_ptr(),
                DnsClass::IN as libc::c_int,
                QueryType::AAAA as libc::c_int,
                Some(query_aaaa_callback::<F>),
                c_arg);
        }
    }

    pub fn process_fd(&mut self, read_fd: io::RawFd, write_fd: io::RawFd) {
        unsafe {
            c_ares_sys::ares_process_fd(self.ares_channel, read_fd, write_fd);
        }
    }
}

impl Drop for Channel {
    fn drop(&mut self) {
        unsafe {
            c_ares_sys::ares_destroy(self.ares_channel);
            c_ares_sys::ares_library_cleanup();
        }
    }
}

pub fn str_error<'a>(code: AresError) -> &'a str {
    let buf = unsafe {
        let ptr = c_ares_sys::ares_strerror(code as libc::c_int);
        CStr::from_ptr(ptr).to_bytes()
    };
    std::str::from_utf8(buf).unwrap()
}

// Convert an error code from the library into a more strongly typed AresError.
fn ares_error(code: libc::c_int) -> AresError {
    match code {
        c_ares_sys::ARES_ENODATA => AresError::ENODATA,
        c_ares_sys::ARES_EFORMERR => AresError::EFORMERR,
        c_ares_sys::ARES_ESERVFAIL => AresError::ESERVFAIL,
        c_ares_sys::ARES_ENOTFOUND => AresError::ENOTFOUND,
        c_ares_sys::ARES_ENOTIMP => AresError::ENOTIMP,
        c_ares_sys::ARES_EREFUSED => AresError::EREFUSED,
        c_ares_sys::ARES_EBADQUERY => AresError::EBADQUERY,
        c_ares_sys::ARES_EBADNAME => AresError::EBADNAME,
        c_ares_sys::ARES_EBADFAMILY => AresError::EBADFAMILY,
        c_ares_sys::ARES_EBADRESP => AresError::EBADRESP,
        c_ares_sys::ARES_ECONNREFUSED => AresError::ECONNREFUSED,
        c_ares_sys::ARES_ETIMEOUT => AresError::ETIMEOUT,
        c_ares_sys::ARES_EOF => AresError::EOF,
        c_ares_sys::ARES_EFILE => AresError::EFILE,
        c_ares_sys::ARES_ENOMEM => AresError::ENOMEM,
        c_ares_sys::ARES_EDESTRUCTION => AresError::EDESTRUCTION,
        c_ares_sys::ARES_EBADSTR => AresError::EBADSTR,
        c_ares_sys::ARES_EBADFLAGS => AresError::EBADFLAGS,
        c_ares_sys::ARES_ENONAME => AresError::ENONAME,
        c_ares_sys::ARES_EBADHINTS => AresError::EBADHINTS,
        c_ares_sys::ARES_ENOTINITIALIZED => AresError::ENOTINITIALIZED,
        c_ares_sys::ARES_ELOADIPHLPAPI => AresError::ELOADIPHLPAPI,
        c_ares_sys::ARES_EADDRGETNETWORKPARAMS => AresError::EADDRGETNETWORKPARAMS,
        c_ares_sys::ARES_ECANCELLED => AresError::ECANCELLED,
        _ => AresError::UNKNOWN,
    }
}

#[repr(C)]
struct hostent {
    h_name: *mut libc::c_char,
    h_aliases: *mut *mut libc::c_char,
    h_addrtype: libc::c_int,
    h_length: libc::c_int,
    h_addr_list: *mut *mut libc::c_char,
}

// See arpa/nameser.h
enum QueryType {
    A = 1,
    AAAA = 28,
}

// See arpa/nameser.h
enum DnsClass {
   IN = 1,
}

extern "C" fn socket_callback<F>(
    data: *mut libc::c_void,
    socket_fd: c_ares_sys::ares_socket_t,
    readable: libc::c_int,
    writable: libc::c_int)
    where F: FnOnce(io::RawFd, bool, bool) + 'static {
    let handler: Box<F> = unsafe { mem::transmute(data) };
    handler(socket_fd as io::RawFd, readable != 0, writable != 0);
}

extern "C" fn query_a_callback<F>(
    arg: *mut libc::c_void,
    status: libc::c_int,
    _timeouts: libc::c_int,
    abuf: *mut libc::c_uchar,
    alen: libc::c_int)
    where F: FnOnce(Result<AResult, AresError>) + 'static {
    let result = if status != c_ares_sys::ARES_SUCCESS {
        Err(ares_error(status))
    } else {
        let mut hostent: *mut hostent = ptr::null_mut();
        let parse_status = unsafe {
            c_ares_sys::ares_parse_a_reply(
                abuf,
                alen,
                &mut hostent as *mut *mut _ as *mut *mut c_ares_sys::Struct_hostent,
                ptr::null_mut(),
                ptr::null_mut())
        };
        if parse_status != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(parse_status))
        } else {
            let mut answers = Vec::new();
            unsafe {
                let mut ptr = (*hostent).h_addr_list;
                while !(*ptr).is_null() {
                    let h_addr = *ptr;
                    let ip_addr = Ipv4Addr::new(
                        *h_addr as u8,
                        *h_addr.offset(1) as u8,
                        *h_addr.offset(2) as u8,
                        *h_addr.offset(3) as u8);
                    answers.push(ip_addr);
                    ptr = ptr.offset(1);
                }
                c_ares_sys::ares_free_hostent(hostent as *mut c_ares_sys::Struct_hostent);
            }
            Ok(AResult { ip_addrs: answers })
        }
    };

    let handler: Box<F> = unsafe { mem::transmute(arg) };
    handler(result);
}

extern "C" fn query_aaaa_callback<F>(
    arg: *mut libc::c_void,
    status: libc::c_int,
    _timeouts: libc::c_int,
    abuf: *mut libc::c_uchar,
    alen: libc::c_int)
    where F: FnOnce(Result<AAAAResult, AresError>) + 'static {
    let result = if status != c_ares_sys::ARES_SUCCESS {
        Err(ares_error(status))
    } else {
        let mut hostent: *mut hostent = ptr::null_mut();
        let parse_status = unsafe {
            c_ares_sys::ares_parse_aaaa_reply(
                abuf,
                alen,
                &mut hostent as *mut *mut _ as *mut *mut c_ares_sys::Struct_hostent,
                ptr::null_mut(),
                ptr::null_mut())
        };
        if parse_status != c_ares_sys::ARES_SUCCESS {
            Err(ares_error(parse_status))
        } else {
            let mut answers = Vec::new();
            unsafe {
                let mut ptr = (*hostent).h_addr_list;
                while !(*ptr).is_null() {
                    let h_addr = *ptr;
                    let ip_addr = Ipv6Addr::new(
                        ((*h_addr as u16) << 8) + *h_addr.offset(1) as u16,
                        ((*h_addr.offset(2) as u16) << 8) + *h_addr.offset(3) as u16,
                        ((*h_addr.offset(4) as u16) << 8) + *h_addr.offset(5) as u16,
                        ((*h_addr.offset(6) as u16) << 8) + *h_addr.offset(7) as u16,
                        ((*h_addr.offset(8) as u16) << 8) + *h_addr.offset(9) as u16,
                        ((*h_addr.offset(10) as u16) << 8) + *h_addr.offset(11) as u16,
                        ((*h_addr.offset(12) as u16) << 8) + *h_addr.offset(13) as u16,
                        ((*h_addr.offset(14) as u16) << 8) + *h_addr.offset(15) as u16);
                    answers.push(ip_addr);
                    ptr = ptr.offset(1);
                }
                c_ares_sys::ares_free_hostent(hostent as *mut c_ares_sys::Struct_hostent);
            }
            Ok(AAAAResult { ip_addrs: answers })
        }
    };

    let handler: Box<F> = unsafe { mem::transmute(arg) };
    handler(result);
}
<|endoftext|>


<commit_before>//! An implementation of [FLAC](https://xiph.org/flac), free lossless audio
//! codec, written in Rust.
//!
//! # Examples
//!
//! Basic decoding from a file.
//!
//! ```
//! match flac::Stream::from_file("path/to/file.flac") {
//!   Ok(mut stream) => {
//!     // Copy of `StreamInfo` to help convert to a different audio format.
//!     let info = stream.info();
//!
//!     for sample in stream.iter() {
//!       // Iterate over each decoded sample
//!     }
//!   }
//!   Err(error)     => println!("{}", error),
//! }
//! ```

#[macro_use]
extern crate nom;

#[macro_use]
mod utility;
pub mod metadata;
pub mod frame;
pub mod subframe;
pub mod stream;

pub use metadata::{metadata_parser, Metadata};
pub use frame::{frame_parser, Frame};
pub use subframe::{subframe_parser, Subframe};
pub use stream::{stream_parser, Stream};
<commit_msg>Remove public export of frame<commit_after>//! An implementation of [FLAC](https://xiph.org/flac), free lossless audio
//! codec, written in Rust.
//!
//! # Examples
//!
//! Basic decoding from a file.
//!
//! ```
//! match flac::Stream::from_file("path/to/file.flac") {
//!   Ok(mut stream) => {
//!     // Copy of `StreamInfo` to help convert to a different audio format.
//!     let info = stream.info();
//!
//!     for sample in stream.iter() {
//!       // Iterate over each decoded sample
//!     }
//!   }
//!   Err(error)     => println!("{}", error),
//! }
//! ```

#[macro_use]
extern crate nom;

#[macro_use]
mod utility;
pub mod metadata;
mod frame;
pub mod subframe;
pub mod stream;

pub use metadata::{metadata_parser, Metadata};
pub use frame::{frame_parser, Frame};
pub use subframe::{subframe_parser, Subframe};
pub use stream::{stream_parser, Stream};
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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
//
//     https://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 bit_field::BitField;
use std::fmt::{self, Debug, Display, Formatter};
use thiserror::Error;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FieldInfo {
    pub name: &'static str,
    pub start: usize,
    pub width: usize,
    pub value: u64,
    pub decoded: Option<Decoded>,
}

impl FieldInfo {
    fn get(register: u64, name: &'static str, start: usize, end: usize) -> Self {
        let value = register.get_bits(start..end);
        Self {
            name,
            start,
            width: end - start,
            value,
            decoded: None,
        }
    }

    fn get_bit(register: u64, name: &'static str, bit: usize) -> Self {
        Self::get(register, name, bit, bit + 1)
    }

    fn with_decoded(self, decoded: Decoded) -> Self {
        Self {
            decoded: Some(decoded),
            ..self
        }
    }

    fn as_bit(&self) -> bool {
        assert!(self.width == 1);
        self.value == 1
    }

    /// Assuming this field has a width of exactly 1, describe it with the given function.
    ///
    /// Panics if `self.width != 1`.
    fn describe_bit<F>(self, describer: F) -> Self
    where
        F: FnOnce(bool) -> &'static str,
    {
        let bit = self.as_bit();
        let description = Some(describer(bit).to_string());
        self.with_decoded(Decoded {
            description,
            fields: vec![],
        })
    }
}

impl Display for FieldInfo {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        if self.width == 1 {
            write!(
                f,
                "{}: {}",
                self.name,
                if self.value == 1 { "true" } else { "false" }
            )
        } else {
            write!(
                f,
                "{}: {:#02$x} {1:#03$b}",
                self.name,
                self.value,
                (self.width + 3) / 4 + 2,
                self.width + 2
            )
        }
    }
}

#[derive(Debug, Error)]
pub enum DecodeError {
    #[error("Invalid ESR, res0 is {res0}")]
    InvalidRes0 { res0: u64 },
    #[error("Invalid EC {ec}")]
    InvalidEc { ec: u64 },
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SyndromeAccessSize {
    Byte = 0b00,
    Halfword = 0b01,
    Word = 0b10,
    Doubleword = 0b11,
}

impl Display for SyndromeAccessSize {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let s = match self {
            Self::Byte => "byte",
            Self::Halfword => "halfword",
            Self::Word => "word",
            Self::Doubleword => "doubleword",
        };
        write!(f, "{}", s)
    }
}

fn describe_sf(sf: bool) -> &'static str {
    if sf {
        "64-bit wide register"
    } else {
        "32-bit wide register"
    }
}

fn describe_ar(ar: bool) -> &'static str {
    if ar {
        "Acquire/release semantics"
    } else {
        "No acquire/release semantics"
    }
}

fn describe_fnv(fnv: bool) -> &'static str {
    if fnv {
        "FAR is not valid, it holds an unknown value"
    } else {
        "FAR is valid"
    }
}

fn describe_wnr(wnr: bool) -> &'static str {
    if wnr {
        "Abort caused by writing to memory"
    } else {
        "Abort caused by reading from memory"
    }
}

fn decode_iss_data_abort(iss: u64) -> Decoded {
    let isv = FieldInfo::get_bit(iss, "ISV", 24);
    let sas = FieldInfo::get(iss, "SAS", 22, 24);
    let sas_value = match sas.value {
        0b00 => SyndromeAccessSize::Byte,
        0b01 => SyndromeAccessSize::Halfword,
        0b10 => SyndromeAccessSize::Word,
        0b11 => SyndromeAccessSize::Doubleword,
        _ => unreachable!(),
    };
    let sas = sas.with_decoded(Decoded {
        description: Some(sas_value.to_string()),
        fields: vec![],
    });
    let sse = FieldInfo::get_bit(iss, "SSE", 21);
    let srt = FieldInfo::get(iss, "SRT", 16, 21);
    let sf = FieldInfo::get_bit(iss, "SF", 15).describe_bit(describe_sf);
    let ar = FieldInfo::get_bit(iss, "AR", 14).describe_bit(describe_ar);
    let vncr = FieldInfo::get_bit(iss, "VNCR", 13);
    let fnv = FieldInfo::get_bit(iss, "FnV", 10).describe_bit(describe_fnv);
    let ea = FieldInfo::get_bit(iss, "EA", 9);
    let cm = FieldInfo::get_bit(iss, "CM", 8);
    let s1ptw = FieldInfo::get_bit(iss, "S1PTW", 7);
    let wnr = FieldInfo::get_bit(iss, "WnR", 6).describe_bit(describe_wnr);
    let dfsc = FieldInfo::get(iss, "DFSC", 0, 6);

    Decoded {
        description: None,
        fields: vec![
            isv, sas, sse, srt, sf, ar, vncr, fnv, ea, cm, s1ptw, wnr, dfsc,
        ],
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Decoded {
    pub description: Option<String>,
    pub fields: Vec<FieldInfo>,
}

pub fn decode(esr: u64) -> Result<Decoded, DecodeError> {
    let res0 = FieldInfo::get(esr, "RES0", 37, 64);
    let iss2 = FieldInfo::get(esr, "ISS2", 32, 37);
    let ec = FieldInfo::get(esr, "EC", 26, 32);
    let il = FieldInfo::get_bit(esr, "IL", 25);
    let iss = FieldInfo::get(esr, "ISS", 0, 25);
    if res0.value != 0 {
        return Err(DecodeError::InvalidRes0 { res0: res0.value });
    }
    let (class, iss_decoded) = match ec.value {
        0b000000 => ("Unknown reason", None),
        0b000001 => ("Wrapped WF* instruction execution", None),
        0b000011 => ("Trapped MCR or MRC access with coproc=0b1111", None),
        0b000100 => ("Trapped MCRR or MRRC access with coproc=0b1111", None),
        0b000101 => ("Trapped MCR or MRC access with coproc=0b1110", None),
        0b000110 => ("Trapped LDC or STC access", None),
        0b000111 => ("Trapped access to SVE, Advanced SIMD or floating point", None),
        0b001010 => ("Trapped execution of an LD64B, ST64B, ST64BV, or ST64BV0 instruction", None),
        0b001100 => ("Trapped MRRC access with (coproc==0b1110)", None),
        0b001101 => ("Branch Target Exception", None),
        0b001110 => ("Illegal Execution state", None),
        0b010001 => ("SVC instruction execution in AArch32 state", None),
        0b010101 => ("SVC instruction execution in AArch64 state", None),
        0b011000 => ("Trapped MSR, MRS or System instruction execution in AArch64 state", None),
        0b011001 => ("Access to SVE functionality trapped as a result of CPACR_EL1.ZEN, CPTR_EL2.ZEN, CPTR_EL2.TZ, or CPTR_EL3.EZ", None),
        0b011100 => ("Exception from a Pointer Authentication instruction authentication failure", None),
        0b100000 => ("Instruction Abort from a lower Exception level", None),
        0b100001 => ("Instruction Abort taken without a change in Exception level", None),
        0b100010 => ("PC alignment fault exception", None),
        0b100100 => ("Data Abort from a lower Exception level", Some(decode_iss_data_abort(iss.value))),
        0b100101 => ("Data Abort taken without a change in Exception level", Some(decode_iss_data_abort(iss.value))),
        0b100110 => ("SP alignment fault exception", None),
        0b101000 => ("Trapped floating-point exception taken from AArch32 state", None),
        0b101100 => ("Trapped floating-point exception taken from AArch64 state", None),
        0b101111 => ("SError interrupt", None),
        0b110000 => ("Breakpoint exception from a lower Exception level", None),
        0b110001 => ("Breakpoint exception taken without a change in Exception level", None),
        0b110010 => ("Software Step exception from a lower Exception level", None),
        0b110011 => ("Software Step exception taken without a change in Exception level", None),
        0b110100 => ("Watchpoint exception from a lower Exception level", None),
        0b110101 => ("Watchpoint exception taken without a change in Exception level", None),
        0b111000 => ("BKPT instruction execution in AArch32 state", None),
        0b111100 => ("BRK instruction execution in AArch64 state", None),
        _ => return Err(DecodeError::InvalidEc { ec: ec.value }),
    };
    let iss = FieldInfo {
        decoded: iss_decoded,
        ..iss
    };
    let ec = ec.with_decoded(Decoded {
        description: Some(class.to_string()),
        fields: vec![],
    });
    Ok(Decoded {
        description: Some(class.to_string()),
        fields: vec![res0, iss2, ec, il, iss],
    })
}
<commit_msg>Describe DFSC.<commit_after>// Copyright 2021 Google LLC
//
// 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
//
//     https://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 bit_field::BitField;
use std::fmt::{self, Debug, Display, Formatter};
use thiserror::Error;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct FieldInfo {
    pub name: &'static str,
    pub start: usize,
    pub width: usize,
    pub value: u64,
    pub decoded: Option<Decoded>,
}

impl FieldInfo {
    fn get(register: u64, name: &'static str, start: usize, end: usize) -> Self {
        let value = register.get_bits(start..end);
        Self {
            name,
            start,
            width: end - start,
            value,
            decoded: None,
        }
    }

    fn get_bit(register: u64, name: &'static str, bit: usize) -> Self {
        Self::get(register, name, bit, bit + 1)
    }

    fn with_decoded(self, decoded: Decoded) -> Self {
        Self {
            decoded: Some(decoded),
            ..self
        }
    }

    fn with_description(self, description: String) -> Self {
        self.with_decoded(Decoded {
            description: Some(description),
            fields: vec![],
        })
    }

    fn as_bit(&self) -> bool {
        assert!(self.width == 1);
        self.value == 1
    }

    /// Assuming this field has a width of exactly 1, describe it with the given function.
    ///
    /// Panics if `self.width != 1`.
    fn describe_bit<F>(self, describer: F) -> Self
    where
        F: FnOnce(bool) -> &'static str,
    {
        let bit = self.as_bit();
        let description = Some(describer(bit).to_string());
        self.with_decoded(Decoded {
            description,
            fields: vec![],
        })
    }
}

impl Display for FieldInfo {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        if self.width == 1 {
            write!(
                f,
                "{}: {}",
                self.name,
                if self.value == 1 { "true" } else { "false" }
            )
        } else {
            write!(
                f,
                "{}: {:#02$x} {1:#03$b}",
                self.name,
                self.value,
                (self.width + 3) / 4 + 2,
                self.width + 2
            )
        }
    }
}

#[derive(Debug, Error)]
pub enum DecodeError {
    #[error("Invalid ESR, res0 is {res0}")]
    InvalidRes0 { res0: u64 },
    #[error("Invalid EC {ec}")]
    InvalidEc { ec: u64 },
    #[error("Invalid DFSC {dfsc}")]
    InvalidDfsc { dfsc: u64 },
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SyndromeAccessSize {
    Byte = 0b00,
    Halfword = 0b01,
    Word = 0b10,
    Doubleword = 0b11,
}

impl Display for SyndromeAccessSize {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let s = match self {
            Self::Byte => "byte",
            Self::Halfword => "halfword",
            Self::Word => "word",
            Self::Doubleword => "doubleword",
        };
        write!(f, "{}", s)
    }
}

fn describe_sf(sf: bool) -> &'static str {
    if sf {
        "64-bit wide register"
    } else {
        "32-bit wide register"
    }
}

fn describe_ar(ar: bool) -> &'static str {
    if ar {
        "Acquire/release semantics"
    } else {
        "No acquire/release semantics"
    }
}

fn describe_fnv(fnv: bool) -> &'static str {
    if fnv {
        "FAR is not valid, it holds an unknown value"
    } else {
        "FAR is valid"
    }
}

fn describe_wnr(wnr: bool) -> &'static str {
    if wnr {
        "Abort caused by writing to memory"
    } else {
        "Abort caused by reading from memory"
    }
}

fn describe_dfsc(dfsc: u64) -> Result<&'static str, DecodeError> {
    let description = match dfsc {
        0b000000 => "Address size fault, level 0 of translation or translation table base register.",
        0b000001 => "Address size fault, level 1.",
        0b000010 => "Address size fault, level 2.",
        0b000011 => "Address size fault, level 3.",
        0b000100 => "Translation fault, level 0.",
        0b000101 => "Translation fault, level 1.",
        0b000110 => "Translation fault, level 2.",
        0b000111 => "Translation fault, level 3.",
        0b001001 => "Access flag fault, level 1.",
        0b001010 => "Access flag fault, level 2.",
        0b001011 => "Access flag fault, level 3.",
        0b001000 => "Access flag fault, level 0.",
        0b001100 => "Permission fault, level 0.",
        0b001101 => "Permission fault, level 1.",
        0b001110 => "Permission fault, level 2.",
        0b001111 => "Permission fault, level 3.",
        0b010000 => "Synchronous External abort, not on translation table walk or hardware update of translation table.",
        0b010001 => "Synchronous Tag Check Fault.",
        0b010011 => "Synchronous External abort on translation table walk or hardware update of translation table, level -1.",
        0b010100 => "Synchronous External abort on translation table walk or hardware update of translation table, level 0.",
        0b010101 => "Synchronous External abort on translation table walk or hardware update of translation table, level 1.",
        0b010110 => "Synchronous External abort on translation table walk or hardware update of translation table, level 2.",
        0b010111 => "Synchronous External abort on translation table walk or hardware update of translation table, level 3.",
        0b011000 => "Synchronous parity or ECC error on memory access, not on translation table walk.",
        0b011011 => "Synchronous parity or ECC error on memory access on translation table walk or hardware update of translation table, level -1.",
        0b011100 => "Synchronous parity or ECC error on memory access on translation table walk or hardware update of translation table, level 0.",
        0b011101 => "Synchronous parity or ECC error on memory access on translation table walk or hardware update of translation table, level 1.",
        0b011110 => "Synchronous parity or ECC error on memory access on translation table walk or hardware update of translation table, level 2.",
        0b011111 => "Synchronous parity or ECC error on memory access on translation table walk or hardware update of translation table, level 3.",
        0b100001 => "Alignment fault.",
        0b101001 => "Address size fault, level -1.",
        0b101011 => "Translation fault, level -1.",
        0b110000 => "TLB conflict abort.",
        0b110001 => "Unsupported atomic hardware update fault.",
        0b110100 => "IMPLEMENTATION DEFINED fault (Lockdown).",
        0b110101 => "IMPLEMENTATION DEFINED fault (Unsupported Exclusive or Atomic access).",
        _ => return Err(DecodeError::InvalidDfsc { dfsc }),
    };
    Ok(description)
}

fn decode_iss_data_abort(iss: u64) -> Result<Decoded, DecodeError> {
    let isv = FieldInfo::get_bit(iss, "ISV", 24);
    let sas = FieldInfo::get(iss, "SAS", 22, 24);
    let sas_value = match sas.value {
        0b00 => SyndromeAccessSize::Byte,
        0b01 => SyndromeAccessSize::Halfword,
        0b10 => SyndromeAccessSize::Word,
        0b11 => SyndromeAccessSize::Doubleword,
        _ => unreachable!(),
    };
    let sas = sas.with_decoded(Decoded {
        description: Some(sas_value.to_string()),
        fields: vec![],
    });
    let sse = FieldInfo::get_bit(iss, "SSE", 21);
    let srt = FieldInfo::get(iss, "SRT", 16, 21);
    let sf = FieldInfo::get_bit(iss, "SF", 15).describe_bit(describe_sf);
    let ar = FieldInfo::get_bit(iss, "AR", 14).describe_bit(describe_ar);
    let vncr = FieldInfo::get_bit(iss, "VNCR", 13);
    let fnv = FieldInfo::get_bit(iss, "FnV", 10).describe_bit(describe_fnv);
    let ea = FieldInfo::get_bit(iss, "EA", 9);
    let cm = FieldInfo::get_bit(iss, "CM", 8);
    let s1ptw = FieldInfo::get_bit(iss, "S1PTW", 7);
    let wnr = FieldInfo::get_bit(iss, "WnR", 6).describe_bit(describe_wnr);
    let dfsc = FieldInfo::get(iss, "DFSC", 0, 6);
    let dfsc_description = describe_dfsc(dfsc.value)?;
    let dfsc = dfsc.with_description(dfsc_description.to_string());

    Ok(Decoded {
        description: None,
        fields: vec![
            isv, sas, sse, srt, sf, ar, vncr, fnv, ea, cm, s1ptw, wnr, dfsc,
        ],
    })
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Decoded {
    pub description: Option<String>,
    pub fields: Vec<FieldInfo>,
}

pub fn decode(esr: u64) -> Result<Decoded, DecodeError> {
    let res0 = FieldInfo::get(esr, "RES0", 37, 64);
    let iss2 = FieldInfo::get(esr, "ISS2", 32, 37);
    let ec = FieldInfo::get(esr, "EC", 26, 32);
    let il = FieldInfo::get_bit(esr, "IL", 25);
    let iss = FieldInfo::get(esr, "ISS", 0, 25);
    if res0.value != 0 {
        return Err(DecodeError::InvalidRes0 { res0: res0.value });
    }
    let (class, iss_decoded) = match ec.value {
        0b000000 => ("Unknown reason", None),
        0b000001 => ("Wrapped WF* instruction execution", None),
        0b000011 => ("Trapped MCR or MRC access with coproc=0b1111", None),
        0b000100 => ("Trapped MCRR or MRRC access with coproc=0b1111", None),
        0b000101 => ("Trapped MCR or MRC access with coproc=0b1110", None),
        0b000110 => ("Trapped LDC or STC access", None),
        0b000111 => ("Trapped access to SVE, Advanced SIMD or floating point", None),
        0b001010 => ("Trapped execution of an LD64B, ST64B, ST64BV, or ST64BV0 instruction", None),
        0b001100 => ("Trapped MRRC access with (coproc==0b1110)", None),
        0b001101 => ("Branch Target Exception", None),
        0b001110 => ("Illegal Execution state", None),
        0b010001 => ("SVC instruction execution in AArch32 state", None),
        0b010101 => ("SVC instruction execution in AArch64 state", None),
        0b011000 => ("Trapped MSR, MRS or System instruction execution in AArch64 state", None),
        0b011001 => ("Access to SVE functionality trapped as a result of CPACR_EL1.ZEN, CPTR_EL2.ZEN, CPTR_EL2.TZ, or CPTR_EL3.EZ", None),
        0b011100 => ("Exception from a Pointer Authentication instruction authentication failure", None),
        0b100000 => ("Instruction Abort from a lower Exception level", None),
        0b100001 => ("Instruction Abort taken without a change in Exception level", None),
        0b100010 => ("PC alignment fault exception", None),
        0b100100 => ("Data Abort from a lower Exception level", Some(decode_iss_data_abort(iss.value)?)),
        0b100101 => ("Data Abort taken without a change in Exception level", Some(decode_iss_data_abort(iss.value)?)),
        0b100110 => ("SP alignment fault exception", None),
        0b101000 => ("Trapped floating-point exception taken from AArch32 state", None),
        0b101100 => ("Trapped floating-point exception taken from AArch64 state", None),
        0b101111 => ("SError interrupt", None),
        0b110000 => ("Breakpoint exception from a lower Exception level", None),
        0b110001 => ("Breakpoint exception taken without a change in Exception level", None),
        0b110010 => ("Software Step exception from a lower Exception level", None),
        0b110011 => ("Software Step exception taken without a change in Exception level", None),
        0b110100 => ("Watchpoint exception from a lower Exception level", None),
        0b110101 => ("Watchpoint exception taken without a change in Exception level", None),
        0b111000 => ("BKPT instruction execution in AArch32 state", None),
        0b111100 => ("BRK instruction execution in AArch64 state", None),
        _ => return Err(DecodeError::InvalidEc { ec: ec.value }),
    };
    let iss = FieldInfo {
        decoded: iss_decoded,
        ..iss
    };
    let ec = ec.with_decoded(Decoded {
        description: Some(class.to_string()),
        fields: vec![],
    });
    Ok(Decoded {
        description: Some(class.to_string()),
        fields: vec![res0, iss2, ec, il, iss],
    })
}
<|endoftext|>


<commit_before>// Copyright 2015 Philipp Oppermann
//
// 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.

#![feature(no_std, lang_items)]
#![feature(const_fn, unique, core_str_ext)]
#![no_std]

extern crate rlibc;
extern crate spin;

mod vga_buffer;

#[no_mangle]
pub extern fn rust_main() {
    // ATTENTION: we have a very small stack and no guard page

    let hello = b"Hello World!";
    let color_byte = 0x1f; // white foreground, blue background

    let mut hello_colored = [color_byte; 24];
    for (i, char_byte) in hello.into_iter().enumerate() {
        hello_colored[i*2] = *char_byte;
    }

    // write `Hello World!` to the center of the VGA text buffer
    let buffer_ptr = (0xb8000 + 1988) as *mut _;
    unsafe { *buffer_ptr = hello_colored };

    loop{}
}

#[cfg(not(test))]
#[lang = "eh_personality"]
extern fn eh_personality() {}

#[cfg(not(test))]
#[lang = "panic_fmt"]
extern fn panic_fmt() -> ! {loop{}}
<commit_msg>Use `println!` to print the `Hello World!<commit_after>// Copyright 2015 Philipp Oppermann
//
// 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.

#![feature(no_std, lang_items)]
#![feature(const_fn, unique, core_str_ext)]
#![no_std]

extern crate rlibc;
extern crate spin;

#[macro_use]
mod vga_buffer;

#[no_mangle]
pub extern fn rust_main() {
    // ATTENTION: we have a very small stack and no guard page
    println!("Hello World{}", "!");

    loop{}
}

#[cfg(not(test))]
#[lang = "eh_personality"]
extern fn eh_personality() {}

#[cfg(not(test))]
#[lang = "panic_fmt"]
extern fn panic_fmt() -> ! {loop{}}
<|endoftext|>


<commit_before>extern crate redis;

use std::time::Duration;

use redis::{Connection, RedisError, cmd};

fn duration_to_millis(d: &Duration) -> u64 {
    (d.subsec_nanos() / 1_000_000) as u64 + d.as_secs()
}

macro_rules! option_arg {
    ($disque: expr, $name: expr, $param: expr) => (
        match $param {
            Some(u) => $disque.arg($name).arg(u),
            None => &mut $disque,
        };
    )
}
pub struct Disque {
    connection: Connection,
}

impl Disque {
    pub fn open(url: &str) -> Result<Disque, RedisError> {
        let client = try!(redis::Client::open(url));
        let connection = try!(client.get_connection());
        Ok(Disque { connection: connection })
    }

    pub fn hello(&self) -> Result<Vec<Vec<u8>>, RedisError> {
        cmd("HELLO").query(&self.connection)
    }

    pub fn addjob(&self, queue_name: &[u8], job: &[u8], timeout: Duration,
            replicate: Option<usize>, delay: Option<Duration>,
            retry: Option<Duration>, ttl: Option<Duration>,
            maxlen: Option<usize>, async: bool,
            ) -> Result<String, RedisError> {
        let mut c = cmd("ADDJOB");
        c
            .arg(queue_name)
            .arg(job)
            .arg(duration_to_millis(&timeout));
        option_arg!(c, "REPLICATE", replicate);
        option_arg!(c, "DELAY", delay.map(|x| x.as_secs()));
        option_arg!(c, "RETRY", retry.map(|x| x.as_secs()));
        option_arg!(c, "TTL", ttl.map(|x| x.as_secs()));
        option_arg!(c, "MAXLEN", maxlen);
        if async { c.arg("ASYNC"); }
        c.query(&self.connection)
    }

    pub fn getjob_count(&self, nohang: bool, timeout: Option<Duration>,
            count: usize, withcounters: bool, queues: &[&[u8]]
            ) -> Result<Vec<Vec<Vec<u8>>>, RedisError> {
        let mut c = cmd("GETJOB");
        if nohang { c.arg("NOHANG"); }
        option_arg!(c, "TIMEOUT", timeout.map(|t| duration_to_millis(&t)));
        c.arg("COUNT").arg(count);
        if withcounters { c.arg("WITHCOUNTERS"); }
        c.arg("FROM");
        for queue in queues { c.arg(*queue); }
        c.query(&self.connection)
    }

    pub fn getjob(&self, nohang: bool, timeout: Option<Duration>,
            withcounters: bool, queues: &[&[u8]]
            ) -> Result<Option<Vec<Vec<u8>>>, RedisError> {
        let mut jobs = try!(self.getjob_count(nohang, timeout, 1, withcounters,
                    queues));
        Ok(jobs.pop())
    }

    pub fn ackjob(&self, jobids: &[&[u8]]) -> Result<bool, RedisError> {
        let mut c = cmd("ACKJOB");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }

    pub fn fastack(&self, jobids: &[&[u8]]) -> Result<usize, RedisError> {
        let mut c = cmd("FASTACK");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }

    pub fn working(&self, jobid: &[u8]) -> Result<Duration, RedisError> {
        let retry = try!(cmd("WORKING").arg(jobid).query(&self.connection));
        Ok(Duration::from_secs(retry))
    }

    pub fn nack(&self, jobids: &[&[u8]]) -> Result<usize, RedisError> {
        let mut c = cmd("NACK");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }
}

#[cfg(test)]
fn conn() -> Disque {
    Disque::open("redis://127.0.0.1:7711/").unwrap()
}

#[test]
fn can_connect() {
    conn();
}

#[test]
fn hello() {
    let disque = conn();
    disque.hello().unwrap();
}

#[test]
fn addjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue", b"job", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert_eq!(jobid.len(), 48);
    assert_eq!(&jobid[..2], "DI");
    assert_eq!(&jobid[46..], "SQ");
}

#[test]
fn getjob_count() {
    let disque = conn();
    let j1 = disque.addjob(b"queue1", b"job1", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j2 = disque.addjob(b"queue2", b"job2", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    disque.addjob(b"queue3", b"job3", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let jobs = disque.getjob_count(false, None, 3, true, &[b"queue1", b"queue2"]).unwrap();
    assert_eq!(jobs.len(), 2);
    assert_eq!(jobs[0][0], b"queue1");
    assert_eq!(jobs[0][1], j1.into_bytes());
    assert_eq!(jobs[0][2], b"job1");
    assert_eq!(jobs[1][0], b"queue2");
    assert_eq!(jobs[1][1], j2.into_bytes());
    assert_eq!(jobs[1][2], b"job2");
}

#[test]
fn getjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue4", b"job4", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let job = disque.getjob(false, None, true, &[b"queue4", b"queue5"]).unwrap().unwrap();
    assert_eq!(job[0], b"queue4");
    assert_eq!(job[1], jobid.into_bytes());
    assert_eq!(job[2], b"job4");
}

#[test]
fn ackjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue6", b"job6", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.ackjob(&[jobid.as_bytes()]).unwrap());
    assert!(!disque.ackjob(&[jobid.as_bytes()]).unwrap());
    assert!(!disque.ackjob(&[jobid.as_bytes()]).unwrap());
}

#[test]
fn fastack() {
    let disque = conn();
    let jobid = disque.addjob(b"queue7", b"job7", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 1);
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 0);
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 0);
}

#[test]
fn working() {
    let disque = conn();
    let jobid = disque.addjob(b"queue8", b"job8", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.working(jobid.as_bytes()).unwrap().as_secs() > 0);
}

#[test]
fn nack() {
    let disque = conn();
    let j1 = disque.addjob(b"queue9", b"job9.1", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j2 = disque.addjob(b"queue9", b"job9.2", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j3 = disque.addjob(b"queue9", b"job9.3", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert_eq!(disque.getjob_count(false, None, 100, true, &[b"queue9"]).unwrap().len(), 3);
    assert_eq!(disque.nack(&[j1.as_bytes(), j2.as_bytes(), j3.as_bytes()]).unwrap(), 3);
    assert_eq!(disque.getjob_count(false, None, 100, true, &[b"queue9"]).unwrap().len(), 3);
}
<commit_msg>Add INFO<commit_after>extern crate redis;

use std::time::Duration;

use redis::{Connection, RedisError, cmd};

fn duration_to_millis(d: &Duration) -> u64 {
    (d.subsec_nanos() / 1_000_000) as u64 + d.as_secs()
}

macro_rules! option_arg {
    ($disque: expr, $name: expr, $param: expr) => (
        match $param {
            Some(u) => $disque.arg($name).arg(u),
            None => &mut $disque,
        };
    )
}
pub struct Disque {
    connection: Connection,
}

impl Disque {
    pub fn open(url: &str) -> Result<Disque, RedisError> {
        let client = try!(redis::Client::open(url));
        let connection = try!(client.get_connection());
        Ok(Disque { connection: connection })
    }

    pub fn hello(&self) -> Result<Vec<Vec<u8>>, RedisError> {
        cmd("HELLO").query(&self.connection)
    }

    pub fn addjob(&self, queue_name: &[u8], job: &[u8], timeout: Duration,
            replicate: Option<usize>, delay: Option<Duration>,
            retry: Option<Duration>, ttl: Option<Duration>,
            maxlen: Option<usize>, async: bool,
            ) -> Result<String, RedisError> {
        let mut c = cmd("ADDJOB");
        c
            .arg(queue_name)
            .arg(job)
            .arg(duration_to_millis(&timeout));
        option_arg!(c, "REPLICATE", replicate);
        option_arg!(c, "DELAY", delay.map(|x| x.as_secs()));
        option_arg!(c, "RETRY", retry.map(|x| x.as_secs()));
        option_arg!(c, "TTL", ttl.map(|x| x.as_secs()));
        option_arg!(c, "MAXLEN", maxlen);
        if async { c.arg("ASYNC"); }
        c.query(&self.connection)
    }

    pub fn getjob_count(&self, nohang: bool, timeout: Option<Duration>,
            count: usize, withcounters: bool, queues: &[&[u8]]
            ) -> Result<Vec<Vec<Vec<u8>>>, RedisError> {
        let mut c = cmd("GETJOB");
        if nohang { c.arg("NOHANG"); }
        option_arg!(c, "TIMEOUT", timeout.map(|t| duration_to_millis(&t)));
        c.arg("COUNT").arg(count);
        if withcounters { c.arg("WITHCOUNTERS"); }
        c.arg("FROM");
        for queue in queues { c.arg(*queue); }
        c.query(&self.connection)
    }

    pub fn getjob(&self, nohang: bool, timeout: Option<Duration>,
            withcounters: bool, queues: &[&[u8]]
            ) -> Result<Option<Vec<Vec<u8>>>, RedisError> {
        let mut jobs = try!(self.getjob_count(nohang, timeout, 1, withcounters,
                    queues));
        Ok(jobs.pop())
    }

    pub fn ackjob(&self, jobids: &[&[u8]]) -> Result<bool, RedisError> {
        let mut c = cmd("ACKJOB");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }

    pub fn fastack(&self, jobids: &[&[u8]]) -> Result<usize, RedisError> {
        let mut c = cmd("FASTACK");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }

    pub fn working(&self, jobid: &[u8]) -> Result<Duration, RedisError> {
        let retry = try!(cmd("WORKING").arg(jobid).query(&self.connection));
        Ok(Duration::from_secs(retry))
    }

    pub fn nack(&self, jobids: &[&[u8]]) -> Result<usize, RedisError> {
        let mut c = cmd("NACK");
        for jobid in jobids { c.arg(*jobid); }
        c.query(&self.connection)
    }

    pub fn info(&self) -> Result<Vec<u8>, RedisError> {
        cmd("INFO").query(&self.connection)
    }
}

#[cfg(test)]
fn conn() -> Disque {
    Disque::open("redis://127.0.0.1:7711/").unwrap()
}

#[test]
fn can_connect() {
    conn();
}

#[test]
fn hello() {
    let disque = conn();
    disque.hello().unwrap();
}

#[test]
fn addjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue", b"job", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert_eq!(jobid.len(), 48);
    assert_eq!(&jobid[..2], "DI");
    assert_eq!(&jobid[46..], "SQ");
}

#[test]
fn getjob_count() {
    let disque = conn();
    let j1 = disque.addjob(b"queue1", b"job1", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j2 = disque.addjob(b"queue2", b"job2", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    disque.addjob(b"queue3", b"job3", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let jobs = disque.getjob_count(false, None, 3, true, &[b"queue1", b"queue2"]).unwrap();
    assert_eq!(jobs.len(), 2);
    assert_eq!(jobs[0][0], b"queue1");
    assert_eq!(jobs[0][1], j1.into_bytes());
    assert_eq!(jobs[0][2], b"job1");
    assert_eq!(jobs[1][0], b"queue2");
    assert_eq!(jobs[1][1], j2.into_bytes());
    assert_eq!(jobs[1][2], b"job2");
}

#[test]
fn getjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue4", b"job4", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let job = disque.getjob(false, None, true, &[b"queue4", b"queue5"]).unwrap().unwrap();
    assert_eq!(job[0], b"queue4");
    assert_eq!(job[1], jobid.into_bytes());
    assert_eq!(job[2], b"job4");
}

#[test]
fn ackjob() {
    let disque = conn();
    let jobid = disque.addjob(b"queue6", b"job6", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.ackjob(&[jobid.as_bytes()]).unwrap());
    assert!(!disque.ackjob(&[jobid.as_bytes()]).unwrap());
    assert!(!disque.ackjob(&[jobid.as_bytes()]).unwrap());
}

#[test]
fn fastack() {
    let disque = conn();
    let jobid = disque.addjob(b"queue7", b"job7", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 1);
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 0);
    assert!(disque.fastack(&[jobid.as_bytes()]).unwrap() == 0);
}

#[test]
fn working() {
    let disque = conn();
    let jobid = disque.addjob(b"queue8", b"job8", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert!(disque.working(jobid.as_bytes()).unwrap().as_secs() > 0);
}

#[test]
fn nack() {
    let disque = conn();
    let j1 = disque.addjob(b"queue9", b"job9.1", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j2 = disque.addjob(b"queue9", b"job9.2", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    let j3 = disque.addjob(b"queue9", b"job9.3", Duration::from_secs(10), None, None, None, None, None, false).unwrap();
    assert_eq!(disque.getjob_count(false, None, 100, true, &[b"queue9"]).unwrap().len(), 3);
    assert_eq!(disque.nack(&[j1.as_bytes(), j2.as_bytes(), j3.as_bytes()]).unwrap(), 3);
    assert_eq!(disque.getjob_count(false, None, 100, true, &[b"queue9"]).unwrap().len(), 3);
}

#[test]
fn info() {
    let disque = conn();
    disque.info().unwrap();
}
<|endoftext|>


<commit_before>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Debug)]
pub struct NodeSimilarityMatrix {
    // number of nodes in graph A
    na: usize,
    // number of nodes in graph B
    nb: usize,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iters: usize,
}

impl NodeSimilarityMatrix {
    pub fn init<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let (na, nb) = (in_a.len(), in_b.len());
        assert!((na, nb) == (out_a.len(), out_b.len()));

        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(na, nb, |i, j| {
            node_color_scale((i, j)) *
            if in_a[i].len() + out_a[i].len() > 0 && in_b[j].len() + out_b[j].len() > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(na, nb);

        NodeSimilarityMatrix {
            na: na,
            nb: nb,
            current: x,
            previous: new_x,
            num_iters: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    fn next<F>(&mut self,
               in_a: &[Vec<Idx>],
               in_b: &[Vec<Idx>],
               out_a: &[Vec<Idx>],
               out_b: &[Vec<Idx>],
               node_color_scale: &F)
        where F: Fn((usize, usize)) -> f32
    {
        next_x(&self.current,
               &mut self.previous,
               in_a,
               in_b,
               out_a,
               out_b,
               node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iters += 1;
    }

    #[inline]
    /// Calculates the similarity matrix for two graphs A and B.
    ///
    /// `in_a`:  Incoming edge list for each node of graph A
    /// `in_b`:  Incoming edge list for each node of graph B
    /// `out_a`: Outgoing edge list for each node of graph A
    /// `out_b`: Outgoing edge list for each node of graph B
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn calc<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   eps: f32,
                   stop_after_iter: usize,
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let mut x = NodeSimilarityMatrix::init(in_a, in_b, out_a, out_b, node_color_scale);

        for _ in 0..stop_after_iter {
            if x.in_eps(eps) {
                break;
            }
            x.next(in_a, in_b, out_a, out_b, node_color_scale);
        }

        return x;
    }

    fn matrix(self) -> DMat<f32> {
        self.current
    }

    pub fn num_iterations(&self) -> usize {
        self.num_iters
    }

    fn optimal_node_assignment(&self, n: usize) -> Vec<(usize, usize)> {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        assignment
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        self.optimal_node_assignment(n).iter().fold(0.0, |acc, &ab| acc + self.current[ab])
    }

    /// Calculate a measure how good the edge weights match up.
    ///
    /// We start by calculating the optimal node assignment between nodes of graph A and
    /// graph B, then compare all outgoing edges of similar nodes by again using an
    /// assignment.
    pub fn score_outgoing_edge_weights(&self) -> f32 {
        // XXX
        0.0
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);
    println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 1, &|_| 1.0);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<commit_msg>Silence warnings. Add assertions<commit_after>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, _n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Debug)]
pub struct NodeSimilarityMatrix {
    // number of nodes in graph A
    na: usize,
    // number of nodes in graph B
    nb: usize,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iters: usize,
}

impl NodeSimilarityMatrix {
    pub fn init<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let (na, nb) = (in_a.len(), in_b.len());
        assert!((na, nb) == (out_a.len(), out_b.len()));

        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(na, nb, |i, j| {
            node_color_scale((i, j)) *
            if in_a[i].len() + out_a[i].len() > 0 && in_b[j].len() + out_b[j].len() > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(na, nb);

        NodeSimilarityMatrix {
            na: na,
            nb: nb,
            current: x,
            previous: new_x,
            num_iters: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    fn next<F>(&mut self,
               in_a: &[Vec<Idx>],
               in_b: &[Vec<Idx>],
               out_a: &[Vec<Idx>],
               out_b: &[Vec<Idx>],
               node_color_scale: &F)
        where F: Fn((usize, usize)) -> f32
    {
        assert!((self.na, self.nb) == (in_a.len(), in_b.len()));
        assert!((self.na, self.nb) == (out_a.len(), out_b.len()));

        next_x(&self.current,
               &mut self.previous,
               in_a,
               in_b,
               out_a,
               out_b,
               node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iters += 1;
    }

    #[inline]
    /// Calculates the similarity matrix for two graphs A and B.
    ///
    /// `in_a`:  Incoming edge list for each node of graph A
    /// `in_b`:  Incoming edge list for each node of graph B
    /// `out_a`: Outgoing edge list for each node of graph A
    /// `out_b`: Outgoing edge list for each node of graph B
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn calc<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   eps: f32,
                   stop_after_iter: usize,
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let mut x = NodeSimilarityMatrix::init(in_a, in_b, out_a, out_b, node_color_scale);

        for _ in 0..stop_after_iter {
            if x.in_eps(eps) {
                break;
            }
            x.next(in_a, in_b, out_a, out_b, node_color_scale);
        }

        return x;
    }

    pub fn matrix(&self) -> &DMat<f32> {
        &self.current
    }

    pub fn num_iterations(&self) -> usize {
        self.num_iters
    }

    fn optimal_node_assignment(&self, n: usize) -> Vec<(usize, usize)> {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        assignment
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        self.optimal_node_assignment(n).iter().fold(0.0, |acc, &ab| acc + self.current[ab])
    }

    /// Calculate a measure how good the edge weights match up.
    ///
    /// We start by calculating the optimal node assignment between nodes of graph A and
    /// graph B, then compare all outgoing edges of similar nodes by again using an
    /// assignment.
    pub fn score_outgoing_edge_weights(&self) -> f32 {
        // XXX
        0.0
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);
    println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 1, &|_| 1.0);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<|endoftext|>