<commit_before>#[macro_use]
extern crate human_panic;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate xxd;

use cli::create_arg_parser;
use failure::*;
use xxd::dump::{dump_iterator, Config, Format};
use xxd::generate::{Render, Template};

use clap::ArgMatches;

use std::fmt::Display;
use std::io::{Read, Write};
use std::process::exit;

mod cli;

fn main() {
    setup_panic!();
    let matches = create_arg_parser().get_matches();
    match run(&matches) {
        Ok(_) => exit(0),
        Err(e) => {
            report_error(&e);
            exit(1)
        }
    }
}

fn run(args: &ArgMatches) -> Result<(), failure::Error> {
    match args.subcommand_name() {
        Some("dump") => dump(args.subcommand_matches("dump")),
        Some("generate") => generate(args.subcommand_matches("generate")),
        _ => bail!(format_err!("{}", args.usage())),
    }
}

pub fn create_reader(path: String) -> Result<Box<std::io::Read>, failure::Error> {
    match path.as_ref() {
        "stdin" => Ok(Box::new(std::io::stdin())),
        _ => {
            let file_reader = std::fs::File::open(path)?;
            Ok(Box::new(file_reader))
        }
    }
}

pub fn create_writer(path: String) -> Result<Box<std::io::Write>, failure::Error> {
    match path.as_ref() {
        "stdout" => Ok(Box::new(std::io::stdout())),
        _ => {
            let mut file_writer = std::fs::File::create(path)?;
            Ok(Box::new(file_writer))
        }
    }
}

fn dump<'a>(args: Option<&ArgMatches<'a>>) -> Result<(), failure::Error> {
    let args = args.ok_or_else(|| format_err!("No arguments available"))?;
    let output_file = args.value_of("outfile").unwrap_or("stdout");
    let input_file = args.value_of("file").unwrap_or("stdin");
    let seek = usize::from_str_radix(args.value_of("seek").unwrap_or("0"), 10)?;
    let length = args.value_of("length");
    let settings = create_dump_settings(args)?;
    let reader = create_reader(input_file.to_string())?;
    let mut writer = create_writer(output_file.to_string())?;
    match length {
        None => dump_iterator(
            reader.bytes().skip(seek).flat_map(|result| result),
            &mut *writer,
            settings,
        ),
        Some(n) => dump_iterator(
            reader
                .bytes()
                .skip(seek)
                .take(usize::from_str_radix(n, 10)?)
                .flat_map(|result| result),
            &mut *writer,
            settings,
        ),
    };
    Ok(())
}

fn create_dump_settings<'a>(args: &ArgMatches<'a>) -> Result<Config, failure::Error> {
    let columns = usize::from_str_radix(args.value_of("columns").unwrap_or("8"), 10)?;
    let format = args.value_of("format").unwrap_or("hex");
    let group_size = usize::from_str_radix(args.value_of("group-size").unwrap_or("2"), 10)?;
    let settings = Config::new()
        .format(Format::from(format.to_string()))
        .group_size(group_size)
        .columns(columns);
    if args.is_present("plain_hexdump") {
        Ok(settings
            .separator(false)
            .show_address(false)
            .show_interpretation(false))
    } else {
        Ok(settings)
    }
}

fn generate<'a>(args: Option<&ArgMatches<'a>>) -> Result<(), failure::Error> {
    let args = args.ok_or_else(|| format_err!("No arguments available"))?;
    let output_file = args.value_of("outfile").unwrap_or("stdout");
    let input_file = args.value_of("file").unwrap_or("stdin");
    let seek = usize::from_str_radix(args.value_of("seek").unwrap_or("0"), 10)?;
    let length = args.value_of("length");
    let reader = create_reader(input_file.to_string())?;
    let mut writer = create_writer(output_file.to_string())?;
    let lang = xxd::generate::Language::from(args.value_of("template").unwrap_or("c"));
    let template = Template::new(lang);
    let data: Vec<u8> = match length {
        None => reader
            .bytes()
            .skip(seek)
            .flat_map(|result| result)
            .collect(),
        Some(n) => reader
            .bytes()
            .skip(seek)
            .take(usize::from_str_radix(n, 10)?)
            .flat_map(|result| result)
            .collect(),
    };
    writer
        .write_fmt(format_args!("{}\n", template.render(&data)))
        .unwrap();
    Ok(())
}

// fn command_not_supported() -> Result<(), failure::Error> {
//     bail!(format_err!("Command not supported yet!"))
// }

fn report_error<T: Display>(error: &T) {
    eprintln!("xxd-rs: {}", error);
}
<commit_msg>correction compiler warning<commit_after>#[macro_use]
extern crate human_panic;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate xxd;

use cli::create_arg_parser;
use failure::*;
use xxd::dump::{dump_iterator, Config, Format};
use xxd::generate::{Render, Template};

use clap::ArgMatches;

use std::fmt::Display;
use std::io::{Read, Write};
use std::process::exit;

mod cli;

fn main() {
    setup_panic!();
    let matches = create_arg_parser().get_matches();
    match run(&matches) {
        Ok(_) => exit(0),
        Err(e) => {
            report_error(&e);
            exit(1)
        }
    }
}

fn run(args: &ArgMatches) -> Result<(), failure::Error> {
    match args.subcommand_name() {
        Some("dump") => dump(args.subcommand_matches("dump")),
        Some("generate") => generate(args.subcommand_matches("generate")),
        _ => bail!(format_err!("{}", args.usage())),
    }
}

pub fn create_reader(path: String) -> Result<Box<std::io::Read>, failure::Error> {
    match path.as_ref() {
        "stdin" => Ok(Box::new(std::io::stdin())),
        _ => {
            let file_reader = std::fs::File::open(path)?;
            Ok(Box::new(file_reader))
        }
    }
}

pub fn create_writer(path: String) -> Result<Box<std::io::Write>, failure::Error> {
    match path.as_ref() {
        "stdout" => Ok(Box::new(std::io::stdout())),
        _ => {
            let mut file_writer = std::fs::File::create(path)?;
            Ok(Box::new(file_writer))
        }
    }
}

fn dump<'a>(args: Option<&ArgMatches<'a>>) -> Result<(), failure::Error> {
    let args = args.ok_or_else(|| format_err!("No arguments available"))?;
    let output_file = args.value_of("outfile").unwrap_or("stdout");
    let input_file = args.value_of("file").unwrap_or("stdin");
    let seek = usize::from_str_radix(args.value_of("seek").unwrap_or("0"), 10)?;
    let length = args.value_of("length");
    let settings = create_dump_settings(args)?;
    let reader = create_reader(input_file.to_string())?;
    let mut writer = create_writer(output_file.to_string())?;
    match length {
        None => dump_iterator(
            reader.bytes().skip(seek).flat_map(|result| result),
            &mut *writer,
            settings,
        ),
        Some(n) => dump_iterator(
            reader
                .bytes()
                .skip(seek)
                .take(usize::from_str_radix(n, 10)?)
                .flat_map(|result| result),
            &mut *writer,
            settings,
        ),
    }
}

fn create_dump_settings<'a>(args: &ArgMatches<'a>) -> Result<Config, failure::Error> {
    let columns = usize::from_str_radix(args.value_of("columns").unwrap_or("8"), 10)?;
    let format = args.value_of("format").unwrap_or("hex");
    let group_size = usize::from_str_radix(args.value_of("group-size").unwrap_or("2"), 10)?;
    let settings = Config::new()
        .format(Format::from(format.to_string()))
        .group_size(group_size)
        .columns(columns);
    if args.is_present("plain_hexdump") {
        Ok(settings
            .separator(false)
            .show_address(false)
            .show_interpretation(false))
    } else {
        Ok(settings)
    }
}

fn generate<'a>(args: Option<&ArgMatches<'a>>) -> Result<(), failure::Error> {
    let args = args.ok_or_else(|| format_err!("No arguments available"))?;
    let output_file = args.value_of("outfile").unwrap_or("stdout");
    let input_file = args.value_of("file").unwrap_or("stdin");
    let seek = usize::from_str_radix(args.value_of("seek").unwrap_or("0"), 10)?;
    let length = args.value_of("length");
    let reader = create_reader(input_file.to_string())?;
    let mut writer = create_writer(output_file.to_string())?;
    let lang = xxd::generate::Language::from(args.value_of("template").unwrap_or("c"));
    let template = Template::new(lang);
    let data: Vec<u8> = match length {
        None => reader
            .bytes()
            .skip(seek)
            .flat_map(|result| result)
            .collect(),
        Some(n) => reader
            .bytes()
            .skip(seek)
            .take(usize::from_str_radix(n, 10)?)
            .flat_map(|result| result)
            .collect(),
    };
    writer
        .write_fmt(format_args!("{}\n", template.render(&data)))
        .unwrap();
    Ok(())
}

// fn command_not_supported() -> Result<(), failure::Error> {
//     bail!(format_err!("Command not supported yet!"))
// }

fn report_error<T: Display>(error: &T) {
    eprintln!("xxd-rs: {}", error);
}
<|endoftext|>


<commit_before>use crate::{
    connection_closer::ConnectionCloser,
    error_handler::ErrorHandler,
    executor::Executor,
    frames::Frames,
    id_sequence::IdSequence,
    internal_rpc::InternalRPCHandle,
    protocol::{AMQPClass, AMQPError, AMQPHardError},
    socket_state::SocketStateHandle,
    BasicProperties, Channel, ChannelState, Configuration, ConnectionState, ConnectionStatus,
    Error, Promise, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use log::{debug, error, log_enabled, trace, Level::Trace};
use parking_lot::Mutex;
use std::{collections::HashMap, fmt, sync::Arc};

#[derive(Clone)]
pub(crate) struct Channels {
    inner: Arc<Mutex<Inner>>,
    connection_status: ConnectionStatus,
    internal_rpc: InternalRPCHandle,
    executor: Arc<dyn Executor>,
    frames: Frames,
    error_handler: ErrorHandler,
}

impl Channels {
    pub(crate) fn new(
        configuration: Configuration,
        connection_status: ConnectionStatus,
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        Self {
            inner: Arc::new(Mutex::new(Inner::new(configuration, waker))),
            connection_status,
            internal_rpc,
            executor,
            frames,
            error_handler: ErrorHandler::default(),
        }
    }

    pub(crate) fn create(&self, connection_closer: Arc<ConnectionCloser>) -> Result<Channel> {
        self.inner.lock().create(
            self.connection_status.clone(),
            self.internal_rpc.clone(),
            self.frames.clone(),
            self.executor.clone(),
            connection_closer,
        )
    }

    pub(crate) fn create_zero(&self) {
        self.inner
            .lock()
            .create_channel(
                0,
                self.connection_status.clone(),
                self.internal_rpc.clone(),
                self.frames.clone(),
                self.executor.clone(),
                None,
            )
            .set_state(ChannelState::Connected);
    }

    pub(crate) fn get(&self, id: u16) -> Option<Channel> {
        self.inner.lock().channels.get(&id).cloned()
    }

    pub(crate) fn remove(&self, id: u16, error: Error) -> Result<()> {
        self.frames.clear_expected_replies(id, error);
        if self.inner.lock().channels.remove(&id).is_some() {
            Ok(())
        } else {
            Err(Error::InvalidChannel(id))
        }
    }

    pub(crate) fn receive_method(&self, id: u16, method: AMQPClass) -> Result<()> {
        self.get(id)
            .map(|channel| channel.receive_method(method))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn handle_content_header_frame(
        &self,
        id: u16,
        class_id: u16,
        size: u64,
        properties: BasicProperties,
    ) -> Result<()> {
        self.get(id)
            .map(|channel| channel.handle_content_header_frame(class_id, size, properties))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn handle_body_frame(&self, id: u16, payload: Vec<u8>) -> Result<()> {
        self.get(id)
            .map(|channel| channel.handle_body_frame(payload))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn set_connection_closing(&self) {
        self.connection_status.set_state(ConnectionState::Closing);
        for channel in self.inner.lock().channels.values() {
            channel.set_state(ChannelState::Closing);
        }
    }

    pub(crate) fn set_connection_closed(&self, error: Error) -> Result<()> {
        self.connection_status.set_state(ConnectionState::Closed);
        self.inner
            .lock()
            .channels
            .drain()
            .map(|(id, channel)| {
                self.frames.clear_expected_replies(id, error.clone());
                channel.set_state(ChannelState::Closed);
                channel.error_publisher_confirms(error.clone());
                channel.cancel_consumers()
            })
            .fold(Ok(()), Result::and)
    }

    pub(crate) fn set_connection_error(&self, error: Error) -> Result<()> {
        if let ConnectionState::Error = self.connection_status.state() {
            return Ok(());
        }

        error!("Connection error");
        self.connection_status.set_state(ConnectionState::Error);
        self.frames.drop_pending(error.clone());
        self.error_handler.on_error(error.clone());
        self.inner
            .lock()
            .channels
            .drain()
            .map(|(id, channel)| {
                self.frames.clear_expected_replies(id, error.clone());
                channel.set_state(ChannelState::Error);
                channel.error_publisher_confirms(error.clone());
                channel.error_consumers(error.clone())
            })
            .fold(Ok(()), Result::and)
    }

    pub(crate) fn flow(&self) -> bool {
        self.inner
            .lock()
            .channels
            .values()
            .all(|c| c.status().flow())
    }

    pub(crate) fn send_heartbeat(&self) -> Result<()> {
        debug!("send heartbeat");

        self.get(0)
            .map(|channel0| {
                let (promise, resolver) = Promise::new();

                if log_enabled!(Trace) {
                    promise.set_marker("Heartbeat".into());
                }

                channel0.send_frame(AMQPFrame::Heartbeat(0), resolver, None);
                self.internal_rpc.register_internal_future(promise)
            })
            .unwrap_or_else(|| {
                Err(Error::InvalidConnectionState(
                    self.connection_status.state(),
                ))
            })
    }

    pub(crate) fn handle_frame(&self, f: AMQPFrame) -> Result<()> {
        if let Err(err) = self.do_handle_frame(f) {
            self.set_connection_error(err.clone())?;
            Err(err)
        } else {
            Ok(())
        }
    }

    fn do_handle_frame(&self, f: AMQPFrame) -> Result<()> {
        trace!("will handle frame: {:?}", f);
        match f {
            AMQPFrame::ProtocolHeader(version) => {
                error!(
                    "we asked for AMQP {} but the server only supports AMQP {}",
                    ProtocolVersion::amqp_0_9_1(),
                    version
                );
                let error = Error::InvalidProtocolVersion(version);
                if let Some(resolver) = self.connection_status.connection_resolver() {
                    resolver.swear(Err(error.clone()));
                }
                return Err(error);
            }
            AMQPFrame::Method(channel_id, method) => {
                self.receive_method(channel_id, method)?;
            }
            AMQPFrame::Heartbeat(channel_id) => {
                if channel_id == 0 {
                    debug!("received heartbeat from server");
                } else {
                    error!("received invalid heartbeat on channel {}", channel_id);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("heartbeat frame received on channel {}", channel_id).into(),
                    );
                    if let Some(Err(error)) = self.get(0).map(|channel0| {
                        self.internal_rpc
                            .register_internal_future(channel0.connection_close(
                                error.get_id(),
                                error.get_message().as_str(),
                                0,
                                0,
                            ))
                    }) {
                        return Err(error);
                    }
                    return Err(Error::ProtocolError(error));
                }
            }
            AMQPFrame::Header(channel_id, class_id, header) => {
                if channel_id == 0 {
                    error!("received content header on channel {}", channel_id);
                    let error = AMQPError::new(
                        AMQPHardError::CHANNELERROR.into(),
                        format!("content header frame received on channel {}", channel_id).into(),
                    );
                    if let Some(Err(error)) = self.get(0).map(|channel0| {
                        self.internal_rpc
                            .register_internal_future(channel0.connection_close(
                                error.get_id(),
                                error.get_message().as_str(),
                                class_id,
                                0,
                            ))
                    }) {
                        return Err(error);
                    }
                    return Err(Error::ProtocolError(error));
                } else {
                    self.handle_content_header_frame(
                        channel_id,
                        class_id,
                        header.body_size,
                        header.properties,
                    )?;
                }
            }
            AMQPFrame::Body(channel_id, payload) => {
                self.handle_body_frame(channel_id, payload)?;
            }
        }
        Ok(())
    }

    pub(crate) fn set_error_handler<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.error_handler.set_handler(handler);
    }
}

impl fmt::Debug for Channels {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Channels");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("channels", &inner.channels.values())
                .field("channel_id", &inner.channel_id)
                .field("configuration", &inner.configuration);
        }
        debug
            .field("frames", &self.frames)
            .field("executor", &self.executor)
            .field("connection_status", &self.connection_status)
            .field("error_handler", &self.error_handler)
            .finish()
    }
}

struct Inner {
    channels: HashMap<u16, Channel>,
    channel_id: IdSequence<u16>,
    configuration: Configuration,
    waker: SocketStateHandle,
}

impl Inner {
    fn new(configuration: Configuration, waker: SocketStateHandle) -> Self {
        Self {
            channels: HashMap::default(),
            channel_id: IdSequence::new(false),
            configuration,
            waker,
        }
    }

    fn create_channel(
        &mut self,
        id: u16,
        connection_status: ConnectionStatus,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
        connection_closer: Option<Arc<ConnectionCloser>>,
    ) -> Channel {
        debug!("create channel with id {}", id);
        let channel = Channel::new(
            id,
            self.configuration.clone(),
            connection_status,
            self.waker.clone(),
            internal_rpc,
            frames,
            executor,
            connection_closer,
        );
        self.channels.insert(id, channel.clone_internal());
        channel
    }

    fn create(
        &mut self,
        connection_status: ConnectionStatus,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
        connection_closer: Arc<ConnectionCloser>,
    ) -> Result<Channel> {
        debug!("create channel");
        self.channel_id.set_max(self.configuration.channel_max());
        let first_id = self.channel_id.next();
        let mut looped = false;
        let mut id = first_id;
        while !looped || id < first_id {
            if id == 1 {
                looped = true;
            }
            if !self.channels.contains_key(&id) {
                return Ok(self.create_channel(
                    id,
                    connection_status,
                    internal_rpc,
                    frames,
                    executor,
                    Some(connection_closer),
                ));
            }
            id = self.channel_id.next();
        }
        Err(Error::ChannelsLimitReached)
    }
}
<commit_msg>fix reusing channel ids<commit_after>use crate::{
    connection_closer::ConnectionCloser,
    error_handler::ErrorHandler,
    executor::Executor,
    frames::Frames,
    id_sequence::IdSequence,
    internal_rpc::InternalRPCHandle,
    protocol::{AMQPClass, AMQPError, AMQPHardError},
    socket_state::SocketStateHandle,
    BasicProperties, Channel, ChannelState, Configuration, ConnectionState, ConnectionStatus,
    Error, Promise, Result,
};
use amq_protocol::frame::{AMQPFrame, ProtocolVersion};
use log::{debug, error, log_enabled, trace, Level::Trace};
use parking_lot::Mutex;
use std::{collections::HashMap, fmt, sync::Arc};

#[derive(Clone)]
pub(crate) struct Channels {
    inner: Arc<Mutex<Inner>>,
    connection_status: ConnectionStatus,
    internal_rpc: InternalRPCHandle,
    executor: Arc<dyn Executor>,
    frames: Frames,
    error_handler: ErrorHandler,
}

impl Channels {
    pub(crate) fn new(
        configuration: Configuration,
        connection_status: ConnectionStatus,
        waker: SocketStateHandle,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
    ) -> Self {
        Self {
            inner: Arc::new(Mutex::new(Inner::new(configuration, waker))),
            connection_status,
            internal_rpc,
            executor,
            frames,
            error_handler: ErrorHandler::default(),
        }
    }

    pub(crate) fn create(&self, connection_closer: Arc<ConnectionCloser>) -> Result<Channel> {
        self.inner.lock().create(
            self.connection_status.clone(),
            self.internal_rpc.clone(),
            self.frames.clone(),
            self.executor.clone(),
            connection_closer,
        )
    }

    pub(crate) fn create_zero(&self) {
        self.inner
            .lock()
            .create_channel(
                0,
                self.connection_status.clone(),
                self.internal_rpc.clone(),
                self.frames.clone(),
                self.executor.clone(),
                None,
            )
            .set_state(ChannelState::Connected);
    }

    pub(crate) fn get(&self, id: u16) -> Option<Channel> {
        self.inner.lock().channels.get(&id).cloned()
    }

    pub(crate) fn remove(&self, id: u16, error: Error) -> Result<()> {
        self.frames.clear_expected_replies(id, error);
        if self.inner.lock().channels.remove(&id).is_some() {
            Ok(())
        } else {
            Err(Error::InvalidChannel(id))
        }
    }

    pub(crate) fn receive_method(&self, id: u16, method: AMQPClass) -> Result<()> {
        self.get(id)
            .map(|channel| channel.receive_method(method))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn handle_content_header_frame(
        &self,
        id: u16,
        class_id: u16,
        size: u64,
        properties: BasicProperties,
    ) -> Result<()> {
        self.get(id)
            .map(|channel| channel.handle_content_header_frame(class_id, size, properties))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn handle_body_frame(&self, id: u16, payload: Vec<u8>) -> Result<()> {
        self.get(id)
            .map(|channel| channel.handle_body_frame(payload))
            .unwrap_or_else(|| Err(Error::InvalidChannel(id)))
    }

    pub(crate) fn set_connection_closing(&self) {
        self.connection_status.set_state(ConnectionState::Closing);
        for channel in self.inner.lock().channels.values() {
            channel.set_state(ChannelState::Closing);
        }
    }

    pub(crate) fn set_connection_closed(&self, error: Error) -> Result<()> {
        self.connection_status.set_state(ConnectionState::Closed);
        self.inner
            .lock()
            .channels
            .drain()
            .map(|(id, channel)| {
                self.frames.clear_expected_replies(id, error.clone());
                channel.set_state(ChannelState::Closed);
                channel.error_publisher_confirms(error.clone());
                channel.cancel_consumers()
            })
            .fold(Ok(()), Result::and)
    }

    pub(crate) fn set_connection_error(&self, error: Error) -> Result<()> {
        if let ConnectionState::Error = self.connection_status.state() {
            return Ok(());
        }

        error!("Connection error");
        self.connection_status.set_state(ConnectionState::Error);
        self.frames.drop_pending(error.clone());
        self.error_handler.on_error(error.clone());
        self.inner
            .lock()
            .channels
            .drain()
            .map(|(id, channel)| {
                self.frames.clear_expected_replies(id, error.clone());
                channel.set_state(ChannelState::Error);
                channel.error_publisher_confirms(error.clone());
                channel.error_consumers(error.clone())
            })
            .fold(Ok(()), Result::and)
    }

    pub(crate) fn flow(&self) -> bool {
        self.inner
            .lock()
            .channels
            .values()
            .all(|c| c.status().flow())
    }

    pub(crate) fn send_heartbeat(&self) -> Result<()> {
        debug!("send heartbeat");

        self.get(0)
            .map(|channel0| {
                let (promise, resolver) = Promise::new();

                if log_enabled!(Trace) {
                    promise.set_marker("Heartbeat".into());
                }

                channel0.send_frame(AMQPFrame::Heartbeat(0), resolver, None);
                self.internal_rpc.register_internal_future(promise)
            })
            .unwrap_or_else(|| {
                Err(Error::InvalidConnectionState(
                    self.connection_status.state(),
                ))
            })
    }

    pub(crate) fn handle_frame(&self, f: AMQPFrame) -> Result<()> {
        if let Err(err) = self.do_handle_frame(f) {
            self.set_connection_error(err.clone())?;
            Err(err)
        } else {
            Ok(())
        }
    }

    fn do_handle_frame(&self, f: AMQPFrame) -> Result<()> {
        trace!("will handle frame: {:?}", f);
        match f {
            AMQPFrame::ProtocolHeader(version) => {
                error!(
                    "we asked for AMQP {} but the server only supports AMQP {}",
                    ProtocolVersion::amqp_0_9_1(),
                    version
                );
                let error = Error::InvalidProtocolVersion(version);
                if let Some(resolver) = self.connection_status.connection_resolver() {
                    resolver.swear(Err(error.clone()));
                }
                return Err(error);
            }
            AMQPFrame::Method(channel_id, method) => {
                self.receive_method(channel_id, method)?;
            }
            AMQPFrame::Heartbeat(channel_id) => {
                if channel_id == 0 {
                    debug!("received heartbeat from server");
                } else {
                    error!("received invalid heartbeat on channel {}", channel_id);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("heartbeat frame received on channel {}", channel_id).into(),
                    );
                    if let Some(Err(error)) = self.get(0).map(|channel0| {
                        self.internal_rpc
                            .register_internal_future(channel0.connection_close(
                                error.get_id(),
                                error.get_message().as_str(),
                                0,
                                0,
                            ))
                    }) {
                        return Err(error);
                    }
                    return Err(Error::ProtocolError(error));
                }
            }
            AMQPFrame::Header(channel_id, class_id, header) => {
                if channel_id == 0 {
                    error!("received content header on channel {}", channel_id);
                    let error = AMQPError::new(
                        AMQPHardError::CHANNELERROR.into(),
                        format!("content header frame received on channel {}", channel_id).into(),
                    );
                    if let Some(Err(error)) = self.get(0).map(|channel0| {
                        self.internal_rpc
                            .register_internal_future(channel0.connection_close(
                                error.get_id(),
                                error.get_message().as_str(),
                                class_id,
                                0,
                            ))
                    }) {
                        return Err(error);
                    }
                    return Err(Error::ProtocolError(error));
                } else {
                    self.handle_content_header_frame(
                        channel_id,
                        class_id,
                        header.body_size,
                        header.properties,
                    )?;
                }
            }
            AMQPFrame::Body(channel_id, payload) => {
                self.handle_body_frame(channel_id, payload)?;
            }
        }
        Ok(())
    }

    pub(crate) fn set_error_handler<E: FnMut(Error) + Send + 'static>(&self, handler: E) {
        self.error_handler.set_handler(handler);
    }
}

impl fmt::Debug for Channels {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Channels");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("channels", &inner.channels.values())
                .field("channel_id", &inner.channel_id)
                .field("configuration", &inner.configuration);
        }
        debug
            .field("frames", &self.frames)
            .field("executor", &self.executor)
            .field("connection_status", &self.connection_status)
            .field("error_handler", &self.error_handler)
            .finish()
    }
}

struct Inner {
    channels: HashMap<u16, Channel>,
    channel_id: IdSequence<u16>,
    configuration: Configuration,
    waker: SocketStateHandle,
}

impl Inner {
    fn new(configuration: Configuration, waker: SocketStateHandle) -> Self {
        Self {
            channels: HashMap::default(),
            channel_id: IdSequence::new(false),
            configuration,
            waker,
        }
    }

    fn create_channel(
        &mut self,
        id: u16,
        connection_status: ConnectionStatus,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
        connection_closer: Option<Arc<ConnectionCloser>>,
    ) -> Channel {
        debug!("create channel with id {}", id);
        let channel = Channel::new(
            id,
            self.configuration.clone(),
            connection_status,
            self.waker.clone(),
            internal_rpc,
            frames,
            executor,
            connection_closer,
        );
        self.channels.insert(id, channel.clone_internal());
        channel
    }

    fn create(
        &mut self,
        connection_status: ConnectionStatus,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        executor: Arc<dyn Executor>,
        connection_closer: Arc<ConnectionCloser>,
    ) -> Result<Channel> {
        debug!("create channel");
        self.channel_id.set_max(self.configuration.channel_max());
        let first_id = self.channel_id.next();
        let mut id = first_id;
        let mut met_first_id = false;
        loop {
            if id == first_id {
                if met_first_id {
                    break;
                }
                met_first_id = true;
            }
            if !self.channels.contains_key(&id) {
                return Ok(self.create_channel(
                    id,
                    connection_status,
                    internal_rpc,
                    frames,
                    executor,
                    Some(connection_closer),
                ));
            }
            id = self.channel_id.next();
        }
        Err(Error::ChannelsLimitReached)
    }
}
<|endoftext|>


<commit_before>use ffi::{pk11, sec};
use std::{ptr, mem};
use libc::{c_uint, c_int, c_void};

#[allow(non_camel_case_types)]
pub enum RSAPadding
{
    PKCS1v15,
    OAEP_MGF1_SHA1,
    OAEP_MGF1_SHA224,
    OAEP_MGF1_SHA256,
    OAEP_MGF1_SHA384,
    OAEP_MGF1_SHA512,
}

enum RSAPaddingParam
{
    NullParam,
    OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS),
}

impl RSAPaddingParam
{
    unsafe fn to_secitem(&self) -> sec::SECItem
    {
        match *self
        {
            NullParam => sec::SECItem::from_struct(&()),
            OAEPParam(ref param) => sec::SECItem::from_struct(param),
        }
    }
}

impl RSAPadding
{
    fn to_ckm(&self) -> pk11::CK_MECHANISM_TYPE
    {
        match *self
        {
            PKCS1v15 => pk11::CKM_RSA_PKCS,
            OAEP_MGF1_SHA1
            | OAEP_MGF1_SHA224
            | OAEP_MGF1_SHA256
            | OAEP_MGF1_SHA384
            | OAEP_MGF1_SHA512 => pk11::CKM_RSA_PKCS_OAEP,
        }
    }

    fn get_param(&self) -> RSAPaddingParam
    {
        match *self
        {
            PKCS1v15 => NullParam,
            OAEP_MGF1_SHA1 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_1)),
            OAEP_MGF1_SHA224 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_224)),
            OAEP_MGF1_SHA256 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_256)),
            OAEP_MGF1_SHA384 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_384)),
            OAEP_MGF1_SHA512 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_512)),
        }
    }
}

pub struct RSAPrivateKey
{
    key: *mut pk11::SECKEYPrivateKey,
}

impl RSAPrivateKey
{
    pub fn load(data: &[u8]) -> ::NSSResult<RSAPrivateKey>
    {
        try!(::nss::init());
        unsafe
        {
            let mut der = sec::SECItem::new(data);
            let slot = try_ptr!(pk11::PK11_GetInternalKeySlot());
            let mut key = ptr::null_mut();

            try!(pk11::PK11_ImportDERPrivateKeyInfoAndReturnKey(slot, &mut der, ptr::null_mut(), ptr::null_mut(), false, true, pk11::KU_ALL, &mut key, ptr::null_mut()).to_result());

            pk11::PK11_FreeSlot(slot);
            Ok(RSAPrivateKey { key: key })
        }
    }

    pub fn gen(key_size_bits: uint) -> ::NSSResult<RSAPrivateKey>
    {
        try!(::nss::init());
        unsafe
        {
            let slot = try_ptr!(pk11::PK11_GetInternalKeySlot());
            let mut param = pk11::PK11RSAGenParams { key_size_bits: key_size_bits as c_int, pe: 65537, };
            let param_ptr = mem::transmute::<_, *mut c_void>(&mut param);
            let mut pubkey = ptr::null_mut();
            let privkey = try_ptr!(pk11::PK11_GenerateKeyPair(slot, pk11::CKM_RSA_PKCS_KEY_PAIR_GEN, param_ptr, &mut pubkey, false, false, ptr::null_mut()));

            pk11::SECKEY_DestroyPublicKey(pubkey);
            pk11::PK11_FreeSlot(slot);

            Ok(RSAPrivateKey { key: privkey })
        }
    }

    pub fn save(&self) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let secitem = try_ptr!(pk11::PK11_ExportDERPrivateKeyInfo(self.key, ptr::null_mut()));
            let result = (&*secitem).copy_buf();
            sec::SECItem::free(secitem);
            Ok(result)
        }
    }

    pub fn key_len(&self) -> uint
    {
        match RSAPublicKey::from_private(self)
        {
            Err(..) => 0,
            Ok(public) => public.key_len(),
        }
    }

    pub fn encrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        let public = try!(RSAPublicKey::from_private(self));
        public.encrypt(padding, data)
    }

    pub fn decrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let mut out = Vec::with_capacity(self.key_len());
            let mut outlen = 0;

            let params = padding.get_param();
            let mut secitem = params.to_secitem();

            try!(pk11::PK11_PrivDecrypt(self.key, padding.to_ckm(), &mut secitem, out.as_mut_ptr(),
                                        &mut outlen, out.capacity() as c_uint, data.as_ptr(), data.len() as c_uint).to_result());
            out.set_len(outlen as uint);
            Ok(out)
        }
    }
}

impl Drop for RSAPrivateKey
{
    fn drop(&mut self)
    {
        unsafe
        {
            pk11::SECKEY_DestroyPrivateKey(self.key);
        }
    }
}

pub struct RSAPublicKey
{
    key: *mut pk11::SECKEYPublicKey,
}

impl RSAPublicKey
{
    pub fn load(data: &[u8]) -> ::NSSResult<RSAPublicKey>
    {
        try!(::nss::init());
        unsafe
        {
            let der = sec::SECItem::new(data);
            let spki = try_ptr!(pk11::SECKEY_DecodeDERSubjectPublicKeyInfo(&der));
            let key = try_ptr!(pk11::SECKEY_ExtractPublicKey(spki as *const pk11::CERTSubjectPublicKeyInfo));

            pk11::SECKEY_DestroySubjectPublicKeyInfo(spki);
            Ok(RSAPublicKey { key: key })
        }
    }

    pub fn save(&self) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let secitem = try_ptr!(pk11::SECKEY_EncodeDERSubjectPublicKeyInfo(self.key as *const _));
            let result = (&*secitem).copy_buf();
            sec::SECItem::free(secitem);
            Ok(result)
        }
    }

    pub fn key_len(&self) -> uint
    {
        unsafe
        {
            pk11::SECKEY_PublicKeyStrength(self.key as *const pk11::SECKEYPublicKey) as uint
        }
    }

    pub fn from_private(input: &RSAPrivateKey) -> ::NSSResult<RSAPublicKey>
    {
        unsafe
        {
            let mypub = try_ptr!(pk11::SECKEY_ConvertToPublicKey(input.key));
            Ok(RSAPublicKey { key: mypub })
        }
    }

    pub fn encrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let mut out = Vec::with_capacity(self.key_len());
            let mut outlen = 0;

            let params = padding.get_param();
            let mut secitem = params.to_secitem();

            try!(pk11::PK11_PubEncrypt(self.key, padding.to_ckm(), &mut secitem, out.as_mut_ptr(),
                                       &mut outlen, out.capacity() as c_uint, data.as_ptr(), data.len() as c_uint,
                                       ptr::null_mut()).to_result());
            out.set_len(outlen as uint);
            Ok(out)
        }
    }
}

impl Drop for RSAPublicKey
{
    fn drop(&mut self)
    {
        unsafe
        {
            pk11::SECKEY_DestroyPublicKey(self.key);
        }
    }
}
<commit_msg>Change parameter to prevent issue with exporting<commit_after>use ffi::{pk11, sec};
use std::{ptr, mem};
use libc::{c_uint, c_int, c_void};

#[allow(non_camel_case_types)]
pub enum RSAPadding
{
    PKCS1v15,
    OAEP_MGF1_SHA1,
    OAEP_MGF1_SHA224,
    OAEP_MGF1_SHA256,
    OAEP_MGF1_SHA384,
    OAEP_MGF1_SHA512,
}

enum RSAPaddingParam
{
    NullParam,
    OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS),
}

impl RSAPaddingParam
{
    unsafe fn to_secitem(&self) -> sec::SECItem
    {
        match *self
        {
            NullParam => sec::SECItem::from_struct(&()),
            OAEPParam(ref param) => sec::SECItem::from_struct(param),
        }
    }
}

impl RSAPadding
{
    fn to_ckm(&self) -> pk11::CK_MECHANISM_TYPE
    {
        match *self
        {
            PKCS1v15 => pk11::CKM_RSA_PKCS,
            OAEP_MGF1_SHA1
            | OAEP_MGF1_SHA224
            | OAEP_MGF1_SHA256
            | OAEP_MGF1_SHA384
            | OAEP_MGF1_SHA512 => pk11::CKM_RSA_PKCS_OAEP,
        }
    }

    fn get_param(&self) -> RSAPaddingParam
    {
        match *self
        {
            PKCS1v15 => NullParam,
            OAEP_MGF1_SHA1 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_1)),
            OAEP_MGF1_SHA224 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_224)),
            OAEP_MGF1_SHA256 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_256)),
            OAEP_MGF1_SHA384 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_384)),
            OAEP_MGF1_SHA512 => OAEPParam(pk11::CK_RSA_PKCS_OAEP_PARAMS::from_algorithm(pk11::CKM_SHA_512)),
        }
    }
}

pub struct RSAPrivateKey
{
    key: *mut pk11::SECKEYPrivateKey,
}

impl RSAPrivateKey
{
    pub fn load(data: &[u8]) -> ::NSSResult<RSAPrivateKey>
    {
        try!(::nss::init());
        unsafe
        {
            let mut der = sec::SECItem::new(data);
            let slot = try_ptr!(pk11::PK11_GetInternalKeySlot());
            let mut key = ptr::null_mut();

            try!(pk11::PK11_ImportDERPrivateKeyInfoAndReturnKey(slot, &mut der, ptr::null_mut(), ptr::null_mut(), false, false, pk11::KU_ALL, &mut key, ptr::null_mut()).to_result());

            pk11::PK11_FreeSlot(slot);
            Ok(RSAPrivateKey { key: key })
        }
    }

    pub fn gen(key_size_bits: uint) -> ::NSSResult<RSAPrivateKey>
    {
        try!(::nss::init());
        unsafe
        {
            let slot = try_ptr!(pk11::PK11_GetInternalKeySlot());
            let mut param = pk11::PK11RSAGenParams { key_size_bits: key_size_bits as c_int, pe: 65537, };
            let param_ptr = mem::transmute::<_, *mut c_void>(&mut param);
            let mut pubkey = ptr::null_mut();
            let privkey = try_ptr!(pk11::PK11_GenerateKeyPair(slot, pk11::CKM_RSA_PKCS_KEY_PAIR_GEN, param_ptr, &mut pubkey, false, false, ptr::null_mut()));

            pk11::SECKEY_DestroyPublicKey(pubkey);
            pk11::PK11_FreeSlot(slot);

            Ok(RSAPrivateKey { key: privkey })
        }
    }

    pub fn save(&self) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let secitem = try_ptr!(pk11::PK11_ExportDERPrivateKeyInfo(self.key, ptr::null_mut()));
            let result = (&*secitem).copy_buf();
            sec::SECItem::free(secitem);
            Ok(result)
        }
    }

    pub fn key_len(&self) -> uint
    {
        match RSAPublicKey::from_private(self)
        {
            Err(..) => 0,
            Ok(public) => public.key_len(),
        }
    }

    pub fn encrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        let public = try!(RSAPublicKey::from_private(self));
        public.encrypt(padding, data)
    }

    pub fn decrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let mut out = Vec::with_capacity(self.key_len());
            let mut outlen = 0;

            let params = padding.get_param();
            let mut secitem = params.to_secitem();

            try!(pk11::PK11_PrivDecrypt(self.key, padding.to_ckm(), &mut secitem, out.as_mut_ptr(),
                                        &mut outlen, out.capacity() as c_uint, data.as_ptr(), data.len() as c_uint).to_result());
            out.set_len(outlen as uint);
            Ok(out)
        }
    }
}

impl Drop for RSAPrivateKey
{
    fn drop(&mut self)
    {
        unsafe
        {
            pk11::SECKEY_DestroyPrivateKey(self.key);
        }
    }
}

pub struct RSAPublicKey
{
    key: *mut pk11::SECKEYPublicKey,
}

impl RSAPublicKey
{
    pub fn load(data: &[u8]) -> ::NSSResult<RSAPublicKey>
    {
        try!(::nss::init());
        unsafe
        {
            let der = sec::SECItem::new(data);
            let spki = try_ptr!(pk11::SECKEY_DecodeDERSubjectPublicKeyInfo(&der));
            let key = try_ptr!(pk11::SECKEY_ExtractPublicKey(spki as *const pk11::CERTSubjectPublicKeyInfo));

            pk11::SECKEY_DestroySubjectPublicKeyInfo(spki);
            Ok(RSAPublicKey { key: key })
        }
    }

    pub fn save(&self) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let secitem = try_ptr!(pk11::SECKEY_EncodeDERSubjectPublicKeyInfo(self.key as *const _));
            let result = (&*secitem).copy_buf();
            sec::SECItem::free(secitem);
            Ok(result)
        }
    }

    pub fn key_len(&self) -> uint
    {
        unsafe
        {
            pk11::SECKEY_PublicKeyStrength(self.key as *const pk11::SECKEYPublicKey) as uint
        }
    }

    pub fn from_private(input: &RSAPrivateKey) -> ::NSSResult<RSAPublicKey>
    {
        unsafe
        {
            let mypub = try_ptr!(pk11::SECKEY_ConvertToPublicKey(input.key));
            Ok(RSAPublicKey { key: mypub })
        }
    }

    pub fn encrypt(&self, padding: RSAPadding, data: &[u8]) -> ::NSSResult<Vec<u8>>
    {
        unsafe
        {
            let mut out = Vec::with_capacity(self.key_len());
            let mut outlen = 0;

            let params = padding.get_param();
            let mut secitem = params.to_secitem();

            try!(pk11::PK11_PubEncrypt(self.key, padding.to_ckm(), &mut secitem, out.as_mut_ptr(),
                                       &mut outlen, out.capacity() as c_uint, data.as_ptr(), data.len() as c_uint,
                                       ptr::null_mut()).to_result());
            out.set_len(outlen as uint);
            Ok(out)
        }
    }
}

impl Drop for RSAPublicKey
{
    fn drop(&mut self)
    {
        unsafe
        {
            pk11::SECKEY_DestroyPublicKey(self.key);
        }
    }
}
<|endoftext|>


<commit_before>use std::mem;
use std::f64;

use itertools::Itertools;
use ndarray::prelude::*;
use bit_vec::BitVec;
use ordered_float::NotNaN;
use ndarray_rand::RandomExt;
use rand::distributions::Range;

use bio::stats::{Prob, LogProb};

use io::codebook::Codebook;

pub type Readout = BitVec;

pub fn estimate<I: Iterator<Item=(String, Readout)>>(
    codebook: &Codebook,
    readouts: I
) -> (Vec<Prob>, Vec<Prob>) {
    let max_diff = 0.001;
    let shape = (codebook.len(), codebook.N as usize);

    let mut observed_01_errors = Array2::from_elem(shape, 0.0);
    let mut observed_10_errors = Array2::from_elem(shape, 0.0);
    let mut feat_count = Array2::from_elem((codebook.len(), 1), 0.0);
    for (feat, readout) in readouts {
        let feature_id = codebook.get_id(&feat);
        let encoding = codebook.record(feature_id).codeword();

        // count features
        feat_count[(feature_id, 0)] += 1.0;

        for k in 0..codebook.N as usize {
            if encoding[k] == false {
                // count 01 errors
                if readout[k] == true {
                    observed_01_errors[(feature_id, k)] += 1.0;
                }
            } else {
                // count 10 errors
                if readout[k] == false {
                    observed_10_errors[(feature_id, k)] += 1.0;
                }
            }
        }
    }

    // calculate observed (joint) error rates
    let observed_p0 = observed_01_errors / &feat_count;
    let observed_p1 = observed_10_errors / &feat_count;

    // run EM algorithm to estimate error rates
    let mut prob_no_error = Array2::from_elem(
        (codebook.len(), codebook.N as usize), LogProb::ln_one()
    );

    let mut p0 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));
    let mut p1 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));

    let mut p0_last = Array1::from_elem(codebook.N as usize, 0.0);
    let mut p1_last = Array1::from_elem(codebook.N as usize, 0.0);
    for i in 0..1000 {
        // E-step: estimate probability to have no errors in rest of readout
        for feature_id in 0..codebook.len() {
            let encoding = codebook.record(feature_id).codeword();
            for k in 0..codebook.N as usize {
                prob_no_error[(feature_id, k)] = LogProb((0..codebook.N as usize).filter_map(|l| {
                    if l != k {
                        let p_err = if encoding[l] == false {
                            p0[k]
                        } else {
                            p1[k]
                        };
                        Some((1.0f64 - p_err).ln())
                    } else { None }
                }).sum());
            }
        }
        //debug!("{:?}", prob_no_error);

        // M-step: find p0 and p1 that maximizes the likelihood
        mem::swap(&mut p0, &mut p0_last);
        mem::swap(&mut p1, &mut p1_last);
        for k in 0..codebook.N as usize {
            let p = |observed_p: &Array2<f64>| {
                (0..codebook.len()).filter_map(|feature_id| {
                    if feat_count[(feature_id, 0)] < 20.0 {
                        None
                    } else {
                        let p = observed_p[(feature_id, k)] / prob_no_error[(feature_id, k)].exp();
                        if p > 1.0 { Some(1.0) } else { Some(p) }
                    }
                }).sum::<f64>() / codebook.len() as f64
            };
            // solve least squares: take mean of all estimates
            // TODO discard features with too low count?
            p0[k] = p(&observed_p0);
            p1[k] = p(&observed_p1);
        }

        // check convergence
        let abs = |d: &f64| NotNaN::new(d.abs()).unwrap();
        debug!("{:?}", p0);
        //debug!("{:?}", p0_last);
        // let convergence = **(&p0 - &p0_last).map(&abs).iter().max().unwrap() <= max_diff &&
        //                   **(&p1 - &p1_last).map(&abs).iter().max().unwrap() <= max_diff;
        let convergence = p0.all_close(&p0_last, max_diff) && p1.all_close(&p1_last, max_diff);

        if convergence && i >= 5 {
            debug!("convergence reached");
            break;
        }

        debug!("EM-algorithm step {} of {}.", i + 1, 100);
    }

    (
        p0.iter().map(|p| Prob(*p)).collect_vec(),
        p1.iter().map(|p| Prob(*p)).collect_vec()
    )
}
<commit_msg>Fixed formula for M-step. Fail if not enough counts available.<commit_after>use std::mem;
use std::f64;

use itertools::Itertools;
use ndarray::prelude::*;
use bit_vec::BitVec;
use ordered_float::NotNaN;
use ndarray_rand::RandomExt;
use rand::distributions::Range;

use bio::stats::{Prob, LogProb};

use io::codebook::Codebook;

pub type Readout = BitVec;

pub fn estimate<I: Iterator<Item=(String, Readout)>>(
    codebook: &Codebook,
    readouts: I
) -> (Vec<Prob>, Vec<Prob>) {
    let max_diff = 0.001;
    let shape = (codebook.len(), codebook.N as usize);

    let mut observed_01_errors = Array2::from_elem(shape, 0.0);
    let mut observed_10_errors = Array2::from_elem(shape, 0.0);
    let mut feat_count = Array2::from_elem((codebook.len(), 1), 0.0);
    for (feat, readout) in readouts {
        let feature_id = codebook.get_id(&feat);
        let encoding = codebook.record(feature_id).codeword();

        // count features
        feat_count[(feature_id, 0)] += 1.0;

        for k in 0..codebook.N as usize {
            if encoding[k] == false {
                // count 01 errors
                if readout[k] == true {
                    observed_01_errors[(feature_id, k)] += 1.0;
                }
            } else {
                // count 10 errors
                if readout[k] == false {
                    observed_10_errors[(feature_id, k)] += 1.0;
                }
            }
        }
    }

    // calculate observed (joint) error rates
    let observed_p0 = observed_01_errors / &feat_count;
    let observed_p1 = observed_10_errors / &feat_count;

    // run EM algorithm to estimate error rates
    let mut prob_no_error = Array2::from_elem(
        (codebook.len(), codebook.N as usize), LogProb::ln_one()
    );

    let mut p0 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));
    let mut p1 = Array1::random(codebook.N as usize, Range::new(0.0, 1.0));

    let mut p0_last = Array1::from_elem(codebook.N as usize, 0.0);
    let mut p1_last = Array1::from_elem(codebook.N as usize, 0.0);
    for i in 0..1000 {
        // E-step: estimate probability to have no errors in rest of readout
        for feature_id in 0..codebook.len() {
            let encoding = codebook.record(feature_id).codeword();
            for k in 0..codebook.N as usize {
                let p = LogProb((0..codebook.N as usize).filter_map(|l| {
                    if l != k {
                        let p_err = if encoding[l] == false {
                            p0[k]
                        } else {
                            p1[k]
                        };
                        Some(*LogProb((p_err as f64).ln()).ln_one_minus_exp())
                    } else { None }
                }).sum());
                prob_no_error[(feature_id, k)] = p;
            }
        }

        // M-step: find p0 and p1 that maximizes the likelihood
        mem::swap(&mut p0, &mut p0_last);
        mem::swap(&mut p1, &mut p1_last);
        for k in 0..codebook.N as usize {
            // solve least squares: take mean of all estimates
            let p = |observed_p: &Array2<f64>, bit: bool| {
                let mut count = 0.0;
                let p = (0..codebook.len()).filter_map(|feature_id| {
                    let encoding = codebook.record(feature_id).codeword();
                    if feat_count[(feature_id, 0)] < 10.0 || encoding[k] != bit {
                        None
                    } else {
                        count += 1.0;
                        let p = observed_p[(feature_id, k)] * prob_no_error[(feature_id, k)].exp();
                        Some(p)
                    }
                }).sum::<f64>() / count;
                if count == 0.0 {
                    panic!("not enough counts to estimate error rate");
                }
                p
            };
            //debug!("----------- p0 ------------");
            p0[k] = p(&observed_p0, false);
            //debug!("DEBUG: p0={}", p0[k]);
            //debug!("----------- p1 ------------");
            p1[k] = p(&observed_p1, true);
            //debug!("DEBUG: p1={}", p1[k]);
        }

        // check convergence
        let abs = |d: &f64| NotNaN::new(d.abs()).unwrap();
        //debug!("{:?}", p0);
        let convergence = p0.all_close(&p0_last, max_diff) && p1.all_close(&p1_last, max_diff);

        if convergence && i >= 5 {
            debug!("convergence reached");
            debug!("{:?}", p1);
            break;
        }

        debug!("EM-algorithm step {} of {}.", i + 1, 100);
    }

    (
        p0.iter().map(|p| Prob(*p)).collect_vec(),
        p1.iter().map(|p| Prob(*p)).collect_vec()
    )
}
<|endoftext|>


<commit_before>/// A type for which the layout is determined by its C++ definition.
///
/// This trait serves the following two related purposes.
///
///

///
/// ## Safely unifying occurrences of the same extern type
///
/// `ExternType` makes it possible for CXX to safely share a consistent Rust
/// type across multiple #\[cxx::bridge\] invocations that refer to a common
/// extern C++ type.
///
/// In the following snippet, two #\[cxx::bridge\] invocations in different
/// files (possibly different crates) both contain function signatures involving
/// the same C++ type `example::Demo`. If both were written just containing
/// `type Demo;`, then both macro expansions would produce their own separate
/// Rust type called `Demo` and thus the compiler wouldn't allow us to take the
/// `Demo` returned by `file1::ffi::create_demo` and pass it as the `Demo`
/// argument accepted by `file2::ffi::take_ref_demo`. Instead, one of the two
/// `Demo`s has been defined as an extern type alias of the other, making them
/// the same type in Rust. The CXX code generator will use an automatically
/// generated `ExternType` impl emitted in file1 to statically verify that in
/// file2 `crate::file1::ffi::Demo` really does refer to the C++ type
/// `example::Demo` as expected in file2.
///
/// ```no_run
/// // file1.rs
/// # mod file1 {
/// #[cxx::bridge(namespace = example)]
/// pub mod ffi {
///     extern "C" {
///         type Demo;
///
///         fn create_demo() -> UniquePtr<Demo>;
///     }
/// }
/// # }
///
/// // file2.rs
/// #[cxx::bridge(namespace = example)]
/// pub mod ffi {
///     extern "C" {
///         type Demo = crate::file1::ffi::Demo;
///
///         fn take_ref_demo(demo: &Demo);
///     }
/// }
/// #
/// # fn main() {}
/// ```
///
///

///
/// ## Integrating with bindgen-generated types
///
/// Handwritten `ExternType` impls make it possible to plug in a data structure
/// emitted by bindgen as the definition of a C++ type emitted by CXX.
///
/// By writing the unsafe `ExternType` impl, the programmer asserts that the C++
/// namespace and type name given in the type id refers to a C++ type that is
/// equivalent to Rust type that is the `Self` type of the impl.
///
/// ```no_run
/// # const _: &str = stringify! {
/// mod folly_sys;  // the bindgen-generated bindings
/// # };
/// # mod folly_sys {
/// #     #[repr(transparent)]
/// #     pub struct StringPiece([usize; 2]);
/// # }
///
/// use cxx::{type_id, ExternType};
///
/// unsafe impl ExternType for folly_sys::StringPiece {
///     type Id = type_id!("folly::StringPiece");
///     type Kind = cxx::kind::Opaque;
/// }
///
/// #[cxx::bridge(namespace = folly)]
/// pub mod ffi {
///     extern "C" {
///         include!("rust_cxx_bindings.h");
///
///         type StringPiece = crate::folly_sys::StringPiece;
///
///         fn print_string_piece(s: &StringPiece);
///     }
/// }
///
/// // Now if we construct a StringPiece or obtain one through one
/// // of the bindgen-generated signatures, we are able to pass it
/// // along to ffi::print_string_piece.
/// #
/// # fn main() {}
/// ```
pub unsafe trait ExternType {
    /// A type-level representation of the type's C++ namespace and type name.
    ///
    /// This will always be defined using `type_id!` in the following form:
    ///
    /// ```
    /// # struct TypeName;
    /// # unsafe impl cxx::ExternType for TypeName {
    /// type Id = cxx::type_id!("name::space::of::TypeName");
    /// #     type Kind = cxx::kind::Opaque;
    /// # }
    /// ```
    type Id;

    /// Either [`cxx::kind::Opaque`] or [`cxx::kind::Trivial`].
    ///
    /// [`cxx::kind::Opaque`]: kind::Opaque
    /// [`cxx::kind::Trivial`]: kind::Trivial
    ///
    /// A C++ type is only okay to hold and pass around by value in Rust if its
    /// [move constructor is trivial] and it has no destructor. In CXX, these
    /// are called Trivial extern C++ types, while types with nontrivial move
    /// behavior or a destructor must be considered Opaque and handled by Rust
    /// only behind an indirection, such as a reference or UniquePtr.
    ///
    /// [move constructor is trivial]: https://en.cppreference.com/w/cpp/types/is_move_constructible
    ///
    /// If you believe your C++ type reflected by this ExternType impl is indeed
    /// trivial, you can specify:
    ///
    /// ```
    /// # struct TypeName;
    /// # unsafe impl cxx::ExternType for TypeName {
    /// #     type Id = cxx::type_id!("name::space::of::TypeName");
    /// type Kind = cxx::kind::Trivial;
    /// # }
    /// ```
    ///
    /// which will enable you to pass it into C++ functions by value, return it
    /// by value, and include it in `struct`s that you have declared to
    /// `cxx::bridge`. Your claim about the triviality of the C++ type will be
    /// checked by a `static_assert` in the generated C++ side of the binding.
    type Kind;
}

/// Marker types identifying Rust's knowledge about an extern C++ type.
///
/// These markers are used in the [`Kind`][ExternType::Kind] associated type in
/// impls of the `ExternType` trait. Refer to the documentation of `Kind` for an
/// overview of their purpose.
pub mod kind {
    /// An opaque type which cannot be passed or held by value within Rust.
    ///
    /// Rust's move semantics are such that every move is equivalent to a
    /// memcpy. This is incompatible in general with C++'s constructor-based
    /// move semantics, so a C++ type which has a destructor or nontrivial move
    /// constructor must never exist by value in Rust. In CXX, such types are
    /// called opaque C++ types.
    ///
    /// When passed across an FFI boundary, an opaque C++ type must be behind an
    /// indirection such as a reference or UniquePtr.
    pub enum Opaque {}

    /// A type with trivial move constructor and no destructor, which can
    /// therefore be owned and moved around in Rust code without requiring
    /// indirection.
    pub enum Trivial {}
}

#[doc(hidden)]
pub fn verify_extern_type<T: ExternType<Id = Id>, Id>() {}

#[doc(hidden)]
pub fn verify_extern_kind<T: ExternType<Kind = Kind>, Kind>() {}
<commit_msg>Add trait bound to ExternType::Kind<commit_after>use self::kind::Kind;

/// A type for which the layout is determined by its C++ definition.
///
/// This trait serves the following two related purposes.
///
///

///
/// ## Safely unifying occurrences of the same extern type
///
/// `ExternType` makes it possible for CXX to safely share a consistent Rust
/// type across multiple #\[cxx::bridge\] invocations that refer to a common
/// extern C++ type.
///
/// In the following snippet, two #\[cxx::bridge\] invocations in different
/// files (possibly different crates) both contain function signatures involving
/// the same C++ type `example::Demo`. If both were written just containing
/// `type Demo;`, then both macro expansions would produce their own separate
/// Rust type called `Demo` and thus the compiler wouldn't allow us to take the
/// `Demo` returned by `file1::ffi::create_demo` and pass it as the `Demo`
/// argument accepted by `file2::ffi::take_ref_demo`. Instead, one of the two
/// `Demo`s has been defined as an extern type alias of the other, making them
/// the same type in Rust. The CXX code generator will use an automatically
/// generated `ExternType` impl emitted in file1 to statically verify that in
/// file2 `crate::file1::ffi::Demo` really does refer to the C++ type
/// `example::Demo` as expected in file2.
///
/// ```no_run
/// // file1.rs
/// # mod file1 {
/// #[cxx::bridge(namespace = example)]
/// pub mod ffi {
///     extern "C" {
///         type Demo;
///
///         fn create_demo() -> UniquePtr<Demo>;
///     }
/// }
/// # }
///
/// // file2.rs
/// #[cxx::bridge(namespace = example)]
/// pub mod ffi {
///     extern "C" {
///         type Demo = crate::file1::ffi::Demo;
///
///         fn take_ref_demo(demo: &Demo);
///     }
/// }
/// #
/// # fn main() {}
/// ```
///
///

///
/// ## Integrating with bindgen-generated types
///
/// Handwritten `ExternType` impls make it possible to plug in a data structure
/// emitted by bindgen as the definition of a C++ type emitted by CXX.
///
/// By writing the unsafe `ExternType` impl, the programmer asserts that the C++
/// namespace and type name given in the type id refers to a C++ type that is
/// equivalent to Rust type that is the `Self` type of the impl.
///
/// ```no_run
/// # const _: &str = stringify! {
/// mod folly_sys;  // the bindgen-generated bindings
/// # };
/// # mod folly_sys {
/// #     #[repr(transparent)]
/// #     pub struct StringPiece([usize; 2]);
/// # }
///
/// use cxx::{type_id, ExternType};
///
/// unsafe impl ExternType for folly_sys::StringPiece {
///     type Id = type_id!("folly::StringPiece");
///     type Kind = cxx::kind::Opaque;
/// }
///
/// #[cxx::bridge(namespace = folly)]
/// pub mod ffi {
///     extern "C" {
///         include!("rust_cxx_bindings.h");
///
///         type StringPiece = crate::folly_sys::StringPiece;
///
///         fn print_string_piece(s: &StringPiece);
///     }
/// }
///
/// // Now if we construct a StringPiece or obtain one through one
/// // of the bindgen-generated signatures, we are able to pass it
/// // along to ffi::print_string_piece.
/// #
/// # fn main() {}
/// ```
pub unsafe trait ExternType {
    /// A type-level representation of the type's C++ namespace and type name.
    ///
    /// This will always be defined using `type_id!` in the following form:
    ///
    /// ```
    /// # struct TypeName;
    /// # unsafe impl cxx::ExternType for TypeName {
    /// type Id = cxx::type_id!("name::space::of::TypeName");
    /// #     type Kind = cxx::kind::Opaque;
    /// # }
    /// ```
    type Id;

    /// Either [`cxx::kind::Opaque`] or [`cxx::kind::Trivial`].
    ///
    /// [`cxx::kind::Opaque`]: kind::Opaque
    /// [`cxx::kind::Trivial`]: kind::Trivial
    ///
    /// A C++ type is only okay to hold and pass around by value in Rust if its
    /// [move constructor is trivial] and it has no destructor. In CXX, these
    /// are called Trivial extern C++ types, while types with nontrivial move
    /// behavior or a destructor must be considered Opaque and handled by Rust
    /// only behind an indirection, such as a reference or UniquePtr.
    ///
    /// [move constructor is trivial]: https://en.cppreference.com/w/cpp/types/is_move_constructible
    ///
    /// If you believe your C++ type reflected by this ExternType impl is indeed
    /// trivial, you can specify:
    ///
    /// ```
    /// # struct TypeName;
    /// # unsafe impl cxx::ExternType for TypeName {
    /// #     type Id = cxx::type_id!("name::space::of::TypeName");
    /// type Kind = cxx::kind::Trivial;
    /// # }
    /// ```
    ///
    /// which will enable you to pass it into C++ functions by value, return it
    /// by value, and include it in `struct`s that you have declared to
    /// `cxx::bridge`. Your claim about the triviality of the C++ type will be
    /// checked by a `static_assert` in the generated C++ side of the binding.
    type Kind: Kind;
}

/// Marker types identifying Rust's knowledge about an extern C++ type.
///
/// These markers are used in the [`Kind`][ExternType::Kind] associated type in
/// impls of the `ExternType` trait. Refer to the documentation of `Kind` for an
/// overview of their purpose.
pub mod kind {
    /// An opaque type which cannot be passed or held by value within Rust.
    ///
    /// Rust's move semantics are such that every move is equivalent to a
    /// memcpy. This is incompatible in general with C++'s constructor-based
    /// move semantics, so a C++ type which has a destructor or nontrivial move
    /// constructor must never exist by value in Rust. In CXX, such types are
    /// called opaque C++ types.
    ///
    /// When passed across an FFI boundary, an opaque C++ type must be behind an
    /// indirection such as a reference or UniquePtr.
    pub enum Opaque {}

    /// A type with trivial move constructor and no destructor, which can
    /// therefore be owned and moved around in Rust code without requiring
    /// indirection.
    pub enum Trivial {}

    pub trait Kind {}
    impl Kind for Opaque {}
    impl Kind for Trivial {}
}

#[doc(hidden)]
pub fn verify_extern_type<T: ExternType<Id = Id>, Id>() {}

#[doc(hidden)]
pub fn verify_extern_kind<T: ExternType<Kind = Kind>, Kind>() {}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 J.C. Moyer
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// corresponds to lauxlib.h
use libc::{c_int, c_long, c_char, c_void, size_t};
use ffi::lua;
use ffi::lua::{lua_State, lua_CFunction, lua_Integer, lua_Number};
use ffi::luaconf::LUAL_BUFFERSIZE;
use std::ptr;

pub use super::glue::LUAL_NUMSIZES;
pub use super::glue::LUA_FILEHANDLE;

// extra error code for 'luaL_load'
pub const LUA_ERRFILE: c_int = lua::LUA_ERRERR + 1;

#[repr(C)]
pub struct luaL_Reg {
  pub name: *const c_char,
  pub func: lua_CFunction,
}

#[inline(always)]
pub unsafe fn luaL_checkversion(L: *mut lua_State) {
  luaL_checkversion_(L, lua::LUA_VERSION_NUM as lua_Number, LUAL_NUMSIZES as size_t)
}

extern {
  pub fn luaL_checkversion_(L: *mut lua_State, ver: lua_Number, sz: size_t);

  pub fn luaL_getmetafield(L: *mut lua_State, obj: c_int, e: *const c_char) -> c_int;
  pub fn luaL_callmeta(L: *mut lua_State, obj: c_int, e: *const c_char) -> c_int;
  pub fn luaL_tolstring(L: *mut lua_State, idx: c_int, len: *mut size_t) -> *const c_char;
  pub fn luaL_argerror(L: *mut lua_State, arg: c_int, l: *const c_char) -> c_int;
  pub fn luaL_checklstring(L: *mut lua_State, arg: c_int, l: *mut size_t) -> *const c_char;
  pub fn luaL_optlstring(L: *mut lua_State, arg: c_int, def: *const c_char, l: *mut size_t) -> *const c_char;
  pub fn luaL_checknumber(L: *mut lua_State, arg: c_int) -> lua_Number;
  pub fn luaL_optnumber(L: *mut lua_State, arg: c_int, def: lua_Number) -> lua_Number;
  pub fn luaL_checkinteger(L: *mut lua_State, arg: c_int) -> lua_Integer;
  pub fn luaL_optinteger(L: *mut lua_State, arg: c_int, def: lua_Integer) -> lua_Integer;

  pub fn luaL_checkstack(L: *mut lua_State, sz: c_int, msg: *const c_char);
  pub fn luaL_checktype(L: *mut lua_State, arg: c_int, t: c_int);
  pub fn luaL_checkany(L: *mut lua_State, arg: c_int);

  pub fn luaL_newmetatable(L: *mut lua_State, tname: *const c_char) -> c_int;
  pub fn luaL_setmetatable(L: *mut lua_State, tname: *const c_char);
  pub fn luaL_testudata(L: *mut lua_State, ud: c_int, tname: *const c_char) -> *mut c_void;
  pub fn luaL_checkudata(L: *mut lua_State, ud: c_int, tname: *const c_char) -> *mut c_void;

  pub fn luaL_where(L: *mut lua_State, lvl: c_int);
  pub fn luaL_error(L: *mut lua_State, fmt: *const c_char, ...) -> c_int;

  // TODO: test this
  pub fn luaL_checkoption(L: *mut lua_State, arg: c_int, def: *const c_char, lst: *const *const c_char) -> c_int;

  pub fn luaL_fileresult(L: *mut lua_State, stat: c_int, fname: *const c_char) -> c_int;
  pub fn luaL_execresult(L: *mut lua_State, stat: c_int) -> c_int;
}

// pre-defined references
pub const LUA_NOREF: c_int = -2;
pub const LUA_REFNIL: c_int = -1;

extern {
  pub fn luaL_ref(L: *mut lua_State, t: c_int) -> c_int;
  pub fn luaL_unref(L: *mut lua_State, t: c_int, r: c_int);

  pub fn luaL_loadfilex(L: *mut lua_State, filename: *const c_char, mode: *const c_char) -> c_int;
}

#[inline(always)]
pub unsafe fn luaL_loadfile(L: *mut lua_State, f: *const c_char) -> c_int {
  luaL_loadfilex(L, f, ptr::null())
}

extern {
  pub fn luaL_loadbufferx(L: *mut lua_State, buff: *const c_char, sz: size_t, name: *const c_char, mode: *const c_char) -> c_int;
  pub fn luaL_loadstring(L: *mut lua_State, s: *const c_char) -> c_int;

  pub fn luaL_newstate() -> *mut lua_State;

  pub fn luaL_len(L: *mut lua_State, idx: c_int) -> lua_Integer;

  pub fn luaL_gsub(L: *mut lua_State, s: *const c_char, p: *const c_char, r: *const c_char) -> *const c_char;

  pub fn luaL_setfuncs(L: *mut lua_State, l: *const luaL_Reg, nup: c_int);

  pub fn luaL_getsubtable(L: *mut lua_State, idx: c_int, fname: *const c_char) -> c_int;

  pub fn luaL_traceback(L: *mut lua_State, L1: *mut lua_State, msg: *const c_char, level: c_int);

  pub fn luaL_requiref(L: *mut lua_State, modname: *const c_char, openf: lua_CFunction, glb: c_int);
}

#[inline(always)]
#[allow(unused_variables)]
pub unsafe fn luaL_newlibtable(L: *mut lua_State, l: *const luaL_Reg) {
  // TODO: figure out how to pass an appropriate hint for the second param
  // this involves correcting the second parameter's type; in C this is
  // sizeof(l)/sizeof(l[0])
  lua::lua_createtable(L, 0, 0)
}

#[inline(always)]
pub unsafe fn luaL_newlib(L: *mut lua_State, l: *const luaL_Reg) {
  luaL_checkversion(L);
  luaL_newlibtable(L, l);
  luaL_setfuncs(L, l, 0)
}

#[inline(always)]
pub unsafe fn luaL_argcheck(L: *mut lua_State, cond: c_int, arg: c_int, extramsg: *const c_char) {
  if cond == 0 {
    luaL_argerror(L, arg, extramsg);
  }
}

#[inline(always)]
pub unsafe fn luaL_checkstring(L: *mut lua_State, n: c_int) -> *const c_char {
  luaL_checklstring(L, n, ptr::null_mut())
}

#[inline(always)]
pub unsafe fn luaL_optstring(L: *mut lua_State, n: c_int, d: *const c_char) -> *const c_char {
  luaL_optlstring(L, n, d, ptr::null_mut())
}

// From 5.3 user manual:
// Macros to project non-default integer types (luaL_checkint, luaL_optint,
// luaL_checklong, luaL_optlong) were deprecated. Use their equivalent over
// lua_Integer with a type cast (or, when possible, use lua_Integer in your
// code).
#[inline(always)]
#[deprecated]
pub unsafe fn luaL_checkint(L: *mut lua_State, n: c_int) -> c_int {
  luaL_checkinteger(L, n) as c_int
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_optint(L: *mut lua_State, n: c_int, d: c_int) -> c_int {
  luaL_optinteger(L, n, d as lua_Integer) as c_int
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_checklong(L: *mut lua_State, n: c_int) -> c_long {
  luaL_checkinteger(L, n) as c_int
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_optlong(L: *mut lua_State, n: c_int, d: c_long) -> c_long {
  luaL_optinteger(L, n, d as lua_Integer) as c_int
}

#[inline(always)]
pub unsafe fn luaL_typename(L: *mut lua_State, i: c_int) -> *const c_char {
  lua::lua_typename(L, lua::lua_type(L, i))
}

#[inline(always)]
pub unsafe fn luaL_dofile(L: *mut lua_State, filename: *const c_char) -> c_int {
  let status = luaL_loadfile(L, filename);
  if status == 0 {
    lua::lua_pcall(L, 0, lua::LUA_MULTRET, 0)
  } else {
    status
  }
}

#[inline(always)]
pub unsafe fn luaL_dostring(L: *mut lua_State, s: *const c_char) -> c_int {
  let status = luaL_loadstring(L, s);
  if status == 0 {
    lua::lua_pcall(L, 0, lua::LUA_MULTRET, 0)
  } else {
    status
  }
}

#[inline(always)]
pub unsafe fn luaL_getmetatable(L: *mut lua_State, n: *const c_char) {
  lua::lua_getfield(L, lua::LUA_REGISTRYINDEX, n);
}

// luaL_opt would be implemented here but it is undocumented, so it's omitted

#[inline(always)]
pub unsafe fn luaL_loadbuffer(L: *mut lua_State, s: *const c_char, sz: size_t, n: *const c_char) -> c_int {
  luaL_loadbufferx(L, s, sz, n, ptr::null())
}

#[repr(C)]
pub struct luaL_Buffer {
  pub b: *mut c_char,
  pub size: size_t,
  pub n: size_t,
  pub L: *mut lua_State,
  pub initb: [c_char, ..LUAL_BUFFERSIZE as uint]
}

// TODO: Test this thoroughly
#[inline(always)]
pub unsafe fn luaL_addchar(B: *mut luaL_Buffer, c: c_char) {
  // (B)->n < (B) -> size || luaL_prepbuffsize((B), 1)
  if (*B).n < (*B).size {
    luaL_prepbuffsize(B, 1);
  }
  // (B)->b[(B)->n++] = (c)
  let offset = (*B).b.offset((*B).n as int);
  ptr::write(offset, c);
  (*B).n += 1;
}

#[inline(always)]
pub unsafe fn luaL_addsize(B: *mut luaL_Buffer, s: size_t) {
  (*B).n += s;
}

extern {
  pub fn luaL_buffinit(L: *mut lua_State, B: *mut luaL_Buffer);
  pub fn luaL_prepbuffsize(B: *mut luaL_Buffer, sz: size_t) -> *mut c_char;
  pub fn luaL_addlstring(B: *mut luaL_Buffer, s: *const c_char, l: size_t);
  pub fn luaL_addstring(B: *mut luaL_Buffer, s: *const c_char);
  pub fn luaL_addvalue(B: *mut luaL_Buffer);
  pub fn luaL_pushresult(B: *mut luaL_Buffer);
  pub fn luaL_pushresultsize(B: *mut luaL_Buffer, sz: size_t);
  pub fn luaL_buffinitsize(L: *mut lua_State, B: *mut luaL_Buffer, sz: size_t) -> *mut c_char;
}

pub unsafe fn luaL_prepbuffer(B: *mut luaL_Buffer) -> *mut c_char {
  luaL_prepbuffsize(B, LUAL_BUFFERSIZE as size_t)
}

#[repr(C)]
pub struct luaL_Stream {
  pub f: *mut ::libc::FILE,
  pub closef: lua_CFunction
}

// omitted: old module system compatibility
<commit_msg>Cast to proper integer type<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 J.C. Moyer
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// corresponds to lauxlib.h
use libc::{c_int, c_long, c_char, c_void, size_t};
use ffi::lua;
use ffi::lua::{lua_State, lua_CFunction, lua_Integer, lua_Number};
use ffi::luaconf::LUAL_BUFFERSIZE;
use std::ptr;

pub use super::glue::LUAL_NUMSIZES;
pub use super::glue::LUA_FILEHANDLE;

// extra error code for 'luaL_load'
pub const LUA_ERRFILE: c_int = lua::LUA_ERRERR + 1;

#[repr(C)]
pub struct luaL_Reg {
  pub name: *const c_char,
  pub func: lua_CFunction,
}

#[inline(always)]
pub unsafe fn luaL_checkversion(L: *mut lua_State) {
  luaL_checkversion_(L, lua::LUA_VERSION_NUM as lua_Number, LUAL_NUMSIZES as size_t)
}

extern {
  pub fn luaL_checkversion_(L: *mut lua_State, ver: lua_Number, sz: size_t);

  pub fn luaL_getmetafield(L: *mut lua_State, obj: c_int, e: *const c_char) -> c_int;
  pub fn luaL_callmeta(L: *mut lua_State, obj: c_int, e: *const c_char) -> c_int;
  pub fn luaL_tolstring(L: *mut lua_State, idx: c_int, len: *mut size_t) -> *const c_char;
  pub fn luaL_argerror(L: *mut lua_State, arg: c_int, l: *const c_char) -> c_int;
  pub fn luaL_checklstring(L: *mut lua_State, arg: c_int, l: *mut size_t) -> *const c_char;
  pub fn luaL_optlstring(L: *mut lua_State, arg: c_int, def: *const c_char, l: *mut size_t) -> *const c_char;
  pub fn luaL_checknumber(L: *mut lua_State, arg: c_int) -> lua_Number;
  pub fn luaL_optnumber(L: *mut lua_State, arg: c_int, def: lua_Number) -> lua_Number;
  pub fn luaL_checkinteger(L: *mut lua_State, arg: c_int) -> lua_Integer;
  pub fn luaL_optinteger(L: *mut lua_State, arg: c_int, def: lua_Integer) -> lua_Integer;

  pub fn luaL_checkstack(L: *mut lua_State, sz: c_int, msg: *const c_char);
  pub fn luaL_checktype(L: *mut lua_State, arg: c_int, t: c_int);
  pub fn luaL_checkany(L: *mut lua_State, arg: c_int);

  pub fn luaL_newmetatable(L: *mut lua_State, tname: *const c_char) -> c_int;
  pub fn luaL_setmetatable(L: *mut lua_State, tname: *const c_char);
  pub fn luaL_testudata(L: *mut lua_State, ud: c_int, tname: *const c_char) -> *mut c_void;
  pub fn luaL_checkudata(L: *mut lua_State, ud: c_int, tname: *const c_char) -> *mut c_void;

  pub fn luaL_where(L: *mut lua_State, lvl: c_int);
  pub fn luaL_error(L: *mut lua_State, fmt: *const c_char, ...) -> c_int;

  // TODO: test this
  pub fn luaL_checkoption(L: *mut lua_State, arg: c_int, def: *const c_char, lst: *const *const c_char) -> c_int;

  pub fn luaL_fileresult(L: *mut lua_State, stat: c_int, fname: *const c_char) -> c_int;
  pub fn luaL_execresult(L: *mut lua_State, stat: c_int) -> c_int;
}

// pre-defined references
pub const LUA_NOREF: c_int = -2;
pub const LUA_REFNIL: c_int = -1;

extern {
  pub fn luaL_ref(L: *mut lua_State, t: c_int) -> c_int;
  pub fn luaL_unref(L: *mut lua_State, t: c_int, r: c_int);

  pub fn luaL_loadfilex(L: *mut lua_State, filename: *const c_char, mode: *const c_char) -> c_int;
}

#[inline(always)]
pub unsafe fn luaL_loadfile(L: *mut lua_State, f: *const c_char) -> c_int {
  luaL_loadfilex(L, f, ptr::null())
}

extern {
  pub fn luaL_loadbufferx(L: *mut lua_State, buff: *const c_char, sz: size_t, name: *const c_char, mode: *const c_char) -> c_int;
  pub fn luaL_loadstring(L: *mut lua_State, s: *const c_char) -> c_int;

  pub fn luaL_newstate() -> *mut lua_State;

  pub fn luaL_len(L: *mut lua_State, idx: c_int) -> lua_Integer;

  pub fn luaL_gsub(L: *mut lua_State, s: *const c_char, p: *const c_char, r: *const c_char) -> *const c_char;

  pub fn luaL_setfuncs(L: *mut lua_State, l: *const luaL_Reg, nup: c_int);

  pub fn luaL_getsubtable(L: *mut lua_State, idx: c_int, fname: *const c_char) -> c_int;

  pub fn luaL_traceback(L: *mut lua_State, L1: *mut lua_State, msg: *const c_char, level: c_int);

  pub fn luaL_requiref(L: *mut lua_State, modname: *const c_char, openf: lua_CFunction, glb: c_int);
}

#[inline(always)]
#[allow(unused_variables)]
pub unsafe fn luaL_newlibtable(L: *mut lua_State, l: *const luaL_Reg) {
  // TODO: figure out how to pass an appropriate hint for the second param
  // this involves correcting the second parameter's type; in C this is
  // sizeof(l)/sizeof(l[0])
  lua::lua_createtable(L, 0, 0)
}

#[inline(always)]
pub unsafe fn luaL_newlib(L: *mut lua_State, l: *const luaL_Reg) {
  luaL_checkversion(L);
  luaL_newlibtable(L, l);
  luaL_setfuncs(L, l, 0)
}

#[inline(always)]
pub unsafe fn luaL_argcheck(L: *mut lua_State, cond: c_int, arg: c_int, extramsg: *const c_char) {
  if cond == 0 {
    luaL_argerror(L, arg, extramsg);
  }
}

#[inline(always)]
pub unsafe fn luaL_checkstring(L: *mut lua_State, n: c_int) -> *const c_char {
  luaL_checklstring(L, n, ptr::null_mut())
}

#[inline(always)]
pub unsafe fn luaL_optstring(L: *mut lua_State, n: c_int, d: *const c_char) -> *const c_char {
  luaL_optlstring(L, n, d, ptr::null_mut())
}

// From 5.3 user manual:
// Macros to project non-default integer types (luaL_checkint, luaL_optint,
// luaL_checklong, luaL_optlong) were deprecated. Use their equivalent over
// lua_Integer with a type cast (or, when possible, use lua_Integer in your
// code).
#[inline(always)]
#[deprecated]
pub unsafe fn luaL_checkint(L: *mut lua_State, n: c_int) -> c_int {
  luaL_checkinteger(L, n) as c_int
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_optint(L: *mut lua_State, n: c_int, d: c_int) -> c_int {
  luaL_optinteger(L, n, d as lua_Integer) as c_int
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_checklong(L: *mut lua_State, n: c_int) -> c_long {
  luaL_checkinteger(L, n) as c_long
}

#[inline(always)]
#[deprecated]
pub unsafe fn luaL_optlong(L: *mut lua_State, n: c_int, d: c_long) -> c_long {
  luaL_optinteger(L, n, d as lua_Integer) as c_long
}

#[inline(always)]
pub unsafe fn luaL_typename(L: *mut lua_State, i: c_int) -> *const c_char {
  lua::lua_typename(L, lua::lua_type(L, i))
}

#[inline(always)]
pub unsafe fn luaL_dofile(L: *mut lua_State, filename: *const c_char) -> c_int {
  let status = luaL_loadfile(L, filename);
  if status == 0 {
    lua::lua_pcall(L, 0, lua::LUA_MULTRET, 0)
  } else {
    status
  }
}

#[inline(always)]
pub unsafe fn luaL_dostring(L: *mut lua_State, s: *const c_char) -> c_int {
  let status = luaL_loadstring(L, s);
  if status == 0 {
    lua::lua_pcall(L, 0, lua::LUA_MULTRET, 0)
  } else {
    status
  }
}

#[inline(always)]
pub unsafe fn luaL_getmetatable(L: *mut lua_State, n: *const c_char) {
  lua::lua_getfield(L, lua::LUA_REGISTRYINDEX, n);
}

// luaL_opt would be implemented here but it is undocumented, so it's omitted

#[inline(always)]
pub unsafe fn luaL_loadbuffer(L: *mut lua_State, s: *const c_char, sz: size_t, n: *const c_char) -> c_int {
  luaL_loadbufferx(L, s, sz, n, ptr::null())
}

#[repr(C)]
pub struct luaL_Buffer {
  pub b: *mut c_char,
  pub size: size_t,
  pub n: size_t,
  pub L: *mut lua_State,
  pub initb: [c_char, ..LUAL_BUFFERSIZE as uint]
}

// TODO: Test this thoroughly
#[inline(always)]
pub unsafe fn luaL_addchar(B: *mut luaL_Buffer, c: c_char) {
  // (B)->n < (B) -> size || luaL_prepbuffsize((B), 1)
  if (*B).n < (*B).size {
    luaL_prepbuffsize(B, 1);
  }
  // (B)->b[(B)->n++] = (c)
  let offset = (*B).b.offset((*B).n as int);
  ptr::write(offset, c);
  (*B).n += 1;
}

#[inline(always)]
pub unsafe fn luaL_addsize(B: *mut luaL_Buffer, s: size_t) {
  (*B).n += s;
}

extern {
  pub fn luaL_buffinit(L: *mut lua_State, B: *mut luaL_Buffer);
  pub fn luaL_prepbuffsize(B: *mut luaL_Buffer, sz: size_t) -> *mut c_char;
  pub fn luaL_addlstring(B: *mut luaL_Buffer, s: *const c_char, l: size_t);
  pub fn luaL_addstring(B: *mut luaL_Buffer, s: *const c_char);
  pub fn luaL_addvalue(B: *mut luaL_Buffer);
  pub fn luaL_pushresult(B: *mut luaL_Buffer);
  pub fn luaL_pushresultsize(B: *mut luaL_Buffer, sz: size_t);
  pub fn luaL_buffinitsize(L: *mut lua_State, B: *mut luaL_Buffer, sz: size_t) -> *mut c_char;
}

pub unsafe fn luaL_prepbuffer(B: *mut luaL_Buffer) -> *mut c_char {
  luaL_prepbuffsize(B, LUAL_BUFFERSIZE as size_t)
}

#[repr(C)]
pub struct luaL_Stream {
  pub f: *mut ::libc::FILE,
  pub closef: lua_CFunction
}

// omitted: old module system compatibility
<|endoftext|>