<commit_before>//! A stably addressed token buffer supporting efficient traversal based on a
//! cheaply copyable cursor.
//!
//! *This module is available only if Syn is built with the `"parsing"` feature.*

// This module is heavily commented as it contains most of the unsafe code in
// Syn, and caution should be used when editing it. The public-facing interface
// is 100% safe but the implementation is fragile internally.

#[cfg(all(
    not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
    feature = "proc-macro"
))]
use crate::proc_macro as pm;
use crate::Lifetime;
use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
use std::marker::PhantomData;
use std::ptr;
use std::slice;

/// Internal type which is used instead of `TokenTree` to represent a token tree
/// within a `TokenBuffer`.
enum Entry {
    // Mimicking types from proc-macro.
    Group(Group, TokenBuffer),
    Ident(Ident),
    Punct(Punct),
    Literal(Literal),
    // End entries contain a raw pointer to the entry from the containing
    // token tree, or null if this is the outermost level.
    End(*const Entry),
}

/// A buffer that can be efficiently traversed multiple times, unlike
/// `TokenStream` which requires a deep copy in order to traverse more than
/// once.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct TokenBuffer {
    // NOTE: Do not implement clone on this - there are raw pointers inside
    // these entries which will be messed up. Moving the `TokenBuffer` itself is
    // safe as the data pointed to won't be moved.
    ptr: *const Entry,
    len: usize,
}

impl Drop for TokenBuffer {
    fn drop(&mut self) {
        unsafe {
            let slice = slice::from_raw_parts_mut(self.ptr as *mut Entry, self.len);
            let _ = Box::from_raw(slice);
        }
    }
}

impl TokenBuffer {
    // NOTE: Do not mutate the Vec returned from this function once it returns;
    // the address of its backing memory must remain stable.
    fn inner_new(stream: TokenStream, up: *const Entry) -> TokenBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut groups = Vec::new();
        for tt in stream {
            match tt {
                TokenTree::Ident(sym) => {
                    entries.push(Entry::Ident(sym));
                }
                TokenTree::Punct(op) => {
                    entries.push(Entry::Punct(op));
                }
                TokenTree::Literal(l) => {
                    entries.push(Entry::Literal(l));
                }
                TokenTree::Group(g) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarily.
                    groups.push((entries.len(), g));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let entries = entries.into_boxed_slice();
        let len = entries.len();
        // Convert our boxed slice into a pointer to the first element early,
        // to avoid invalidating pointers into this slice when we move the Box
        // see https://github.com/rust-lang/unsafe-code-guidelines/issues/326
        let entries = Box::into_raw(entries) as *mut Entry;
        for (idx, group) in groups {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let group_up = unsafe { entries.add(idx + 1) };

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(group.stream(), group_up);
            unsafe { *entries.add(idx) = Entry::Group(group, inner) };
        }

        TokenBuffer { ptr: entries, len }
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro::TokenStream`.
    ///
    /// *This method is available only if Syn is built with both the `"parsing"` and
    /// `"proc-macro"` features.*
    #[cfg(all(
        not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
        feature = "proc-macro"
    ))]
    pub fn new(stream: pm::TokenStream) -> Self {
        Self::new2(stream.into())
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro2::TokenStream`.
    pub fn new2(stream: TokenStream) -> Self {
        Self::inner_new(stream, ptr::null())
    }

    /// Creates a cursor referencing the first token in the buffer and able to
    /// traverse until the end of the buffer.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(self.ptr, self.ptr.add(self.len - 1)) }
    }
}

/// A cheaply copyable cursor into a `TokenBuffer`.
///
/// This cursor holds a shared reference into the immutable data which is used
/// internally to represent a `TokenStream`, and can be efficiently manipulated
/// and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `TokenBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct Cursor<'a> {
    // The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    // This is the only `Entry::End(..)` object which this cursor is allowed to
    // point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    // Cursor is covariant in 'a. This field ensures that our pointers are still
    // valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Creates a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Ident` is a reference into a thread-local table). This is because
        // this entry never includes a `Ident` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr,
            scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// While the cursor is looking at a `None`-delimited group, move it to look
    /// at the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        while let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == Delimiter::None {
                // NOTE: We call `Cursor::create` here to make sure that
                // situations where we should immediately exit the span after
                // entering it are handled correctly.
                unsafe {
                    *self = Cursor::create(buf.ptr, self.scope);
                }
            } else {
                break;
            }
        }
    }

    /// Checks whether the cursor is currently pointing at the end of its valid
    /// scope.
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a `Group` with the given delimiter, returns
    /// a cursor into that group and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == delim {
                return Some((buf.begin(), group.span(), unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a `Ident`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn ident(mut self) -> Option<(Ident, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Ident(ident) => Some((ident.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Punct`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn punct(mut self) -> Option<(Punct, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() != '\'' => Some((op.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Literal`, return it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Literal, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Literal(lit) => Some((lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Lifetime`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    pub fn lifetime(mut self) -> Option<(Lifetime, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.ident() {
                    Some((ident, rest)) => {
                        let lifetime = Lifetime {
                            apostrophe: op.span(),
                            ident,
                        };
                        Some((lifetime, rest))
                    }
                    None => None,
                }
            }
            _ => None,
        }
    }

    /// Copies all remaining tokens visible from this cursor into a
    /// `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is pointing at a `TokenTree`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    ///
    /// Returns `None` if the cursor has reached the end of its stream.
    ///
    /// This method does not treat `None`-delimited groups as transparent, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match self.entry() {
            Entry::Group(group, _) => group.clone().into(),
            Entry::Literal(lit) => lit.clone().into(),
            Entry::Ident(ident) => ident.clone().into(),
            Entry::Punct(op) => op.clone().into(),
            Entry::End(..) => return None,
        };

        Some((tree, unsafe { self.bump() }))
    }

    /// Returns the `Span` of the current token, or `Span::call_site()` if this
    /// cursor points to eof.
    pub fn span(self) -> Span {
        match self.entry() {
            Entry::Group(group, _) => group.span(),
            Entry::Literal(l) => l.span(),
            Entry::Ident(t) => t.span(),
            Entry::Punct(o) => o.span(),
            Entry::End(..) => Span::call_site(),
        }
    }

    /// Skip over the next token without cloning it. Returns `None` if this
    /// cursor points to eof.
    ///
    /// This method treats `'lifetimes` as a single token.
    pub(crate) fn skip(self) -> Option<Cursor<'a>> {
        match self.entry() {
            Entry::End(..) => None,

            // Treat lifetimes as a single tt for the purposes of 'skip'.
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.entry() {
                    Entry::Ident(_) => Some(unsafe { next.bump() }),
                    _ => Some(next),
                }
            }
            _ => Some(unsafe { self.bump() }),
        }
    }
}

impl<'a> Copy for Cursor<'a> {}

impl<'a> Clone for Cursor<'a> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a> Eq for Cursor<'a> {}

impl<'a> PartialEq for Cursor<'a> {
    fn eq(&self, other: &Self) -> bool {
        let Cursor { ptr, scope, marker } = self;
        let _ = marker;
        *ptr == other.ptr && *scope == other.scope
    }
}

pub(crate) fn same_scope(a: Cursor, b: Cursor) -> bool {
    a.scope == b.scope
}

pub(crate) fn open_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_open(),
        _ => cursor.span(),
    }
}

pub(crate) fn close_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_close(),
        _ => cursor.span(),
    }
}
<commit_msg>Punctuate Box->ptr comment in TokenBuffer constructor<commit_after>//! A stably addressed token buffer supporting efficient traversal based on a
//! cheaply copyable cursor.
//!
//! *This module is available only if Syn is built with the `"parsing"` feature.*

// This module is heavily commented as it contains most of the unsafe code in
// Syn, and caution should be used when editing it. The public-facing interface
// is 100% safe but the implementation is fragile internally.

#[cfg(all(
    not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
    feature = "proc-macro"
))]
use crate::proc_macro as pm;
use crate::Lifetime;
use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
use std::marker::PhantomData;
use std::ptr;
use std::slice;

/// Internal type which is used instead of `TokenTree` to represent a token tree
/// within a `TokenBuffer`.
enum Entry {
    // Mimicking types from proc-macro.
    Group(Group, TokenBuffer),
    Ident(Ident),
    Punct(Punct),
    Literal(Literal),
    // End entries contain a raw pointer to the entry from the containing
    // token tree, or null if this is the outermost level.
    End(*const Entry),
}

/// A buffer that can be efficiently traversed multiple times, unlike
/// `TokenStream` which requires a deep copy in order to traverse more than
/// once.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct TokenBuffer {
    // NOTE: Do not implement clone on this - there are raw pointers inside
    // these entries which will be messed up. Moving the `TokenBuffer` itself is
    // safe as the data pointed to won't be moved.
    ptr: *const Entry,
    len: usize,
}

impl Drop for TokenBuffer {
    fn drop(&mut self) {
        unsafe {
            let slice = slice::from_raw_parts_mut(self.ptr as *mut Entry, self.len);
            let _ = Box::from_raw(slice);
        }
    }
}

impl TokenBuffer {
    // NOTE: Do not mutate the Vec returned from this function once it returns;
    // the address of its backing memory must remain stable.
    fn inner_new(stream: TokenStream, up: *const Entry) -> TokenBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut groups = Vec::new();
        for tt in stream {
            match tt {
                TokenTree::Ident(sym) => {
                    entries.push(Entry::Ident(sym));
                }
                TokenTree::Punct(op) => {
                    entries.push(Entry::Punct(op));
                }
                TokenTree::Literal(l) => {
                    entries.push(Entry::Literal(l));
                }
                TokenTree::Group(g) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarily.
                    groups.push((entries.len(), g));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let entries = entries.into_boxed_slice();
        let len = entries.len();
        // Convert boxed slice into a pointer to the first element early, to
        // avoid invalidating pointers into this slice when we move the Box.
        // See https://github.com/rust-lang/unsafe-code-guidelines/issues/326
        let entries = Box::into_raw(entries) as *mut Entry;
        for (idx, group) in groups {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let group_up = unsafe { entries.add(idx + 1) };

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(group.stream(), group_up);
            unsafe { *entries.add(idx) = Entry::Group(group, inner) };
        }

        TokenBuffer { ptr: entries, len }
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro::TokenStream`.
    ///
    /// *This method is available only if Syn is built with both the `"parsing"` and
    /// `"proc-macro"` features.*
    #[cfg(all(
        not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
        feature = "proc-macro"
    ))]
    pub fn new(stream: pm::TokenStream) -> Self {
        Self::new2(stream.into())
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro2::TokenStream`.
    pub fn new2(stream: TokenStream) -> Self {
        Self::inner_new(stream, ptr::null())
    }

    /// Creates a cursor referencing the first token in the buffer and able to
    /// traverse until the end of the buffer.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(self.ptr, self.ptr.add(self.len - 1)) }
    }
}

/// A cheaply copyable cursor into a `TokenBuffer`.
///
/// This cursor holds a shared reference into the immutable data which is used
/// internally to represent a `TokenStream`, and can be efficiently manipulated
/// and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `TokenBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct Cursor<'a> {
    // The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    // This is the only `Entry::End(..)` object which this cursor is allowed to
    // point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    // Cursor is covariant in 'a. This field ensures that our pointers are still
    // valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Creates a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Ident` is a reference into a thread-local table). This is because
        // this entry never includes a `Ident` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr,
            scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// While the cursor is looking at a `None`-delimited group, move it to look
    /// at the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        while let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == Delimiter::None {
                // NOTE: We call `Cursor::create` here to make sure that
                // situations where we should immediately exit the span after
                // entering it are handled correctly.
                unsafe {
                    *self = Cursor::create(buf.ptr, self.scope);
                }
            } else {
                break;
            }
        }
    }

    /// Checks whether the cursor is currently pointing at the end of its valid
    /// scope.
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a `Group` with the given delimiter, returns
    /// a cursor into that group and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == delim {
                return Some((buf.begin(), group.span(), unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a `Ident`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn ident(mut self) -> Option<(Ident, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Ident(ident) => Some((ident.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Punct`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn punct(mut self) -> Option<(Punct, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() != '\'' => Some((op.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Literal`, return it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Literal, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Literal(lit) => Some((lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Lifetime`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    pub fn lifetime(mut self) -> Option<(Lifetime, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.ident() {
                    Some((ident, rest)) => {
                        let lifetime = Lifetime {
                            apostrophe: op.span(),
                            ident,
                        };
                        Some((lifetime, rest))
                    }
                    None => None,
                }
            }
            _ => None,
        }
    }

    /// Copies all remaining tokens visible from this cursor into a
    /// `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is pointing at a `TokenTree`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    ///
    /// Returns `None` if the cursor has reached the end of its stream.
    ///
    /// This method does not treat `None`-delimited groups as transparent, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match self.entry() {
            Entry::Group(group, _) => group.clone().into(),
            Entry::Literal(lit) => lit.clone().into(),
            Entry::Ident(ident) => ident.clone().into(),
            Entry::Punct(op) => op.clone().into(),
            Entry::End(..) => return None,
        };

        Some((tree, unsafe { self.bump() }))
    }

    /// Returns the `Span` of the current token, or `Span::call_site()` if this
    /// cursor points to eof.
    pub fn span(self) -> Span {
        match self.entry() {
            Entry::Group(group, _) => group.span(),
            Entry::Literal(l) => l.span(),
            Entry::Ident(t) => t.span(),
            Entry::Punct(o) => o.span(),
            Entry::End(..) => Span::call_site(),
        }
    }

    /// Skip over the next token without cloning it. Returns `None` if this
    /// cursor points to eof.
    ///
    /// This method treats `'lifetimes` as a single token.
    pub(crate) fn skip(self) -> Option<Cursor<'a>> {
        match self.entry() {
            Entry::End(..) => None,

            // Treat lifetimes as a single tt for the purposes of 'skip'.
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.entry() {
                    Entry::Ident(_) => Some(unsafe { next.bump() }),
                    _ => Some(next),
                }
            }
            _ => Some(unsafe { self.bump() }),
        }
    }
}

impl<'a> Copy for Cursor<'a> {}

impl<'a> Clone for Cursor<'a> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a> Eq for Cursor<'a> {}

impl<'a> PartialEq for Cursor<'a> {
    fn eq(&self, other: &Self) -> bool {
        let Cursor { ptr, scope, marker } = self;
        let _ = marker;
        *ptr == other.ptr && *scope == other.scope
    }
}

pub(crate) fn same_scope(a: Cursor, b: Cursor) -> bool {
    a.scope == b.scope
}

pub(crate) fn open_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_open(),
        _ => cursor.span(),
    }
}

pub(crate) fn close_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_close(),
        _ => cursor.span(),
    }
}
<|endoftext|>


<commit_before>use std::ops::{Add, Div};
use std::time;

use euclid::{TypedPoint2D, TypedRect, TypedSize2D, TypedTransform2D};

use defs::*;
use math;
use math::*;
use shape::*;

// TODO: Do I need a custom AATransform that manages precision
// somehow?

// TODO: think about how movement will work: this is tricky. When a
// chunk is broken in to pieces, we need to make sure our
// current_shape and target_shape are switched to being ones that are
// in the correct chunk

#[derive(Debug)]
pub struct CameraState {
    pub camera_bounds: TypedRect<f32, DrawSpace>,

    pub current_chunk: TopChunk,
    pub current_neutral_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,
    pub current_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,

    pub target_chunk: TopChunk, // This should always be the parent chunk of the player chunk
    pub target_neutral_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,

    pub current_to_target_path: MonotonePath, // This should stay in sync with the above...
}

impl CameraState {
    fn center<T, U>(rect: &TypedRect<T, U>) -> TypedPoint2D<T, U>
    where
        T: Copy + Add<T, Output = T> + Div<T, Output = T> + From<f32>,
    {
        TypedPoint2D::new(
            rect.origin.x + rect.size.width / T::from(2.0),
            rect.origin.y + rect.size.height / T::from(2.0),
        )
    }

    fn intended_target_chunk(logical_state: &LogicalState) -> TopChunk {
        logical_state
            .universe
            .parent_of(&logical_state.player_chunk)
    }

    fn target_transform_for(
        chunk: &TopChunk,
        logical_state: &LogicalState,
        bounds: TypedRect<f32, DrawSpace>,
    ) -> TypedTransform2D<f32, UniverseSpace, DrawSpace> {
        let chunk_bounds = logical_state
            .universe
            .region_of_chunk(chunk)
            .bounding_box(&logical_state.universe);

        transform::fit_rect_in(
            &chunk_bounds
                .to_f32()
                .inflate(DRAW_PARENT_MARGIN, DRAW_PARENT_MARGIN),
            &bounds,
        )
    }

    pub fn new(
        resolution: &TypedSize2D<u32, ScreenSpace>,
        logical_state: &LogicalState,
    ) -> CameraState {
        let camera_bounds = TypedRect::new(
            TypedPoint2D::zero(),
            TypedSize2D::new(resolution.width as f32, resolution.height as f32),
        );

        let start_chunk = CameraState::intended_target_chunk(&logical_state);
        let transform =
            CameraState::target_transform_for(&start_chunk, logical_state, camera_bounds);

        CameraState {
            camera_bounds: camera_bounds,

            current_chunk: start_chunk.clone(),
            current_transform: transform,
            current_neutral_transform: transform,

            target_chunk: start_chunk.clone(),
            target_neutral_transform: transform,

            current_to_target_path: MonotonePath::Zero,
        }
    }

    pub fn do_zoom(&mut self, logical_state: &LogicalState) -> () {
        let new_target = CameraState::intended_target_chunk(logical_state);
        self.target_chunk = new_target.clone();
        self.target_neutral_transform =
            CameraState::target_transform_for(&new_target, logical_state, self.camera_bounds);

        self.current_to_target_path = self.current_to_target_path.up_target();
    }

    pub fn scale_from_neutral(&self) -> f32 {
        transform::scale(
            &self
                .current_transform
                .post_mul(&self.current_neutral_transform.inverse().unwrap()),
        )
    }

    pub fn current_visual_level(&self) -> f32 {
        self.scale_from_neutral().log(ZOOM_SCALE as f32)
    }

    fn target_to_current_transform(
        &self,
        logical_state: &LogicalState,
    ) -> TypedTransform2D<f32, UniverseSpace, UniverseSpace> {
        self.current_to_target_path
            .as_delta_from(&logical_state.universe, &self.current_chunk)
            .0
            .to_scale_transform()
    }

    fn true_target_transform(
        &self,
        logical_state: &LogicalState,
    ) -> TypedTransform2D<f32, UniverseSpace, DrawSpace> {
        self.target_to_current_transform(logical_state)
            .post_mul(&self.target_neutral_transform)
    }

    pub fn update(&mut self, logical_state: &LogicalState, time_delta: time::Duration) -> () {
        // TODO: This is currently busted but somehow I get away with it?

        self.current_transform = transform::lerp(
            &self.current_transform,
            &self.true_target_transform(logical_state),
            1.0 - (-math::time::duration_to_secs(time_delta) * CAMERA_LERP_SPEED).exp(),
        );

        self.normalise(logical_state);
    }

    fn normalise(&mut self, logical_state: &LogicalState) -> () {
        let scale_from_neutral = self.scale_from_neutral();

        // MUST TODO: This will break if the individual chunks get too
        // big, need to scale by the size of the target/current chunk?
        // println!("{:#?}", scale_from_neutral);

        let need_normalisation = (scale_from_neutral > CAMERA_UPPER_NORMALISE_TRIGGER
            && self.current_to_target_path.zdelta() < 0)
            || (scale_from_neutral < CAMERA_LOWER_NORMALISE_TRIGGER
                && self.current_to_target_path.zdelta() > 0);

        if need_normalisation {
            let (adjustment, new_current_chunk) = self
                .current_to_target_path
                .take(1)
                .as_delta_from(&logical_state.universe, &self.current_chunk);

            self.current_neutral_transform = CameraState::target_transform_for(
                &new_current_chunk,
                logical_state,
                self.camera_bounds,
            );
            self.current_chunk = new_current_chunk;

            self.current_transform = self
                .current_transform
                .pre_mul(&adjustment.invert().to_scale_transform());

            self.current_to_target_path = self.current_to_target_path.drop(1);
        }
    }

    // MUST TODO: This will start to matter when chunks are made
    pub fn recenter(&mut self, logical_state: &LogicalState) -> () {
        // unimplemented!();

        // let old_target_chunk = &self.target_chunk;
        // let new_target_chunk = CameraState::intended_target(logical_state);

        // self.target_chunk = new_target_chunk;

        // let old_current_chunk = &self.current_chunk;
        // let new_current_chunk = CameraState::intended_target(logical_state);

        // self.current_chunk = new_current_chunk;

        //unimplemented!();
    }
}
<commit_msg>Fix transform calculation in camera<commit_after>use std::ops::{Add, Div};
use std::time;

use euclid::{TypedPoint2D, TypedRect, TypedSize2D, TypedTransform2D};

use defs::*;
use math;
use math::*;
use shape::*;

// TODO: Do I need a custom AATransform that manages precision
// somehow?

// TODO: think about how movement will work: this is tricky. When a
// chunk is broken in to pieces, we need to make sure our
// current_shape and target_shape are switched to being ones that are
// in the correct chunk

#[derive(Debug)]
pub struct CameraState {
    pub camera_bounds: TypedRect<f32, DrawSpace>,

    pub current_chunk: TopChunk,
    pub current_neutral_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,
    pub current_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,

    pub target_chunk: TopChunk, // This should always be the parent chunk of the player chunk
    pub target_neutral_transform: TypedTransform2D<f32, UniverseSpace, DrawSpace>,

    pub current_to_target_path: MonotonePath, // This should stay in sync with the above...
}

impl CameraState {
    fn center<T, U>(rect: &TypedRect<T, U>) -> TypedPoint2D<T, U>
    where
        T: Copy + Add<T, Output = T> + Div<T, Output = T> + From<f32>,
    {
        TypedPoint2D::new(
            rect.origin.x + rect.size.width / T::from(2.0),
            rect.origin.y + rect.size.height / T::from(2.0),
        )
    }

    fn intended_target_chunk(logical_state: &LogicalState) -> TopChunk {
        logical_state
            .universe
            .parent_of(&logical_state.player_chunk)
    }

    fn target_transform_for(
        chunk: &TopChunk,
        logical_state: &LogicalState,
        bounds: TypedRect<f32, DrawSpace>,
    ) -> TypedTransform2D<f32, UniverseSpace, DrawSpace> {
        let chunk_bounds = logical_state
            .universe
            .region_of_chunk(chunk)
            .bounding_box(&logical_state.universe);

        transform::fit_rect_in(
            &chunk_bounds
                .to_f32()
                .inflate(DRAW_PARENT_MARGIN, DRAW_PARENT_MARGIN),
            &bounds,
        )
    }

    pub fn new(
        resolution: &TypedSize2D<u32, ScreenSpace>,
        logical_state: &LogicalState,
    ) -> CameraState {
        let camera_bounds = TypedRect::new(
            TypedPoint2D::zero(),
            TypedSize2D::new(resolution.width as f32, resolution.height as f32),
        );

        let start_chunk = CameraState::intended_target_chunk(&logical_state);
        let transform =
            CameraState::target_transform_for(&start_chunk, logical_state, camera_bounds);

        CameraState {
            camera_bounds: camera_bounds,

            current_chunk: start_chunk.clone(),
            current_transform: transform,
            current_neutral_transform: transform,

            target_chunk: start_chunk.clone(),
            target_neutral_transform: transform,

            current_to_target_path: MonotonePath::Zero,
        }
    }

    pub fn do_zoom(&mut self, logical_state: &LogicalState) -> () {
        let new_target = CameraState::intended_target_chunk(logical_state);
        self.target_chunk = new_target.clone();
        self.target_neutral_transform =
            CameraState::target_transform_for(&new_target, logical_state, self.camera_bounds);

        self.current_to_target_path = self.current_to_target_path.up_target();
    }

    pub fn scale_from_neutral(&self) -> f32 {
        transform::scale(
            &self
                .current_transform
                .post_mul(&self.current_neutral_transform.inverse().unwrap()),
        )
    }

    pub fn current_visual_level(&self) -> f32 {
        self.scale_from_neutral().log(ZOOM_SCALE as f32)
    }

    fn target_to_current_transform(
        &self,
        logical_state: &LogicalState,
    ) -> TypedTransform2D<f32, UniverseSpace, UniverseSpace> {
        self.current_to_target_path
            .as_delta_from(&logical_state.universe, &self.current_chunk)
            .0
            .invert()
            .to_scale_transform()
    }

    fn true_target_transform(
        &self,
        logical_state: &LogicalState,
    ) -> TypedTransform2D<f32, UniverseSpace, DrawSpace> {
        self.target_to_current_transform(logical_state)
            .post_mul(&self.target_neutral_transform)
    }

    pub fn update(&mut self, logical_state: &LogicalState, time_delta: time::Duration) -> () {
        // TODO: This is currently busted but somehow I get away with it?

        self.current_transform = transform::lerp(
            &self.current_transform,
            &self.true_target_transform(logical_state),
            1.0 - (-math::time::duration_to_secs(time_delta) * CAMERA_LERP_SPEED).exp(),
        );

        self.normalise(logical_state);
    }

    fn normalise(&mut self, logical_state: &LogicalState) -> () {
        let scale_from_neutral = self.scale_from_neutral();

        // MUST TODO: This will break if the individual chunks get too
        // big, need to scale by the size of the target/current chunk?
        // println!("{:#?}", scale_from_neutral);

        let need_normalisation = (scale_from_neutral > CAMERA_UPPER_NORMALISE_TRIGGER
            && self.current_to_target_path.zdelta() < 0)
            || (scale_from_neutral < CAMERA_LOWER_NORMALISE_TRIGGER
                && self.current_to_target_path.zdelta() > 0);

        if need_normalisation {
            let (adjustment, new_current_chunk) = self
                .current_to_target_path
                .take(1)
                .as_delta_from(&logical_state.universe, &self.current_chunk);

            self.current_neutral_transform = CameraState::target_transform_for(
                &new_current_chunk,
                logical_state,
                self.camera_bounds,
            );
            self.current_chunk = new_current_chunk;

            self.current_transform = self
                .current_transform
                .pre_mul(&adjustment.to_scale_transform());

            self.current_to_target_path = self.current_to_target_path.drop(1);
        }
    }

    // MUST TODO: This will start to matter when chunks are made
    pub fn recenter(&mut self, logical_state: &LogicalState) -> () {
        // unimplemented!();

        // let old_target_chunk = &self.target_chunk;
        // let new_target_chunk = CameraState::intended_target(logical_state);

        // self.target_chunk = new_target_chunk;

        // let old_current_chunk = &self.current_chunk;
        // let new_current_chunk = CameraState::intended_target(logical_state);

        // self.current_chunk = new_current_chunk;

        //unimplemented!();
    }
}
<|endoftext|>


<commit_before>use nalgebra::{dot, BaseFloat, Col, Mat4, Vec3, Vec4};

pub struct Camera {
    fov: f32,
    near: f32,
    far: f32,
    aspect_ratio: f32,
    transform: Mat4<f32>,
}

impl Camera {
    pub fn new(pos: Vec3<f32>, aspect_ratio: f32) -> Self {
        let transform = Mat4::new(1., 0., 0., pos.x,
                                  0., 1., 0., pos.y,
                                  0., 0., 1., pos.z,
                                  0., 0., 0., 1.);
        Camera {
            fov: BaseFloat::frac_pi_2(),
            near: 0.1,
            far: 1024.,
            aspect_ratio: aspect_ratio,
            transform: transform,
        }
    }

    pub fn fov(&self) -> f32 {
        self.fov
    }

    pub fn set_fov(&mut self, fov: f32) {
        self.fov = fov;
        debug!("Camera fov set to {:?}", fov);
    }

    pub fn set_pos(&mut self, pos: &Vec3<f32>) {
        self.transform.set_col(3, Vec4::new(pos.x, pos.y, pos.z, 1.));
        debug!("Camera position set to {:?}", pos);
    }

    pub fn set_aspect_ratio(&mut self, aspect_ratio: f32) {
        self.aspect_ratio = aspect_ratio;
    }

    pub fn projection_matrix(&self) -> Mat4<f32> {
        let n = self.near;
        let f = self.far;

        let y = 1. / (self.fov / 2.).tan();
        let x = y / self.aspect_ratio;
        let a = (f + n) / (n - f);
        let b = (2. * f * n) / (n - f);

        Mat4::new(x,  0., 0.,  0.,
                  0., y,  0.,  0.,
                  0., 0., a,   b,
                  0., 0., -1., 0.)
    }

    pub fn view_matrix(&self) -> Mat4<f32> {
        let t = &self.transform;

        // This is the camera position applied after the rotation
        let p = (-dot(&t.col(3), &t.col(0)),
                 -dot(&t.col(3), &t.col(1)),
                 -dot(&t.col(3), &t.col(2)));

        // We take the inverse of the rotational part of the transform matrix
        // Because this is an orthogonal matrix, we can just take the transpose
        Mat4::new(t[(0, 0)], t[(1, 0)], t[(2, 0)], p.0,
                  t[(0, 1)], t[(1, 1)], t[(2, 1)], p.1,
                  t[(0, 2)], t[(1, 2)], t[(2, 2)], p.2,
                  0.,        0.,        0.,        1.)
    }

    pub fn set_abs_rotation(&mut self, pitch: f32, yaw: f32) {
        let pitch_mat = Mat4::new(1., 0.,           0.,          0.,
                                  0., pitch.cos(), -pitch.sin(), 0.,
                                  0., pitch.sin(),  pitch.cos(), 0.,
                                  0., 0.,           0.,          1.);
        let yaw_mat = Mat4::new(yaw.cos(),  0., yaw.sin(), 0.,
                                0.,         1., 0.,        0.,
                                -yaw.sin(), 0., yaw.cos(), 0.,
                                0.,         0., 0.,        1.);

        let trans_row = self.transform.col(3);
        self.transform = pitch_mat * yaw_mat;
        self.transform.set_col(3, trans_row);

        debug!("Transform set to {:?}", self.transform);
    }

    pub fn rotate(&mut self, pitch: f32, yaw: f32) {
        let pitch_mat = Mat4::new(1., 0.,           0.,          0.,
                                  0., pitch.cos(), -pitch.sin(), 0.,
                                  0., pitch.sin(),  pitch.cos(), 0.,
                                  0., 0.,           0.,          1.);
        let yaw_mat = Mat4::new(yaw.cos(),  0., yaw.sin(), 0.,
                                0.,         1., 0.,        0.,
                                -yaw.sin(), 0., yaw.cos(), 0.,
                                0.,         0., 0.,        1.);

        self.transform = self.transform * pitch_mat * yaw_mat;

        debug!("Transform set to {:?}", self.transform);
    }
}
<commit_msg>Cache view/projection matrices<commit_after>use std::cell::Cell;

use nalgebra::{self, dot, BaseFloat, Col, Mat4, Vec3, Vec4};

pub struct Camera {
    fov: f32,
    near: f32,
    far: f32,
    aspect_ratio: f32,
    transform: Mat4<f32>,
    view_matrix: Cell<Mat4<f32>>,
    proj_matrix: Cell<Mat4<f32>>,
    view_dirty: Cell<bool>,
    proj_dirty: Cell<bool>,
}

impl Camera {
    pub fn new(pos: Vec3<f32>, aspect_ratio: f32) -> Self {
        let transform = Mat4::new(1., 0., 0., pos.x,
                                  0., 1., 0., pos.y,
                                  0., 0., 1., pos.z,
                                  0., 0., 0., 1.);
        Camera {
            fov: BaseFloat::frac_pi_2(),
            near: 0.1,
            far: 1024.,
            aspect_ratio: aspect_ratio,
            transform: transform,
            view_matrix: Cell::new(nalgebra::new_identity(4)),
            proj_matrix: Cell::new(nalgebra::new_identity(4)),
            view_dirty: Cell::new(true),
            proj_dirty: Cell::new(true),
        }
    }

    pub fn fov(&self) -> f32 {
        self.fov
    }

    pub fn set_fov(&mut self, fov: f32) {
        self.proj_dirty.set(true);

        self.fov = fov;
        debug!("Camera fov set to {:?}", fov);
    }

    pub fn set_pos(&mut self, pos: &Vec3<f32>) {
        self.view_dirty.set(true);

        self.transform.set_col(3, Vec4::new(pos.x, pos.y, pos.z, 1.));
        debug!("Camera position set to {:?}", pos);
    }

    pub fn set_aspect_ratio(&mut self, aspect_ratio: f32) {
        self.proj_dirty.set(true);

        self.aspect_ratio = aspect_ratio;
    }

    pub fn set_abs_rotation(&mut self, pitch: f32, yaw: f32) {
        self.view_dirty.set(true);

        let pitch_mat = Mat4::new(1., 0.,           0.,          0.,
                                  0., pitch.cos(), -pitch.sin(), 0.,
                                  0., pitch.sin(),  pitch.cos(), 0.,
                                  0., 0.,           0.,          1.);
        let yaw_mat = Mat4::new(yaw.cos(),  0., yaw.sin(), 0.,
                                0.,         1., 0.,        0.,
                                -yaw.sin(), 0., yaw.cos(), 0.,
                                0.,         0., 0.,        1.);

        let trans_row = self.transform.col(3);
        self.transform = pitch_mat * yaw_mat;
        self.transform.set_col(3, trans_row);

        debug!("Transform set to {:?}", self.transform);
    }

    pub fn rotate(&mut self, pitch: f32, yaw: f32) {
        self.view_dirty.set(true);

        let pitch_mat = Mat4::new(1., 0.,           0.,          0.,
                                  0., pitch.cos(), -pitch.sin(), 0.,
                                  0., pitch.sin(),  pitch.cos(), 0.,
                                  0., 0.,           0.,          1.);
        let yaw_mat = Mat4::new(yaw.cos(),  0., yaw.sin(), 0.,
                                0.,         1., 0.,        0.,
                                -yaw.sin(), 0., yaw.cos(), 0.,
                                0.,         0., 0.,        1.);

        self.transform = self.transform * pitch_mat * yaw_mat;

        debug!("Transform set to {:?}", self.transform);
    }

    pub fn projection_matrix(&self) -> Mat4<f32> {
        if self.proj_dirty.get() {
            let n = self.near;
            let f = self.far;

            let y = 1. / (self.fov / 2.).tan();
            let x = y / self.aspect_ratio;
            let a = (f + n) / (n - f);
            let b = (2. * f * n) / (n - f);

            self.proj_matrix.set(Mat4::new(x,  0., 0.,  0.,
                                           0., y,  0.,  0.,
                                           0., 0., a,   b,
                                           0., 0., -1., 0.))
        }
        self.proj_matrix.get()
    }

    pub fn view_matrix(&self) -> Mat4<f32> {
        if self.view_dirty.get() {
            let t = &self.transform;

            // This is the camera position applied after the rotation
            let p = (-dot(&t.col(3), &t.col(0)),
                     -dot(&t.col(3), &t.col(1)),
                     -dot(&t.col(3), &t.col(2)));

            // We take the inverse of the rotational part of the transform matrix
            // Because this is an orthogonal matrix, we can just take the transpose
            self.view_matrix.set(Mat4::new(t[(0, 0)], t[(1, 0)], t[(2, 0)], p.0,
                                           t[(0, 1)], t[(1, 1)], t[(2, 1)], p.1,
                                           t[(0, 2)], t[(1, 2)], t[(2, 2)], p.2,
                                           0.,        0.,        0.,        1.))
        }
        self.view_matrix.get()
    }
}
<|endoftext|>


<commit_before>extern crate rustc_serialize;
extern crate websocket;

use transport::{WampSender, WampConnector, WebSocket, SerializerType, Serializer};
use message::{WampEncodable, MessageType, EventMessage, EventJoin, new_event_id};
use options::{Options, Details};
use rustc_serialize::Encodable;
use websocket::Message;

use super::{WampResult};

pub struct Client {
    url: String,
    realm: String,
}

pub struct Session <S: WampSender>{
    sender: S,
}

impl <S: WampSender> Session<S> {
    pub fn join(&mut self, realm: String) -> WampResult<()> {
        let join_msg = EventJoin {
            message_type: MessageType::HELLO,
            realm: realm.clone(),
            details: Details::new(),
        };

       self.sender.send(&join_msg)
    }

    pub fn publish<A, K>(&mut self, topic: &str, args: Vec<WampEncodable<A>>, kwargs: K) where A: Encodable, K: Encodable {
        let msg = EventMessage {
            message_type: MessageType::PUBLISH,
            id: new_event_id(),
            topic: topic.to_string(),
            options: Options {id: 1}, // other stuff goes here
            args: args,
            kwargs: kwargs,
        };

        self.sender.send(&msg);
    }
}

impl Client {
    fn new(url: &str, realm: &str) -> Self {
        Client {url: String::from(url), realm: String::from(realm)}
    }

    fn connect(&self) -> WampResult<Session<WebSocket>> {
        println!("starting...");
        let on_message = move |message: Message| (println!("{:?}", message));
        let transport = try!(WebSocket::connect(self.url.clone(),
                                                Serializer::new(SerializerType::JSON),
                                                on_message));
        let mut session = Session {sender: transport};
        try!(session.join(self.realm.clone()));
        Ok(session)
    }
}

#[test]
fn client_naive_connect() {
    let mut session = Client::new("ws://localhost:8080/ws", "realm1").connect().unwrap();
    session.publish::<(), WampEncodable<()>>("com.myapp.topic1",
                                             vec![WampEncodable::i32(5),
                                             WampEncodable::String("hello".to_string())],
                                             WampEncodable::None);
}
<commit_msg>add mutable state to session<commit_after>extern crate rustc_serialize;
extern crate websocket;

use transport::{WampSender, WampConnector, WebSocket, SerializerType, Serializer};
use message::{WampEncodable, MessageType, EventMessage, EventJoin, new_event_id};
use options::{Options, Details};
use rustc_serialize::Encodable;
use websocket::Message;

use super::{WampResult};

pub struct Client {
    url: String,
    realm: String,
}

enum SessionState {
    NotConnected,
    Connected,
}

pub struct Session <S: WampSender>{
    sender: S,
    state: SessionState,
}

impl <S: WampSender> Session<S> {
    pub fn join(&self, realm: String) -> WampResult<()> {
        let join_msg = EventJoin {
            message_type: MessageType::HELLO,
            realm: realm.clone(),
            details: Details::new(),
        };

       self.sender.send(&join_msg)
    }

    pub fn publish<A, K>(&self, topic: &str, args: Vec<WampEncodable<A>>, kwargs: K) where A: Encodable, K: Encodable {
        let msg = EventMessage {
            message_type: MessageType::PUBLISH,
            id: new_event_id(),
            topic: topic.to_string(),
            options: Options {id: 1}, // other stuff goes here
            args: args,
            kwargs: kwargs,
        };

        self.sender.send(&msg);
    }
}

impl Client {
    fn new(url: &str, realm: &str) -> Self {
        Client {url: String::from(url), realm: String::from(realm)}
    }

    fn connect(&self) -> WampResult<Session<WebSocket>> {
        println!("starting...");

        let mut state = SessionState::NotConnected;
        let on_message = |message: Message| {
        };

        let transport = try!(WebSocket::connect(self.url.clone(),
                                                Serializer::new(SerializerType::JSON),
                                                on_message));
        let mut session = Session {sender: transport, state: state};
        try!(session.join(self.realm.clone()));
        Ok(session)
    }
}

#[test]
fn client_naive_connect() {
    let mut session = Client::new("ws://localhost:8080/ws", "realm1").connect().unwrap();
    session.publish::<(), WampEncodable<()>>("com.myapp.topic1",
                                             vec![WampEncodable::i32(5),
                                             WampEncodable::String("hello".to_string())],
                                             WampEncodable::None);
}
<|endoftext|>


<commit_before>//! The HTTP Client that perform query

use std::io::{self, Write};
use std::fmt::Arguments;

use super::request::Request;
use super::http;
use super::response::{Response, ResponseBuilder};
use super::results::CabotResult;

/// Perform the http query
pub struct Client {
    verbose: bool
}

impl Client {

    /// Construct a new `Client`
    pub fn new() -> Self {
        Client{verbose: false}
    }

    /// Execute the query [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html).
    pub fn execute(&self, request: &Request) -> CabotResult<Response> {
        let mut out = CabotLibWrite::new();
        http::http_query(&request, &mut out, self.verbose)?;
        out.response()
    }
}

struct CabotLibWrite {
    response_builder: ResponseBuilder
}

impl CabotLibWrite {

    pub fn new() -> Self {
        CabotLibWrite{response_builder: ResponseBuilder::new()}
    }

    pub fn response(&self) -> CabotResult<Response> {
        self.response_builder.build()
    }
}

impl Write for CabotLibWrite {

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        info!("Parsing http response");
        let mut response = Vec::with_capacity(buf.len());
        response.extend_from_slice(buf);
        let response = unsafe { String::from_utf8_unchecked(response) };
        let response: Vec<&str> = response.splitn(2, "\r\n\r\n").collect();
        let header_len = response.get(0).unwrap().len();
        let mut headers: Vec<&str> = response.get(0).unwrap().split("\r\n").collect();
        let mut builder = ResponseBuilder::new();
        let status_line = headers.remove(0);
        debug!("Adding status line {}", status_line);
        builder = builder.set_status_line(status_line);
        let mut iter_header = headers.iter().peekable();
        let header = iter_header.next();
        if header.is_some() {
            let buf = header.unwrap();
            let mut header = String::with_capacity(buf.len() * 2);
            header.push_str(buf);
            loop {
                {
                    let buf = iter_header.peek();
                    if buf.is_none() {
                        builder = builder.add_header(header.as_str());
                        break;
                    }
                    let buf = buf.unwrap();
                    if buf.starts_with(" ") || buf.starts_with("\t") {
                        debug!("Obsolete line folded header reveived in {}",
                               header);
                        header.push_str(" ");
                        header.push_str(buf.trim_left());
                    }
                    else {
                        debug!("Adding header {}", header);
                        builder = builder.add_header(header.as_str());
                        header.clear();
                        header.push_str(buf);
                    }
                }
                let _  = iter_header.next();
            }
        }
        let body = if (header_len + 4) < buf.len() {
            &buf[(header_len + 4)..buf.len()]
        }
        else {
            &buf[..]
        };
        //debug!("Adding body {:?}", body);
        builder = builder.set_body(body);
        self.response_builder = builder;
        // debug!("Response Builder - {:?}", self.response_builder);
        Ok(())
    }

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

    // Don't implemented unused method

    fn write(&mut self, _: &[u8]) -> io::Result<usize> {
        Err(io::Error::new(io::ErrorKind::Other, "Not Implemented"))
    }

    fn write_fmt(&mut self, _: Arguments) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "Not Implemented"))
    }

}

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

    #[test]
    fn test_build_http_response_from_string() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "Content-Type: text/plain",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &["Content-Type: text/plain", "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

    #[test]
    fn test_build_http_header_obsolete_line_folding() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "ows: https://tools.ietf.org/html/rfc7230",
                            "  #section-3.2.4",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

    #[test]
    fn test_build_http_header_obsolete_line_folding_tab() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "ows: https://tools.ietf.org/html/rfc7230",
                            "\t#section-3.2.4",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

}<commit_msg>Fix response body in case of no content<commit_after>//! The HTTP Client that perform query

use std::io::{self, Write};
use std::fmt::Arguments;

use super::request::Request;
use super::http;
use super::response::{Response, ResponseBuilder};
use super::results::CabotResult;

/// Perform the http query
pub struct Client {
    verbose: bool
}

impl Client {

    /// Construct a new `Client`
    pub fn new() -> Self {
        Client{verbose: false}
    }

    /// Execute the query [Request](../request/struct.Request.html) and
    /// return the associate [Response](../response/struct.Response.html).
    pub fn execute(&self, request: &Request) -> CabotResult<Response> {
        let mut out = CabotLibWrite::new();
        http::http_query(&request, &mut out, self.verbose)?;
        out.response()
    }
}

struct CabotLibWrite {
    response_builder: ResponseBuilder
}

impl CabotLibWrite {

    pub fn new() -> Self {
        CabotLibWrite{response_builder: ResponseBuilder::new()}
    }

    pub fn response(&self) -> CabotResult<Response> {
        self.response_builder.build()
    }
}

impl Write for CabotLibWrite {

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        info!("Parsing http response");
        let mut response = Vec::with_capacity(buf.len());
        response.extend_from_slice(buf);
        let response = unsafe { String::from_utf8_unchecked(response) };
        let response: Vec<&str> = response.splitn(2, "\r\n\r\n").collect();
        let header_len = response.get(0).unwrap().len();
        let mut headers: Vec<&str> = response.get(0).unwrap().split("\r\n").collect();
        let mut builder = ResponseBuilder::new();
        let status_line = headers.remove(0);
        debug!("Adding status line {}", status_line);
        builder = builder.set_status_line(status_line);
        let mut iter_header = headers.iter().peekable();
        let header = iter_header.next();
        if header.is_some() {
            let buf = header.unwrap();
            let mut header = String::with_capacity(buf.len() * 2);
            header.push_str(buf);
            loop {
                {
                    let buf = iter_header.peek();
                    if buf.is_none() {
                        builder = builder.add_header(header.as_str());
                        break;
                    }
                    let buf = buf.unwrap();
                    if buf.starts_with(" ") || buf.starts_with("\t") {
                        debug!("Obsolete line folded header reveived in {}",
                               header);
                        header.push_str(" ");
                        header.push_str(buf.trim_left());
                    }
                    else {
                        debug!("Adding header {}", header);
                        builder = builder.add_header(header.as_str());
                        header.clear();
                        header.push_str(buf);
                    }
                }
                let _  = iter_header.next();
            }
        }
        let body = if (header_len + 4) < buf.len() {
            &buf[(header_len + 4)..buf.len()]
        }
        else {
            &[]
        };
        //debug!("Adding body {:?}", body);
        builder = builder.set_body(body);
        self.response_builder = builder;
        // debug!("Response Builder - {:?}", self.response_builder);
        Ok(())
    }

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

    // Don't implemented unused method

    fn write(&mut self, _: &[u8]) -> io::Result<usize> {
        Err(io::Error::new(io::ErrorKind::Other, "Not Implemented"))
    }

    fn write_fmt(&mut self, _: Arguments) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "Not Implemented"))
    }

}

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

    #[test]
    fn test_build_http_response_from_string() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "Content-Type: text/plain",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &["Content-Type: text/plain", "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

    #[test]
    fn test_build_http_header_obsolete_line_folding() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "ows: https://tools.ietf.org/html/rfc7230",
                            "  #section-3.2.4",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

    #[test]
    fn test_build_http_header_obsolete_line_folding_tab() {
        let response = vec!["HTTP/1.1 200 Ok",
                            "ows: https://tools.ietf.org/html/rfc7230",
                            "\t#section-3.2.4",
                            "Content-Length: 12",
                            "",
                            "Hello World!"];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 200);
        assert_eq!(response.status_line(), "200 Ok");
        let headers: &[&str] = &[
            "ows: https://tools.ietf.org/html/rfc7230 #section-3.2.4",
            "Content-Length: 12"];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "Hello World!".to_owned());

    }

    #[test]
    fn test_build_http_no_response_body() {
        let response = vec!["HTTP/1.1 302 Moved",
                            "Location: https://tools.ietf.org/html/rfc7230#section-3.3",
                            ];
        let response = response.join("\r\n");

        let mut out = CabotLibWrite::new();
        out.write_all(response.as_bytes()).unwrap();
        let response = out.response().unwrap();
        assert_eq!(response.http_version(), "HTTP/1.1");
        assert_eq!(response.status_code(), 302);
        assert_eq!(response.status_line(), "302 Moved");
        let headers: &[&str] = &[
            "Location: https://tools.ietf.org/html/rfc7230#section-3.3",
            ];
        assert_eq!(response.headers(), headers);
        assert_eq!(response.body_as_string().unwrap(), "");
    }
}<|endoftext|>


<commit_before>//! Methods to connect to an WebSocket as a client.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::result::Result as StdResult;

use log::*;
use url::Url;

use crate::handshake::client::Response;
use crate::protocol::WebSocketConfig;

#[cfg(feature = "tls")]
mod encryption {
    pub use native_tls::TlsStream;
    use native_tls::{HandshakeError as TlsHandshakeError, TlsConnector};
    use std::net::TcpStream;

    pub use crate::stream::Stream as StreamSwitcher;
    /// TCP stream switcher (plain/TLS).
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    use crate::error::Result;
    use crate::stream::Mode;

    pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
            Mode::Tls => {
                let connector = TlsConnector::builder().build()?;
                connector
                    .connect(domain, stream)
                    .map_err(|e| match e {
                        TlsHandshakeError::Failure(f) => f.into(),
                        TlsHandshakeError::WouldBlock(_) => {
                            panic!("Bug: TLS handshake not blocked")
                        }
                    })
                    .map(StreamSwitcher::Tls)
            }
        }
    }
}

#[cfg(not(feature = "tls"))]
mod encryption {
    use std::net::TcpStream;

    use error::{Error, Result};
    use stream::Mode;

    /// TLS support is nod compiled in, this is just standard `TcpStream`.
    pub type AutoStream = TcpStream;

    pub fn wrap_stream(stream: TcpStream, _domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(stream),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::wrap_stream;
pub use self::encryption::AutoStream;

use crate::error::{Error, Result};
use crate::handshake::client::{ClientHandshake, Request};
use crate::handshake::HandshakeError;
use crate::protocol::WebSocket;
use crate::stream::{Mode, NoDelay};

/// Connect to the given WebSocket in blocking mode.
///
/// Uses a websocket configuration passed as an argument to the function. Calling it with `None` is
/// equal to calling `connect()` function.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect_with_config<'t, Req: Into<Request<'t>>>(
    request: Req,
    config: Option<WebSocketConfig>,
) -> Result<(WebSocket<AutoStream>, Response)> {
    let request: Request = request.into();
    let mode = url_mode(&request.url)?;
    let host = request
        .url
        .host()
        .ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    let port = request
        .url
        .port_or_known_default()
        .ok_or_else(|| Error::Url("No port number in the URL".into()))?;
    let addrs;
    let addr;
    let addrs = match host {
        url::Host::Domain(domain) => {
            addrs = (domain, port).to_socket_addrs()?;
            addrs.as_slice()
        }
        url::Host::Ipv4(ip) => {
            addr = (ip, port).into();
            std::slice::from_ref(&addr)
        }
        url::Host::Ipv6(ip) => {
            addr = (ip, port).into();
            std::slice::from_ref(&addr)
        }
    };
    let mut stream = connect_to_some(addrs, &request.url, mode)?;
    NoDelay::set_nodelay(&mut stream, true)?;
    client_with_config(request, stream, config).map_err(|e| match e {
        HandshakeError::Failure(f) => f,
        HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
    })
}

/// Connect to the given WebSocket in blocking mode.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect<'t, Req: Into<Request<'t>>>(
    request: Req,
) -> Result<(WebSocket<AutoStream>, Response)> {
    connect_with_config(request, None)
}

fn connect_to_some(addrs: &[SocketAddr], url: &Url, mode: Mode) -> Result<AutoStream> {
    let domain = url
        .host_str()
        .ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    for addr in addrs {
        debug!("Trying to contact {} at {}...", url, addr);
        if let Ok(raw_stream) = TcpStream::connect(addr) {
            if let Ok(stream) = wrap_stream(raw_stream, domain, mode) {
                return Ok(stream);
            }
        }
    }
    Err(Error::Url(format!("Unable to connect to {}", url).into()))
}

/// Get the mode of the given URL.
///
/// This function may be used to ease the creation of custom TLS streams
/// in non-blocking algorithmss or for use with TLS libraries other than `native_tls`.
pub fn url_mode(url: &Url) -> Result<Mode> {
    match url.scheme() {
        "ws" => Ok(Mode::Plain),
        "wss" => Ok(Mode::Tls),
        _ => Err(Error::Url("URL scheme not supported".into())),
    }
}

/// Do the client handshake over the given stream given a web socket configuration. Passing `None`
/// as configuration is equal to calling `client()` function.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client_with_config<'t, Stream, Req>(
    request: Req,
    stream: Stream,
    config: Option<WebSocketConfig>,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: Into<Request<'t>>,
{
    ClientHandshake::start(stream, request.into(), config).handshake()
}

/// Do the client handshake over the given stream.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client<'t, Stream, Req>(
    request: Req,
    stream: Stream,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: Into<Request<'t>>,
{
    client_with_config(request, stream, None)
}
<commit_msg>Fix compilation error without default features<commit_after>//! Methods to connect to an WebSocket as a client.

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::result::Result as StdResult;

use log::*;
use url::Url;

use crate::handshake::client::Response;
use crate::protocol::WebSocketConfig;

#[cfg(feature = "tls")]
mod encryption {
    pub use native_tls::TlsStream;
    use native_tls::{HandshakeError as TlsHandshakeError, TlsConnector};
    use std::net::TcpStream;

    pub use crate::stream::Stream as StreamSwitcher;
    /// TCP stream switcher (plain/TLS).
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    use crate::error::Result;
    use crate::stream::Mode;

    pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
            Mode::Tls => {
                let connector = TlsConnector::builder().build()?;
                connector
                    .connect(domain, stream)
                    .map_err(|e| match e {
                        TlsHandshakeError::Failure(f) => f.into(),
                        TlsHandshakeError::WouldBlock(_) => {
                            panic!("Bug: TLS handshake not blocked")
                        }
                    })
                    .map(StreamSwitcher::Tls)
            }
        }
    }
}

#[cfg(not(feature = "tls"))]
mod encryption {
    use std::net::TcpStream;

    use crate::error::{Error, Result};
    use crate::stream::Mode;

    /// TLS support is nod compiled in, this is just standard `TcpStream`.
    pub type AutoStream = TcpStream;

    pub fn wrap_stream(stream: TcpStream, _domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(stream),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::wrap_stream;
pub use self::encryption::AutoStream;

use crate::error::{Error, Result};
use crate::handshake::client::{ClientHandshake, Request};
use crate::handshake::HandshakeError;
use crate::protocol::WebSocket;
use crate::stream::{Mode, NoDelay};

/// Connect to the given WebSocket in blocking mode.
///
/// Uses a websocket configuration passed as an argument to the function. Calling it with `None` is
/// equal to calling `connect()` function.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect_with_config<'t, Req: Into<Request<'t>>>(
    request: Req,
    config: Option<WebSocketConfig>,
) -> Result<(WebSocket<AutoStream>, Response)> {
    let request: Request = request.into();
    let mode = url_mode(&request.url)?;
    let host = request
        .url
        .host()
        .ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    let port = request
        .url
        .port_or_known_default()
        .ok_or_else(|| Error::Url("No port number in the URL".into()))?;
    let addrs;
    let addr;
    let addrs = match host {
        url::Host::Domain(domain) => {
            addrs = (domain, port).to_socket_addrs()?;
            addrs.as_slice()
        }
        url::Host::Ipv4(ip) => {
            addr = (ip, port).into();
            std::slice::from_ref(&addr)
        }
        url::Host::Ipv6(ip) => {
            addr = (ip, port).into();
            std::slice::from_ref(&addr)
        }
    };
    let mut stream = connect_to_some(addrs, &request.url, mode)?;
    NoDelay::set_nodelay(&mut stream, true)?;
    client_with_config(request, stream, config).map_err(|e| match e {
        HandshakeError::Failure(f) => f,
        HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
    })
}

/// Connect to the given WebSocket in blocking mode.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect<'t, Req: Into<Request<'t>>>(
    request: Req,
) -> Result<(WebSocket<AutoStream>, Response)> {
    connect_with_config(request, None)
}

fn connect_to_some(addrs: &[SocketAddr], url: &Url, mode: Mode) -> Result<AutoStream> {
    let domain = url
        .host_str()
        .ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    for addr in addrs {
        debug!("Trying to contact {} at {}...", url, addr);
        if let Ok(raw_stream) = TcpStream::connect(addr) {
            if let Ok(stream) = wrap_stream(raw_stream, domain, mode) {
                return Ok(stream);
            }
        }
    }
    Err(Error::Url(format!("Unable to connect to {}", url).into()))
}

/// Get the mode of the given URL.
///
/// This function may be used to ease the creation of custom TLS streams
/// in non-blocking algorithmss or for use with TLS libraries other than `native_tls`.
pub fn url_mode(url: &Url) -> Result<Mode> {
    match url.scheme() {
        "ws" => Ok(Mode::Plain),
        "wss" => Ok(Mode::Tls),
        _ => Err(Error::Url("URL scheme not supported".into())),
    }
}

/// Do the client handshake over the given stream given a web socket configuration. Passing `None`
/// as configuration is equal to calling `client()` function.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client_with_config<'t, Stream, Req>(
    request: Req,
    stream: Stream,
    config: Option<WebSocketConfig>,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: Into<Request<'t>>,
{
    ClientHandshake::start(stream, request.into(), config).handshake()
}

/// Do the client handshake over the given stream.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client<'t, Stream, Req>(
    request: Req,
    stream: Stream,
) -> StdResult<(WebSocket<Stream>, Response), HandshakeError<ClientHandshake<Stream>>>
where
    Stream: Read + Write,
    Req: Into<Request<'t>>,
{
    client_with_config(request, stream, None)
}
<|endoftext|>