<commit_before>//! Primitives for styled text.
//!
//! A terminal UI is at its root a lot of strings. In order to make it accessible and stylish,
//! those strings may be associated to a set of styles. `tui` has three ways to represent them:
//! - A single line string where all graphemes have the same style is represented by a [`Span`].
//! - A single line string where each grapheme may have its own style is represented by [`Spans`].
//! - A multiple line string where each grapheme may have its own style is represented by a
//! [`Text`].
//!
//! These types form a hierarchy: [`Spans`] is a collection of [`Span`] and each line of [`Text`]
//! is a [`Spans`].
//!
//! Keep it mind that a lot of widgets will use those types to advertise what kind of string is
//! supported for their properties. Moreover, `tui` provides convenient `From` implementations so
//! that you can start by using simple `String` or `&str` and then promote them to the previous
//! primitives when you need additional styling capabilities.
//!
//! For example, for the [`crate::widgets::Block`] widget, all the following calls are valid to set
//! its `title` property (which is a [`Spans`] under the hood):
//!
//! ```rust
//! # use tui::widgets::Block;
//! # use tui::text::{Span, Spans};
//! # use tui::style::{Color, Style};
//! // A simple string with no styling.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My title"), style: Style { .. } }
//! // ])
//! let block = Block::default().title("My title");
//!
//! // A simple string with a unique style.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My title"), style: Style { fg: Some(Color::Yellow), .. }
//! // ])
//! let block = Block::default().title(
//!     Span::styled("My title", Style::default().fg(Color::Yellow))
//! );
//!
//! // A string with multiple styles.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My"), style: Style { fg: Some(Color::Yellow), .. } },
//! //   Span { content: Cow::Borrowed(" title"), .. }
//! // ])
//! let block = Block::default().title(vec![
//!     Span::styled("My", Style::default().fg(Color::Yellow)),
//!     Span::raw(" title"),
//! ]);
//! ```
use crate::style::Style;
use std::borrow::Cow;
use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

/// A grapheme associated to a style.
#[derive(Debug, Clone, PartialEq)]
pub struct StyledGrapheme<'a> {
    pub symbol: &'a str,
    pub style: Style,
}

/// A string where all graphemes have the same style.
#[derive(Debug, Clone, PartialEq)]
pub struct Span<'a> {
    pub content: Cow<'a, str>,
    pub style: Style,
}

impl<'a> Span<'a> {
    /// Create a span with no style.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::Span;
    /// Span::raw("My text");
    /// Span::raw(String::from("My text"));
    /// ```
    pub fn raw<T>(content: T) -> Span<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Span {
            content: content.into(),
            style: Style::default(),
        }
    }

    /// Create a span with a style.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Span;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// Span::styled("My text", style);
    /// Span::styled(String::from("My text"), style);
    /// ```
    pub fn styled<T>(content: T, style: Style) -> Span<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Span {
            content: content.into(),
            style,
        }
    }

    /// Returns the width of the content held by this span.
    pub fn width(&self) -> usize {
        self.content.width()
    }

    /// Returns an iterator over the graphemes held by this span.
    ///
    /// `base_style` is the [`Style`] that will be patched with each grapheme [`Style`] to get
    /// the resulting [`Style`].
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::{Span, StyledGrapheme};
    /// # use tui::style::{Color, Modifier, Style};
    /// # use std::iter::Iterator;
    /// let style = Style::default().fg(Color::Yellow);
    /// let span = Span::styled("Text", style);
    /// let style = Style::default().fg(Color::Green).bg(Color::Black);
    /// let styled_graphemes = span.styled_graphemes(style);
    /// assert_eq!(
    ///     vec![
    ///         StyledGrapheme {
    ///             symbol: "T",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "e",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "x",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "t",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///     ],
    ///     styled_graphemes.collect::<Vec<StyledGrapheme>>()
    /// );
    /// ```
    pub fn styled_graphemes(
        &'a self,
        base_style: Style,
    ) -> impl Iterator<Item = StyledGrapheme<'a>> {
        UnicodeSegmentation::graphemes(self.content.as_ref(), true)
            .map(move |g| StyledGrapheme {
                symbol: g,
                style: base_style.patch(self.style),
            })
            .filter(|s| s.symbol != "\n")
    }
}

impl<'a> From<String> for Span<'a> {
    fn from(s: String) -> Span<'a> {
        Span::raw(s)
    }
}

impl<'a> From<&'a str> for Span<'a> {
    fn from(s: &'a str) -> Span<'a> {
        Span::raw(s)
    }
}

/// A string composed of clusters of graphemes, each with their own style.
#[derive(Debug, Clone, PartialEq)]
pub struct Spans<'a>(pub Vec<Span<'a>>);

impl<'a> Default for Spans<'a> {
    fn default() -> Spans<'a> {
        Spans(Vec::new())
    }
}

impl<'a> Spans<'a> {
    /// Returns the width of the underlying string.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::{Span, Spans};
    /// # use tui::style::{Color, Style};
    /// let spans = Spans::from(vec![
    ///     Span::styled("My", Style::default().fg(Color::Yellow)),
    ///     Span::raw(" text"),
    /// ]);
    /// assert_eq!(7, spans.width());
    /// ```
    pub fn width(&self) -> usize {
        self.0.iter().map(Span::width).sum()
    }
}

impl<'a> From<String> for Spans<'a> {
    fn from(s: String) -> Spans<'a> {
        Spans(vec![Span::from(s)])
    }
}

impl<'a> From<&'a str> for Spans<'a> {
    fn from(s: &'a str) -> Spans<'a> {
        Spans(vec![Span::from(s)])
    }
}

impl<'a> From<Vec<Span<'a>>> for Spans<'a> {
    fn from(spans: Vec<Span<'a>>) -> Spans<'a> {
        Spans(spans)
    }
}

impl<'a> From<Span<'a>> for Spans<'a> {
    fn from(span: Span<'a>) -> Spans<'a> {
        Spans(vec![span])
    }
}

impl<'a> From<Spans<'a>> for String {
    fn from(line: Spans<'a>) -> String {
        line.0.iter().fold(String::new(), |mut acc, s| {
            acc.push_str(s.content.as_ref());
            acc
        })
    }
}

/// A string split over multiple lines where each line is composed of several clusters, each with
/// their own style.
///
/// A [`Text`], like a [`Span`], can be constructed using one of the many `From` implementations
/// or via the [`Text::raw`] and [`Text::styled`] methods. Helpfully, [`Text`] also implements
/// [`core::iter::Extend`] which enables the concatenation of several [`Text`] blocks.
///
/// ```rust
/// # use tui::text::Text;
/// # use tui::style::{Color, Modifier, Style};
/// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
///
/// // An initial two lines of `Text` built from a `&str`
/// let mut text = Text::from("The first line\nThe second line");
/// assert_eq!(2, text.height());
///
/// // Adding two more unstyled lines
/// text.extend(Text::raw("These are two\nmore lines!"));
/// assert_eq!(4, text.height());
///
/// // Adding a final two styled lines
/// text.extend(Text::styled("Some more lines\nnow with more style!", style));
/// assert_eq!(6, text.height());
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Text<'a> {
    pub lines: Vec<Spans<'a>>,
}

impl<'a> Default for Text<'a> {
    fn default() -> Text<'a> {
        Text { lines: Vec::new() }
    }
}

impl<'a> Text<'a> {
    /// Create some text (potentially multiple lines) with no style.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// Text::raw("The first line\nThe second line");
    /// Text::raw(String::from("The first line\nThe second line"));
    /// ```
    pub fn raw<T>(content: T) -> Text<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Text {
            lines: match content.into() {
                Cow::Borrowed(s) => s.lines().map(Spans::from).collect(),
                Cow::Owned(s) => s.lines().map(|l| Spans::from(l.to_owned())).collect(),
            },
        }
    }

    /// Create some text (potentially multiple lines) with a style.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// Text::styled("The first line\nThe second line", style);
    /// Text::styled(String::from("The first line\nThe second line"), style);
    /// ```
    pub fn styled<T>(content: T, style: Style) -> Text<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        let mut text = Text::raw(content);
        text.patch_style(style);
        text
    }

    /// Returns the max width of all the lines.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// use tui::text::Text;
    /// let text = Text::from("The first line\nThe second line");
    /// assert_eq!(15, text.width());
    /// ```
    pub fn width(&self) -> usize {
        self.lines
            .iter()
            .map(Spans::width)
            .max()
            .unwrap_or_default()
    }

    /// Returns the height.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// use tui::text::Text;
    /// let text = Text::from("The first line\nThe second line");
    /// assert_eq!(2, text.height());
    /// ```
    pub fn height(&self) -> usize {
        self.lines.len()
    }

    /// Apply a new style to existing text.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// let mut raw_text = Text::raw("The first line\nThe second line");
    /// let styled_text = Text::styled(String::from("The first line\nThe second line"), style);
    /// assert_ne!(raw_text, styled_text);
    ///
    /// raw_text.patch_style(style);
    /// assert_eq!(raw_text, styled_text);
    /// ```
    pub fn patch_style(&mut self, style: Style) {
        for line in &mut self.lines {
            for span in &mut line.0 {
                span.style = span.style.patch(style);
            }
        }
    }
}

impl<'a> From<String> for Text<'a> {
    fn from(s: String) -> Text<'a> {
        Text::raw(s)
    }
}

impl<'a> From<&'a str> for Text<'a> {
    fn from(s: &'a str) -> Text<'a> {
        Text::raw(s)
    }
}

impl<'a> From<Span<'a>> for Text<'a> {
    fn from(span: Span<'a>) -> Text<'a> {
        Text {
            lines: vec![Spans::from(span)],
        }
    }
}

impl<'a> From<Spans<'a>> for Text<'a> {
    fn from(spans: Spans<'a>) -> Text<'a> {
        Text { lines: vec![spans] }
    }
}

impl<'a> From<Vec<Spans<'a>>> for Text<'a> {
    fn from(lines: Vec<Spans<'a>>) -> Text<'a> {
        Text { lines }
    }
}

impl<'a> IntoIterator for Text<'a> {
    type Item = Spans<'a>;
    type IntoIter = std::vec::IntoIter<Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.lines.into_iter()
    }
}

impl<'a> Extend<Spans<'a>> for Text<'a> {
    fn extend<T: IntoIterator<Item = Spans<'a>>>(&mut self, iter: T) {
        self.lines.extend(iter);
    }
}
<commit_msg>feat(Text): Add a From<Cow<str>> impl for Text<commit_after>//! Primitives for styled text.
//!
//! A terminal UI is at its root a lot of strings. In order to make it accessible and stylish,
//! those strings may be associated to a set of styles. `tui` has three ways to represent them:
//! - A single line string where all graphemes have the same style is represented by a [`Span`].
//! - A single line string where each grapheme may have its own style is represented by [`Spans`].
//! - A multiple line string where each grapheme may have its own style is represented by a
//! [`Text`].
//!
//! These types form a hierarchy: [`Spans`] is a collection of [`Span`] and each line of [`Text`]
//! is a [`Spans`].
//!
//! Keep it mind that a lot of widgets will use those types to advertise what kind of string is
//! supported for their properties. Moreover, `tui` provides convenient `From` implementations so
//! that you can start by using simple `String` or `&str` and then promote them to the previous
//! primitives when you need additional styling capabilities.
//!
//! For example, for the [`crate::widgets::Block`] widget, all the following calls are valid to set
//! its `title` property (which is a [`Spans`] under the hood):
//!
//! ```rust
//! # use tui::widgets::Block;
//! # use tui::text::{Span, Spans};
//! # use tui::style::{Color, Style};
//! // A simple string with no styling.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My title"), style: Style { .. } }
//! // ])
//! let block = Block::default().title("My title");
//!
//! // A simple string with a unique style.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My title"), style: Style { fg: Some(Color::Yellow), .. }
//! // ])
//! let block = Block::default().title(
//!     Span::styled("My title", Style::default().fg(Color::Yellow))
//! );
//!
//! // A string with multiple styles.
//! // Converted to Spans(vec![
//! //   Span { content: Cow::Borrowed("My"), style: Style { fg: Some(Color::Yellow), .. } },
//! //   Span { content: Cow::Borrowed(" title"), .. }
//! // ])
//! let block = Block::default().title(vec![
//!     Span::styled("My", Style::default().fg(Color::Yellow)),
//!     Span::raw(" title"),
//! ]);
//! ```
use crate::style::Style;
use std::borrow::Cow;
use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

/// A grapheme associated to a style.
#[derive(Debug, Clone, PartialEq)]
pub struct StyledGrapheme<'a> {
    pub symbol: &'a str,
    pub style: Style,
}

/// A string where all graphemes have the same style.
#[derive(Debug, Clone, PartialEq)]
pub struct Span<'a> {
    pub content: Cow<'a, str>,
    pub style: Style,
}

impl<'a> Span<'a> {
    /// Create a span with no style.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::Span;
    /// Span::raw("My text");
    /// Span::raw(String::from("My text"));
    /// ```
    pub fn raw<T>(content: T) -> Span<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Span {
            content: content.into(),
            style: Style::default(),
        }
    }

    /// Create a span with a style.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Span;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// Span::styled("My text", style);
    /// Span::styled(String::from("My text"), style);
    /// ```
    pub fn styled<T>(content: T, style: Style) -> Span<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Span {
            content: content.into(),
            style,
        }
    }

    /// Returns the width of the content held by this span.
    pub fn width(&self) -> usize {
        self.content.width()
    }

    /// Returns an iterator over the graphemes held by this span.
    ///
    /// `base_style` is the [`Style`] that will be patched with each grapheme [`Style`] to get
    /// the resulting [`Style`].
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::{Span, StyledGrapheme};
    /// # use tui::style::{Color, Modifier, Style};
    /// # use std::iter::Iterator;
    /// let style = Style::default().fg(Color::Yellow);
    /// let span = Span::styled("Text", style);
    /// let style = Style::default().fg(Color::Green).bg(Color::Black);
    /// let styled_graphemes = span.styled_graphemes(style);
    /// assert_eq!(
    ///     vec![
    ///         StyledGrapheme {
    ///             symbol: "T",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "e",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "x",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///         StyledGrapheme {
    ///             symbol: "t",
    ///             style: Style {
    ///                 fg: Some(Color::Yellow),
    ///                 bg: Some(Color::Black),
    ///                 add_modifier: Modifier::empty(),
    ///                 sub_modifier: Modifier::empty(),
    ///             },
    ///         },
    ///     ],
    ///     styled_graphemes.collect::<Vec<StyledGrapheme>>()
    /// );
    /// ```
    pub fn styled_graphemes(
        &'a self,
        base_style: Style,
    ) -> impl Iterator<Item = StyledGrapheme<'a>> {
        UnicodeSegmentation::graphemes(self.content.as_ref(), true)
            .map(move |g| StyledGrapheme {
                symbol: g,
                style: base_style.patch(self.style),
            })
            .filter(|s| s.symbol != "\n")
    }
}

impl<'a> From<String> for Span<'a> {
    fn from(s: String) -> Span<'a> {
        Span::raw(s)
    }
}

impl<'a> From<&'a str> for Span<'a> {
    fn from(s: &'a str) -> Span<'a> {
        Span::raw(s)
    }
}

/// A string composed of clusters of graphemes, each with their own style.
#[derive(Debug, Clone, PartialEq)]
pub struct Spans<'a>(pub Vec<Span<'a>>);

impl<'a> Default for Spans<'a> {
    fn default() -> Spans<'a> {
        Spans(Vec::new())
    }
}

impl<'a> Spans<'a> {
    /// Returns the width of the underlying string.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::{Span, Spans};
    /// # use tui::style::{Color, Style};
    /// let spans = Spans::from(vec![
    ///     Span::styled("My", Style::default().fg(Color::Yellow)),
    ///     Span::raw(" text"),
    /// ]);
    /// assert_eq!(7, spans.width());
    /// ```
    pub fn width(&self) -> usize {
        self.0.iter().map(Span::width).sum()
    }
}

impl<'a> From<String> for Spans<'a> {
    fn from(s: String) -> Spans<'a> {
        Spans(vec![Span::from(s)])
    }
}

impl<'a> From<&'a str> for Spans<'a> {
    fn from(s: &'a str) -> Spans<'a> {
        Spans(vec![Span::from(s)])
    }
}

impl<'a> From<Vec<Span<'a>>> for Spans<'a> {
    fn from(spans: Vec<Span<'a>>) -> Spans<'a> {
        Spans(spans)
    }
}

impl<'a> From<Span<'a>> for Spans<'a> {
    fn from(span: Span<'a>) -> Spans<'a> {
        Spans(vec![span])
    }
}

impl<'a> From<Spans<'a>> for String {
    fn from(line: Spans<'a>) -> String {
        line.0.iter().fold(String::new(), |mut acc, s| {
            acc.push_str(s.content.as_ref());
            acc
        })
    }
}

/// A string split over multiple lines where each line is composed of several clusters, each with
/// their own style.
///
/// A [`Text`], like a [`Span`], can be constructed using one of the many `From` implementations
/// or via the [`Text::raw`] and [`Text::styled`] methods. Helpfully, [`Text`] also implements
/// [`core::iter::Extend`] which enables the concatenation of several [`Text`] blocks.
///
/// ```rust
/// # use tui::text::Text;
/// # use tui::style::{Color, Modifier, Style};
/// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
///
/// // An initial two lines of `Text` built from a `&str`
/// let mut text = Text::from("The first line\nThe second line");
/// assert_eq!(2, text.height());
///
/// // Adding two more unstyled lines
/// text.extend(Text::raw("These are two\nmore lines!"));
/// assert_eq!(4, text.height());
///
/// // Adding a final two styled lines
/// text.extend(Text::styled("Some more lines\nnow with more style!", style));
/// assert_eq!(6, text.height());
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Text<'a> {
    pub lines: Vec<Spans<'a>>,
}

impl<'a> Default for Text<'a> {
    fn default() -> Text<'a> {
        Text { lines: Vec::new() }
    }
}

impl<'a> Text<'a> {
    /// Create some text (potentially multiple lines) with no style.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// Text::raw("The first line\nThe second line");
    /// Text::raw(String::from("The first line\nThe second line"));
    /// ```
    pub fn raw<T>(content: T) -> Text<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        Text {
            lines: match content.into() {
                Cow::Borrowed(s) => s.lines().map(Spans::from).collect(),
                Cow::Owned(s) => s.lines().map(|l| Spans::from(l.to_owned())).collect(),
            },
        }
    }

    /// Create some text (potentially multiple lines) with a style.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// Text::styled("The first line\nThe second line", style);
    /// Text::styled(String::from("The first line\nThe second line"), style);
    /// ```
    pub fn styled<T>(content: T, style: Style) -> Text<'a>
    where
        T: Into<Cow<'a, str>>,
    {
        let mut text = Text::raw(content);
        text.patch_style(style);
        text
    }

    /// Returns the max width of all the lines.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// use tui::text::Text;
    /// let text = Text::from("The first line\nThe second line");
    /// assert_eq!(15, text.width());
    /// ```
    pub fn width(&self) -> usize {
        self.lines
            .iter()
            .map(Spans::width)
            .max()
            .unwrap_or_default()
    }

    /// Returns the height.
    ///
    /// ## Examples
    ///
    /// ```rust
    /// use tui::text::Text;
    /// let text = Text::from("The first line\nThe second line");
    /// assert_eq!(2, text.height());
    /// ```
    pub fn height(&self) -> usize {
        self.lines.len()
    }

    /// Apply a new style to existing text.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tui::text::Text;
    /// # use tui::style::{Color, Modifier, Style};
    /// let style = Style::default().fg(Color::Yellow).add_modifier(Modifier::ITALIC);
    /// let mut raw_text = Text::raw("The first line\nThe second line");
    /// let styled_text = Text::styled(String::from("The first line\nThe second line"), style);
    /// assert_ne!(raw_text, styled_text);
    ///
    /// raw_text.patch_style(style);
    /// assert_eq!(raw_text, styled_text);
    /// ```
    pub fn patch_style(&mut self, style: Style) {
        for line in &mut self.lines {
            for span in &mut line.0 {
                span.style = span.style.patch(style);
            }
        }
    }
}

impl<'a> From<String> for Text<'a> {
    fn from(s: String) -> Text<'a> {
        Text::raw(s)
    }
}

impl<'a> From<&'a str> for Text<'a> {
    fn from(s: &'a str) -> Text<'a> {
        Text::raw(s)
    }
}

impl<'a> From<Cow<'a, str>> for Text<'a> {
    fn from(s: Cow<'a, str>) -> Text<'a> {
        Text::raw(s)
    }
}

impl<'a> From<Span<'a>> for Text<'a> {
    fn from(span: Span<'a>) -> Text<'a> {
        Text {
            lines: vec![Spans::from(span)],
        }
    }
}

impl<'a> From<Spans<'a>> for Text<'a> {
    fn from(spans: Spans<'a>) -> Text<'a> {
        Text { lines: vec![spans] }
    }
}

impl<'a> From<Vec<Spans<'a>>> for Text<'a> {
    fn from(lines: Vec<Spans<'a>>) -> Text<'a> {
        Text { lines }
    }
}

impl<'a> IntoIterator for Text<'a> {
    type Item = Spans<'a>;
    type IntoIter = std::vec::IntoIter<Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        self.lines.into_iter()
    }
}

impl<'a> Extend<Spans<'a>> for Text<'a> {
    fn extend<T: IntoIterator<Item = Spans<'a>>>(&mut self, iter: T) {
        self.lines.extend(iter);
    }
}
<|endoftext|>


<commit_before>use std::fmt::Debug;
use std::hash::{Hash, Hasher, SipHasher};

use adapton::bitstring::*;
use adapton::engine::{Art, Name, force};

/// Probablistically Balanced Trie
/// Rough implementation of probabilistic tries from OOPSLA 2015 paper.
///
/// See also: [Tries in OCaml](http://github.com/plum-umd/adapton.ocaml)
#[derive(Debug,PartialEq,Eq,Clone)]
pub enum Trie<X> {
    Nil(BS),
    Leaf(BS, X),
    Bin(BS, Box<Trie<X>>, Box<Trie<X>>),
    Root(Meta, Box<Trie<X>>),
    Name(Name, Box<Trie<X>>),
    Art(Art<Trie<X>>),
}

#[derive(Debug,PartialEq,Eq,Clone)]
pub enum IFreq {
    Never,
    Depth(u64, u64),
    First(u64),
    Const(u64),
}

impl Hash for IFreq {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            IFreq::Never => {
                "Adapton.Trie.Meta.Freq#Never".hash(state);
            }
            IFreq::Depth(i, j) => {
                "Adapton.Trie.Meta.Freq#Depth".hash(state);
                i.hash(state);
                j.hash(state);
            }
            IFreq::First(n) => {
                "Adapton.Trie.Meta.Freq#First".hash(state);
                n.hash(state);
            }
            IFreq::Const(i) => {
                "Adapton.Trie.Meta.Freq#Const".hash(state);
                i.hash(state);
            }
        }
    }
}

/// Metadata held by the root node.
#[derive(Debug,PartialEq,Eq,Hash,Clone)]
pub struct Meta {
    min_depth: u64,
    ifreq: IFreq,
}

pub trait MetaT {
    fn hash_seeded(&self, u64);
}

impl MetaT for Meta {
    fn hash_seeded(&self, seed: u64) {
        let mut hasher = SipHasher::new_with_keys(0, seed);
        "Adapton.Trie.Meta".hash(&mut hasher);
        self.min_depth.hash(&mut hasher);
        self.ifreq.hash(&mut hasher);
    }
}

pub trait TrieIntro<X>: Debug + Hash + PartialEq + Eq + Clone + 'static {
    fn nil() -> Self;
    fn leaf(BS, X) -> Self;
    fn bin(BS, Self, Self) -> Self;
    fn root(Meta, Self) -> Self;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name(Name, Self) -> Self;
    fn art(Art<Self>) -> Self;
}

pub trait TrieElim<X>: Debug + Hash + PartialEq + Eq + Clone + 'static {
    fn find(Self, i64) -> Option<X>;

    fn elim<Res, NilC, LeafC, BinC, RootC, NameC>(Self, NilC, LeafC, BinC, RootC, NameC) -> Res
        where NilC: FnOnce(BS) -> Res,
              LeafC: FnOnce(BS, X) -> Res,
              BinC: FnOnce(BS, Self, Self) -> Res,
              RootC: FnOnce(Meta, Self) -> Res,
              NameC: FnOnce(Name, Self) -> Res;

    fn elim_ref<Res, NilC, LeafC, BinC, RootC, NameC>(&Self,
                                                      NilC,
                                                      LeafC,
                                                      BinC,
                                                      RootC,
                                                      NameC)
                                                      -> Res
        where NilC: FnOnce(&BS) -> Res,
              LeafC: FnOnce(&BS, &X) -> Res,
              BinC: FnOnce(&BS, &Self, &Self) -> Res,
              RootC: FnOnce(&Meta, &Self) -> Res,
              NameC: FnOnce(&Name, &Self) -> Res;
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> TrieIntro<X> for Trie<X> {
    fn nil() -> Self {
        Trie::Nil(BS {
            length: 0,
            value: 0,
        })
    }
    fn leaf(bs: BS, x: X) -> Self {
        Trie::Leaf(bs, x)
    }
    fn bin(bs: BS, l: Self, r: Self) -> Self {
        Trie::Bin(bs, Box::new(l), Box::new(r))
    }
    fn root(meta: Meta, trie: Self) -> Self {
        Trie::Root(meta, Box::new(trie))
    }
    fn name(nm: Name, trie: Self) -> Self {
        Trie::Name(nm, Box::new(trie))
    }
    fn art(art: Art<Self>) -> Self {
        Trie::Art(art)
    }
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> Hash for Trie<X> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            Trie::Nil(bs) => bs.hash(state),
            Trie::Leaf(bs, ref x) => {
                x.hash(state);
                bs.hash(state)
            }
            Trie::Bin(bs, ref left, ref right) => {
                right.hash(state);
                left.hash(state);
                bs.hash(state)
            }
            Trie::Root(ref md, ref t) => {
                t.hash(state);
                md.hash_seeded(state.finish());
            }
            Trie::Name(ref nm, ref t) => {
                t.hash(state);
                nm.hash(state)
            }
            Trie::Art(ref art_t) => art_t.hash(state),
        }
    }
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> TrieElim<X> for Trie<X> {
    fn find(trie: Self, i: i64) -> Option<X> {
        fn _loop<X: Debug + Hash + PartialEq + Eq + Clone + 'static>(h: i64,
                                                                     t: Trie<X>)
                                                                     -> Option<X> {
            match t {
                Trie::Nil(_) => None,
                Trie::Leaf(_, x) => Some(x),
                Trie::Bin(_, left, right) => _loop(h >> 1, if h % 2 == 0 { *left } else { *right }),
                Trie::Root(_, t) => _loop(h, *t),
                Trie::Name(_, t) => _loop(h, *t),
                Trie::Art(art_t) => _loop(h, force(&art_t)),
            }
        };
        _loop(i, trie)
    }

    fn elim<Res, NilC, LeafC, BinC, RootC, NameC>(trie: Self,
                                                  nil: NilC,
                                                  leaf: LeafC,
                                                  bin: BinC,
                                                  root: RootC,
                                                  name: NameC)
                                                  -> Res
        where NilC: FnOnce(BS) -> Res,
              LeafC: FnOnce(BS, X) -> Res,
              BinC: FnOnce(BS, Self, Self) -> Res,
              RootC: FnOnce(Meta, Self) -> Res,
              NameC: FnOnce(Name, Self) -> Res
    {
        match trie {
            Trie::Nil(bs) => nil(bs),
            Trie::Leaf(bs, x) => leaf(bs, x),
            Trie::Bin(bs, l, r) => bin(bs, *l, *r),
            Trie::Name(nm, t) => name(nm, *t),
            Trie::Root(meta, t) => root(meta, *t),
            Trie::Art(art) => {
                let trie = force(&art);
                Self::elim(trie, nil, leaf, bin, root, name)
            }
        }
    }

    fn elim_ref<Res, NilC, LeafC, BinC, RootC, NameC>(trie: &Self,
                                                      nil: NilC,
                                                      leaf: LeafC,
                                                      bin: BinC,
                                                      root: RootC,
                                                      name: NameC)
                                                      -> Res
        where NilC: FnOnce(&BS) -> Res,
              LeafC: FnOnce(&BS, &X) -> Res,
              BinC: FnOnce(&BS, &Self, &Self) -> Res,
              RootC: FnOnce(&Meta, &Self) -> Res,
              NameC: FnOnce(&Name, &Self) -> Res
    {
        match *trie {
            Trie::Nil(ref bs) => nil(bs),
            Trie::Leaf(ref bs, ref x) => leaf(bs, x),
            Trie::Bin(ref bs, ref l, ref r) => bin(bs, &*l, &*r),
            Trie::Name(ref nm, ref t) => name(nm, &*t),
            Trie::Root(ref meta, ref t) => root(meta, &*t),
            Trie::Art(ref art) => {
                let trie = force(&art);
                Self::elim_ref(&trie, nil, leaf, bin, root, name)
            }
        }
    }
}
<commit_msg>Implement specialized eq fn for tries<commit_after>use std::fmt::Debug;
use std::hash::{Hash, Hasher, SipHasher};

use adapton::bitstring::*;
use adapton::engine::{Art, Name, force};

/// Probablistically Balanced Trie
/// Rough implementation of probabilistic tries from OOPSLA 2015 paper.
///
/// See also: [Tries in OCaml](http://github.com/plum-umd/adapton.ocaml)
#[derive(Debug,Clone)]
pub enum Trie<X> {
    Nil(BS),
    Leaf(BS, X),
    Bin(BS, Box<Trie<X>>, Box<Trie<X>>),
    Root(Meta, Box<Trie<X>>),
    Name(Name, Box<Trie<X>>),
    Art(Art<Trie<X>>),
}

#[derive(Debug,PartialEq,Eq,Clone)]
pub enum IFreq {
    Never,
    Depth(u64, u64),
    First(u64),
    Const(u64),
}

impl Hash for IFreq {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            IFreq::Never => {
                "Adapton.Trie.Meta.Freq#Never".hash(state);
            }
            IFreq::Depth(i, j) => {
                "Adapton.Trie.Meta.Freq#Depth".hash(state);
                i.hash(state);
                j.hash(state);
            }
            IFreq::First(n) => {
                "Adapton.Trie.Meta.Freq#First".hash(state);
                n.hash(state);
            }
            IFreq::Const(i) => {
                "Adapton.Trie.Meta.Freq#Const".hash(state);
                i.hash(state);
            }
        }
    }
}

/// Metadata held by the root node.
#[derive(Debug,PartialEq,Eq,Hash,Clone)]
pub struct Meta {
    min_depth: u64,
    ifreq: IFreq,
}

pub trait MetaT {
    fn hash_seeded(&self, u64);
}

impl MetaT for Meta {
    fn hash_seeded(&self, seed: u64) {
        let mut hasher = SipHasher::new_with_keys(0, seed);
        "Adapton.Trie.Meta".hash(&mut hasher);
        self.min_depth.hash(&mut hasher);
        self.ifreq.hash(&mut hasher);
    }
}

pub trait TrieIntro<X>: Debug + Hash + PartialEq + Eq + Clone + 'static {
    fn nil() -> Self;
    fn leaf(BS, X) -> Self;
    fn bin(BS, Self, Self) -> Self;
    fn root(Meta, Self) -> Self;

    // requisite "adaptonic" constructors: `name` and `art`:
    fn name(Name, Self) -> Self;
    fn art(Art<Self>) -> Self;
}

pub trait TrieElim<X>: Debug + Hash + PartialEq + Eq + Clone + 'static {
    fn find(Self, i64) -> Option<X>;

    fn elim<Res, NilC, LeafC, BinC, RootC, NameC>(Self, NilC, LeafC, BinC, RootC, NameC) -> Res
        where NilC: FnOnce(BS) -> Res,
              LeafC: FnOnce(BS, X) -> Res,
              BinC: FnOnce(BS, Self, Self) -> Res,
              RootC: FnOnce(Meta, Self) -> Res,
              NameC: FnOnce(Name, Self) -> Res;

    fn elim_ref<Res, NilC, LeafC, BinC, RootC, NameC>(&Self,
                                                      NilC,
                                                      LeafC,
                                                      BinC,
                                                      RootC,
                                                      NameC)
                                                      -> Res
        where NilC: FnOnce(&BS) -> Res,
              LeafC: FnOnce(&BS, &X) -> Res,
              BinC: FnOnce(&BS, &Self, &Self) -> Res,
              RootC: FnOnce(&Meta, &Self) -> Res,
              NameC: FnOnce(&Name, &Self) -> Res;
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> TrieIntro<X> for Trie<X> {
    fn nil() -> Self {
        Trie::Nil(BS {
            length: 0,
            value: 0,
        })
    }
    fn leaf(bs: BS, x: X) -> Self {
        Trie::Leaf(bs, x)
    }
    fn bin(bs: BS, l: Self, r: Self) -> Self {
        Trie::Bin(bs, Box::new(l), Box::new(r))
    }
    fn root(meta: Meta, trie: Self) -> Self {
        Trie::Root(meta, Box::new(trie))
    }
    fn name(nm: Name, trie: Self) -> Self {
        Trie::Name(nm, Box::new(trie))
    }
    fn art(art: Art<Self>) -> Self {
        Trie::Art(art)
    }
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> PartialEq for Trie<X> {
    fn eq(&self, other: &Trie<X>) -> bool {
        match (self, other) {
            (&Trie::Nil(bs_self), &Trie::Nil(bs_other)) => bs_self == bs_other,
            (&Trie::Leaf(bs_self, ref e_self), &Trie::Leaf(bs_other, ref e_other)) => {
                let bs_equal = bs_self == bs_other;
                bs_equal && e_self == e_other
            }
            (&Trie::Bin(bs, ref left, ref right),
             &Trie::Bin(bs_other, ref left_other, ref right_other)) => {
                bs == bs_other && left == left_other && right == right_other
            }
            (&Trie::Root(ref md, ref t), &Trie::Root(ref md_other, ref t_other)) => {
                md == md_other && t == t_other
            }
            (&Trie::Name(ref nm, ref t), &Trie::Name(ref nm_other, ref t_other)) => {
                nm == nm_other && t == t_other
            }
            _ => false,
        }
    }
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> Eq for Trie<X> {}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> Hash for Trie<X> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match *self {
            Trie::Nil(bs) => bs.hash(state),
            Trie::Leaf(bs, ref x) => {
                x.hash(state);
                bs.hash(state)
            }
            Trie::Bin(bs, ref left, ref right) => {
                right.hash(state);
                left.hash(state);
                bs.hash(state)
            }
            Trie::Root(ref md, ref t) => {
                t.hash(state);
                md.hash_seeded(state.finish());
            }
            Trie::Name(ref nm, ref t) => {
                t.hash(state);
                nm.hash(state)
            }
            Trie::Art(ref art_t) => art_t.hash(state),
        }
    }
}

impl<X: Debug + Hash + PartialEq + Eq + Clone + 'static> TrieElim<X> for Trie<X> {
    fn find(trie: Self, i: i64) -> Option<X> {
        fn _loop<X: Debug + Hash + PartialEq + Eq + Clone + 'static>(h: i64,
                                                                     t: Trie<X>)
                                                                     -> Option<X> {
            match t {
                Trie::Nil(_) => None,
                Trie::Leaf(_, x) => Some(x),
                Trie::Bin(_, left, right) => _loop(h >> 1, if h % 2 == 0 { *left } else { *right }),
                Trie::Root(_, t) => _loop(h, *t),
                Trie::Name(_, t) => _loop(h, *t),
                Trie::Art(art_t) => _loop(h, force(&art_t)),
            }
        };
        _loop(i, trie)
    }

    fn elim<Res, NilC, LeafC, BinC, RootC, NameC>(trie: Self,
                                                  nil: NilC,
                                                  leaf: LeafC,
                                                  bin: BinC,
                                                  root: RootC,
                                                  name: NameC)
                                                  -> Res
        where NilC: FnOnce(BS) -> Res,
              LeafC: FnOnce(BS, X) -> Res,
              BinC: FnOnce(BS, Self, Self) -> Res,
              RootC: FnOnce(Meta, Self) -> Res,
              NameC: FnOnce(Name, Self) -> Res
    {
        match trie {
            Trie::Nil(bs) => nil(bs),
            Trie::Leaf(bs, x) => leaf(bs, x),
            Trie::Bin(bs, l, r) => bin(bs, *l, *r),
            Trie::Name(nm, t) => name(nm, *t),
            Trie::Root(meta, t) => root(meta, *t),
            Trie::Art(art) => {
                let trie = force(&art);
                Self::elim(trie, nil, leaf, bin, root, name)
            }
        }
    }

    fn elim_ref<Res, NilC, LeafC, BinC, RootC, NameC>(trie: &Self,
                                                      nil: NilC,
                                                      leaf: LeafC,
                                                      bin: BinC,
                                                      root: RootC,
                                                      name: NameC)
                                                      -> Res
        where NilC: FnOnce(&BS) -> Res,
              LeafC: FnOnce(&BS, &X) -> Res,
              BinC: FnOnce(&BS, &Self, &Self) -> Res,
              RootC: FnOnce(&Meta, &Self) -> Res,
              NameC: FnOnce(&Name, &Self) -> Res
    {
        match *trie {
            Trie::Nil(ref bs) => nil(bs),
            Trie::Leaf(ref bs, ref x) => leaf(bs, x),
            Trie::Bin(ref bs, ref l, ref r) => bin(bs, &*l, &*r),
            Trie::Name(ref nm, ref t) => name(nm, &*t),
            Trie::Root(ref meta, ref t) => root(meta, &*t),
            Trie::Art(ref art) => {
                let trie = force(&art);
                Self::elim_ref(&trie, nil, leaf, bin, root, name)
            }
        }
    }
}
<|endoftext|>


<commit_before>// rust-xmpp
// Copyright (c) 2014 Florian Zeitz
//
// This project is MIT licensed.
// Please see the COPYING file for more information.

use std::io;
use std::io::BufRead;
use std::str;

// https://tools.ietf.org/html/rfc3629
static UTF8_CHAR_WIDTH: [u8; 256] = [
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x1F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x3F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x5F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x7F
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, // 0x9F
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, // 0xBF
0,0,2,2,2,2,2,2,2,2,2,2,2,2,2,2,
2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2, // 0xDF
3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3, // 0xEF
4,4,4,4,4,0,0,0,0,0,0,0,0,0,0,0, // 0xFF
];

/// Given a first byte, determine how many bytes are in this UTF-8 character
#[inline]
fn utf8_char_width(b: u8) -> usize {
    return UTF8_CHAR_WIDTH[b as usize] as usize;
}

pub trait ReadString {
    fn read_str(&mut self) -> io::Result<String>;
}

impl<T: BufRead> ReadString for T {
    fn read_str(&mut self) -> io::Result<String> {
        let (result, last) = {
            let available = try!(self.fill_buf());
            let len = available.len();
            let mut last = if len < 3 { 0 } else { len - 3 };
            while last < len {
                let width = utf8_char_width(available[last]);
                if width == 0 {
                    last += 1;
                    continue;
                }
                if last+width <= len {
                    last += width;
                } else {
                    break;
                }
            }
            (str::from_utf8(&available[..last]).unwrap().to_string(), last)
        };
        self.consume(last);

        Ok(result)
    }
}
<commit_msg>ReadString: Add proper error handling<commit_after>// rust-xmpp
// Copyright (c) 2014 Florian Zeitz
//
// This project is MIT licensed.
// Please see the COPYING file for more information.

use std::io;
use std::io::BufRead;
use std::str;

// https://tools.ietf.org/html/rfc3629
static UTF8_CHAR_WIDTH: [u8; 256] = [
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x1F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x3F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x5F
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, // 0x7F
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, // 0x9F
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, // 0xBF
0,0,2,2,2,2,2,2,2,2,2,2,2,2,2,2,
2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2, // 0xDF
3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3, // 0xEF
4,4,4,4,4,0,0,0,0,0,0,0,0,0,0,0, // 0xFF
];

/// Given a first byte, determine how many bytes are in this UTF-8 character
#[inline]
fn utf8_char_width(b: u8) -> usize {
    return UTF8_CHAR_WIDTH[b as usize] as usize;
}

pub trait ReadString {
    fn read_str(&mut self) -> io::Result<String>;
}

impl<T: BufRead> ReadString for T {
    fn read_str(&mut self) -> io::Result<String> {
        let (result, last) = {
            let available = try!(self.fill_buf());
            let len = available.len();
            let mut last = if len < 3 { 0 } else { len - 3 };
            while last < len {
                let width = utf8_char_width(available[last]);
                if width == 0 {
                    last += 1;
                    continue;
                }
                if last+width <= len {
                    last += width;
                } else {
                    break;
                }
            }
            let res = str::from_utf8(&available[..last]);
            (res.map(|x| x.to_string())
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData,
                                            "stream did not contain valid UTF-8")), last)
        };
        self.consume(last);

        result
    }
}
<|endoftext|>


<commit_before>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/
#![allow(non_snake_case)]
//! Testing physical properties of a Lennard-Jones gaz of Helium.
extern crate cymbalum;
use self::cymbalum::*;

use std::sync::{Once, ONCE_INIT};
static START: Once = ONCE_INIT;

use std::path::Path;

fn get_universe() -> Universe {
    let data_dir = Path::new(file!()).parent().unwrap();
    let configuration = data_dir.join("data").join("helium.xyz");
    let mut universe = Universe::from_file(configuration.to_str().unwrap()).unwrap();
    universe.set_cell(UnitCell::cubic(10.0));
    universe.add_pair_interaction("He", "He",
        Box::new(LennardJones{
            sigma: units::from(2.0, "A").unwrap(),
            epsilon: units::from(0.2, "kJ/mol").unwrap()
        })
    );

    let mut velocities = BoltzmanVelocities::new(units::from(300.0, "K").unwrap());
    velocities.init(&mut universe);
    return universe;
}

#[test]
fn constant_energy_velocity_verlet() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );

    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();

    assert!(f64::abs(E_initial - E_final) < 1e-5);
}

#[test]
fn constant_energy_verlet() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            Verlet::new(units::from(1.0, "fs").unwrap())
        )
    );
    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    assert!(f64::abs(E_initial - E_final) < 1e-4);
}

#[test]
fn constant_energy_leap_frog() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            LeapFrog::new(units::from(1.0, "fs").unwrap())
        )
    );
    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    assert!(f64::abs(E_initial - E_final) < 1e-5);
}

#[test]
fn perfect_gaz() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );

    // dilating the universe!
    for particle in universe.iter_mut() {
        let pos = particle.position().clone();
        particle.set_position(10.0 * pos);
    }
    universe.set_cell(UnitCell::cubic(100.0));

    simulation.run(&mut universe, 1000);
    let P = universe.pressure();
    let V = universe.volume();
    let T = universe.temperature();
    let N = universe.size() as f64;

    assert!(f64::abs(P * V - N * constants::K_BOLTZMANN * T) < 1e-3);
}

#[test]
fn berendsen_barostat() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();
    let mut md = MolecularDynamics::from_integrator(
        BerendsenBarostat::with_timestep(
            units::from(1.0, "fs").unwrap(),
            units::from(5000.0, "bar").unwrap(),
            100.0
        )
    );
    md.add_control(BerendsenThermostat::new(units::from(300.0, "K").unwrap(), 10.0));
    let mut simulation = Simulation::new(md);

    simulation.run(&mut universe, 5000);
    let P = universe.pressure();
    let T = universe.temperature();

    assert!(f64::abs(P - units::from(5000.0, "bar").unwrap())/P < 8e-2);
    assert!(f64::abs(T - units::from(300.0, "K").unwrap())/T < 1e-2);
}
<commit_msg>Add test for potential computation<commit_after>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/
#![allow(non_snake_case)]
//! Testing physical properties of a Lennard-Jones gaz of Helium.
extern crate cymbalum;
use self::cymbalum::*;

use std::sync::{Once, ONCE_INIT};
static START: Once = ONCE_INIT;

use std::path::Path;

fn get_universe() -> Universe {
    let data_dir = Path::new(file!()).parent().unwrap();
    let configuration = data_dir.join("data").join("helium.xyz");
    let mut universe = Universe::from_file(configuration.to_str().unwrap()).unwrap();
    universe.set_cell(UnitCell::cubic(10.0));

    let mut velocities = BoltzmanVelocities::new(units::from(300.0, "K").unwrap());
    velocities.init(&mut universe);
    return universe;
}

fn get_universe_with_interaction() -> Universe {
    let mut universe = get_universe();
    universe.add_pair_interaction("He", "He",
        Box::new(LennardJones{
            sigma: units::from(2.0, "A").unwrap(),
            epsilon: units::from(0.2, "kJ/mol").unwrap()
        })
    );

    return universe;
}

#[test]
fn constant_energy_velocity_verlet() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe_with_interaction();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );
    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    assert!(f64::abs(E_initial - E_final)/E_final < 1e-3);
}

#[test]
fn constant_energy_verlet() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe_with_interaction();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            Verlet::new(units::from(1.0, "fs").unwrap())
        )
    );
    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    assert!(f64::abs(E_initial - E_final)/E_final < 1e-2);
}

#[test]
fn constant_energy_leap_frog() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe_with_interaction();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            LeapFrog::new(units::from(1.0, "fs").unwrap())
        )
    );
    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    assert!(f64::abs(E_initial - E_final)/E_final < 1e-3);
}

#[test]
fn perfect_gaz() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe_with_interaction();
    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );

    // dilating the universe!
    for particle in universe.iter_mut() {
        let pos = particle.position().clone();
        particle.set_position(10.0 * pos);
    }
    universe.set_cell(UnitCell::cubic(100.0));

    simulation.run(&mut universe, 1000);
    let P = universe.pressure();
    let V = universe.volume();
    let T = universe.temperature();
    let N = universe.size() as f64;

    assert!(f64::abs(P * V - N * constants::K_BOLTZMANN * T) < 1e-3);
}

#[test]
fn berendsen_barostat() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe_with_interaction();
    let mut md = MolecularDynamics::from_integrator(
        BerendsenBarostat::with_timestep(
            units::from(1.0, "fs").unwrap(),
            units::from(5000.0, "bar").unwrap(),
            100.0
        )
    );
    md.add_control(BerendsenThermostat::new(units::from(300.0, "K").unwrap(), 10.0));
    let mut simulation = Simulation::new(md);

    simulation.run(&mut universe, 5000);
    let P = universe.pressure();
    let T = universe.temperature();

    assert!(f64::abs(P - units::from(5000.0, "bar").unwrap())/P < 8e-2);
    assert!(f64::abs(T - units::from(300.0, "K").unwrap())/T < 1e-2);
}

#[test]
fn cutoff_computation() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();

    universe.add_pair_interaction("He", "He",
        Box::new(CutoffComputation::new(
            Box::new(LennardJones{
                sigma: units::from(2.0, "A").unwrap(),
                epsilon: units::from(0.2, "kJ/mol").unwrap()
            }),
            units::from(7.0, "A").unwrap()
        ))
    );

    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );
    simulation.run(&mut universe, 100);

    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    println!("{}", f64::abs(E_initial - E_final)/E_final);
    assert!(f64::abs(E_initial - E_final)/E_final < 2e-1);
}

#[test]
fn table_computation() {
    START.call_once(|| {Logger::stdout();});
    let mut universe = get_universe();

    universe.add_pair_interaction("He", "He",
        Box::new(TableComputation::new(
            Box::new(LennardJones{
                sigma: units::from(2.0, "A").unwrap(),
                epsilon: units::from(0.2, "kJ/mol").unwrap()
            }),
            1000,
            units::from(7.0, "A").unwrap()
        ))
    );

    let mut simulation = Simulation::new(
        MolecularDynamics::from_integrator(
            VelocityVerlet::new(units::from(1.0, "fs").unwrap())
        )
    );
    simulation.run(&mut universe, 100);

    let E_initial = universe.total_energy();
    simulation.run(&mut universe, 1000);
    let E_final = universe.total_energy();
    println!("{}", f64::abs(E_initial - E_final)/E_final);
    assert!(f64::abs(E_initial - E_final)/E_final < 2e-1);
}
<|endoftext|>


<commit_before>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::Keyword;

named!(fits<&[u8], ((Vec<Keyword>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >), Vec<&[u8]>) >,
       pair!(primary_header, many0!( take!(2880) )));

named!(primary_header<&[u8], (Vec<Keyword>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >)>,
       tuple!(
           many0!(keyword_record),
           end_record,
           many0!(blank_record)
       ));

named!(keyword_record<&[u8], Keyword>,
       map!( // TODO should use map_res!
           tuple!(
               take!(8),
               tag!("="),
               take!(71)
           ), |(slice, _, _) : (&[u8], &[u8], &[u8])| {
               match str::from_utf8(slice) {
                   Ok(s) => {
                       match Keyword::from_str(s) {
                           Ok(keyword) => keyword,
                           Err(_) => Keyword::SIMPLE // TODO this is not correct
                       }
                   }
                   Err(_) => Keyword::SIMPLE // TODO this is not correct
               }
           }));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(end_record<&[u8],(&[u8], Vec<&[u8]>)>, pair!(tag!("END"), count!(tag!(" "), 77)));

named!(blank_record<&[u8],Vec<&[u8]> >, count!(tag!(" "), 80));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::Keyword;
    use super::{fits, primary_header, keyword_record, keyword, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, (header, blocks)) => {
                assert_eq!(blocks.len(), 3675);
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, _) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<commit_msg>Check for keyword that is parsed<commit_after>//! The parser module is responsible for parsing FITS files.

use std::str;
use std::str::FromStr;
use super::types::Keyword;

named!(fits<&[u8], ((Vec<Keyword>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >), Vec<&[u8]>) >,
       pair!(primary_header, many0!( take!(2880) )));

named!(primary_header<&[u8], (Vec<Keyword>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >)>,
       tuple!(
           many0!(keyword_record),
           end_record,
           many0!(blank_record)
       ));

named!(keyword_record<&[u8], Keyword>,
       map!( // TODO should use map_res!
           tuple!(
               take!(8),
               tag!("="),
               take!(71)
           ), |(slice, _, _) : (&[u8], &[u8], &[u8])| {
               match str::from_utf8(slice) {
                   Ok(s) => {
                       match Keyword::from_str(s) {
                           Ok(keyword) => keyword,
                           Err(_) => Keyword::SIMPLE // TODO this is not correct
                       }
                   }
                   Err(_) => Keyword::SIMPLE // TODO this is not correct
               }
           }));

named!(keyword<&[u8], Keyword>,
       map_res!(
           map_res!(
               take!(8),
               str::from_utf8),
           Keyword::from_str
       ));

named!(end_record<&[u8],(&[u8], Vec<&[u8]>)>, pair!(tag!("END"), count!(tag!(" "), 77)));

named!(blank_record<&[u8],Vec<&[u8]> >, count!(tag!(" "), 80));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::super::types::Keyword;
    use super::{fits, primary_header, keyword_record, keyword, end_record, blank_record};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, (header, blocks)) => {
                assert_eq!(blocks.len(), 3675);
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, _) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_record_should_parse_a_keyword_record(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = keyword_record(data);

        match result {
            IResult::Done(_,k) => assert_eq!(k, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn keyword_should_parse_a_keyword(){
        let data = "OBJECT  "
            .as_bytes();

        let result = keyword(data);

        match result {
            IResult::Done(_, keyword) => assert_eq!(keyword, Keyword::OBJECT),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn end_record_should_parse_an_END_record(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_record(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn blank_record_should_parse_a_BLANK_record(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_record(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<|endoftext|>


<commit_before>//! Plugin (SIP003)
//!
//! ```plain
//! +------------+                    +---------------------------+
//! |  SS Client +-- Local Loopback --+  Plugin Client (Tunnel)   +--+
//! +------------+                    +---------------------------+  |
//!                                                                  |
//!             Public Internet (Obfuscated/Transformed traffic) ==> |
//!                                                                  |
//! +------------+                    +---------------------------+  |
//! |  SS Server +-- Local Loopback --+  Plugin Server (Tunnel)   +--+
//! +------------+                    +---------------------------+
//! ```

use std::{
    future::Future,
    io::{self, Error, ErrorKind},
    net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use futures::future;
use log::{debug, error, info, trace};
use tokio::{net::TcpStream, process::Child, time};

use crate::config::{Config, ServerAddr};

mod obfs_proxy;
mod ss_plugin;

/// Config for plugin
#[derive(Debug, Clone)]
pub struct PluginConfig {
    pub plugin: String,
    pub plugin_opt: Option<String>,
}

/// Mode of Plugin
#[derive(Debug, Clone, Copy)]
pub enum PluginMode {
    Server,
    Client,
}

/// Started plugins' subprocesses carrier
pub struct Plugins {
    plugins: Vec<Child>,
}

impl Drop for Plugins {
    #[cfg(not(unix))]
    fn drop(&mut self) {
        for plugin in &mut self.plugins {
            debug!("killing plugin process {}", plugin.id());
            let _ = plugin.kill();
        }
    }

    #[cfg(unix)]
    fn drop(&mut self) {
        // Step.1 Send SIGTERM to let them exit gracefully
        for plugin in &mut self.plugins {
            debug!("terminating plugin process {}", plugin.id());

            unsafe {
                let ret = libc::kill(plugin.id() as libc::pid_t, libc::SIGTERM);
                if ret != 0 {
                    let err = io::Error::last_os_error();
                    error!("terminating plugin process {}, error: {}", plugin.id(), err);
                }
            }
        }

        // Step.2 Sit and tight. Let plugins to exit gracefully
        std::thread::sleep(Duration::from_millis(500));

        // Step.3 SIGKILL. Kill all of them forcibly
        for plugin in &mut self.plugins {
            if let Ok(..) = plugin.kill() {
                debug!("killed plugin process {}", plugin.id());
            }
        }
    }
}

impl Plugins {
    /// Launch plugins in configuration.
    ///
    /// Will modify servers' listen addresses to plugins' listen addresses.
    pub fn launch_plugins(config: &mut Config, mode: PluginMode) -> io::Result<Plugins> {
        let mut plugins = Vec::new();

        for svr in &mut config.server {
            let mut svr_addr_opt = None;

            if let Some(c) = svr.plugin() {
                let loop_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
                let local_addr = SocketAddr::new(loop_ip, get_local_port()?);

                let svr_addr = match start_plugin(c, svr.addr(), &local_addr, mode) {
                    Err(err) => {
                        error!(
                            "failed to start plugin \"{}\" for server {}, err: {}",
                            c.plugin,
                            svr.addr(),
                            err
                        );
                        return Err(err);
                    }
                    Ok(process) => {
                        let svr_addr = ServerAddr::SocketAddr(local_addr);
                        plugins.push(process);

                        // Replace addr with plugin
                        svr_addr
                    }
                };

                match mode {
                    PluginMode::Client => info!("started plugin \"{}\" on {} <-> {}", c.plugin, local_addr, svr.addr()),
                    PluginMode::Server => info!("started plugin \"{}\" on {} <-> {}", c.plugin, svr.addr(), local_addr),
                }

                svr_addr_opt = Some(svr_addr); // Fuck borrow checker
            }

            if let Some(svr_addr) = svr_addr_opt {
                svr.set_plugin_addr(svr_addr);
            }
        }

        if plugins.is_empty() {
            panic!("didn't find any plugins to start");
        }

        Ok(Plugins { plugins })
    }

    /// Check plugin started
    ///
    /// This future won't resolve until all plugins are started
    pub async fn check_plugins_started(config: &Config) -> io::Result<()> {
        if !config.has_server_plugins() {
            return Ok(());
        }

        let mut v = Vec::new();

        for svr in &config.server {
            if let Some(ref saddr) = svr.plugin_addr() {
                let addr = match *saddr {
                    ServerAddr::SocketAddr(a) => a,
                    ServerAddr::DomainName(..) => unreachable!("impossible, plugin_addr shouldn't be domain name"),
                };

                v.push(async move {
                    // Try to connect plugin 10 times (nearly 2 seconds)
                    for r in 0..10 {
                        if let Ok(..) = TcpStream::connect(&addr).await {
                            debug!(
                                "plugin \"{}\" for {} listening on {} is started",
                                svr.plugin().as_ref().unwrap().plugin,
                                svr.addr(),
                                addr
                            );

                            return Ok(());
                        }

                        trace!(
                            "plugin \"{}\" for {} listening on {} isn't started yet, tried {} times",
                            svr.plugin().as_ref().unwrap().plugin,
                            svr.addr(),
                            addr,
                            r
                        );

                        time::delay_for(Duration::from_millis(200)).await;
                    }

                    let err = Error::new(ErrorKind::Other, format!("failed to connect plugin \"{}\"", addr));
                    Err(err)
                });
            }
        }

        future::try_join_all(v).await.map(|_| ())
    }
}

impl Future for Plugins {
    type Output = io::Result<()>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        for p in &mut self.plugins {
            match Pin::new(p).poll(cx) {
                Poll::Ready(Ok(exit_status)) => {
                    let msg = format!("plugin exited unexpectedly with {}", exit_status);
                    return Poll::Ready(Err(Error::new(io::ErrorKind::Other, msg)));
                }
                Poll::Ready(Err(err)) => {
                    error!("error while waiting for plugin subprocess: {}", err);
                    return Poll::Ready(Err(err));
                }
                Poll::Pending => {}
            }
        }
        Poll::Pending
    }
}

fn start_plugin(plugin: &PluginConfig, remote: &ServerAddr, local: &SocketAddr, mode: PluginMode) -> io::Result<Child> {
    let mut cmd = if plugin.plugin == "obfsproxy" {
        obfs_proxy::plugin_cmd(plugin, remote, local, mode)
    } else {
        ss_plugin::plugin_cmd(plugin, remote, local, mode)
    };
    cmd.spawn()
}

fn get_local_port() -> io::Result<u16> {
    let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0))?;
    let addr = listener.local_addr()?;
    Ok(addr.port())
}

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

    #[test]
    fn generate_random_port() {
        let port = get_local_port().unwrap();
        println!("{:?}", port);
    }
}
<commit_msg>[#234] Actively wait for plugin subprocesses to exit gracefully<commit_after>//! Plugin (SIP003)
//!
//! ```plain
//! +------------+                    +---------------------------+
//! |  SS Client +-- Local Loopback --+  Plugin Client (Tunnel)   +--+
//! +------------+                    +---------------------------+  |
//!                                                                  |
//!             Public Internet (Obfuscated/Transformed traffic) ==> |
//!                                                                  |
//! +------------+                    +---------------------------+  |
//! |  SS Server +-- Local Loopback --+  Plugin Server (Tunnel)   +--+
//! +------------+                    +---------------------------+
//! ```

use std::{
    future::Future,
    io::{self, Error, ErrorKind},
    net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use futures::future;
use log::{debug, error, info, trace};
use tokio::{net::TcpStream, process::Child, time};

use crate::config::{Config, ServerAddr};

mod obfs_proxy;
mod ss_plugin;

/// Config for plugin
#[derive(Debug, Clone)]
pub struct PluginConfig {
    pub plugin: String,
    pub plugin_opt: Option<String>,
}

/// Mode of Plugin
#[derive(Debug, Clone, Copy)]
pub enum PluginMode {
    Server,
    Client,
}

/// Started plugins' subprocesses carrier
pub struct Plugins {
    plugins: Vec<Child>,
}

impl Drop for Plugins {
    #[cfg(not(unix))]
    fn drop(&mut self) {
        for plugin in &mut self.plugins {
            debug!("killing plugin process {}", plugin.id());
            let _ = plugin.kill();
        }
    }

    #[cfg(unix)]
    fn drop(&mut self) {
        use std::{collections::BTreeSet, time::Instant};

        let mut exited = BTreeSet::new();

        // Step.1 Send SIGTERM to let them exit gracefully
        for plugin in &self.plugins {
            debug!("terminating plugin process {}", plugin.id());

            unsafe {
                let ret = libc::kill(plugin.id() as libc::pid_t, libc::SIGTERM);
                if ret != 0 {
                    let err = io::Error::last_os_error();
                    error!("terminating plugin process {}, error: {}", plugin.id(), err);
                }
            }
        }

        // Step.2 Waits for gracefully exit
        for plugin in &self.plugins {
            const MAX_WAIT_DURATION: Duration = Duration::from_millis(10);

            let start = Instant::now();

            loop {
                unsafe {
                    let mut status: libc::c_int = 0;
                    let ret = libc::waitpid(plugin.id() as libc::pid_t, &mut status, libc::WNOHANG);
                    if ret < 0 {
                        let err = io::Error::last_os_error();
                        error!("waitpid({}) error: {}", plugin.id(), err);
                        break;
                    } else if ret > 0 {
                        // subprocess is finished
                        debug!("plugin process {} is terminated gracefully", plugin.id());
                        exited.insert(plugin.id());
                        break;
                    }
                }

                let elapsed = Instant::now() - start;
                if elapsed > MAX_WAIT_DURATION {
                    debug!(
                        "plugin process {} isn't terminated in {:?}",
                        plugin.id(),
                        MAX_WAIT_DURATION
                    );
                    break;
                }

                std::thread::yield_now();
            }
        }

        // Step.3 SIGKILL. Kill all of them forcibly
        for plugin in &mut self.plugins {
            if exited.contains(&plugin.id()) {
                continue;
            }

            if let Ok(..) = plugin.kill() {
                debug!("killed plugin process {}", plugin.id());
            }
        }
    }
}

impl Plugins {
    /// Launch plugins in configuration.
    ///
    /// Will modify servers' listen addresses to plugins' listen addresses.
    pub fn launch_plugins(config: &mut Config, mode: PluginMode) -> io::Result<Plugins> {
        let mut plugins = Vec::new();

        for svr in &mut config.server {
            let mut svr_addr_opt = None;

            if let Some(c) = svr.plugin() {
                let loop_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
                let local_addr = SocketAddr::new(loop_ip, get_local_port()?);

                let svr_addr = match start_plugin(c, svr.addr(), &local_addr, mode) {
                    Err(err) => {
                        error!(
                            "failed to start plugin \"{}\" for server {}, err: {}",
                            c.plugin,
                            svr.addr(),
                            err
                        );
                        return Err(err);
                    }
                    Ok(process) => {
                        let svr_addr = ServerAddr::SocketAddr(local_addr);
                        plugins.push(process);

                        // Replace addr with plugin
                        svr_addr
                    }
                };

                match mode {
                    PluginMode::Client => info!("started plugin \"{}\" on {} <-> {}", c.plugin, local_addr, svr.addr()),
                    PluginMode::Server => info!("started plugin \"{}\" on {} <-> {}", c.plugin, svr.addr(), local_addr),
                }

                svr_addr_opt = Some(svr_addr); // Fuck borrow checker
            }

            if let Some(svr_addr) = svr_addr_opt {
                svr.set_plugin_addr(svr_addr);
            }
        }

        if plugins.is_empty() {
            panic!("didn't find any plugins to start");
        }

        Ok(Plugins { plugins })
    }

    /// Check plugin started
    ///
    /// This future won't resolve until all plugins are started
    pub async fn check_plugins_started(config: &Config) -> io::Result<()> {
        if !config.has_server_plugins() {
            return Ok(());
        }

        let mut v = Vec::new();

        for svr in &config.server {
            if let Some(ref saddr) = svr.plugin_addr() {
                let addr = match *saddr {
                    ServerAddr::SocketAddr(a) => a,
                    ServerAddr::DomainName(..) => unreachable!("impossible, plugin_addr shouldn't be domain name"),
                };

                v.push(async move {
                    // Try to connect plugin 10 times (nearly 2 seconds)
                    for r in 0..10 {
                        if let Ok(..) = TcpStream::connect(&addr).await {
                            debug!(
                                "plugin \"{}\" for {} listening on {} is started",
                                svr.plugin().as_ref().unwrap().plugin,
                                svr.addr(),
                                addr
                            );

                            return Ok(());
                        }

                        trace!(
                            "plugin \"{}\" for {} listening on {} isn't started yet, tried {} times",
                            svr.plugin().as_ref().unwrap().plugin,
                            svr.addr(),
                            addr,
                            r
                        );

                        time::delay_for(Duration::from_millis(200)).await;
                    }

                    let err = Error::new(ErrorKind::Other, format!("failed to connect plugin \"{}\"", addr));
                    Err(err)
                });
            }
        }

        future::try_join_all(v).await.map(|_| ())
    }
}

impl Future for Plugins {
    type Output = io::Result<()>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        for p in &mut self.plugins {
            match Pin::new(p).poll(cx) {
                Poll::Ready(Ok(exit_status)) => {
                    let msg = format!("plugin exited unexpectedly with {}", exit_status);
                    return Poll::Ready(Err(Error::new(io::ErrorKind::Other, msg)));
                }
                Poll::Ready(Err(err)) => {
                    error!("error while waiting for plugin subprocess: {}", err);
                    return Poll::Ready(Err(err));
                }
                Poll::Pending => {}
            }
        }
        Poll::Pending
    }
}

fn start_plugin(plugin: &PluginConfig, remote: &ServerAddr, local: &SocketAddr, mode: PluginMode) -> io::Result<Child> {
    let mut cmd = if plugin.plugin == "obfsproxy" {
        obfs_proxy::plugin_cmd(plugin, remote, local, mode)
    } else {
        ss_plugin::plugin_cmd(plugin, remote, local, mode)
    };
    cmd.spawn()
}

fn get_local_port() -> io::Result<u16> {
    let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0))?;
    let addr = listener.local_addr()?;
    Ok(addr.port())
}

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

    #[test]
    fn generate_random_port() {
        let port = get_local_port().unwrap();
        println!("{:?}", port);
    }
}
<|endoftext|>


<commit_before>use cursor::Direction;
use buffer::Buffer;
use cursor::Cursor;

use utils;

/// A View is an abstract Window (into a Buffer).
///
/// It draws a portion of a Buffer to a UIBuffer which in turn is drawn to the
/// screen. It maintains the status bar for the current view, the "dirty status"
/// which is whether the buffer has been modified or not and a number of other
/// pieces of information.
pub struct View<'v> {
    pub buffer: Buffer<'v>,
    pub top_line_num: uint,
    pub cursor: Cursor<'v>,
}

impl<'v> View<'v> {
    pub fn new(path: &Path) -> View<'v> {
        let buffer = Buffer::new_from_file(path);
        let mut cursor = Cursor::new();
        cursor.set_line(Some(&buffer.lines[0]));

        View {
            buffer: buffer,
            top_line_num: 0,
            cursor: cursor,
        }
    }

    /// Get the height of the view in which content can be drawn
    ///
    /// Excludes the status bar height
    pub fn get_internal_height(&self) -> uint {
        let term_height = utils::get_term_height();

        term_height - 2
    }

    pub fn draw(&self) {
        let height = self.get_internal_height();

        let num_lines = self.buffer.lines.len();
        let lines_to_draw = self.buffer.lines.slice(self.top_line_num, num_lines);

        for (index, line) in lines_to_draw.iter().enumerate() {
            if index <= height {
                let ln = line.borrow();
                utils::draw(index, ln.data.clone());
            }
        }
    }

    pub fn draw_status(&self) {
        let buffer_status = self.buffer.get_status_text();
        let cursor_status = self.cursor.get_status_text();
        let term_height = utils::get_term_height();

        let status_text = format!("{} {}", buffer_status, cursor_status);

        utils::draw(term_height-1, status_text);
    }

    pub fn move_cursor(&mut self, direction: Direction) {
        let y = self.cursor.get_linenum();
        match direction {
            Direction::Up    => { self.move_cursor_to(y-1); },
            Direction::Down  => { self.move_cursor_to(y+1); },
            Direction::Right => { self.cursor.move_right(); },
            Direction::Left  => { self.cursor.move_left(); },
        }
    }

    pub fn move_cursor_to(&mut self, line_num: uint) {
        let internal_height = self.get_internal_height();
        let total_view_heigh = internal_height + self.top_line_num;

        let num_lines = self.buffer.lines.len() - 1;
        if line_num > num_lines { return }

        // if the desired line is below the top of the view
        // and above the bottom of the view
        if line_num >= self.top_line_num && line_num <= total_view_heigh {
            let line = &self.buffer.lines[line_num];
            self.cursor.set_line(Some(line));
            self.cursor.set_linenum(line_num);
        }

    }

    pub fn delete_char(&mut self) {
        let (offset, line_num) = self.cursor.get_position();

        if offset == 0 {
            let (offset, line_num) = self.buffer.join_line_with_previous(offset, line_num);
            self.move_cursor_to(line_num);
            self.cursor.set_offset(offset);
            return
        }

        self.cursor.delete_char()
    }

    pub fn insert_char(&mut self, ch: char) {
        self.cursor.insert_char(ch);
    }

    pub fn insert_line(&mut self) {
        let (offset, line_num) = self.cursor.get_position();
        let (offset, line_num) = self.buffer.insert_line(offset, line_num);

        self.move_cursor_to(line_num);
        self.cursor.set_offset(offset);
    }
}
<commit_msg>small cleanup<commit_after>use cursor::Direction;
use buffer::Buffer;
use cursor::Cursor;

use utils;

/// A View is an abstract Window (into a Buffer).
///
/// It draws a portion of a Buffer to a UIBuffer which in turn is drawn to the
/// screen. It maintains the status bar for the current view, the "dirty status"
/// which is whether the buffer has been modified or not and a number of other
/// pieces of information.
pub struct View<'v> {
    pub buffer: Buffer<'v>,
    pub top_line_num: uint,
    pub cursor: Cursor<'v>,
}

impl<'v> View<'v> {
    pub fn new(path: &Path) -> View<'v> {
        let buffer = Buffer::new_from_file(path);
        let mut cursor = Cursor::new();
        cursor.set_line(Some(&buffer.lines[0]));

        View {
            buffer: buffer,
            top_line_num: 0,
            cursor: cursor,
        }
    }

    /// Get the height of the view in which content can be drawn
    ///
    /// Excludes the status bar height
    pub fn get_internal_height(&self) -> uint {
        let term_height = utils::get_term_height();

        term_height - 2
    }

    pub fn draw(&self) {
        let height = self.get_internal_height();

        let num_lines = self.buffer.lines.len();
        let lines_to_draw = self.buffer.lines.slice(self.top_line_num, num_lines);

        for (index, line) in lines_to_draw.iter().enumerate() {
            if index <= height {
                let ln = line.borrow();
                utils::draw(index, ln.data.clone());
            }
        }
    }

    pub fn draw_status(&self) {
        let buffer_status = self.buffer.get_status_text();
        let cursor_status = self.cursor.get_status_text();
        let term_height = utils::get_term_height();

        let status_text = format!("{} {} {} {}",
                                  buffer_status,
                                  cursor_status,
                                  self.top_line_num,
                                  self.get_internal_height());

        utils::draw(term_height-1, status_text);
    }

    pub fn move_cursor(&mut self, direction: Direction) {
        let y = self.cursor.get_linenum();
        match direction {
            Direction::Up    => { self.move_cursor_to(y-1); },
            Direction::Down  => { self.move_cursor_to(y+1); },
            Direction::Right => { self.cursor.move_right(); },
            Direction::Left  => { self.cursor.move_left(); },
        }
    }

    pub fn move_cursor_to(&mut self, line_num: uint) {
        let internal_height = self.get_internal_height();
        let calculated_height = internal_height + self.top_line_num;

        let num_lines = self.buffer.lines.len() - 1;

        if line_num <= calculated_height {
            let line = &self.buffer.lines[line_num];
            self.cursor.set_line(Some(line));
            self.cursor.set_linenum(line_num);
        }

    }

    pub fn delete_char(&mut self) {
        let (offset, line_num) = self.cursor.get_position();

        if offset == 0 {
            let (offset, line_num) = self.buffer.join_line_with_previous(offset, line_num);
            self.move_cursor_to(line_num);
            self.cursor.set_offset(offset);
            return
        }

        self.cursor.delete_char()
    }

    pub fn insert_char(&mut self, ch: char) {
        self.cursor.insert_char(ch);
    }

    pub fn insert_line(&mut self) {
        let (offset, line_num) = self.cursor.get_position();
        let (offset, line_num) = self.buffer.insert_line(offset, line_num);

        self.move_cursor_to(line_num);
        self.cursor.set_offset(offset);
    }
}
<|endoftext|>