<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> {
        let ctx = self.0.builder_mut();
        // FIXME clear chain certs to clean up if called multiple times
        try!(ctx.set_certificate(&pkcs12.0.cert));
        try!(ctx.set_private_key(&pkcs12.0.pkey));
        try!(ctx.check_private_key());
        for cert in pkcs12.0.chain {
            try!(ctx.add_extra_chain_cert(cert));
        }
        Ok(())
    }

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

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, self.0.builder_mut());
        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 {
    /// 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 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 {
    /// 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 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>Allow initialization for openssl extension builders<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> {
        let ctx = self.0.builder_mut();
        // FIXME clear chain certs to clean up if called multiple times
        try!(ctx.set_certificate(&pkcs12.0.cert));
        try!(ctx.set_private_key(&pkcs12.0.pkey));
        try!(ctx.check_private_key());
        for cert in pkcs12.0.chain {
            try!(ctx.add_extra_chain_cert(cert));
        }
        Ok(())
    }

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

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        supported_protocols(protocols, self.0.builder_mut());
        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>/**
 * An atomically reference counted wrapper that can be used to
 * share immutable data between tasks.
 */

import sys::methods;

export arc, get, clone;

export exclusive, methods;

#[abi = "cdecl"]
extern mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: &mut libc::intptr_t)
        -> libc::intptr_t;
}

type arc_data<T> = {
    mut count: libc::intptr_t,
    data: T
};

class arc_destruct<T> {
  let data: *libc::c_void;
  new(data: *libc::c_void) { self.data = data; }
  drop unsafe {
     let data: ~arc_data<T> = unsafe::reinterpret_cast(self.data);
     let new_count = rustrt::rust_atomic_decrement(&mut data.count);
     assert new_count >= 0;
     if new_count == 0 {
         // drop glue takes over.
     } else {
       unsafe::forget(data);
     }
  }
}

type arc<T: const send> = arc_destruct<T>;

/// Create an atomically reference counted wrapper.
fn arc<T: const send>(-data: T) -> arc<T> {
    let data = ~{mut count: 1, data: data};
    unsafe {
        let ptr = unsafe::transmute(data);
        arc_destruct(ptr)
    }
}

/**
 * Access the underlying data in an atomically reference counted
 * wrapper.
 */
fn get<T: const send>(rc: &arc<T>) -> &T {
    unsafe {
        let ptr: ~arc_data<T> = unsafe::reinterpret_cast((*rc).data);
        // Cast us back into the correct region
        let r = unsafe::reinterpret_cast(&ptr.data);
        unsafe::forget(ptr);
        return r;
    }
}

/**
 * Duplicate an atomically reference counted wrapper.
 *
 * The resulting two `arc` objects will point to the same underlying data
 * object. However, one of the `arc` objects can be sent to another task,
 * allowing them to share the underlying data.
 */
fn clone<T: const send>(rc: &arc<T>) -> arc<T> {
    unsafe {
        let ptr: ~arc_data<T> = unsafe::reinterpret_cast((*rc).data);
        let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
        assert new_count >= 2;
        unsafe::forget(ptr);
    }
    arc_destruct((*rc).data)
}

// An arc over mutable data that is protected by a lock.
type ex_data<T: send> = {lock: sys::lock_and_signal, mut data: T};
type exclusive<T: send> = arc_destruct<ex_data<T>>;

fn exclusive<T:send >(-data: T) -> exclusive<T> {
    let data = ~{mut count: 1, data: {lock: sys::lock_and_signal(),
                                      data: data}};
    unsafe {
        let ptr = unsafe::reinterpret_cast(data);
        unsafe::forget(data);
        arc_destruct(ptr)
    }
}

impl methods<T: send> for exclusive<T> {
    fn clone() -> exclusive<T> {
        unsafe {
            // this makes me nervous...
            let ptr: ~arc_data<ex_data<T>> =
                  unsafe::reinterpret_cast(self.data);
            let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
            assert new_count > 1;
            unsafe::forget(ptr);
        }
        arc_destruct(self.data)
    }

    unsafe fn with<U>(f: fn(sys::condition, x: &mut T) -> U) -> U {
        let ptr: ~arc_data<ex_data<T>> =
            unsafe::reinterpret_cast(self.data);
        assert ptr.count > 0;
        let r = {
            let rec: &ex_data<T> = &(*ptr).data;
            rec.lock.lock_cond(|c| f(c, &mut rec.data))
        };
        unsafe::forget(ptr);
        r
    }
}

#[cfg(test)]
mod tests {
    import comm::*;
    import future::extensions;

    #[test]
    fn manually_share_arc() {
        let v = ~[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let arc_v = arc::arc(v);

        let p = port();
        let c = chan(p);

        do task::spawn() {
            let p = port();
            c.send(chan(p));

            let arc_v = p.recv();

            let v = *arc::get::<~[int]>(&arc_v);
            assert v[3] == 4;
        };

        let c = p.recv();
        c.send(arc::clone(&arc_v));

        assert (*arc::get(&arc_v))[2] == 3;

        log(info, arc_v);
    }

    #[test]
    #[ignore] // this can probably infinite loop too.
    fn exclusive_arc() {
        let mut futures = ~[];

        let num_tasks = 10u;
        let count = 1000u;

        let total = exclusive(~mut 0u);

        for uint::range(0u, num_tasks) |_i| {
            let total = total.clone();
            futures += ~[future::spawn(|| {
                for uint::range(0u, count) |_i| {
                    do total.with |_cond, count| {
                        **count += 1u;
                    }
                }
            })];
        };

        for futures.each |f| { f.get() }

        do total.with |_cond, total| {
            assert **total == num_tasks * count
        };
    }
}
<commit_msg>Document arc::exclusive.<commit_after>/**
 * An atomically reference counted wrapper that can be used to
 * share immutable data between tasks.
 */

import sys::methods;

export arc, get, clone;

export exclusive, methods;

#[abi = "cdecl"]
extern mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: &mut libc::intptr_t)
        -> libc::intptr_t;
}

type arc_data<T> = {
    mut count: libc::intptr_t,
    data: T
};

class arc_destruct<T> {
  let data: *libc::c_void;
  new(data: *libc::c_void) { self.data = data; }
  drop unsafe {
     let data: ~arc_data<T> = unsafe::reinterpret_cast(self.data);
     let new_count = rustrt::rust_atomic_decrement(&mut data.count);
     assert new_count >= 0;
     if new_count == 0 {
         // drop glue takes over.
     } else {
       unsafe::forget(data);
     }
  }
}

type arc<T: const send> = arc_destruct<T>;

/// Create an atomically reference counted wrapper.
fn arc<T: const send>(-data: T) -> arc<T> {
    let data = ~{mut count: 1, data: data};
    unsafe {
        let ptr = unsafe::transmute(data);
        arc_destruct(ptr)
    }
}

/**
 * Access the underlying data in an atomically reference counted
 * wrapper.
 */
fn get<T: const send>(rc: &arc<T>) -> &T {
    unsafe {
        let ptr: ~arc_data<T> = unsafe::reinterpret_cast((*rc).data);
        // Cast us back into the correct region
        let r = unsafe::reinterpret_cast(&ptr.data);
        unsafe::forget(ptr);
        return r;
    }
}

/**
 * Duplicate an atomically reference counted wrapper.
 *
 * The resulting two `arc` objects will point to the same underlying data
 * object. However, one of the `arc` objects can be sent to another task,
 * allowing them to share the underlying data.
 */
fn clone<T: const send>(rc: &arc<T>) -> arc<T> {
    unsafe {
        let ptr: ~arc_data<T> = unsafe::reinterpret_cast((*rc).data);
        let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
        assert new_count >= 2;
        unsafe::forget(ptr);
    }
    arc_destruct((*rc).data)
}

// An arc over mutable data that is protected by a lock.
type ex_data<T: send> = {lock: sys::lock_and_signal, mut data: T};
type exclusive<T: send> = arc_destruct<ex_data<T>>;

fn exclusive<T:send >(-data: T) -> exclusive<T> {
    let data = ~{mut count: 1, data: {lock: sys::lock_and_signal(),
                                      data: data}};
    unsafe {
        let ptr = unsafe::reinterpret_cast(data);
        unsafe::forget(data);
        arc_destruct(ptr)
    }
}

impl methods<T: send> for exclusive<T> {
    /// Duplicate an exclusive ARC. See arc::clone.
    fn clone() -> exclusive<T> {
        unsafe {
            // this makes me nervous...
            let ptr: ~arc_data<ex_data<T>> =
                  unsafe::reinterpret_cast(self.data);
            let new_count = rustrt::rust_atomic_increment(&mut ptr.count);
            assert new_count > 1;
            unsafe::forget(ptr);
        }
        arc_destruct(self.data)
    }

    /**
     * Access the underlying mutable data with mutual exclusion from other
     * tasks. The argument closure will be run with the mutex locked; all
     * other tasks wishing to access the data will block until the closure
     * finishes running.
     *
     * Currently, scheduling operations (i.e., yielding, receiving on a pipe,
     * accessing the provided condition variable) are prohibited while inside
     * the exclusive. Supporting that is a work in progress.
     *
     * The reason this function is 'unsafe' is because it is possible to
     * construct a circular reference among multiple ARCs by mutating the
     * underlying data. This creates potential for deadlock, but worse, this
     * will guarantee a memory leak of all involved ARCs. Using exclusive
     * ARCs inside of other ARCs is safe in absence of circular references.
     */
    unsafe fn with<U>(f: fn(sys::condition, x: &mut T) -> U) -> U {
        let ptr: ~arc_data<ex_data<T>> =
            unsafe::reinterpret_cast(self.data);
        assert ptr.count > 0;
        let r = {
            let rec: &ex_data<T> = &(*ptr).data;
            rec.lock.lock_cond(|c| f(c, &mut rec.data))
        };
        unsafe::forget(ptr);
        r
    }
}

#[cfg(test)]
mod tests {
    import comm::*;
    import future::extensions;

    #[test]
    fn manually_share_arc() {
        let v = ~[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let arc_v = arc::arc(v);

        let p = port();
        let c = chan(p);

        do task::spawn() {
            let p = port();
            c.send(chan(p));

            let arc_v = p.recv();

            let v = *arc::get::<~[int]>(&arc_v);
            assert v[3] == 4;
        };

        let c = p.recv();
        c.send(arc::clone(&arc_v));

        assert (*arc::get(&arc_v))[2] == 3;

        log(info, arc_v);
    }

    #[test]
    #[ignore] // this can probably infinite loop too.
    fn exclusive_arc() {
        let mut futures = ~[];

        let num_tasks = 10u;
        let count = 1000u;

        let total = exclusive(~mut 0u);

        for uint::range(0u, num_tasks) |_i| {
            let total = total.clone();
            futures += ~[future::spawn(|| {
                for uint::range(0u, count) |_i| {
                    do total.with |_cond, count| {
                        **count += 1u;
                    }
                }
            })];
        };

        for futures.each |f| { f.get() }

        do total.with |_cond, total| {
            assert **total == num_tasks * count
        };
    }
}
<|endoftext|>


<commit_before>use crate::error::{Error, Result};
use askama_escape::{Escaper, MarkupDisplay};
use serde::Serialize;

/// Serialize to JSON (requires `serde_json` feature)
///
/// ## Errors
///
/// This will panic if `S`'s implementation of `Serialize` decides to fail,
/// or if `T` contains a map with non-string keys.
pub fn yaml<E: Escaper, S: Serialize>(e: E, s: S) -> Result<MarkupDisplay<E, String>> {
    match serde_yaml::to_string(&s) {
        Ok(s) => Ok(MarkupDisplay::new_safe(s, e)),
        Err(e) => Err(Error::from(e)),
    }
}

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

    #[test]
    fn test_yaml() {
        assert_eq!(yaml(Html, true).unwrap().to_string(), "---\ntrue");
        assert_eq!(yaml(Html, "foo").unwrap().to_string(), "---\nfoo");
        assert_eq!(yaml(Html, &true).unwrap().to_string(), "---\ntrue");
        assert_eq!(yaml(Html, &"foo").unwrap().to_string(), "---\nfoo");
        assert_eq!(
            yaml(Html, &vec!["foo", "bar"]).unwrap().to_string(),
            "---\n- foo\n- bar"
        );
    }
}
<commit_msg>fix typo in docstring<commit_after>use crate::error::{Error, Result};
use askama_escape::{Escaper, MarkupDisplay};
use serde::Serialize;

/// Serialize to YAML (requires `serde_yaml` feature)
///
/// ## Errors
///
/// This will panic if `S`'s implementation of `Serialize` decides to fail,
/// or if `T` contains a map with non-string keys.
pub fn yaml<E: Escaper, S: Serialize>(e: E, s: S) -> Result<MarkupDisplay<E, String>> {
    match serde_yaml::to_string(&s) {
        Ok(s) => Ok(MarkupDisplay::new_safe(s, e)),
        Err(e) => Err(Error::from(e)),
    }
}

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

    #[test]
    fn test_yaml() {
        assert_eq!(yaml(Html, true).unwrap().to_string(), "---\ntrue");
        assert_eq!(yaml(Html, "foo").unwrap().to_string(), "---\nfoo");
        assert_eq!(yaml(Html, &true).unwrap().to_string(), "---\ntrue");
        assert_eq!(yaml(Html, &"foo").unwrap().to_string(), "---\nfoo");
        assert_eq!(
            yaml(Html, &vec!["foo", "bar"]).unwrap().to_string(),
            "---\n- foo\n- bar"
        );
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate genomic_fx;
extern crate bio;

use bio::utils::Strand;
use genomic_fx::TBuilder;
use test::Bencher;

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

    #[bench]
    fn tbuilder_single_exon_no_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(vec![(100, 10000)])
                .build()
        });
    }

    #[bench]
    fn tbuilder_single_exon_single_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(vec![(100, 10000)])
                .coding_coord(200, 9500)
                .build()
        });
    }

    #[bench]
    fn tbuilder_mult_exons_no_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(
                    vec![
                        (100, 300), (400, 500), (700, 1000),
                        (1100, 1300), (1400, 1500), (1700, 2000),
                        (2100, 2300), (2400, 2500), (2700, 3000),
                        (3000, 6000), (7000, 8000), (9000, 10000),
                    ])
                .build()
        });
    }

    #[bench]
    fn tbuilder_mult_exons_mult_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(
                    vec![
                        (100, 300), (400, 500), (700, 1000),
                        (1100, 1300), (1400, 1500), (1700, 2000),
                        (2100, 2300), (2400, 2500), (2700, 3000),
                        (3000, 6000), (7000, 8000), (9000, 10000),
                    ])
                .coding_coord(200, 9500)
                .build()
        });
    }
}
<commit_msg>Add benches for EBuilder<commit_after>#![feature(test)]

extern crate test;
extern crate genomic_fx;
extern crate bio;

use bio::utils::{Interval, Strand};
use test::Bencher;

use genomic_fx::{EBuilder, TBuilder, ExonFeature, ExonFeatureKind};
use ExonFeatureKind::*;

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

    #[bench]
    fn ebuilder_no_features(b: &mut Bencher) {
        b.iter(|| {
            EBuilder::new("chrT", 100, 500)
                .strand(Strand::Forward)
                .build()
        });
    }

    #[bench]
    fn ebuiler_mult_features(b: &mut Bencher) {
        b.iter(|| {
            EBuilder::new("chrT", 100, 300)
                .strand(Strand::Forward)
                .features(vec![
                    ExonFeature {
                        interval: Interval::new(100..150).unwrap(),
                        kind: UTR5,
                    },
                    ExonFeature {
                        interval: Interval::new(150..153).unwrap(),
                        kind: StartCodon { frame: Some(0) },
                    },
                    ExonFeature {
                        interval: Interval::new(150..300).unwrap(),
                        kind: CDS { frame: Some(0) },
                    },
                    ExonFeature {
                        interval: Interval::new(300..303).unwrap(),
                        kind: StopCodon { frame: Some(0) },
                    },
                    ExonFeature {
                        interval: Interval::new(300..500).unwrap(),
                        kind: UTR3,
                    },
                ])
                .build()
        });
    }

    #[bench]
    fn tbuilder_single_exon_no_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(vec![(100, 10000)])
                .build()
        });
    }

    #[bench]
    fn tbuilder_single_exon_single_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(vec![(100, 10000)])
                .coding_coord(200, 9500)
                .build()
        });
    }

    #[bench]
    fn tbuilder_mult_exons_no_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(
                    vec![
                        (100, 300), (400, 500), (700, 1000),
                        (1100, 1300), (1400, 1500), (1700, 2000),
                        (2100, 2300), (2400, 2500), (2700, 3000),
                        (3000, 6000), (7000, 8000), (9000, 10000),
                    ])
                .build()
        });
    }

    #[bench]
    fn tbuilder_mult_exons_mult_cds(b: &mut Bencher) {
        b.iter(|| {
            TBuilder::new("chrT", 100, 10000)
                .strand(Strand::Forward)
                .id("trx01")
                .exon_coords(
                    vec![
                        (100, 300), (400, 500), (700, 1000),
                        (1100, 1300), (1400, 1500), (1700, 2000),
                        (2100, 2300), (2400, 2500), (2700, 3000),
                        (3000, 6000), (7000, 8000), (9000, 10000),
                    ])
                .coding_coord(200, 9500)
                .build()
        });
    }
}
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! codegen_init {
    { [ $($class:tt)* ] } => {
        #[allow(non_snake_case)]
        #[no_mangle]
        pub extern "C" fn Init_native() {
            $crate::sys::check_version();

            $(
                codegen_class_binding!($class, $class);
            )*
        }
    }
}

#[macro_export]
macro_rules! codegen_class_binding {
    { $class:tt, {
        type: class,
        name: $name:tt,
        meta: { pub: $pub:tt, reopen: false },
        struct: (),
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;
        let def = $crate::ClassDefinition::new(cstr!(stringify!($name)));

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $name = transmute(def.class) };
    });

    { $class:tt, {
        type: class,
        name: $name:tt,
        meta: { pub: $pub:tt, reopen: true },
        struct: (),
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;
        let def = $crate::ClassDefinition::reopen(cstr!(stringify!($name)));

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $name = transmute(def.class) };
    });

    { $class:tt, {
        type: class,
        name: $cls:tt,
        meta: { pub: $pub:tt, reopen: $reopen:tt },
        struct: { $($struct:tt)* },
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;

        extern "C" fn __alloc__(_klass: $crate::sys::VALUE) -> $crate::sys::VALUE {
            $cls::__alloc_with__(None)
        }

        let def = $crate::ClassDefinition::wrapped(cstr!(stringify!($cls)), __alloc__);

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $cls = transmute(def.class) }
    });

}

#[macro_export]
macro_rules! codegen_define_method {
    ($def:tt, {
        type: class,
        name: $cls:tt,
        $($rest:tt)*
    }, {
        type: class_method,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: (),
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        use $crate::sys::{VALUE, SPRINTF_TO_S, Qnil, rb_raise};

        #[repr(C)]
        #[derive(Debug)]
        struct CallResult {
            error_klass: VALUE,
            value: VALUE
        }

        extern "C" fn __ruby_method__(_: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __rust_method__($($arg),*);

            if result.error_klass == unsafe { Qnil } {
                result.value
            } else {
                unsafe { rb_raise(result.error_klass, SPRINTF_TO_S, result.value) }
            }
        }

        #[inline]
        fn __rust_method__($($arg : $crate::sys::VALUE),*) -> CallResult {
            let checked = __checked_call__($($arg),*);

            match checked {
                Ok(val) => CallResult { error_klass: unsafe { Qnil }, value: $crate::ToRuby::to_ruby(val) },
                Err(err) => CallResult { error_klass: err.exception.inner(), value: err.message }
            }
        }

        #[inline]
        fn __checked_call__($($arg : $crate::sys::VALUE),*) -> Result<$($ret)*, $crate::ExceptionInfo> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            $(
                let $arg = match $crate::UncheckedValue::<$argty>::to_checked($arg) {
                    Ok(v) => v,
                    Err(e) => return Err($crate::ExceptionInfo::type_error(e))
                };
            )*

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            handle_exception! {
                $cls::$rust_name($($arg),*)
            }
        }

        let name = cstr!($($ruby_name)*);
        let method = __ruby_method__ as *const $crate::libc::c_void;
        let arity = method_arity!($($arg)*);

        $def.define_method($crate::MethodDefinition::class(name, method, arity));
    });

    ($def:tt, {
        type: class,
        name: $cls:tt,
        meta: $meta:tt,
        struct: $struct:tt,
        $($rest:tt)*
    }, {
        type: instance_method,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: { ownership: { $($ownership:tt)* }, name: $self:tt },
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        use $crate::sys::{VALUE, SPRINTF_TO_S, Qnil, rb_raise};

        #[repr(C)]
        struct CallResult {
            error_klass: VALUE,
            value: VALUE
        }

        extern "C" fn __ruby_method__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __rust_method__(rb_self, $($arg),*);

            if result.error_klass == unsafe { Qnil } {
                result.value
            } else {
                unsafe { rb_raise(result.error_klass, SPRINTF_TO_S, result.value) }
            }
        }

        #[inline]
        fn __rust_method__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> CallResult {
            let checked = __checked_call__(rb_self, $($arg),*);

            match checked {
                Ok(val) => CallResult { error_klass: unsafe { Qnil }, value: $crate::ToRuby::to_ruby(val) },
                Err(err) => CallResult { error_klass: err.exception.inner(), value: err.message }
            }
        }

        #[inline]
        fn __checked_call__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> Result<$($ret)*, $crate::ExceptionInfo> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            let rust_self = match $crate::UncheckedValue::<codegen_self_pointer_type! { struct: $struct, ownership: { $($ownership)* }, type: $cls }>::to_checked(rb_self) {
                Ok(v)  => v,
                Err(e) => return Err($crate::ExceptionInfo::with_message(e))
            };

            $(
                let $arg = match $crate::UncheckedValue::<$argty>::to_checked($arg) {
                    Ok(v) => v,
                    Err(e) => return Err($crate::ExceptionInfo::type_error(e))
                };
            )*

            let rust_self = rust_self.to_rust();

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            handle_exception! {
                rust_self.$rust_name($($arg),*)
            }
        }

        let name = cstr!($($ruby_name)*);
        let method = __ruby_method__ as *const $crate::libc::c_void;
        let arity = method_arity!($($arg)*);

        $def.define_method($crate::MethodDefinition::instance(name, method, arity))
    });

    ($def:tt, {
        type: class,
        name: $cls:tt,
        meta: $meta:tt,
        struct: $struct:tt,
        $($rest:tt)*
    }, {
        type: initializer,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: { ownership: {}, name: $self:tt },
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        impl $cls {
            pub fn new($($arg : $argty),*) -> $($ret)* {
                $cls::$rust_name(unsafe { $crate::sys::Qnil } , $($arg),*)
            }
        }

        extern "C" fn __initialize__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __checked_initialize__(rb_self $(, $arg)*);

            match result {
                Ok(rust_self) => {
                    let data = Box::new(rust_self);
                    unsafe { $crate::sys::Data_Set_Struct_Value(rb_self, ::std::mem::transmute(data)) };
                }
                Err(err) => { println!("TYPE ERROR: {:?}", err); }
            }

            rb_self
        }

        fn __checked_initialize__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> Result<$cls, String> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            $(
                let $arg = try!($crate::UncheckedValue::<$argty>::to_checked($arg));
            )*

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            Ok($cls::initialize(rb_self, $($arg),*))
        }

        let arity = method_arity!($($arg)*);
        let method = __initialize__ as *const $crate::libc::c_void;

        $def.define_method($crate::MethodDefinition::instance(cstr!($($ruby_name)*), method, arity));
    });
}

#[macro_export]
macro_rules! codegen_self_pointer_type {
    {
        struct: (),
        ownership: $ownership:tt,
        type: $type:tt
    } => {
        $type
    };

    {
        struct: $struct:tt,
        ownership: { $($ownership:tt)* },
        type: $type:tt
    } => {
        $($ownership)* $type
    };
}

#[macro_export]
macro_rules! method_arity {
  ( $($arg:tt)* ) => {
    { 0isize $(+ replace_expr!($arg 1isize))* }
  }
}

#[macro_export]
macro_rules! replace_expr {
    ($_t:tt $sub:expr) => {$sub};
}

#[doc(hidden)]
#[macro_export]
macro_rules! cstr {
    ($s:expr) => (
        concat!($s, "\0") as *const str as *const [::std::os::raw::c_char] as *const ::std::os::raw::c_char
    )
}

#[macro_export]
macro_rules! handle_exception {
    { $($body:tt)* } => {
        let hide_err = ::std::env::var("RUST_BACKTRACE").is_err();
        if hide_err {
            ::std::panic::set_hook(Box::new(|_| {
                // Silence
            }));
        }

        let res = ::std::panic::catch_unwind(|| {
            $($body)*
        });

        if hide_err {
            let _ = ::std::panic::take_hook();
        }

        res.map_err(|e| $crate::ExceptionInfo::from_any(e))
    }
}
<commit_msg>Raise TypeError on bad initialize<commit_after>#[macro_export]
macro_rules! codegen_init {
    { [ $($class:tt)* ] } => {
        #[allow(non_snake_case)]
        #[no_mangle]
        pub extern "C" fn Init_native() {
            $crate::sys::check_version();

            $(
                codegen_class_binding!($class, $class);
            )*
        }
    }
}

#[macro_export]
macro_rules! codegen_class_binding {
    { $class:tt, {
        type: class,
        name: $name:tt,
        meta: { pub: $pub:tt, reopen: false },
        struct: (),
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;
        let def = $crate::ClassDefinition::new(cstr!(stringify!($name)));

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $name = transmute(def.class) };
    });

    { $class:tt, {
        type: class,
        name: $name:tt,
        meta: { pub: $pub:tt, reopen: true },
        struct: (),
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;
        let def = $crate::ClassDefinition::reopen(cstr!(stringify!($name)));

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $name = transmute(def.class) };
    });

    { $class:tt, {
        type: class,
        name: $cls:tt,
        meta: { pub: $pub:tt, reopen: $reopen:tt },
        struct: { $($struct:tt)* },
        methods: [ $($method:tt)* ]
    } } => ({
        use ::std::mem::transmute;

        extern "C" fn __alloc__(_klass: $crate::sys::VALUE) -> $crate::sys::VALUE {
            $cls::__alloc_with__(None)
        }

        let def = $crate::ClassDefinition::wrapped(cstr!(stringify!($cls)), __alloc__);

        $(
            codegen_define_method!(def, $class, $method);
        )*

        unsafe { $cls = transmute(def.class) }
    });

}

#[macro_export]
macro_rules! codegen_define_method {
    ($def:tt, {
        type: class,
        name: $cls:tt,
        $($rest:tt)*
    }, {
        type: class_method,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: (),
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        use $crate::sys::{VALUE, SPRINTF_TO_S, Qnil, rb_raise};

        #[repr(C)]
        #[derive(Debug)]
        struct CallResult {
            error_klass: VALUE,
            value: VALUE
        }

        extern "C" fn __ruby_method__(_: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __rust_method__($($arg),*);

            if result.error_klass == unsafe { Qnil } {
                result.value
            } else {
                unsafe { rb_raise(result.error_klass, SPRINTF_TO_S, result.value) }
            }
        }

        #[inline]
        fn __rust_method__($($arg : $crate::sys::VALUE),*) -> CallResult {
            let checked = __checked_call__($($arg),*);

            match checked {
                Ok(val) => CallResult { error_klass: unsafe { Qnil }, value: $crate::ToRuby::to_ruby(val) },
                Err(err) => CallResult { error_klass: err.exception.inner(), value: err.message }
            }
        }

        #[inline]
        fn __checked_call__($($arg : $crate::sys::VALUE),*) -> Result<$($ret)*, $crate::ExceptionInfo> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            $(
                let $arg = match $crate::UncheckedValue::<$argty>::to_checked($arg) {
                    Ok(v) => v,
                    Err(e) => return Err($crate::ExceptionInfo::type_error(e))
                };
            )*

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            handle_exception! {
                $cls::$rust_name($($arg),*)
            }
        }

        let name = cstr!($($ruby_name)*);
        let method = __ruby_method__ as *const $crate::libc::c_void;
        let arity = method_arity!($($arg)*);

        $def.define_method($crate::MethodDefinition::class(name, method, arity));
    });

    ($def:tt, {
        type: class,
        name: $cls:tt,
        meta: $meta:tt,
        struct: $struct:tt,
        $($rest:tt)*
    }, {
        type: instance_method,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: { ownership: { $($ownership:tt)* }, name: $self:tt },
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        use $crate::sys::{VALUE, SPRINTF_TO_S, Qnil, rb_raise};

        #[repr(C)]
        struct CallResult {
            error_klass: VALUE,
            value: VALUE
        }

        extern "C" fn __ruby_method__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __rust_method__(rb_self, $($arg),*);

            if result.error_klass == unsafe { Qnil } {
                result.value
            } else {
                unsafe { rb_raise(result.error_klass, SPRINTF_TO_S, result.value) }
            }
        }

        #[inline]
        fn __rust_method__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> CallResult {
            let checked = __checked_call__(rb_self, $($arg),*);

            match checked {
                Ok(val) => CallResult { error_klass: unsafe { Qnil }, value: $crate::ToRuby::to_ruby(val) },
                Err(err) => CallResult { error_klass: err.exception.inner(), value: err.message }
            }
        }

        #[inline]
        fn __checked_call__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> Result<$($ret)*, $crate::ExceptionInfo> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            let rust_self = match $crate::UncheckedValue::<codegen_self_pointer_type! { struct: $struct, ownership: { $($ownership)* }, type: $cls }>::to_checked(rb_self) {
                Ok(v)  => v,
                Err(e) => return Err($crate::ExceptionInfo::with_message(e))
            };

            $(
                let $arg = match $crate::UncheckedValue::<$argty>::to_checked($arg) {
                    Ok(v) => v,
                    Err(e) => return Err($crate::ExceptionInfo::type_error(e))
                };
            )*

            let rust_self = rust_self.to_rust();

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            handle_exception! {
                rust_self.$rust_name($($arg),*)
            }
        }

        let name = cstr!($($ruby_name)*);
        let method = __ruby_method__ as *const $crate::libc::c_void;
        let arity = method_arity!($($arg)*);

        $def.define_method($crate::MethodDefinition::instance(name, method, arity))
    });

    ($def:tt, {
        type: class,
        name: $cls:tt,
        meta: $meta:tt,
        struct: $struct:tt,
        $($rest:tt)*
    }, {
        type: initializer,
        rust_name: $rust_name:tt,
        ruby_name: { $($ruby_name:tt)* },
        self: { ownership: {}, name: $self:tt },
        args: [ $($arg:tt : $argty:ty),* ],
        ret: { $($ret:tt)* },
        body: $body:tt
    }) => ({
        impl $cls {
            pub fn new($($arg : $argty),*) -> $($ret)* {
                $cls::$rust_name(unsafe { $crate::sys::Qnil } , $($arg),*)
            }
        }

        extern "C" fn __initialize__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> $crate::sys::VALUE {
            let result = __checked_initialize__(rb_self $(, $arg)*);

            match result {
                Ok(rust_self) => {
                    let data = Box::new(rust_self);
                    unsafe { $crate::sys::Data_Set_Struct_Value(rb_self, ::std::mem::transmute(data)) };
                }
                Err(e) => $crate::ExceptionInfo::type_error(e).raise()
            }

            rb_self
        }

        fn __checked_initialize__(rb_self: $crate::sys::VALUE, $($arg : $crate::sys::VALUE),*) -> Result<$cls, String> {
            #[allow(unused_imports)]
            use $crate::{ToRust};

            $(
                let $arg = try!($crate::UncheckedValue::<$argty>::to_checked($arg));
            )*

            $(
                let $arg = $crate::ToRust::to_rust($arg);
            )*

            Ok($cls::initialize(rb_self, $($arg),*))
        }

        let arity = method_arity!($($arg)*);
        let method = __initialize__ as *const $crate::libc::c_void;

        $def.define_method($crate::MethodDefinition::instance(cstr!($($ruby_name)*), method, arity));
    });
}

#[macro_export]
macro_rules! codegen_self_pointer_type {
    {
        struct: (),
        ownership: $ownership:tt,
        type: $type:tt
    } => {
        $type
    };

    {
        struct: $struct:tt,
        ownership: { $($ownership:tt)* },
        type: $type:tt
    } => {
        $($ownership)* $type
    };
}

#[macro_export]
macro_rules! method_arity {
  ( $($arg:tt)* ) => {
    { 0isize $(+ replace_expr!($arg 1isize))* }
  }
}

#[macro_export]
macro_rules! replace_expr {
    ($_t:tt $sub:expr) => {$sub};
}

#[doc(hidden)]
#[macro_export]
macro_rules! cstr {
    ($s:expr) => (
        concat!($s, "\0") as *const str as *const [::std::os::raw::c_char] as *const ::std::os::raw::c_char
    )
}

#[macro_export]
macro_rules! handle_exception {
    { $($body:tt)* } => {
        let hide_err = ::std::env::var("RUST_BACKTRACE").is_err();
        if hide_err {
            ::std::panic::set_hook(Box::new(|_| {
                // Silence
            }));
        }

        let res = ::std::panic::catch_unwind(|| {
            $($body)*
        });

        if hide_err {
            let _ = ::std::panic::take_hook();
        }

        res.map_err(|e| $crate::ExceptionInfo::from_any(e))
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use fnv::{FnvHashMap, FnvHashSet};
use interner::StringKey;
use serde::{Deserialize, Serialize};

#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields, rename_all = "lowercase")]
pub enum TypegenLanguage {
    Flow,
    TypeScript,
}

impl Default for TypegenLanguage {
    fn default() -> Self {
        Self::Flow
    }
}

#[derive(Debug, Serialize, Deserialize, Default)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct TypegenConfig {
    /// The desired output language, "flow" or "typescript".
    #[serde(default)]
    pub language: TypegenLanguage,

    /// # For Flow type generation
    /// When set, enum values are imported from a module with this suffix.
    /// For example, an enum Foo and this property set to ".test" would be
    /// imported from "Foo.test".
    /// Note: an empty string is allowed and different from not setting the
    /// value, in the example above it would just import from "Foo".
    pub enum_module_suffix: Option<String>,

    /// # For Flow type generation
    /// Generate enum files using Flow Enums instead of string unions except
    /// for the denylist or enums that contain lowercase first letters which
    /// are invalid Flow Enum values.
    #[serde(default)]
    pub legacy_enum_style: FnvHashSet<StringKey>,

    /// # For Flow type generation
    /// When set, generated input types will have the listed fields optional
    /// even if the schema defines them as required.
    #[serde(default)]
    pub optional_input_fields: Vec<StringKey>,

    /// # For Typescript type generation
    /// Whether to use the `import type` syntax introduced in Typescript
    /// version 3.8. This will prevent warnings from `importsNotUsedAsValues`.
    #[serde(default)]
    pub use_import_type_syntax: bool,

    /// A map from GraphQL scalar types to a custom JS type, example:
    /// { "Url": "String" }
    #[serde(default)]
    pub custom_scalar_types: FnvHashMap<StringKey, StringKey>,

    /// Require all GraphQL scalar types mapping to be defined, will throw
    /// if a GraphQL scalar type doesn't have a JS type
    #[serde(default)]
    pub require_custom_scalar_types: bool,
}
<commit_msg>use stable data structure in config<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use fnv::FnvBuildHasher;
use indexmap::{IndexMap, IndexSet};
use interner::StringKey;
use serde::{Deserialize, Serialize};

type FnvIndexMap<K, V> = IndexMap<K, V, FnvBuildHasher>;
type FnvIndexSet<T> = IndexSet<T, FnvBuildHasher>;

#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq)]
#[serde(deny_unknown_fields, rename_all = "lowercase")]
pub enum TypegenLanguage {
    Flow,
    TypeScript,
}

impl Default for TypegenLanguage {
    fn default() -> Self {
        Self::Flow
    }
}

#[derive(Debug, Serialize, Deserialize, Default)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct TypegenConfig {
    /// The desired output language, "flow" or "typescript".
    #[serde(default)]
    pub language: TypegenLanguage,

    /// # For Flow type generation
    /// When set, enum values are imported from a module with this suffix.
    /// For example, an enum Foo and this property set to ".test" would be
    /// imported from "Foo.test".
    /// Note: an empty string is allowed and different from not setting the
    /// value, in the example above it would just import from "Foo".
    pub enum_module_suffix: Option<String>,

    /// # For Flow type generation
    /// Generate enum files using Flow Enums instead of string unions except
    /// for the denylist or enums that contain lowercase first letters which
    /// are invalid Flow Enum values.
    #[serde(default)]
    pub legacy_enum_style: FnvIndexSet<StringKey>,

    /// # For Flow type generation
    /// When set, generated input types will have the listed fields optional
    /// even if the schema defines them as required.
    #[serde(default)]
    pub optional_input_fields: Vec<StringKey>,

    /// # For Typescript type generation
    /// Whether to use the `import type` syntax introduced in Typescript
    /// version 3.8. This will prevent warnings from `importsNotUsedAsValues`.
    #[serde(default)]
    pub use_import_type_syntax: bool,

    /// A map from GraphQL scalar types to a custom JS type, example:
    /// { "Url": "String" }
    #[serde(default)]
    pub custom_scalar_types: FnvIndexMap<StringKey, StringKey>,

    /// Require all GraphQL scalar types mapping to be defined, will throw
    /// if a GraphQL scalar type doesn't have a JS type
    #[serde(default)]
    pub require_custom_scalar_types: bool,
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use command::studio::enter::ARTIFACT_PATH_ENVVAR;
use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use hcore::os::process;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;
    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }
    if args[0] == OsString::from("rm") {
        return rm_container(&docker_cmd);
    }
    return enter(&docker_cmd, &args);
}

fn enter(docker_cmd: &Path, args: &Vec<OsString>) -> Result<()> {
    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd)
        && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![
        String::from("DEBUG"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_ORIGIN"),
        String::from("HAB_ORIGIN_KEYS"),
        String::from("HAB_STUDIO_BACKLINE_PKG"),
        String::from("HAB_STUDIO_NOSTUDIORC"),
        String::from("HAB_STUDIO_SUP"),
        String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
        String::from("http_proxy"),
        String::from("https_proxy"),
        String::from("RUST_LOG"),
    ];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    if container_exists(&docker_cmd) {
        return start_container(&docker_cmd);
    }
    run_container(&docker_cmd, args.to_vec(), volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images")
        .arg(&image_identifier(docker_cmd))
        .arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd
        .spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image, result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
        && (stderr.as_str().contains("Mounts denied")
            || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn container_exists(docker_cmd: &Path) -> bool {
    let name = path_to_name().expect("Unable to load name from path");
    let cmd_args: Vec<OsString> = vec!["container".into(), "inspect".into(), name.into()];
    let output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");
    if output.status.success() {
        return true;
    }
    return false;
}

#[cfg(not(target_os = "windows"))]
fn path_to_name() -> Result<String> {
    let mut cwd = env::current_dir()?;
    cwd = cwd.strip_prefix("/")?.to_path_buf();
    let pathstr = cwd.to_str().expect("Path to be parseable");
    return Ok(pathstr.replace("/", "--"));
}

#[cfg(target_os = "windows")]
fn path_to_name() -> Result<String> {
    let cwd = env::current_dir()?;
    let pathstr = cwd.to_str().expect("Path to be parseable");
    let stripped_path: String = pathstr.chars().skip(3).take(pathstr.len() - 3).collect();
    return Ok(stripped_path.replace("\\", "--"));
}

fn rm_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["rm".into()];
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn run_container<I, J, S, T>(
    docker_cmd: &Path,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    let name = path_to_name()?;
    cmd_args.push("--name".into());
    cmd_args.push(name.into());
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR, opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn start_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["start".into()];
    cmd_args.push("--interactive".into());
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<commit_msg>studio build and run should not be interactive and remove the container on exit<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use command::studio::enter::ARTIFACT_PATH_ENVVAR;
use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use hcore::os::process;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;
    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }
    if args[0] == OsString::from("rm") {
        return rm_container(&docker_cmd);
    }
    return enter(&docker_cmd, &args);
}

fn enter(docker_cmd: &Path, args: &Vec<OsString>) -> Result<()> {
    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd)
        && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![
        String::from("DEBUG"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_ORIGIN"),
        String::from("HAB_ORIGIN_KEYS"),
        String::from("HAB_STUDIO_BACKLINE_PKG"),
        String::from("HAB_STUDIO_NOSTUDIORC"),
        String::from("HAB_STUDIO_SUP"),
        String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
        String::from("http_proxy"),
        String::from("https_proxy"),
        String::from("RUST_LOG"),
    ];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    if container_exists(&docker_cmd) {
        return start_container(&docker_cmd);
    }
    run_container(&docker_cmd, args.to_vec(), volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images")
        .arg(&image_identifier(docker_cmd))
        .arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd
        .spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image, result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
        && (stderr.as_str().contains("Mounts denied")
            || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn container_exists(docker_cmd: &Path) -> bool {
    let name = path_to_name().expect("Unable to load name from path");
    let cmd_args: Vec<OsString> = vec!["container".into(), "inspect".into(), name.into()];
    let output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");
    if output.status.success() {
        return true;
    }
    return false;
}

#[cfg(not(target_os = "windows"))]
fn path_to_name() -> Result<String> {
    let mut cwd = env::current_dir()?;
    cwd = cwd.strip_prefix("/")?.to_path_buf();
    let pathstr = cwd.to_str().expect("Path to be parseable");
    return Ok(pathstr.replace("/", "--"));
}

#[cfg(target_os = "windows")]
fn path_to_name() -> Result<String> {
    let cwd = env::current_dir()?;
    let pathstr = cwd.to_str().expect("Path to be parseable");
    let stripped_path: String = pathstr.chars().skip(3).take(pathstr.len() - 3).collect();
    return Ok(stripped_path.replace("\\", "--"));
}

fn rm_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["rm".into()];
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn run_container<I, J, S, T>(
    docker_cmd: &Path,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    let name = path_to_name()?;
    cmd_args.push("--name".into());
    cmd_args.push(name.into());
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") | Some("run") => {
            cmd_args.push("--rm".into());
        }
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR, opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn start_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["start".into()];
    cmd_args.push("--interactive".into());
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<|endoftext|>


<commit_before>use nalgebra::{Point3, Scalar};
use num::Float;
use num::traits::FloatConst;
use std::cmp;
use std::marker::PhantomData;

use super::generate::{ConjointPointGenerator, IndexPolygonGenerator, PolygonGenerator};
use super::primitive::{Polygon, Triangle, Quad};

#[derive(Clone)]
pub struct UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    nu: usize, // Meridians.
    nv: usize, // Parallels.
    phantom: PhantomData<T>,
}

impl<T> UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    pub fn with_unit_radius(nu: usize, nv: usize) -> Self {
        let nu = cmp::max(3, nu);
        let nv = cmp::max(2, nv);
        UVSphere {
            nu: nu,
            nv: nv,
            phantom: PhantomData,
        }
    }

    fn point(&self, u: usize, v: usize) -> Point3<T> {
        let u = (T::from(u).unwrap() / T::from(self.nu).unwrap()) * T::PI() * (T::one() + T::one());
        let v = (T::from(v).unwrap() / T::from(self.nv).unwrap()) * T::PI();
        Point3::new(u.cos() * v.sin(), u.sin() * v.sin(), v.cos())
    }

    fn index_point(&self, u: usize, v: usize) -> usize {
        if v == 0 {
            0
        }
        else if v == self.nv {
            ((self.nv - 1) * self.nu) + 1
        }
        else {
            ((v - 1) * self.nu) + (u % self.nu) + 1
        }
    }

    fn map_polygon_index(&self, index: usize) -> (usize, usize) {
        (index % self.nu, index / self.nv)
    }
}

impl<T> ConjointPointGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Point3<T>;

    fn conjoint_point(&self, index: usize) -> Self::Output {
        if index == 0 {
            self.point(0, 0)
        }
        else if index == self.conjoint_point_count() - 1 {
            self.point(0, self.nv)
        }
        else {
            let index = index - 1;
            self.point(index % self.nu, (index / self.nv) + 1)
        }
    }

    fn conjoint_point_count(&self) -> usize {
        (self.nv - 1) * self.nu + 2
    }
}

impl<T> PolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<Point3<T>>;

    fn polygon(&self, index: usize) -> Self::Output {
        let (u, v) = self.map_polygon_index(index);

        // Generate the points at the requested meridian and parallel. The
        // upper and lower bounds of (u, v) are always used, so generate them
        // in advance (`low` and `high`). Emit triangles at the poles,
        // otherwise quads.
        let low = self.point(u, v);
        let high = self.point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.point(u, v + 1),
                high,
                self.point(u + 1, v),
            ))
        }
    }

    fn polygon_count(&self) -> usize {
        self.nu * self.nv
    }
}

impl<T> IndexPolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<usize>;

    fn index_polygon(&self, index: usize) -> <Self as IndexPolygonGenerator>::Output {
        let (u, v) = self.map_polygon_index(index);

        let low = self.index_point(u, v);
        let high = self.index_point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.index_point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.index_point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.index_point(u, v + 1),
                high,
                self.index_point(u + 1, v),
            ))
        }
    }
}
<commit_msg>Fix index to uv mapping.<commit_after>use nalgebra::{Point3, Scalar};
use num::Float;
use num::traits::FloatConst;
use std::cmp;
use std::marker::PhantomData;

use super::generate::{ConjointPointGenerator, IndexPolygonGenerator, PolygonGenerator};
use super::primitive::{Polygon, Triangle, Quad};

#[derive(Clone)]
pub struct UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    nu: usize, // Meridians.
    nv: usize, // Parallels.
    phantom: PhantomData<T>,
}

impl<T> UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    pub fn with_unit_radius(nu: usize, nv: usize) -> Self {
        let nu = cmp::max(3, nu);
        let nv = cmp::max(2, nv);
        UVSphere {
            nu: nu,
            nv: nv,
            phantom: PhantomData,
        }
    }

    fn point(&self, u: usize, v: usize) -> Point3<T> {
        let u = (T::from(u).unwrap() / T::from(self.nu).unwrap()) * T::PI() * (T::one() + T::one());
        let v = (T::from(v).unwrap() / T::from(self.nv).unwrap()) * T::PI();
        Point3::new(u.cos() * v.sin(), u.sin() * v.sin(), v.cos())
    }

    fn index_point(&self, u: usize, v: usize) -> usize {
        if v == 0 {
            0
        }
        else if v == self.nv {
            ((self.nv - 1) * self.nu) + 1
        }
        else {
            ((v - 1) * self.nu) + (u % self.nu) + 1
        }
    }

    fn map_polygon_index(&self, index: usize) -> (usize, usize) {
        (index % self.nu, index / self.nu)
    }
}

impl<T> ConjointPointGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Point3<T>;

    fn conjoint_point(&self, index: usize) -> Self::Output {
        if index == 0 {
            self.point(0, 0)
        }
        else if index == self.conjoint_point_count() - 1 {
            self.point(0, self.nv)
        }
        else {
            let index = index - 1;
            self.point(index % self.nu, (index / self.nv) + 1)
        }
    }

    fn conjoint_point_count(&self) -> usize {
        (self.nv - 1) * self.nu + 2
    }
}

impl<T> PolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<Point3<T>>;

    fn polygon(&self, index: usize) -> Self::Output {
        let (u, v) = self.map_polygon_index(index);

        // Generate the points at the requested meridian and parallel. The
        // upper and lower bounds of (u, v) are always used, so generate them
        // in advance (`low` and `high`). Emit triangles at the poles,
        // otherwise quads.
        let low = self.point(u, v);
        let high = self.point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.point(u, v + 1),
                high,
                self.point(u + 1, v),
            ))
        }
    }

    fn polygon_count(&self) -> usize {
        self.nu * self.nv
    }
}

impl<T> IndexPolygonGenerator for UVSphere<T>
where
    T: Float + FloatConst + Scalar,
{
    type Output = Polygon<usize>;

    fn index_polygon(&self, index: usize) -> <Self as IndexPolygonGenerator>::Output {
        let (u, v) = self.map_polygon_index(index);

        let low = self.index_point(u, v);
        let high = self.index_point(u + 1, v + 1);
        if v == 0 {
            Polygon::Triangle(Triangle::new(low, self.index_point(u, v + 1), high))
        }
        else if v == self.nv - 1 {
            Polygon::Triangle(Triangle::new(high, self.index_point(u + 1, v), low))
        }
        else {
            Polygon::Quad(Quad::new(
                low,
                self.index_point(u, v + 1),
                high,
                self.index_point(u + 1, v),
            ))
        }
    }
}
<|endoftext|>