<commit_before>use std::path::PathBuf;
use std::path::Path;
use std::borrow::Borrow;
use std::ops::Deref;

use semver::Version;
use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;

use error::StoreErrorKind as SEK;
use store::Result;
use store::Store;

/// The Index into the Store
#[derive(Debug, Clone, PartialEq, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    store_location: PathBuf,
    id: PathBuf,
}

impl StoreId {

    pub fn storified(self, store: &Store) -> StoreId {
        if self.starts_with(store.path()) {
            debug!("Not storifying {:?}, because it is already.", self);
            self
        } else {
            debug!("Create new store id out of: {:?} and {:?}", store.path(), self.id);

            let new_id = StoreId { store_location: store.path().clone(), self.id };

            debug!("Created: '{:?}'", new_id);
            new_id
        }
    }

    pub fn exists(&self) -> bool {
        let pb : PathBuf = self.clone().into();
        pb.exists()
    }

    pub fn is_file(&self) -> bool {
        true
    }

    pub fn is_dir(&self) -> bool {
        false
    }

}

impl Into<PathBuf> for StoreId {

    fn into(self) -> PathBuf {
        let mut base = self.store_location;
        base.push(self.id);
        base
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

impl From<PathBuf> for StoreId {

    fn from(pb: PathBuf) -> StoreId {
        StoreId(pb)
    }

}

impl From<String> for StoreId {

    fn from(string: String) -> StoreId {
        StoreId(string.into())
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> StoreId;
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> StoreId {
        StoreId(self)
    }
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> StoreId {
        self
    }
}

pub fn build_entry_path(store: &Store, path_elem: &str) -> Result<PathBuf> {
    debug!("Checking path element for version");
    if path_elem.split('~').last().map_or(false, |v| Version::parse(v).is_err()) {
        debug!("Version cannot be parsed from {:?}", path_elem);
        debug!("Path does not contain version!");
        return Err(SEK::StorePathLacksVersion.into());
    }
    debug!("Version checking succeeded");

    debug!("Building path from {:?}", path_elem);
    let mut path = store.path().clone();

    if path_elem.starts_with('/') {
        path.push(&path_elem[1..]);
    } else {
        path.push(path_elem);
    }

    Ok(path)
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr, $version:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use semver::Version;
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let version = Version::parse($version).unwrap();
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(format!("{}~{}",
                                               name,
                                               version));
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> $crate::storeid::StoreId {
                    StoreId::from(self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {

    use storeid::IntoStoreId;

    module_entry_path_mod!("test", "0.2.0-alpha+leet1337");

    #[test]
    fn correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().to_str().unwrap(), "test/test~0.2.0-alpha+leet1337");
    }

}
<commit_msg>Remove From<String> for StoreId<commit_after>use std::path::PathBuf;
use std::path::Path;
use std::borrow::Borrow;
use std::ops::Deref;

use semver::Version;
use std::fmt::{Display, Debug, Formatter};
use std::fmt::Error as FmtError;
use std::result::Result as RResult;

use error::StoreErrorKind as SEK;
use store::Result;
use store::Store;

/// The Index into the Store
#[derive(Debug, Clone, PartialEq, Hash, Eq, PartialOrd, Ord)]
pub struct StoreId {
    store_location: PathBuf,
    id: PathBuf,
}

impl StoreId {

    pub fn storified(self, store: &Store) -> StoreId {
        if self.starts_with(store.path()) {
            debug!("Not storifying {:?}, because it is already.", self);
            self
        } else {
            debug!("Create new store id out of: {:?} and {:?}", store.path(), self.id);

            let new_id = StoreId { store_location: store.path().clone(), self.id };

            debug!("Created: '{:?}'", new_id);
            new_id
        }
    }

    pub fn exists(&self) -> bool {
        let pb : PathBuf = self.clone().into();
        pb.exists()
    }

    pub fn is_file(&self) -> bool {
        true
    }

    pub fn is_dir(&self) -> bool {
        false
    }

}

impl Into<PathBuf> for StoreId {

    fn into(self) -> PathBuf {
        let mut base = self.store_location;
        base.push(self.id);
        base
    }

}

impl Display for StoreId {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        match self.id.to_str() {
            Some(s) => write!(fmt, "{}", s),
            None    => write!(fmt, "{}", self.id.to_string_lossy()),
        }
    }

}

impl From<PathBuf> for StoreId {

    fn from(pb: PathBuf) -> StoreId {
        StoreId(pb)
    }

}

/// This Trait allows you to convert various representations to a single one
/// suitable for usage in the Store
pub trait IntoStoreId {
    fn into_storeid(self) -> StoreId;
}

impl IntoStoreId for PathBuf {
    fn into_storeid(self) -> StoreId {
        StoreId(self)
    }
}

impl IntoStoreId for StoreId {
    fn into_storeid(self) -> StoreId {
        self
    }
}

pub fn build_entry_path(store: &Store, path_elem: &str) -> Result<PathBuf> {
    debug!("Checking path element for version");
    if path_elem.split('~').last().map_or(false, |v| Version::parse(v).is_err()) {
        debug!("Version cannot be parsed from {:?}", path_elem);
        debug!("Path does not contain version!");
        return Err(SEK::StorePathLacksVersion.into());
    }
    debug!("Version checking succeeded");

    debug!("Building path from {:?}", path_elem);
    let mut path = store.path().clone();

    if path_elem.starts_with('/') {
        path.push(&path_elem[1..]);
    } else {
        path.push(path_elem);
    }

    Ok(path)
}

#[macro_export]
macro_rules! module_entry_path_mod {
    ($name:expr, $version:expr) => (
        #[deny(missing_docs,
                missing_copy_implementations,
                trivial_casts, trivial_numeric_casts,
                unsafe_code,
                unstable_features,
                unused_import_braces, unused_qualifications,
                unused_imports)]
        /// A helper module to create valid module entry paths
        pub mod module_path {
            use semver::Version;
            use std::convert::AsRef;
            use std::path::Path;
            use std::path::PathBuf;

            use $crate::storeid::StoreId;

            /// A Struct giving you the ability to choose store entries assigned
            /// to it.
            ///
            /// It is created through a call to `new`.
            pub struct ModuleEntryPath(PathBuf);

            impl ModuleEntryPath {
                /// Path has to be a valid UTF-8 string or this will panic!
                pub fn new<P: AsRef<Path>>(pa: P) -> ModuleEntryPath {
                    let mut path = PathBuf::new();
                    path.push(format!("{}", $name));
                    path.push(pa.as_ref().clone());
                    let version = Version::parse($version).unwrap();
                    let name = pa.as_ref().file_name().unwrap()
                        .to_str().unwrap();
                    path.set_file_name(format!("{}~{}",
                                               name,
                                               version));
                    ModuleEntryPath(path)
                }
            }

            impl $crate::storeid::IntoStoreId for ModuleEntryPath {
                fn into_storeid(self) -> $crate::storeid::StoreId {
                    StoreId::from(self.0)
                }
            }
        }
    )
}

pub struct StoreIdIterator {
    iter: Box<Iterator<Item = StoreId>>,
}

impl Debug for StoreIdIterator {

    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "StoreIdIterator")
    }

}

impl StoreIdIterator {

    pub fn new(iter: Box<Iterator<Item = StoreId>>) -> StoreIdIterator {
        StoreIdIterator {
            iter: iter,
        }
    }

}

impl Iterator for StoreIdIterator {
    type Item = StoreId;

    fn next(&mut self) -> Option<StoreId> {
        self.iter.next()
    }

}

#[cfg(test)]
mod test {

    use storeid::IntoStoreId;

    module_entry_path_mod!("test", "0.2.0-alpha+leet1337");

    #[test]
    fn correct_path() {
        let p = module_path::ModuleEntryPath::new("test");

        assert_eq!(p.into_storeid().to_str().unwrap(), "test/test~0.2.0-alpha+leet1337");
    }

}
<|endoftext|>


<commit_before>use std::thread::{Thread, JoinGuard};
use std::sync::Arc;
use std::sync::mpsc;
use net::NetworkAcceptor;

pub struct AcceptorPool<A: NetworkAcceptor> {
    acceptor: A
}

impl<A: NetworkAcceptor + 'static> AcceptorPool<A> {
    /// Create a thread pool to manage the acceptor.
    pub fn new(acceptor: A) -> AcceptorPool<A> {
        AcceptorPool { acceptor: acceptor }
    }

    /// Runs the acceptor pool. Blocks until the acceptors are closed.
    ///
    /// ## Panics
    ///
    /// Panics if threads == 0.
    pub fn accept<F>(self, work: F, threads: usize) -> JoinGuard<'static, ()>
        where F: Fn(A::Stream) + Send + Sync + 'static {
        assert!(threads != 0, "Can't accept on 0 threads.");

        // Replace with &F when Send changes land.
        let work = Arc::new(work);

        let (super_tx, supervisor_rx) = mpsc::channel();

        let spawn =
            move || spawn_with(super_tx.clone(), work.clone(), self.acceptor.clone());

        // Go
        for _ in 0..threads { spawn() }

        // Spawn the supervisor
        Thread::scoped(move || for () in supervisor_rx.iter() { spawn() })
    }
}

fn spawn_with<A, F>(supervisor: mpsc::Sender<()>, work: Arc<F>, mut acceptor: A)
where A: NetworkAcceptor + 'static,
      F: Fn(<A as NetworkAcceptor>::Stream) + Send + Sync + 'static {
    use std::old_io::EndOfFile;

    Thread::spawn(move || {
        let sentinel = Sentinel::new(supervisor, ());

        loop {
            match acceptor.accept() {
                Ok(stream) => work(stream),
                Err(ref e) if e.kind == EndOfFile => {
                    debug!("Server closed.");
                    sentinel.cancel();
                    return;
                },

                Err(e) => {
                    error!("Connection failed: {}", e);
                }
            }
        }
    });
}

struct Sentinel<T: Send> {
    value: Option<T>,
    supervisor: mpsc::Sender<T>,
    active: bool
}

impl<T: Send> Sentinel<T> {
    fn new(channel: mpsc::Sender<T>, data: T) -> Sentinel<T> {
        Sentinel {
            value: Some(data),
            supervisor: channel,
            active: true
        }
    }

    fn cancel(mut self) { self.active = false; }
}

#[unsafe_destructor]
impl<T: Send + 'static> Drop for Sentinel<T> {
    fn drop(&mut self) {
        // If we were cancelled, get out of here.
        if !self.active { return; }

        // Respawn ourselves
        let _ = self.supervisor.send(self.value.take().unwrap());
    }
}

<commit_msg>fix(rustup): use module-level thread functions<commit_after>use std::thread::{self, JoinGuard};
use std::sync::Arc;
use std::sync::mpsc;
use net::NetworkAcceptor;

pub struct AcceptorPool<A: NetworkAcceptor> {
    acceptor: A
}

impl<A: NetworkAcceptor + 'static> AcceptorPool<A> {
    /// Create a thread pool to manage the acceptor.
    pub fn new(acceptor: A) -> AcceptorPool<A> {
        AcceptorPool { acceptor: acceptor }
    }

    /// Runs the acceptor pool. Blocks until the acceptors are closed.
    ///
    /// ## Panics
    ///
    /// Panics if threads == 0.
    pub fn accept<F>(self, work: F, threads: usize) -> JoinGuard<'static, ()>
        where F: Fn(A::Stream) + Send + Sync + 'static {
        assert!(threads != 0, "Can't accept on 0 threads.");

        // Replace with &F when Send changes land.
        let work = Arc::new(work);

        let (super_tx, supervisor_rx) = mpsc::channel();

        let spawn =
            move || spawn_with(super_tx.clone(), work.clone(), self.acceptor.clone());

        // Go
        for _ in 0..threads { spawn() }

        // Spawn the supervisor
        thread::scoped(move || for () in supervisor_rx.iter() { spawn() })
    }
}

fn spawn_with<A, F>(supervisor: mpsc::Sender<()>, work: Arc<F>, mut acceptor: A)
where A: NetworkAcceptor + 'static,
      F: Fn(<A as NetworkAcceptor>::Stream) + Send + Sync + 'static {
    use std::old_io::EndOfFile;

    thread::spawn(move || {
        let sentinel = Sentinel::new(supervisor, ());

        loop {
            match acceptor.accept() {
                Ok(stream) => work(stream),
                Err(ref e) if e.kind == EndOfFile => {
                    debug!("Server closed.");
                    sentinel.cancel();
                    return;
                },

                Err(e) => {
                    error!("Connection failed: {}", e);
                }
            }
        }
    });
}

struct Sentinel<T: Send> {
    value: Option<T>,
    supervisor: mpsc::Sender<T>,
    active: bool
}

impl<T: Send> Sentinel<T> {
    fn new(channel: mpsc::Sender<T>, data: T) -> Sentinel<T> {
        Sentinel {
            value: Some(data),
            supervisor: channel,
            active: true
        }
    }

    fn cancel(mut self) { self.active = false; }
}

#[unsafe_destructor]
impl<T: Send + 'static> Drop for Sentinel<T> {
    fn drop(&mut self) {
        // If we were cancelled, get out of here.
        if !self.active { return; }

        // Respawn ourselves
        let _ = self.supervisor.send(self.value.take().unwrap());
    }
}

<|endoftext|>


<commit_before>// Code for creating styled buffers

use crate::snippet::{Style, StyledString};

#[derive(Debug)]
pub struct StyledBuffer {
    text: Vec<Vec<StyledChar>>,
}

#[derive(Debug)]
struct StyledChar {
    chr: char,
    style: Style,
}

impl StyledChar {
    fn new(chr: char, style: Style) -> Self {
        StyledChar { chr, style }
    }
}

impl Default for StyledChar {
    fn default() -> Self {
        StyledChar::new(' ', Style::NoStyle)
    }
}

impl StyledBuffer {
    pub fn new() -> StyledBuffer {
        StyledBuffer { text: vec![] }
    }

    /// Returns content of `StyledBuffer` splitted by lines and line styles
    pub fn render(&self) -> Vec<Vec<StyledString>> {
        // Tabs are assumed to have been replaced by spaces in calling code.
        debug_assert!(self.text.iter().all(|r| !r.iter().any(|sc| sc.chr == '\t')));

        let mut output: Vec<Vec<StyledString>> = vec![];
        let mut styled_vec: Vec<StyledString> = vec![];

        for styled_row in &self.text {
            let mut current_style = Style::NoStyle;
            let mut current_text = String::new();

            for sc in styled_row {
                if sc.style != current_style {
                    if !current_text.is_empty() {
                        styled_vec.push(StyledString { text: current_text, style: current_style });
                    }
                    current_style = sc.style;
                    current_text = String::new();
                }
                current_text.push(sc.chr);
            }
            if !current_text.is_empty() {
                styled_vec.push(StyledString { text: current_text, style: current_style });
            }

            // We're done with the row, push and keep going
            output.push(styled_vec);

            styled_vec = vec![];
        }

        output
    }

    fn ensure_lines(&mut self, line: usize) {
        while line >= self.text.len() {
            self.text.push(vec![]);
        }
    }

    /// Sets `chr` with `style` for given `line`, `col`.
    /// If line not exist in `StyledBuffer`, adds lines up to given
    /// and fills last line with spaces and `Style::NoStyle` style
    pub fn putc(&mut self, line: usize, col: usize, chr: char, style: Style) {
        self.ensure_lines(line);
        if col < self.text[line].len() {
            self.text[line][col] = StyledChar::new(chr, style);
        } else {
            let mut i = self.text[line].len();
            while i < col {
                self.text[line].push(StyledChar::default());
                i += 1;
            }
            self.text[line].push(StyledChar::new(chr, style));
        }
    }

    /// Sets `string` with `style` for given `line`, starting from `col`.
    /// If line not exist in `StyledBuffer`, adds lines up to given
    /// and fills last line with spaces and `Style::NoStyle` style
    pub fn puts(&mut self, line: usize, col: usize, string: &str, style: Style) {
        let mut n = col;
        for c in string.chars() {
            self.putc(line, n, c, style);
            n += 1;
        }
    }

    /// For given `line` inserts `string` with `style` before old content of that line,
    /// adding lines if needed
    pub fn prepend(&mut self, line: usize, string: &str, style: Style) {
        self.ensure_lines(line);
        let string_len = string.chars().count();

        if !self.text[line].is_empty() {
            // Push the old content over to make room for new content
            for _ in 0..string_len {
                self.text[line].insert(0, StyledChar::default());
            }
        }

        self.puts(line, 0, string, style);
    }

    /// For given `line` inserts `string` with `style` after old content of that line,
    /// adding lines if needed
    pub fn append(&mut self, line: usize, string: &str, style: Style) {
        if line >= self.text.len() {
            self.puts(line, 0, string, style);
        } else {
            let col = self.text[line].len();
            self.puts(line, col, string, style);
        }
    }

    pub fn num_lines(&self) -> usize {
        self.text.len()
    }

    /// Set `style` for `line`, `col_start..col_end` range if:
    /// 1. That line and column range exist in `StyledBuffer`
    /// 2. `overwrite` is `true` or existing style is `Style::NoStyle` or `Style::Quotation`
    pub fn set_style_range(
        &mut self,
        line: usize,
        col_start: usize,
        col_end: usize,
        style: Style,
        overwrite: bool,
    ) {
        for col in col_start..col_end {
            self.set_style(line, col, style, overwrite);
        }
    }

    /// Set `style` for `line`, `col` if:
    /// 1. That line and column exist in `StyledBuffer`
    /// 2. Existing style is `Style::NoStyle` or `Style::Quotation` or `overwrite` is `true`
    pub fn set_style(&mut self, line: usize, col: usize, style: Style, overwrite: bool) {
        if let Some(ref mut line) = self.text.get_mut(line) {
            if let Some(StyledChar { style: s, .. }) = line.get_mut(col) {
                if *s == Style::NoStyle || *s == Style::Quotation || overwrite {
                    *s = style;
                }
            }
        }
    }
}
<commit_msg>StyledBuffer::set_style: check overwrite first<commit_after>// Code for creating styled buffers

use crate::snippet::{Style, StyledString};

#[derive(Debug)]
pub struct StyledBuffer {
    text: Vec<Vec<StyledChar>>,
}

#[derive(Debug)]
struct StyledChar {
    chr: char,
    style: Style,
}

impl StyledChar {
    fn new(chr: char, style: Style) -> Self {
        StyledChar { chr, style }
    }
}

impl Default for StyledChar {
    fn default() -> Self {
        StyledChar::new(' ', Style::NoStyle)
    }
}

impl StyledBuffer {
    pub fn new() -> StyledBuffer {
        StyledBuffer { text: vec![] }
    }

    /// Returns content of `StyledBuffer` splitted by lines and line styles
    pub fn render(&self) -> Vec<Vec<StyledString>> {
        // Tabs are assumed to have been replaced by spaces in calling code.
        debug_assert!(self.text.iter().all(|r| !r.iter().any(|sc| sc.chr == '\t')));

        let mut output: Vec<Vec<StyledString>> = vec![];
        let mut styled_vec: Vec<StyledString> = vec![];

        for styled_row in &self.text {
            let mut current_style = Style::NoStyle;
            let mut current_text = String::new();

            for sc in styled_row {
                if sc.style != current_style {
                    if !current_text.is_empty() {
                        styled_vec.push(StyledString { text: current_text, style: current_style });
                    }
                    current_style = sc.style;
                    current_text = String::new();
                }
                current_text.push(sc.chr);
            }
            if !current_text.is_empty() {
                styled_vec.push(StyledString { text: current_text, style: current_style });
            }

            // We're done with the row, push and keep going
            output.push(styled_vec);

            styled_vec = vec![];
        }

        output
    }

    fn ensure_lines(&mut self, line: usize) {
        while line >= self.text.len() {
            self.text.push(vec![]);
        }
    }

    /// Sets `chr` with `style` for given `line`, `col`.
    /// If line not exist in `StyledBuffer`, adds lines up to given
    /// and fills last line with spaces and `Style::NoStyle` style
    pub fn putc(&mut self, line: usize, col: usize, chr: char, style: Style) {
        self.ensure_lines(line);
        if col < self.text[line].len() {
            self.text[line][col] = StyledChar::new(chr, style);
        } else {
            let mut i = self.text[line].len();
            while i < col {
                self.text[line].push(StyledChar::default());
                i += 1;
            }
            self.text[line].push(StyledChar::new(chr, style));
        }
    }

    /// Sets `string` with `style` for given `line`, starting from `col`.
    /// If line not exist in `StyledBuffer`, adds lines up to given
    /// and fills last line with spaces and `Style::NoStyle` style
    pub fn puts(&mut self, line: usize, col: usize, string: &str, style: Style) {
        let mut n = col;
        for c in string.chars() {
            self.putc(line, n, c, style);
            n += 1;
        }
    }

    /// For given `line` inserts `string` with `style` before old content of that line,
    /// adding lines if needed
    pub fn prepend(&mut self, line: usize, string: &str, style: Style) {
        self.ensure_lines(line);
        let string_len = string.chars().count();

        if !self.text[line].is_empty() {
            // Push the old content over to make room for new content
            for _ in 0..string_len {
                self.text[line].insert(0, StyledChar::default());
            }
        }

        self.puts(line, 0, string, style);
    }

    /// For given `line` inserts `string` with `style` after old content of that line,
    /// adding lines if needed
    pub fn append(&mut self, line: usize, string: &str, style: Style) {
        if line >= self.text.len() {
            self.puts(line, 0, string, style);
        } else {
            let col = self.text[line].len();
            self.puts(line, col, string, style);
        }
    }

    pub fn num_lines(&self) -> usize {
        self.text.len()
    }

    /// Set `style` for `line`, `col_start..col_end` range if:
    /// 1. That line and column range exist in `StyledBuffer`
    /// 2. `overwrite` is `true` or existing style is `Style::NoStyle` or `Style::Quotation`
    pub fn set_style_range(
        &mut self,
        line: usize,
        col_start: usize,
        col_end: usize,
        style: Style,
        overwrite: bool,
    ) {
        for col in col_start..col_end {
            self.set_style(line, col, style, overwrite);
        }
    }

    /// Set `style` for `line`, `col` if:
    /// 1. That line and column exist in `StyledBuffer`
    /// 2. `overwrite` is `true` or existing style is `Style::NoStyle` or `Style::Quotation`
    pub fn set_style(&mut self, line: usize, col: usize, style: Style, overwrite: bool) {
        if let Some(ref mut line) = self.text.get_mut(line) {
            if let Some(StyledChar { style: s, .. }) = line.get_mut(col) {
                if overwrite || *s == Style::NoStyle || *s == Style::Quotation {
                    *s = style;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use rustc_middle::thir::{self, *};
use rustc_middle::ty::Const;

pub trait Visitor<'a, 'tcx: 'a>: Sized {
    fn thir(&self) -> &'a Thir<'tcx>;

    fn visit_expr(&mut self, expr: &Expr<'tcx>) {
        walk_expr(self, expr);
    }

    fn visit_stmt(&mut self, stmt: &Stmt<'tcx>) {
        walk_stmt(self, stmt);
    }

    fn visit_block(&mut self, block: &Block) {
        walk_block(self, block);
    }

    fn visit_arm(&mut self, arm: &Arm<'tcx>) {
        walk_arm(self, arm);
    }

    fn visit_const(&mut self, _cnst: &'tcx Const<'tcx>) {}
}

pub fn walk_expr<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, expr: &Expr<'tcx>) {
    use ExprKind::*;
    match expr.kind {
        Scope { value, region_scope: _, lint_level: _ } => {
            visitor.visit_expr(&visitor.thir()[value])
        }
        Box { value } => visitor.visit_expr(&visitor.thir()[value]),
        If { cond, then, else_opt } => {
            visitor.visit_expr(&visitor.thir()[cond]);
            visitor.visit_expr(&visitor.thir()[then]);
            if let Some(else_expr) = else_opt {
                visitor.visit_expr(&visitor.thir()[else_expr]);
            }
        }
        Call { fun, ref args, ty: _, from_hir_call: _, fn_span: _ } => {
            visitor.visit_expr(&visitor.thir()[fun]);
            for &arg in &**args {
                visitor.visit_expr(&visitor.thir()[arg]);
            }
        }
        Deref { arg } => visitor.visit_expr(&visitor.thir()[arg]),
        Binary { lhs, rhs, op: _ } | LogicalOp { lhs, rhs, op: _ } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[rhs]);
        }
        Unary { arg, op: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        Cast { source } => visitor.visit_expr(&visitor.thir()[source]),
        Use { source } => visitor.visit_expr(&visitor.thir()[source]),
        NeverToAny { source } => visitor.visit_expr(&visitor.thir()[source]),
        Pointer { source, cast: _ } => visitor.visit_expr(&visitor.thir()[source]),
        Loop { body } => visitor.visit_expr(&visitor.thir()[body]),
        Match { scrutinee, ref arms } => {
            visitor.visit_expr(&visitor.thir()[scrutinee]);
            for &arm in &**arms {
                visitor.visit_arm(&visitor.thir()[arm]);
            }
        }
        Block { ref body } => visitor.visit_block(body),
        Assign { lhs, rhs } | AssignOp { lhs, rhs, op: _ } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[rhs]);
        }
        Field { lhs, name: _ } => visitor.visit_expr(&visitor.thir()[lhs]),
        Index { lhs, index } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[index]);
        }
        VarRef { id: _ } | UpvarRef { closure_def_id: _, var_hir_id: _ } => {}
        Borrow { arg, borrow_kind: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        AddressOf { arg, mutability: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        Break { value, label: _ } => {
            if let Some(value) = value {
                visitor.visit_expr(&visitor.thir()[value])
            }
        }
        Continue { label: _ } => {}
        Return { value } => {
            if let Some(value) = value {
                visitor.visit_expr(&visitor.thir()[value])
            }
        }
        ConstBlock { value } => visitor.visit_const(value),
        Repeat { value, count } => {
            visitor.visit_expr(&visitor.thir()[value]);
            visitor.visit_const(count);
        }
        Array { ref fields } | Tuple { ref fields } => {
            for &field in &**fields {
                visitor.visit_expr(&visitor.thir()[field]);
            }
        }
        Adt(box thir::Adt {
            ref fields,
            ref base,
            adt_def: _,
            variant_index: _,
            substs: _,
            user_ty: _,
        }) => {
            for field in &**fields {
                visitor.visit_expr(&visitor.thir()[field.expr]);
            }
            if let Some(base) = base {
                visitor.visit_expr(&visitor.thir()[base.base]);
            }
        }
        PlaceTypeAscription { source, user_ty: _ } | ValueTypeAscription { source, user_ty: _ } => {
            visitor.visit_expr(&visitor.thir()[source])
        }
        Closure { closure_id: _, substs: _, upvars: _, movability: _, fake_reads: _ } => {}
        Literal { literal, user_ty: _, const_id: _ } => visitor.visit_const(literal),
        StaticRef { literal, def_id: _ } => visitor.visit_const(literal),
        InlineAsm { ref operands, template: _, options: _, line_spans: _ } => {
            for op in &**operands {
                use InlineAsmOperand::*;
                match op {
                    In { expr, reg: _ }
                    | Out { expr: Some(expr), reg: _, late: _ }
                    | InOut { expr, reg: _, late: _ }
                    | SymFn { expr } => visitor.visit_expr(&visitor.thir()[*expr]),
                    SplitInOut { in_expr, out_expr, reg: _, late: _ } => {
                        visitor.visit_expr(&visitor.thir()[*in_expr]);
                        if let Some(out_expr) = out_expr {
                            visitor.visit_expr(&visitor.thir()[*out_expr]);
                        }
                    }
                    Out { expr: None, reg: _, late: _ }
                    | Const { value: _, span: _ }
                    | SymStatic { def_id: _ } => {}
                }
            }
        }
        ThreadLocalRef(_) => {}
        LlvmInlineAsm { ref outputs, ref inputs, asm: _ } => {
            for &out_expr in &**outputs {
                visitor.visit_expr(&visitor.thir()[out_expr]);
            }
            for &in_expr in &**inputs {
                visitor.visit_expr(&visitor.thir()[in_expr]);
            }
        }
        Yield { value } => visitor.visit_expr(&visitor.thir()[value]),
    }
}

pub fn walk_stmt<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, stmt: &Stmt<'tcx>) {
    match stmt.kind {
        StmtKind::Expr { expr, scope: _ } => visitor.visit_expr(&visitor.thir()[expr]),
        StmtKind::Let {
            initializer,
            remainder_scope: _,
            init_scope: _,
            pattern: _,
            lint_level: _,
        } => {
            if let Some(init) = initializer {
                visitor.visit_expr(&visitor.thir()[init]);
            }
        }
    }
}

pub fn walk_block<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, block: &Block) {
    for &stmt in &*block.stmts {
        visitor.visit_stmt(&visitor.thir()[stmt]);
    }
    if let Some(expr) = block.expr {
        visitor.visit_expr(&visitor.thir()[expr]);
    }
}

pub fn walk_arm<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, arm: &Arm<'tcx>) {
    match arm.guard {
        Some(Guard::If(expr)) => visitor.visit_expr(&visitor.thir()[expr]),
        Some(Guard::IfLet(ref _pat, expr)) => {
            visitor.visit_expr(&visitor.thir()[expr]);
        }
        None => {}
    }
    visitor.visit_expr(&visitor.thir()[arm.body]);
}
<commit_msg>Auto merge of #86378 - Smittyvb:thir-walker-pat, r=LeSeulArtichaut<commit_after>use rustc_middle::thir::{self, *};
use rustc_middle::ty::Const;

pub trait Visitor<'a, 'tcx: 'a>: Sized {
    fn thir(&self) -> &'a Thir<'tcx>;

    fn visit_expr(&mut self, expr: &Expr<'tcx>) {
        walk_expr(self, expr);
    }

    fn visit_stmt(&mut self, stmt: &Stmt<'tcx>) {
        walk_stmt(self, stmt);
    }

    fn visit_block(&mut self, block: &Block) {
        walk_block(self, block);
    }

    fn visit_arm(&mut self, arm: &Arm<'tcx>) {
        walk_arm(self, arm);
    }

    fn visit_pat(&mut self, pat: &Pat<'tcx>) {
        walk_pat(self, pat);
    }

    fn visit_const(&mut self, _cnst: &'tcx Const<'tcx>) {}
}

pub fn walk_expr<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, expr: &Expr<'tcx>) {
    use ExprKind::*;
    match expr.kind {
        Scope { value, region_scope: _, lint_level: _ } => {
            visitor.visit_expr(&visitor.thir()[value])
        }
        Box { value } => visitor.visit_expr(&visitor.thir()[value]),
        If { cond, then, else_opt } => {
            visitor.visit_expr(&visitor.thir()[cond]);
            visitor.visit_expr(&visitor.thir()[then]);
            if let Some(else_expr) = else_opt {
                visitor.visit_expr(&visitor.thir()[else_expr]);
            }
        }
        Call { fun, ref args, ty: _, from_hir_call: _, fn_span: _ } => {
            visitor.visit_expr(&visitor.thir()[fun]);
            for &arg in &**args {
                visitor.visit_expr(&visitor.thir()[arg]);
            }
        }
        Deref { arg } => visitor.visit_expr(&visitor.thir()[arg]),
        Binary { lhs, rhs, op: _ } | LogicalOp { lhs, rhs, op: _ } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[rhs]);
        }
        Unary { arg, op: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        Cast { source } => visitor.visit_expr(&visitor.thir()[source]),
        Use { source } => visitor.visit_expr(&visitor.thir()[source]),
        NeverToAny { source } => visitor.visit_expr(&visitor.thir()[source]),
        Pointer { source, cast: _ } => visitor.visit_expr(&visitor.thir()[source]),
        Loop { body } => visitor.visit_expr(&visitor.thir()[body]),
        Match { scrutinee, ref arms } => {
            visitor.visit_expr(&visitor.thir()[scrutinee]);
            for &arm in &**arms {
                visitor.visit_arm(&visitor.thir()[arm]);
            }
        }
        Block { ref body } => visitor.visit_block(body),
        Assign { lhs, rhs } | AssignOp { lhs, rhs, op: _ } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[rhs]);
        }
        Field { lhs, name: _ } => visitor.visit_expr(&visitor.thir()[lhs]),
        Index { lhs, index } => {
            visitor.visit_expr(&visitor.thir()[lhs]);
            visitor.visit_expr(&visitor.thir()[index]);
        }
        VarRef { id: _ } | UpvarRef { closure_def_id: _, var_hir_id: _ } => {}
        Borrow { arg, borrow_kind: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        AddressOf { arg, mutability: _ } => visitor.visit_expr(&visitor.thir()[arg]),
        Break { value, label: _ } => {
            if let Some(value) = value {
                visitor.visit_expr(&visitor.thir()[value])
            }
        }
        Continue { label: _ } => {}
        Return { value } => {
            if let Some(value) = value {
                visitor.visit_expr(&visitor.thir()[value])
            }
        }
        ConstBlock { value } => visitor.visit_const(value),
        Repeat { value, count } => {
            visitor.visit_expr(&visitor.thir()[value]);
            visitor.visit_const(count);
        }
        Array { ref fields } | Tuple { ref fields } => {
            for &field in &**fields {
                visitor.visit_expr(&visitor.thir()[field]);
            }
        }
        Adt(box thir::Adt {
            ref fields,
            ref base,
            adt_def: _,
            variant_index: _,
            substs: _,
            user_ty: _,
        }) => {
            for field in &**fields {
                visitor.visit_expr(&visitor.thir()[field.expr]);
            }
            if let Some(base) = base {
                visitor.visit_expr(&visitor.thir()[base.base]);
            }
        }
        PlaceTypeAscription { source, user_ty: _ } | ValueTypeAscription { source, user_ty: _ } => {
            visitor.visit_expr(&visitor.thir()[source])
        }
        Closure { closure_id: _, substs: _, upvars: _, movability: _, fake_reads: _ } => {}
        Literal { literal, user_ty: _, const_id: _ } => visitor.visit_const(literal),
        StaticRef { literal, def_id: _ } => visitor.visit_const(literal),
        InlineAsm { ref operands, template: _, options: _, line_spans: _ } => {
            for op in &**operands {
                use InlineAsmOperand::*;
                match op {
                    In { expr, reg: _ }
                    | Out { expr: Some(expr), reg: _, late: _ }
                    | InOut { expr, reg: _, late: _ }
                    | SymFn { expr } => visitor.visit_expr(&visitor.thir()[*expr]),
                    SplitInOut { in_expr, out_expr, reg: _, late: _ } => {
                        visitor.visit_expr(&visitor.thir()[*in_expr]);
                        if let Some(out_expr) = out_expr {
                            visitor.visit_expr(&visitor.thir()[*out_expr]);
                        }
                    }
                    Out { expr: None, reg: _, late: _ }
                    | Const { value: _, span: _ }
                    | SymStatic { def_id: _ } => {}
                }
            }
        }
        ThreadLocalRef(_) => {}
        LlvmInlineAsm { ref outputs, ref inputs, asm: _ } => {
            for &out_expr in &**outputs {
                visitor.visit_expr(&visitor.thir()[out_expr]);
            }
            for &in_expr in &**inputs {
                visitor.visit_expr(&visitor.thir()[in_expr]);
            }
        }
        Yield { value } => visitor.visit_expr(&visitor.thir()[value]),
    }
}

pub fn walk_stmt<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, stmt: &Stmt<'tcx>) {
    match stmt.kind {
        StmtKind::Expr { expr, scope: _ } => visitor.visit_expr(&visitor.thir()[expr]),
        StmtKind::Let {
            initializer,
            remainder_scope: _,
            init_scope: _,
            ref pattern,
            lint_level: _,
        } => {
            if let Some(init) = initializer {
                visitor.visit_expr(&visitor.thir()[init]);
            }
            visitor.visit_pat(pattern);
        }
    }
}

pub fn walk_block<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, block: &Block) {
    for &stmt in &*block.stmts {
        visitor.visit_stmt(&visitor.thir()[stmt]);
    }
    if let Some(expr) = block.expr {
        visitor.visit_expr(&visitor.thir()[expr]);
    }
}

pub fn walk_arm<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, arm: &Arm<'tcx>) {
    match arm.guard {
        Some(Guard::If(expr)) => visitor.visit_expr(&visitor.thir()[expr]),
        Some(Guard::IfLet(ref pat, expr)) => {
            visitor.visit_pat(pat);
            visitor.visit_expr(&visitor.thir()[expr]);
        }
        None => {}
    }
    visitor.visit_pat(&arm.pattern);
    visitor.visit_expr(&visitor.thir()[arm.body]);
}

pub fn walk_pat<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, pat: &Pat<'tcx>) {
    use PatKind::*;
    match pat.kind.as_ref() {
        AscribeUserType { subpattern, ascription: _ }
        | Deref { subpattern }
        | Binding {
            subpattern: Some(subpattern),
            mutability: _,
            mode: _,
            var: _,
            ty: _,
            is_primary: _,
            name: _,
        } => visitor.visit_pat(&subpattern),
        Binding { .. } | Wild => {}
        Variant { subpatterns, adt_def: _, substs: _, variant_index: _ } | Leaf { subpatterns } => {
            for subpattern in subpatterns {
                visitor.visit_pat(&subpattern.pattern);
            }
        }
        Constant { value } => visitor.visit_const(value),
        Range(range) => {
            visitor.visit_const(range.lo);
            visitor.visit_const(range.hi);
        }
        Slice { prefix, slice, suffix } | Array { prefix, slice, suffix } => {
            for subpattern in prefix {
                visitor.visit_pat(&subpattern);
            }
            if let Some(pat) = slice {
                visitor.visit_pat(pat);
            }
            for subpattern in suffix {
                visitor.visit_pat(&subpattern);
            }
        }
        Or { pats } => {
            for pat in pats {
                visitor.visit_pat(&pat);
            }
        }
    };
}
<|endoftext|>


<commit_before>//! SMTP client

use crate::smtp::authentication::{Credentials, Mechanism};
use crate::smtp::client::net::ClientTlsParameters;
use crate::smtp::client::net::{Connector, NetworkStream, Timeout};
use crate::smtp::commands::*;
use crate::smtp::error::{Error, SmtpResult};
use crate::smtp::response::Response;
use bufstream::BufStream;
use log::debug;
use std::fmt::{Debug, Display};
use std::io::{self, BufRead, BufReader, Read, Write};
use std::net::ToSocketAddrs;
use std::string::String;
use std::time::Duration;

pub mod mock;
pub mod net;

/// The codec used for transparency
#[derive(Default, Clone, Copy, Debug)]
#[cfg_attr(feature = "serde-impls", derive(serde::Serialize, serde::Deserialize))]
pub struct ClientCodec {
    escape_count: u8,
}

impl ClientCodec {
    /// Creates a new client codec
    pub fn new() -> Self {
        ClientCodec::default()
    }

    /// Adds transparency
    /// TODO: replace CR and LF by CRLF
    fn encode(&mut self, frame: &[u8], buf: &mut Vec<u8>) -> Result<(), Error> {
        match frame.len() {
            0 => {
                match self.escape_count {
                    0 => buf.write_all(b"\r\n.\r\n")?,
                    1 => buf.write_all(b"\n.\r\n")?,
                    2 => buf.write_all(b".\r\n")?,
                    _ => unreachable!(),
                }
                self.escape_count = 0;
                Ok(())
            }
            _ => {
                let mut start = 0;
                for (idx, byte) in frame.iter().enumerate() {
                    match self.escape_count {
                        0 => self.escape_count = if *byte == b'\r' { 1 } else { 0 },
                        1 => self.escape_count = if *byte == b'\n' { 2 } else { 0 },
                        2 => self.escape_count = if *byte == b'.' { 3 } else { 0 },
                        _ => unreachable!(),
                    }
                    if self.escape_count == 3 {
                        self.escape_count = 0;
                        buf.write_all(&frame[start..idx])?;
                        buf.write_all(b".")?;
                        start = idx;
                    }
                }
                buf.write_all(&frame[start..])?;
                Ok(())
            }
        }
    }
}

/// Returns the string replacing all the CRLF with "\<CRLF\>"
/// Used for debug displays
fn escape_crlf(string: &str) -> String {
    string.replace("\r\n", "<CRLF>")
}

/// Structure that implements the SMTP client
#[derive(Default)]
pub struct InnerClient<S: Write + Read = NetworkStream> {
    /// TCP stream between client and server
    /// Value is None before connection
    stream: Option<BufStream<S>>,
}

macro_rules! return_err (
    ($err: expr, $client: ident) => ({
        return Err(From::from($err))
    })
);

impl<S: Write + Read> InnerClient<S> {
    /// Creates a new SMTP client
    ///
    /// It does not connects to the server, but only creates the `Client`
    pub fn new() -> InnerClient<S> {
        InnerClient { stream: None }
    }
}

impl<S: Connector + Write + Read + Timeout> InnerClient<S> {
    /// Closes the SMTP transaction if possible
    pub fn close(&mut self) {
        let _ = self.command(QuitCommand);
        self.stream = None;
    }

    /// Sets the underlying stream
    pub fn set_stream(&mut self, stream: S) {
        self.stream = Some(BufStream::new(stream));
    }

    /// Upgrades the underlying connection to SSL/TLS
    #[cfg(any(feature = "native-tls", feature = "rustls"))]
    pub fn upgrade_tls_stream(
        &mut self,
        tls_parameters: &ClientTlsParameters,
    ) -> Result<(), Error> {
        match self.stream {
            Some(ref mut stream) => stream.get_mut().upgrade_tls(tls_parameters),
            None => Ok(()),
        }
    }

    /// Tells if the underlying stream is currently encrypted
    pub fn is_encrypted(&self) -> bool {
        self.stream
            .as_ref()
            .map(|s| s.get_ref().is_encrypted())
            .unwrap_or(false)
    }

    /// Set timeout
    pub fn set_timeout(&mut self, duration: Option<Duration>) -> io::Result<()> {
        if let Some(ref mut stream) = self.stream {
            stream.get_mut().set_read_timeout(duration)?;
            stream.get_mut().set_write_timeout(duration)?;
        }
        Ok(())
    }

    /// Connects to the configured server
    pub fn connect<A: ToSocketAddrs>(
        &mut self,
        addr: &A,
        timeout: Option<Duration>,
        tls_parameters: Option<&ClientTlsParameters>,
    ) -> Result<(), Error> {
        // Connect should not be called when the client is already connected
        if self.stream.is_some() {
            return_err!("The connection is already established", self);
        }

        let mut addresses = addr.to_socket_addrs()?;

        let server_addr = match addresses.next() {
            Some(addr) => addr,
            None => return_err!("Could not resolve hostname", self),
        };

        debug!("connecting to {}", server_addr);

        // Try to connect
        self.set_stream(Connector::connect(&server_addr, timeout, tls_parameters)?);

        Ok(())
    }

    /// Checks if the server is connected using the NOOP SMTP command
    #[cfg_attr(feature = "cargo-clippy", allow(clippy::wrong_self_convention))]
    pub fn is_connected(&mut self) -> bool {
        self.stream.is_some() && self.command(NoopCommand).is_ok()
    }

    /// Sends an AUTH command with the given mechanism, and handles challenge if needed
    pub fn auth(&mut self, mechanism: Mechanism, credentials: &Credentials) -> SmtpResult {
        // TODO
        let mut challenges = 10;
        let mut response = self.command(AuthCommand::new(mechanism, credentials.clone(), None)?)?;

        while challenges > 0 && response.has_code(334) {
            challenges -= 1;
            response = self.command(AuthCommand::new_from_response(
                mechanism,
                credentials.clone(),
                &response,
            )?)?;
        }

        if challenges == 0 {
            Err(Error::ResponseParsing("Unexpected number of challenges"))
        } else {
            Ok(response)
        }
    }

    /// Sends the message content
    pub fn message(&mut self, message: Box<dyn Read>) -> SmtpResult {
        let mut out_buf: Vec<u8> = vec![];
        let mut codec = ClientCodec::new();

        let mut message_reader = BufReader::new(message);

        loop {
            out_buf.clear();

            let consumed = match message_reader.fill_buf() {
                Ok(bytes) => {
                    codec.encode(bytes, &mut out_buf)?;
                    bytes.len()
                }
                Err(ref err) => panic!("Failed with: {}", err),
            };
            message_reader.consume(consumed);

            if consumed == 0 {
                break;
            }

            self.write(out_buf.as_slice())?;
        }

        self.write(b"\r\n.\r\n")?;
        self.read_response()
    }

    /// Sends an SMTP command
    pub fn command<C: Display>(&mut self, command: C) -> SmtpResult {
        self.write(command.to_string().as_bytes())?;
        self.read_response()
    }

    /// Writes a string to the server
    fn write(&mut self, string: &[u8]) -> Result<(), Error> {
        if self.stream.is_none() {
            return Err(From::from("Connection closed"));
        }

        self.stream.as_mut().unwrap().write_all(string)?;
        self.stream.as_mut().unwrap().flush()?;

        debug!(
            "Wrote: {}",
            escape_crlf(String::from_utf8_lossy(string).as_ref())
        );
        Ok(())
    }

    /// Gets the SMTP response
    pub fn read_response(&mut self) -> SmtpResult {
        let mut raw_response = String::new();
        let mut response = raw_response.parse::<Response>();

        while response.is_err() {
            if let Error::Parsing(nom::error::ErrorKind::Complete) =
                response.as_ref().err().unwrap()
            {
                break;
            }
            // TODO read more than one line
            let read_count = self.stream.as_mut().unwrap().read_line(&mut raw_response)?;

            // EOF is reached
            if read_count == 0 {
                break;
            }

            response = raw_response.parse::<Response>();
        }

        debug!("Read: {}", escape_crlf(raw_response.as_ref()));

        let final_response = response?;

        if final_response.is_positive() {
            Ok(final_response)
        } else {
            Err(From::from(final_response))
        }
    }
}

#[cfg(test)]
mod test {
    use super::{escape_crlf, ClientCodec};

    #[test]
    fn test_codec() {
        let mut codec = ClientCodec::new();
        let mut buf: Vec<u8> = vec![];

        assert!(codec.encode(b"test\r\n", &mut buf).is_ok());
        assert!(codec.encode(b".\r\n", &mut buf).is_ok());
        assert!(codec.encode(b"\r\ntest", &mut buf).is_ok());
        assert!(codec.encode(b"te\r\n.\r\nst", &mut buf).is_ok());
        assert!(codec.encode(b"test", &mut buf).is_ok());
        assert!(codec.encode(b"test.", &mut buf).is_ok());
        assert!(codec.encode(b"test\n", &mut buf).is_ok());
        assert!(codec.encode(b".test\n", &mut buf).is_ok());
        assert!(codec.encode(b"test", &mut buf).is_ok());
        assert_eq!(
            String::from_utf8(buf).unwrap(),
            "test\r\n..\r\n\r\ntestte\r\n..\r\nsttesttest.test\n.test\ntest"
        );
    }

    #[test]
    fn test_escape_crlf() {
        assert_eq!(escape_crlf("\r\n"), "<CRLF>");
        assert_eq!(escape_crlf("EHLO my_name\r\n"), "EHLO my_name<CRLF>");
        assert_eq!(
            escape_crlf("EHLO my_name\r\nSIZE 42\r\n"),
            "EHLO my_name<CRLF>SIZE 42<CRLF>"
        );
    }
}
<commit_msg>fix(all): Fix conditional use for Debug<commit_after>//! SMTP client

use crate::smtp::authentication::{Credentials, Mechanism};
use crate::smtp::client::net::ClientTlsParameters;
use crate::smtp::client::net::{Connector, NetworkStream, Timeout};
use crate::smtp::commands::*;
use crate::smtp::error::{Error, SmtpResult};
use crate::smtp::response::Response;
use bufstream::BufStream;
use log::debug;
#[cfg(feature = "serde-impls")]
use std::fmt::Debug;
use std::fmt::Display;
use std::io::{self, BufRead, BufReader, Read, Write};
use std::net::ToSocketAddrs;
use std::string::String;
use std::time::Duration;

pub mod mock;
pub mod net;

/// The codec used for transparency
#[derive(Default, Clone, Copy, Debug)]
#[cfg_attr(feature = "serde-impls", derive(serde::Serialize, serde::Deserialize))]
pub struct ClientCodec {
    escape_count: u8,
}

impl ClientCodec {
    /// Creates a new client codec
    pub fn new() -> Self {
        ClientCodec::default()
    }

    /// Adds transparency
    /// TODO: replace CR and LF by CRLF
    fn encode(&mut self, frame: &[u8], buf: &mut Vec<u8>) -> Result<(), Error> {
        match frame.len() {
            0 => {
                match self.escape_count {
                    0 => buf.write_all(b"\r\n.\r\n")?,
                    1 => buf.write_all(b"\n.\r\n")?,
                    2 => buf.write_all(b".\r\n")?,
                    _ => unreachable!(),
                }
                self.escape_count = 0;
                Ok(())
            }
            _ => {
                let mut start = 0;
                for (idx, byte) in frame.iter().enumerate() {
                    match self.escape_count {
                        0 => self.escape_count = if *byte == b'\r' { 1 } else { 0 },
                        1 => self.escape_count = if *byte == b'\n' { 2 } else { 0 },
                        2 => self.escape_count = if *byte == b'.' { 3 } else { 0 },
                        _ => unreachable!(),
                    }
                    if self.escape_count == 3 {
                        self.escape_count = 0;
                        buf.write_all(&frame[start..idx])?;
                        buf.write_all(b".")?;
                        start = idx;
                    }
                }
                buf.write_all(&frame[start..])?;
                Ok(())
            }
        }
    }
}

/// Returns the string replacing all the CRLF with "\<CRLF\>"
/// Used for debug displays
fn escape_crlf(string: &str) -> String {
    string.replace("\r\n", "<CRLF>")
}

/// Structure that implements the SMTP client
#[derive(Default)]
pub struct InnerClient<S: Write + Read = NetworkStream> {
    /// TCP stream between client and server
    /// Value is None before connection
    stream: Option<BufStream<S>>,
}

macro_rules! return_err (
    ($err: expr, $client: ident) => ({
        return Err(From::from($err))
    })
);

impl<S: Write + Read> InnerClient<S> {
    /// Creates a new SMTP client
    ///
    /// It does not connects to the server, but only creates the `Client`
    pub fn new() -> InnerClient<S> {
        InnerClient { stream: None }
    }
}

impl<S: Connector + Write + Read + Timeout> InnerClient<S> {
    /// Closes the SMTP transaction if possible
    pub fn close(&mut self) {
        let _ = self.command(QuitCommand);
        self.stream = None;
    }

    /// Sets the underlying stream
    pub fn set_stream(&mut self, stream: S) {
        self.stream = Some(BufStream::new(stream));
    }

    /// Upgrades the underlying connection to SSL/TLS
    #[cfg(any(feature = "native-tls", feature = "rustls"))]
    pub fn upgrade_tls_stream(
        &mut self,
        tls_parameters: &ClientTlsParameters,
    ) -> Result<(), Error> {
        match self.stream {
            Some(ref mut stream) => stream.get_mut().upgrade_tls(tls_parameters),
            None => Ok(()),
        }
    }

    /// Tells if the underlying stream is currently encrypted
    pub fn is_encrypted(&self) -> bool {
        self.stream
            .as_ref()
            .map(|s| s.get_ref().is_encrypted())
            .unwrap_or(false)
    }

    /// Set timeout
    pub fn set_timeout(&mut self, duration: Option<Duration>) -> io::Result<()> {
        if let Some(ref mut stream) = self.stream {
            stream.get_mut().set_read_timeout(duration)?;
            stream.get_mut().set_write_timeout(duration)?;
        }
        Ok(())
    }

    /// Connects to the configured server
    pub fn connect<A: ToSocketAddrs>(
        &mut self,
        addr: &A,
        timeout: Option<Duration>,
        tls_parameters: Option<&ClientTlsParameters>,
    ) -> Result<(), Error> {
        // Connect should not be called when the client is already connected
        if self.stream.is_some() {
            return_err!("The connection is already established", self);
        }

        let mut addresses = addr.to_socket_addrs()?;

        let server_addr = match addresses.next() {
            Some(addr) => addr,
            None => return_err!("Could not resolve hostname", self),
        };

        debug!("connecting to {}", server_addr);

        // Try to connect
        self.set_stream(Connector::connect(&server_addr, timeout, tls_parameters)?);

        Ok(())
    }

    /// Checks if the server is connected using the NOOP SMTP command
    #[cfg_attr(feature = "cargo-clippy", allow(clippy::wrong_self_convention))]
    pub fn is_connected(&mut self) -> bool {
        self.stream.is_some() && self.command(NoopCommand).is_ok()
    }

    /// Sends an AUTH command with the given mechanism, and handles challenge if needed
    pub fn auth(&mut self, mechanism: Mechanism, credentials: &Credentials) -> SmtpResult {
        // TODO
        let mut challenges = 10;
        let mut response = self.command(AuthCommand::new(mechanism, credentials.clone(), None)?)?;

        while challenges > 0 && response.has_code(334) {
            challenges -= 1;
            response = self.command(AuthCommand::new_from_response(
                mechanism,
                credentials.clone(),
                &response,
            )?)?;
        }

        if challenges == 0 {
            Err(Error::ResponseParsing("Unexpected number of challenges"))
        } else {
            Ok(response)
        }
    }

    /// Sends the message content
    pub fn message(&mut self, message: Box<dyn Read>) -> SmtpResult {
        let mut out_buf: Vec<u8> = vec![];
        let mut codec = ClientCodec::new();

        let mut message_reader = BufReader::new(message);

        loop {
            out_buf.clear();

            let consumed = match message_reader.fill_buf() {
                Ok(bytes) => {
                    codec.encode(bytes, &mut out_buf)?;
                    bytes.len()
                }
                Err(ref err) => panic!("Failed with: {}", err),
            };
            message_reader.consume(consumed);

            if consumed == 0 {
                break;
            }

            self.write(out_buf.as_slice())?;
        }

        self.write(b"\r\n.\r\n")?;
        self.read_response()
    }

    /// Sends an SMTP command
    pub fn command<C: Display>(&mut self, command: C) -> SmtpResult {
        self.write(command.to_string().as_bytes())?;
        self.read_response()
    }

    /// Writes a string to the server
    fn write(&mut self, string: &[u8]) -> Result<(), Error> {
        if self.stream.is_none() {
            return Err(From::from("Connection closed"));
        }

        self.stream.as_mut().unwrap().write_all(string)?;
        self.stream.as_mut().unwrap().flush()?;

        debug!(
            "Wrote: {}",
            escape_crlf(String::from_utf8_lossy(string).as_ref())
        );
        Ok(())
    }

    /// Gets the SMTP response
    pub fn read_response(&mut self) -> SmtpResult {
        let mut raw_response = String::new();
        let mut response = raw_response.parse::<Response>();

        while response.is_err() {
            if let Error::Parsing(nom::error::ErrorKind::Complete) =
                response.as_ref().err().unwrap()
            {
                break;
            }
            // TODO read more than one line
            let read_count = self.stream.as_mut().unwrap().read_line(&mut raw_response)?;

            // EOF is reached
            if read_count == 0 {
                break;
            }

            response = raw_response.parse::<Response>();
        }

        debug!("Read: {}", escape_crlf(raw_response.as_ref()));

        let final_response = response?;

        if final_response.is_positive() {
            Ok(final_response)
        } else {
            Err(From::from(final_response))
        }
    }
}

#[cfg(test)]
mod test {
    use super::{escape_crlf, ClientCodec};

    #[test]
    fn test_codec() {
        let mut codec = ClientCodec::new();
        let mut buf: Vec<u8> = vec![];

        assert!(codec.encode(b"test\r\n", &mut buf).is_ok());
        assert!(codec.encode(b".\r\n", &mut buf).is_ok());
        assert!(codec.encode(b"\r\ntest", &mut buf).is_ok());
        assert!(codec.encode(b"te\r\n.\r\nst", &mut buf).is_ok());
        assert!(codec.encode(b"test", &mut buf).is_ok());
        assert!(codec.encode(b"test.", &mut buf).is_ok());
        assert!(codec.encode(b"test\n", &mut buf).is_ok());
        assert!(codec.encode(b".test\n", &mut buf).is_ok());
        assert!(codec.encode(b"test", &mut buf).is_ok());
        assert_eq!(
            String::from_utf8(buf).unwrap(),
            "test\r\n..\r\n\r\ntestte\r\n..\r\nsttesttest.test\n.test\ntest"
        );
    }

    #[test]
    fn test_escape_crlf() {
        assert_eq!(escape_crlf("\r\n"), "<CRLF>");
        assert_eq!(escape_crlf("EHLO my_name\r\n"), "EHLO my_name<CRLF>");
        assert_eq!(
            escape_crlf("EHLO my_name\r\nSIZE 42\r\n"),
            "EHLO my_name<CRLF>SIZE 42<CRLF>"
        );
    }
}
<|endoftext|>


<commit_before>/*!
Rust bindings to the [sodium library](https://github.com/jedisct1/libsodium).

Sodium is a portable implementation of Dan Bernsteins [NaCl: Networking and
Cryptography library](http://nacl.cr.yp.to)

For most users, if you want public-key (asymmetric) cryptography you should use
the functions in `crypto::asymmetricbox` for encryption/decryption.

If you want secret-key (symmetric) cryptography you should be using the
functions in `crypto::secretbox` for encryption/decryption.

For public-key signatures you should use the functions in `crypto::sign` for
signature creation and verification.

Unless you know what you're doing you most certainly don't want to use the
functions in `crypto::scalarmult`, `crypto::stream`, `crypto::auth` and
`crypto::onetimeauth`.

## Thread Safety
All functions in this library are thread-safe provided that the `init()`
function has been called during program execution.

If `init()` hasn't been called then all functions except the random-number
generation functions and the key-generation functions are thread-safe.

# Public-key cryptography
 `crypto::asymmetricbox`

 `crypto::sign`

# Secret-key cryptography
 `crypto::secretbox`

 `crypto::stream`

 `crypto::auth`

 `crypto::onetimeauth`

# Low-level functions
 `crypto::hash`

 `crypto::verify`

 `crypto::shorthash`
 */
#![crate_name = "sodiumoxide"]
#![comment = "Fast cryptographic library"]
#![license = "MIT"]
#![crate_type = "lib"]
#![warn(missing_doc)]
#![warn(non_uppercase_statics)]
#![warn(non_camel_case_types)]
#![warn(managed_heap_memory)]
#![warn(unnecessary_qualification)]
#![feature(globs)]
#![feature(macro_rules)]
extern crate libc;
use libc::c_int;

#[link(name = "sodium")]
extern {
    fn sodium_init() -> c_int;
}

/**
 * `init()` initializes the sodium library and chooses faster versions of
 * the primitives if possible. `init()` also makes the random number generation
 * functions (`gen_key`, `gen_keypair`, `gen_nonce`, `randombytes`, `randombytes_into`)
 * thread-safe
 */
pub fn init() -> bool {
    unsafe {
        sodium_init() == 0
    }
}

/**
 * Cryptographic functions
 */
pub mod crypto {
    pub mod asymmetricbox;
    pub mod sign;
    pub mod scalarmult;
    pub mod auth;
    pub mod hash;
    pub mod secretbox;
    pub mod onetimeauth;
    pub mod stream;
    pub mod shorthash;

    pub mod verify;
}

pub mod randombytes;
mod utils;
<commit_msg>removed lint for managed_heap_memory<commit_after>/*!
Rust bindings to the [sodium library](https://github.com/jedisct1/libsodium).

Sodium is a portable implementation of Dan Bernsteins [NaCl: Networking and
Cryptography library](http://nacl.cr.yp.to)

For most users, if you want public-key (asymmetric) cryptography you should use
the functions in `crypto::asymmetricbox` for encryption/decryption.

If you want secret-key (symmetric) cryptography you should be using the
functions in `crypto::secretbox` for encryption/decryption.

For public-key signatures you should use the functions in `crypto::sign` for
signature creation and verification.

Unless you know what you're doing you most certainly don't want to use the
functions in `crypto::scalarmult`, `crypto::stream`, `crypto::auth` and
`crypto::onetimeauth`.

## Thread Safety
All functions in this library are thread-safe provided that the `init()`
function has been called during program execution.

If `init()` hasn't been called then all functions except the random-number
generation functions and the key-generation functions are thread-safe.

# Public-key cryptography
 `crypto::asymmetricbox`

 `crypto::sign`

# Secret-key cryptography
 `crypto::secretbox`

 `crypto::stream`

 `crypto::auth`

 `crypto::onetimeauth`

# Low-level functions
 `crypto::hash`

 `crypto::verify`

 `crypto::shorthash`
 */
#![crate_name = "sodiumoxide"]
#![comment = "Fast cryptographic library"]
#![license = "MIT"]
#![crate_type = "lib"]
#![warn(missing_doc)]
#![warn(non_uppercase_statics)]
#![warn(non_camel_case_types)]
#![warn(unnecessary_qualification)]
#![feature(globs)]
#![feature(macro_rules)]
extern crate libc;
use libc::c_int;

#[link(name = "sodium")]
extern {
    fn sodium_init() -> c_int;
}

/**
 * `init()` initializes the sodium library and chooses faster versions of
 * the primitives if possible. `init()` also makes the random number generation
 * functions (`gen_key`, `gen_keypair`, `gen_nonce`, `randombytes`, `randombytes_into`)
 * thread-safe
 */
pub fn init() -> bool {
    unsafe {
        sodium_init() == 0
    }
}

/**
 * Cryptographic functions
 */
pub mod crypto {
    pub mod asymmetricbox;
    pub mod sign;
    pub mod scalarmult;
    pub mod auth;
    pub mod hash;
    pub mod secretbox;
    pub mod onetimeauth;
    pub mod stream;
    pub mod shorthash;

    pub mod verify;
}

pub mod randombytes;
mod utils;
<|endoftext|>


<commit_before>use crate::constants::*;

#[derive(Clone, Debug, PartialEq)]
pub enum Style {
    Old,
    OldInverse,
    New,
    NewInverse,
    Error,
}

#[derive(Clone, Debug, PartialEq)]
pub struct StyledToken {
    token: String,
    style: Style,
}

impl StyledToken {
    pub fn new(token: String, style: Style) -> StyledToken {
        return StyledToken { token, style };
    }

    pub fn is_whitespace(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        if chars_iterator.next().is_some() {
            // Multiple chars found in this token, but whitespace will only be
            // one per token.
            return false;
        }

        return first_char.is_whitespace();
    }

    pub fn is_word(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        let second_char = chars_iterator.next();
        if second_char.is_some() {
            // We consist of multiple characters, that means we are a word
            return true;
        }

        // If we get here, it means our token consists of one character only. If
        // that single character is alphanumeric, we are a word, otherwise not.
        return first_char.is_alphanumeric();
    }
}

pub struct TokenCollector {
    line_prefix: StyledToken,
    tokens: Vec<StyledToken>,
    rendered: bool,
}

impl Style {
    #[must_use]
    pub fn is_inverse(&self) -> bool {
        match self {
            Style::OldInverse | Style::NewInverse | Style::Error => {
                return true;
            }
            _ => {
                return false;
            }
        }
    }

    pub fn inverted(&self) -> Style {
        return match self {
            Style::Old => Style::OldInverse,
            Style::New => Style::NewInverse,
            Style::OldInverse => Style::OldInverse,
            Style::NewInverse => Style::NewInverse,
            Style::Error => Style::Error,
        };
    }

    #[must_use]
    pub fn color<'a>(&self) -> &'a str {
        match self {
            Style::Old => {
                return OLD;
            }
            Style::OldInverse => {
                return OLD;
            }
            Style::New => {
                return NEW;
            }
            Style::NewInverse => {
                return NEW;
            }
            Style::Error => {
                return ERROR;
            }
        }
    }
}

impl TokenCollector {
    #[must_use]
    pub fn create(line_prefix: StyledToken) -> Self {
        return TokenCollector {
            line_prefix,
            tokens: Vec::new(),
            rendered: false,
        };
    }

    pub fn push(&mut self, token: StyledToken) {
        self.tokens.push(token);
    }

    #[must_use]
    fn render_row(&self, row: &mut [StyledToken]) -> String {
        let mut rendered = String::new();

        if self.line_prefix.style == Style::New {
            highlight_trailing_whitespace(row);
            highlight_nonleading_tab(row);
        }
        highlight_space_between_words(row);

        // Set inverse from prefix
        let mut is_inverse = self.line_prefix.style.is_inverse();
        if is_inverse {
            rendered.push_str(INVERSE_VIDEO);
        }

        // Set line color from prefix
        let mut color = self.line_prefix.style.color();
        rendered.push_str(self.line_prefix.style.color());

        // Render prefix
        rendered.push_str(&self.line_prefix.token);

        for token in row {
            if token.style.is_inverse() && !is_inverse {
                rendered.push_str(INVERSE_VIDEO);
            }
            if is_inverse && !token.style.is_inverse() {
                rendered.push_str(NOT_INVERSE_VIDEO);
            }
            is_inverse = token.style.is_inverse();

            if token.style.color() != color {
                rendered.push_str(token.style.color());
                color = token.style.color();
            }

            rendered.push_str(&token.token);
        }

        rendered.push_str(NORMAL);

        return rendered;
    }

    /// Render all the tokens into a (most of the time multiline) string
    #[must_use]
    pub fn render(&mut self) -> String {
        assert!(!self.rendered);
        let mut current_row: Vec<StyledToken> = Vec::new();
        let mut rendered = String::new();

        let mut tokens = std::mem::take(&mut self.tokens);

        // FIXME: Do highlight_space_between_words() first?
        censor_multi_line_highlights(&mut tokens);

        for token in tokens {
            if token.token == "\n" {
                let rendered_row = &self.render_row(&mut current_row);
                rendered.push_str(rendered_row);
                rendered.push('\n');
                current_row.clear();
                continue;
            }

            current_row.push(token);
        }

        if !current_row.is_empty() {
            let rendered_row = &self.render_row(&mut current_row);
            rendered.push_str(rendered_row);
        }

        self.rendered = true;
        return rendered;
    }
}

fn censor_multi_line_highlights(_: &mut [StyledToken]) {
    // FIXME: Code missing here
}

fn highlight_trailing_whitespace(row: &mut [StyledToken]) {
    for token in row.iter_mut().rev() {
        if !token.is_whitespace() {
            return;
        }

        token.style = Style::Error;
    }
}

fn highlight_nonleading_tab(row: &mut [StyledToken]) {
    let mut token_iter = row.iter_mut();

    // Skip leading TABs
    loop {
        let next = token_iter.next();
        if next.is_none() {
            // Done!
            return;
        }

        let token = next.unwrap();
        if token.token != "\t" {
            // Not a TAB, this means we're out of skipping the leading TABs
            break;
        }
    }

    // Scan the rest of the line for non-leading TABs
    for token in token_iter {
        if token.token != "\t" {
            // Not a TAB, never mind
            continue;
        }

        // Non-leading TAB, mark it!
        token.style = Style::Error;
    }
}

/// Highlight single space between two highlighted words
fn highlight_space_between_words(row: &mut [StyledToken]) {
    enum FoundState {
        Nothing,
        HighlightedWord,
        WordSpace,
    }

    let mut found_state = FoundState::Nothing;
    let mut previous_token: Option<&mut StyledToken> = None;
    for token in row.iter_mut() {
        match found_state {
            FoundState::Nothing => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey"
                    found_state = FoundState::HighlightedWord;
                }
            }

            FoundState::HighlightedWord => {
                if token.is_whitespace() {
                    // Found "Monkey " (note trailing space)
                    found_state = FoundState::WordSpace;
                } else if token.style.is_inverse() && token.is_word() {
                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }

            FoundState::WordSpace => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey Dance"
                    if let Some(_previous_token) = previous_token {
                        _previous_token.style = _previous_token.style.inverted();
                    }

                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }
        }

        previous_token = Some(token);
    }
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    #[test]
    fn test_basic() {
        let mut test_me = TokenCollector::create(StyledToken {
            token: "+".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "hej".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "\n".to_string(),
            style: Style::New,
        });

        let rendered = test_me.render();
        assert_eq!(rendered, format!("{}+hej{}\n", NEW, NORMAL));
    }

    #[test]
    fn test_add_trailing_whitespace() {
        // Just a whitespace
        let mut row = [StyledToken::new(" ".to_string(), Style::New)];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(row, [StyledToken::new(" ".to_string(), Style::Error)]);

        // Trailing whitespace
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new(" ".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new(" ".to_string(), Style::Error),
            ]
        );

        // Leading whitespace
        let mut row = [
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new(" ".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );
    }

    #[test]
    fn test_removed_trailing_whitespace() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new(" ".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}- {}", OLD, NORMAL));
    }

    #[test]
    fn test_add_nonleading_tab() {
        // Trailing TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
            ]
        );

        // Middle TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("y".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
                StyledToken::new("y".to_string(), Style::New),
            ]
        );

        // Leading TAB (don't highlight)
        let mut row = [
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("\t".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );

        // Single TAB (don't highlight because it is leading)
        let mut row = [StyledToken::new("\t".to_string(), Style::New)];
        highlight_nonleading_tab(&mut row);
        assert_eq!(row, [StyledToken::new("\t".to_string(), Style::New),]);
    }

    #[test]
    fn test_removed_nonleading_tab() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new("x".to_string(), Style::Old));
        test_me.push(StyledToken::new("\t".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}-x\t{}", OLD, NORMAL));
    }

    #[test]
    fn test_highlight_space_between_words() {
        let mut row = [
            StyledToken::new("Monkey".to_string(), Style::NewInverse),
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("Dance".to_string(), Style::NewInverse),
        ];

        highlight_space_between_words(&mut row);

        assert_eq!(
            row,
            [
                StyledToken::new("Monkey".to_string(), Style::NewInverse),
                StyledToken::new(" ".to_string(), Style::NewInverse),
                StyledToken::new("Dance".to_string(), Style::NewInverse),
            ]
        );
    }

    /// "_" is an unhighlighted token, "." is a highlighted token, "n" is an
    /// unhighlighted newline and "N" is a highlighted newline.
    fn test_censor_multiline_highlighting(input: &str, expected_output: &str) {
        let mut row = Vec::new();
        for c in input.chars() {
            let mut token = "x".to_string();
            if c == 'n' || c == 'N' {
                token = '\n'.to_string();
            }

            let mut style = Style::New;
            if c == '.' || c == 'N' {
                style = Style::NewInverse;
            }

            row.push(StyledToken { token, style });
        }

        censor_multi_line_highlights(&mut row);

        let mut actual_output = String::new();
        for token in row {
            if token.token != "\n" && !token.style.is_inverse() {
                actual_output.push('_');
            } else if token.token != "\n" && token.style.is_inverse() {
                actual_output.push('.');
            } else if token.token == "\n" && !token.style.is_inverse() {
                actual_output.push('n');
            } else if token.token == "\n" && token.style.is_inverse() {
                actual_output.push('N');
            } else {
                panic!("How did we get here?");
            }
        }

        assert_eq!(actual_output, expected_output);
    }

    #[test]
    fn test_censor_multiline_highlights_start_at_beginning() {
        test_censor_multiline_highlighting("...n___", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_at_end() {
        test_censor_multiline_highlighting("___n...", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_midline() {
        test_censor_multiline_highlighting("___n...n___", "___n___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_edgecases() {
        test_censor_multiline_highlighting("n.n", "n_n");
    }

    #[test]
    fn test_censor_multiline_highlights_dont_highlight() {
        // Highlighted line parts should not be censored
        test_censor_multiline_highlighting("_.._n.__.", "_.._n.__.");
    }
}
<commit_msg>Add two more test cases<commit_after>use crate::constants::*;

#[derive(Clone, Debug, PartialEq)]
pub enum Style {
    Old,
    OldInverse,
    New,
    NewInverse,
    Error,
}

#[derive(Clone, Debug, PartialEq)]
pub struct StyledToken {
    token: String,
    style: Style,
}

impl StyledToken {
    pub fn new(token: String, style: Style) -> StyledToken {
        return StyledToken { token, style };
    }

    pub fn is_whitespace(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        if chars_iterator.next().is_some() {
            // Multiple chars found in this token, but whitespace will only be
            // one per token.
            return false;
        }

        return first_char.is_whitespace();
    }

    pub fn is_word(&self) -> bool {
        let mut chars_iterator = self.token.chars();
        let first_char = chars_iterator.next().unwrap();
        let second_char = chars_iterator.next();
        if second_char.is_some() {
            // We consist of multiple characters, that means we are a word
            return true;
        }

        // If we get here, it means our token consists of one character only. If
        // that single character is alphanumeric, we are a word, otherwise not.
        return first_char.is_alphanumeric();
    }
}

pub struct TokenCollector {
    line_prefix: StyledToken,
    tokens: Vec<StyledToken>,
    rendered: bool,
}

impl Style {
    #[must_use]
    pub fn is_inverse(&self) -> bool {
        match self {
            Style::OldInverse | Style::NewInverse | Style::Error => {
                return true;
            }
            _ => {
                return false;
            }
        }
    }

    pub fn inverted(&self) -> Style {
        return match self {
            Style::Old => Style::OldInverse,
            Style::New => Style::NewInverse,
            Style::OldInverse => Style::OldInverse,
            Style::NewInverse => Style::NewInverse,
            Style::Error => Style::Error,
        };
    }

    #[must_use]
    pub fn color<'a>(&self) -> &'a str {
        match self {
            Style::Old => {
                return OLD;
            }
            Style::OldInverse => {
                return OLD;
            }
            Style::New => {
                return NEW;
            }
            Style::NewInverse => {
                return NEW;
            }
            Style::Error => {
                return ERROR;
            }
        }
    }
}

impl TokenCollector {
    #[must_use]
    pub fn create(line_prefix: StyledToken) -> Self {
        return TokenCollector {
            line_prefix,
            tokens: Vec::new(),
            rendered: false,
        };
    }

    pub fn push(&mut self, token: StyledToken) {
        self.tokens.push(token);
    }

    #[must_use]
    fn render_row(&self, row: &mut [StyledToken]) -> String {
        let mut rendered = String::new();

        if self.line_prefix.style == Style::New {
            highlight_trailing_whitespace(row);
            highlight_nonleading_tab(row);
        }
        highlight_space_between_words(row);

        // Set inverse from prefix
        let mut is_inverse = self.line_prefix.style.is_inverse();
        if is_inverse {
            rendered.push_str(INVERSE_VIDEO);
        }

        // Set line color from prefix
        let mut color = self.line_prefix.style.color();
        rendered.push_str(self.line_prefix.style.color());

        // Render prefix
        rendered.push_str(&self.line_prefix.token);

        for token in row {
            if token.style.is_inverse() && !is_inverse {
                rendered.push_str(INVERSE_VIDEO);
            }
            if is_inverse && !token.style.is_inverse() {
                rendered.push_str(NOT_INVERSE_VIDEO);
            }
            is_inverse = token.style.is_inverse();

            if token.style.color() != color {
                rendered.push_str(token.style.color());
                color = token.style.color();
            }

            rendered.push_str(&token.token);
        }

        rendered.push_str(NORMAL);

        return rendered;
    }

    /// Render all the tokens into a (most of the time multiline) string
    #[must_use]
    pub fn render(&mut self) -> String {
        assert!(!self.rendered);
        let mut current_row: Vec<StyledToken> = Vec::new();
        let mut rendered = String::new();

        let mut tokens = std::mem::take(&mut self.tokens);

        // FIXME: Do highlight_space_between_words() first?
        censor_multi_line_highlights(&mut tokens);

        for token in tokens {
            if token.token == "\n" {
                let rendered_row = &self.render_row(&mut current_row);
                rendered.push_str(rendered_row);
                rendered.push('\n');
                current_row.clear();
                continue;
            }

            current_row.push(token);
        }

        if !current_row.is_empty() {
            let rendered_row = &self.render_row(&mut current_row);
            rendered.push_str(rendered_row);
        }

        self.rendered = true;
        return rendered;
    }
}

fn censor_multi_line_highlights(_: &mut [StyledToken]) {
    // FIXME: Code missing here
}

fn highlight_trailing_whitespace(row: &mut [StyledToken]) {
    for token in row.iter_mut().rev() {
        if !token.is_whitespace() {
            return;
        }

        token.style = Style::Error;
    }
}

fn highlight_nonleading_tab(row: &mut [StyledToken]) {
    let mut token_iter = row.iter_mut();

    // Skip leading TABs
    loop {
        let next = token_iter.next();
        if next.is_none() {
            // Done!
            return;
        }

        let token = next.unwrap();
        if token.token != "\t" {
            // Not a TAB, this means we're out of skipping the leading TABs
            break;
        }
    }

    // Scan the rest of the line for non-leading TABs
    for token in token_iter {
        if token.token != "\t" {
            // Not a TAB, never mind
            continue;
        }

        // Non-leading TAB, mark it!
        token.style = Style::Error;
    }
}

/// Highlight single space between two highlighted words
fn highlight_space_between_words(row: &mut [StyledToken]) {
    enum FoundState {
        Nothing,
        HighlightedWord,
        WordSpace,
    }

    let mut found_state = FoundState::Nothing;
    let mut previous_token: Option<&mut StyledToken> = None;
    for token in row.iter_mut() {
        match found_state {
            FoundState::Nothing => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey"
                    found_state = FoundState::HighlightedWord;
                }
            }

            FoundState::HighlightedWord => {
                if token.is_whitespace() {
                    // Found "Monkey " (note trailing space)
                    found_state = FoundState::WordSpace;
                } else if token.style.is_inverse() && token.is_word() {
                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }

            FoundState::WordSpace => {
                if token.style.is_inverse() && token.is_word() {
                    // Found "Monkey Dance"
                    if let Some(_previous_token) = previous_token {
                        _previous_token.style = _previous_token.style.inverted();
                    }

                    found_state = FoundState::HighlightedWord;
                } else {
                    found_state = FoundState::Nothing;
                }
            }
        }

        previous_token = Some(token);
    }
}

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

    #[cfg(test)]
    use pretty_assertions::assert_eq;

    #[test]
    fn test_basic() {
        let mut test_me = TokenCollector::create(StyledToken {
            token: "+".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "hej".to_string(),
            style: Style::New,
        });

        test_me.push(StyledToken {
            token: "\n".to_string(),
            style: Style::New,
        });

        let rendered = test_me.render();
        assert_eq!(rendered, format!("{}+hej{}\n", NEW, NORMAL));
    }

    #[test]
    fn test_add_trailing_whitespace() {
        // Just a whitespace
        let mut row = [StyledToken::new(" ".to_string(), Style::New)];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(row, [StyledToken::new(" ".to_string(), Style::Error)]);

        // Trailing whitespace
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new(" ".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new(" ".to_string(), Style::Error),
            ]
        );

        // Leading whitespace
        let mut row = [
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_trailing_whitespace(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new(" ".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );
    }

    #[test]
    fn test_removed_trailing_whitespace() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new(" ".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}- {}", OLD, NORMAL));
    }

    #[test]
    fn test_add_nonleading_tab() {
        // Trailing TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
            ]
        );

        // Middle TAB
        let mut row = [
            StyledToken::new("x".to_string(), Style::New),
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("y".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("x".to_string(), Style::New),
                StyledToken::new("\t".to_string(), Style::Error),
                StyledToken::new("y".to_string(), Style::New),
            ]
        );

        // Leading TAB (don't highlight)
        let mut row = [
            StyledToken::new("\t".to_string(), Style::New),
            StyledToken::new("x".to_string(), Style::New),
        ];
        highlight_nonleading_tab(&mut row);
        assert_eq!(
            row,
            [
                StyledToken::new("\t".to_string(), Style::New),
                StyledToken::new("x".to_string(), Style::New),
            ]
        );

        // Single TAB (don't highlight because it is leading)
        let mut row = [StyledToken::new("\t".to_string(), Style::New)];
        highlight_nonleading_tab(&mut row);
        assert_eq!(row, [StyledToken::new("\t".to_string(), Style::New),]);
    }

    #[test]
    fn test_removed_nonleading_tab() {
        // It shouldn't be highlighted, just added ones should
        let mut test_me = TokenCollector::create(StyledToken::new("-".to_string(), Style::Old));
        test_me.push(StyledToken::new("x".to_string(), Style::Old));
        test_me.push(StyledToken::new("\t".to_string(), Style::Old));
        let actual = test_me.render();

        assert_eq!(actual, format!("{}-x\t{}", OLD, NORMAL));
    }

    #[test]
    fn test_highlight_space_between_words() {
        let mut row = [
            StyledToken::new("Monkey".to_string(), Style::NewInverse),
            StyledToken::new(" ".to_string(), Style::New),
            StyledToken::new("Dance".to_string(), Style::NewInverse),
        ];

        highlight_space_between_words(&mut row);

        assert_eq!(
            row,
            [
                StyledToken::new("Monkey".to_string(), Style::NewInverse),
                StyledToken::new(" ".to_string(), Style::NewInverse),
                StyledToken::new("Dance".to_string(), Style::NewInverse),
            ]
        );
    }

    /// "_" is an unhighlighted token, "." is a highlighted token, "n" is an
    /// unhighlighted newline and "N" is a highlighted newline.
    fn test_censor_multiline_highlighting(input: &str, expected_output: &str) {
        let mut row = Vec::new();
        for c in input.chars() {
            let mut token = "x".to_string();
            if c == 'n' || c == 'N' {
                token = '\n'.to_string();
            }

            let mut style = Style::New;
            if c == '.' || c == 'N' {
                style = Style::NewInverse;
            }

            row.push(StyledToken { token, style });
        }

        censor_multi_line_highlights(&mut row);

        let mut actual_output = String::new();
        for token in row {
            if token.token != "\n" && !token.style.is_inverse() {
                actual_output.push('_');
            } else if token.token != "\n" && token.style.is_inverse() {
                actual_output.push('.');
            } else if token.token == "\n" && !token.style.is_inverse() {
                actual_output.push('n');
            } else if token.token == "\n" && token.style.is_inverse() {
                actual_output.push('N');
            } else {
                panic!("How did we get here?");
            }
        }

        assert_eq!(actual_output, expected_output);
    }

    #[test]
    fn test_censor_multiline_highlights_start_at_beginning() {
        test_censor_multiline_highlighting("...n___", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_at_end() {
        test_censor_multiline_highlighting("___n...", "___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_midline() {
        test_censor_multiline_highlighting("___n...n___", "___n___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_spans_multiline() {
        test_censor_multiline_highlighting("__.n...n.__", "___n___n___");
    }

    #[test]
    fn test_censor_multiline_highlights_edgecases() {
        test_censor_multiline_highlighting("n.n", "n_n");
        test_censor_multiline_highlighting("...", "___");
    }

    #[test]
    fn test_censor_multiline_highlights_dont_highlight() {
        // Highlighted line parts should not be censored
        test_censor_multiline_highlighting("_.._n.__.", "_.._n.__.");
    }
}
<|endoftext|>