<commit_before>//! Single client connection

use std::io;
use std::result::Result as std_Result;
use std::sync::Arc;

use crate::error;
use crate::error::Error;
use crate::result;
use crate::AnySocketAddr;

use crate::solicit::end_stream::EndStream;
use crate::solicit::header::*;

use tls_api::TlsConnector;

use tls_api;

use crate::solicit_async::*;

use crate::assert_types::assert_send_future;
use crate::client::req::ClientRequest;

use crate::client::stream_handler::ClientStreamCreatedHandler;
use crate::client::types::ClientTypes;
use crate::client::ClientInterface;
use crate::client_died_error_holder::ConnDiedType;
use crate::common::conn::Conn;
use crate::common::conn::ConnStateSnapshot;
use crate::common::conn::SideSpecific;
use crate::common::conn_read::ConnReadSideCustom;
use crate::common::conn_write::CommonToWriteMessage;
use crate::common::conn_write::ConnWriteSideCustom;
use crate::common::death_aware_channel::DeathAwareSender;
use crate::common::death_aware_channel::ErrorAwareDrop;
use crate::common::sender::CommonSender;
use crate::common::stream::HttpStreamCommon;
use crate::common::stream::HttpStreamData;
use crate::common::stream::HttpStreamDataSpecific;
use crate::common::stream::InMessageStage;
use crate::common::stream_handler::StreamHandlerInternal;
use crate::common::stream_map::HttpStreamRef;
use crate::data_or_headers::DataOrHeaders;
use crate::headers_place::HeadersPlace;
use crate::net::connect::ToClientStream;
use crate::net::socket::SocketStream;
use crate::req_resp::RequestOrResponse;
use crate::solicit::stream_id::StreamId;
use crate::ClientConf;
use crate::ClientTlsOption;
use crate::ErrorCode;
use bytes::Bytes;
use futures::channel::oneshot;
use futures::TryFutureExt;
use std::pin::Pin;

use crate::client::resp::ClientResponse;
use std::future::Future;
use tokio::runtime::Handle;
use tokio::time;

pub struct ClientStreamData {}

impl HttpStreamDataSpecific for ClientStreamData {}

pub(crate) type ClientStream = HttpStreamCommon<ClientTypes>;

impl HttpStreamData for ClientStream {
    type Types = ClientTypes;
}

pub struct ClientConnData {
    _callbacks: Box<dyn ClientConnCallbacks>,
}

impl SideSpecific for ClientConnData {}

pub struct ClientConn {
    write_tx: DeathAwareSender<ClientToWriteMessage>,
}

unsafe impl Sync for ClientConn {}

pub(crate) struct StartRequestMessage {
    pub headers: Headers,
    pub body: Option<Bytes>,
    pub trailers: Option<Headers>,
    pub end_stream: bool,
    pub stream_handler: Box<dyn ClientStreamCreatedHandler>,
}

pub struct ClientStartRequestMessage {
    start: StartRequestMessage,
    write_tx: DeathAwareSender<ClientToWriteMessage>,
}

pub(crate) enum ClientToWriteMessage {
    Start(ClientStartRequestMessage),
    WaitForHandshake(oneshot::Sender<result::Result<()>>),
    Common(CommonToWriteMessage),
}

impl ErrorAwareDrop for ClientToWriteMessage {
    type DiedType = ConnDiedType;

    fn drop_with_error(self, error: Error) {
        match self {
            ClientToWriteMessage::Start(start) => {
                start.start.stream_handler.error(error);
            }
            ClientToWriteMessage::WaitForHandshake(_) => {
                // TODO: error
            }
            ClientToWriteMessage::Common(_) => {
                // TODO: error
            }
        }
    }
}

impl From<CommonToWriteMessage> for ClientToWriteMessage {
    fn from(m: CommonToWriteMessage) -> Self {
        ClientToWriteMessage::Common(m)
    }
}

impl<I> ConnWriteSideCustom for Conn<ClientTypes, I>
where
    I: SocketStream,
{
    type Types = ClientTypes;

    fn process_message(&mut self, message: ClientToWriteMessage) -> result::Result<()> {
        match message {
            ClientToWriteMessage::Start(start) => self.process_start(start),
            ClientToWriteMessage::Common(common) => self.process_common_message(common),
            ClientToWriteMessage::WaitForHandshake(tx) => {
                // ignore error
                drop(tx.send(Ok(())));
                Ok(())
            }
        }
    }
}

impl<I> Conn<ClientTypes, I>
where
    I: SocketStream,
{
    fn process_start(&mut self, start: ClientStartRequestMessage) -> result::Result<()> {
        let ClientStartRequestMessage {
            start:
                StartRequestMessage {
                    headers,
                    body,
                    trailers,
                    end_stream,
                    stream_handler,
                },
            write_tx,
        } = start;

        let stream_id = self.next_local_stream_id();

        {
            let (_, out_window) = self.new_stream_data(
                stream_id,
                None,
                InMessageStage::Initial,
                ClientStreamData {},
            );

            let in_window_size = self
                .streams
                .get_mut(stream_id)
                .unwrap()
                .stream()
                .in_window_size
                .size() as u32;

            let req = ClientRequest {
                common: if end_stream {
                    CommonSender::new_done(stream_id)
                } else {
                    CommonSender::new(stream_id, write_tx, out_window, true)
                },
                drop_callback: None,
            };

            let mut handler = None;
            let resp = ClientResponse {
                stream_handler: &mut handler,
                in_window_size,
                stream_id,
                to_write_tx: &self.to_write_tx,
            };

            match stream_handler.request_created(req, resp) {
                Err(e) => {
                    warn!("client cancelled request: {:?}", e);
                    // Should be fine to cancel before start, but TODO: check
                    self.streams
                        .get_mut(stream_id)
                        .unwrap()
                        .close_outgoing(ErrorCode::InternalError);
                }
                Ok(()) => {
                    let mut stream = self.streams.get_mut(stream_id).unwrap();
                    stream.stream().peer_tx = handler;

                    stream.push_back(DataOrHeaders::Headers(headers));
                    if let Some(body) = body {
                        stream.push_back(DataOrHeaders::Data(body));
                    }
                    if let Some(trailers) = trailers {
                        stream.push_back(DataOrHeaders::Headers(trailers));
                    }
                    if end_stream {
                        stream.close_outgoing(ErrorCode::NoError);
                    }
                }
            };
        }

        // Also opens latch if necessary
        self.buffer_outg_conn()?;
        Ok(())
    }
}

pub trait ClientConnCallbacks: Send + 'static {
    // called at most once
    fn goaway(&self, stream_id: StreamId, raw_error_code: u32);
}

impl ClientConn {
    fn spawn_connected<I, C>(
        lh: Handle,
        connect: impl Future<Output = crate::Result<I>> + Send + 'static,
        peer_addr: AnySocketAddr,
        conf: ClientConf,
        callbacks: C,
    ) -> Self
    where
        I: SocketStream,
        C: ClientConnCallbacks,
    {
        let (future, write_tx) = Conn::<ClientTypes, _>::new(
            lh.clone(),
            ClientConnData {
                _callbacks: Box::new(callbacks),
            },
            conf.common,
            connect,
            peer_addr,
        );

        lh.spawn(future);

        ClientConn { write_tx }
    }

    pub fn spawn<H, C>(
        lh: Handle,
        addr: Pin<Box<dyn ToClientStream + Send>>,
        tls: ClientTlsOption<C>,
        conf: ClientConf,
        callbacks: H,
    ) -> Self
    where
        H: ClientConnCallbacks,
        C: TlsConnector + Sync,
    {
        match tls {
            ClientTlsOption::Plain => ClientConn::spawn_plain(lh.clone(), addr, conf, callbacks),
            ClientTlsOption::Tls(domain, connector) => {
                ClientConn::spawn_tls(lh.clone(), &domain, connector, addr, conf, callbacks)
            }
        }
    }

    pub fn spawn_plain<C>(
        lh: Handle,
        addr: Pin<Box<dyn ToClientStream>>,
        conf: ClientConf,
        callbacks: C,
    ) -> Self
    where
        C: ClientConnCallbacks,
    {
        let addr_struct = addr.socket_addr();

        let no_delay = conf.no_delay.unwrap_or(true);
        let connect = TryFutureExt::map_err(addr.connect(&lh), error::Error::from);

        let connect_timeout = conf.connection_timeout;

        let addr_copy = addr_struct.clone();
        let connect = async move {
            let socket = if let Some(timeout) = connect_timeout {
                match time::timeout(timeout, connect).await {
                    Ok(r) => r?,
                    Err(_) => return Err(error::Error::ConnectionTimeout),
                }
            } else {
                connect.await?
            };

            info!("connected to {}", addr_copy);

            if socket.is_tcp() {
                socket.set_tcp_nodelay(no_delay)?;
            }

            Ok(socket)
        };

        ClientConn::spawn_connected(lh, connect, addr_struct, conf, callbacks)
    }

    pub fn spawn_tls<H, C>(
        lh: Handle,
        domain: &str,
        connector: Arc<C>,
        addr: Pin<Box<dyn ToClientStream + Send>>,
        conf: ClientConf,
        callbacks: H,
    ) -> Self
    where
        H: ClientConnCallbacks,
        C: TlsConnector + Sync,
    {
        let addr_struct = addr.socket_addr();

        let domain = domain.to_owned();

        let no_delay = conf.no_delay.unwrap_or(true);
        let connect = addr
            .connect(&lh)
            .map_ok(move |socket| {
                info!("connected to {}", addr);

                if socket.is_tcp() {
                    socket
                        .set_tcp_nodelay(no_delay)
                        .expect("failed to set TCP_NODELAY");
                }
                socket
            })
            .map_err(|e| error::Error::from(e));

        let connect = assert_send_future(connect);

        let tls_conn = connect.and_then(move |conn| async move {
            Ok(connector.connect_with_socket(&domain, conn).await?)
        });

        let tls_conn = assert_send_future(tls_conn);

        let tls_conn = tls_conn.map_err(Error::from);

        let tls_conn = assert_send_future(tls_conn);

        ClientConn::spawn_connected(lh, Box::pin(tls_conn), addr_struct, conf, callbacks)
    }

    pub(crate) fn start_request_with_resp_sender(
        &self,
        start: StartRequestMessage,
    ) -> Result<(), (StartRequestMessage, error::Error)> {
        let client_start = ClientStartRequestMessage {
            start,
            write_tx: self.write_tx.clone(),
        };

        self.write_tx
            .unbounded_send_recover(ClientToWriteMessage::Start(client_start))
            .map_err(|(sent_message, e)| match sent_message {
                ClientToWriteMessage::Start(start) => (start.start, e),
                _ => unreachable!(),
            })
    }

    pub fn dump_state_with_resp_sender(&self, tx: oneshot::Sender<ConnStateSnapshot>) {
        let message = ClientToWriteMessage::Common(CommonToWriteMessage::DumpState(tx));
        // ignore error
        drop(self.write_tx.unbounded_send(message));
    }

    /// For tests
    #[doc(hidden)]
    pub fn _dump_state(&self) -> HttpFutureSend<ConnStateSnapshot> {
        let (tx, rx) = oneshot::channel();

        self.dump_state_with_resp_sender(tx);

        let rx =
            rx.map_err(|_| Error::from(io::Error::new(io::ErrorKind::Other, "oneshot canceled")));

        Box::pin(rx)
    }

    pub fn wait_for_connect_with_resp_sender(
        &self,
        tx: oneshot::Sender<result::Result<()>>,
    ) -> std_Result<(), oneshot::Sender<result::Result<()>>> {
        self.write_tx
            .unbounded_send_recover(ClientToWriteMessage::WaitForHandshake(tx))
            .map_err(|(send_message, _)| match send_message {
                ClientToWriteMessage::WaitForHandshake(tx) => tx,
                _ => unreachable!(),
            })
    }
}

impl ClientInterface for ClientConn {
    fn start_request_low_level(
        &self,
        headers: Headers,
        body: Option<Bytes>,
        trailers: Option<Headers>,
        end_stream: bool,
        stream_handler: Box<dyn ClientStreamCreatedHandler>,
    ) -> result::Result<()> {
        let start = StartRequestMessage {
            headers,
            body,
            trailers,
            end_stream,
            stream_handler,
        };

        if let Err((_, e)) = self.start_request_with_resp_sender(start) {
            return Err(e);
        }

        Ok(())
    }
}

impl<I> ConnReadSideCustom for Conn<ClientTypes, I>
where
    I: SocketStream,
{
    type Types = ClientTypes;

    fn process_headers(
        &mut self,
        stream_id: StreamId,
        end_stream: EndStream,
        headers: Headers,
    ) -> result::Result<Option<HttpStreamRef<ClientTypes>>> {
        let existing_stream = self
            .get_stream_for_headers_maybe_send_error(stream_id)?
            .is_some();
        if !existing_stream {
            return Ok(None);
        }

        let in_message_stage = self
            .streams
            .get_mut(stream_id)
            .unwrap()
            .stream()
            .in_message_stage;

        let headers_place = match in_message_stage {
            InMessageStage::Initial => HeadersPlace::Initial,
            InMessageStage::AfterInitialHeaders => HeadersPlace::Trailing,
            InMessageStage::AfterTrailingHeaders => {
                return Err(error::Error::InternalError(format!(
                    "closed stream must be handled before"
                )));
            }
        };

        if let Err(e) = headers.validate(RequestOrResponse::Response, headers_place) {
            warn!("invalid headers: {:?}: {:?}", e, headers);
            self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
            return Ok(None);
        }

        let status_1xx = match headers_place {
            HeadersPlace::Initial => {
                let status = headers.status();

                let status_1xx = status >= 100 && status <= 199;
                if status_1xx && end_stream == EndStream::Yes {
                    warn!("1xx headers and end stream: {}", stream_id);
                    self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
                    return Ok(None);
                }
                status_1xx
            }
            HeadersPlace::Trailing => {
                if end_stream == EndStream::No {
                    warn!("headers without end stream after data: {}", stream_id);
                    self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
                    return Ok(None);
                }
                false
            }
        };

        let mut stream = self.streams.get_mut(stream_id).unwrap();
        if let Some(in_rem_content_length) = headers.content_length() {
            stream.stream().in_rem_content_length = Some(in_rem_content_length);
        }

        stream.stream().in_message_stage = match (headers_place, status_1xx) {
            (HeadersPlace::Initial, false) => InMessageStage::AfterInitialHeaders,
            (HeadersPlace::Initial, true) => InMessageStage::Initial,
            (HeadersPlace::Trailing, _) => InMessageStage::AfterTrailingHeaders,
        };

        // Ignore 1xx headers
        if !status_1xx {
            if let Some(ref mut response_handler) = stream.stream().peer_tx {
                // TODO: reset stream on error
                drop(match headers_place {
                    HeadersPlace::Initial => response_handler
                        .0
                        .headers(headers, end_stream == EndStream::Yes),
                    HeadersPlace::Trailing => {
                        assert_eq!(EndStream::Yes, end_stream);
                        response_handler.trailers(headers)
                    }
                });
            } else {
                // TODO: reset stream
            }
        }

        Ok(Some(stream))
    }
}
<commit_msg>Rewrite tls connect to async<commit_after>//! Single client connection

use std::io;
use std::result::Result as std_Result;
use std::sync::Arc;

use crate::error;
use crate::error::Error;
use crate::result;
use crate::AnySocketAddr;

use crate::solicit::end_stream::EndStream;
use crate::solicit::header::*;

use tls_api::TlsConnector;

use tls_api;

use crate::solicit_async::*;

use crate::client::req::ClientRequest;
use crate::client::stream_handler::ClientStreamCreatedHandler;
use crate::client::types::ClientTypes;
use crate::client::ClientInterface;
use crate::client_died_error_holder::ConnDiedType;
use crate::common::conn::Conn;
use crate::common::conn::ConnStateSnapshot;
use crate::common::conn::SideSpecific;
use crate::common::conn_read::ConnReadSideCustom;
use crate::common::conn_write::CommonToWriteMessage;
use crate::common::conn_write::ConnWriteSideCustom;
use crate::common::death_aware_channel::DeathAwareSender;
use crate::common::death_aware_channel::ErrorAwareDrop;
use crate::common::sender::CommonSender;
use crate::common::stream::HttpStreamCommon;
use crate::common::stream::HttpStreamData;
use crate::common::stream::HttpStreamDataSpecific;
use crate::common::stream::InMessageStage;
use crate::common::stream_handler::StreamHandlerInternal;
use crate::common::stream_map::HttpStreamRef;
use crate::data_or_headers::DataOrHeaders;
use crate::headers_place::HeadersPlace;
use crate::net::connect::ToClientStream;
use crate::net::socket::SocketStream;
use crate::req_resp::RequestOrResponse;
use crate::solicit::stream_id::StreamId;
use crate::ClientConf;
use crate::ClientTlsOption;
use crate::ErrorCode;
use bytes::Bytes;
use futures::channel::oneshot;
use futures::TryFutureExt;
use std::pin::Pin;

use crate::client::resp::ClientResponse;
use std::future::Future;
use tokio::runtime::Handle;
use tokio::time;

pub struct ClientStreamData {}

impl HttpStreamDataSpecific for ClientStreamData {}

pub(crate) type ClientStream = HttpStreamCommon<ClientTypes>;

impl HttpStreamData for ClientStream {
    type Types = ClientTypes;
}

pub struct ClientConnData {
    _callbacks: Box<dyn ClientConnCallbacks>,
}

impl SideSpecific for ClientConnData {}

pub struct ClientConn {
    write_tx: DeathAwareSender<ClientToWriteMessage>,
}

unsafe impl Sync for ClientConn {}

pub(crate) struct StartRequestMessage {
    pub headers: Headers,
    pub body: Option<Bytes>,
    pub trailers: Option<Headers>,
    pub end_stream: bool,
    pub stream_handler: Box<dyn ClientStreamCreatedHandler>,
}

pub struct ClientStartRequestMessage {
    start: StartRequestMessage,
    write_tx: DeathAwareSender<ClientToWriteMessage>,
}

pub(crate) enum ClientToWriteMessage {
    Start(ClientStartRequestMessage),
    WaitForHandshake(oneshot::Sender<result::Result<()>>),
    Common(CommonToWriteMessage),
}

impl ErrorAwareDrop for ClientToWriteMessage {
    type DiedType = ConnDiedType;

    fn drop_with_error(self, error: Error) {
        match self {
            ClientToWriteMessage::Start(start) => {
                start.start.stream_handler.error(error);
            }
            ClientToWriteMessage::WaitForHandshake(_) => {
                // TODO: error
            }
            ClientToWriteMessage::Common(_) => {
                // TODO: error
            }
        }
    }
}

impl From<CommonToWriteMessage> for ClientToWriteMessage {
    fn from(m: CommonToWriteMessage) -> Self {
        ClientToWriteMessage::Common(m)
    }
}

impl<I> ConnWriteSideCustom for Conn<ClientTypes, I>
where
    I: SocketStream,
{
    type Types = ClientTypes;

    fn process_message(&mut self, message: ClientToWriteMessage) -> result::Result<()> {
        match message {
            ClientToWriteMessage::Start(start) => self.process_start(start),
            ClientToWriteMessage::Common(common) => self.process_common_message(common),
            ClientToWriteMessage::WaitForHandshake(tx) => {
                // ignore error
                drop(tx.send(Ok(())));
                Ok(())
            }
        }
    }
}

impl<I> Conn<ClientTypes, I>
where
    I: SocketStream,
{
    fn process_start(&mut self, start: ClientStartRequestMessage) -> result::Result<()> {
        let ClientStartRequestMessage {
            start:
                StartRequestMessage {
                    headers,
                    body,
                    trailers,
                    end_stream,
                    stream_handler,
                },
            write_tx,
        } = start;

        let stream_id = self.next_local_stream_id();

        {
            let (_, out_window) = self.new_stream_data(
                stream_id,
                None,
                InMessageStage::Initial,
                ClientStreamData {},
            );

            let in_window_size = self
                .streams
                .get_mut(stream_id)
                .unwrap()
                .stream()
                .in_window_size
                .size() as u32;

            let req = ClientRequest {
                common: if end_stream {
                    CommonSender::new_done(stream_id)
                } else {
                    CommonSender::new(stream_id, write_tx, out_window, true)
                },
                drop_callback: None,
            };

            let mut handler = None;
            let resp = ClientResponse {
                stream_handler: &mut handler,
                in_window_size,
                stream_id,
                to_write_tx: &self.to_write_tx,
            };

            match stream_handler.request_created(req, resp) {
                Err(e) => {
                    warn!("client cancelled request: {:?}", e);
                    // Should be fine to cancel before start, but TODO: check
                    self.streams
                        .get_mut(stream_id)
                        .unwrap()
                        .close_outgoing(ErrorCode::InternalError);
                }
                Ok(()) => {
                    let mut stream = self.streams.get_mut(stream_id).unwrap();
                    stream.stream().peer_tx = handler;

                    stream.push_back(DataOrHeaders::Headers(headers));
                    if let Some(body) = body {
                        stream.push_back(DataOrHeaders::Data(body));
                    }
                    if let Some(trailers) = trailers {
                        stream.push_back(DataOrHeaders::Headers(trailers));
                    }
                    if end_stream {
                        stream.close_outgoing(ErrorCode::NoError);
                    }
                }
            };
        }

        // Also opens latch if necessary
        self.buffer_outg_conn()?;
        Ok(())
    }
}

pub trait ClientConnCallbacks: Send + 'static {
    // called at most once
    fn goaway(&self, stream_id: StreamId, raw_error_code: u32);
}

impl ClientConn {
    fn spawn_connected<I, C>(
        lh: Handle,
        connect: impl Future<Output = crate::Result<I>> + Send + 'static,
        peer_addr: AnySocketAddr,
        conf: ClientConf,
        callbacks: C,
    ) -> Self
    where
        I: SocketStream,
        C: ClientConnCallbacks,
    {
        let (future, write_tx) = Conn::<ClientTypes, _>::new(
            lh.clone(),
            ClientConnData {
                _callbacks: Box::new(callbacks),
            },
            conf.common,
            connect,
            peer_addr,
        );

        lh.spawn(future);

        ClientConn { write_tx }
    }

    pub fn spawn<H, C>(
        lh: Handle,
        addr: Pin<Box<dyn ToClientStream + Send>>,
        tls: ClientTlsOption<C>,
        conf: ClientConf,
        callbacks: H,
    ) -> Self
    where
        H: ClientConnCallbacks,
        C: TlsConnector + Sync,
    {
        match tls {
            ClientTlsOption::Plain => ClientConn::spawn_plain(lh.clone(), addr, conf, callbacks),
            ClientTlsOption::Tls(domain, connector) => {
                ClientConn::spawn_tls(lh.clone(), &domain, connector, addr, conf, callbacks)
            }
        }
    }

    pub fn spawn_plain<C>(
        lh: Handle,
        addr: Pin<Box<dyn ToClientStream>>,
        conf: ClientConf,
        callbacks: C,
    ) -> Self
    where
        C: ClientConnCallbacks,
    {
        let addr_struct = addr.socket_addr();

        let no_delay = conf.no_delay.unwrap_or(true);
        let connect = TryFutureExt::map_err(addr.connect(&lh), error::Error::from);

        let connect_timeout = conf.connection_timeout;

        let addr_copy = addr_struct.clone();
        let connect = async move {
            let socket = if let Some(timeout) = connect_timeout {
                match time::timeout(timeout, connect).await {
                    Ok(r) => r?,
                    Err(_) => return Err(error::Error::ConnectionTimeout),
                }
            } else {
                connect.await?
            };

            info!("connected to {}", addr_copy);

            if socket.is_tcp() {
                socket.set_tcp_nodelay(no_delay)?;
            }

            Ok(socket)
        };

        ClientConn::spawn_connected(lh, connect, addr_struct, conf, callbacks)
    }

    pub fn spawn_tls<H, C>(
        lh: Handle,
        domain: &str,
        connector: Arc<C>,
        addr: Pin<Box<dyn ToClientStream + Send>>,
        conf: ClientConf,
        callbacks: H,
    ) -> Self
    where
        H: ClientConnCallbacks,
        C: TlsConnector + Sync,
    {
        let addr_struct = addr.socket_addr();
        let domain = domain.to_owned();
        let no_delay = conf.no_delay.unwrap_or(true);
        let lh_copy = lh.clone();
        let tls_conn = async move {
            // TODO: connect timeout
            let socket = addr.connect(&lh_copy).await?;
            info!("connected to {}", addr);

            if socket.is_tcp() {
                socket.set_tcp_nodelay(no_delay)?;
            }

            connector
                .connect_with_socket(&domain, socket)
                .await
                .map_err(crate::Error::from)
        };

        ClientConn::spawn_connected(lh, tls_conn, addr_struct, conf, callbacks)
    }

    pub(crate) fn start_request_with_resp_sender(
        &self,
        start: StartRequestMessage,
    ) -> Result<(), (StartRequestMessage, error::Error)> {
        let client_start = ClientStartRequestMessage {
            start,
            write_tx: self.write_tx.clone(),
        };

        self.write_tx
            .unbounded_send_recover(ClientToWriteMessage::Start(client_start))
            .map_err(|(sent_message, e)| match sent_message {
                ClientToWriteMessage::Start(start) => (start.start, e),
                _ => unreachable!(),
            })
    }

    pub fn dump_state_with_resp_sender(&self, tx: oneshot::Sender<ConnStateSnapshot>) {
        let message = ClientToWriteMessage::Common(CommonToWriteMessage::DumpState(tx));
        // ignore error
        drop(self.write_tx.unbounded_send(message));
    }

    /// For tests
    #[doc(hidden)]
    pub fn _dump_state(&self) -> HttpFutureSend<ConnStateSnapshot> {
        let (tx, rx) = oneshot::channel();

        self.dump_state_with_resp_sender(tx);

        let rx =
            rx.map_err(|_| Error::from(io::Error::new(io::ErrorKind::Other, "oneshot canceled")));

        Box::pin(rx)
    }

    pub fn wait_for_connect_with_resp_sender(
        &self,
        tx: oneshot::Sender<result::Result<()>>,
    ) -> std_Result<(), oneshot::Sender<result::Result<()>>> {
        self.write_tx
            .unbounded_send_recover(ClientToWriteMessage::WaitForHandshake(tx))
            .map_err(|(send_message, _)| match send_message {
                ClientToWriteMessage::WaitForHandshake(tx) => tx,
                _ => unreachable!(),
            })
    }
}

impl ClientInterface for ClientConn {
    fn start_request_low_level(
        &self,
        headers: Headers,
        body: Option<Bytes>,
        trailers: Option<Headers>,
        end_stream: bool,
        stream_handler: Box<dyn ClientStreamCreatedHandler>,
    ) -> result::Result<()> {
        let start = StartRequestMessage {
            headers,
            body,
            trailers,
            end_stream,
            stream_handler,
        };

        if let Err((_, e)) = self.start_request_with_resp_sender(start) {
            return Err(e);
        }

        Ok(())
    }
}

impl<I> ConnReadSideCustom for Conn<ClientTypes, I>
where
    I: SocketStream,
{
    type Types = ClientTypes;

    fn process_headers(
        &mut self,
        stream_id: StreamId,
        end_stream: EndStream,
        headers: Headers,
    ) -> result::Result<Option<HttpStreamRef<ClientTypes>>> {
        let existing_stream = self
            .get_stream_for_headers_maybe_send_error(stream_id)?
            .is_some();
        if !existing_stream {
            return Ok(None);
        }

        let in_message_stage = self
            .streams
            .get_mut(stream_id)
            .unwrap()
            .stream()
            .in_message_stage;

        let headers_place = match in_message_stage {
            InMessageStage::Initial => HeadersPlace::Initial,
            InMessageStage::AfterInitialHeaders => HeadersPlace::Trailing,
            InMessageStage::AfterTrailingHeaders => {
                return Err(error::Error::InternalError(format!(
                    "closed stream must be handled before"
                )));
            }
        };

        if let Err(e) = headers.validate(RequestOrResponse::Response, headers_place) {
            warn!("invalid headers: {:?}: {:?}", e, headers);
            self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
            return Ok(None);
        }

        let status_1xx = match headers_place {
            HeadersPlace::Initial => {
                let status = headers.status();

                let status_1xx = status >= 100 && status <= 199;
                if status_1xx && end_stream == EndStream::Yes {
                    warn!("1xx headers and end stream: {}", stream_id);
                    self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
                    return Ok(None);
                }
                status_1xx
            }
            HeadersPlace::Trailing => {
                if end_stream == EndStream::No {
                    warn!("headers without end stream after data: {}", stream_id);
                    self.send_rst_stream(stream_id, ErrorCode::ProtocolError)?;
                    return Ok(None);
                }
                false
            }
        };

        let mut stream = self.streams.get_mut(stream_id).unwrap();
        if let Some(in_rem_content_length) = headers.content_length() {
            stream.stream().in_rem_content_length = Some(in_rem_content_length);
        }

        stream.stream().in_message_stage = match (headers_place, status_1xx) {
            (HeadersPlace::Initial, false) => InMessageStage::AfterInitialHeaders,
            (HeadersPlace::Initial, true) => InMessageStage::Initial,
            (HeadersPlace::Trailing, _) => InMessageStage::AfterTrailingHeaders,
        };

        // Ignore 1xx headers
        if !status_1xx {
            if let Some(ref mut response_handler) = stream.stream().peer_tx {
                // TODO: reset stream on error
                drop(match headers_place {
                    HeadersPlace::Initial => response_handler
                        .0
                        .headers(headers, end_stream == EndStream::Yes),
                    HeadersPlace::Trailing => {
                        assert_eq!(EndStream::Yes, end_stream);
                        response_handler.trailers(headers)
                    }
                });
            } else {
                // TODO: reset stream
            }
        }

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


<commit_before>//! Statistics types for two-player games.

use crate::game;

use std::cmp;
use std::fmt;
use std::marker::PhantomData;
use std::sync::atomic;

use syncbox::atomic::AtomicU64;

/// Generic player designations for a two-player game.
///
/// To use the payoff and statistics types defined in this module, implement
/// [PlayerMapping](struct.PlayerMapping.html) to map your game's player type to
/// and from this one.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum Player {
  One,
  Two,
}

/// Provides a bidirectional mapping between an arbitrary game-specific type and
/// `Player`.
pub trait PlayerMapping: fmt::Debug {
  fn player_one() -> Self;
  fn player_two() -> Self;
  fn resolve_player(&self) -> Player;
}

impl PlayerMapping for Player {
  fn player_one() -> Self { Player::One }
  fn player_two() -> Self { Player::Two }
  fn resolve_player(&self) -> Player { *self }
}

/// Generic game payoff for a two-player game where each player gets a
/// whole-number score at the end of the game.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ScoredPayoff {
  pub visits: u32,
  pub score_one: u32,
  pub score_two: u32,
}

/// Atomically mutable game statistics for a two-player game where each player
/// gets a whole-number score at the end of the game. This type counts the
/// number of games that have been observed and the sum of the final score of
/// each game for each player.
///
/// This type's fields have limited range. See constants in this module for
/// their limiting values.
pub struct ScoredStatistics<M: PlayerMapping> {
  packed: AtomicU64,
  player: PhantomData<M>,
}

const VISITS_MASK: u64  // Upper 20 bits.
  = 0xFFFFF00000000000;
const ONE_SCORE_MASK: u64  // Middle 22 bits.
  = 0x00000FFFFFC00000;
const TWO_SCORE_MASK: u64  // Lower 22 bits.
  = 0x00000000003FFFFF;
/// The maximum value of the `visits` field of a
/// [ScoredStatistics](struct.ScoredStatistics.html).
pub const VISITS_MAX: u32 = (VISITS_MASK >> 44) as u32;
/// The maximum value of the two `score` fields of a
/// [ScoredStatistics](struct.ScoredStatistics.html).
pub const SCORE_MAX: u32 = TWO_SCORE_MASK as u32;

fn pack_scores(visits: u32, score_one: u32, score_two: u32) -> u64 {
  ((visits as u64) << 44)  // Visits.
    | (((score_one as u64) << 22) & ONE_SCORE_MASK)
    | ((score_two as u64) & TWO_SCORE_MASK)
}

fn unpack_scores(packed: u64) -> (u32, u32, u32) {
  (
    ((packed & VISITS_MASK) >> 44) as u32,
    ((packed & ONE_SCORE_MASK) >> 22) as u32,
    (packed & TWO_SCORE_MASK) as u32,
  )
}

impl<M: PlayerMapping> ScoredStatistics<M> {
  /// Creates statistics with no observed outcomes and starting scores of 0 for
  /// each player.
  pub fn new() -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(0),
      player: PhantomData,
    }
  }

  /// Creates statistics for the given number of observed outcomes (`visits`)
  /// and the sum of the final scores for each player.
  pub fn from_scores(visits: u32, score_one: u32, score_two: u32) -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(pack_scores(visits, score_one, score_two)),
      player: PhantomData,
    }
  }

  /// Returns the number of outcomes that have been recorded.
  pub fn visits(&self) -> u32 {
    let packed = self.packed.load(atomic::Ordering::SeqCst);
    ((packed & VISITS_MASK) >> 44) as u32
  }

  /// Returns the score for `player`.
  pub fn score(&self, player: Player) -> u32 {
    let packed = self.packed.load(atomic::Ordering::SeqCst);
    match player {
      Player::One => ((packed & ONE_SCORE_MASK) >> 22) as u32,
      Player::Two => (packed & TWO_SCORE_MASK) as u32,
    }
  }

  /// Increments the number of outcomes observed by 1 and adds each score to the
  /// running total for its respective player.
  ///
  /// If any field becomes saturated, it will stay at its maximum value.
  pub fn record_final_score(&self, score_one: u32, score_two: u32) {
    let mut success = false;
    // CAS loop because we have multiple fields to check for saturation.
    while !success {
      let old_packed = self.packed.load(atomic::Ordering::SeqCst);
      let (old_visits, old_score_one, old_score_two) = unpack_scores(old_packed);
      let visits = cmp::min(old_visits + 1, VISITS_MAX);
      let score_one = cmp::min(old_score_one + score_one, SCORE_MAX);
      let score_two = cmp::min(old_score_two + score_two, SCORE_MAX);
      success = self.packed.compare_and_swap(
        old_packed,
        pack_scores(visits, score_one, score_two),
        atomic::Ordering::SeqCst,
      ) == old_packed;
    }
  }
}

impl<M: PlayerMapping> Clone for ScoredStatistics<M> {
  fn clone(&self) -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(self.packed.load(atomic::Ordering::SeqCst)),
      player: PhantomData,
    }
  }
}

impl<M: PlayerMapping> Default for ScoredStatistics<M> {
  fn default() -> Self {
    ScoredStatistics::new()
  }
}

impl<M: PlayerMapping> fmt::Debug for ScoredStatistics<M> {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "Statistics(visits: {}, {:?}: {}, {:?}: {})",
      self.visits(),
      M::player_one(),
      self.score(Player::One),
      M::player_two(),
      self.score(Player::Two)
    )
  }
}

impl<S, M> game::Statistics<S, ScoredPayoff> for ScoredStatistics<M>
where
  S: game::State<PlayerId = M>,
  M: PlayerMapping,
{
  fn increment(&self, payoff: &ScoredPayoff) {
    self.record_final_score(payoff.score_one, payoff.score_two)
  }

  fn visits(&self) -> u32 {
    self.visits()
  }

  fn score(&self, player: &S::PlayerId) -> f32 {
    self.score(player.resolve_player()) as f32
  }
}

#[cfg(test)]
mod test {
  use super::{Player, ScoredStatistics};
  use std::u32;

  #[test]
  fn new_statistics_zero() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    assert_eq!(0, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_visits() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(0, 0);
    assert_eq!(1, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(3, 0);
    assert_eq!(1, stats.visits());
    assert_eq!(3, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(0, 5);
    assert_eq!(1, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(5, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(10, 15);
    assert_eq!(1, stats.visits());
    assert_eq!(10, stats.score(Player::One));
    assert_eq!(15, stats.score(Player::Two));

    stats.record_final_score(3, 3);
    assert_eq!(2, stats.visits());
    assert_eq!(13, stats.score(Player::One));
    assert_eq!(18, stats.score(Player::Two));

    stats.record_final_score(1000, 1000);
    assert_eq!(3, stats.visits());
    assert_eq!(1013, stats.score(Player::One));
    assert_eq!(1018, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_truncate() {
    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(u32::MAX, u32::MAX, u32::MAX);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_overflow() {
    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(0xFFFFFu32 - 1, 0x3FFFFFu32 - 1, 0x3FFFFF - 1);
    assert_eq!(0xFFFFF - 1, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));
    stats.record_final_score(1, 0);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));
    stats.record_final_score(1, 0);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(0xFFFFFu32 - 1, 0x3FFFFFu32 - 1, 0x3FFFFF - 1);
    stats.record_final_score(0, 1);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
    stats.record_final_score(0, 1);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
  }
}
<commit_msg>Use net score instead of absolute score so that MCTS agent doesn't resign immediately.<commit_after>//! Statistics types for two-player games.

use crate::game;

use std::cmp;
use std::fmt;
use std::marker::PhantomData;
use std::sync::atomic;

use syncbox::atomic::AtomicU64;

/// Generic player designations for a two-player game.
///
/// To use the payoff and statistics types defined in this module, implement
/// [PlayerMapping](struct.PlayerMapping.html) to map your game's player type to
/// and from this one.
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum Player {
  One,
  Two,
}

/// Provides a bidirectional mapping between an arbitrary game-specific type and
/// `Player`.
pub trait PlayerMapping: fmt::Debug {
  fn player_one() -> Self;
  fn player_two() -> Self;
  fn resolve_player(&self) -> Player;
}

impl PlayerMapping for Player {
  fn player_one() -> Self { Player::One }
  fn player_two() -> Self { Player::Two }
  fn resolve_player(&self) -> Player { *self }
}

/// Generic game payoff for a two-player game where each player gets a
/// whole-number score at the end of the game.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ScoredPayoff {
  pub visits: u32,
  pub score_one: u32,
  pub score_two: u32,
}

/// Atomically mutable game statistics for a two-player game where each player
/// gets a whole-number score at the end of the game. This type counts the
/// number of games that have been observed and the sum of the final score of
/// each game for each player.
///
/// This type's fields have limited range. See constants in this module for
/// their limiting values.
pub struct ScoredStatistics<M: PlayerMapping> {
  packed: AtomicU64,
  player: PhantomData<M>,
}

const VISITS_MASK: u64  // Upper 20 bits.
  = 0xFFFFF00000000000;
const ONE_SCORE_MASK: u64  // Middle 22 bits.
  = 0x00000FFFFFC00000;
const TWO_SCORE_MASK: u64  // Lower 22 bits.
  = 0x00000000003FFFFF;
/// The maximum value of the `visits` field of a
/// [ScoredStatistics](struct.ScoredStatistics.html).
pub const VISITS_MAX: u32 = (VISITS_MASK >> 44) as u32;
/// The maximum value of the two `score` fields of a
/// [ScoredStatistics](struct.ScoredStatistics.html).
pub const SCORE_MAX: u32 = TWO_SCORE_MASK as u32;

fn pack_scores(visits: u32, score_one: u32, score_two: u32) -> u64 {
  ((visits as u64) << 44)  // Visits.
    | (((score_one as u64) << 22) & ONE_SCORE_MASK)
    | ((score_two as u64) & TWO_SCORE_MASK)
}

fn unpack_scores(packed: u64) -> (u32, u32, u32) {
  (
    ((packed & VISITS_MASK) >> 44) as u32,
    ((packed & ONE_SCORE_MASK) >> 22) as u32,
    (packed & TWO_SCORE_MASK) as u32,
  )
}

impl<M: PlayerMapping> ScoredStatistics<M> {
  /// Creates statistics with no observed outcomes and starting scores of 0 for
  /// each player.
  pub fn new() -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(0),
      player: PhantomData,
    }
  }

  /// Creates statistics for the given number of observed outcomes (`visits`)
  /// and the sum of the final scores for each player.
  pub fn from_scores(visits: u32, score_one: u32, score_two: u32) -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(pack_scores(visits, score_one, score_two)),
      player: PhantomData,
    }
  }

  /// Returns the number of outcomes that have been recorded.
  pub fn visits(&self) -> u32 {
    let packed = self.packed.load(atomic::Ordering::SeqCst);
    ((packed & VISITS_MASK) >> 44) as u32
  }

  /// Returns the net score for `player`, which is equal to the difference
  /// between `player`'s score and that of the opponent.
  pub fn net_score(&self, player: Player) -> i32 {
    let packed = self.packed.load(atomic::Ordering::SeqCst);
    let score_one = ((packed & ONE_SCORE_MASK) >> 22) as i32;
    let score_two = (packed & TWO_SCORE_MASK) as i32;
    match player {
      Player::One => score_one - score_two,
      Player::Two => score_two - score_one,
    }
  }

  /// Returns the score for `player`.
  pub fn score(&self, player: Player) -> u32 {
    let packed = self.packed.load(atomic::Ordering::SeqCst);
    match player {
      Player::One => ((packed & ONE_SCORE_MASK) >> 22) as u32,
      Player::Two => (packed & TWO_SCORE_MASK) as u32,
    }
  }

  /// Increments the number of outcomes observed by 1 and adds each score to the
  /// running total for its respective player.
  ///
  /// If any field becomes saturated, it will stay at its maximum value.
  pub fn record_final_score(&self, score_one: u32, score_two: u32) {
    let mut success = false;
    // CAS loop because we have multiple fields to check for saturation.
    while !success {
      let old_packed = self.packed.load(atomic::Ordering::SeqCst);
      let (old_visits, old_score_one, old_score_two) = unpack_scores(old_packed);
      let visits = cmp::min(old_visits + 1, VISITS_MAX);
      let score_one = cmp::min(old_score_one + score_one, SCORE_MAX);
      let score_two = cmp::min(old_score_two + score_two, SCORE_MAX);
      success = self.packed.compare_and_swap(
        old_packed,
        pack_scores(visits, score_one, score_two),
        atomic::Ordering::SeqCst,
      ) == old_packed;
    }
  }
}

impl<M: PlayerMapping> Clone for ScoredStatistics<M> {
  fn clone(&self) -> Self {
    ScoredStatistics {
      packed: AtomicU64::new(self.packed.load(atomic::Ordering::SeqCst)),
      player: PhantomData,
    }
  }
}

impl<M: PlayerMapping> Default for ScoredStatistics<M> {
  fn default() -> Self {
    ScoredStatistics::new()
  }
}

impl<M: PlayerMapping> fmt::Debug for ScoredStatistics<M> {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    write!(
      f,
      "Statistics(visits: {}, {:?}: {}, {:?}: {})",
      self.visits(),
      M::player_one(),
      self.score(Player::One),
      M::player_two(),
      self.score(Player::Two)
    )
  }
}

impl<S, M> game::Statistics<S, ScoredPayoff> for ScoredStatistics<M>
where
  S: game::State<PlayerId = M>,
  M: PlayerMapping,
{
  fn increment(&self, payoff: &ScoredPayoff) {
    self.record_final_score(payoff.score_one, payoff.score_two)
  }

  fn visits(&self) -> u32 {
    self.visits()
  }

  fn score(&self, player: &S::PlayerId) -> f32 {
    self.net_score(player.resolve_player()) as f32
  }
}

#[cfg(test)]
mod test {
  use super::{Player, ScoredStatistics};
  use std::u32;

  #[test]
  fn new_statistics_zero() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    assert_eq!(0, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_visits() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(0, 0);
    assert_eq!(1, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum() {
    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(3, 0);
    assert_eq!(1, stats.visits());
    assert_eq!(3, stats.score(Player::One));
    assert_eq!(0, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(0, 5);
    assert_eq!(1, stats.visits());
    assert_eq!(0, stats.score(Player::One));
    assert_eq!(5, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> = ScoredStatistics::new();
    stats.record_final_score(10, 15);
    assert_eq!(1, stats.visits());
    assert_eq!(10, stats.score(Player::One));
    assert_eq!(15, stats.score(Player::Two));

    stats.record_final_score(3, 3);
    assert_eq!(2, stats.visits());
    assert_eq!(13, stats.score(Player::One));
    assert_eq!(18, stats.score(Player::Two));

    stats.record_final_score(1000, 1000);
    assert_eq!(3, stats.visits());
    assert_eq!(1013, stats.score(Player::One));
    assert_eq!(1018, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_truncate() {
    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(u32::MAX, u32::MAX, u32::MAX);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
  }

  #[test]
  fn statistics_sum_overflow() {
    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(0xFFFFFu32 - 1, 0x3FFFFFu32 - 1, 0x3FFFFF - 1);
    assert_eq!(0xFFFFF - 1, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));
    stats.record_final_score(1, 0);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));
    stats.record_final_score(1, 0);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF, stats.score(Player::One));
    assert_eq!(0x3FFFFF - 1, stats.score(Player::Two));

    let stats: ScoredStatistics<Player> =
      ScoredStatistics::from_scores(0xFFFFFu32 - 1, 0x3FFFFFu32 - 1, 0x3FFFFF - 1);
    stats.record_final_score(0, 1);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
    stats.record_final_score(0, 1);
    assert_eq!(0xFFFFF, stats.visits());
    assert_eq!(0x3FFFFF - 1, stats.score(Player::One));
    assert_eq!(0x3FFFFF, stats.score(Player::Two));
  }
}
<|endoftext|>


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

use std::collections::HashSet;

use crate::{
    conversion::{api::ApiName, convert_error::ErrorContext, ConvertError},
    types::QualifiedName,
};

use super::api::{AnalysisPhase, Api};

/// Newtype wrapper for a list of APIs, which enforced the invariant
/// that each API has a unique name.
pub(crate) struct ApiVec<P: AnalysisPhase> {
    apis: Vec<Api<P>>,
    names: HashSet<QualifiedName>,
}

impl<P: AnalysisPhase> ApiVec<P> {
    pub(crate) fn push(&mut self, api: Api<P>) {
        let name = api.name();
        let already_contains = self.already_contains(name);
        if already_contains {
            if api.discard_duplicates() {
                log::info!("Discarding duplicate API for {}", name);
            } else {
                panic!(
                    "Already have an API with that name: {}. API was {:?}",
                    name, api
                );
            }
        }
        self.names.insert(name.clone());
        self.apis.push(api)
    }

    pub(crate) fn push_eliminating_duplicates(&mut self, api: Api<P>) {
        let name = api.name();
        let already_contains = self.already_contains(name);
        if already_contains {
            log::info!(
                "Duplicate API for {} - removing all of them and replacing with an IgnoredItem.",
                name
            );
            self.retain(|api| api.name() != name);
            self.push(Api::IgnoredItem {
                name: ApiName::new_from_qualified_name(name.clone()),
                err: ConvertError::DuplicateItemsFoundInParsing,
                ctx: ErrorContext::Item(name.get_final_ident()),
            })
        } else {
            self.names.insert(name.clone());
            self.apis.push(api)
        }
    }

    fn already_contains(&self, name: &QualifiedName) -> bool {
        self.names.contains(name)
    }

    pub(crate) fn new() -> Self {
        Self::default()
    }

    pub(crate) fn append(&mut self, more: &mut ApiVec<P>) {
        self.extend(more.apis.drain(..))
    }

    pub(crate) fn extend(&mut self, it: impl Iterator<Item = Api<P>>) {
        // Could be optimized in future
        for api in it {
            self.push(api)
        }
    }

    pub(crate) fn iter(&self) -> impl Iterator<Item = &Api<P>> {
        self.apis.iter()
    }

    pub(crate) fn into_iter(self) -> impl Iterator<Item = Api<P>> {
        self.apis.into_iter()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.apis.is_empty()
    }

    pub fn retain<F>(&mut self, f: F)
    where
        F: FnMut(&Api<P>) -> bool,
    {
        self.apis.retain(f);
        self.names.clear();
        self.names
            .extend(self.apis.iter().map(|api| api.name()).cloned());
    }
}

impl<P: AnalysisPhase> Default for ApiVec<P> {
    fn default() -> Self {
        Self {
            apis: Default::default(),
            names: Default::default(),
        }
    }
}

impl<P: AnalysisPhase> FromIterator<Api<P>> for ApiVec<P> {
    fn from_iter<I: IntoIterator<Item = Api<P>>>(iter: I) -> Self {
        let mut this = ApiVec::new();
        for i in iter {
            // Could be optimized in future
            this.push(i);
        }
        this
    }
}
<commit_msg>Add a comment.<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use crate::{
    conversion::{api::ApiName, convert_error::ErrorContext, ConvertError},
    types::QualifiedName,
};

use super::api::{AnalysisPhase, Api};

/// Newtype wrapper for a list of APIs, which enforces the invariant
/// that each API has a unique name.
///
/// Specifically, each API should have a unique [`QualifiedName`] which is kept
/// within an [`ApiName`]. The [`QualifiedName`] is used to refer to this API
/// from others, e.g. to represent edges in the graph used for garbage collection,
/// so that's why this uniqueness is so important.
///
/// At present, this type also refuses to allow mutation of an API once it
/// has been added to a set. This is because the autocxx engine is
/// fundamentally organized into lots of analysis phases, each one _adding_
/// fields rather than mutating earlier fields. The idea here is that it's
/// impossible for stupid future maintainers (i.e. me) to make errors by
/// referring to fields before they're filled in. If a field exists, it's
/// correct.
///
/// While this is currently the case, it's possible that in future we could
/// see legitimate reasons to break this latter invariant and allow mutation
/// of APIs within an existing `ApiVec`. But it's extremely important that
/// the naming-uniqueness-invariant remains, so any such mutation should
/// allow mutation only of other fields, not the name.
pub(crate) struct ApiVec<P: AnalysisPhase> {
    apis: Vec<Api<P>>,
    names: HashSet<QualifiedName>,
}

impl<P: AnalysisPhase> ApiVec<P> {
    pub(crate) fn push(&mut self, api: Api<P>) {
        let name = api.name();
        let already_contains = self.already_contains(name);
        if already_contains {
            if api.discard_duplicates() {
                log::info!("Discarding duplicate API for {}", name);
            } else {
                panic!(
                    "Already have an API with that name: {}. API was {:?}",
                    name, api
                );
            }
        }
        self.names.insert(name.clone());
        self.apis.push(api)
    }

    pub(crate) fn push_eliminating_duplicates(&mut self, api: Api<P>) {
        let name = api.name();
        let already_contains = self.already_contains(name);
        if already_contains {
            log::info!(
                "Duplicate API for {} - removing all of them and replacing with an IgnoredItem.",
                name
            );
            self.retain(|api| api.name() != name);
            self.push(Api::IgnoredItem {
                name: ApiName::new_from_qualified_name(name.clone()),
                err: ConvertError::DuplicateItemsFoundInParsing,
                ctx: ErrorContext::Item(name.get_final_ident()),
            })
        } else {
            self.names.insert(name.clone());
            self.apis.push(api)
        }
    }

    fn already_contains(&self, name: &QualifiedName) -> bool {
        self.names.contains(name)
    }

    pub(crate) fn new() -> Self {
        Self::default()
    }

    pub(crate) fn append(&mut self, more: &mut ApiVec<P>) {
        self.extend(more.apis.drain(..))
    }

    pub(crate) fn extend(&mut self, it: impl Iterator<Item = Api<P>>) {
        // Could be optimized in future
        for api in it {
            self.push(api)
        }
    }

    pub(crate) fn iter(&self) -> impl Iterator<Item = &Api<P>> {
        self.apis.iter()
    }

    pub(crate) fn into_iter(self) -> impl Iterator<Item = Api<P>> {
        self.apis.into_iter()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.apis.is_empty()
    }

    pub fn retain<F>(&mut self, f: F)
    where
        F: FnMut(&Api<P>) -> bool,
    {
        self.apis.retain(f);
        self.names.clear();
        self.names
            .extend(self.apis.iter().map(|api| api.name()).cloned());
    }
}

impl<P: AnalysisPhase> Default for ApiVec<P> {
    fn default() -> Self {
        Self {
            apis: Default::default(),
            names: Default::default(),
        }
    }
}

impl<P: AnalysisPhase> FromIterator<Api<P>> for ApiVec<P> {
    fn from_iter<I: IntoIterator<Item = Api<P>>>(iter: I) -> Self {
        let mut this = ApiVec::new();
        for i in iter {
            // Could be optimized in future
            this.push(i);
        }
        this
    }
}
<|endoftext|>


<commit_before>// rustfmt/main.rs

use std::io;
use std::str;
use syntax::parse::lexer::{StringReader, TokenAndSpan};
use syntax::parse::lexer;
use syntax::parse::token::Token;
use syntax::parse::token::keywords;
use syntax::parse::token;
use syntax::parse;

static TAB_WIDTH: i32 = 4;

enum ProductionToParse {
    MatchProduction,
    UseProduction,
    BracesProduction,
    ParenthesesProduction,
}

struct LineToken {
    token_and_span: TokenAndSpan,
    x_pos: i32,
}

impl LineToken {
    fn new(token_and_span: TokenAndSpan) -> LineToken {
        LineToken {
            token_and_span: token_and_span,
            x_pos: 0,
        }
    }

    fn whitespace_needed_after(&self, next: &LineToken) -> bool {
        match (&self.token_and_span.tok, &next.token_and_span.tok) {
            (&token::IDENT(..), &token::IDENT(..)) => true,
            (&token::IDENT(..), &token::NOT)
                    if !token::is_any_keyword(&self.token_and_span.tok) => {
                // Macros.
                false
            }

            (&token::IDENT(..), _) if
                    token::is_keyword(keywords::If, &self.token_and_span.tok) ||
                    token::is_keyword(keywords::As, &self.token_and_span.tok) ||
                    token::is_keyword(keywords::Match, &self.token_and_span.tok) => {
                true
            }
            (_, &token::IDENT(..))
                    if token::is_keyword(keywords::If, &next.token_and_span.tok) => {
                true
            }

            (&token::COLON, _) => true,
            (&token::COMMA, _) => true,
            (&token::EQ, _) | (_, &token::EQ) => true,
            (&token::LT, _) | (_, &token::LT) => true,
            (&token::LE, _) | (_, &token::LE) => true,
            (&token::EQEQ, _) | (_, &token::EQEQ) => true,
            (&token::NE, _) | (_, &token::NE) => true,
            (&token::GE, _) | (_, &token::GE) => true,
            (&token::GT, _) | (_, &token::GT) => true,
            (&token::ANDAND, _) | (_, &token::ANDAND) => true,
            (&token::OROR, _) | (_, &token::OROR) => true,
            (&token::TILDE, _) | (_, &token::TILDE) => true,

            (&token::LPAREN, _) => false,
            (_, &token::RPAREN) => false,
            (&token::BINOP(token::AND), _) => false,

            (&token::BINOP(_), _) | (_, &token::BINOP(_)) => true,
            (&token::BINOPEQ(_), _) | (_, &token::BINOPEQ(_)) => true,

            (&token::MOD_SEP, _) | (_, &token::MOD_SEP) => false,

            (&token::RARROW, _) | (_, &token::RARROW) => true,
            (&token::FAT_ARROW, _) | (_, &token::FAT_ARROW) => true,
            (&token::LBRACE, _) | (_, &token::LBRACE) => true,
            (&token::RBRACE, _) | (_, &token::RBRACE) => true,
            _ => false,
        }
    }

    fn length(&self) -> i32 {
        token::to_str(&self.token_and_span.tok).len() as i32
    }

    fn preindentation(&self) -> i32 {
        match self.token_and_span.tok {
            token::RBRACE => -TAB_WIDTH,
            _ => 0,
        }
    }
}

struct LogicalLine {
    tokens: Vec<LineToken>,
}

impl LogicalLine {
    fn new() -> LogicalLine {
        LogicalLine {
            tokens: Vec::new(),
        }
    }

    fn layout(&mut self, mut x_pos: i32) {
        if self.tokens.len() == 0 {
            return
        }

        for i in range(0, self.tokens.len()) {
            self.tokens.get_mut(i).x_pos = x_pos;
            x_pos += self.tokens.get(i).length();

            if i < self.tokens.len() - 1 &&
                    self.tokens.get(i).whitespace_needed_after(self.tokens.get(i + 1)) {
                x_pos += 1;
            }
        }
    }

    fn whitespace_after(&self, index: uint) -> i32 {
        if self.tokens.len() <= 1 || index >= self.tokens.len() - 1 {
            return 0
        }

        self.tokens.get(index + 1).x_pos - (self.tokens.get(index).x_pos +
                                            self.tokens.get(index).length())
    }

    fn postindentation(&self) -> i32 {
        match self.tokens.as_slice().last() {
            None => 0,
            Some(line_token) => {
                match line_token.token_and_span.tok {
                    token::LBRACE => TAB_WIDTH,
                    _ => 0,
                }
            }
        }
    }
}

struct Formatter<'a> {
    lexer: StringReader<'a>,
    indent: i32,
    logical_line: LogicalLine,
    last_token: Token,
    newline_after_comma: bool,
    newline_after_brace: bool,
}

impl<'a> Formatter<'a> {
    fn new<'a>(lexer: StringReader<'a>) -> Formatter<'a> {
        Formatter {
            lexer: lexer,
            indent: 0,
            logical_line: LogicalLine::new(),
            last_token: token::SEMI,
            newline_after_comma: false,
            newline_after_brace: true,
        }
    }

    fn token_ends_logical_line(&self, line_token: &LineToken) -> bool {
        match line_token.token_and_span.tok {
            token::SEMI | token::RBRACE => true,
            token::COMMA => self.newline_after_comma,
            token::LBRACE => self.newline_after_brace,
            _ => false,
        }
    }

    fn token_starts_logical_line(&self, line_token: &LineToken) -> bool {
        match line_token.token_and_span.tok {
            token::RBRACE => self.newline_after_brace,
            _ => false,
        }
    }

    fn parse_tokens_up_to(&mut self, pred: |&token::Token| -> bool) -> bool {
        while self.next_token() {
            if pred(&self.last_token) {
                return true;
            }
        }
        return false;
    }

    fn parse_productions_up_to(&mut self, pred: |&token::Token| -> bool) -> bool {
        while self.next_token() {
            if pred(&self.last_token) {
                return true;
            }
            self.parse_production();
        }
        return false;
    }

    fn parse_match(&mut self) -> bool {
        // We've already parsed the keyword. Parse until we find a `{`.
        if !self.parse_tokens_up_to(|token| *token == token::LBRACE) {
            return false;
        }

        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = true;

        if !self.parse_productions_up_to(|token| *token == token::RBRACE) {
            return false;
        }

        self.newline_after_comma = old_newline_after_comma_setting;
        return true;
    }

    fn parse_use(&mut self) -> bool {
        let old_newline_after_brace_setting = self.newline_after_brace;
        self.newline_after_brace = false;

        // We've already parsed the keyword. Parse until we find a `{`.
        if !self.parse_tokens_up_to(|token| *token == token::LBRACE || *token == token::SEMI) {
            return false;
        }

        if self.last_token == token::LBRACE {
            let old_newline_after_comma_setting = self.newline_after_comma;
            self.newline_after_comma = false;

            if !self.parse_productions_up_to(|token| *token == token::RBRACE) {
                return false;
            }

            self.newline_after_comma = old_newline_after_comma_setting;
        }

        self.newline_after_brace = old_newline_after_brace_setting;
        return true;
    }

    fn parse_braces(&mut self) -> bool {
        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = true;

        // We've already parsed the '{'. Parse until we find a '}'.
        let result = self.parse_productions_up_to(|token| *token == token::RBRACE);

        self.newline_after_comma = old_newline_after_comma_setting;
        return result;
    }

    fn parse_parentheses(&mut self) -> bool {
        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = false;

        // We've already parsed the '('. Parse until we find a ')'.
        let result = self.parse_productions_up_to(|token| *token == token::RPAREN);

        self.newline_after_comma = old_newline_after_comma_setting;
        return result;
    }

    fn parse_production(&mut self) -> bool {
        let production_to_parse;
        match self.last_token {
            token::IDENT(..) if token::is_keyword(keywords::Match, &self.last_token) => {
                production_to_parse = MatchProduction;
            }
            token::IDENT(..) if token::is_keyword(keywords::Use, &self.last_token) => {
                production_to_parse = UseProduction;
            }
            token::LBRACE => production_to_parse = BracesProduction,
            token::LPAREN => production_to_parse = ParenthesesProduction,
            _ => return true,
        }

        match production_to_parse {
            MatchProduction => return self.parse_match(),
            UseProduction => return self.parse_use(),
            BracesProduction => return self.parse_braces(),
            ParenthesesProduction => return self.parse_parentheses(),
        }
    }

    fn next_token(&mut self) -> bool {
        use syntax::parse::lexer::Reader;

        loop {
            if self.lexer.is_eof() {
                return false;
            }

            let last_token = self.lexer.peek();
            self.last_token = last_token.tok.clone();
            let line_token = LineToken::new(last_token);
            if self.token_starts_logical_line(&line_token) && self.logical_line.tokens.len() > 0 {
                self.flush_line();
                continue;
            }

            if self.logical_line.tokens.len() == 0 {
                self.indent += line_token.preindentation();
            }

            drop(self.lexer.next_token());
            let token_ends_logical_line = self.token_ends_logical_line(&line_token);
            self.logical_line.tokens.push(line_token);
            if token_ends_logical_line {
                self.flush_line();
            }

            return true;
        }
    }

    fn flush_line(&mut self) {
        self.logical_line.layout(self.indent);

        for _ in range(0, self.indent) {
            print!(" ");
        }
        for i in range(0, self.logical_line.tokens.len()) {
            print!("{}", token::to_str(&self.logical_line.tokens.get(i).token_and_span.tok));
            for _ in range(0, self.logical_line.whitespace_after(i)) {
                print!(" ");
            }
        }
        println!("");

        self.indent += self.logical_line.postindentation();
        self.logical_line = LogicalLine::new();
    }
}

#[main]
pub fn main() {
    let source = io::stdin().read_to_end().unwrap();
    let source = str::from_utf8(source.as_slice()).unwrap();

    let session = parse::new_parse_sess();
    let filemap = parse::string_to_filemap(&session, source.to_strbuf(), "<stdin>".to_strbuf());
    let lexer = lexer::new_string_reader(&session.span_diagnostic, filemap);
    let mut formatter = Formatter::new(lexer);

    while formatter.next_token() {
        formatter.parse_production();
    }
}

<commit_msg>get it building on latest nightly<commit_after>// rustfmt/main.rs

use std::io;
use std::str;
use syntax::parse::lexer::{StringReader, TokenAndSpan};
use syntax::parse::lexer;
use syntax::parse::token::Token;
use syntax::parse::token::keywords;
use syntax::parse::token;
use syntax::parse;

static TAB_WIDTH: i32 = 4;

enum ProductionToParse {
    MatchProduction,
    UseProduction,
    BracesProduction,
    ParenthesesProduction,
}

struct LineToken {
    token_and_span: TokenAndSpan,
    x_pos: i32,
}

impl LineToken {
    fn new(token_and_span: TokenAndSpan) -> LineToken {
        LineToken {
            token_and_span: token_and_span,
            x_pos: 0,
        }
    }

    fn whitespace_needed_after(&self, next: &LineToken) -> bool {
        match (&self.token_and_span.tok, &next.token_and_span.tok) {
            (&token::IDENT(..), &token::IDENT(..)) => true,
            (&token::IDENT(..), &token::NOT)
                    if !token::is_any_keyword(&self.token_and_span.tok) => {
                // Macros.
                false
            }

            (&token::IDENT(..), _) if
                    token::is_keyword(keywords::If, &self.token_and_span.tok) ||
                    token::is_keyword(keywords::As, &self.token_and_span.tok) ||
                    token::is_keyword(keywords::Match, &self.token_and_span.tok) => {
                true
            }
            (_, &token::IDENT(..))
                    if token::is_keyword(keywords::If, &next.token_and_span.tok) => {
                true
            }

            (&token::COLON, _) => true,
            (&token::COMMA, _) => true,
            (&token::EQ, _) | (_, &token::EQ) => true,
            (&token::LT, _) | (_, &token::LT) => true,
            (&token::LE, _) | (_, &token::LE) => true,
            (&token::EQEQ, _) | (_, &token::EQEQ) => true,
            (&token::NE, _) | (_, &token::NE) => true,
            (&token::GE, _) | (_, &token::GE) => true,
            (&token::GT, _) | (_, &token::GT) => true,
            (&token::ANDAND, _) | (_, &token::ANDAND) => true,
            (&token::OROR, _) | (_, &token::OROR) => true,
            (&token::TILDE, _) | (_, &token::TILDE) => true,

            (&token::LPAREN, _) => false,
            (_, &token::RPAREN) => false,
            (&token::BINOP(token::AND), _) => false,

            (&token::BINOP(_), _) | (_, &token::BINOP(_)) => true,
            (&token::BINOPEQ(_), _) | (_, &token::BINOPEQ(_)) => true,

            (&token::MOD_SEP, _) | (_, &token::MOD_SEP) => false,

            (&token::RARROW, _) | (_, &token::RARROW) => true,
            (&token::FAT_ARROW, _) | (_, &token::FAT_ARROW) => true,
            (&token::LBRACE, _) | (_, &token::LBRACE) => true,
            (&token::RBRACE, _) | (_, &token::RBRACE) => true,
            _ => false,
        }
    }

    fn length(&self) -> i32 {
        token::to_str(&self.token_and_span.tok).len() as i32
    }

    fn preindentation(&self) -> i32 {
        match self.token_and_span.tok {
            token::RBRACE => -TAB_WIDTH,
            _ => 0,
        }
    }
}

struct LogicalLine {
    tokens: Vec<LineToken>,
}

impl LogicalLine {
    fn new() -> LogicalLine {
        LogicalLine {
            tokens: Vec::new(),
        }
    }

    fn layout(&mut self, mut x_pos: i32) {
        if self.tokens.len() == 0 {
            return
        }

        for i in range(0, self.tokens.len()) {
            self.tokens.get_mut(i).x_pos = x_pos;
            x_pos += self.tokens.get(i).length();

            if i < self.tokens.len() - 1 &&
                    self.tokens.get(i).whitespace_needed_after(self.tokens.get(i + 1)) {
                x_pos += 1;
            }
        }
    }

    fn whitespace_after(&self, index: uint) -> i32 {
        if self.tokens.len() <= 1 || index >= self.tokens.len() - 1 {
            return 0
        }

        self.tokens.get(index + 1).x_pos - (self.tokens.get(index).x_pos +
                                            self.tokens.get(index).length())
    }

    fn postindentation(&self) -> i32 {
        match self.tokens.as_slice().last() {
            None => 0,
            Some(line_token) => {
                match line_token.token_and_span.tok {
                    token::LBRACE => TAB_WIDTH,
                    _ => 0,
                }
            }
        }
    }
}

struct Formatter<'a> {
    lexer: StringReader<'a>,
    indent: i32,
    logical_line: LogicalLine,
    last_token: Token,
    newline_after_comma: bool,
    newline_after_brace: bool,
}

impl<'a> Formatter<'a> {
    fn new<'a>(lexer: StringReader<'a>) -> Formatter<'a> {
        Formatter {
            lexer: lexer,
            indent: 0,
            logical_line: LogicalLine::new(),
            last_token: token::SEMI,
            newline_after_comma: false,
            newline_after_brace: true,
        }
    }

    fn token_ends_logical_line(&self, line_token: &LineToken) -> bool {
        match line_token.token_and_span.tok {
            token::SEMI | token::RBRACE => true,
            token::COMMA => self.newline_after_comma,
            token::LBRACE => self.newline_after_brace,
            _ => false,
        }
    }

    fn token_starts_logical_line(&self, line_token: &LineToken) -> bool {
        match line_token.token_and_span.tok {
            token::RBRACE => self.newline_after_brace,
            _ => false,
        }
    }

    fn parse_tokens_up_to(&mut self, pred: |&token::Token| -> bool) -> bool {
        while self.next_token() {
            if pred(&self.last_token) {
                return true;
            }
        }
        return false;
    }

    fn parse_productions_up_to(&mut self, pred: |&token::Token| -> bool) -> bool {
        while self.next_token() {
            if pred(&self.last_token) {
                return true;
            }
            self.parse_production();
        }
        return false;
    }

    fn parse_match(&mut self) -> bool {
        // We've already parsed the keyword. Parse until we find a `{`.
        if !self.parse_tokens_up_to(|token| *token == token::LBRACE) {
            return false;
        }

        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = true;

        if !self.parse_productions_up_to(|token| *token == token::RBRACE) {
            return false;
        }

        self.newline_after_comma = old_newline_after_comma_setting;
        return true;
    }

    fn parse_use(&mut self) -> bool {
        let old_newline_after_brace_setting = self.newline_after_brace;
        self.newline_after_brace = false;

        // We've already parsed the keyword. Parse until we find a `{`.
        if !self.parse_tokens_up_to(|token| *token == token::LBRACE || *token == token::SEMI) {
            return false;
        }

        if self.last_token == token::LBRACE {
            let old_newline_after_comma_setting = self.newline_after_comma;
            self.newline_after_comma = false;

            if !self.parse_productions_up_to(|token| *token == token::RBRACE) {
                return false;
            }

            self.newline_after_comma = old_newline_after_comma_setting;
        }

        self.newline_after_brace = old_newline_after_brace_setting;
        return true;
    }

    fn parse_braces(&mut self) -> bool {
        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = true;

        // We've already parsed the '{'. Parse until we find a '}'.
        let result = self.parse_productions_up_to(|token| *token == token::RBRACE);

        self.newline_after_comma = old_newline_after_comma_setting;
        return result;
    }

    fn parse_parentheses(&mut self) -> bool {
        let old_newline_after_comma_setting = self.newline_after_comma;
        self.newline_after_comma = false;

        // We've already parsed the '('. Parse until we find a ')'.
        let result = self.parse_productions_up_to(|token| *token == token::RPAREN);

        self.newline_after_comma = old_newline_after_comma_setting;
        return result;
    }

    fn parse_production(&mut self) -> bool {
        let production_to_parse;
        match self.last_token {
            token::IDENT(..) if token::is_keyword(keywords::Match, &self.last_token) => {
                production_to_parse = MatchProduction;
            }
            token::IDENT(..) if token::is_keyword(keywords::Use, &self.last_token) => {
                production_to_parse = UseProduction;
            }
            token::LBRACE => production_to_parse = BracesProduction,
            token::LPAREN => production_to_parse = ParenthesesProduction,
            _ => return true,
        }

        match production_to_parse {
            MatchProduction => return self.parse_match(),
            UseProduction => return self.parse_use(),
            BracesProduction => return self.parse_braces(),
            ParenthesesProduction => return self.parse_parentheses(),
        }
    }

    fn next_token(&mut self) -> bool {
        use syntax::parse::lexer::Reader;

        loop {
            if self.lexer.is_eof() {
                return false;
            }

            let last_token = self.lexer.peek();
            self.last_token = last_token.tok.clone();
            let line_token = LineToken::new(last_token);
            if self.token_starts_logical_line(&line_token) && self.logical_line.tokens.len() > 0 {
                self.flush_line();
                continue;
            }

            if self.logical_line.tokens.len() == 0 {
                self.indent += line_token.preindentation();
            }

            drop(self.lexer.next_token());
            let token_ends_logical_line = self.token_ends_logical_line(&line_token);
            self.logical_line.tokens.push(line_token);
            if token_ends_logical_line {
                self.flush_line();
            }

            return true;
        }
    }

    fn flush_line(&mut self) {
        self.logical_line.layout(self.indent);

        for _ in range(0, self.indent) {
            print!(" ");
        }
        for i in range(0, self.logical_line.tokens.len()) {
            print!("{}", token::to_str(&self.logical_line.tokens.get(i).token_and_span.tok));
            for _ in range(0, self.logical_line.whitespace_after(i)) {
                print!(" ");
            }
        }
        println!("");

        self.indent += self.logical_line.postindentation();
        self.logical_line = LogicalLine::new();
    }
}

#[main]
pub fn main() {
    let source = io::stdin().read_to_end().unwrap();
    let source = str::from_utf8(source.as_slice()).unwrap();

    let session = parse::new_parse_sess();
    let filemap = parse::string_to_filemap(&session, source.to_string(), "<stdin>".to_string());
    let lexer = lexer::StringReader::new(&session.span_diagnostic, filemap);
    let mut formatter = Formatter::new(lexer);

    while formatter.next_token() {
        formatter.parse_production();
    }
}

<|endoftext|>