<commit_before>mod parsers;
mod parsing_types;
pub mod preprocessing;
mod utils;

const LINESEP: &str = "\n";
const LATEX_TO_INDENTEX_FACTOR: f64 = 1.5;
const DO_NOT_EDIT_NOTICE: &str =
    "\
     % ============================================================== %\n\
     %                                                                %\n\
     % THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY. %\n\
     %                                                                %\n\
     % ============================================================== %\n";

pub struct TranspileOptions {
    pub prepend_do_not_edit_notice: bool,
}

pub fn transpile(mut lines: Vec<String>, options: &TranspileOptions) -> String {
    use crate::parsers::process_line;
    use crate::parsing_types::Environment;
    use crate::parsing_types::Hashline::{OpenEnv, PlainLine};

    // The number of environments is not known beforehand
    let mut env_stack: Vec<Environment> = Vec::new();

    // Input size is the sum of all line lengths plus the number of lines (for lineseps)
    let input_size = lines.iter().fold(0, |sum, l| sum + l.len()) + lines.len();
    // We do not know how much larger the transpiled LaTeX file will be, but we can guess...
    let indentex_size = (LATEX_TO_INDENTEX_FACTOR * (input_size as f64)).round() as usize;
    let mut transpiled = String::with_capacity(indentex_size);

    let adjusted_indents = crate::utils::scan_indents(lines.as_ref());

    if options.prepend_do_not_edit_notice {
        transpiled.push_str(DO_NOT_EDIT_NOTICE);
    }

    for (line_num, line) in lines.drain(..).enumerate() {
        let list_like_active = match env_stack.last() {
            None => false, // No environment is active at all
            Some(d) => d.is_list_like(),
        };

        let tl = match process_line(line, list_like_active) {
            PlainLine(l) => l,
            OpenEnv(e) => {
                let tag_begin = e.latex_begin();
                env_stack.push(e);
                tag_begin
            }
        };
        transpiled.push_str(&tl);
        transpiled.push_str(LINESEP);

        // Check if we are in an environment and close as many as needed
        while match env_stack.last() {
            None => false,
            Some(d) => d.indent_depth() >= adjusted_indents[line_num + 1],
        } {
            // `unwrap()` is safe here since we have already checked if the stack is empty
            let tag_end = env_stack.pop().unwrap().latex_end();
            transpiled.push_str(&tag_end);
            transpiled.push_str(LINESEP);
        }
    }

    transpiled
}

// LCOV_EXCL_START
#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod transpilation_spec {
        use super::super::{transpile, TranspileOptions};

        #[test]
        fn vanilla_latex_is_not_modified() {
            let input = (vec![
                r"\newcommand{\foo}[1]{foo #1}",
                r"  \begin{equation}1+1\end{equation}",
                r" $ 1 + \frac{1}{2}$",
            ])
            .into_iter()
            .map(|s| s.to_string())
            .collect();

            let expected = "\
                            \\newcommand{\\foo}[1]{foo #1}\
                            \n  \\begin{equation}1+1\\end{equation}\
                            \n $ 1 + \\frac{1}{2}$\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn single_line_commands() {
            let input = (vec!["# foo: bar % qux", "  # foo[opts]: bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\foo{bar} % qux\n  \\foo[opts]{bar}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly() {
            let input = (vec!["# foo:", "  bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  bar\n\\end{foo}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly_if_empty_body() {
            let input = (vec!["# foo:", "  # bar:"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  \\begin{bar}\n  \\end{bar}\n\\end{foo}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn do_not_edit_notice() {
            let to = TranspileOptions {
                prepend_do_not_edit_notice: true,
            };
            assert_eq!(
                &transpile(vec!["".to_string()], &to),
                &format!("{}\n", super::super::DO_NOT_EDIT_NOTICE)
            );
        }
    }
}
// LCOV_EXCL_STOP
<commit_msg>Apply cargo fmt<commit_after>mod parsers;
mod parsing_types;
pub mod preprocessing;
mod utils;

const LINESEP: &str = "\n";
const LATEX_TO_INDENTEX_FACTOR: f64 = 1.5;
const DO_NOT_EDIT_NOTICE: &str = "\
     % ============================================================== %\n\
     %                                                                %\n\
     % THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY. %\n\
     %                                                                %\n\
     % ============================================================== %\n";

pub struct TranspileOptions {
    pub prepend_do_not_edit_notice: bool,
}

pub fn transpile(mut lines: Vec<String>, options: &TranspileOptions) -> String {
    use crate::parsers::process_line;
    use crate::parsing_types::Environment;
    use crate::parsing_types::Hashline::{OpenEnv, PlainLine};

    // The number of environments is not known beforehand
    let mut env_stack: Vec<Environment> = Vec::new();

    // Input size is the sum of all line lengths plus the number of lines (for lineseps)
    let input_size = lines.iter().fold(0, |sum, l| sum + l.len()) + lines.len();
    // We do not know how much larger the transpiled LaTeX file will be, but we can guess...
    let indentex_size = (LATEX_TO_INDENTEX_FACTOR * (input_size as f64)).round() as usize;
    let mut transpiled = String::with_capacity(indentex_size);

    let adjusted_indents = crate::utils::scan_indents(lines.as_ref());

    if options.prepend_do_not_edit_notice {
        transpiled.push_str(DO_NOT_EDIT_NOTICE);
    }

    for (line_num, line) in lines.drain(..).enumerate() {
        let list_like_active = match env_stack.last() {
            None => false, // No environment is active at all
            Some(d) => d.is_list_like(),
        };

        let tl = match process_line(line, list_like_active) {
            PlainLine(l) => l,
            OpenEnv(e) => {
                let tag_begin = e.latex_begin();
                env_stack.push(e);
                tag_begin
            }
        };
        transpiled.push_str(&tl);
        transpiled.push_str(LINESEP);

        // Check if we are in an environment and close as many as needed
        while match env_stack.last() {
            None => false,
            Some(d) => d.indent_depth() >= adjusted_indents[line_num + 1],
        } {
            // `unwrap()` is safe here since we have already checked if the stack is empty
            let tag_end = env_stack.pop().unwrap().latex_end();
            transpiled.push_str(&tag_end);
            transpiled.push_str(LINESEP);
        }
    }

    transpiled
}

// LCOV_EXCL_START
#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod transpilation_spec {
        use super::super::{transpile, TranspileOptions};

        #[test]
        fn vanilla_latex_is_not_modified() {
            let input = (vec![
                r"\newcommand{\foo}[1]{foo #1}",
                r"  \begin{equation}1+1\end{equation}",
                r" $ 1 + \frac{1}{2}$",
            ])
            .into_iter()
            .map(|s| s.to_string())
            .collect();

            let expected = "\
                            \\newcommand{\\foo}[1]{foo #1}\
                            \n  \\begin{equation}1+1\\end{equation}\
                            \n $ 1 + \\frac{1}{2}$\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn single_line_commands() {
            let input = (vec!["# foo: bar % qux", "  # foo[opts]: bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\foo{bar} % qux\n  \\foo[opts]{bar}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly() {
            let input = (vec!["# foo:", "  bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  bar\n\\end{foo}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly_if_empty_body() {
            let input = (vec!["# foo:", "  # bar:"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  \\begin{bar}\n  \\end{bar}\n\\end{foo}\n";
            let to = TranspileOptions {
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn do_not_edit_notice() {
            let to = TranspileOptions {
                prepend_do_not_edit_notice: true,
            };
            assert_eq!(
                &transpile(vec!["".to_string()], &to),
                &format!("{}\n", super::super::DO_NOT_EDIT_NOTICE)
            );
        }
    }
}
// LCOV_EXCL_STOP
<|endoftext|>


<commit_before>use libc;
use glib::translate::*;
use itertools::Itertools;

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum XmlSpace {
    Default,
    Preserve
}

/// Implements `xml:space` handling per the SVG spec
///
/// Normalizes a string as it comes out of the XML parser's handler
/// for character data according to the SVG rules in
/// https://www.w3.org/TR/SVG/text.html#WhiteSpace
pub fn xml_space_normalize(mode: XmlSpace, s: &str) -> String {
    match mode {
        XmlSpace::Default => normalize_default(s),
        XmlSpace::Preserve => normalize_preserve(s)
    }
}

// From https://www.w3.org/TR/SVG/text.html#WhiteSpace
//
// When xml:space="default", the SVG user agent will do the following
// using a copy of the original character data content. First, it will
// remove all newline characters. Then it will convert all tab
// characters into space characters. Then, it will strip off all
// leading and trailing space characters. Then, all contiguous space
// characters will be consolidated.
fn normalize_default(s: &str) -> String {
    s.chars()
        .filter(|ch| *ch != '\n')
        .collect::<String>()
        // split at whitespace, also trims whitespace.
        .split_whitespace()
        .join(" ")
}

// From https://www.w3.org/TR/SVG/text.html#WhiteSpace
//
// When xml:space="preserve", the SVG user agent will do the following
// using a copy of the original character data content. It will
// convert all newline and tab characters into space characters. Then,
// it will draw all space characters, including leading, trailing and
// multiple contiguous space characters. Thus, when drawn with
// xml:space="preserve", the string "a   b" (three spaces between "a"
// and "b") will produce a larger separation between "a" and "b" than
// "a b" (one space between "a" and "b").
fn normalize_preserve(s: &str) -> String {
    let s = s.replace("\n", " ");
    s.replace("\t", " ")
}

#[no_mangle]
pub extern fn rsvg_xml_space_normalize(mode: XmlSpace, s: *const libc::c_char) -> *const libc::c_char {
    let rs = unsafe { String::from_glib_none (s) };

    xml_space_normalize(mode, &rs).to_glib_full()
}

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

    #[test]
    fn xml_space_default() {
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n    WS example\n    indented lines\n  "),
                   "WS example indented lines");
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "WS example indented lines");

        assert_eq!(xml_space_normalize(XmlSpace::Default, "\nWS example\nnon-indented lines\n  "),
                   "WS examplenon-indented lines");
    }

    #[test]
    fn xml_space_preserve() {
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n    WS example\n    indented lines\n  "),
                   "     WS example     indented lines   ");
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "       WS    example      indented lines       ");

    }
}
<commit_msg>Change the signature of the whitespace normalizer functions.<commit_after>use libc;
use glib::translate::*;
use itertools::Itertools;
use std::borrow::Cow;

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum XmlSpace {
    Default,
    Preserve
}

/// Implements `xml:space` handling per the SVG spec
///
/// Normalizes a string as it comes out of the XML parser's handler
/// for character data according to the SVG rules in
/// https://www.w3.org/TR/SVG/text.html#WhiteSpace
pub fn xml_space_normalize(mode: XmlSpace, s: &str) -> String {
    match mode {
        XmlSpace::Default => normalize_default(s),
        XmlSpace::Preserve => normalize_preserve(s)
    }
}

// From https://www.w3.org/TR/SVG/text.html#WhiteSpace
//
// When xml:space="default", the SVG user agent will do the following
// using a copy of the original character data content. First, it will
// remove all newline characters. Then it will convert all tab
// characters into space characters. Then, it will strip off all
// leading and trailing space characters. Then, all contiguous space
// characters will be consolidated.
fn normalize_default<'a, S: Into<Cow<'a, str>>>(s: S) -> String {
    let s = s.into();
    s.chars()
        .filter(|ch| *ch != '\n')
        .collect::<String>()
        // split at whitespace, also trims whitespace.
        .split_whitespace()
        .join(" ")
}

// From https://www.w3.org/TR/SVG/text.html#WhiteSpace
//
// When xml:space="preserve", the SVG user agent will do the following
// using a copy of the original character data content. It will
// convert all newline and tab characters into space characters. Then,
// it will draw all space characters, including leading, trailing and
// multiple contiguous space characters. Thus, when drawn with
// xml:space="preserve", the string "a   b" (three spaces between "a"
// and "b") will produce a larger separation between "a" and "b" than
// "a b" (one space between "a" and "b").
fn normalize_preserve<'a, S: Into<Cow<'a, str>>>(s: S) -> String {
    let s = s.into();
    let s = s.replace("\n", " ");
    s.replace("\t", " ")
}

#[no_mangle]
pub extern fn rsvg_xml_space_normalize(mode: XmlSpace, s: *const libc::c_char) -> *const libc::c_char {
    let rs = unsafe { String::from_glib_none (s) };

    xml_space_normalize(mode, &rs).to_glib_full()
}

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

    #[test]
    fn xml_space_default() {
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n    WS example\n    indented lines\n  "),
                   "WS example indented lines");
        assert_eq!(xml_space_normalize(XmlSpace::Default, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "WS example indented lines");

        assert_eq!(xml_space_normalize(XmlSpace::Default, "\nWS example\nnon-indented lines\n  "),
                   "WS examplenon-indented lines");
    }

    #[test]
    fn xml_space_preserve() {
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n    WS example\n    indented lines\n  "),
                   "     WS example     indented lines   ");
        assert_eq!(xml_space_normalize(XmlSpace::Preserve, "\n  \t  \tWS \t\t\texample\n  \t  indented lines\t\t  \n  "),
                   "       WS    example      indented lines       ");

    }
}
<|endoftext|>


<commit_before>#![allow(clippy::new_without_default)]

//! Interprets the terminal events we care about (keyboard input).
//!
//! Keyboard events are read from `stdin`. We look for byte strings of key combinations that we
//! care about. E.g. Alt-arrow keys, C-w etc.
//!
//! NOTE: Make sure to either set `stdin` to non-blocking mode or enable canonical input. For the
//! latter, see:
//!
//! - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
//! - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
//!
//! Also make sure to set `VMIN` and `VTIME` to zero in non-canonical mode.

#[cfg(test)]
mod tests;

use std::char;
use std::collections::VecDeque;
use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin;
use std::task::{Context, Poll};

use tokio::io::unix::AsyncFd;
use tokio::stream::Stream;

use term_input_macros::byte_seq_parser;

////////////////////////////////////////////////////////////////////////////////////////////////////
// Public types
////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Key {
    AltArrow(Arrow),
    AltChar(char),
    Arrow(Arrow),
    Backspace,
    Char(char),
    Ctrl(char),
    CtrlArrow(Arrow),
    Del,
    End,
    Esc,
    Home,
    PageDown,
    PageUp,
    ShiftDown,
    ShiftUp,
    Tab,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arrow {
    Left,
    Right,
    Up,
    Down,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Event {
    /// A single key input.
    Key(Key),

    /// Usually a paste.
    String(String),

    // These are not generated anymore. I needed these for a feature, but never really got around
    // implementing it.
    // FocusGained,
    // FocusLost,
    /// An unknown sequence of bytes (probably for a key combination that we don't care about).
    Unknown(Vec<u8>),
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Byte sequences of key presses we want to capture
////////////////////////////////////////////////////////////////////////////////////////////////////

byte_seq_parser! {
    parse_key_bytes -> Key, // Function name. Generation function type is
                            // fn(&[u8]) -> Option<(Key, usize)>.

    [27, 91, 49, 59, 51, 66] => Key::AltArrow(Arrow::Down),
    [27, 91, 49, 59, 51, 68] => Key::AltArrow(Arrow::Left),
    [27, 91, 49, 59, 51, 67] => Key::AltArrow(Arrow::Right),
    [27, 91, 49, 59, 51, 65] => Key::AltArrow(Arrow::Up),
    [27, 91, 66] => Key::Arrow(Arrow::Down),
    [27, 79, 66] => Key::Arrow(Arrow::Down),
    [27, 91, 68] => Key::Arrow(Arrow::Left),
    [27, 79, 68] => Key::Arrow(Arrow::Left),
    [27, 91, 67] => Key::Arrow(Arrow::Right),
    [27, 79, 67] => Key::Arrow(Arrow::Right),
    [27, 91, 65] => Key::Arrow(Arrow::Up),
    [27, 79, 65] => Key::Arrow(Arrow::Up),
    [27, 91, 49, 59, 53, 66] => Key::CtrlArrow(Arrow::Down),
    [27, 91, 49, 59, 53, 68] => Key::CtrlArrow(Arrow::Left),
    [27, 91, 49, 59, 53, 67] => Key::CtrlArrow(Arrow::Right),
    [27, 91, 49, 59, 53, 65] => Key::CtrlArrow(Arrow::Up),
    [27, 91, 51, 126] => Key::Del,
    [27, 91, 54, 126] => Key::PageDown,
    [27, 91, 53, 126] => Key::PageUp,
    [27, 91, 49, 59, 50, 65] => Key::ShiftUp,
    [27, 91, 49, 59, 50, 66] => Key::ShiftDown,
    [27, 91, 72] => Key::Home,
    [27, 91, 70] => Key::End,
    [27, 79, 72] => Key::Home,
    [27, 79, 70] => Key::End,
    [27, 91, 52, 126] => Key::End,
    [9] => Key::Tab,
    [127] => Key::Backspace,
    [1] => Key::Ctrl('a'),
    [5] => Key::Ctrl('e'),
    [23] => Key::Ctrl('w'),
    [11] => Key::Ctrl('k'),
    [4] => Key::Ctrl('d'),
    [3] => Key::Ctrl('c'),
    [17] => Key::Ctrl('q'),
    [16] => Key::Ctrl('p'),
    [14] => Key::Ctrl('n'),
    [21] => Key::Ctrl('u'),
    [24] => Key::Ctrl('x'),
}

// static XTERM_FOCUS_GAINED: [u8; 3] = [27, 91, 73];
// static XTERM_FOCUS_LOST: [u8; 3] = [27, 91, 79];

////////////////////////////////////////////////////////////////////////////////////////////////////

struct RawFd_(RawFd);

// TODO: Remove this when `AsRawFd for RawFd` is stable.
impl AsRawFd for RawFd_ {
    fn as_raw_fd(&self) -> RawFd {
        self.0
    }
}

pub struct Input {
    /// Queue of events waiting to be polled.
    evs: VecDeque<Event>,

    /// Used when reading from stdin.
    buf: Vec<u8>,

    stdin: AsyncFd<RawFd_>,
}

impl Input {
    /// Create a new input handler. Requires a `tokio` reactor to be running. Make sure to either
    /// set `stdin` to non-blocking mode or enable non-canonical input, otherwise the `Stream`
    /// implementation will not work. If you are using this with `termbox` it does the latter so no
    /// need to do anything. See module documentation for links on how to enable non-canonical
    /// input.
    pub fn new() -> Input {
        Input {
            evs: VecDeque::new(),
            buf: Vec::with_capacity(100),
            stdin: AsyncFd::new(RawFd_(libc::STDIN_FILENO)).unwrap(),
        }
    }

    fn parse_buffer(&mut self) {
        let mut buf_slice: &[u8] = &self.buf;

        while !buf_slice.is_empty() {
            // Special treatment for 127 (backspace, 0x1B) and 13 ('\r', 0xD)
            let fst = buf_slice[0];
            let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
                parse_key_comb
            } else {
                parse_chars
            };

            match parse_fn(buf_slice) {
                Some((ev, used)) => {
                    buf_slice = &buf_slice[used..];
                    self.evs.push_back(ev);
                }
                None => {
                    self.evs.push_back(Event::Unknown(buf_slice.to_owned()));
                    break;
                }
            }
        }

        self.buf.clear();
    }
}

impl Stream for Input {
    type Item = std::io::Result<Event>;

    fn poll_next(mut self: Pin<&mut Input>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let self_: &mut Input = &mut *self;

        'main: loop {
            // Try to parse any bytes in the input buffer from the last poll
            self_.parse_buffer();

            // Yield pending events
            if let Some(ev) = self_.evs.pop_front() {
                return Poll::Ready(Some(Ok(ev)));
            }

            // Otherwise read stdin and loop if successful
            let mut poll_ret = self_.stdin.poll_read_ready(cx);
            loop {
                match poll_ret {
                    Poll::Ready(Ok(mut ready)) => {
                        let read_ret = read_stdin(&mut self_.buf);

                        // read_stdin reads until stdin is empty, poll again for readiness
                        ready.clear_ready();
                        poll_ret = self_.stdin.poll_read_ready(cx);

                        match read_ret {
                            Ok(()) => {}
                            Err(err) => {
                                // NOTE: `poll_ret` is ignored here but I think that's OK?
                                let err =
                                    std::io::Error::from(err.as_errno().expect("Weird nix error"));
                                return Poll::Ready(Some(Err(err)));
                            }
                        }

                        if self_.buf.is_empty() {
                            continue;
                        } else {
                            continue 'main;
                        }
                    }
                    Poll::Ready(Err(err)) => {
                        return Poll::Ready(Some(Err(err)));
                    }
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub(crate) fn parse_single_event(buf: &[u8]) -> Event {
    let fst = buf[0];
    let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
        parse_key_comb
    } else {
        parse_chars
    };

    let (ev, used) = match parse_fn(buf) {
        Some((ev, used)) => (ev, used),
        None => (Event::Unknown(buf.to_owned()), buf.len()),
    };

    assert_eq!(buf.len(), used);

    ev
}

fn parse_chars(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    // Use a fast path for the common case: Single utf-8 character.
    // (TODO: What about other encodings?)

    utf8_char_len(buf[0]).map(|char_len| {
        if char_len as usize == buf.len() {
            // Fast path: single character
            let ev = Event::Key(Key::Char(get_utf8_char(buf, char_len)));
            (ev, char_len as usize)
        } else {
            // Probably a paste: allocate a string and collect chars
            let mut string = String::with_capacity(1000);
            let mut start_idx = 0;
            loop {
                if start_idx == buf.len() {
                    break;
                }
                match utf8_char_len(buf[start_idx]) {
                    Some(char_len) => {
                        string.push(get_utf8_char(&buf[start_idx..], char_len));
                        start_idx += char_len as usize;
                    }
                    None => {
                        break;
                    }
                }
            }
            let ev = Event::String(string);
            (ev, start_idx)
        }
    })
}

fn parse_key_comb(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    if let Some((key, used)) = parse_key_bytes(buf) {
        return Some((Event::Key(key), used));
    }

    if buf[0] == 27 {
        // 0x1B, ESC
        // 27 not followed by anything is an actual ESC
        if buf.len() == 1 {
            return Some((Event::Key(Key::Esc), 1));
        }
        // otherwise it's probably alt + key
        else {
            debug_assert!(buf.len() >= 2);
            return utf8_char_len(buf[1]).map(|char_len| {
                let ev = Event::Key(Key::AltChar(get_utf8_char(&buf[1..], char_len)));
                (ev, char_len as usize + 1)
            });
        }
    }

    None
}

fn utf8_char_len(byte: u8) -> Option<u8> {
    if byte >> 7 == 0b0 {
        Some(1)
    } else if byte >> 5 == 0b110 {
        Some(2)
    } else if byte >> 4 == 0b1110 {
        Some(3)
    } else if byte >> 3 == 0b11110 {
        Some(4)
    } else {
        None
    }
}

fn get_utf8_char(buf: &[u8], len: u8) -> char {
    let codepoint: u32 = {
        if len == 1 {
            u32::from(buf[0] & 0b0111_1111)
        } else if len == 2 {
            ((u32::from(buf[0] & 0b0001_1111)) << 6) + (u32::from(buf[1] & 0b0011_1111))
        } else if len == 3 {
            ((u32::from(buf[0] & 0b0000_1111)) << 12)
                + ((u32::from(buf[1] & 0b0011_1111)) << 6)
                + (u32::from(buf[2] & 0b0011_1111))
        } else {
            debug_assert!(len == 4);
            ((u32::from(buf[0] & 0b0000_0111)) << 18)
                + ((u32::from(buf[1] & 0b0011_1111)) << 12)
                + ((u32::from(buf[2] & 0b0011_1111)) << 6)
                + (u32::from(buf[3] & 0b0011_1111))
        }
    };

    char::from_u32(codepoint).unwrap()
}

/// Read `stdin` until `read` fails with `EWOULDBLOCK` (happens in canonical mode, when `stdin` is
/// set to non-blocking mode) or returns 0 (happens in non-canonical mode when `VMIN` and `VTIME`
/// are 0). If you are using `term_input` with `termbox` then you don't need to set `stdin` to
/// non-blocking mode as `termbox` enables non-canonical mode.
///
/// See also:
///
/// - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
/// - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
pub fn read_stdin(buf: &mut Vec<u8>) -> Result<(), nix::Error> {
    loop {
        let old_len = buf.len();
        buf.reserve(100);
        unsafe {
            buf.set_len(old_len + 100);
        }

        match nix::unistd::read(libc::STDIN_FILENO, &mut buf[old_len..]) {
            Ok(n_read) => {
                unsafe { buf.set_len(old_len + n_read) };
                if n_read == 0 {
                    // We're in non-canonical mode, or stdin is closed. We can't distinguish the
                    // two here but I think it's fine to return OK when stdin is closed.
                    return Ok(());
                }
            }
            Err(err) => {
                unsafe { buf.set_len(old_len) };
                match err {
                    nix::Error::Sys(nix::errno::EWOULDBLOCK) => {
                        return Ok(());
                    }
                    _ => {
                        return Err(err);
                    }
                }
            }
        }
    }
}
<commit_msg>term_input: Fix a potential race, more comments<commit_after>#![allow(clippy::new_without_default)]

//! Interprets the terminal events we care about (keyboard input).
//!
//! Keyboard events are read from `stdin`. We look for byte strings of key combinations that we
//! care about. E.g. Alt-arrow keys, C-w etc.
//!
//! NOTE: Make sure to either set `stdin` to non-blocking mode or enable canonical input. For the
//! latter, see:
//!
//! - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
//! - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
//!
//! Also make sure to set `VMIN` and `VTIME` to zero in non-canonical mode.

#[cfg(test)]
mod tests;

use std::char;
use std::collections::VecDeque;
use std::os::unix::io::{AsRawFd, RawFd};
use std::pin::Pin;
use std::task::{Context, Poll};

use tokio::io::unix::AsyncFd;
use tokio::stream::Stream;

use term_input_macros::byte_seq_parser;

////////////////////////////////////////////////////////////////////////////////////////////////////
// Public types
////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Key {
    AltArrow(Arrow),
    AltChar(char),
    Arrow(Arrow),
    Backspace,
    Char(char),
    Ctrl(char),
    CtrlArrow(Arrow),
    Del,
    End,
    Esc,
    Home,
    PageDown,
    PageUp,
    ShiftDown,
    ShiftUp,
    Tab,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arrow {
    Left,
    Right,
    Up,
    Down,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Event {
    /// A single key input.
    Key(Key),

    /// Usually a paste.
    String(String),

    // These are not generated anymore. I needed these for a feature, but never really got around
    // implementing it.
    // FocusGained,
    // FocusLost,
    /// An unknown sequence of bytes (probably for a key combination that we don't care about).
    Unknown(Vec<u8>),
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Byte sequences of key presses we want to capture
////////////////////////////////////////////////////////////////////////////////////////////////////

byte_seq_parser! {
    parse_key_bytes -> Key, // Function name. Generation function type is
                            // fn(&[u8]) -> Option<(Key, usize)>.

    [27, 91, 49, 59, 51, 66] => Key::AltArrow(Arrow::Down),
    [27, 91, 49, 59, 51, 68] => Key::AltArrow(Arrow::Left),
    [27, 91, 49, 59, 51, 67] => Key::AltArrow(Arrow::Right),
    [27, 91, 49, 59, 51, 65] => Key::AltArrow(Arrow::Up),
    [27, 91, 66] => Key::Arrow(Arrow::Down),
    [27, 79, 66] => Key::Arrow(Arrow::Down),
    [27, 91, 68] => Key::Arrow(Arrow::Left),
    [27, 79, 68] => Key::Arrow(Arrow::Left),
    [27, 91, 67] => Key::Arrow(Arrow::Right),
    [27, 79, 67] => Key::Arrow(Arrow::Right),
    [27, 91, 65] => Key::Arrow(Arrow::Up),
    [27, 79, 65] => Key::Arrow(Arrow::Up),
    [27, 91, 49, 59, 53, 66] => Key::CtrlArrow(Arrow::Down),
    [27, 91, 49, 59, 53, 68] => Key::CtrlArrow(Arrow::Left),
    [27, 91, 49, 59, 53, 67] => Key::CtrlArrow(Arrow::Right),
    [27, 91, 49, 59, 53, 65] => Key::CtrlArrow(Arrow::Up),
    [27, 91, 51, 126] => Key::Del,
    [27, 91, 54, 126] => Key::PageDown,
    [27, 91, 53, 126] => Key::PageUp,
    [27, 91, 49, 59, 50, 65] => Key::ShiftUp,
    [27, 91, 49, 59, 50, 66] => Key::ShiftDown,
    [27, 91, 72] => Key::Home,
    [27, 91, 70] => Key::End,
    [27, 79, 72] => Key::Home,
    [27, 79, 70] => Key::End,
    [27, 91, 52, 126] => Key::End,
    [9] => Key::Tab,
    [127] => Key::Backspace,
    [1] => Key::Ctrl('a'),
    [5] => Key::Ctrl('e'),
    [23] => Key::Ctrl('w'),
    [11] => Key::Ctrl('k'),
    [4] => Key::Ctrl('d'),
    [3] => Key::Ctrl('c'),
    [17] => Key::Ctrl('q'),
    [16] => Key::Ctrl('p'),
    [14] => Key::Ctrl('n'),
    [21] => Key::Ctrl('u'),
    [24] => Key::Ctrl('x'),
}

// static XTERM_FOCUS_GAINED: [u8; 3] = [27, 91, 73];
// static XTERM_FOCUS_LOST: [u8; 3] = [27, 91, 79];

////////////////////////////////////////////////////////////////////////////////////////////////////

struct RawFd_(RawFd);

// TODO: Remove this when `AsRawFd for RawFd` is stable.
impl AsRawFd for RawFd_ {
    fn as_raw_fd(&self) -> RawFd {
        self.0
    }
}

pub struct Input {
    /// Queue of events waiting to be polled.
    evs: VecDeque<Event>,

    /// Used when reading from stdin.
    buf: Vec<u8>,

    stdin: AsyncFd<RawFd_>,
}

impl Input {
    /// Create a new input handler. Requires a `tokio` reactor to be running. Make sure to either
    /// set `stdin` to non-blocking mode or enable non-canonical input, otherwise the `Stream`
    /// implementation will not work. If you are using this with `termbox` it does the latter so no
    /// need to do anything. See module documentation for links on how to enable non-canonical
    /// input.
    pub fn new() -> Input {
        Input {
            evs: VecDeque::new(),
            buf: Vec::with_capacity(100),
            stdin: AsyncFd::new(RawFd_(libc::STDIN_FILENO)).unwrap(),
        }
    }

    fn parse_buffer(&mut self) {
        let mut buf_slice: &[u8] = &self.buf;

        while !buf_slice.is_empty() {
            // Special treatment for 127 (backspace, 0x1B) and 13 ('\r', 0xD)
            let fst = buf_slice[0];
            let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
                parse_key_comb
            } else {
                parse_chars
            };

            match parse_fn(buf_slice) {
                Some((ev, used)) => {
                    buf_slice = &buf_slice[used..];
                    self.evs.push_back(ev);
                }
                None => {
                    self.evs.push_back(Event::Unknown(buf_slice.to_owned()));
                    break;
                }
            }
        }

        self.buf.clear();
    }
}

impl Stream for Input {
    type Item = std::io::Result<Event>;

    fn poll_next(mut self: Pin<&mut Input>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let self_: &mut Input = &mut *self;

        'main: loop {
            // Try to parse any bytes in the input buffer from the last poll
            self_.parse_buffer();

            // Yield pending events
            if let Some(ev) = self_.evs.pop_front() {
                return Poll::Ready(Some(Ok(ev)));
            }

            // Otherwise read stdin and loop if successful
            let mut poll_ret = self_.stdin.poll_read_ready(cx);
            loop {
                match poll_ret {
                    Poll::Ready(Ok(mut ready)) => {
                        // We will read stdin until it's empty, clear the readiness state now
                        // (before reading) to avoid a potential race that would occur if stdin
                        // becomes available again between the last read and clear_ready.
                        ready.clear_ready();

                        // Read stdin until it fails with EWOULDBLOCK or returns 0
                        let read_ret = read_stdin(&mut self_.buf);

                        // Poll again so that we will be notified when stdin becomes ready again
                        poll_ret = self_.stdin.poll_read_ready(cx);

                        match read_ret {
                            Ok(()) => {}
                            Err(err) => {
                                // NOTE: `poll_ret` is ignored here but I think that's OK?
                                // TODO: Contents of `buf` is ignored here until `poll_next` is
                                // called again.
                                let err =
                                    std::io::Error::from(err.as_errno().expect("Weird nix error"));
                                return Poll::Ready(Some(Err(err)));
                            }
                        }

                        if self_.buf.is_empty() {
                            continue;
                        } else {
                            continue 'main;
                        }
                    }
                    Poll::Ready(Err(err)) => {
                        return Poll::Ready(Some(Err(err)));
                    }
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub(crate) fn parse_single_event(buf: &[u8]) -> Event {
    let fst = buf[0];
    let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
        parse_key_comb
    } else {
        parse_chars
    };

    let (ev, used) = match parse_fn(buf) {
        Some((ev, used)) => (ev, used),
        None => (Event::Unknown(buf.to_owned()), buf.len()),
    };

    assert_eq!(buf.len(), used);

    ev
}

fn parse_chars(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    // Use a fast path for the common case: Single utf-8 character.
    // (TODO: What about other encodings?)

    utf8_char_len(buf[0]).map(|char_len| {
        if char_len as usize == buf.len() {
            // Fast path: single character
            let ev = Event::Key(Key::Char(get_utf8_char(buf, char_len)));
            (ev, char_len as usize)
        } else {
            // Probably a paste: allocate a string and collect chars
            let mut string = String::with_capacity(1000);
            let mut start_idx = 0;
            loop {
                if start_idx == buf.len() {
                    break;
                }
                match utf8_char_len(buf[start_idx]) {
                    Some(char_len) => {
                        string.push(get_utf8_char(&buf[start_idx..], char_len));
                        start_idx += char_len as usize;
                    }
                    None => {
                        break;
                    }
                }
            }
            let ev = Event::String(string);
            (ev, start_idx)
        }
    })
}

fn parse_key_comb(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    if let Some((key, used)) = parse_key_bytes(buf) {
        return Some((Event::Key(key), used));
    }

    if buf[0] == 27 {
        // 0x1B, ESC
        // 27 not followed by anything is an actual ESC
        if buf.len() == 1 {
            return Some((Event::Key(Key::Esc), 1));
        }
        // otherwise it's probably alt + key
        else {
            debug_assert!(buf.len() >= 2);
            return utf8_char_len(buf[1]).map(|char_len| {
                let ev = Event::Key(Key::AltChar(get_utf8_char(&buf[1..], char_len)));
                (ev, char_len as usize + 1)
            });
        }
    }

    None
}

fn utf8_char_len(byte: u8) -> Option<u8> {
    if byte >> 7 == 0b0 {
        Some(1)
    } else if byte >> 5 == 0b110 {
        Some(2)
    } else if byte >> 4 == 0b1110 {
        Some(3)
    } else if byte >> 3 == 0b11110 {
        Some(4)
    } else {
        None
    }
}

fn get_utf8_char(buf: &[u8], len: u8) -> char {
    let codepoint: u32 = {
        if len == 1 {
            u32::from(buf[0] & 0b0111_1111)
        } else if len == 2 {
            ((u32::from(buf[0] & 0b0001_1111)) << 6) + (u32::from(buf[1] & 0b0011_1111))
        } else if len == 3 {
            ((u32::from(buf[0] & 0b0000_1111)) << 12)
                + ((u32::from(buf[1] & 0b0011_1111)) << 6)
                + (u32::from(buf[2] & 0b0011_1111))
        } else {
            debug_assert!(len == 4);
            ((u32::from(buf[0] & 0b0000_0111)) << 18)
                + ((u32::from(buf[1] & 0b0011_1111)) << 12)
                + ((u32::from(buf[2] & 0b0011_1111)) << 6)
                + (u32::from(buf[3] & 0b0011_1111))
        }
    };

    char::from_u32(codepoint).unwrap()
}

/// Read `stdin` until `read` fails with `EWOULDBLOCK` (happens in canonical mode, when `stdin` is
/// set to non-blocking mode) or returns 0 (happens in non-canonical mode when `VMIN` and `VTIME`
/// are 0). If you are using `term_input` with `termbox` then you don't need to set `stdin` to
/// non-blocking mode as `termbox` enables non-canonical mode.
///
/// See also:
///
/// - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
/// - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
pub fn read_stdin(buf: &mut Vec<u8>) -> Result<(), nix::Error> {
    loop {
        let old_len = buf.len();
        buf.reserve(100);
        unsafe {
            buf.set_len(old_len + 100);
        }

        match nix::unistd::read(libc::STDIN_FILENO, &mut buf[old_len..]) {
            Ok(n_read) => {
                unsafe { buf.set_len(old_len + n_read) };
                if n_read == 0 {
                    // We're in non-canonical mode, or stdin is closed. We can't distinguish the
                    // two here but I think it's fine to return OK when stdin is closed.
                    return Ok(());
                }
            }
            Err(err) => {
                unsafe { buf.set_len(old_len) };
                match err {
                    nix::Error::Sys(nix::errno::EWOULDBLOCK) => {
                        return Ok(());
                    }
                    _ => {
                        return Err(err);
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// aux-build:issue-15562.rs

// pretty-expanded FIXME #23616

extern crate issue_15562 as i;

pub fn main() {
    extern {
        fn transmute();
    }
    unsafe {
        transmute();
        i::transmute();
    }
}

// We declare this so we don't run into unresolved symbol errors
// The above extern is NOT `extern "rust-intrinsic"` and thus
// means it'll try to find a corresponding symbol to link to.
#[no_mangle]
pub extern fn transmute() {}
<commit_msg>Fix a broken test<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// aux-build:issue-15562.rs

// pretty-expanded FIXME #23616

extern crate issue_15562 as i;

pub fn main() {
    unsafe {
        transmute();
        i::transmute();
    }
}

// We declare this so we don't run into unresolved symbol errors
// The above extern is NOT `extern "rust-intrinsic"` and thus
// means it'll try to find a corresponding symbol to link to.
#[no_mangle]
pub extern fn transmute() {}
<|endoftext|>


<commit_before>// run-pass
// ignore-wasm32-bare no libc to test ffi with
// ignore-sgx no libc
#![feature(rustc_private)]

extern crate libc;

use std::mem;
use libc::{c_double, c_int};

fn to_c_int(v: &mut isize) -> &mut c_int {
    unsafe {
        mem::transmute_copy(&v)
    }
}

fn lgamma(n: c_double, value: &mut isize) -> c_double {
    unsafe {
        return m::lgamma(n, to_c_int(value));
    }
}

mod m {
    use libc::{c_double, c_int};

    #[link(name = "m")]
    extern {
        #[cfg(any(all(unix, not(target_os = "vxworks")), target_os = "cloudabi"))]
        #[link_name="lgamma_r"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
        #[cfg(windows)]
        #[link_name="lgamma"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
        #[cfg(target_os = "vxworks")]
        #[link_name="lgamma"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
    }
}

pub fn main() {
  let mut y: isize = 5;
  let x: &mut isize = &mut y;
  assert_eq!(lgamma(1.0 as c_double, x), 0.0 as c_double);
}
<commit_msg>Fix broken test on MSVC<commit_after>// run-pass
// ignore-wasm32-bare no libc to test ffi with
// ignore-sgx no libc
#![feature(rustc_private)]

extern crate libc;

use std::mem;
use libc::{c_double, c_int};

fn to_c_int(v: &mut isize) -> &mut c_int {
    unsafe {
        mem::transmute_copy(&v)
    }
}

fn lgamma(n: c_double, value: &mut isize) -> c_double {
    unsafe {
        return m::lgamma(n, to_c_int(value));
    }
}

mod m {
    use libc::{c_double, c_int};

    extern {
        #[cfg(any(all(unix, not(target_os = "vxworks")), target_os = "cloudabi"))]
        #[link_name="lgamma_r"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
        #[cfg(windows)]
        #[link_name="lgamma"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
        #[cfg(target_os = "vxworks")]
        #[link_name="lgamma"]
        pub fn lgamma(n: c_double, sign: &mut c_int) -> c_double;
    }
}

pub fn main() {
  let mut y: isize = 5;
  let x: &mut isize = &mut y;
  assert_eq!(lgamma(1.0 as c_double, x), 0.0 as c_double);
}
<|endoftext|>


<commit_before>#[cfg(unix)]
extern crate nix;
#[cfg(windows)]
extern crate libc;

use std::io::{Error, ErrorKind, Result};
use std::net::{UdpSocket, SocketAddr};

/// A trait to make time-limited reads from socket-like objects.
pub trait WithReadTimeout {
    /// Receives data from the object, blocking for at most the specified number of milliseconds.
    /// On success, returns the number of bytes read and the address from whence the data came.  If
    /// the timeout expires, it returns `ErrorKind::WouldBlock`.
    fn recv_timeout(&mut self, &mut [u8], i64) -> Result<(usize, SocketAddr)>;
}

impl WithReadTimeout for UdpSocket {
    #[cfg(unix)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use nix::sys::socket::{SockLevel, sockopt, setsockopt};
        use nix::sys::time::TimeVal;
        use std::os::unix::io::AsRawFd;

        setsockopt(self.as_raw_fd(),
                   SockLevel::Socket,
                   sockopt::ReceiveTimeout,
                   &TimeVal::milliseconds(timeout)).unwrap();

        fn map_os_error(e: Error) -> Error {
            // TODO: Replace with constant from libc
            const EAGAIN: i32 = 35;

            match e.raw_os_error() {
                Some(EAGAIN) => Error::new(ErrorKind::WouldBlock, ""),
                _ => e
            }
        }
        self.recv_from(buf).map_err(map_os_error)
    }

    #[cfg(windows)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use select::fd_set;
        use std::os::windows::io::AsRawSocket;
        use libc;

        // Initialize relevant data structures
        let mut readfds = fd_set::new();
        let null = std::ptr::null_mut();

        fd_set(&mut readfds, self.as_raw_socket());

        // Set timeout
        let mut tv = libc::timeval {
            tv_sec: timeout as i32 / 1000,
            tv_usec: (timeout as i32 % 1000) * 1000,
        };

        // In Windows, the first argument to `select` is ignored.
        let retval = unsafe { select::select(0, &mut readfds, null, null, &mut tv) };
        if retval == 0 {
            return Err(Error::new(ErrorKind::TimedOut, "Time limit expired"));
        } else if retval < 0 {
            return Err(Error::last_os_error());
        }

        self.recv_from(buf)
    }
}

// Most of the following was copied from 'rust/src/libstd/sys/windows/net.rs'
#[cfg(windows)]
mod select {
    use libc;

    pub const FD_SETSIZE: usize = 64;

    #[repr(C)]
    pub struct fd_set {
        fd_count: libc::c_uint,
        fd_array: [libc::SOCKET; FD_SETSIZE],
    }

    pub fn fd_set(set: &mut fd_set, s: libc::SOCKET) {
        set.fd_array[set.fd_count as usize] = s;
        set.fd_count += 1;
    }

    impl fd_set {
        pub fn new() -> fd_set {
            fd_set {
                fd_count: 0,
                fd_array: [0; FD_SETSIZE],
            }
        }
    }

    #[link(name = "ws2_32")]
    extern "system" {
        pub fn select(nfds: libc::c_int,
                      readfds: *mut fd_set,
                      writefds: *mut fd_set,
                      exceptfds: *mut fd_set,
                      timeout: *mut libc::timeval) -> libc::c_int;
    }
}

#[test]
fn test_socket_timeout() {
    let mut socket = UdpSocket::bind("0.0.0.0:0").unwrap();
    let mut buf = [0; 10];
    assert!(socket.recv_timeout(&mut buf, 100).is_err());
}
<commit_msg>Don't call `select` if the timeout is non-positive.<commit_after>#[cfg(unix)]
extern crate nix;
#[cfg(windows)]
extern crate libc;

use std::io::{Error, ErrorKind, Result};
use std::net::{UdpSocket, SocketAddr};

/// A trait to make time-limited reads from socket-like objects.
pub trait WithReadTimeout {
    /// Receives data from the object, blocking for at most the specified number of milliseconds.
    /// On success, returns the number of bytes read and the address from whence the data came.  If
    /// the timeout expires, it returns `ErrorKind::WouldBlock`.
    fn recv_timeout(&mut self, &mut [u8], i64) -> Result<(usize, SocketAddr)>;
}

impl WithReadTimeout for UdpSocket {
    #[cfg(unix)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use nix::sys::socket::{SockLevel, sockopt, setsockopt};
        use nix::sys::time::TimeVal;
        use std::os::unix::io::AsRawFd;

        if timeout > 0 {
            setsockopt(self.as_raw_fd(),
                       SockLevel::Socket,
                       sockopt::ReceiveTimeout,
                       &TimeVal::milliseconds(timeout)).unwrap();
        }

        fn map_os_error(e: Error) -> Error {
            // TODO: Replace with constant from libc
            const EAGAIN: i32 = 35;

            match e.raw_os_error() {
                Some(EAGAIN) => Error::new(ErrorKind::WouldBlock, ""),
                _ => e
            }
        }

        self.recv_from(buf).map_err(map_os_error)
    }

    #[cfg(windows)]
    fn recv_timeout(&mut self, buf: &mut [u8], timeout: i64) -> Result<(usize, SocketAddr)> {
        use select::fd_set;
        use std::os::windows::io::AsRawSocket;
        use libc;

        if timeout > 0 {
            // Initialize relevant data structures
            let mut readfds = fd_set::new();
            let null = std::ptr::null_mut();

            fd_set(&mut readfds, self.as_raw_socket());

            // Set timeout
            let mut tv = libc::timeval {
                tv_sec: timeout as i32 / 1000,
                tv_usec: (timeout as i32 % 1000) * 1000,
            };

            // In Windows, the first argument to `select` is ignored.
            let retval = unsafe { select::select(0, &mut readfds, null, null, &mut tv) };
            if retval == 0 {
                return Err(Error::new(ErrorKind::TimedOut, "Time limit expired"));
            } else if retval < 0 {
                return Err(Error::last_os_error());
            }
        }

        self.recv_from(buf)
    }
}

// Most of the following was copied from 'rust/src/libstd/sys/windows/net.rs'
#[cfg(windows)]
mod select {
    use libc;

    pub const FD_SETSIZE: usize = 64;

    #[repr(C)]
    pub struct fd_set {
        fd_count: libc::c_uint,
        fd_array: [libc::SOCKET; FD_SETSIZE],
    }

    pub fn fd_set(set: &mut fd_set, s: libc::SOCKET) {
        set.fd_array[set.fd_count as usize] = s;
        set.fd_count += 1;
    }

    impl fd_set {
        pub fn new() -> fd_set {
            fd_set {
                fd_count: 0,
                fd_array: [0; FD_SETSIZE],
            }
        }
    }

    #[link(name = "ws2_32")]
    extern "system" {
        pub fn select(nfds: libc::c_int,
                      readfds: *mut fd_set,
                      writefds: *mut fd_set,
                      exceptfds: *mut fd_set,
                      timeout: *mut libc::timeval) -> libc::c_int;
    }
}

#[test]
fn test_socket_timeout() {
    let mut socket = UdpSocket::bind("0.0.0.0:0").unwrap();
    let mut buf = [0; 10];
    assert!(socket.recv_timeout(&mut buf, 100).is_err());
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use error::{
    Result as MyResult,
};
use error::Error::DriverError;
use error::DriverError::MixedParams;

enum ParserState {
    TopLevel,
    // (string_delimiter, last_char)
    InStringLiteral(char, char),
    MaybeInNamedParam,
    InNamedParam,
}

use self::ParserState::*;

pub fn parse_named_params<'a>(query: &'a str) -> MyResult<(Option<Vec<String>>, Cow<'a, str>)> {
    let mut state = TopLevel;
    let mut have_positional = false;
    let mut cur_param = 0;
    // Vec<(start_offset, end_offset, name)>
    let mut params = Vec::new();
    for (i, c) in query.chars().enumerate() {
        let mut rematch = false;
        match state {
            TopLevel => {
                match c {
                    ':' => state = MaybeInNamedParam,
                    '\'' => state = InStringLiteral('\'', '\''),
                    '"' => state = InStringLiteral('"', '"'),
                    '?' => have_positional = true,
                    _ => (),
                }
            },
            InStringLiteral(separator, prev_char) => {
                match c {
                    x if x == separator && prev_char != '\\' => state = TopLevel,
                    x => state = InStringLiteral(separator, x),
                }
            },
            MaybeInNamedParam => {
                match c {
                    'a' ... 'z' | '_' => {
                        params.push((i - 1, 0, String::with_capacity(16)));
                        params[cur_param].2.push(c);
                        state = InNamedParam;
                    },
                    _ => rematch = true,
                }
            },
            InNamedParam => {
                match c {
                    'a' ... 'z' | '_' => params[cur_param].2.push(c),
                    _ => {
                        params[cur_param].1 = i;
                        cur_param += 1;
                        rematch = true;
                    }
                }
            },
        }
        if rematch {
            match c {
                ':' => state = MaybeInNamedParam,
                '\'' => state = InStringLiteral('\'', '\''),
                '"' => state = InStringLiteral('"', '"'),
                _ => state = TopLevel,
            }
        }
    }
    match state {
        InNamedParam => params[cur_param].1 = query.len(),
        _ => (),
    }
    if params.len() > 0 {
        if have_positional {
            return Err(DriverError(MixedParams));
        }
        let mut real_query = String::with_capacity(query.len());
        let mut last = 0;
        let mut out_params = Vec::with_capacity(params.len());
        for (start, end, name) in params.into_iter() {
            real_query.push_str(&query[last..start]);
            real_query.push('?');
            last = end;
            out_params.push(name);
        }
        real_query.push_str(&query[last..]);
        Ok((Some(out_params), real_query.into()))
    } else {
        Ok((None, query.into()))
    }
}

#[cfg(test)]
mod test {
    use error::Error::DriverError;
    use error::DriverError::MixedParams;
    use named_params::parse_named_params;

    #[test]
    fn should_parse_named_params() {
        let result = parse_named_params(":a :b").unwrap();
        assert_eq!((Some(vec!["a".to_string(), "b".into()]), "? ?".into()), result);

        let result = parse_named_params("SELECT (:a-10)").unwrap();
        assert_eq!((Some(vec!["a".to_string()]), "SELECT (?-10)".into()), result);

        let result = parse_named_params(r#"SELECT '"\':a' "'\"':c" :b"#).unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r#"SELECT '"\':a' "'\"':c" ?"#.into()), result);

        let result = parse_named_params(r":a_Aa:b").unwrap();
        assert_eq!((Some(vec!["a_".to_string(), "b".into()]), r"?Aa?".into()), result);

        let result = parse_named_params(r"::b").unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r":?".into()), result);

        let err = parse_named_params(r":a ?").unwrap_err();
        match err {
            DriverError(MixedParams) => (),
            _ => unreachable!(),
        }
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use test;
        use named_params::parse_named_params;

        #[bench]
        fn parse_ten_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r#"
                SELECT :one, :two, :three, :four, :five, :six, :seven, :eight, :nine, :ten
                "#);
                test::black_box(result);
            });
        }

        #[bench]
        fn parse_zero_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r"
                SELECT one, two, three, four, five, six, seven, eight, nine, ten
                ");
                test::black_box(result);
            });
        }
    }
}
<commit_msg>bench: use unused results<commit_after>use std::borrow::Cow;
use error::{
    Result as MyResult,
};
use error::Error::DriverError;
use error::DriverError::MixedParams;

enum ParserState {
    TopLevel,
    // (string_delimiter, last_char)
    InStringLiteral(char, char),
    MaybeInNamedParam,
    InNamedParam,
}

use self::ParserState::*;

pub fn parse_named_params<'a>(query: &'a str) -> MyResult<(Option<Vec<String>>, Cow<'a, str>)> {
    let mut state = TopLevel;
    let mut have_positional = false;
    let mut cur_param = 0;
    // Vec<(start_offset, end_offset, name)>
    let mut params = Vec::new();
    for (i, c) in query.chars().enumerate() {
        let mut rematch = false;
        match state {
            TopLevel => {
                match c {
                    ':' => state = MaybeInNamedParam,
                    '\'' => state = InStringLiteral('\'', '\''),
                    '"' => state = InStringLiteral('"', '"'),
                    '?' => have_positional = true,
                    _ => (),
                }
            },
            InStringLiteral(separator, prev_char) => {
                match c {
                    x if x == separator && prev_char != '\\' => state = TopLevel,
                    x => state = InStringLiteral(separator, x),
                }
            },
            MaybeInNamedParam => {
                match c {
                    'a' ... 'z' | '_' => {
                        params.push((i - 1, 0, String::with_capacity(16)));
                        params[cur_param].2.push(c);
                        state = InNamedParam;
                    },
                    _ => rematch = true,
                }
            },
            InNamedParam => {
                match c {
                    'a' ... 'z' | '_' => params[cur_param].2.push(c),
                    _ => {
                        params[cur_param].1 = i;
                        cur_param += 1;
                        rematch = true;
                    }
                }
            },
        }
        if rematch {
            match c {
                ':' => state = MaybeInNamedParam,
                '\'' => state = InStringLiteral('\'', '\''),
                '"' => state = InStringLiteral('"', '"'),
                _ => state = TopLevel,
            }
        }
    }
    match state {
        InNamedParam => params[cur_param].1 = query.len(),
        _ => (),
    }
    if params.len() > 0 {
        if have_positional {
            return Err(DriverError(MixedParams));
        }
        let mut real_query = String::with_capacity(query.len());
        let mut last = 0;
        let mut out_params = Vec::with_capacity(params.len());
        for (start, end, name) in params.into_iter() {
            real_query.push_str(&query[last..start]);
            real_query.push('?');
            last = end;
            out_params.push(name);
        }
        real_query.push_str(&query[last..]);
        Ok((Some(out_params), real_query.into()))
    } else {
        Ok((None, query.into()))
    }
}

#[cfg(test)]
mod test {
    use error::Error::DriverError;
    use error::DriverError::MixedParams;
    use named_params::parse_named_params;

    #[test]
    fn should_parse_named_params() {
        let result = parse_named_params(":a :b").unwrap();
        assert_eq!((Some(vec!["a".to_string(), "b".into()]), "? ?".into()), result);

        let result = parse_named_params("SELECT (:a-10)").unwrap();
        assert_eq!((Some(vec!["a".to_string()]), "SELECT (?-10)".into()), result);

        let result = parse_named_params(r#"SELECT '"\':a' "'\"':c" :b"#).unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r#"SELECT '"\':a' "'\"':c" ?"#.into()), result);

        let result = parse_named_params(r":a_Aa:b").unwrap();
        assert_eq!((Some(vec!["a_".to_string(), "b".into()]), r"?Aa?".into()), result);

        let result = parse_named_params(r"::b").unwrap();
        assert_eq!((Some(vec!["b".to_string()]), r":?".into()), result);

        let err = parse_named_params(r":a ?").unwrap_err();
        match err {
            DriverError(MixedParams) => (),
            _ => unreachable!(),
        }
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use test;
        use named_params::parse_named_params;

        #[bench]
        fn parse_ten_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r#"
                SELECT :one, :two, :three, :four, :five, :six, :seven, :eight, :nine, :ten
                "#).unwrap();
                test::black_box(result);
            });
        }

        #[bench]
        fn parse_zero_named_params(bencher: &mut test::Bencher) {
            bencher.iter(|| {
                let result = parse_named_params(r"
                SELECT one, two, three, four, five, six, seven, eight, nine, ten
                ").unwrap();
                test::black_box(result);
            });
        }
    }
}
<|endoftext|>


<commit_before>use std::{iter,str};
use std::fmt::Debug;

pub type Key = Vec<u8>;
pub type KeyValue<K,V> = (K,V);

const BRANCH_FACTOR: usize = 4;

#[derive(Debug,PartialEq)]
pub struct TrieNode<V> {
  partial_key: Key,
  key_value:   Option<KeyValue<Key,V>>,
  children:    Vec<TrieNode<V>>,
}

#[derive(Debug,PartialEq)]
pub enum InsertResult {
  Ok,
  Existing,
  Failed
}

#[derive(Debug,PartialEq)]
pub enum RemoveResult {
  Ok,
  NotFound,
}

impl<V:Debug> TrieNode<V> {
  pub fn new(partial: &[u8], key: Key, value: V) -> TrieNode<V> {
    TrieNode {
      partial_key:    Vec::from(partial),
      key_value:      Some((key, value)),
      children:       vec!(),
    }
  }

  pub fn root() -> TrieNode<V> {
    TrieNode {
      partial_key:    vec!(),
      key_value:      None,
      children:       vec!(),
    }
  }

  pub fn split(&mut self, index:usize) {
    let key_value = self.key_value.take();
    let children  = self.children.drain(..).collect();

    let child = TrieNode {
      partial_key: (&self.partial_key[index..]).to_vec(),
      key_value:   key_value,
      children:    children,
    };

    self.children.push(child);
    self.partial_key = (&self.partial_key[..index]).to_vec();
  }

  pub fn insert(&mut self, key: Key, value: V) -> InsertResult {
    println!("adding {}", str::from_utf8(&key).unwrap());
    let res = self.insert_recursive(&key, &key, value);
    assert_ne!(res, InsertResult::Failed);
    res
  }

  pub fn insert_recursive(&mut self, partial_key: &[u8], key: &Key, value: V) -> InsertResult {
    assert_ne!(partial_key, &b""[..]);

    println!("at level: {}, testing {}", str::from_utf8(&self.partial_key).unwrap(),
      str::from_utf8(partial_key).unwrap());

    // checking directly the children
    for (index, child) in self.children.iter_mut().enumerate() {
      let pos = partial_key.iter().zip(child.partial_key.iter()).position(|(&a,&b)| a != b);
      match pos {
        Some(0) => continue,
        Some(i) => {
          child.split(i);
          let new_child = TrieNode {
            partial_key: (&partial_key[i..]).to_vec(),
            key_value:   Some((key.clone(), value)),
            children:    vec!(),
          };

          child.children.push(new_child);
          return InsertResult::Ok;
        },
        None    => {
          if partial_key.len() > child.partial_key.len()  {
            let i = child.partial_key.len();
            return child.insert_recursive(&partial_key[i..], key, value);
          } else if partial_key.len() == child.partial_key.len() {
            if child.key_value.is_some() {
              return InsertResult::Existing;
            } else {
              child.key_value = Some((key.clone(), value));
              return InsertResult::Ok;
            }
          } else {
            // the partial key is smaller, insert as parent
            child.split(partial_key.len());
            child.key_value = Some((key.clone(), value));
            return InsertResult::Ok;
          }
        }
      };
    }

    let new_child = TrieNode {
      partial_key: partial_key.to_vec(),
      key_value:   Some((key.clone(), value)),
      children:    vec!(),
    };

    self.children.push(new_child);
    InsertResult::Ok
  }

  pub fn remove(&mut self, key: &Key) -> RemoveResult {
    self.remove_recursive(key)
  }

  pub fn remove_recursive(&mut self, partial_key: &[u8]) -> RemoveResult {
    assert_ne!(partial_key, &b""[..]);
    let mut found_child:Option<usize> = None;

    // checking directly the children
    for (index, child) in self.children.iter_mut().enumerate() {

      let pos = partial_key.iter().zip(child.partial_key.iter()).position(|(&a,&b)| a != b);
      match pos {
        Some(i) => continue,
        None    => {
          if partial_key.len() > child.partial_key.len()  {
            let i = child.partial_key.len();
            return child.remove_recursive(&partial_key[i..]);
          } else if partial_key.len() == child.partial_key.len() {
            found_child = Some(index);
            break;
          } else {
            continue
          }
        }
      };
    }

    if let Some(index) = found_child {
      if self.children[index].children.len() > 0 && self.children[index].key_value.is_some() {
        self.children[index].key_value = None;
        return RemoveResult::Ok;
      } else {
        self.children.remove(index);
        if self.key_value.is_some() {
          return RemoveResult::Ok;
        } else {
          //merging with the child
          if self.children.len() == 1 {
            let mut ch     = self.children.remove(0);
            self.key_value = ch.key_value.take();
            self.children  = ch.children;
            self.partial_key.extend(ch.partial_key);
          }
          // not handling the case of empty children vec
          // this case should only happen if it is the root node
          // otherwise, when we get to one last node and there is
          // no key_value, it is merged with the child node
          return RemoveResult::Ok
        }
      }
    } else {
      RemoveResult::NotFound
    }
  }

  pub fn lookup(&self, partial_key: &[u8]) -> Option<&KeyValue<Key,V>> {
    assert_ne!(partial_key, &b""[..]);

    if partial_key.starts_with(&self.partial_key) {
      if partial_key.len() == self.partial_key.len() {
        return self.key_value.as_ref();
      } else {
        for child in self.children.iter() {
          let res = child.lookup(&partial_key[self.partial_key.len()..]);
          if res.is_some() {
            return res
          }
        }
        None
      }
    } else {
      None
    }
  }

  // specific version that will handle wildcard domains
  pub fn domain_insert(&mut self, key: Key, value: V) -> InsertResult {
    let mut partial_key = key.clone();
    partial_key.reverse();
    self.insert_recursive(&partial_key, &key, value)
  }

  // specific version that will handle wildcard domains
  pub fn domain_remove(&mut self, key: &Key) -> RemoveResult {
    let mut partial_key = key.clone();
    partial_key.reverse();
    self.remove_recursive(&partial_key)
  }

  // specific version that will handle wildcard domains
  pub fn domain_lookup(&self, key: &[u8]) -> Option<&KeyValue<Key,V>> {
    let mut partial_key = key.to_vec();
    partial_key.reverse();
    self.domain_lookup_recursive(&partial_key)
  }

  // specific version that will handle wildcard domains
  pub fn domain_lookup_recursive(&self, partial_key: &[u8]) -> Option<&KeyValue<Key,V>> {
    assert_ne!(partial_key, &b""[..]);
    let pos = partial_key.iter().zip(self.partial_key.iter()).position(|(&a,&b)| a != b);

    match pos {
      Some(0) => return None,
      Some(i) => {
        // check for wildcard
        if i+1 == self.partial_key.len() && self.partial_key[i] == '*' as u8 {
          let c = '.' as u8;
          if (&partial_key[i..]).contains(&c) {
            None
          } else {
            self.key_value.as_ref()
          }
        } else {
          None
        }
      },
      None    => {
        if partial_key.len() > self.partial_key.len() {
          for child in self.children.iter() {
            let res = child.domain_lookup_recursive(&partial_key[self.partial_key.len()..]);
            if res.is_some() {
              return res
            }
          }
          None
        } else if partial_key.len() == self.partial_key.len() {
          self.key_value.as_ref()
        } else {
          None
        }

      }
    }
  }

  pub fn print(&self) {
    self.print_recursive(0)
  }

  pub fn print_recursive(&self, indent:u8) {
    let raw_prefix:Vec<u8> = iter::repeat(' ' as u8).take(2*indent as usize).collect();
    let prefix = str::from_utf8(&raw_prefix).unwrap();

    if let Some((ref key, ref value)) = self.key_value {
    println!("{}{}: ({},{:?})", prefix, str::from_utf8(&self.partial_key).unwrap(),
      str::from_utf8(&key).unwrap(), value);
    } else {
    println!("{}{}: None", prefix, str::from_utf8(&self.partial_key).unwrap());
    }
    for child in self.children.iter() {
      child.print_recursive(indent+1);
    }
  }
}

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

  #[test]
  fn insert() {
    let mut root: TrieNode<u8> = TrieNode::root();
    root.print();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    root.print();
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    root.print();
    assert_eq!(root.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);
    root.print();

    assert_eq!(root.lookup(&b"abce"[..]), Some(&((&b"abce"[..]).to_vec(), 2)));
    //assert!(false);
  }

  #[test]
  fn remove() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);

    let mut root2: TrieNode<u8> = TrieNode::root();

    assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);

    println!("before remove");
    root.print();
    assert_eq!(root.remove(&Vec::from(&b"abce"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    println!("expected");
    root2.print();
    assert_eq!(root, root2);

    assert_eq!(root.remove(&Vec::from(&b"abgh"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    println!("expected");
    let mut root3: TrieNode<u8> = TrieNode::root();
    assert_eq!(root3.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    root3.print();
    assert_eq!(root, root3);
  }

  #[test]
  fn add_child_to_leaf() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);

    root.print();

    let mut root2: TrieNode<u8> = TrieNode::root();

    assert_eq!(root2.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);

    root2.print();
    assert_eq!(root2.remove(&Vec::from(&b"abc"[..])), RemoveResult::Ok);

    let mut expected: TrieNode<u8> = TrieNode::root();

    assert_eq!(expected.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(expected.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);

    println!("after remove");
    root2.print();
    println!("expected");
    expected.print();
    assert_eq!(root2, expected);

  }

  #[test]
  fn domains() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.domain_insert(Vec::from(&b"www.example.com"[..]), 1), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"test.example.com"[..]), 2), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"*.alldomains.org"[..]), 3), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"alldomains.org"[..]), 4), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"hello.com"[..]), 5), InsertResult::Ok);
    root.print();

    assert_eq!(root.domain_lookup(&b"example.com"[..]), None);
    assert_eq!(root.domain_lookup(&b"blah.test.example.com"[..]), None);
    assert_eq!(root.domain_lookup(&b"www.example.com"[..]), Some(&((&b"www.example.com"[..]).to_vec(), 1)));
    assert_eq!(root.domain_lookup(&b"alldomains.org"[..]), Some(&((&b"alldomains.org"[..]).to_vec(), 4)));
    assert_eq!(root.domain_lookup(&b"test.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"hello.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"blah.test.alldomains.org"[..]), None);

    assert_eq!(root.domain_remove(&Vec::from(&b"alldomains.org"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    assert_eq!(root.domain_lookup(&b"alldomains.org"[..]), None);
    assert_eq!(root.domain_lookup(&b"test.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"hello.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"blah.test.alldomains.org"[..]), None);
  }
}
<commit_msg>fix wildcard domain lookup<commit_after>use std::{iter,str};
use std::fmt::Debug;

pub type Key = Vec<u8>;
pub type KeyValue<K,V> = (K,V);

const BRANCH_FACTOR: usize = 4;

#[derive(Debug,PartialEq)]
pub struct TrieNode<V> {
  partial_key: Key,
  key_value:   Option<KeyValue<Key,V>>,
  children:    Vec<TrieNode<V>>,
}

#[derive(Debug,PartialEq)]
pub enum InsertResult {
  Ok,
  Existing,
  Failed
}

#[derive(Debug,PartialEq)]
pub enum RemoveResult {
  Ok,
  NotFound,
}

impl<V:Debug> TrieNode<V> {
  pub fn new(partial: &[u8], key: Key, value: V) -> TrieNode<V> {
    TrieNode {
      partial_key:    Vec::from(partial),
      key_value:      Some((key, value)),
      children:       vec!(),
    }
  }

  pub fn root() -> TrieNode<V> {
    TrieNode {
      partial_key:    vec!(),
      key_value:      None,
      children:       vec!(),
    }
  }

  pub fn split(&mut self, index:usize) {
    let key_value = self.key_value.take();
    let children  = self.children.drain(..).collect();

    let child = TrieNode {
      partial_key: (&self.partial_key[index..]).to_vec(),
      key_value:   key_value,
      children:    children,
    };

    self.children.push(child);
    self.partial_key = (&self.partial_key[..index]).to_vec();
  }

  pub fn insert(&mut self, key: Key, value: V) -> InsertResult {
    println!("adding {}", str::from_utf8(&key).unwrap());
    let res = self.insert_recursive(&key, &key, value);
    assert_ne!(res, InsertResult::Failed);
    res
  }

  pub fn insert_recursive(&mut self, partial_key: &[u8], key: &Key, value: V) -> InsertResult {
    assert_ne!(partial_key, &b""[..]);

    // checking directly the children
    for (index, child) in self.children.iter_mut().enumerate() {
      let pos = partial_key.iter().zip(child.partial_key.iter()).position(|(&a,&b)| a != b);
      match pos {
        Some(0) => continue,
        Some(i) => {
          child.split(i);
          let new_child = TrieNode {
            partial_key: (&partial_key[i..]).to_vec(),
            key_value:   Some((key.clone(), value)),
            children:    vec!(),
          };

          child.children.push(new_child);
          return InsertResult::Ok;
        },
        None    => {
          if partial_key.len() > child.partial_key.len()  {
            let i = child.partial_key.len();
            return child.insert_recursive(&partial_key[i..], key, value);
          } else if partial_key.len() == child.partial_key.len() {
            if child.key_value.is_some() {
              return InsertResult::Existing;
            } else {
              child.key_value = Some((key.clone(), value));
              return InsertResult::Ok;
            }
          } else {
            // the partial key is smaller, insert as parent
            child.split(partial_key.len());
            child.key_value = Some((key.clone(), value));
            return InsertResult::Ok;
          }
        }
      };
    }

    let new_child = TrieNode {
      partial_key: partial_key.to_vec(),
      key_value:   Some((key.clone(), value)),
      children:    vec!(),
    };

    self.children.push(new_child);
    InsertResult::Ok
  }

  pub fn remove(&mut self, key: &Key) -> RemoveResult {
    self.remove_recursive(key)
  }

  pub fn remove_recursive(&mut self, partial_key: &[u8]) -> RemoveResult {
    assert_ne!(partial_key, &b""[..]);
    let mut found_child:Option<usize> = None;

    // checking directly the children
    for (index, child) in self.children.iter_mut().enumerate() {

      let pos = partial_key.iter().zip(child.partial_key.iter()).position(|(&a,&b)| a != b);
      match pos {
        Some(i) => continue,
        None    => {
          if partial_key.len() > child.partial_key.len()  {
            let i = child.partial_key.len();
            return child.remove_recursive(&partial_key[i..]);
          } else if partial_key.len() == child.partial_key.len() {
            found_child = Some(index);
            break;
          } else {
            continue
          }
        }
      };
    }

    if let Some(index) = found_child {
      if self.children[index].children.len() > 0 && self.children[index].key_value.is_some() {
        self.children[index].key_value = None;
        return RemoveResult::Ok;
      } else {
        self.children.remove(index);
        if self.key_value.is_some() {
          return RemoveResult::Ok;
        } else {
          //merging with the child
          if self.children.len() == 1 {
            let mut ch     = self.children.remove(0);
            self.key_value = ch.key_value.take();
            self.children  = ch.children;
            self.partial_key.extend(ch.partial_key);
          }
          // not handling the case of empty children vec
          // this case should only happen if it is the root node
          // otherwise, when we get to one last node and there is
          // no key_value, it is merged with the child node
          return RemoveResult::Ok
        }
      }
    } else {
      RemoveResult::NotFound
    }
  }

  pub fn lookup(&self, partial_key: &[u8]) -> Option<&KeyValue<Key,V>> {
    assert_ne!(partial_key, &b""[..]);

    if partial_key.starts_with(&self.partial_key) {
      if partial_key.len() == self.partial_key.len() {
        return self.key_value.as_ref();
      } else {
        for child in self.children.iter() {
          let res = child.lookup(&partial_key[self.partial_key.len()..]);
          if res.is_some() {
            return res
          }
        }
        None
      }
    } else {
      None
    }
  }

  // specific version that will handle wildcard domains
  pub fn domain_insert(&mut self, key: Key, value: V) -> InsertResult {
    let mut partial_key = key.clone();
    partial_key.reverse();
    self.insert_recursive(&partial_key, &key, value)
  }

  // specific version that will handle wildcard domains
  pub fn domain_remove(&mut self, key: &Key) -> RemoveResult {
    let mut partial_key = key.clone();
    partial_key.reverse();
    self.remove_recursive(&partial_key)
  }

  // specific version that will handle wildcard domains
  pub fn domain_lookup(&self, key: &[u8]) -> Option<&KeyValue<Key,V>> {
    let mut partial_key = key.to_vec();
    partial_key.reverse();
    self.domain_lookup_recursive(&partial_key)
  }

  // specific version that will handle wildcard domains
  pub fn domain_lookup_recursive(&self, partial_key: &[u8]) -> Option<&KeyValue<Key,V>> {
    assert_ne!(partial_key, &b""[..]);
    let pos = partial_key.iter().zip(self.partial_key.iter()).position(|(&a,&b)| a != b);
    //info!("lookup at level: {}, testing {}", str::from_utf8(&self.partial_key).unwrap(),
    //  str::from_utf8(partial_key).unwrap());

    match pos {
      Some(i) => {
        // check for wildcard
        if i+1 == self.partial_key.len() && self.partial_key[i] == '*' as u8 {
          let c = '.' as u8;
          if (&partial_key[i..]).contains(&c) {
            None
          } else {
            self.key_value.as_ref()
          }
        } else {
          None
        }
      },
      None    => {
        if partial_key.len() > self.partial_key.len() {
          for child in self.children.iter() {
            let res = child.domain_lookup_recursive(&partial_key[self.partial_key.len()..]);
            if res.is_some() {
              return res
            }
          }
          None
        } else if partial_key.len() == self.partial_key.len() {
          self.key_value.as_ref()
        } else {
          None
        }

      }
    }
  }

  pub fn print(&self) {
    self.print_recursive(0)
  }

  pub fn print_recursive(&self, indent:u8) {
    let raw_prefix:Vec<u8> = iter::repeat(' ' as u8).take(2*indent as usize).collect();
    let prefix = str::from_utf8(&raw_prefix).unwrap();

    if let Some((ref key, ref value)) = self.key_value {
    println!("{}{}: ({},{:?})", prefix, str::from_utf8(&self.partial_key).unwrap(),
      str::from_utf8(&key).unwrap(), value);
    } else {
    println!("{}{}: None", prefix, str::from_utf8(&self.partial_key).unwrap());
    }
    for child in self.children.iter() {
      child.print_recursive(indent+1);
    }
  }
}

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

  #[test]
  fn insert() {
    let mut root: TrieNode<u8> = TrieNode::root();
    root.print();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    root.print();
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    root.print();
    assert_eq!(root.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);
    root.print();

    assert_eq!(root.lookup(&b"abce"[..]), Some(&((&b"abce"[..]).to_vec(), 2)));
    //assert!(false);
  }

  #[test]
  fn remove() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);

    let mut root2: TrieNode<u8> = TrieNode::root();

    assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);

    println!("before remove");
    root.print();
    assert_eq!(root.remove(&Vec::from(&b"abce"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    println!("expected");
    root2.print();
    assert_eq!(root, root2);

    assert_eq!(root.remove(&Vec::from(&b"abgh"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    println!("expected");
    let mut root3: TrieNode<u8> = TrieNode::root();
    assert_eq!(root3.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    root3.print();
    assert_eq!(root, root3);
  }

  #[test]
  fn add_child_to_leaf() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
    assert_eq!(root.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);

    root.print();

    let mut root2: TrieNode<u8> = TrieNode::root();

    assert_eq!(root2.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(root2.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);

    root2.print();
    assert_eq!(root2.remove(&Vec::from(&b"abc"[..])), RemoveResult::Ok);

    let mut expected: TrieNode<u8> = TrieNode::root();

    assert_eq!(expected.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
    assert_eq!(expected.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);

    println!("after remove");
    root2.print();
    println!("expected");
    expected.print();
    assert_eq!(root2, expected);

  }

  #[test]
  fn domains() {
    let mut root: TrieNode<u8> = TrieNode::root();

    assert_eq!(root.domain_insert(Vec::from(&b"www.example.com"[..]), 1), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"test.example.com"[..]), 2), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"*.alldomains.org"[..]), 3), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"alldomains.org"[..]), 4), InsertResult::Ok);
    root.print();
    assert_eq!(root.domain_insert(Vec::from(&b"hello.com"[..]), 5), InsertResult::Ok);
    root.print();

    assert_eq!(root.domain_lookup(&b"example.com"[..]), None);
    assert_eq!(root.domain_lookup(&b"blah.test.example.com"[..]), None);
    assert_eq!(root.domain_lookup(&b"www.example.com"[..]), Some(&((&b"www.example.com"[..]).to_vec(), 1)));
    assert_eq!(root.domain_lookup(&b"alldomains.org"[..]), Some(&((&b"alldomains.org"[..]).to_vec(), 4)));
    assert_eq!(root.domain_lookup(&b"test.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"hello.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"blah.test.alldomains.org"[..]), None);

    assert_eq!(root.domain_remove(&Vec::from(&b"alldomains.org"[..])), RemoveResult::Ok);
    println!("after remove");
    root.print();
    assert_eq!(root.domain_lookup(&b"alldomains.org"[..]), None);
    assert_eq!(root.domain_lookup(&b"test.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"hello.alldomains.org"[..]), Some(&((&b"*.alldomains.org"[..]).to_vec(), 3)));
    assert_eq!(root.domain_lookup(&b"blah.test.alldomains.org"[..]), None);
  }
}
<|endoftext|>


<commit_before>use std::rc::Rc;

use gtk::prelude::*;
use gtk::{Window, FileChooser, Button, Entry, ComboBoxText};

use cargo::util::config::Config;
use cargo::ops::{self, VersionControl, NewOptions};

use util::PreContext;
use pages::Page;

/// An interface for creating packages.
#[derive(Clone)]
pub struct NewPackagePage {
    pub window: Window,
    pub config: Rc<Config>,
    pub file: FileChooser,
    pub new: Button,
    pub name: Entry,
    pub ty: ComboBoxText,
    pub vcs: ComboBoxText,
}
impl Page for NewPackagePage {
    fn new(c: &PreContext) -> NewPackagePage {
        let b = &c.builder;
        NewPackagePage {
            window: c.window.clone(),
            config: c.config.clone(),
            file: b.get_object("package-file").unwrap(),
            name: b.get_object("package-name").unwrap(),
            vcs: b.get_object("package-vcs").unwrap(),
            ty: b.get_object("package-type").unwrap(),
            new: b.get_object("package-new").unwrap()
        }
    }
    fn bind_listeners(&self) {
        let name = self.name.clone();
        let ty = self.ty.clone();
        let vcs = self.vcs.clone();
        let window = self.window.clone();
        let config = self.config.clone();
        ::bind_file_button(&self.file, &self.new, move |file| {
            let text = name.get_text();
            let opts = NewOptions {
                path: file.to_str().unwrap(),
                name: text.as_ref().map(String::as_str),
                bin: ty.get_active_id().as_ref().map(String::as_str) == Some("bin"),
                version_control: vcs.get_active_id().as_ref().map(String::as_str).map(|v| match v {
                    "git" => VersionControl::Git,
                    "mercurial" => VersionControl::Hg,
                    _ => VersionControl::NoVcs
                })
            };
            if ops::new(opts, &config).is_ok() {
                ::info(Some(&window), "Created crate successfully");
            } else {
                ::error(Some(&window), "Failed to create crate");
            }
        });
    }
}
<commit_msg>Revert "Use `ops::new` instead of `ops::init`"<commit_after>use std::rc::Rc;

use gtk::prelude::*;
use gtk::{Window, FileChooser, Button, Entry, ComboBoxText};

use cargo::util::config::Config;
use cargo::ops::{self, VersionControl, NewOptions};

use util::PreContext;
use pages::Page;

/// An interface for creating packages.
#[derive(Clone)]
pub struct NewPackagePage {
    pub window: Window,
    pub config: Rc<Config>,
    pub file: FileChooser,
    pub new: Button,
    pub name: Entry,
    pub ty: ComboBoxText,
    pub vcs: ComboBoxText,
}
impl Page for NewPackagePage {
    fn new(c: &PreContext) -> NewPackagePage {
        let b = &c.builder;
        NewPackagePage {
            window: c.window.clone(),
            config: c.config.clone(),
            file: b.get_object("package-file").unwrap(),
            name: b.get_object("package-name").unwrap(),
            vcs: b.get_object("package-vcs").unwrap(),
            ty: b.get_object("package-type").unwrap(),
            new: b.get_object("package-new").unwrap()
        }
    }
    fn bind_listeners(&self) {
        let name = self.name.clone();
        let ty = self.ty.clone();
        let vcs = self.vcs.clone();
        let window = self.window.clone();
        let config = self.config.clone();
        ::bind_file_button(&self.file, &self.new, move |file| {
            let text = name.get_text();
            let opts = NewOptions {
                path: file.to_str().unwrap(),
                name: text.as_ref().map(String::as_str),
                bin: ty.get_active_id().as_ref().map(String::as_str) == Some("bin"),
                version_control: vcs.get_active_id().as_ref().map(String::as_str).map(|v| match v {
                    "git" => VersionControl::Git,
                    "mercurial" => VersionControl::Hg,
                    _ => VersionControl::NoVcs
                })
            };
            if ops::init(opts, &config).is_ok() {
                ::info(Some(&window), "Created crate successfully");
            } else {
                ::error(Some(&window), "Failed to create crate");
            }
        });
    }
}
<|endoftext|>