<commit_before>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<ssl::Error> for Error {
    fn from(err: ssl::Error) -> Error {
        Error(err)
    }
}

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(X509::from_der(buf));
        Ok(Certificate(cert))
    }
}

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    Interrupted(MidHandshakeTlsStream<S>),
}

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<ErrorStack> for HandshakeError<S> {
    fn from(e: ErrorStack) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        // FIXME clear chain certs to clean up if called multiple times
        try!(self.0.set_certificate(&pkcs12.0.cert));
        try!(self.0.set_private_key(&pkcs12.0.pkey));
        try!(self.0.check_private_key());
        for cert in pkcs12.0.chain {
            try!(self.0.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.cert_store_mut().add_cert(cert.0));
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, self.0);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(TlsConnector(self.0.build()))
    }
}

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, self.0.builder_mut());
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(TlsAcceptor(self.0.build()))
    }
}

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsStream<S>(ssl::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

/// OpenSSL-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the OpenSSL `SslStream`.
    fn raw_stream(&self) -> &ssl::SslStream<S>;

    /// Returns a mutable reference to the OpenSSL `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut ssl::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &ssl::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut ssl::SslStream<S> {
        &mut (self.0).0
    }
}

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

impl ErrorExt for ::Error {
    fn openssl_error(&self) -> &ssl::Error {
        &(self.0).0
    }
}
<commit_msg>Fix build<commit_after>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(&self.0, fmt)
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<ssl::Error> for Error {
    fn from(err: ssl::Error) -> Error {
        Error(err)
    }
}

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(X509::from_der(buf));
        Ok(Certificate(cert))
    }
}

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub enum HandshakeError<S> {
    Failure(Error),
    Interrupted(MidHandshakeTlsStream<S>),
}

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

impl<S> From<ErrorStack> for HandshakeError<S> {
    fn from(e: ErrorStack) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        // FIXME clear chain certs to clean up if called multiple times
        try!(self.0.set_certificate(&pkcs12.0.cert));
        try!(self.0.set_private_key(&pkcs12.0.pkey));
        try!(self.0.check_private_key());
        for cert in pkcs12.0.chain {
            try!(self.0.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.cert_store_mut().add_cert(cert.0));
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, &mut self.0);
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(TlsConnector(self.0.build()))
    }
}

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, self.0.builder_mut());
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(TlsAcceptor(self.0.build()))
    }
}

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsStream<S>(ssl::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

/// OpenSSL-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the OpenSSL `SslStream`.
    fn raw_stream(&self) -> &ssl::SslStream<S>;

    /// Returns a mutable reference to the OpenSSL `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut ssl::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &ssl::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut ssl::SslStream<S> {
        &mut (self.0).0
    }
}

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

impl ErrorExt for ::Error {
    fn openssl_error(&self) -> &ssl::Error {
        &(self.0).0
    }
}
<|endoftext|>


<commit_before>use pancurses::{Input, Window};
use errext::{ErrExt, PResult};
use command::run_expression;

fn print_buffer(screen: &Window, buffer: &str) -> PResult {
    let height = (screen.get_max_y() - 1) as usize;
    let width = screen.get_max_x() as usize;
    screen.mv(0, 0).as_err()?;
    for line in buffer.lines().take(height) {
        let max = ::std::cmp::min(width, line.len());
        screen.addstr(&line[..max]).as_err()?;
        screen.addch('\n').as_err()?;
    }
    Ok(())
}

fn print_command(screen: &Window, buffer: &str) -> PResult {
    let height = screen.get_max_y() - 1;
    let width = screen.get_max_x() - 1;
    screen.mv(height - 1, 0).as_err()?;
    screen.hline(0, width).as_err()?;
    screen.mv(height, 0).as_err()?;
    let max = ::std::cmp::min(width as usize, buffer.len());
    screen.addstr(&buffer[..max]).as_err()?;
    Ok(())
}

pub fn input_loop(screen: Window) -> PResult {
    let mut output_buffer = String::from("");
    let mut command_buffer = String::new();
    // let (mut x, mut y) = (0, 0);
    loop {
        screen.clear().as_err()?;
        print_buffer(&screen, &output_buffer)?;
        print_command(&screen, &command_buffer)?;
        screen.refresh().as_err()?;
        let next = screen.getch().unwrap();
        match next {
            Input::KeyBackspace |
            Input::Character('\x08') |
            Input::Character('\x7f') => {
                command_buffer.pop().unwrap();
            }
            Input::Character('\n') => {
                output_buffer.push_str("$ ");
                let output = run_expression(&command_buffer)?;
                output_buffer.push_str(&output.command);
                output_buffer.push_str("\n");
                output_buffer.push_str(&output.interleaved);
                output_buffer.push_str("\n");
                command_buffer.clear();
            }
            Input::Character(chr) => {
                println!("char: {:?}", chr);
                command_buffer.push(chr);
            }
            _ => {}
        };
    }
}
<commit_msg>Remove unnecessary unwrap<commit_after>use pancurses::{Input, Window};
use errext::{ErrExt, PResult};
use command::run_expression;

fn print_buffer(screen: &Window, buffer: &str) -> PResult {
    let height = (screen.get_max_y() - 1) as usize;
    let width = screen.get_max_x() as usize;
    screen.mv(0, 0).as_err()?;
    for line in buffer.lines().take(height) {
        let max = ::std::cmp::min(width, line.len());
        screen.addstr(&line[..max]).as_err()?;
        screen.addch('\n').as_err()?;
    }
    Ok(())
}

fn print_command(screen: &Window, buffer: &str) -> PResult {
    let height = screen.get_max_y() - 1;
    let width = screen.get_max_x() - 1;
    screen.mv(height - 1, 0).as_err()?;
    screen.hline(0, width).as_err()?;
    screen.mv(height, 0).as_err()?;
    let max = ::std::cmp::min(width as usize, buffer.len());
    screen.addstr(&buffer[..max]).as_err()?;
    Ok(())
}

pub fn input_loop(screen: Window) -> PResult {
    let mut output_buffer = String::from("");
    let mut command_buffer = String::new();
    // let (mut x, mut y) = (0, 0);
    loop {
        screen.clear().as_err()?;
        print_buffer(&screen, &output_buffer)?;
        print_command(&screen, &command_buffer)?;
        screen.refresh().as_err()?;
        let next = screen.getch().unwrap();
        match next {
            Input::KeyBackspace |
            Input::Character('\x08') |
            Input::Character('\x7f') => {
                command_buffer.pop();
            }
            Input::Character('\n') => {
                output_buffer.push_str("$ ");
                let output = run_expression(&command_buffer)?;
                output_buffer.push_str(&output.command);
                output_buffer.push_str("\n");
                output_buffer.push_str(&output.interleaved);
                output_buffer.push_str("\n");
                command_buffer.clear();
            }
            Input::Character(chr) => {
                println!("char: {:?}", chr);
                command_buffer.push(chr);
            }
            _ => {}
        };
    }
}
<|endoftext|>


<commit_before>use actix::actors::signal;
use actix::prelude::*;
use bytes::BytesMut;
use godcoin_p2p::{cmd, session::*, Network, Payload};
use log::{error, info};
use std::{
    cell::RefCell,
    collections::HashSet,
    rc::Rc,
    time::{Duration, Instant},
};
use tokio::{prelude::*, timer::Delay};

struct Signals;

impl Actor for Signals {
    type Context = Context<Self>;
}

impl Handler<signal::Signal> for Signals {
    type Result = ();

    fn handle(&mut self, msg: signal::Signal, _: &mut Context<Self>) {
        if let signal::SignalType::Int = msg.0 {
            info!("SIGINT received, exiting");
            System::current().stop();
        }
    }
}

struct NetState {
    net_id: usize,
    messages: HashSet<BytesMut>,
    msg_counter: Rc<RefCell<usize>>,
    msg_threshold: usize,
}

impl NetState {
    pub fn new(net_id: usize, msg_counter: Rc<RefCell<usize>>, msg_threshold: usize) -> Self {
        NetState {
            net_id,
            messages: HashSet::new(),
            msg_counter,
            msg_threshold,
        }
    }
}

fn connected(state: &mut NetState, ses: SessionInfo) {
    match ses.conn_type {
        ConnectionType::Inbound => {
            info!(
                "[net:{}] Accepted connection -> {}",
                state.net_id, ses.peer_addr
            );
        }
        ConnectionType::Outbound => {
            info!(
                "[net:{}] Connected to node -> {}",
                state.net_id, ses.peer_addr
            );
        }
    }
}

fn disconnected(state: &mut NetState, ses: SessionInfo) {
    info!(
        "[net:{}] Connection disconnected -> {}",
        state.net_id, ses.peer_addr
    );
}

fn message(state: &mut NetState, id: SessionId, payload: &Payload) -> bool {
    info!(
        "[net:{}] Received message from {} with: {:?}",
        state.net_id, id, payload
    );
    let broadcast = state.messages.contains(&payload.id);
    *state.msg_counter.borrow_mut() += 1;
    if *state.msg_counter.borrow() == state.msg_threshold {
        info!("Threshold reached => evicting cached messages");
        state.messages.clear();
    }
    broadcast
}

fn main() {
    let env = env_logger::Env::new()
        .filter_or(env_logger::DEFAULT_FILTER_ENV, "godcoin_simulation=debug");
    env_logger::init_from_env(env);

    let sys = System::new("simulation");

    {
        let sig_addr = Signals.start();
        let addr = signal::ProcessSignals::from_registry();
        addr.do_send(signal::Subscribe(sig_addr.recipient()));
    }

    let nets = {
        let net_count = 3;
        let port = 7777;
        let mut nets = Vec::with_capacity(net_count);
        let msg_counter = Rc::new(RefCell::new(0));
        // The threshold is always one less to exclude the current network from being counted
        let threshold = net_count - 1;
        for net_id in 0..net_count {
            let msg_counter = Rc::clone(&msg_counter);
            let net = Network::new(NetState::new(net_id, msg_counter, threshold), message)
                .on_connect(connected)
                .on_disconnect(disconnected)
                .start();
            net.do_send(cmd::Listen(
                format!("127.0.0.1:{}", port + net_id).parse().unwrap(),
            ));
            info!("[net:{}] Accepting connections on 127.0.0.1:7777", net_id);
            nets.push(net);
        }
        nets
    };

    nets[1].do_send(cmd::Connect("127.0.0.1:7777".parse().unwrap()));
    nets[2].do_send(cmd::Connect("127.0.0.1:7777".parse().unwrap()));
    nets[2].do_send(cmd::Connect("127.0.0.1:7778".parse().unwrap()));

    let deadline = Instant::now() + Duration::from_secs(1);
    Arbiter::spawn(
        Delay::new(deadline)
            .and_then(move |_| {
                let payload = Payload {
                    id: BytesMut::from(vec![1, 2, 3]),
                    msg: BytesMut::from(vec![4, 5, 6]),
                };
                println!();
                info!("[net:2] Broadcasting message: {:?}", &payload);
                nets[2].do_send(cmd::Broadcast(payload));
                Ok(())
            })
            .map_err(|e| {
                error!("Timer failed: {:?}", e);
            }),
    );

    sys.run();
}
<commit_msg>Fix accepting connections message<commit_after>use actix::actors::signal;
use actix::prelude::*;
use bytes::BytesMut;
use godcoin_p2p::{cmd, session::*, Network, Payload};
use log::{error, info};
use std::{
    cell::RefCell,
    collections::HashSet,
    rc::Rc,
    time::{Duration, Instant},
};
use tokio::{prelude::*, timer::Delay};

struct Signals;

impl Actor for Signals {
    type Context = Context<Self>;
}

impl Handler<signal::Signal> for Signals {
    type Result = ();

    fn handle(&mut self, msg: signal::Signal, _: &mut Context<Self>) {
        if let signal::SignalType::Int = msg.0 {
            info!("SIGINT received, exiting");
            System::current().stop();
        }
    }
}

struct NetState {
    net_id: usize,
    messages: HashSet<BytesMut>,
    msg_counter: Rc<RefCell<usize>>,
    msg_threshold: usize,
}

impl NetState {
    pub fn new(net_id: usize, msg_counter: Rc<RefCell<usize>>, msg_threshold: usize) -> Self {
        NetState {
            net_id,
            messages: HashSet::new(),
            msg_counter,
            msg_threshold,
        }
    }
}

fn connected(state: &mut NetState, ses: SessionInfo) {
    match ses.conn_type {
        ConnectionType::Inbound => {
            info!(
                "[net:{}] Accepted connection -> {}",
                state.net_id, ses.peer_addr
            );
        }
        ConnectionType::Outbound => {
            info!(
                "[net:{}] Connected to node -> {}",
                state.net_id, ses.peer_addr
            );
        }
    }
}

fn disconnected(state: &mut NetState, ses: SessionInfo) {
    info!(
        "[net:{}] Connection disconnected -> {}",
        state.net_id, ses.peer_addr
    );
}

fn message(state: &mut NetState, id: SessionId, payload: &Payload) -> bool {
    info!(
        "[net:{}] Received message from {} with: {:?}",
        state.net_id, id, payload
    );
    let broadcast = state.messages.contains(&payload.id);
    *state.msg_counter.borrow_mut() += 1;
    if *state.msg_counter.borrow() == state.msg_threshold {
        info!("Threshold reached => evicting cached messages");
        state.messages.clear();
    }
    broadcast
}

fn main() {
    let env = env_logger::Env::new()
        .filter_or(env_logger::DEFAULT_FILTER_ENV, "godcoin_simulation=debug");
    env_logger::init_from_env(env);

    let sys = System::new("simulation");

    {
        let sig_addr = Signals.start();
        let addr = signal::ProcessSignals::from_registry();
        addr.do_send(signal::Subscribe(sig_addr.recipient()));
    }

    let nets = {
        let net_count = 3;
        let port = 7777;
        let mut nets = Vec::with_capacity(net_count);
        let msg_counter = Rc::new(RefCell::new(0));
        // The threshold is always one less to exclude the current network from being counted
        let threshold = net_count - 1;
        for net_id in 0..net_count {
            let msg_counter = Rc::clone(&msg_counter);
            let net = Network::new(NetState::new(net_id, msg_counter, threshold), message)
                .on_connect(connected)
                .on_disconnect(disconnected)
                .start();

            let port = port + net_id;
            net.do_send(cmd::Listen(
                format!("127.0.0.1:{}", port).parse().unwrap(),
            ));
            info!("[net:{}] Accepting connections on 127.0.0.1:{}", net_id, port);
            nets.push(net);
        }
        nets
    };

    nets[1].do_send(cmd::Connect("127.0.0.1:7777".parse().unwrap()));
    nets[2].do_send(cmd::Connect("127.0.0.1:7777".parse().unwrap()));
    nets[2].do_send(cmd::Connect("127.0.0.1:7778".parse().unwrap()));

    let deadline = Instant::now() + Duration::from_secs(1);
    Arbiter::spawn(
        Delay::new(deadline)
            .and_then(move |_| {
                let payload = Payload {
                    id: BytesMut::from(vec![1, 2, 3]),
                    msg: BytesMut::from(vec![4, 5, 6]),
                };
                println!();
                info!("[net:2] Broadcasting message: {:?}", &payload);
                nets[2].do_send(cmd::Broadcast(payload));
                Ok(())
            })
            .map_err(|e| {
                error!("Timer failed: {:?}", e);
            }),
    );

    sys.run();
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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

//! Optimizer rule to switch build and probe order of hash join
//! based on statistics of a `TableProvider`. If the number of
//! rows of both sources is known, the order can be switched
//! for a faster hash join.

use std::sync::Arc;

use crate::logical_plan::LogicalPlan;
use crate::optimizer::optimizer::OptimizerRule;
use crate::{error::Result, prelude::JoinType};

use super::utils;

/// BuildProbeOrder reorders the build and probe phase of
/// hash joins. This uses the amount of rows that a datasource has.
/// The rule optimizes the order such that the left (build) side of the join
/// is the smallest.
/// If the information is not available, the order stays the same,
/// so that it could be optimized manually in a query.
pub struct HashBuildProbeOrder {}

// Gets exact number of rows, if known by the statistics of the underlying
fn get_num_rows(logical_plan: &LogicalPlan) -> Option<usize> {
    match logical_plan {
        LogicalPlan::Projection { input, .. } => get_num_rows(input),
        LogicalPlan::Sort { input, .. } => get_num_rows(input),
        LogicalPlan::TableScan { source, .. } => source.statistics().num_rows,
        LogicalPlan::EmptyRelation {
            produce_one_row, ..
        } => {
            if *produce_one_row {
                Some(1)
            } else {
                Some(0)
            }
        }
        LogicalPlan::Limit { n: limit, input } => {
            let num_rows_input = get_num_rows(input);
            num_rows_input.map(|rows| std::cmp::min(*limit, rows))
        }
        _ => None,
    }
}

// Finds out whether to swap left vs right order based on statistics
fn should_swap_join_order(left: &LogicalPlan, right: &LogicalPlan) -> bool {
    let left_rows = get_num_rows(left);
    let right_rows = get_num_rows(right);

    match (left_rows, right_rows) {
        (Some(l), Some(r)) => l > r,
        _ => false,
    }
}

impl OptimizerRule for HashBuildProbeOrder {
    fn name(&self) -> &str {
        "hash_build_probe_order"
    }

    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        match plan {
            // Main optimization rule, swaps order of left and right
            // based on number of rows in each table
            LogicalPlan::Join {
                left,
                right,
                on,
                join_type,
                schema,
            } => {
                let left = self.optimize(left)?;
                let right = self.optimize(right)?;
                if should_swap_join_order(&left, &right) {
                    // Swap left and right, change join type and (equi-)join key order
                    Ok(LogicalPlan::Join {
                        left: Arc::new(right),
                        right: Arc::new(left),
                        on: on
                            .iter()
                            .map(|(l, r)| (r.to_string(), l.to_string()))
                            .collect(),
                        join_type: swap_join_type(*join_type),
                        schema: schema.clone(),
                    })
                } else {
                    // Keep join as is
                    Ok(LogicalPlan::Join {
                        left: Arc::new(left),
                        right: Arc::new(right),
                        on: on.clone(),
                        join_type: *join_type,
                        schema: schema.clone(),
                    })
                }
            }
            // Rest: recurse into plan, apply optimization where possible
            LogicalPlan::Projection { .. }
            | LogicalPlan::Aggregate { .. }
            | LogicalPlan::TableScan { .. }
            | LogicalPlan::Limit { .. }
            | LogicalPlan::Filter { .. }
            | LogicalPlan::Repartition { .. }
            | LogicalPlan::EmptyRelation { .. }
            | LogicalPlan::Sort { .. }
            | LogicalPlan::CreateExternalTable { .. }
            | LogicalPlan::Explain { .. }
            | LogicalPlan::Extension { .. } => {
                let expr = utils::expressions(plan);

                // apply the optimization to all inputs of the plan
                let inputs = utils::inputs(plan);
                let new_inputs = inputs
                    .iter()
                    .map(|plan| self.optimize(plan))
                    .collect::<Result<Vec<_>>>()?;

                utils::from_plan(plan, &expr, &new_inputs)
            }
        }
    }
}

impl HashBuildProbeOrder {
    #[allow(missing_docs)]
    pub fn new() -> Self {
        Self {}
    }
}

fn swap_join_type(join_type: JoinType) -> JoinType {
    match join_type {
        JoinType::Inner => JoinType::Inner,
        JoinType::Left => JoinType::Right,
        JoinType::Right => JoinType::Left,
    }
}

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

    use crate::{
        datasource::{datasource::Statistics, TableProvider},
        logical_plan::{DFSchema, Expr},
        test::*,
    };

    struct TestTableProvider {
        num_rows: usize,
    }

    impl TableProvider for TestTableProvider {
        fn as_any(&self) -> &dyn std::any::Any {
            unimplemented!()
        }
        fn schema(&self) -> arrow::datatypes::SchemaRef {
            unimplemented!()
        }

        fn scan(
            &self,
            _projection: &Option<Vec<usize>>,
            _batch_size: usize,
            _filters: &[Expr],
        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
            unimplemented!()
        }
        fn statistics(&self) -> crate::datasource::datasource::Statistics {
            Statistics {
                num_rows: Some(self.num_rows),
                total_byte_size: None,
                column_statistics: None,
            }
        }
    }

    #[test]
    fn test_num_rows() -> Result<()> {
        let table_scan = test_table_scan()?;

        assert_eq!(get_num_rows(&table_scan), Some(0));

        Ok(())
    }

    #[test]
    fn test_swap_order() -> Result<()> {
        let lp_left = LogicalPlan::TableScan {
            table_name: "left".to_string(),
            projection: None,
            source: Arc::new(TestTableProvider { num_rows: 1000 }),
            projected_schema: Arc::new(DFSchema::empty()),
            filters: vec![],
        };

        let lp_right = LogicalPlan::TableScan {
            table_name: "right".to_string(),
            projection: None,
            source: Arc::new(TestTableProvider { num_rows: 100 }),
            projected_schema: Arc::new(DFSchema::empty()),
            filters: vec![],
        };

        assert!(should_swap_join_order(&lp_left, &lp_right));
        assert!(!should_swap_join_order(&lp_right, &lp_left));

        Ok(())
    }
}
<commit_msg>ARROW-11029: [Rust] [DataFusion] Add documentation for code that determines number of rows per operator<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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

//! Optimizer rule to switch build and probe order of hash join
//! based on statistics of a `TableProvider`. If the number of
//! rows of both sources is known, the order can be switched
//! for a faster hash join.

use std::sync::Arc;

use crate::logical_plan::LogicalPlan;
use crate::optimizer::optimizer::OptimizerRule;
use crate::{error::Result, prelude::JoinType};

use super::utils;

/// BuildProbeOrder reorders the build and probe phase of
/// hash joins. This uses the amount of rows that a datasource has.
/// The rule optimizes the order such that the left (build) side of the join
/// is the smallest.
/// If the information is not available, the order stays the same,
/// so that it could be optimized manually in a query.
pub struct HashBuildProbeOrder {}

// Gets exact number of rows, if known by the statistics of the underlying
fn get_num_rows(logical_plan: &LogicalPlan) -> Option<usize> {
    match logical_plan {
        LogicalPlan::TableScan { source, .. } => source.statistics().num_rows,
        LogicalPlan::EmptyRelation {
            produce_one_row, ..
        } => {
            if *produce_one_row {
                Some(1)
            } else {
                Some(0)
            }
        }
        LogicalPlan::Limit { n: limit, input } => {
            let num_rows_input = get_num_rows(input);
            num_rows_input.map(|rows| std::cmp::min(*limit, rows))
        }
        LogicalPlan::Aggregate { .. } => {
            // we cannot yet predict how many rows will be produced by an aggregate because
            // we do not know the cardinality of the grouping keys
            None
        }
        LogicalPlan::Filter { .. } => {
            // we cannot yet predict how many rows will be produced by a filter because
            // we don't know how selective it is (how many rows it will filter out)
            None
        }
        LogicalPlan::Join { .. } => {
            // we cannot predict the cardinality of the join output
            None
        }
        LogicalPlan::Repartition { .. } => {
            // we cannot predict how rows will be repartitioned
            None
        }
        // the following operators are special cases and not querying data
        LogicalPlan::CreateExternalTable { .. } => None,
        LogicalPlan::Explain { .. } => None,
        // we do not support estimating rows with extensions yet
        LogicalPlan::Extension { .. } => None,
        // the following operators do not modify row count in any way
        LogicalPlan::Projection { input, .. } => get_num_rows(input),
        LogicalPlan::Sort { input, .. } => get_num_rows(input),
    }
}

// Finds out whether to swap left vs right order based on statistics
fn should_swap_join_order(left: &LogicalPlan, right: &LogicalPlan) -> bool {
    let left_rows = get_num_rows(left);
    let right_rows = get_num_rows(right);

    match (left_rows, right_rows) {
        (Some(l), Some(r)) => l > r,
        _ => false,
    }
}

impl OptimizerRule for HashBuildProbeOrder {
    fn name(&self) -> &str {
        "hash_build_probe_order"
    }

    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        match plan {
            // Main optimization rule, swaps order of left and right
            // based on number of rows in each table
            LogicalPlan::Join {
                left,
                right,
                on,
                join_type,
                schema,
            } => {
                let left = self.optimize(left)?;
                let right = self.optimize(right)?;
                if should_swap_join_order(&left, &right) {
                    // Swap left and right, change join type and (equi-)join key order
                    Ok(LogicalPlan::Join {
                        left: Arc::new(right),
                        right: Arc::new(left),
                        on: on
                            .iter()
                            .map(|(l, r)| (r.to_string(), l.to_string()))
                            .collect(),
                        join_type: swap_join_type(*join_type),
                        schema: schema.clone(),
                    })
                } else {
                    // Keep join as is
                    Ok(LogicalPlan::Join {
                        left: Arc::new(left),
                        right: Arc::new(right),
                        on: on.clone(),
                        join_type: *join_type,
                        schema: schema.clone(),
                    })
                }
            }
            // Rest: recurse into plan, apply optimization where possible
            LogicalPlan::Projection { .. }
            | LogicalPlan::Aggregate { .. }
            | LogicalPlan::TableScan { .. }
            | LogicalPlan::Limit { .. }
            | LogicalPlan::Filter { .. }
            | LogicalPlan::Repartition { .. }
            | LogicalPlan::EmptyRelation { .. }
            | LogicalPlan::Sort { .. }
            | LogicalPlan::CreateExternalTable { .. }
            | LogicalPlan::Explain { .. }
            | LogicalPlan::Extension { .. } => {
                let expr = utils::expressions(plan);

                // apply the optimization to all inputs of the plan
                let inputs = utils::inputs(plan);
                let new_inputs = inputs
                    .iter()
                    .map(|plan| self.optimize(plan))
                    .collect::<Result<Vec<_>>>()?;

                utils::from_plan(plan, &expr, &new_inputs)
            }
        }
    }
}

impl HashBuildProbeOrder {
    #[allow(missing_docs)]
    pub fn new() -> Self {
        Self {}
    }
}

fn swap_join_type(join_type: JoinType) -> JoinType {
    match join_type {
        JoinType::Inner => JoinType::Inner,
        JoinType::Left => JoinType::Right,
        JoinType::Right => JoinType::Left,
    }
}

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

    use crate::{
        datasource::{datasource::Statistics, TableProvider},
        logical_plan::{DFSchema, Expr},
        test::*,
    };

    struct TestTableProvider {
        num_rows: usize,
    }

    impl TableProvider for TestTableProvider {
        fn as_any(&self) -> &dyn std::any::Any {
            unimplemented!()
        }
        fn schema(&self) -> arrow::datatypes::SchemaRef {
            unimplemented!()
        }

        fn scan(
            &self,
            _projection: &Option<Vec<usize>>,
            _batch_size: usize,
            _filters: &[Expr],
        ) -> Result<std::sync::Arc<dyn crate::physical_plan::ExecutionPlan>> {
            unimplemented!()
        }
        fn statistics(&self) -> crate::datasource::datasource::Statistics {
            Statistics {
                num_rows: Some(self.num_rows),
                total_byte_size: None,
                column_statistics: None,
            }
        }
    }

    #[test]
    fn test_num_rows() -> Result<()> {
        let table_scan = test_table_scan()?;

        assert_eq!(get_num_rows(&table_scan), Some(0));

        Ok(())
    }

    #[test]
    fn test_swap_order() -> Result<()> {
        let lp_left = LogicalPlan::TableScan {
            table_name: "left".to_string(),
            projection: None,
            source: Arc::new(TestTableProvider { num_rows: 1000 }),
            projected_schema: Arc::new(DFSchema::empty()),
            filters: vec![],
        };

        let lp_right = LogicalPlan::TableScan {
            table_name: "right".to_string(),
            projection: None,
            source: Arc::new(TestTableProvider { num_rows: 100 }),
            projected_schema: Arc::new(DFSchema::empty()),
            filters: vec![],
        };

        assert!(should_swap_join_order(&lp_left, &lp_right));
        assert!(!should_swap_join_order(&lp_right, &lp_left));

        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # The Rust Core Library
//!
//! The Rust Core Library is the dependency-free[^free] foundation of [The
//! Rust Standard Library](../std/index.html). It is the portable glue
//! between the language and its libraries, defining the intrinsic and
//! primitive building blocks of all Rust code. It links to no
//! upstream libraries, no system libraries, and no libc.
//!
//! [^free]: Strictly speaking, there are some symbols which are needed but
//!          they aren't always necessary.
//!
//! The core library is *minimal*: it isn't even aware of heap allocation,
//! nor does it provide concurrency or I/O. These things require
//! platform integration, and this library is platform-agnostic.
//!
//! *It is not recommended to use the core library*. The stable
//! functionality of libcore is reexported from the
//! [standard library](../std/index.html). The composition of this library is
//! subject to change over time; only the interface exposed through libstd is
//! intended to be stable.
//!
//! # How to use the core library
//!
// FIXME: Fill me in with more detail when the interface settles
//! This library is built on the assumption of a few existing symbols:
//!
//! * `memcpy`, `memcmp`, `memset` - These are core memory routines which are
//!   often generated by LLVM. Additionally, this library can make explicit
//!   calls to these functions. Their signatures are the same as found in C.
//!   These functions are often provided by the system libc, but can also be
//!   provided by the [rlibc crate](https://crates.io/crates/rlibc).
//!
//! * `rust_begin_unwind` - This function takes three arguments, a
//!   `fmt::Arguments`, a `&str`, and a `u32`. These three arguments dictate
//!   the panic message, the file at which panic was invoked, and the line.
//!   It is up to consumers of this core library to define this panic
//!   function; it is only required to never return.

// Since libcore defines many fundamental lang items, all tests live in a
// separate crate, libcoretest, to avoid bizarre issues.

// Do not remove on snapshot creation. Needed for bootstrap. (Issue #22364)
#![cfg_attr(stage0, feature(custom_attribute))]
#![crate_name = "core"]
#![stable(feature = "core", since = "1.6.0")]
#![cfg_attr(stage0, staged_api)]
#![crate_type = "rlib"]
#![doc(html_logo_url = "https://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "https://doc.rust-lang.org/favicon.ico",
       html_root_url = "https://doc.rust-lang.org/nightly/",
       html_playground_url = "https://play.rust-lang.org/",
       issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
       test(no_crate_inject, attr(deny(warnings))),
       test(attr(allow(dead_code, deprecated, unused_variables, unused_mut))))]

#![no_core]
#![deny(missing_docs)]

#![cfg_attr(stage0, feature(rustc_attrs))]
#![cfg_attr(stage0, allow(unused_attributes))]
#![feature(allow_internal_unstable)]
#![feature(associated_type_defaults)]
#![feature(concat_idents)]
#![feature(const_fn)]
#![feature(custom_attribute)]
#![feature(fundamental)]
#![feature(intrinsics)]
#![feature(lang_items)]
#![feature(no_core)]
#![feature(on_unimplemented)]
#![feature(optin_builtin_traits)]
#![feature(reflect)]
#![feature(unwind_attributes)]
#![cfg_attr(stage0, feature(simd))]
#![cfg_attr(not(stage0), feature(repr_simd, platform_intrinsics))]
#![feature(staged_api)]
#![feature(unboxed_closures)]

#[macro_use]
mod macros;

#[path = "num/float_macros.rs"]
#[macro_use]
mod float_macros;

#[path = "num/int_macros.rs"]
#[macro_use]
mod int_macros;

#[path = "num/uint_macros.rs"]
#[macro_use]
mod uint_macros;

#[path = "num/isize.rs"]  pub mod isize;
#[path = "num/i8.rs"]   pub mod i8;
#[path = "num/i16.rs"]  pub mod i16;
#[path = "num/i32.rs"]  pub mod i32;
#[path = "num/i64.rs"]  pub mod i64;

#[path = "num/usize.rs"] pub mod usize;
#[path = "num/u8.rs"]   pub mod u8;
#[path = "num/u16.rs"]  pub mod u16;
#[path = "num/u32.rs"]  pub mod u32;
#[path = "num/u64.rs"]  pub mod u64;

#[path = "num/f32.rs"]   pub mod f32;
#[path = "num/f64.rs"]   pub mod f64;

#[macro_use]
pub mod num;

/* The libcore prelude, not as all-encompassing as the libstd prelude */

pub mod prelude;

/* Core modules for ownership management */

pub mod intrinsics;
pub mod mem;
pub mod nonzero;
pub mod ptr;

/* Core language traits */

pub mod marker;
pub mod ops;
pub mod cmp;
pub mod clone;
pub mod default;
pub mod convert;
pub mod borrow;

/* Core types and methods on primitives */

pub mod any;
pub mod array;
pub mod sync;
pub mod cell;
pub mod char;
pub mod panicking;
pub mod iter;
pub mod option;
pub mod raw;
pub mod result;

pub mod slice;
pub mod str;
pub mod hash;
pub mod fmt;

// note: does not need to be public
mod tuple;
<commit_msg>Rollup merge of #30415 - steveklabnik:remove_bad_stability_note, r=alexcrichton<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # The Rust Core Library
//!
//! The Rust Core Library is the dependency-free[^free] foundation of [The
//! Rust Standard Library](../std/index.html). It is the portable glue
//! between the language and its libraries, defining the intrinsic and
//! primitive building blocks of all Rust code. It links to no
//! upstream libraries, no system libraries, and no libc.
//!
//! [^free]: Strictly speaking, there are some symbols which are needed but
//!          they aren't always necessary.
//!
//! The core library is *minimal*: it isn't even aware of heap allocation,
//! nor does it provide concurrency or I/O. These things require
//! platform integration, and this library is platform-agnostic.
//!
//! # How to use the core library
//!
// FIXME: Fill me in with more detail when the interface settles
//! This library is built on the assumption of a few existing symbols:
//!
//! * `memcpy`, `memcmp`, `memset` - These are core memory routines which are
//!   often generated by LLVM. Additionally, this library can make explicit
//!   calls to these functions. Their signatures are the same as found in C.
//!   These functions are often provided by the system libc, but can also be
//!   provided by the [rlibc crate](https://crates.io/crates/rlibc).
//!
//! * `rust_begin_unwind` - This function takes three arguments, a
//!   `fmt::Arguments`, a `&str`, and a `u32`. These three arguments dictate
//!   the panic message, the file at which panic was invoked, and the line.
//!   It is up to consumers of this core library to define this panic
//!   function; it is only required to never return.

// Since libcore defines many fundamental lang items, all tests live in a
// separate crate, libcoretest, to avoid bizarre issues.

// Do not remove on snapshot creation. Needed for bootstrap. (Issue #22364)
#![cfg_attr(stage0, feature(custom_attribute))]
#![crate_name = "core"]
#![stable(feature = "core", since = "1.6.0")]
#![cfg_attr(stage0, staged_api)]
#![crate_type = "rlib"]
#![doc(html_logo_url = "https://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "https://doc.rust-lang.org/favicon.ico",
       html_root_url = "https://doc.rust-lang.org/nightly/",
       html_playground_url = "https://play.rust-lang.org/",
       issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
       test(no_crate_inject, attr(deny(warnings))),
       test(attr(allow(dead_code, deprecated, unused_variables, unused_mut))))]

#![no_core]
#![deny(missing_docs)]

#![cfg_attr(stage0, feature(rustc_attrs))]
#![cfg_attr(stage0, allow(unused_attributes))]
#![feature(allow_internal_unstable)]
#![feature(associated_type_defaults)]
#![feature(concat_idents)]
#![feature(const_fn)]
#![feature(custom_attribute)]
#![feature(fundamental)]
#![feature(intrinsics)]
#![feature(lang_items)]
#![feature(no_core)]
#![feature(on_unimplemented)]
#![feature(optin_builtin_traits)]
#![feature(reflect)]
#![feature(unwind_attributes)]
#![cfg_attr(stage0, feature(simd))]
#![cfg_attr(not(stage0), feature(repr_simd, platform_intrinsics))]
#![feature(staged_api)]
#![feature(unboxed_closures)]

#[macro_use]
mod macros;

#[path = "num/float_macros.rs"]
#[macro_use]
mod float_macros;

#[path = "num/int_macros.rs"]
#[macro_use]
mod int_macros;

#[path = "num/uint_macros.rs"]
#[macro_use]
mod uint_macros;

#[path = "num/isize.rs"]  pub mod isize;
#[path = "num/i8.rs"]   pub mod i8;
#[path = "num/i16.rs"]  pub mod i16;
#[path = "num/i32.rs"]  pub mod i32;
#[path = "num/i64.rs"]  pub mod i64;

#[path = "num/usize.rs"] pub mod usize;
#[path = "num/u8.rs"]   pub mod u8;
#[path = "num/u16.rs"]  pub mod u16;
#[path = "num/u32.rs"]  pub mod u32;
#[path = "num/u64.rs"]  pub mod u64;

#[path = "num/f32.rs"]   pub mod f32;
#[path = "num/f64.rs"]   pub mod f64;

#[macro_use]
pub mod num;

/* The libcore prelude, not as all-encompassing as the libstd prelude */

pub mod prelude;

/* Core modules for ownership management */

pub mod intrinsics;
pub mod mem;
pub mod nonzero;
pub mod ptr;

/* Core language traits */

pub mod marker;
pub mod ops;
pub mod cmp;
pub mod clone;
pub mod default;
pub mod convert;
pub mod borrow;

/* Core types and methods on primitives */

pub mod any;
pub mod array;
pub mod sync;
pub mod cell;
pub mod char;
pub mod panicking;
pub mod iter;
pub mod option;
pub mod raw;
pub mod result;

pub mod slice;
pub mod str;
pub mod hash;
pub mod fmt;

// note: does not need to be public
mod tuple;
<|endoftext|>


<commit_before>#[link(name = "knob",
       vers = "1.0",
       uuid = "3299B6C2-99DE-44FD-8867-8EE7304959D7",
       url = "http://github.com/skade/knob")];

#[crate_type = "lib"];
#[comment = "A convenient Rust settings system"];
#[license = "MIT"];

//! A convenient structure to store and load settings.
//!
//! It is meant for items that are rarely read and stored, like command line flags or
//! application configuration.
//!
//! knob::Settings expects all values to implement ToStr + FromStr and stores them as
//! strings internally. This allows sideloading of settings through multiple means, e.g.
//! the command line or a simple config file. knob is not meant for structured data. If
//! you want to load such data, store the location of the data as a Knob setting and do
//! the loading parsing yourself.
//!
//! knob is typesafe in a sense that it will return you the type you wanted back if possible and
//! fails otherwise.
//!
//! knob allows you to decorate the Settings structure yourself for convenience (see examples).
//!
//! # Storing and fetching values
//!
//! The following example shows you how to store an ip in knob:
//!
//! ~~~{.rust}
//! extern mod knob;
//! use knob::*;
//! use std::rt::io::net::ip::IpAddr;
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set("ip", "0.0.0.0:4567");
//!   let socket: IpAddr = settings.socket();
//!   assert_eq!(socket.to_str(), ~"0.0.0.0:4567")
//! }
//! ~~~
//!
//! This works the same for IPv6 addresses:
//!
//! ~~~{.rust}
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set("ip", "::0.0.0.1");
//!   let socket: IpAddr = settings.fetch("ip");
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! # Providing your own keys
//!
//! You can use enums as keys, as long as they implement ToStr:
//!
//! ~~~{.rust}
//! enum Keys {
//!   Port,
//!   Ip,
//!   Addr,
//! }
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set(Ip, "::0.0.0.1");
//!   let socket: IpAddr = settings.fetch(Ip);
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! # Registering command line options
//!
//! Knob allows you to register command line options to read from the command line later.
//!
//! The options are getopts option groups.
//!
//! If the loading of the command line args fails, an error will be returned.
//!
//! ~~~{.rust}
//! use extra::getopts::groups::*;
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.opt(optopt("p", "port", "the port to bind to", "4000"))
//!   settings.opt(reqopt("e", "environment", "the environment to run in", ""));;
//!   let errors = settings.load_os_args();
//!   if errors.is_some() {
//!     println(settings.usage("Try one of these:"))
//!   }
//! }
//! ~~~
//!
//! # Decorating the settings struct
//!
//! To make matters more convenient, you can implement a decorator
//! to implement your own loading behaviour.
//!
//! ~~~{.rust}
//! enum Keys {
//!   Port,
//!   Ip,
//!   Addr,
//! }
//!
//! pub trait SocketSettings {
//!   fn socket(&self) -> SocketAddr;
//!   fn port(&self) -> u16;
//!   fn ip(&self) -> IpAddr;
//! }
//!
//! impl SocketSettings for Settings {
//!   fn socket(&self) -> SocketAddr {
//!     do self.fetch_with(Addr) |addr| {
//!       match addr {
//!         Some(socket_addr) => { socket_addr },
//!         None => {
//!           let port: u16 = self.port();
//!           let ip: IpAddr = self.ip();
//!           SocketAddr { ip: ip, port: port }
//!         }
//!       }
//!     }
//!   }
//!
//!   fn port(&self) -> u16 {
//!     self.fetch(Port).unwrap_or(8080)
//!   }
//!
//!   fn ip(&self) -> IpAddr {
//!     self.fetch(Ip).unwrap_or(Ipv4Addr(127,0,0,1))
//!   }
//! }
//! enum Keys {
//!   Port,
//!   Ip,
//! }
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set(Ip, "::0.0.0.1");
//!   let socket: IpAddr = settings.fetch(Ip);
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! knob goes up to 11.

extern mod extra;

use std::hashmap::*;
use std::os;

use extra::getopts::groups::{OptGroup,getopts,usage};
use extra::getopts::Fail_;

/// The settings structure we save the options and settings in.
#[deriving(Clone)]
pub struct Settings {
  priv store: HashMap<~str,~str>,
  priv options: ~[OptGroup],
}

impl Settings {
  /// Create a new Settings struct.
  pub fn new() -> Settings {
    Settings { store: HashMap::new(), options: ~[] }
  }

  /// Set a settings key to a value. The value will be serialized.
  pub fn set<A: ToStr, T: ToStr>(&mut self, setting: A, value: T) {
    self.store.swap(setting.to_str(), value.to_str());
  }

  /// Set a value using an Option struct. The value will only be set if the
  /// Some value is given. This way, you can avoid unwrapping the result of a
  /// previous operation by yourself.
  pub fn set_opt<A: ToStr, T: ToStr>(&mut self, setting: A, value: Option<T>) {
    if !value.is_none() {
      self.store.swap(setting.to_str(), value.unwrap().to_str());
    }
  }

  /// Fetch a setting for a key. Fails if the setting is present but could not be
  /// parsed.
  pub fn fetch<A: ToStr, T: FromStr>(&self, setting: A) -> Option<T> {
    match self.store.find(&setting.to_str()) {
      Some(string) => {
        let value = from_str(string.to_owned());
        if value.is_none() {
          fail!("setting could not be parsed:" + setting.to_str())
        }
        value
      },
      None => { None }
    }
  }

  /// Fetch a setting for a key and pass it to given function. The result of the function
  /// will be returned.
  pub fn fetch_with<A: ToStr, T: FromStr>(&self, setting: A, f: &fn(Option<T>) -> T) -> T {
    let value = self.fetch(setting.to_str());
    f(value)
  }

  /// Register a commandline for laster use with load_args.
  ///
  /// Currently, only optopt and reqopt are properly supported.
  pub fn opt(&mut self, opt: OptGroup) {
    self.options.push(opt);
  }

  /// Load the command line argument given by the OS.
  ///
  /// Optionally returns failures.
  pub fn load_os_args(&mut self) -> Option<Fail_> {
    let args = os::args();
    self.load_args(args)
  }

  /// Load a list of command line arguments.
  ///
  /// Automatically sets "knob.progname" to the name of the program.
  ///
  /// Optionally returns failures.
  pub fn load_args(&mut self, args: ~[~str]) -> Option<Fail_> {
    let ref prog_name = args[0];

    self.set("knob.progname", prog_name.clone());

    let matches = match getopts(args.tail(), self.options) {
      Ok(m) => { m }
      Err(fail) => { return Some(fail) }
    };

    let given_options = self.options.clone();
    for opt in given_options.iter() {
      let opt_strings = &[opt.short_name.clone(), opt.long_name.clone()];
      self.set_opt(opt.long_name.clone(), matches.opts_str(opt_strings))
    };
    None
  }

  /// Returns the usage string for the stored OptGroups. Pass `brief`
  /// to have it included.
  pub fn usage(&self, brief: &str) -> ~str {
    usage(brief, self.options)
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::rt::io::net::ip::{SocketAddr, IpAddr, Ipv4Addr};
  use extra::getopts::groups::*;

  #[deriving(ToStr)]
  enum Keys {
    Ip,
    Port,
    Addr
  }

  pub trait SocketSettings {
    fn socket(&self) -> SocketAddr;
    fn port(&self) -> u16;
    fn ip(&self) -> IpAddr;
  }

  impl SocketSettings for Settings {
    fn socket(&self) -> SocketAddr {
      do self.fetch_with(Addr) |addr| {
        match addr {
          Some(socket_addr) => { socket_addr },
          None => {
            let port: u16 = self.port();
            let ip: IpAddr = self.ip();
            SocketAddr { ip: ip, port: port }
          }
        }
      }
    }

    fn port(&self) -> u16 {
      self.fetch(Port).unwrap_or(8080)
    }

    fn ip(&self) -> IpAddr {
      self.fetch(Ip).unwrap_or(Ipv4Addr(127,0,0,1))
    }
  }

  #[test]
  fn test_simple_conversion() {
    let mut settings = Settings::new();
    settings.set("port", 12345);
    let settings = settings.fetch("port");
    assert_eq!(Some(12345), settings)
  }

  #[test]
  #[should_fail]
  fn test_should_fail_on_garbage() {
    let mut settings = Settings::new();
    settings.set("port", "foobar");
    let settings: Option<int> = settings.fetch("port");
    assert_eq!(Some(12345), settings)
  }

  #[test]
  fn test_enum() {
    let mut settings = Settings::new();
    settings.set(Port, 12345);
    let settings = settings.fetch(Port);
    assert_eq!(Some(12345), settings)
  }

  #[test]
  fn test_compound_socket_settings() {
    let mut settings = Settings::new();
    settings.set(Port, "12345");
    settings.set(Ip, "127.0.0.1");
    let socket = settings.socket();
    assert_eq!(socket.to_str(), ~"127.0.0.1:12345")
  }

  #[test]
  fn test_ipv6_addr() {
    let mut settings = Settings::new();
    settings.set(Ip, "::0.0.0.1");
    let ip = settings.ip();
    assert_eq!(ip.to_str(), ~"::0.0.0.1")
  }

  #[test]
  fn test_socket_overrides_port() {
    let mut settings = Settings::new();
    settings.set(Port, "12345");
    settings.set(Ip, "127.0.0.1");
    settings.set(Addr, "0.0.0.0:4567");
    let socket = settings.socket();
    assert_eq!(socket.to_str(), ~"0.0.0.0:4567")
  }

  #[test]
  fn test_opt_parse_optional() {
    let mut settings = Settings::new();
    settings.opt(optopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog", ~"-p", ~"3000"];
    let error = settings.load_args(args);

    assert!(error.is_none());
    assert_eq!(settings.fetch("port"), Some(3000))
  }

  #[test]
  fn test_opt_parse_req_with_given() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog", ~"-p", ~"3000"];
    let error = settings.load_args(args);

    assert!(error.is_none());
    assert_eq!(settings.fetch("port"), Some(3000))
  }

  #[test]
  fn test_opt_parse_req_without_given() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog"];
    let error = settings.load_args(args);

    assert!(error.is_some());
    let port: Option<int> = settings.fetch("port");
    assert_eq!(port, None)
  }

  #[test]
  fn test_usage() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));
    let usage = settings.usage(&"this is how it works");

    assert!(usage.contains("this is how it works"))
    assert!(usage.contains("--port"))
  }

  #[test]
  fn test_goes_up_to_eleven() {
    let mut settings = Settings::new();
    settings.set("knob", 11);
    assert_eq!(settings.fetch("knob"), Some(11));
  }
}<commit_msg>One example in the docs completely missed the point<commit_after>#[link(name = "knob",
       vers = "1.0",
       uuid = "3299B6C2-99DE-44FD-8867-8EE7304959D7",
       url = "http://github.com/skade/knob")];

#[crate_type = "lib"];
#[comment = "A convenient Rust settings system"];
#[license = "MIT"];

//! A convenient structure to store and load settings.
//!
//! It is meant for items that are rarely read and stored, like command line flags or
//! application configuration.
//!
//! knob::Settings expects all values to implement ToStr + FromStr and stores them as
//! strings internally. This allows sideloading of settings through multiple means, e.g.
//! the command line or a simple config file. knob is not meant for structured data. If
//! you want to load such data, store the location of the data as a Knob setting and do
//! the loading parsing yourself.
//!
//! knob is typesafe in a sense that it will return you the type you wanted back if possible and
//! fails otherwise.
//!
//! knob allows you to decorate the Settings structure yourself for convenience (see examples).
//!
//! # Storing and fetching values
//!
//! The following example shows you how to store an ip in knob:
//!
//! ~~~{.rust}
//! extern mod knob;
//! use knob::*;
//! use std::rt::io::net::ip::IpAddr;
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set("ip", "0.0.0.0:4567");
//!   let socket: IpAddr = settings.socket();
//!   assert_eq!(socket.to_str(), ~"0.0.0.0:4567")
//! }
//! ~~~
//!
//! This works the same for IPv6 addresses:
//!
//! ~~~{.rust}
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set("ip", "::0.0.0.1");
//!   let socket: IpAddr = settings.fetch("ip");
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! # Providing your own keys
//!
//! You can use enums as keys, as long as they implement ToStr:
//!
//! ~~~{.rust}
//! enum Keys {
//!   Port,
//!   Ip,
//!   Addr,
//! }
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set(Ip, "::0.0.0.1");
//!   let socket: IpAddr = settings.fetch(Ip);
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! # Registering command line options
//!
//! Knob allows you to register command line options to read from the command line later.
//!
//! The options are getopts option groups.
//!
//! If the loading of the command line args fails, an error will be returned.
//!
//! ~~~{.rust}
//! use extra::getopts::groups::*;
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.opt(optopt("p", "port", "the port to bind to", "4000"))
//!   settings.opt(reqopt("e", "environment", "the environment to run in", ""));;
//!   let errors = settings.load_os_args();
//!   if errors.is_some() {
//!     println(settings.usage("Try one of these:"))
//!   }
//! }
//! ~~~
//!
//! # Decorating the settings struct
//!
//! To make matters more convenient, you can implement a decorator
//! to implement your own loading behaviour.
//!
//! ~~~{.rust}
//! enum Keys {
//!   Port,
//!   Ip,
//!   Addr,
//! }
//!
//! pub trait SocketSettings {
//!   fn socket(&self) -> SocketAddr;
//!   fn port(&self) -> u16;
//!   fn ip(&self) -> IpAddr;
//! }
//!
//! impl SocketSettings for Settings {
//!   fn socket(&self) -> SocketAddr {
//!     do self.fetch_with(Addr) |addr| {
//!       match addr {
//!         Some(socket_addr) => { socket_addr },
//!         None => {
//!           let port: u16 = self.port();
//!           let ip: IpAddr = self.ip();
//!           SocketAddr { ip: ip, port: port }
//!         }
//!       }
//!     }
//!   }
//!
//!   fn port(&self) -> u16 {
//!     self.fetch(Port).unwrap_or(8080)
//!   }
//!
//!   fn ip(&self) -> IpAddr {
//!     self.fetch(Ip).unwrap_or(Ipv4Addr(127,0,0,1))
//!   }
//! }
//! enum Keys {
//!   Port,
//!   Ip,
//! }
//!
//! fn main() {
//!   let settings = Settings::new();
//!   settings.set(Ip, "::0.0.0.1");
//!   let socket: IpAddr = settings.ip();
//!   assert_eq!(socket.to_str(), ~"::0.0.0.1")
//! }
//! ~~~
//!
//! knob goes up to 11.

extern mod extra;

use std::hashmap::*;
use std::os;

use extra::getopts::groups::{OptGroup,getopts,usage};
use extra::getopts::Fail_;

/// The settings structure we save the options and settings in.
#[deriving(Clone)]
pub struct Settings {
  priv store: HashMap<~str,~str>,
  priv options: ~[OptGroup],
}

impl Settings {
  /// Create a new Settings struct.
  pub fn new() -> Settings {
    Settings { store: HashMap::new(), options: ~[] }
  }

  /// Set a settings key to a value. The value will be serialized.
  pub fn set<A: ToStr, T: ToStr>(&mut self, setting: A, value: T) {
    self.store.swap(setting.to_str(), value.to_str());
  }

  /// Set a value using an Option struct. The value will only be set if the
  /// Some value is given. This way, you can avoid unwrapping the result of a
  /// previous operation by yourself.
  pub fn set_opt<A: ToStr, T: ToStr>(&mut self, setting: A, value: Option<T>) {
    if !value.is_none() {
      self.store.swap(setting.to_str(), value.unwrap().to_str());
    }
  }

  /// Fetch a setting for a key. Fails if the setting is present but could not be
  /// parsed.
  pub fn fetch<A: ToStr, T: FromStr>(&self, setting: A) -> Option<T> {
    match self.store.find(&setting.to_str()) {
      Some(string) => {
        let value = from_str(string.to_owned());
        if value.is_none() {
          fail!("setting could not be parsed:" + setting.to_str())
        }
        value
      },
      None => { None }
    }
  }

  /// Fetch a setting for a key and pass it to given function. The result of the function
  /// will be returned.
  pub fn fetch_with<A: ToStr, T: FromStr>(&self, setting: A, f: &fn(Option<T>) -> T) -> T {
    let value = self.fetch(setting.to_str());
    f(value)
  }

  /// Register a commandline for laster use with load_args.
  ///
  /// Currently, only optopt and reqopt are properly supported.
  pub fn opt(&mut self, opt: OptGroup) {
    self.options.push(opt);
  }

  /// Load the command line argument given by the OS.
  ///
  /// Optionally returns failures.
  pub fn load_os_args(&mut self) -> Option<Fail_> {
    let args = os::args();
    self.load_args(args)
  }

  /// Load a list of command line arguments.
  ///
  /// Automatically sets "knob.progname" to the name of the program.
  ///
  /// Optionally returns failures.
  pub fn load_args(&mut self, args: ~[~str]) -> Option<Fail_> {
    let ref prog_name = args[0];

    self.set("knob.progname", prog_name.clone());

    let matches = match getopts(args.tail(), self.options) {
      Ok(m) => { m }
      Err(fail) => { return Some(fail) }
    };

    let given_options = self.options.clone();
    for opt in given_options.iter() {
      let opt_strings = &[opt.short_name.clone(), opt.long_name.clone()];
      self.set_opt(opt.long_name.clone(), matches.opts_str(opt_strings))
    };
    None
  }

  /// Returns the usage string for the stored OptGroups. Pass `brief`
  /// to have it included.
  pub fn usage(&self, brief: &str) -> ~str {
    usage(brief, self.options)
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::rt::io::net::ip::{SocketAddr, IpAddr, Ipv4Addr};
  use extra::getopts::groups::*;

  #[deriving(ToStr)]
  enum Keys {
    Ip,
    Port,
    Addr
  }

  pub trait SocketSettings {
    fn socket(&self) -> SocketAddr;
    fn port(&self) -> u16;
    fn ip(&self) -> IpAddr;
  }

  impl SocketSettings for Settings {
    fn socket(&self) -> SocketAddr {
      do self.fetch_with(Addr) |addr| {
        match addr {
          Some(socket_addr) => { socket_addr },
          None => {
            let port: u16 = self.port();
            let ip: IpAddr = self.ip();
            SocketAddr { ip: ip, port: port }
          }
        }
      }
    }

    fn port(&self) -> u16 {
      self.fetch(Port).unwrap_or(8080)
    }

    fn ip(&self) -> IpAddr {
      self.fetch(Ip).unwrap_or(Ipv4Addr(127,0,0,1))
    }
  }

  #[test]
  fn test_simple_conversion() {
    let mut settings = Settings::new();
    settings.set("port", 12345);
    let settings = settings.fetch("port");
    assert_eq!(Some(12345), settings)
  }

  #[test]
  #[should_fail]
  fn test_should_fail_on_garbage() {
    let mut settings = Settings::new();
    settings.set("port", "foobar");
    let settings: Option<int> = settings.fetch("port");
    assert_eq!(Some(12345), settings)
  }

  #[test]
  fn test_enum() {
    let mut settings = Settings::new();
    settings.set(Port, 12345);
    let settings = settings.fetch(Port);
    assert_eq!(Some(12345), settings)
  }

  #[test]
  fn test_compound_socket_settings() {
    let mut settings = Settings::new();
    settings.set(Port, "12345");
    settings.set(Ip, "127.0.0.1");
    let socket = settings.socket();
    assert_eq!(socket.to_str(), ~"127.0.0.1:12345")
  }

  #[test]
  fn test_ipv6_addr() {
    let mut settings = Settings::new();
    settings.set(Ip, "::0.0.0.1");
    let ip = settings.ip();
    assert_eq!(ip.to_str(), ~"::0.0.0.1")
  }

  #[test]
  fn test_socket_overrides_port() {
    let mut settings = Settings::new();
    settings.set(Port, "12345");
    settings.set(Ip, "127.0.0.1");
    settings.set(Addr, "0.0.0.0:4567");
    let socket = settings.socket();
    assert_eq!(socket.to_str(), ~"0.0.0.0:4567")
  }

  #[test]
  fn test_opt_parse_optional() {
    let mut settings = Settings::new();
    settings.opt(optopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog", ~"-p", ~"3000"];
    let error = settings.load_args(args);

    assert!(error.is_none());
    assert_eq!(settings.fetch("port"), Some(3000))
  }

  #[test]
  fn test_opt_parse_req_with_given() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog", ~"-p", ~"3000"];
    let error = settings.load_args(args);

    assert!(error.is_none());
    assert_eq!(settings.fetch("port"), Some(3000))
  }

  #[test]
  fn test_opt_parse_req_without_given() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));

    let args = ~[~"myprog"];
    let error = settings.load_args(args);

    assert!(error.is_some());
    let port: Option<int> = settings.fetch("port");
    assert_eq!(port, None)
  }

  #[test]
  fn test_usage() {
    let mut settings = Settings::new();
    settings.opt(reqopt("p", "port", "The port to bind to", "eg: 4000"));
    let usage = settings.usage(&"this is how it works");

    assert!(usage.contains("this is how it works"))
    assert!(usage.contains("--port"))
  }

  #[test]
  fn test_goes_up_to_eleven() {
    let mut settings = Settings::new();
    settings.set("knob", 11);
    assert_eq!(settings.fetch("knob"), Some(11));
  }
}<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A "once initialization" primitive
//!
//! This primitive is meant to be used to run one-time initialization. An
//! example use case would be for initializing an FFI library.

use core::prelude::*;

use core::int;
use core::atomics;

use mutex::{StaticMutex, MUTEX_INIT};

/// A type which can be used to run a one-time global initialization. This type
/// is *unsafe* to use because it is built on top of the `Mutex` in this module.
/// It does not know whether the currently running task is in a green or native
/// context, and a blocking mutex should *not* be used under normal
/// circumstances on a green task.
///
/// Despite its unsafety, it is often useful to have a one-time initialization
/// routine run for FFI bindings or related external functionality. This type
/// can only be statically constructed with the `ONCE_INIT` value.
///
/// # Example
///
/// ```rust
/// use sync::one::{Once, ONCE_INIT};
///
/// static mut START: Once = ONCE_INIT;
/// unsafe {
///     START.doit(|| {
///         // run initialization here
///     });
/// }
/// ```
pub struct Once {
    mutex: StaticMutex,
    cnt: atomics::AtomicInt,
    lock_cnt: atomics::AtomicInt,
}

/// Initialization value for static `Once` values.
pub static ONCE_INIT: Once = Once {
    mutex: MUTEX_INIT,
    cnt: atomics::INIT_ATOMIC_INT,
    lock_cnt: atomics::INIT_ATOMIC_INT,
};

impl Once {
    /// Perform an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `doit` has been called, and
    /// otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling *os thread* if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified).
    pub fn doit(&self, f: ||) {
        // Optimize common path: load is much cheaper than fetch_add.
        if self.cnt.load(atomics::SeqCst) < 0 {
            return
        }

        // Implementation-wise, this would seem like a fairly trivial primitive.
        // The stickler part is where our mutexes currently require an
        // allocation, and usage of a `Once` shouldn't leak this allocation.
        //
        // This means that there must be a deterministic destroyer of the mutex
        // contained within (because it's not needed after the initialization
        // has run).
        //
        // The general scheme here is to gate all future threads once
        // initialization has completed with a "very negative" count, and to
        // allow through threads to lock the mutex if they see a non negative
        // count. For all threads grabbing the mutex, exactly one of them should
        // be responsible for unlocking the mutex, and this should only be done
        // once everyone else is done with the mutex.
        //
        // This atomicity is achieved by swapping a very negative value into the
        // shared count when the initialization routine has completed. This will
        // read the number of threads which will at some point attempt to
        // acquire the mutex. This count is then squirreled away in a separate
        // variable, and the last person on the way out of the mutex is then
        // responsible for destroying the mutex.
        //
        // It is crucial that the negative value is swapped in *after* the
        // initialization routine has completed because otherwise new threads
        // calling `doit` will return immediately before the initialization has
        // completed.

        let prev = self.cnt.fetch_add(1, atomics::SeqCst);
        if prev < 0 {
            // Make sure we never overflow, we'll never have int::MIN
            // simultaneous calls to `doit` to make this value go back to 0
            self.cnt.store(int::MIN, atomics::SeqCst);
            return
        }

        // If the count is negative, then someone else finished the job,
        // otherwise we run the job and record how many people will try to grab
        // this lock
        let guard = self.mutex.lock();
        if self.cnt.load(atomics::SeqCst) > 0 {
            f();
            let prev = self.cnt.swap(int::MIN, atomics::SeqCst);
            self.lock_cnt.store(prev, atomics::SeqCst);
        }
        drop(guard);

        // Last one out cleans up after everyone else, no leaks!
        if self.lock_cnt.fetch_add(-1, atomics::SeqCst) == 1 {
            unsafe { self.mutex.destroy() }
        }
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;
    use std::task;
    use super::{ONCE_INIT, Once};

    #[test]
    fn smoke_once() {
        static mut o: Once = ONCE_INIT;
        let mut a = 0;
        unsafe { o.doit(|| a += 1); }
        assert_eq!(a, 1);
        unsafe { o.doit(|| a += 1); }
        assert_eq!(a, 1);
    }

    #[test]
    fn stampede_once() {
        static mut o: Once = ONCE_INIT;
        static mut run: bool = false;

        let (tx, rx) = channel();
        for _ in range(0, 10) {
            let tx = tx.clone();
            spawn(proc() {
                for _ in range(0, 4) { task::deschedule() }
                unsafe {
                    o.doit(|| {
                        assert!(!run);
                        run = true;
                    });
                    assert!(run);
                }
                tx.send(());
            });
        }

        unsafe {
            o.doit(|| {
                assert!(!run);
                run = true;
            });
            assert!(run);
        }

        for _ in range(0, 10) {
            rx.recv();
        }
    }
}
<commit_msg>auto merge of #14813 : cmr/rust/once-docs-unsafe, r=alexcrichton<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A "once initialization" primitive
//!
//! This primitive is meant to be used to run one-time initialization. An
//! example use case would be for initializing an FFI library.

use core::prelude::*;

use core::int;
use core::atomics;

use mutex::{StaticMutex, MUTEX_INIT};

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Useful for one-time initialization for FFI or related
/// functionality. This type can only be constructed with the `ONCE_INIT`
/// value.
///
/// # Example
///
/// ```rust
/// use sync::one::{Once, ONCE_INIT};
///
/// static mut START: Once = ONCE_INIT;
///
/// unsafe {
///     START.doit(|| {
///         // run initialization here
///     });
/// }
/// ```
pub struct Once {
    mutex: StaticMutex,
    cnt: atomics::AtomicInt,
    lock_cnt: atomics::AtomicInt,
}

/// Initialization value for static `Once` values.
pub static ONCE_INIT: Once = Once {
    mutex: MUTEX_INIT,
    cnt: atomics::INIT_ATOMIC_INT,
    lock_cnt: atomics::INIT_ATOMIC_INT,
};

impl Once {
    /// Perform an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `doit` has been called, and
    /// otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling task if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified).
    pub fn doit(&self, f: ||) {
        // Optimize common path: load is much cheaper than fetch_add.
        if self.cnt.load(atomics::SeqCst) < 0 {
            return
        }

        // Implementation-wise, this would seem like a fairly trivial primitive.
        // The stickler part is where our mutexes currently require an
        // allocation, and usage of a `Once` shouldn't leak this allocation.
        //
        // This means that there must be a deterministic destroyer of the mutex
        // contained within (because it's not needed after the initialization
        // has run).
        //
        // The general scheme here is to gate all future threads once
        // initialization has completed with a "very negative" count, and to
        // allow through threads to lock the mutex if they see a non negative
        // count. For all threads grabbing the mutex, exactly one of them should
        // be responsible for unlocking the mutex, and this should only be done
        // once everyone else is done with the mutex.
        //
        // This atomicity is achieved by swapping a very negative value into the
        // shared count when the initialization routine has completed. This will
        // read the number of threads which will at some point attempt to
        // acquire the mutex. This count is then squirreled away in a separate
        // variable, and the last person on the way out of the mutex is then
        // responsible for destroying the mutex.
        //
        // It is crucial that the negative value is swapped in *after* the
        // initialization routine has completed because otherwise new threads
        // calling `doit` will return immediately before the initialization has
        // completed.

        let prev = self.cnt.fetch_add(1, atomics::SeqCst);
        if prev < 0 {
            // Make sure we never overflow, we'll never have int::MIN
            // simultaneous calls to `doit` to make this value go back to 0
            self.cnt.store(int::MIN, atomics::SeqCst);
            return
        }

        // If the count is negative, then someone else finished the job,
        // otherwise we run the job and record how many people will try to grab
        // this lock
        let guard = self.mutex.lock();
        if self.cnt.load(atomics::SeqCst) > 0 {
            f();
            let prev = self.cnt.swap(int::MIN, atomics::SeqCst);
            self.lock_cnt.store(prev, atomics::SeqCst);
        }
        drop(guard);

        // Last one out cleans up after everyone else, no leaks!
        if self.lock_cnt.fetch_add(-1, atomics::SeqCst) == 1 {
            unsafe { self.mutex.destroy() }
        }
    }
}

#[cfg(test)]
mod test {
    use std::prelude::*;
    use std::task;
    use super::{ONCE_INIT, Once};

    #[test]
    fn smoke_once() {
        static mut o: Once = ONCE_INIT;
        let mut a = 0;
        unsafe { o.doit(|| a += 1); }
        assert_eq!(a, 1);
        unsafe { o.doit(|| a += 1); }
        assert_eq!(a, 1);
    }

    #[test]
    fn stampede_once() {
        static mut o: Once = ONCE_INIT;
        static mut run: bool = false;

        let (tx, rx) = channel();
        for _ in range(0, 10) {
            let tx = tx.clone();
            spawn(proc() {
                for _ in range(0, 4) { task::deschedule() }
                unsafe {
                    o.doit(|| {
                        assert!(!run);
                        run = true;
                    });
                    assert!(run);
                }
                tx.send(());
            });
        }

        unsafe {
            o.doit(|| {
                assert!(!run);
                run = true;
            });
            assert!(run);
        }

        for _ in range(0, 10) {
            rx.recv();
        }
    }
}
<|endoftext|>


<commit_before>use square_matrix::SquareMatrix;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Mark {
    None,
    Star,
    Prime,
}

#[derive(Debug)]
pub struct MarkMatrix {
    marks: SquareMatrix<Mark>,
}

// XXX: Use two bitmatrices. Each row has a fixed number of u64 integers. Use
// bit-search

impl MarkMatrix {
    pub fn new(n: usize) -> MarkMatrix {
        MarkMatrix {
            marks: SquareMatrix::from_row_vec(n, (0..n * n).map(|_| Mark::None).collect()),
        }
    }

    pub fn n(&self) -> usize {
        self.marks.n()
    }

    pub fn toggle_star(&mut self, pos: (usize, usize)) {
        if self.is_star(pos) {
            self.unmark(pos);
        } else {
            self.star(pos);
        }
    }

    pub fn unmark(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::None;
    }

    pub fn star(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::Star;
    }

    pub fn prime(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::Prime;
    }

    pub fn is_star(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::Star => true,
            _ => false,
        }
    }

    pub fn is_prime(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::Prime => true,
            _ => false,
        }
    }

    #[cfg(test)]
    pub fn is_none(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::None => true,
            _ => false,
        }
    }

    pub fn find_first_star_in_row(&self, row: usize) -> Option<usize> {
        for col in 0..self.n() {
            if self.is_star((row, col)) {
                return Some(col);
            }
        }
        return None;
    }

    pub fn find_first_prime_in_row(&self, row: usize) -> Option<usize> {
        for col in 0..self.n() {
            if self.is_prime((row, col)) {
                return Some(col);
            }
        }
        return None;
    }

    pub fn find_first_star_in_col(&self, col: usize) -> Option<usize> {
        for row in 0..self.n() {
            if self.is_star((row, col)) {
                return Some(row);
            }
        }
        return None;
    }

    pub fn clear_primes(&mut self) {
        for row in 0..self.n() {
            for col in 0..self.n() {
                if self.is_prime((row, col)) {
                    self.unmark((row, col));
                }
            }
        }
    }
}
<commit_msg>Use more compact representation of Mark<commit_after>use square_matrix::SquareMatrix;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[repr(u8)]
enum Mark {
    None,
    Star,
    Prime,
}

#[derive(Debug)]
pub struct MarkMatrix {
    marks: SquareMatrix<Mark>,
}

impl MarkMatrix {
    pub fn new(n: usize) -> MarkMatrix {
        MarkMatrix {
            marks: SquareMatrix::from_row_vec(n, (0..n * n).map(|_| Mark::None).collect()),
        }
    }

    pub fn n(&self) -> usize {
        self.marks.n()
    }

    pub fn toggle_star(&mut self, pos: (usize, usize)) {
        if self.is_star(pos) {
            self.unmark(pos);
        } else {
            self.star(pos);
        }
    }

    pub fn unmark(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::None;
    }

    pub fn star(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::Star;
    }

    pub fn prime(&mut self, pos: (usize, usize)) {
        self.marks[pos] = Mark::Prime;
    }

    pub fn is_star(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::Star => true,
            _ => false,
        }
    }

    pub fn is_prime(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::Prime => true,
            _ => false,
        }
    }

    #[cfg(test)]
    pub fn is_none(&self, pos: (usize, usize)) -> bool {
        match self.marks[pos] {
            Mark::None => true,
            _ => false,
        }
    }

    pub fn find_first_star_in_row(&self, row: usize) -> Option<usize> {
        for col in 0..self.n() {
            if self.is_star((row, col)) {
                return Some(col);
            }
        }
        return None;
    }

    pub fn find_first_prime_in_row(&self, row: usize) -> Option<usize> {
        for col in 0..self.n() {
            if self.is_prime((row, col)) {
                return Some(col);
            }
        }
        return None;
    }

    pub fn find_first_star_in_col(&self, col: usize) -> Option<usize> {
        for row in 0..self.n() {
            if self.is_star((row, col)) {
                return Some(row);
            }
        }
        return None;
    }

    pub fn clear_primes(&mut self) {
        for row in 0..self.n() {
            for col in 0..self.n() {
                if self.is_prime((row, col)) {
                    self.unmark((row, col));
                }
            }
        }
    }
}
<|endoftext|>