<commit_before>use std::collections::HashMap;

use futures::stream::Stream;
use futures::Future;
use futures;

use tokio_core::io::ReadHalf;
use tokio_core::io::Io;

use solicit::http::session::StreamState;
use solicit::http::frame::*;
use solicit::http::StaticHeader;
use solicit::http::StreamId;
use solicit::http::HttpError;
use solicit::http::WindowSize;
use solicit::http::HttpScheme;
use solicit::http::INITIAL_CONNECTION_WINDOW_SIZE;
use solicit::http::connection::HttpConnection;

use futures_misc::*;

use solicit_misc::*;
use solicit_async::*;

#[derive(Debug)]
pub enum HttpStreamPartContent {
    Headers(Vec<StaticHeader>),
    Data(Vec<u8>),
}

pub struct HttpStreamPart {
    pub content: HttpStreamPartContent,
    /// END_STREAM
    pub last: bool,
}

impl HttpStreamPart {
    pub fn last_headers(header: Vec<StaticHeader>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Headers(header),
            last: true,
        }
    }

    pub fn intermediate_headers(headers: Vec<StaticHeader>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Headers(headers),
            last: false,
        }
    }

    pub fn intermediate_data(data: Vec<u8>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Data(data),
            last: false,
        }
    }

    pub fn last_data(data: Vec<u8>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Data(data),
            last: true,
        }
    }
}

pub type HttpStreamStream = Box<Stream<Item=HttpStreamPart, Error=HttpError>>;
pub type HttpStreamStreamSend = Box<Stream<Item=HttpStreamPart, Error=HttpError> + Send>;

pub trait HttpService: Send + 'static {
    fn new_request(&mut self, headers: Vec<StaticHeader>, req: HttpStreamStreamSend) -> HttpStreamStreamSend;
}

pub struct GrpcHttpStreamCommon {
    pub state: StreamState,
    pub out_window_size: WindowSize,
    pub in_window_size: WindowSize,
}

impl GrpcHttpStreamCommon {
    pub fn new() -> GrpcHttpStreamCommon {
        GrpcHttpStreamCommon {
            state: StreamState::Open,
            in_window_size: WindowSize::new(INITIAL_CONNECTION_WINDOW_SIZE),
            out_window_size: WindowSize::new(INITIAL_CONNECTION_WINDOW_SIZE),
        }
    }
}

pub trait GrpcHttpStream {
    fn common(&self) -> &GrpcHttpStreamCommon;
    fn common_mut(&mut self) -> &mut GrpcHttpStreamCommon;
    fn new_data_chunk(&mut self, data: &[u8], last: bool);
    fn close_remote(&mut self);
}

pub struct LoopInnerCommon<S>
    where S : GrpcHttpStream,
{
    pub conn: HttpConnection,
    pub streams: HashMap<StreamId, S>,
}

impl<S> LoopInnerCommon<S>
    where S : GrpcHttpStream,
{
    pub fn new(scheme: HttpScheme) -> LoopInnerCommon<S> {
        LoopInnerCommon {
            conn: HttpConnection::new(scheme),
            streams: HashMap::new(),
        }
    }

    pub fn get_stream_mut(&mut self, stream_id: StreamId) -> Option<&mut S> {
        self.streams.get_mut(&stream_id)
    }

    pub fn remove_stream(&mut self, stream_id: StreamId) -> Option<S> {
        self.streams.remove(&stream_id)
    }

    pub fn remove_stream_if_closed(&mut self, stream_id: StreamId) {
        if self.get_stream_mut(stream_id).expect("unknown stream").common().state == StreamState::Closed {
            self.remove_stream(stream_id);
        }
    }
}

pub trait HttpReadLoopInner : 'static {
    type LoopHttpStream : GrpcHttpStream;

    fn common(&mut self) -> &mut LoopInnerCommon<Self::LoopHttpStream>;

    /// Send a frame back to the network
    /// Must not be data frame
    fn send_frame<R : FrameIR>(&mut self, frame: R);
    fn out_window_increased(&mut self, stream_id: Option<StreamId>);

    fn process_headers_frame(&mut self, frame: HeadersFrame);

    fn process_settings_global(&mut self, _frame: SettingsFrame) {
        // TODO: apply settings
        // TODO: send ack
    }

    fn process_stream_window_update_frame(&mut self, frame: WindowUpdateFrame) {
        {
            let stream = self.common().get_stream_mut(frame.get_stream_id()).expect("stream not found");
            stream.common_mut().out_window_size.try_increase(frame.increment())
                .expect("failed to increment stream window");
        }
        self.out_window_increased(Some(frame.get_stream_id()));
    }

    fn process_conn_window_update(&mut self, frame: WindowUpdateFrame) {
        self.common().conn.out_window_size.try_increase(frame.increment())
            .expect("failed to increment conn window");
        self.out_window_increased(None);
    }

    fn process_rst_stream_frame(&mut self, _frame: RstStreamFrame) {
    }

    fn process_data_frame(&mut self, frame: DataFrame) {
        let stream_id = frame.get_stream_id();

        self.common().conn.decrease_in_window(frame.payload_len())
            .expect("failed to decrease conn win");

        let increment_conn =
            if self.common().conn.in_window_size() < INITIAL_CONNECTION_WINDOW_SIZE / 2 {
                let increment = INITIAL_CONNECTION_WINDOW_SIZE as u32;
                self.common().conn.in_window_size.try_increase(increment).expect("failed to increase");

                Some(increment)
            } else {
                None
            };

        let increment_stream = {
            let stream = self.common().get_stream_mut(frame.get_stream_id()).expect("stream not found");

            stream.common_mut().in_window_size.try_decrease(frame.payload_len() as i32)
                .expect("failed to decrease stream win");

            let increment_stream =
                if stream.common_mut().in_window_size.size() < INITIAL_CONNECTION_WINDOW_SIZE / 2 {
                    let increment = INITIAL_CONNECTION_WINDOW_SIZE as u32;
                    stream.common_mut().in_window_size.try_increase(increment).expect("failed to increase");

                    Some(increment)
                } else {
                    None
                };

            stream.new_data_chunk(&frame.data.as_ref(), frame.is_end_of_stream());

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }

            // TODO: drop stream if closed on both ends

            increment_stream
        };

        if let Some(increment_conn) = increment_conn {
            self.send_frame(WindowUpdateFrame::for_connection(increment_conn));
        }

        if let Some(increment_stream) = increment_stream {
            self.send_frame(WindowUpdateFrame::for_stream(stream_id, increment_stream));
        }
    }

    fn process_ping(&mut self, frame: PingFrame) {
        if frame.is_ack() {

        } else {
            self.send_frame(PingFrame::new_ack(frame.opaque_data()));
        }
    }

    fn process_conn_frame(&mut self, frame: HttpFrameConn) {
        match frame {
            HttpFrameConn::Settings(f) => self.process_settings_global(f),
            HttpFrameConn::Ping(f) => self.process_ping(f),
            HttpFrameConn::Goaway(_f) => panic!("TODO"),
            HttpFrameConn::WindowUpdate(f) => self.process_conn_window_update(f),
        }
    }

    fn process_stream_frame(&mut self, frame: HttpFrameStream) {
        let stream_id = frame.get_stream_id();
        let end_of_stream = frame.is_end_of_stream();
        match frame {
            HttpFrameStream::Data(data) => self.process_data_frame(data),
            HttpFrameStream::Headers(headers) => self.process_headers_frame(headers),
            HttpFrameStream::RstStream(rst) => self.process_rst_stream_frame(rst),
            HttpFrameStream::WindowUpdate(window_update) => self.process_stream_window_update_frame(window_update),
        };
        if end_of_stream {
            self.close_remote(stream_id);
        }
    }

    fn process_raw_frame(&mut self, raw_frame: RawFrame) {
        let frame = HttpFrameClassified::from_raw(&raw_frame).unwrap();
        debug!("received frame: {:?}", frame);
        match frame {
            HttpFrameClassified::Conn(f) => self.process_conn_frame(f),
            HttpFrameClassified::Stream(f) => self.process_stream_frame(f),
            HttpFrameClassified::Unknown(_f) => {},
        }
    }

    fn close_remote(&mut self, stream_id: StreamId) {
        debug!("close remote: {}", stream_id);

        let remove = {
            let mut stream = self.common().get_stream_mut(stream_id).expect("stream not found");
            stream.close_remote();
            stream.common().state == StreamState::Closed
        };
        if remove {
            self.common().remove_stream(stream_id);
        }
    }
}

pub struct HttpReadLoopData<I, N>
    where
        I : Io + 'static,
        N : HttpReadLoopInner,
{
    pub read: ReadHalf<I>,
    pub inner: TaskRcMut<N>,
}

impl<I, N> HttpReadLoopData<I, N>
    where
        I : Io + Send + 'static,
        N : HttpReadLoopInner,
{
    /// Recv a frame from the network
    fn recv_raw_frame(self) -> HttpFuture<(Self, RawFrame<'static>)> {
        let HttpReadLoopData { read, inner } = self;
        Box::new(recv_raw_frame(read)
            .map(|(read, frame)| (HttpReadLoopData { read: read, inner: inner }, frame))
            .map_err(HttpError::from))
    }

    fn read_process_frame(self) -> HttpFuture<Self> {
        Box::new(self.recv_raw_frame()
            .and_then(move |(lp, frame)| lp.process_raw_frame(frame)))
    }

    pub fn run(self) -> HttpFuture<()> {
        let stream = stream_repeat(());

        let future = stream.fold(self, |lp, _| {
            lp.read_process_frame()
        });

        Box::new(future.map(|_| ()))
    }

    fn process_raw_frame(self, frame: RawFrame) -> HttpFuture<Self> {
        self.inner.with(move |inner| {
            inner.process_raw_frame(frame);
        });
        Box::new(futures::finished(self))
    }

}
<commit_msg>Properly handle WINDOW_UPDATE on closed streams<commit_after>use std::collections::HashMap;

use futures::stream::Stream;
use futures::Future;
use futures;

use tokio_core::io::ReadHalf;
use tokio_core::io::Io;

use solicit::http::session::StreamState;
use solicit::http::frame::*;
use solicit::http::StaticHeader;
use solicit::http::StreamId;
use solicit::http::HttpError;
use solicit::http::WindowSize;
use solicit::http::HttpScheme;
use solicit::http::INITIAL_CONNECTION_WINDOW_SIZE;
use solicit::http::connection::HttpConnection;

use futures_misc::*;

use solicit_misc::*;
use solicit_async::*;

#[derive(Debug)]
pub enum HttpStreamPartContent {
    Headers(Vec<StaticHeader>),
    Data(Vec<u8>),
}

pub struct HttpStreamPart {
    pub content: HttpStreamPartContent,
    /// END_STREAM
    pub last: bool,
}

impl HttpStreamPart {
    pub fn last_headers(header: Vec<StaticHeader>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Headers(header),
            last: true,
        }
    }

    pub fn intermediate_headers(headers: Vec<StaticHeader>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Headers(headers),
            last: false,
        }
    }

    pub fn intermediate_data(data: Vec<u8>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Data(data),
            last: false,
        }
    }

    pub fn last_data(data: Vec<u8>) -> Self {
        HttpStreamPart {
            content: HttpStreamPartContent::Data(data),
            last: true,
        }
    }
}

pub type HttpStreamStream = Box<Stream<Item=HttpStreamPart, Error=HttpError>>;
pub type HttpStreamStreamSend = Box<Stream<Item=HttpStreamPart, Error=HttpError> + Send>;

pub trait HttpService: Send + 'static {
    fn new_request(&mut self, headers: Vec<StaticHeader>, req: HttpStreamStreamSend) -> HttpStreamStreamSend;
}

pub struct GrpcHttpStreamCommon {
    pub state: StreamState,
    pub out_window_size: WindowSize,
    pub in_window_size: WindowSize,
}

impl GrpcHttpStreamCommon {
    pub fn new() -> GrpcHttpStreamCommon {
        GrpcHttpStreamCommon {
            state: StreamState::Open,
            in_window_size: WindowSize::new(INITIAL_CONNECTION_WINDOW_SIZE),
            out_window_size: WindowSize::new(INITIAL_CONNECTION_WINDOW_SIZE),
        }
    }
}

pub trait GrpcHttpStream {
    fn common(&self) -> &GrpcHttpStreamCommon;
    fn common_mut(&mut self) -> &mut GrpcHttpStreamCommon;
    fn new_data_chunk(&mut self, data: &[u8], last: bool);
    fn close_remote(&mut self);
}

pub struct LoopInnerCommon<S>
    where S : GrpcHttpStream,
{
    pub conn: HttpConnection,
    pub streams: HashMap<StreamId, S>,
}

impl<S> LoopInnerCommon<S>
    where S : GrpcHttpStream,
{
    pub fn new(scheme: HttpScheme) -> LoopInnerCommon<S> {
        LoopInnerCommon {
            conn: HttpConnection::new(scheme),
            streams: HashMap::new(),
        }
    }

    pub fn get_stream_mut(&mut self, stream_id: StreamId) -> Option<&mut S> {
        self.streams.get_mut(&stream_id)
    }

    pub fn remove_stream(&mut self, stream_id: StreamId) -> Option<S> {
        self.streams.remove(&stream_id)
    }

    pub fn remove_stream_if_closed(&mut self, stream_id: StreamId) {
        if self.get_stream_mut(stream_id).expect("unknown stream").common().state == StreamState::Closed {
            self.remove_stream(stream_id);
        }
    }
}

pub trait HttpReadLoopInner : 'static {
    type LoopHttpStream : GrpcHttpStream;

    fn common(&mut self) -> &mut LoopInnerCommon<Self::LoopHttpStream>;

    /// Send a frame back to the network
    /// Must not be data frame
    fn send_frame<R : FrameIR>(&mut self, frame: R);
    fn out_window_increased(&mut self, stream_id: Option<StreamId>);

    fn process_headers_frame(&mut self, frame: HeadersFrame);

    fn process_settings_global(&mut self, _frame: SettingsFrame) {
        // TODO: apply settings
        // TODO: send ack
    }

    fn process_stream_window_update_frame(&mut self, frame: WindowUpdateFrame) {
        {
            match self.common().get_stream_mut(frame.get_stream_id()) {
                Some(stream) => {
                    stream.common_mut().out_window_size.try_increase(frame.increment())
                        .expect("failed to increment stream window");
                }
                None => {
                    // 6.9
                    // WINDOW_UPDATE can be sent by a peer that has sent a frame bearing the
                    // END_STREAM flag.  This means that a receiver could receive a
                    // WINDOW_UPDATE frame on a "half-closed (remote)" or "closed" stream.
                    // A receiver MUST NOT treat this as an error (see Section 5.1).
                    debug!("WINDOW_UPDATE of unknown stream: {}", frame.get_stream_id());
                }
            }
        }
        self.out_window_increased(Some(frame.get_stream_id()));
    }

    fn process_conn_window_update(&mut self, frame: WindowUpdateFrame) {
        self.common().conn.out_window_size.try_increase(frame.increment())
            .expect("failed to increment conn window");
        self.out_window_increased(None);
    }

    fn process_rst_stream_frame(&mut self, _frame: RstStreamFrame) {
    }

    fn process_data_frame(&mut self, frame: DataFrame) {
        let stream_id = frame.get_stream_id();

        self.common().conn.decrease_in_window(frame.payload_len())
            .expect("failed to decrease conn win");

        let increment_conn =
            if self.common().conn.in_window_size() < INITIAL_CONNECTION_WINDOW_SIZE / 2 {
                let increment = INITIAL_CONNECTION_WINDOW_SIZE as u32;
                self.common().conn.in_window_size.try_increase(increment).expect("failed to increase");

                Some(increment)
            } else {
                None
            };

        let increment_stream = {
            let stream = self.common().get_stream_mut(frame.get_stream_id())
                .expect(&format!("stream not found: {}", frame.get_stream_id()));

            stream.common_mut().in_window_size.try_decrease(frame.payload_len() as i32)
                .expect("failed to decrease stream win");

            let increment_stream =
                if stream.common_mut().in_window_size.size() < INITIAL_CONNECTION_WINDOW_SIZE / 2 {
                    let increment = INITIAL_CONNECTION_WINDOW_SIZE as u32;
                    stream.common_mut().in_window_size.try_increase(increment).expect("failed to increase");

                    Some(increment)
                } else {
                    None
                };

            stream.new_data_chunk(&frame.data.as_ref(), frame.is_end_of_stream());

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }

            // TODO: drop stream if closed on both ends

            increment_stream
        };

        if let Some(increment_conn) = increment_conn {
            self.send_frame(WindowUpdateFrame::for_connection(increment_conn));
        }

        if let Some(increment_stream) = increment_stream {
            self.send_frame(WindowUpdateFrame::for_stream(stream_id, increment_stream));
        }
    }

    fn process_ping(&mut self, frame: PingFrame) {
        if frame.is_ack() {

        } else {
            self.send_frame(PingFrame::new_ack(frame.opaque_data()));
        }
    }

    fn process_conn_frame(&mut self, frame: HttpFrameConn) {
        match frame {
            HttpFrameConn::Settings(f) => self.process_settings_global(f),
            HttpFrameConn::Ping(f) => self.process_ping(f),
            HttpFrameConn::Goaway(_f) => panic!("TODO"),
            HttpFrameConn::WindowUpdate(f) => self.process_conn_window_update(f),
        }
    }

    fn process_stream_frame(&mut self, frame: HttpFrameStream) {
        let stream_id = frame.get_stream_id();
        let end_of_stream = frame.is_end_of_stream();
        match frame {
            HttpFrameStream::Data(data) => self.process_data_frame(data),
            HttpFrameStream::Headers(headers) => self.process_headers_frame(headers),
            HttpFrameStream::RstStream(rst) => self.process_rst_stream_frame(rst),
            HttpFrameStream::WindowUpdate(window_update) => self.process_stream_window_update_frame(window_update),
        };
        if end_of_stream {
            self.close_remote(stream_id);
        }
    }

    fn process_raw_frame(&mut self, raw_frame: RawFrame) {
        let frame = HttpFrameClassified::from_raw(&raw_frame).unwrap();
        debug!("received frame: {:?}", frame);
        match frame {
            HttpFrameClassified::Conn(f) => self.process_conn_frame(f),
            HttpFrameClassified::Stream(f) => self.process_stream_frame(f),
            HttpFrameClassified::Unknown(_f) => {},
        }
    }

    fn close_remote(&mut self, stream_id: StreamId) {
        debug!("close remote: {}", stream_id);

        let remove = {
            let mut stream = self.common().get_stream_mut(stream_id)
                .expect(&format!("stream not found: {}", stream_id));
            stream.close_remote();
            stream.common().state == StreamState::Closed
        };
        if remove {
            self.common().remove_stream(stream_id);
        }
    }
}

pub struct HttpReadLoopData<I, N>
    where
        I : Io + 'static,
        N : HttpReadLoopInner,
{
    pub read: ReadHalf<I>,
    pub inner: TaskRcMut<N>,
}

impl<I, N> HttpReadLoopData<I, N>
    where
        I : Io + Send + 'static,
        N : HttpReadLoopInner,
{
    /// Recv a frame from the network
    fn recv_raw_frame(self) -> HttpFuture<(Self, RawFrame<'static>)> {
        let HttpReadLoopData { read, inner } = self;
        Box::new(recv_raw_frame(read)
            .map(|(read, frame)| (HttpReadLoopData { read: read, inner: inner }, frame))
            .map_err(HttpError::from))
    }

    fn read_process_frame(self) -> HttpFuture<Self> {
        Box::new(self.recv_raw_frame()
            .and_then(move |(lp, frame)| lp.process_raw_frame(frame)))
    }

    pub fn run(self) -> HttpFuture<()> {
        let stream = stream_repeat(());

        let future = stream.fold(self, |lp, _| {
            lp.read_process_frame()
        });

        Box::new(future.map(|_| ()))
    }

    fn process_raw_frame(self, frame: RawFrame) -> HttpFuture<Self> {
        self.inner.with(move |inner| {
            inner.process_raw_frame(frame);
        });
        Box::new(futures::finished(self))
    }

}
<|endoftext|>


<commit_before>extern crate openssl;
extern crate openssl_verify;

use std::io;
use std::fmt;
use std::error;
use self::openssl::crypto::pkey::PKey;
use self::openssl::crypto::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslContext, SslMethod, SSL_VERIFY_PEER, IntoSsl, SSL_OP_NO_SSLV2,
                         SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION, MidHandshakeSslStream};
use self::openssl::x509::X509;
use self::openssl_verify::verify_callback;

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 {
    cert: X509,
    pkey: PKey,
    chain: Vec<X509>,
}

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 {
            cert: parsed.cert,
            pkey: parsed.pkey,
            chain: parsed.chain.into_iter().collect(),
        })
    }
}

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::Failure(e) => HandshakeError::Failure(Error(e)),
            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())
    }
}

fn ctx() -> Result<SslContext, Error> {
    let mut ctx = try!(SslContext::new(SslMethod::Sslv23));
    ctx.set_options(SSL_OP_NO_SSLV2 | SSL_OP_NO_SSLV3 | SSL_OP_NO_COMPRESSION);
    try!(ctx.set_default_verify_paths());
    try!(ctx.set_cipher_list("ALL!EXPORT!EXPORT40!EXPORT56!aNULL!LOW!RC4@STRENGTH"));
    Ok(ctx)
}

pub struct ClientBuilder(SslContext);

impl ClientBuilder {
    pub fn new() -> Result<ClientBuilder, Error> {
        ctx().map(ClientBuilder)
    }

    pub fn handshake<S>(&self,
                        domain: &str,
                        stream: S)
                        -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ssl = try!(self.0.into_ssl());
        try!(ssl.set_hostname(domain));
        let domain = domain.to_owned();
        ssl.set_verify_callback(SSL_VERIFY_PEER, move |p, x| verify_callback(&domain, p, x));

        let s = try!(ssl::SslStream::connect(ssl, stream));
        Ok(TlsStream(s))
    }
}

pub struct ServerBuilder(SslContext);

impl ServerBuilder {
    pub fn new(pkcs12: Pkcs12) -> Result<ServerBuilder, Error> {
        let mut ctx = try!(ctx());
        try!(ctx.set_certificate(&pkcs12.cert));
        try!(ctx.set_private_key(&pkcs12.pkey));
        try!(ctx.check_private_key());
        for cert in &pkcs12.chain {
            try!(ctx.add_extra_chain_cert(&cert));
        }
        Ok(ServerBuilder(ctx))
    }

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

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 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
    }
}
<commit_msg>Remove cipher suite and options setup from openssl<commit_after>extern crate openssl;
extern crate openssl_verify;

use std::io;
use std::fmt;
use std::error;
use self::openssl::crypto::pkey::PKey;
use self::openssl::crypto::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslContext, SslMethod, SSL_VERIFY_PEER, IntoSsl, SSL_OP_NO_SSLV2,
                         SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION, MidHandshakeSslStream};
use self::openssl::x509::X509;
use self::openssl_verify::verify_callback;

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 {
    cert: X509,
    pkey: PKey,
    chain: Vec<X509>,
}

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 {
            cert: parsed.cert,
            pkey: parsed.pkey,
            chain: parsed.chain.into_iter().collect(),
        })
    }
}

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::Failure(e) => HandshakeError::Failure(Error(e)),
            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())
    }
}

fn ctx() -> Result<SslContext, Error> {
    let mut ctx = try!(SslContext::new(SslMethod::Sslv23));
    try!(ctx.set_default_verify_paths());
    Ok(ctx)
}

pub struct ClientBuilder(SslContext);

impl ClientBuilder {
    pub fn new() -> Result<ClientBuilder, Error> {
        ctx().map(ClientBuilder)
    }

    pub fn handshake<S>(&self,
                        domain: &str,
                        stream: S)
                        -> Result<TlsStream<S>, HandshakeError<S>>
        where S: io::Read + io::Write
    {
        let mut ssl = try!(self.0.into_ssl());
        try!(ssl.set_hostname(domain));
        let domain = domain.to_owned();
        ssl.set_verify_callback(SSL_VERIFY_PEER, move |p, x| verify_callback(&domain, p, x));

        let s = try!(ssl::SslStream::connect(ssl, stream));
        Ok(TlsStream(s))
    }
}

pub struct ServerBuilder(SslContext);

impl ServerBuilder {
    pub fn new(pkcs12: Pkcs12) -> Result<ServerBuilder, Error> {
        let mut ctx = try!(ctx());
        try!(ctx.set_certificate(&pkcs12.cert));
        try!(ctx.set_private_key(&pkcs12.pkey));
        try!(ctx.check_private_key());
        for cert in &pkcs12.chain {
            try!(ctx.add_extra_chain_cert(&cert));
        }
        Ok(ServerBuilder(ctx))
    }

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

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 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
    }
}
<|endoftext|>


<commit_before>//! Simple basic model of the C-AST
//!
//! This is the skeleton of the final decompiled output. However, further
//! stages maybe added to
//! make the decompiled output easier to read and add more sugaring.

use petgraph::graph::{Graph, NodeIndex};
use petgraph::EdgeDirection;
use std::default;
use std::iter;

pub type Ty = String;

#[derive(Clone, Debug)]
enum CASTNode {
    FunctionHeader,
    Conditional,
    Declaration(Ty),
    Loop,
    Goto(String),
    Label(String),
    Break,
    ExpressionNode(Expr),
    Var(String),
}

#[derive(Clone, Debug)]
pub enum Expr {
    Eq,
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Shr,
    Shl,
}

#[derive(Clone, Debug)]
enum ASTEdge {
    StatementOrd(u64),
    OpOrd(u8),
}

#[derive(Clone, Debug)]
pub struct CAST {
    ast: Graph<CASTNode, ASTEdge>,
    eidx: u64,
}

impl default::Default for CAST {
    fn default() -> CAST {
        let mut ast = CAST {
            ast: Graph::new(),
            eidx: 0,
        };
        ast.ast.add_node(CASTNode::FunctionHeader);
        ast
    }
}

lazy_static! {
    static ref FN_HEAD: NodeIndex = NodeIndex::new(0);
}

fn format_with_indent(string: &str, depth: usize) -> String {
    iter::repeat(' ').take(depth).collect::<String>() + string
}

impl CAST {
    fn next_edge_idx(&mut self) -> u64 {
        self.eidx += 1;
        self.eidx - 1
    }

    pub fn expr(&mut self, operator: Expr, operands: &[NodeIndex]) -> NodeIndex {
        let operator = self.ast.add_node(CASTNode::ExpressionNode(operator));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, operator, ASTEdge::StatementOrd(idx));
        for (i, n) in operands.iter().enumerate() {
            self.ast.add_edge(operator, *n, ASTEdge::OpOrd(i as u8));
        }
        operator
    }

    pub fn new_loop(&mut self, loop_header: NodeIndex, loop_body: NodeIndex) -> NodeIndex {
        let e1 = self.ast.find_edge(*FN_HEAD, loop_header).expect("This cannot be `None`");
        let e2 = self.ast.find_edge(*FN_HEAD, loop_body).expect("This cannot be `None`");
        let loop_h = self.ast.add_node(CASTNode::Loop);
        let idx = if let ASTEdge::StatementOrd(idx) = self.ast[e1] {
            idx
        } else {
            panic!()
        };
        self.ast.remove_edge(e2);
        self.ast.remove_edge(e1);
        self.ast.add_edge(*FN_HEAD, loop_h, ASTEdge::StatementOrd(idx));
        self.ast.add_edge(loop_h, loop_header, ASTEdge::OpOrd(0));
        self.ast.add_edge(loop_h, loop_body, ASTEdge::OpOrd(1));
        loop_h
    }

    pub fn new_conditional(&mut self,
                           condition: NodeIndex,
                           body: NodeIndex,
                           else_condition: Option<NodeIndex>)
                           -> NodeIndex {
        let e1 = self.ast.find_edge(*FN_HEAD, condition).expect("This cannot be `None`");
        let e2 = self.ast.find_edge(*FN_HEAD, body).expect("This cannot be `None`");
        let e3 = if let Some(else_c) = else_condition {
            self.ast.find_edge(*FN_HEAD, else_c)
        } else {
            None
        };
        let idx = if let ASTEdge::StatementOrd(idx) = self.ast[e1] {
            idx
        } else {
            panic!()
        };

        let conditional = self.ast.add_node(CASTNode::Conditional);
        self.ast.remove_edge(e1);
        self.ast.remove_edge(e2);
        if let Some(e3) = e3 {
            self.ast.remove_edge(e3);
        }

        self.ast.add_edge(*FN_HEAD, conditional, ASTEdge::StatementOrd(idx));
        self.ast.add_edge(conditional, condition, ASTEdge::OpOrd(0));
        self.ast.add_edge(conditional, body, ASTEdge::OpOrd(1));
        if let Some(else_condition) = else_condition {
            self.ast.add_edge(conditional, else_condition, ASTEdge::OpOrd(2));
        }
        conditional
    }

    pub fn goto(&mut self, label: String) {
        let goto_n = self.ast.add_node(CASTNode::Goto(label));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, goto_n, ASTEdge::StatementOrd(idx));
    }

    pub fn label(&mut self, label: String) {
        let label = self.ast.add_node(CASTNode::Label(label));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, label, ASTEdge::StatementOrd(idx));
    }

    pub fn insert_break(&mut self, after: NodeIndex) {
        let break_n = self.ast.add_node(CASTNode::Break);
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, break_n, ASTEdge::StatementOrd(idx));
    }

    fn emit_c(&self, node: &NodeIndex, indent: usize) -> String {
        match self.ast[*node] {
            CASTNode::FunctionHeader => unimplemented!(),
            CASTNode::Conditional => {
                // Get the arguments -> condition, body, else branch.
                let arg1: NodeIndex = panic!();
                let arg2: NodeIndex = panic!();
                let arg3: Option<NodeIndex> = None;
                let condition = format!("{}\n{}",
                                        format_with_indent("if ( {} ) {{", indent),
                                        self.emit_c(&arg1, indent));
                let true_body = self.emit_c(&arg2, indent + 1);
                let false_body = if let Some(arg3) = arg3 {
                    let fbody = self.emit_c(&arg3, indent + 1);
                    if let CASTNode::Conditional = self.ast[arg3] {
                        // Else-If case
                        format_with_indent("}} else ", indent) + &fbody
                    } else {
                        // Plain Else case.
                        format_with_indent("}} else {{", indent) + &fbody
                    }
                } else {
                    "".to_owned()
                };

                format!("{}\n{}\n{}\n{}",
                        condition,
                        true_body,
                        false_body,
                        format_with_indent("}", indent))
            }
            CASTNode::Declaration(ref ty) => {
                unimplemented!()
            }
            CASTNode::Loop => {
                unimplemented!()
            }
            CASTNode::Goto(ref label) => {
                unimplemented!()
            }
            CASTNode::Label(ref label) => {
                unimplemented!()
            }
            CASTNode::Break => {
                unimplemented!()
            }
            CASTNode::ExpressionNode(ref expr) => {
                unimplemented!()
            }
            CASTNode::Var(ref ident) => {
                unimplemented!()
            }
        }
    }

    pub fn print(&self) -> String {
        // Get all the edges from the function header.
        // Take them in order
        // Traverse the subtree and print out accordingly.
        let mut result = String::new();
        let mut edges = self.ast
                            .edges_directed(*FN_HEAD, EdgeDirection::Outgoing)
                            .collect::<Vec<_>>();
        edges.sort_by(|a, b| {
            let idx1 = if let ASTEdge::StatementOrd(idx) = *a.1 {
                idx
            } else {
                u64::max_value()
            };

            let idx2 = if let ASTEdge::StatementOrd(idx) = *b.1 {
                idx
            } else {
                u64::max_value()
            };

            idx1.cmp(&idx2)
        });

        for &(ref node, _) in &edges {
            result.push_str(&self.emit_c(node, 0));
        }

        result
    }
}

#[cfg(test)]
mod test {

    #[test]
    fn c_ast_basic_test() {
        unimplemented!()
    }
}
<commit_msg>Completes phase 1 of simple C-AST<commit_after>//! Simple basic model of the C-AST
//!
//! This is the skeleton of the final decompiled output. However, further
//! stages maybe added to
//! make the decompiled output easier to read and add more sugaring.

use petgraph::graph::{Graph, NodeIndex};
use petgraph::EdgeDirection;
use std::default;
use std::iter;

pub type Ty = String;

#[derive(Clone, Debug)]
enum CASTNode {
    FunctionHeader(String),
    Conditional,
    Declaration(Ty),
    Loop,
    Goto(String),
    Label(String),
    Break,
    ExpressionNode(Expr),
    Var(String),
}

#[derive(Clone, Debug)]
pub enum Expr {
    Eq,
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Shr,
    Shl,
    Xor,
    Or,
    And,
    Not,
    Neg,
    Gt,
    GtEq,
    Lt,
    LtEq,
    Cmp,
}

#[derive(Clone, Debug)]
enum ASTEdge {
    StatementOrd(u64),
    OpOrd(u8),
}

#[derive(Clone, Debug)]
pub struct CAST {
    ast: Graph<CASTNode, ASTEdge>,
    eidx: u64,
}

impl default::Default for CAST {
    fn default() -> CAST {
        let mut ast = CAST {
            ast: Graph::new(),
            eidx: 0,
        };
        ast.ast.add_node(CASTNode::FunctionHeader("Unknown".to_owned()));
        ast
    }
}

lazy_static! {
    static ref FN_HEAD: NodeIndex = NodeIndex::new(0);
}

fn format_with_indent(string: &str, depth: usize) -> String {
    iter::repeat(' ').take(depth * 4).collect::<String>() + string
}

impl CAST {
    fn next_edge_idx(&mut self) -> u64 {
        self.eidx += 1;
        self.eidx - 1
    }

    pub fn expr(&mut self, operator: Expr, operands: &[NodeIndex]) -> NodeIndex {
        let operator = self.ast.add_node(CASTNode::ExpressionNode(operator));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, operator, ASTEdge::StatementOrd(idx));
        for (i, n) in operands.iter().enumerate() {
            if let Some(edx) = self.ast.find_edge(*FN_HEAD, *n) {
                self.ast.remove_edge(edx);
            }
            self.ast.add_edge(operator, *n, ASTEdge::OpOrd(i as u8));
        }
        operator
    }

    pub fn new_loop(&mut self, loop_header: NodeIndex, loop_body: NodeIndex) -> NodeIndex {
        let e1 = self.ast.find_edge(*FN_HEAD, loop_header).expect("This cannot be `None`");
        let e2 = self.ast.find_edge(*FN_HEAD, loop_body).expect("This cannot be `None`");
        let loop_h = self.ast.add_node(CASTNode::Loop);
        let idx = if let ASTEdge::StatementOrd(idx) = self.ast[e1] {
            idx
        } else {
            panic!()
        };
        self.ast.remove_edge(e2);
        self.ast.remove_edge(e1);
        self.ast.add_edge(*FN_HEAD, loop_h, ASTEdge::StatementOrd(idx));
        self.ast.add_edge(loop_h, loop_header, ASTEdge::OpOrd(0));
        self.ast.add_edge(loop_h, loop_body, ASTEdge::OpOrd(1));
        loop_h
    }

    pub fn new_conditional(&mut self,
                           condition: NodeIndex,
                           body: NodeIndex,
                           else_condition: Option<NodeIndex>)
                           -> NodeIndex {
        let e3 = if let Some(else_c) = else_condition {
            self.ast.find_edge(*FN_HEAD, else_c)
        } else {
            None
        };

        if let Some(e3) = e3 {
            self.ast.remove_edge(e3);
        }

        let conditional = self.ast.add_node(CASTNode::Conditional);
        let e1 = self.ast.find_edge(*FN_HEAD, condition).expect("This cannot be `None`");
        let idx = if let ASTEdge::StatementOrd(idx) = self.ast[e1] {
            idx
        } else {
            panic!()
        };
        self.ast.remove_edge(e1);

        let e2 = self.ast.find_edge(*FN_HEAD, body).expect("This cannot be `None`");
        self.ast.remove_edge(e2);

        self.ast.add_edge(*FN_HEAD, conditional, ASTEdge::StatementOrd(idx));
        self.ast.add_edge(conditional, condition, ASTEdge::OpOrd(0));
        self.ast.add_edge(conditional, body, ASTEdge::OpOrd(1));
        if let Some(else_condition) = else_condition {
            self.ast.add_edge(conditional, else_condition, ASTEdge::OpOrd(2));
        }
        conditional
    }

    pub fn goto(&mut self, label: String) {
        let goto_n = self.ast.add_node(CASTNode::Goto(label));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, goto_n, ASTEdge::StatementOrd(idx));
    }

    pub fn label(&mut self, label: String) {
        let label = self.ast.add_node(CASTNode::Label(label));
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, label, ASTEdge::StatementOrd(idx));
    }

    pub fn insert_break(&mut self, after: NodeIndex) {
        let break_n = self.ast.add_node(CASTNode::Break);
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, break_n, ASTEdge::StatementOrd(idx));
    }

    pub fn declare_vars(&mut self, ty: String, vars: &[String]) -> Vec<NodeIndex> {
        let decl = self.ast.add_node(CASTNode::Declaration(ty));
        let mut var_decls = Vec::new();
        for (i, v) in vars.iter().enumerate() {
            let vn = self.ast.add_node(CASTNode::Var(v.clone()));
            self.ast.add_edge(decl, vn, ASTEdge::OpOrd(i as u8));
            var_decls.push(vn);
        }
        let idx = self.next_edge_idx();
        self.ast.add_edge(*FN_HEAD, decl, ASTEdge::StatementOrd(idx));
        var_decls
    }

    fn emit_c(&self, node: &NodeIndex, indent: usize) -> String {
        match self.ast[*node] {
            CASTNode::FunctionHeader(_) => unimplemented!(),
            CASTNode::Conditional => {
                // Get the arguments -> condition, body, else branch.
                let mut args = self.ast
                                   .edges_directed(*node, EdgeDirection::Outgoing)
                                   .collect::<Vec<_>>();
                args.sort_by(|a, b| {
                    let idx1 = if let ASTEdge::OpOrd(idx) = *a.1 {
                        idx
                    } else {
                        u8::max_value()
                    };

                    let idx2 = if let ASTEdge::OpOrd(idx) = *b.1 {
                        idx
                    } else {
                        u8::max_value()
                    };

                    idx1.cmp(&idx2)
                });
                let arg1 = args[0].0;
                let arg2 = args[1].0;
                let arg3 = args.get(2).map(|x| x.0);
                let condition = format!("{} {} {{\n",
                                        format_with_indent("if", indent),
                                        self.emit_c(&arg1, 0));
                let true_body = self.emit_c(&arg2, indent + 1);
                let false_body = if let Some(arg3) = arg3 {
                    let fbody = self.emit_c(&arg3, indent + 1);
                    if let CASTNode::Conditional = self.ast[arg3] {
                        // Else-If case
                        format_with_indent("\n}} else ", indent) + &fbody
                    } else {
                        // Plain Else case.
                        format_with_indent("\n}} else {{", indent) + &fbody
                    }
                } else {
                    "".to_owned()
                };

                format!("{}{}{}\n{}",
                        condition,
                        true_body,
                        false_body,
                        format_with_indent("}", indent))
            }
            CASTNode::Declaration(ref ty) => {
                let mut ty = format_with_indent(ty, indent);
                let mut vars = String::new();
                for (ref op, _) in self.ast.edges_directed(*node, EdgeDirection::Outgoing) {
                    if let CASTNode::Var(ref name) = self.ast[*op] {
                        if vars.is_empty() {
                            vars = vars + name;
                        } else {
                            vars = vars + ", " + name;
                        }
                    }
                }
                format!("{} {};", ty, vars)
            }
            CASTNode::Loop => {
                // Get the arguments -> loop header/check condition, loop body.
                let mut args = self.ast
                                   .edges_directed(*node, EdgeDirection::Outgoing)
                                   .collect::<Vec<_>>();
                args.sort_by(|a, b| {
                    let idx1 = if let ASTEdge::OpOrd(idx) = *a.1 {
                        idx
                    } else {
                        u8::max_value()
                    };
                    let idx2 = if let ASTEdge::OpOrd(idx) = *b.1 {
                        idx
                    } else {
                        u8::max_value()
                    };
                    idx1.cmp(&idx2)
                });
                let loop_header = self.emit_c(&args[0].0, 0);
                let loop_body = self.emit_c(&args[1].0, indent + 1);
                format!("{} ({}) {{\n{}\n{}",
                        format_with_indent("while", indent),
                        loop_header,
                        loop_body,
                        format_with_indent("}", indent))
            }
            CASTNode::Goto(ref label) => {
                format_with_indent(&format!("goto {}", label), indent)
            }
            CASTNode::Label(ref label) => {
                format_with_indent(&format!("label {}:", label), indent)
            }
            CASTNode::Break => {
                format_with_indent("break;", indent)
            }
            CASTNode::ExpressionNode(ref expr) => {
                let mut operands = self.ast
                                       .edges_directed(*node, EdgeDirection::Outgoing)
                                       .collect::<Vec<_>>();
                operands.sort_by(|a, b| {
                    let idx1 = if let ASTEdge::OpOrd(idx) = *a.1 {
                        idx
                    } else {
                        u8::max_value()
                    };
                    let idx2 = if let ASTEdge::OpOrd(idx) = *b.1 {
                        idx
                    } else {
                        u8::max_value()
                    };
                    idx1.cmp(&idx2)
                });

                let op_str = operands.iter().map(|x| self.emit_c(&x.0, 0)).collect::<Vec<_>>();
                match *expr {
                    Expr::Eq => format!("{} = {}",
                                        format_with_indent(&op_str[0], indent),
                                        &op_str[1]),
                    Expr::Add => format!("({} + {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Sub => format!("({} - {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Mul => format!("({} * {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Div => format!("({} / {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Mod => format!("({} % {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Shr => format!("({} >> {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Shl => format!("({} << {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Xor => format!("({} ^ {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Or => format!("({} | {})",
                                        format_with_indent(&op_str[0], indent),
                                        &op_str[1]),
                    Expr::And => format!("({} & {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                    Expr::Not => format!("(!{})", format_with_indent(&op_str[0], indent)),
                    Expr::Neg => format!("(~{})", format_with_indent(&op_str[0], indent)),
                    Expr::Gt => format!("({} > {})",
                                        format_with_indent(&op_str[0], indent),
                                        &op_str[1]),
                    Expr::GtEq => format!("({} >= {})",
                                          format_with_indent(&op_str[0], indent),
                                          &op_str[1]),
                    Expr::Lt => format!("({} < {})",
                                        format_with_indent(&op_str[0], indent),
                                        &op_str[1]),
                    Expr::LtEq => format!("({} <= {})",
                                          format_with_indent(&op_str[0], indent),
                                          &op_str[1]),
                    Expr::Cmp => format!("({} == {})",
                                         format_with_indent(&op_str[0], indent),
                                         &op_str[1]),
                }
            }
            CASTNode::Var(ref ident) => {
                ident.clone()
            }
        }
    }

    pub fn print(&self) -> String {
        // Get all the edges from the function header.
        // Take them in order
        // Traverse the subtree and print out accordingly.
        let mut result = String::new();
        if let CASTNode::FunctionHeader(ref named) = self.ast[*FN_HEAD] {
            result.push_str(&format!("fn {} {{\n", named));
        }
        let mut edges = self.ast
                            .edges_directed(*FN_HEAD, EdgeDirection::Outgoing)
                            .collect::<Vec<_>>();
        edges.sort_by(|a, b| {
            let idx1 = if let ASTEdge::StatementOrd(idx) = *a.1 {
                idx
            } else {
                u64::max_value()
            };

            let idx2 = if let ASTEdge::StatementOrd(idx) = *b.1 {
                idx
            } else {
                u64::max_value()
            };

            idx1.cmp(&idx2)
        });

        for &(ref node, _) in &edges {
            println!("In loop");
            result.push_str(&(self.emit_c(node, 1) + "\n"));
        }

        result.push_str("}");
        result
    }
}

#[cfg(test)]
mod test {

    use super::*;

    #[test]
    fn c_ast_basic_test() {
        let mut c_ast = CAST::default();
        let vars = c_ast.declare_vars("int".to_owned(), &["i".to_owned(), "j".to_owned()]);
        let cmp = c_ast.expr(Expr::Cmp, &vars);
        let increment = c_ast.expr(Expr::Add, &vars);
        let assignment = c_ast.expr(Expr::Eq, &[vars[0], increment]);
        c_ast.new_conditional(cmp, assignment, None);
        println!("{}", c_ast.print());
    }
}
<|endoftext|>


<commit_before>use std::boxed::Box;
use std::collections::HashMap as Interior;
use std::hash::Hash;
use std::cmp::Eq;
use std::borrow::Borrow;
use std::cell::UnsafeCell;

pub struct HashMap<K: Eq + Hash, V> {
  data: UnsafeCell<Interior<K, Box<V>>>,
}

impl<K: Eq + Hash, V> HashMap<K, V> {
  pub fn new() -> HashMap<K, V> {
    HashMap { data: UnsafeCell::new(Interior::new()) }
  }
  pub fn get_default<F>(&self, key: K, default_function: F) -> &V
    where F: FnOnce() -> V
  {
    let entry = unsafe { (*self.data.get()).entry(key) };
    (*entry.or_insert_with(|| Box::new(default_function()))).borrow()
  }
}
<commit_msg>Protect against doing a silly unsafe thing with insert only map<commit_after>use std::boxed::Box;
use std::collections::HashMap as Interior;
use std::hash::Hash;
use std::cmp::Eq;
use std::borrow::Borrow;
use std::cell::UnsafeCell;

pub struct HashMap<K: Eq + Hash, V> {
  data: UnsafeCell<Interior<K, Box<V>>>,
  inserting: bool,
}

impl<K: Eq + Hash, V> HashMap<K, V> {
  pub fn new() -> HashMap<K, V> {
    HashMap { data: UnsafeCell::new(Interior::new()), inserting: false }
  }
  pub fn get_default<F>(&self, key: K, default_function: F) -> &V
    where F: FnOnce() -> V
  {
    assert!(!self.inserting, "Attempt to call get_default on a insert_only::HashMap within the default_function callback for another get_default on the same map");
    self.inserting = true;
    let entry = unsafe { (*self.data.get()).entry(key) };
    let result = (*entry.or_insert_with(|| Box::new(default_function()))).borrow();
    self.inserting = false;
    result
  }
}
<|endoftext|>


<commit_before>use std::mem;

pub struct Profile {
    units: usize,
    steps: usize,
    time: f64,
    time_step: f64,
    data: Vec<f64>,
}

impl Profile {
    pub fn new(units: usize, time_step: f64) -> Profile {
        Profile {
            units: units,
            steps: 0,
            time: 0.0,
            time_step: time_step,
            data: Vec::with_capacity(units),
        }
    }

    pub fn accumulate(&mut self, unit: usize, time: f64, time_step: f64, data: &[f64]) {
        debug_assert!(unit < self.units);
        debug_assert!(time >= self.time);

        let (t1, t2) = (self.time, time);
        let (d1, d2) = (self.time_step, time_step);

        let s2 = data.len();
        let s1 = ((t2 - t1 + (s2 as f64) * d2) / d1).ceil() as usize;

        if s1 > self.steps {
            self.data.extend(vec![0.0; (s1 - self.steps) * self.units]);
            self.steps = s1;
        }

        let mut j1 = ((t2 - t1) / d1) as usize;
        let mut j2 = 0;

        macro_rules! add(
            ($weight:expr) => (self.data[j1 * self.units + unit] += $weight * data[j2]);
        );

        while j1 < s1 && j2 < s2 {
            let l1 = t1 + (j1 as f64) * d1;
            let l2 = t2 + (j2 as f64) * d2;

            let r1 = l1 + d1;
            let r2 = l2 + d2;

            if l1 < l2 {
                if r2 < r1 {
                    add!(1.0);
                    j2 += 1;
                } else {
                    add!((r1 - l2) / d2);
                    j1 += 1;
                }
            } else {
                if r1 < r2 {
                    add!(d1 / d2);
                    j1 += 1;
                } else {
                    add!((r2 - l1) / d2);
                    j2 += 1;
                }
            }
        }
    }

    pub fn discharge(&mut self, time: f64) -> Vec<f64> {
        debug_assert!(time >= self.time);
        let steps = ((time - self.time) / self.time_step).floor() as usize;
        debug_assert!(self.steps >= steps);

        let tail = self.data[(steps * self.units)..].to_vec();
        let mut data = mem::replace(&mut self.data, tail);
        data.truncate(steps * self.units);

        self.steps = self.steps - steps;
        self.time = (time / self.time_step).floor() * self.time_step;

        data
    }
}

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

    #[test]
    fn accumulate_padding() {
        let mut profile = Profile::new(2, 0.5);

        profile.accumulate(0, 4.0, 1.0, &[]);
        assert_eq!(profile.steps, 8);
        assert_eq!(&profile.data, &vec![0.0; 2 * 8]);

        profile.accumulate(0, 2.5, 1.0, &[]);
        assert_eq!(profile.steps, 8);
        assert_eq!(&profile.data, &vec![0.0; 2 * 8]);

        profile.accumulate(0, 6.5, 1.0, &[]);
        assert_eq!(profile.steps, 13);
        assert_eq!(&profile.data, &vec![0.0; 2 * 13]);

        profile.accumulate(0, 6.55, 1.0, &[]);
        assert_eq!(profile.steps, 14);
        assert_eq!(&profile.data, &vec![0.0; 2 * 14]);
    }

    #[test]
    fn accumulate_synchronous() {
        let mut profile = Profile::new(2, 1.0);

        profile.accumulate(0, 1.0, 1.0, &[1.0, 2.0]);
        assert_eq!(profile.steps, 3);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           1.0, 0.0,
           2.0, 0.0,
        ]);

        profile.accumulate(0, 1.0, 1.0, &[1.0, 2.0, 3.0]);
        assert_eq!(profile.steps, 4);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           2.0, 0.0,
           4.0, 0.0,
           3.0, 0.0,
        ]);
    }

    #[test]
    fn accumulate_asynchronous() {
        let mut profile = Profile::new(2, 1.0);

        profile.accumulate(1, 1.5, 1.0, &[1.0, 2.0, 3.0]);
        assert_eq!(profile.steps, 5);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           0.0, 0.5,
           0.0, 1.5,
           0.0, 2.5,
           0.0, 1.5,
        ]);

        profile.accumulate(0, 0.5, 0.25, &[1.0, 2.0, 3.0, 1.0, 3.0]);
        assert_eq!(profile.steps, 5);
        assert_eq!(&profile.data, &vec![
           3.0, 0.0,
           7.0, 0.5,
           0.0, 1.5,
           0.0, 2.5,
           0.0, 1.5,
        ]);

        profile.accumulate(0, 1.25, 1.0, &[1.0, 2.0, 3.0, 0.0, 4.0]);
        assert_eq!(profile.steps, 7);
        assert_eq!(&profile.data, &vec![
           3.00, 0.0,
           7.75, 0.5,
           1.75, 1.5,
           2.75, 2.5,
           0.75, 1.5,
           3.00, 0.0,
           1.00, 0.0,
        ]);
    }

    #[test]
    fn discharge() {
        let mut profile = Profile::new(2, 1.0);
        profile.accumulate(0, 0.0, 1.0, &vec![42.0; 42]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(0.0), vec![]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(0.75), vec![]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(1.0), vec![42.0, 0.0]);
        assert_eq!(profile.time, 1.0);
        assert_eq!(profile.steps, 41);

        assert_eq!(profile.discharge(1.5), vec![]);
        assert_eq!(profile.time, 1.0);
        assert_eq!(profile.steps, 41);

        assert_eq!(profile.discharge(3.5), vec![42.0, 0.0, 42.0, 0.0]);
        assert_eq!(profile.time, 3.0);
        assert_eq!(profile.steps, 39);
    }
}
<commit_msg>profile: make discharge produce a Profile<commit_after>use std::mem;

pub struct Profile {
    units: usize,
    steps: usize,
    time: f64,
    time_step: f64,
    data: Vec<f64>,
}

impl Profile {
    pub fn new(units: usize, time_step: f64) -> Profile {
        Profile {
            units: units,
            steps: 0,
            time: 0.0,
            time_step: time_step,
            data: Vec::with_capacity(units),
        }
    }

    pub fn accumulate(&mut self, unit: usize, time: f64, time_step: f64, data: &[f64]) {
        debug_assert!(unit < self.units);
        debug_assert!(time >= self.time);

        let (t1, t2) = (self.time, time);
        let (d1, d2) = (self.time_step, time_step);

        let s2 = data.len();
        let s1 = ((t2 - t1 + (s2 as f64) * d2) / d1).ceil() as usize;

        if s1 > self.steps {
            self.data.extend(vec![0.0; (s1 - self.steps) * self.units]);
            self.steps = s1;
        }

        let mut j1 = ((t2 - t1) / d1) as usize;
        let mut j2 = 0;

        macro_rules! add(
            ($weight:expr) => (self.data[j1 * self.units + unit] += $weight * data[j2]);
        );

        while j1 < s1 && j2 < s2 {
            let l1 = t1 + (j1 as f64) * d1;
            let l2 = t2 + (j2 as f64) * d2;

            let r1 = l1 + d1;
            let r2 = l2 + d2;

            if l1 < l2 {
                if r2 < r1 {
                    add!(1.0);
                    j2 += 1;
                } else {
                    add!((r1 - l2) / d2);
                    j1 += 1;
                }
            } else {
                if r1 < r2 {
                    add!(d1 / d2);
                    j1 += 1;
                } else {
                    add!((r2 - l1) / d2);
                    j2 += 1;
                }
            }
        }
    }

    pub fn discharge(&mut self, time: f64) -> Self {
        debug_assert!(time >= self.time);
        let steps = ((time - self.time) / self.time_step).floor() as usize;
        debug_assert!(self.steps >= steps);

        let mut profile = Profile {
            units: self.units,
            steps: self.steps - steps,
            time: (time / self.time_step).floor() * self.time_step,
            time_step: self.time_step,
            data: self.data[(steps * self.units)..].to_vec(),
        };

        mem::swap(self, &mut profile);

        profile.steps = steps;
        profile.data.truncate(steps * self.units);

        profile
    }
}

impl Into<Vec<f64>> for Profile {
    #[inline]
    fn into(self) -> Vec<f64> {
        self.data
    }
}

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

    #[test]
    fn accumulate_padding() {
        let mut profile = Profile::new(2, 0.5);

        profile.accumulate(0, 4.0, 1.0, &[]);
        assert_eq!(profile.steps, 8);
        assert_eq!(&profile.data, &vec![0.0; 2 * 8]);

        profile.accumulate(0, 2.5, 1.0, &[]);
        assert_eq!(profile.steps, 8);
        assert_eq!(&profile.data, &vec![0.0; 2 * 8]);

        profile.accumulate(0, 6.5, 1.0, &[]);
        assert_eq!(profile.steps, 13);
        assert_eq!(&profile.data, &vec![0.0; 2 * 13]);

        profile.accumulate(0, 6.55, 1.0, &[]);
        assert_eq!(profile.steps, 14);
        assert_eq!(&profile.data, &vec![0.0; 2 * 14]);
    }

    #[test]
    fn accumulate_synchronous() {
        let mut profile = Profile::new(2, 1.0);

        profile.accumulate(0, 1.0, 1.0, &[1.0, 2.0]);
        assert_eq!(profile.steps, 3);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           1.0, 0.0,
           2.0, 0.0,
        ]);

        profile.accumulate(0, 1.0, 1.0, &[1.0, 2.0, 3.0]);
        assert_eq!(profile.steps, 4);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           2.0, 0.0,
           4.0, 0.0,
           3.0, 0.0,
        ]);
    }

    #[test]
    fn accumulate_asynchronous() {
        let mut profile = Profile::new(2, 1.0);

        profile.accumulate(1, 1.5, 1.0, &[1.0, 2.0, 3.0]);
        assert_eq!(profile.steps, 5);
        assert_eq!(&profile.data, &vec![
           0.0, 0.0,
           0.0, 0.5,
           0.0, 1.5,
           0.0, 2.5,
           0.0, 1.5,
        ]);

        profile.accumulate(0, 0.5, 0.25, &[1.0, 2.0, 3.0, 1.0, 3.0]);
        assert_eq!(profile.steps, 5);
        assert_eq!(&profile.data, &vec![
           3.0, 0.0,
           7.0, 0.5,
           0.0, 1.5,
           0.0, 2.5,
           0.0, 1.5,
        ]);

        profile.accumulate(0, 1.25, 1.0, &[1.0, 2.0, 3.0, 0.0, 4.0]);
        assert_eq!(profile.steps, 7);
        assert_eq!(&profile.data, &vec![
           3.00, 0.0,
           7.75, 0.5,
           1.75, 1.5,
           2.75, 2.5,
           0.75, 1.5,
           3.00, 0.0,
           1.00, 0.0,
        ]);
    }

    #[test]
    fn discharge() {
        let mut profile = Profile::new(2, 1.0);
        profile.accumulate(0, 0.0, 1.0, &vec![42.0; 42]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(0.0).data, vec![]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(0.75).data, vec![]);
        assert_eq!(profile.time, 0.0);
        assert_eq!(profile.steps, 42);

        assert_eq!(profile.discharge(1.0).data, vec![42.0, 0.0]);
        assert_eq!(profile.time, 1.0);
        assert_eq!(profile.steps, 41);

        assert_eq!(profile.discharge(1.5).data, vec![]);
        assert_eq!(profile.time, 1.0);
        assert_eq!(profile.steps, 41);

        assert_eq!(profile.discharge(3.5).data, vec![42.0, 0.0, 42.0, 0.0]);
        assert_eq!(profile.time, 3.0);
        assert_eq!(profile.steps, 39);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hashmap;
use std::default::Default;
use std::slice;
use core::wire_format;
use clear::Clear;

#[deriving(Show)]
pub enum UnknownValue {
    UnknownFixed32(u32),
    UnknownFixed64(u64),
    UnknownVarint(u64),
    UnknownLengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownFixed32(fixed32) => UnknownFixed32Ref(fixed32),
            UnknownFixed64(fixed64) => UnknownFixed64Ref(fixed64),
            UnknownVarint(varint) => UnknownVarintRef(varint),
            UnknownLengthDelimited(ref bytes) => UnknownLengthDelimitedRef(bytes.as_slice()),
        }
    }
}

pub enum UnknownValueRef<'o> {
    UnknownFixed32Ref(u32),
    UnknownFixed64Ref(u64),
    UnknownVarintRef(u64),
    UnknownLengthDelimitedRef(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownFixed32Ref(_) => wire_format::WireTypeFixed32,
            UnknownFixed64Ref(_) => wire_format::WireTypeFixed64,
            UnknownVarintRef(_) => wire_format::WireTypeVarint,
            UnknownLengthDelimitedRef(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownFixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownFixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownVarint(varint) => self.varint.push(varint),
            UnknownLengthDelimited(length_delimited) => self.length_delimited.push(length_delimited),
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Items<'o, u32>,
    fixed64: slice::Items<'o, u64>,
    varint: slice::Items<'o, u64>,
    length_delimited: slice::Items<'o, Vec<u8>>,
}

impl<'o> Iterator<UnknownValueRef<'o>> for UnknownValuesIter<'o> {
    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownFixed32Ref(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownFixed64Ref(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownVarintRef(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownLengthDelimitedRef(length_delimited.get_ref().as_slice()))
        }
        None
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownFields {
    // option is needed, because HashMap constructor performs allocation,
    // and very expensive
    pub fields: Option<HashMap<u32, UnknownValues>>,
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: u32) -> &'a mut UnknownValues {
        self.init_map();

        self.fields.get_mut_ref()
            .find_or_insert_with(number, |_| Default::default())
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(number).length_delimited.push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldIter<'s> {
        UnknownFieldIter {
            entries: self.fields.as_ref().map(|m| m.iter())
        }
    }
}

impl Clear for UnknownFields {
    fn clear(&mut self) {
        if self.fields.is_some() {
            self.fields.get_mut_ref().clear();
        }
    }
}

pub struct UnknownFieldIter<'s> {
    entries: Option<hashmap::Entries<'s, u32, UnknownValues>>,
}

impl<'s> Iterator<(u32, &'s UnknownValues)> for UnknownFieldIter<'s> {
    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        if self.entries.is_none() {
            None
        } else {
            self.entries.get_mut_ref().next().map(|(&number, values)| (number, values))
        }
    }
}
<commit_msg>Put HashMap of UnknownFields into a box<commit_after>use std::collections::HashMap;
use std::collections::hashmap;
use std::default::Default;
use std::slice;
use core::wire_format;
use clear::Clear;

#[deriving(Show)]
pub enum UnknownValue {
    UnknownFixed32(u32),
    UnknownFixed64(u64),
    UnknownVarint(u64),
    UnknownLengthDelimited(Vec<u8>),
}

impl UnknownValue {
    pub fn wire_type(&self) -> wire_format::WireType {
        self.get_ref().wire_type()
    }

    pub fn get_ref<'s>(&'s self) -> UnknownValueRef<'s> {
        match *self {
            UnknownFixed32(fixed32) => UnknownFixed32Ref(fixed32),
            UnknownFixed64(fixed64) => UnknownFixed64Ref(fixed64),
            UnknownVarint(varint) => UnknownVarintRef(varint),
            UnknownLengthDelimited(ref bytes) => UnknownLengthDelimitedRef(bytes.as_slice()),
        }
    }
}

pub enum UnknownValueRef<'o> {
    UnknownFixed32Ref(u32),
    UnknownFixed64Ref(u64),
    UnknownVarintRef(u64),
    UnknownLengthDelimitedRef(&'o [u8]),
}

impl<'o> UnknownValueRef<'o> {
    pub fn wire_type(&self) -> wire_format::WireType {
        match *self {
            UnknownFixed32Ref(_) => wire_format::WireTypeFixed32,
            UnknownFixed64Ref(_) => wire_format::WireTypeFixed64,
            UnknownVarintRef(_) => wire_format::WireTypeVarint,
            UnknownLengthDelimitedRef(_) => wire_format::WireTypeLengthDelimited,
        }
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownValues {
    pub fixed32: Vec<u32>,
    pub fixed64: Vec<u64>,
    pub varint: Vec<u64>,
    pub length_delimited: Vec<Vec<u8>>,
}

impl UnknownValues {
    pub fn add_value(&mut self, value: UnknownValue) {
        match value {
            UnknownFixed64(fixed64) => self.fixed64.push(fixed64),
            UnknownFixed32(fixed32) => self.fixed32.push(fixed32),
            UnknownVarint(varint) => self.varint.push(varint),
            UnknownLengthDelimited(length_delimited) => self.length_delimited.push(length_delimited),
        };
    }

    pub fn iter<'s>(&'s self) -> UnknownValuesIter<'s> {
        UnknownValuesIter {
            fixed32: self.fixed32.iter(),
            fixed64: self.fixed64.iter(),
            varint: self.varint.iter(),
            length_delimited: self.length_delimited.iter(),
        }
    }
}

pub struct UnknownValuesIter<'o> {
    fixed32: slice::Items<'o, u32>,
    fixed64: slice::Items<'o, u64>,
    varint: slice::Items<'o, u64>,
    length_delimited: slice::Items<'o, Vec<u8>>,
}

impl<'o> Iterator<UnknownValueRef<'o>> for UnknownValuesIter<'o> {
    fn next(&mut self) -> Option<UnknownValueRef<'o>> {
        let fixed32 = self.fixed32.next();
        if fixed32.is_some() {
            return Some(UnknownFixed32Ref(*fixed32.unwrap()));
        }
        let fixed64 = self.fixed64.next();
        if fixed64.is_some() {
            return Some(UnknownFixed64Ref(*fixed64.unwrap()));
        }
        let varint = self.varint.next();
        if varint.is_some() {
            return Some(UnknownVarintRef(*varint.unwrap()));
        }
        let length_delimited = self.length_delimited.next();
        if length_delimited.is_some() {
            return Some(UnknownLengthDelimitedRef(length_delimited.get_ref().as_slice()))
        }
        None
    }
}

#[deriving(Clone,PartialEq,Eq,Show,Default)]
pub struct UnknownFields {
    // option is needed, because HashMap constructor performs allocation,
    // and very expensive
    pub fields: Option<Box<HashMap<u32, UnknownValues>>>,
}

impl UnknownFields {
    pub fn new() -> UnknownFields {
        Default::default()
    }

    fn init_map(&mut self) {
        if self.fields.is_none() {
            self.fields = Some(box Default::default());
        }
    }

    fn find_field<'a>(&'a mut self, number: u32) -> &'a mut UnknownValues {
        self.init_map();

        self.fields.get_mut_ref()
            .find_or_insert_with(number, |_| Default::default())
    }

    pub fn add_fixed32(&mut self, number: u32, fixed32: u32) {
        self.find_field(number).fixed32.push(fixed32);
    }

    pub fn add_fixed64(&mut self, number: u32, fixed64: u64) {
        self.find_field(number).fixed64.push(fixed64);
    }

    pub fn add_varint(&mut self, number: u32, varint: u64) {
        self.find_field(number).varint.push(varint);
    }

    pub fn add_length_delimited(&mut self, number: u32, length_delimited: Vec<u8>) {
        self.find_field(number).length_delimited.push(length_delimited);
    }

    pub fn add_value(&mut self, number: u32, value: UnknownValue) {
        self.find_field(number).add_value(value);
    }

    pub fn iter<'s>(&'s self) -> UnknownFieldIter<'s> {
        UnknownFieldIter {
            entries: self.fields.as_ref().map(|m| m.iter())
        }
    }
}

impl Clear for UnknownFields {
    fn clear(&mut self) {
        if self.fields.is_some() {
            self.fields.get_mut_ref().clear();
        }
    }
}

pub struct UnknownFieldIter<'s> {
    entries: Option<hashmap::Entries<'s, u32, UnknownValues>>,
}

impl<'s> Iterator<(u32, &'s UnknownValues)> for UnknownFieldIter<'s> {
    fn next(&mut self) -> Option<(u32, &'s UnknownValues)> {
        if self.entries.is_none() {
            None
        } else {
            self.entries.get_mut_ref().next().map(|(&number, values)| (number, values))
        }
    }
}
<|endoftext|>