<commit_before>use auto_enums::auto_enum;
use glob::{glob_with, MatchOptions};
use ion_shell::{
    expansion::{unescape, Expander},
    Shell,
};
use liner::{Completer, CursorPosition, Event, EventKind};
use std::{env, iter, path::PathBuf, str};

pub struct IonCompleter<'a, 'b> {
    shell:      &'b Shell<'a>,
    completion: CompletionType,
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' | b' ' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

enum CompletionType {
    Nothing,
    Command,
    VariableAndFiles,
}

impl<'a, 'b> IonCompleter<'a, 'b> {
    pub fn new(shell: &'b Shell<'a>) -> Self {
        IonCompleter { shell, completion: CompletionType::Nothing }
    }
}

impl<'a, 'b> Completer for IonCompleter<'a, 'b> {
    fn completions(&mut self, start: &str) -> Vec<String> {
        let mut completions = IonFileCompleter::new(None, &self.shell, false).completions(start);
        let vars = self.shell.variables();

        match self.completion {
            CompletionType::VariableAndFiles => {
                // Initialize a new completer from the definitions collected.
                // Creates a list of definitions from the shell environment that
                // will be used
                // in the creation of a custom completer.
                if start.is_empty() {
                    completions.extend(vars.string_vars().map(|(s, _)| format!("${}", s)));
                    completions.extend(vars.arrays().map(|(s, _)| format!("@{}", s)));
                } else if start.starts_with('$') {
                    completions.extend(
                        // Add the list of available variables to the completer's
                        // definitions. TODO: We should make
                        // it free to do String->SmallString
                        //       and mostly free to go back (free if allocated)
                        vars.string_vars()
                            .filter(|(s, _)| s.starts_with(&start[1..]))
                            .map(|(s, _)| format!("${}", &s)),
                    );
                } else if start.starts_with('@') {
                    completions.extend(
                        vars.arrays()
                            .filter(|(s, _)| s.starts_with(&start[1..]))
                            .map(|(s, _)| format!("@{}", &s)),
                    );
                }
            }
            CompletionType::Command => {
                // Initialize a new completer from the definitions collected.
                // Creates a list of definitions from the shell environment that
                // will be used
                // in the creation of a custom completer.
                completions.extend(
                    self.shell
                        .builtins()
                        .keys()
                        // Add built-in commands to the completer's definitions.
                        .map(ToString::to_string)
                        // Add the aliases to the completer's definitions.
                        .chain(vars.aliases().map(|(key, _)| key.to_string()))
                        // Add the list of available functions to the completer's
                        // definitions.
                        .chain(vars.functions().map(|(key, _)| key.to_string()))
                        .filter(|s| s.starts_with(start)),
                );
                // Creates completers containing definitions from all directories
                // listed
                // in the environment's **$PATH** variable.
                let file_completers: Vec<_> = if let Some(paths) = env::var_os("PATH") {
                    env::split_paths(&paths)
                        .map(|s| {
                            let s = if !s.to_string_lossy().ends_with("/") {
                                PathBuf::from(format!("{}/", s.to_string_lossy()))
                            } else {
                                s
                            };
                            IonFileCompleter::new(Some(s), &self.shell, true)
                        })
                        .collect()
                } else {
                    vec![IonFileCompleter::new(Some("/bin/".into()), &self.shell, true)]
                };
                // Merge the collected definitions with the file path definitions.
                completions.extend(MultiCompleter::new(file_completers).completions(start));
            }
            CompletionType::Nothing => (),
        }

        completions
    }

    fn on_event<W: std::io::Write>(&mut self, event: Event<'_, '_, W>) {
        if let EventKind::BeforeComplete = event.kind {
            let (words, pos) = event.editor.get_words_and_cursor_position();
            self.completion = match pos {
                _ if words.is_empty() => CompletionType::Nothing,
                CursorPosition::InWord(0) => CompletionType::Command,
                CursorPosition::OnWordRightEdge(index) => {
                    if index == 0 {
                        CompletionType::Command
                    } else {
                        let is_pipe = words
                            .into_iter()
                            .nth(index - 1)
                            .map(|(start, end)| event.editor.current_buffer().range(start, end))
                            .filter(|filename| {
                                filename.ends_with("|")
                                    || filename.ends_with("&")
                                    || filename.ends_with(";")
                            })
                            .is_some();
                        if is_pipe {
                            CompletionType::Command
                        } else {
                            CompletionType::VariableAndFiles
                        }
                    }
                }
                _ => CompletionType::VariableAndFiles,
            };
        }
    }
}

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub struct IonFileCompleter<'a, 'b> {
    shell: &'b Shell<'a>,
    /// The directory the expansion takes place in
    path: PathBuf,
    for_command: bool,
}

impl<'a, 'b> IonFileCompleter<'a, 'b> {
    pub fn new(path: Option<PathBuf>, shell: &'b Shell<'a>, for_command: bool) -> Self {
        let path = path.unwrap_or_default();
        IonFileCompleter { shell, path, for_command }
    }
}

impl<'a, 'b> Completer for IonFileCompleter<'a, 'b> {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&mut self, start: &str) -> Vec<String> {
        // Dereferencing the raw pointers here should be entirely safe, theoretically,
        // because no changes will occur to either of the underlying references in the
        // duration between creation of the completers and execution of their
        // completions.
        let expanded = match self.shell.tilde(start) {
            Ok(expanded) => expanded,
            Err(why) => {
                eprintln!("ion: {}", why);
                return vec![start.into()];
            }
        };
        // Now we obtain completions for the `expanded` form of the `start` value.
        let completions = filename_completion(&expanded, &self.path);
        if expanded == start {
            return if self.for_command {
                completions.map(|s| s.rsplit('/').nth(0).unwrap_or(&s).to_string()).collect()
            } else {
                completions.collect()
            };
        }
        // We can do that by obtaining the index position where the tilde character
        // ends. We don't search with `~` because we also want to
        // handle other tilde variants.
        let t_index = start.find('/').unwrap_or(1);
        // `tilde` is the tilde pattern, and `search` is the pattern that follows.
        let (tilde, search) = start.split_at(t_index);

        if search.len() < 2 {
            // If the length of the search pattern is less than 2, the search pattern is
            // empty, and thus the completions actually contain files and directories in
            // the home directory.

            // The tilde pattern will actually be our `start` command in itself,
            // and the completed form will be all of the characters beyond the length of
            // the expanded form of the tilde pattern.
            completions.map(|completion| [start, &completion[expanded.len()..]].concat()).collect()
        // To save processing time, we should get obtain the index position where our
        // search pattern begins, and re-use that index to slice the completions so
        // that we may re-add the tilde character with the completion that follows.
        } else if let Some(e_index) = expanded.rfind(search) {
            // And then we will need to take those completions and remove the expanded form
            // of the tilde pattern and replace it with that pattern yet again.
            completions
                .map(|completion| escape(&[tilde, &completion[e_index..]].concat()))
                .collect()
        } else {
            Vec::new()
        }
    }
}

#[auto_enum]
fn filename_completion<'a>(start: &'a str, path: &'a PathBuf) -> impl Iterator<Item = String> + 'a {
    let unescaped_start = unescape(start);

    let mut split_start = unescaped_start.split('/');
    let mut string = String::with_capacity(128);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    if unescaped_start.starts_with('/') {
        split_start.next();
        string.push('/');
    } else {
        string.push_str(&path.to_string_lossy());
    }

    for element in split_start {
        string.push_str(element);
        if element != "." && element != ".." {
            string.push('*');
        }
        string.push('/');
    }

    string.pop(); // pop out the last '/' character
    if string.ends_with('.') {
        string.push('*')
    }
    let globs = glob_with(
        &string,
        MatchOptions {
            case_sensitive:              true,
            require_literal_separator:   true,
            require_literal_leading_dot: false,
        },
    )
    .ok()
    .map(|completions| {
        completions.filter_map(Result::ok).filter_map(move |file| {
            let out = file.to_str()?;
            let mut joined = String::with_capacity(out.len() + 3); // worst case senario
            if unescaped_start.starts_with("./") {
                joined.push_str("./");
            }
            joined.push_str(out);
            if file.is_dir() {
                joined.push('/');
            }
            Some(escape(&joined))
        })
    });

    #[auto_enum(Iterator)]
    match globs {
        Some(iter) => iter,
        None => iter::once(escape(start)),
    }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MultiCompleter<A>(Vec<A>);

impl<A> MultiCompleter<A> {
    pub fn new(completions: Vec<A>) -> Self { MultiCompleter(completions) }
}

impl<A> Completer for MultiCompleter<A>
where
    A: Completer,
{
    fn completions(&mut self, start: &str) -> Vec<String> {
        self.0.iter_mut().flat_map(|comp| comp.completions(start)).collect()
    }
}

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

    #[test]
    fn filename_completion() {
        let shell = Shell::default();
        let mut completer = IonFileCompleter::new(None, &shell);
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(completer.completions("testing/file"), vec!["testing/file_with_text"]);
        assert_eq!(completer.completions("~"), vec!["~/"]);
        assert_eq!(completer.completions("tes/fil"), vec!["testing/file_with_text"]);
    }
}
<commit_msg>Fix completer test.<commit_after>use auto_enums::auto_enum;
use glob::{glob_with, MatchOptions};
use ion_shell::{
    expansion::{unescape, Expander},
    Shell,
};
use liner::{Completer, CursorPosition, Event, EventKind};
use std::{env, iter, path::PathBuf, str};

pub struct IonCompleter<'a, 'b> {
    shell:      &'b Shell<'a>,
    completion: CompletionType,
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' | b' ' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

enum CompletionType {
    Nothing,
    Command,
    VariableAndFiles,
}

impl<'a, 'b> IonCompleter<'a, 'b> {
    pub fn new(shell: &'b Shell<'a>) -> Self {
        IonCompleter { shell, completion: CompletionType::Nothing }
    }
}

impl<'a, 'b> Completer for IonCompleter<'a, 'b> {
    fn completions(&mut self, start: &str) -> Vec<String> {
        let mut completions = IonFileCompleter::new(None, &self.shell, false).completions(start);
        let vars = self.shell.variables();

        match self.completion {
            CompletionType::VariableAndFiles => {
                // Initialize a new completer from the definitions collected.
                // Creates a list of definitions from the shell environment that
                // will be used
                // in the creation of a custom completer.
                if start.is_empty() {
                    completions.extend(vars.string_vars().map(|(s, _)| format!("${}", s)));
                    completions.extend(vars.arrays().map(|(s, _)| format!("@{}", s)));
                } else if start.starts_with('$') {
                    completions.extend(
                        // Add the list of available variables to the completer's
                        // definitions. TODO: We should make
                        // it free to do String->SmallString
                        //       and mostly free to go back (free if allocated)
                        vars.string_vars()
                            .filter(|(s, _)| s.starts_with(&start[1..]))
                            .map(|(s, _)| format!("${}", &s)),
                    );
                } else if start.starts_with('@') {
                    completions.extend(
                        vars.arrays()
                            .filter(|(s, _)| s.starts_with(&start[1..]))
                            .map(|(s, _)| format!("@{}", &s)),
                    );
                }
            }
            CompletionType::Command => {
                // Initialize a new completer from the definitions collected.
                // Creates a list of definitions from the shell environment that
                // will be used
                // in the creation of a custom completer.
                completions.extend(
                    self.shell
                        .builtins()
                        .keys()
                        // Add built-in commands to the completer's definitions.
                        .map(ToString::to_string)
                        // Add the aliases to the completer's definitions.
                        .chain(vars.aliases().map(|(key, _)| key.to_string()))
                        // Add the list of available functions to the completer's
                        // definitions.
                        .chain(vars.functions().map(|(key, _)| key.to_string()))
                        .filter(|s| s.starts_with(start)),
                );
                // Creates completers containing definitions from all directories
                // listed
                // in the environment's **$PATH** variable.
                let file_completers: Vec<_> = if let Some(paths) = env::var_os("PATH") {
                    env::split_paths(&paths)
                        .map(|s| {
                            let s = if !s.to_string_lossy().ends_with("/") {
                                PathBuf::from(format!("{}/", s.to_string_lossy()))
                            } else {
                                s
                            };
                            IonFileCompleter::new(Some(s), &self.shell, true)
                        })
                        .collect()
                } else {
                    vec![IonFileCompleter::new(Some("/bin/".into()), &self.shell, true)]
                };
                // Merge the collected definitions with the file path definitions.
                completions.extend(MultiCompleter::new(file_completers).completions(start));
            }
            CompletionType::Nothing => (),
        }

        completions
    }

    fn on_event<W: std::io::Write>(&mut self, event: Event<'_, '_, W>) {
        if let EventKind::BeforeComplete = event.kind {
            let (words, pos) = event.editor.get_words_and_cursor_position();
            self.completion = match pos {
                _ if words.is_empty() => CompletionType::Nothing,
                CursorPosition::InWord(0) => CompletionType::Command,
                CursorPosition::OnWordRightEdge(index) => {
                    if index == 0 {
                        CompletionType::Command
                    } else {
                        let is_pipe = words
                            .into_iter()
                            .nth(index - 1)
                            .map(|(start, end)| event.editor.current_buffer().range(start, end))
                            .filter(|filename| {
                                filename.ends_with("|")
                                    || filename.ends_with("&")
                                    || filename.ends_with(";")
                            })
                            .is_some();
                        if is_pipe {
                            CompletionType::Command
                        } else {
                            CompletionType::VariableAndFiles
                        }
                    }
                }
                _ => CompletionType::VariableAndFiles,
            };
        }
    }
}

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub struct IonFileCompleter<'a, 'b> {
    shell: &'b Shell<'a>,
    /// The directory the expansion takes place in
    path: PathBuf,
    for_command: bool,
}

impl<'a, 'b> IonFileCompleter<'a, 'b> {
    pub fn new(path: Option<PathBuf>, shell: &'b Shell<'a>, for_command: bool) -> Self {
        let path = path.unwrap_or_default();
        IonFileCompleter { shell, path, for_command }
    }
}

impl<'a, 'b> Completer for IonFileCompleter<'a, 'b> {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&mut self, start: &str) -> Vec<String> {
        // Dereferencing the raw pointers here should be entirely safe, theoretically,
        // because no changes will occur to either of the underlying references in the
        // duration between creation of the completers and execution of their
        // completions.
        let expanded = match self.shell.tilde(start) {
            Ok(expanded) => expanded,
            Err(why) => {
                eprintln!("ion: {}", why);
                return vec![start.into()];
            }
        };
        // Now we obtain completions for the `expanded` form of the `start` value.
        let completions = filename_completion(&expanded, &self.path);
        if expanded == start {
            return if self.for_command {
                completions.map(|s| s.rsplit('/').nth(0).unwrap_or(&s).to_string()).collect()
            } else {
                completions.collect()
            };
        }
        // We can do that by obtaining the index position where the tilde character
        // ends. We don't search with `~` because we also want to
        // handle other tilde variants.
        let t_index = start.find('/').unwrap_or(1);
        // `tilde` is the tilde pattern, and `search` is the pattern that follows.
        let (tilde, search) = start.split_at(t_index);

        if search.len() < 2 {
            // If the length of the search pattern is less than 2, the search pattern is
            // empty, and thus the completions actually contain files and directories in
            // the home directory.

            // The tilde pattern will actually be our `start` command in itself,
            // and the completed form will be all of the characters beyond the length of
            // the expanded form of the tilde pattern.
            completions.map(|completion| [start, &completion[expanded.len()..]].concat()).collect()
        // To save processing time, we should get obtain the index position where our
        // search pattern begins, and re-use that index to slice the completions so
        // that we may re-add the tilde character with the completion that follows.
        } else if let Some(e_index) = expanded.rfind(search) {
            // And then we will need to take those completions and remove the expanded form
            // of the tilde pattern and replace it with that pattern yet again.
            completions
                .map(|completion| escape(&[tilde, &completion[e_index..]].concat()))
                .collect()
        } else {
            Vec::new()
        }
    }
}

#[auto_enum]
fn filename_completion<'a>(start: &'a str, path: &'a PathBuf) -> impl Iterator<Item = String> + 'a {
    let unescaped_start = unescape(start);

    let mut split_start = unescaped_start.split('/');
    let mut string = String::with_capacity(128);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    if unescaped_start.starts_with('/') {
        split_start.next();
        string.push('/');
    } else {
        string.push_str(&path.to_string_lossy());
    }

    for element in split_start {
        string.push_str(element);
        if element != "." && element != ".." {
            string.push('*');
        }
        string.push('/');
    }

    string.pop(); // pop out the last '/' character
    if string.ends_with('.') {
        string.push('*')
    }
    let globs = glob_with(
        &string,
        MatchOptions {
            case_sensitive:              true,
            require_literal_separator:   true,
            require_literal_leading_dot: false,
        },
    )
    .ok()
    .map(|completions| {
        completions.filter_map(Result::ok).filter_map(move |file| {
            let out = file.to_str()?;
            let mut joined = String::with_capacity(out.len() + 3); // worst case senario
            if unescaped_start.starts_with("./") {
                joined.push_str("./");
            }
            joined.push_str(out);
            if file.is_dir() {
                joined.push('/');
            }
            Some(escape(&joined))
        })
    });

    #[auto_enum(Iterator)]
    match globs {
        Some(iter) => iter,
        None => iter::once(escape(start)),
    }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MultiCompleter<A>(Vec<A>);

impl<A> MultiCompleter<A> {
    pub fn new(completions: Vec<A>) -> Self { MultiCompleter(completions) }
}

impl<A> Completer for MultiCompleter<A>
where
    A: Completer,
{
    fn completions(&mut self, start: &str) -> Vec<String> {
        self.0.iter_mut().flat_map(|comp| comp.completions(start)).collect()
    }
}

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

    #[test]
    fn filename_completion() {
        let shell = Shell::default();
        let mut completer = IonFileCompleter::new(None, &shell, false);
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(completer.completions("testing/file"), vec!["testing/file_with_text"]);
        assert_eq!(completer.completions("~"), vec!["~/"]);
        assert_eq!(completer.completions("tes/fil"), vec!["testing/file_with_text"]);
    }
}
<|endoftext|>


<commit_before><commit_msg>GC most of ChainReplMsg.<commit_after><|endoftext|>


<commit_before>//! atty is a simple utility that answers one question
//! > is this a tty?
//!
//! usage is just as simple
//!
//! ```
//! if atty::is(atty::Stream::Stdout) {
//!   println!("i'm a tty")
//! }
//! ```
//!
//! ```
//! if atty::isnt(atty::Stream::Stdout) {
//!   println!("i'm not a tty")
//! }
//! ```

extern crate libc;

/// possible stream sources
pub enum Stream {
    Stdout,
    Stderr,
    Stdin,
}

/// returns true if this is a tty
#[cfg(unix)]
pub fn is(stream: Stream) -> bool {
    let fd = match stream {
        Stream::Stdout => libc::STDOUT_FILENO,
        Stream::Stderr => libc::STDERR_FILENO,
        Stream::Stdin => libc::STDIN_FILENO,
    };
    unsafe { libc::isatty(fd) != 0 }
}

/// returns true if this is a tty
#[cfg(windows)]
pub fn is(stream: Stream) -> bool {
    extern crate kernel32;
    extern crate winapi;

    let handle = match stream {
        Stream::Stdout => winapi::STD_OUTPUT_HANDLE,
        Stream::Stderr => winapi::STD_ERROR_HANDLE,
        Stream::Stdin => winapi::STD_INPUT_HANDLE,
    };

    unsafe {
        match stream {
            Stream::Stdin => {
                let mut out = 0;
                kernel32::GetConsoleMode(kernel32::GetStdHandle(winapi::STD_INPUT_HANDLE),
                                         &mut out) != 0
            }
            out => {
                let name = match out {
                    Stream::Stdout => b"CONOUT$\0",
                    Stream::Stderr => b"CONERR$\0",
                    _ => unreachable!(),
                };
                let handle = kernel32::CreateFileA(name.as_ptr() as *const i8,
                                                   winapi::GENERIC_READ | winapi::GENERIC_WRITE,
                                                   winapi::FILE_SHARE_WRITE,
                                                   ::std::ptr::null_mut(),
                                                   winapi::OPEN_EXISTING,
                                                   0,
                                                   ::std::ptr::null_mut());
                // https://msdn.microsoft.com/en-us/library/windows/desktop/ms683171(v=vs.85).aspx
                let mut buffer_info = ::std::mem::uninitialized();
                let ret = kernel32::GetConsoleScreenBufferInfo(handle, &mut buffer_info);
                let last_err = kernel32::GetLastError();
                panic!("is invalid? {:#?}  result {:#?} last err {:#?}",
                       winapi::INVALID_HANDLE_VALUE == handle,
                       ret,
                       last_err);
                ret != 0
            }
        }

    }
}

/// returns true if this is _not_ a tty
pub fn isnt(stream: Stream) -> bool {
    !is(stream)
}

#[cfg(test)]
mod tests {
    use super::{is, Stream};

    #[test]
    fn is_err() {
        assert!(is(Stream::Stderr))
    }

    #[test]
    fn is_out() {
        assert!(is(Stream::Stdout))
    }

    #[test]
    fn is_in() {
        assert!(is(Stream::Stdin))
    }
}
<commit_msg>winapi makes me sad<commit_after>//! atty is a simple utility that answers one question
//! > is this a tty?
//!
//! usage is just as simple
//!
//! ```
//! if atty::is(atty::Stream::Stdout) {
//!   println!("i'm a tty")
//! }
//! ```
//!
//! ```
//! if atty::isnt(atty::Stream::Stdout) {
//!   println!("i'm not a tty")
//! }
//! ```

extern crate libc;

/// possible stream sources
pub enum Stream {
    Stdout,
    Stderr,
    Stdin,
}

/// returns true if this is a tty
#[cfg(unix)]
pub fn is(stream: Stream) -> bool {
    let fd = match stream {
        Stream::Stdout => libc::STDOUT_FILENO,
        Stream::Stderr => libc::STDERR_FILENO,
        Stream::Stdin => libc::STDIN_FILENO,
    };
    unsafe { libc::isatty(fd) != 0 }
}

/// returns true if this is a tty
#[cfg(windows)]
pub fn is(stream: Stream) -> bool {
    extern crate kernel32;
    extern crate winapi;

    let handle = match stream {
        Stream::Stdout => winapi::STD_OUTPUT_HANDLE,
        Stream::Stderr => winapi::STD_ERROR_HANDLE,
        Stream::Stdin => winapi::STD_INPUT_HANDLE,
    };

    unsafe {
        match stream {
            Stream::Stdin => {
                let mut out = 0;
                kernel32::GetConsoleMode(kernel32::GetStdHandle(winapi::STD_INPUT_HANDLE),
                                         &mut out) != 0
            }
            out => {
                // note: there is no CONERR, only CONOUT
                let handle = kernel32::CreateFileA(b"CONOUT$\0".as_ptr() as *const i8,
                                                   winapi::GENERIC_READ | winapi::GENERIC_WRITE,
                                                   winapi::FILE_SHARE_WRITE,
                                                   ::std::ptr::null_mut(),
                                                   winapi::OPEN_EXISTING,
                                                   0,
                                                   ::std::ptr::null_mut());
                // https://msdn.microsoft.com/en-us/library/windows/desktop/ms683171(v=vs.85).aspx
                let mut buffer_info = ::std::mem::uninitialized();
                let ret = kernel32::GetConsoleScreenBufferInfo(handle, &mut buffer_info);
                let last_err = kernel32::GetLastError();
                panic!("is invalid? {:#?}  result {:#?} last err {:#?}",
                       winapi::INVALID_HANDLE_VALUE == handle,
                       ret,
                       last_err);
                ret != 0
            }
        }

    }
}

/// returns true if this is _not_ a tty
pub fn isnt(stream: Stream) -> bool {
    !is(stream)
}

#[cfg(test)]
mod tests {
    use super::{is, Stream};

    #[test]
    fn is_err() {
        assert!(is(Stream::Stderr))
    }

    #[test]
    fn is_out() {
        assert!(is(Stream::Stdout))
    }

    #[test]
    fn is_in() {
        assert!(is(Stream::Stdin))
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate serde_json;

pub mod cmd;

pub mod ui;

pub mod hooks;

use std::collections::HashMap;
use std::str::FromStr;
use std::io::{
  Read,
  Write,
  BufReader,
  BufWriter,
};

use std::path::Path;
use std::fs::{
  File,
  OpenOptions,
};

use std::fmt::{
  Display,
  Formatter
};

#[derive(Debug)]
pub enum KVError {
  IoError(std::io::Error),
  EncodingError(serde_json::Error),
  UnknownError(String),
}

impl Display for KVError {
  fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
    match *self {
      KVError::IoError(ref e) => write!(f, "{}", e.to_string()),
      KVError::EncodingError(ref e) => write!(f, "{}", e.to_string()),
      KVError::UnknownError(ref msg) => write!(f, "{}", msg),
    }
  }
}

type Result<V> = std::result::Result<V, KVError>;

impl From<serde_json::Error> for KVError {
  fn from(e: serde_json::Error) -> Self {
    KVError::EncodingError(e)
  }
}

impl From<std::io::Error> for KVError {
  fn from(e: std::io::Error) -> Self {
    KVError::IoError(e)
  }
}

pub fn init<P: AsRef<Path>>(path: P) -> Result<KVStore> {
  let file = File::create(path)?;
  let mut writer = BufWriter::new(file);

  let kvs = KVStore::new();
  kvs.write(&mut writer)?;

  Ok(kvs)
}

#[derive(Serialize, Deserialize, Debug)]
pub struct KVStore {
  content: HashMap<String, Value>
}

impl FromStr for KVStore {
  type Err = KVError;

  fn from_str(s: &str) -> Result<Self> {
    Ok(serde_json::from_str(s)?)
  }
}

impl ToString for KVStore {
  fn to_string(&self) -> String {
    self.serialize().unwrap()
  }
}

impl KVStore {
  pub fn new() -> KVStore {
    KVStore {
      content: HashMap::new()
    }
  }

  pub fn read<R: Read>(r: R) -> Result<KVStore> {
    Ok(serde_json::from_reader(r)?)
  }

  pub fn read_from_file<P: AsRef<Path>>(path: P) -> Result<KVStore> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);

    KVStore::read(reader)
  }

  pub fn serialize(&self) -> Result<String> {
    Ok(serde_json::to_string(self)?)
  }

  pub fn write<W: Write>(&self, w: &mut W) -> Result<()> {
    Ok(serde_json::to_writer(w, self)?)
  }

  pub fn write_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
    let file = OpenOptions::new().write(true).truncate(true).open(path)?;
    let mut writer = BufWriter::new(file);

    self.write(&mut writer)
  }

  fn put_value<S: ToString>(&mut self, key: S, value: Value) -> Option<Value> {
    self.content.insert(key.to_string(), value)
  }

  fn get_mut<S: ToString>(&mut self, key: S) -> Result<&mut Value> {
    self.content.get_mut(&key.to_string()).ok_or(KVError::UnknownError("no key".to_string()))
  }

  fn get_mut_list<S: ToString>(&mut self, key: S) -> Result<&mut Vec<String>> {
    let key = key.to_string();
    let kv_value = self.get_mut(&key)?;

    if let &mut Value::ListValue(ref mut list) = kv_value {
      Ok(list)
    } else {
      Err(KVError::UnknownError(format!("value at {} not a list", &key)))
    }
  }

  pub fn get<S: ToString>(&self, key: &S) -> Option<&Value> {
    self.content.get(&key.to_string())
  }

  pub fn put<KS: ToString, VS: ToString>(&mut self, key: KS, value: VS) {
    let kv_value = Value::StringValue(value.to_string());
    self.put_value(key, kv_value);
  }

  pub fn put_empty_list<S: ToString>(&mut self, key: S) -> Option<Value> {
    let empty_list = Vec::new();
    let value = Value::ListValue(empty_list);

    self.put_value(key.to_string(), value)
  }

  pub fn push_value<KS: ToString, VS: ToString>(&mut self, key: KS, value: VS) -> Result<()> {
    let list = self.get_mut_list(key)?;
    Ok(list.push(value.to_string()))
  }

  pub fn push_all_values<KS: ToString, VS: ToString>(&mut self, key: KS, values: Vec<VS>) -> Result<()> {
    let list = self.get_mut_list(key)?;
    let mut string_values = values.iter().map(ToString::to_string).collect();
    Ok(list.append(&mut string_values))
  }

  pub fn pop_value<KS: ToString>(&mut self, key: KS) -> Result<String> {
    let list_value = self.get_mut_list(key)?;

    list_value.pop().ok_or(KVError::UnknownError("list is empty".to_string()))
  }

  pub fn get_keys(&self) -> Vec<&String> {
    self.content.keys().collect()
  }

  pub fn drop<KS: ToString>(&mut self, key: KS) -> Option<Value> {
    self.content.remove(&key.to_string())
  }

  pub fn has_key<KS: ToString>(&self, key: KS) -> bool {
    self.content.contains_key(&key.to_string())
  }

  pub fn get_value_type<KS: ToString>(&self, key: &KS) -> Option<ValueType> {
    self.get(key).map(Value::get_type)
  }
}

pub enum ValueType {
  String,
  List,
}

#[derive(Serialize, Deserialize, Debug)]
pub enum Value {
  StringValue(String),
  ListValue(Vec<String>)
}

impl Value {
  pub fn get_type(&self) -> ValueType {
    match *self {
      Value::StringValue(_) => ValueType::String,
      Value::ListValue(_) => ValueType::List,
    }
  }
}

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

  #[test]
  fn test_create() {
    KVStore::new();
  }
}
<commit_msg>write prtty json<commit_after>#[macro_use]
extern crate serde_derive;
extern crate serde_json;

pub mod cmd;

pub mod ui;

pub mod hooks;

use std::collections::HashMap;
use std::str::FromStr;
use std::io::{
  Read,
  Write,
  BufReader,
  BufWriter,
};

use std::path::Path;
use std::fs::{
  File,
  OpenOptions,
};

use std::fmt::{
  Display,
  Formatter
};

#[derive(Debug)]
pub enum KVError {
  IoError(std::io::Error),
  EncodingError(serde_json::Error),
  UnknownError(String),
}

impl Display for KVError {
  fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
    match *self {
      KVError::IoError(ref e) => write!(f, "{}", e.to_string()),
      KVError::EncodingError(ref e) => write!(f, "{}", e.to_string()),
      KVError::UnknownError(ref msg) => write!(f, "{}", msg),
    }
  }
}

type Result<V> = std::result::Result<V, KVError>;

impl From<serde_json::Error> for KVError {
  fn from(e: serde_json::Error) -> Self {
    KVError::EncodingError(e)
  }
}

impl From<std::io::Error> for KVError {
  fn from(e: std::io::Error) -> Self {
    KVError::IoError(e)
  }
}

pub fn init<P: AsRef<Path>>(path: P) -> Result<KVStore> {
  let file = File::create(path)?;
  let mut writer = BufWriter::new(file);

  let kvs = KVStore::new();
  kvs.write(&mut writer)?;

  Ok(kvs)
}

#[derive(Serialize, Deserialize, Debug)]
pub struct KVStore {
  content: HashMap<String, Value>
}

impl FromStr for KVStore {
  type Err = KVError;

  fn from_str(s: &str) -> Result<Self> {
    Ok(serde_json::from_str(s)?)
  }
}

impl ToString for KVStore {
  fn to_string(&self) -> String {
    self.serialize().unwrap()
  }
}

impl KVStore {
  pub fn new() -> KVStore {
    KVStore {
      content: HashMap::new()
    }
  }

  pub fn read<R: Read>(r: R) -> Result<KVStore> {
    Ok(serde_json::from_reader(r)?)
  }

  pub fn read_from_file<P: AsRef<Path>>(path: P) -> Result<KVStore> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);

    KVStore::read(reader)
  }

  pub fn serialize(&self) -> Result<String> {
    Ok(serde_json::to_string(self)?)
  }

  pub fn write<W: Write>(&self, w: &mut W) -> Result<()> {
    Ok(serde_json::to_writer_pretty(w, self)?)
  }

  pub fn write_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
    let file = OpenOptions::new().write(true).truncate(true).open(path)?;
    let mut writer = BufWriter::new(file);

    self.write(&mut writer)
  }

  fn put_value<S: ToString>(&mut self, key: S, value: Value) -> Option<Value> {
    self.content.insert(key.to_string(), value)
  }

  fn get_mut<S: ToString>(&mut self, key: S) -> Result<&mut Value> {
    self.content.get_mut(&key.to_string()).ok_or(KVError::UnknownError("no key".to_string()))
  }

  fn get_mut_list<S: ToString>(&mut self, key: S) -> Result<&mut Vec<String>> {
    let key = key.to_string();
    let kv_value = self.get_mut(&key)?;

    if let &mut Value::ListValue(ref mut list) = kv_value {
      Ok(list)
    } else {
      Err(KVError::UnknownError(format!("value at {} not a list", &key)))
    }
  }

  pub fn get<S: ToString>(&self, key: &S) -> Option<&Value> {
    self.content.get(&key.to_string())
  }

  pub fn put<KS: ToString, VS: ToString>(&mut self, key: KS, value: VS) {
    let kv_value = Value::StringValue(value.to_string());
    self.put_value(key, kv_value);
  }

  pub fn put_empty_list<S: ToString>(&mut self, key: S) -> Option<Value> {
    let empty_list = Vec::new();
    let value = Value::ListValue(empty_list);

    self.put_value(key.to_string(), value)
  }

  pub fn push_value<KS: ToString, VS: ToString>(&mut self, key: KS, value: VS) -> Result<()> {
    let list = self.get_mut_list(key)?;
    Ok(list.push(value.to_string()))
  }

  pub fn push_all_values<KS: ToString, VS: ToString>(&mut self, key: KS, values: Vec<VS>) -> Result<()> {
    let list = self.get_mut_list(key)?;
    let mut string_values = values.iter().map(ToString::to_string).collect();
    Ok(list.append(&mut string_values))
  }

  pub fn pop_value<KS: ToString>(&mut self, key: KS) -> Result<String> {
    let list_value = self.get_mut_list(key)?;

    list_value.pop().ok_or(KVError::UnknownError("list is empty".to_string()))
  }

  pub fn get_keys(&self) -> Vec<&String> {
    self.content.keys().collect()
  }

  pub fn drop<KS: ToString>(&mut self, key: KS) -> Option<Value> {
    self.content.remove(&key.to_string())
  }

  pub fn has_key<KS: ToString>(&self, key: KS) -> bool {
    self.content.contains_key(&key.to_string())
  }

  pub fn get_value_type<KS: ToString>(&self, key: &KS) -> Option<ValueType> {
    self.get(key).map(Value::get_type)
  }
}

pub enum ValueType {
  String,
  List,
}

#[derive(Serialize, Deserialize, Debug)]
pub enum Value {
  StringValue(String),
  ListValue(Vec<String>)
}

impl Value {
  pub fn get_type(&self) -> ValueType {
    match *self {
      Value::StringValue(_) => ValueType::String,
      Value::ListValue(_) => ValueType::List,
    }
  }
}

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

  #[test]
  fn test_create() {
    KVStore::new();
  }
}
<|endoftext|>


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

//! Vec2D is a very simple 2D container for storing rectangular data

#![deny(missing_docs)]

/// 2D coordinate
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Coord {
    /// X component
    pub x: usize,

    /// Y component
    pub y: usize
}

/// Rectangle defined by inclusive minimum and maximum coordinates
#[derive(Clone, Copy, Eq, Debug, PartialEq)]
pub struct Rect {
    /// Minimum coordinate (inclusive)
    min_coord: Coord,

    /// Maximum coordinate (inclusive)
    max_coord: Coord
}

/// Rectangle dimensions
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Size {
    /// Width of rectangle
    pub width: usize,
    /// Height of rectangle
    pub height: usize
}

/// Container for 2D data
#[derive(Clone, Debug)]
pub struct Vec2D<T> {
    elems: Vec<T>,
    size: Size
}

/// Iterator over a rectangle within a Vec2D
pub struct RectIter<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *const Elem,
    cur_coord: Coord,
    stride: isize
}

/// Mutable iterator over a rectangle within a Vec2D
pub struct RectIterMut<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a mut Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *mut Elem,
    cur_coord: Coord,
    stride: isize
}

impl Coord {
    /// Create a coordinate at (x, y)
    pub fn new(x: usize, y: usize) -> Coord {
        Coord {
            x: x,
            y: y
        }
    }
}

impl std::ops::Add for Coord {
    type Output = Coord;

    fn add(self, other: Coord) -> Coord {
        Coord::new(self.x + other.x,
                   self.y + other.y)
    }
}

impl Rect {
    /// Calculate rectangle width
    pub fn width(&self) -> usize {
        return self.max_coord.x - self.min_coord.x + 1;
    }

    /// Calculate rectangle height
    pub fn height(&self) -> usize {
        return self.max_coord.y - self.min_coord.y + 1;
    }

    /// Calculate rectangle size
    pub fn size(&self) -> Size {
        Size::new(self.width(), self.height())
    }

    /// Return true if the coordinate is between `min_coord` and
    /// `max_coord` (inclusive).
    pub fn contains_coord(&self, coord: Coord) -> bool {
        (coord.x >= self.min_coord.x && coord.x <= self.max_coord.x &&
         coord.y >= self.min_coord.y && coord.y <= self.max_coord.y)
    }
}

impl Size {
    /// Create a 2D size of (width, height)
    pub fn new(width: usize, height: usize) -> Size {
        Size {
            width: width,
            height: height
        }
    }

    /// width * height
    pub fn area(&self) -> usize {
        self.width * self.height
    }

    /// Return true if the coordinate fits within self's width and
    /// height, false otherwise.
    pub fn contains_coord(&self, coord: Coord) -> bool {
        coord.x < self.width && coord.y < self.height
    }

    /// Create a rectangle starting at (0, 0) with `self`'s size.
    pub fn rect(&self) -> Rect {
        Rect {
            min_coord: Coord::new(0, 0),
            max_coord: Coord::new(self.width - 1, self.height - 1)
        }
    }
}

impl<Elem: Copy> Vec2D<Elem> {
    /// Create a Vec2D with the given `size`. All elements are
    /// initialized as copies of the `example` element.
    pub fn from_example(size: Size, example: &Elem) -> Vec2D<Elem> {
        Vec2D {
            elems: vec![*example; size.area()],
            size: size
        }
    }

    /// Create a Vec2D with the given `size`. The contents are set to
    /// `src`. None is returned if the `size` does not match the
    /// length of `src`.
    pub fn from_vec(size: Size, src: Vec<Elem>) -> Option<Vec2D<Elem>> {
        if size.area() == src.len() {
            Some(Vec2D {
                elems: src,
                size: size
            })
        }
        else {
            None
        }
    }

    /// Shortcut for self.size.rect()
    pub fn rect(&self) -> Rect {
        self.size.rect()
    }

    fn stride(&self, rect: &Rect) -> isize {
        (self.size.width + 1 - rect.width()) as isize
    }

    /// Calculate pointer offset for `start` element.
    fn start_offset(&self, start: Coord) -> isize {
        debug_assert_eq!(self.size.contains_coord(start), true);
        (start.y * self.size.width + start.x) as isize
    }

    /// Iterator over the entire Vec2D.
    pub fn iter<'a>(&'a self) -> RectIter<'a, Elem> {
        self.rect_iter(self.size.rect()).unwrap()
    }

    /// Create an iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter<'a>(&'a self, rect: Rect) -> Option<RectIter<'a, Elem>> {
        self.rect_iter_at(rect, rect.min_coord)
    }

    /// Create an iterator over a rectangular region of the Vec2D with
    /// the `start` coord. None is returned if the given `rect` does
    /// not fit entirely within the Vec2D or if the `start` coord is
    /// not within `rect`.
    pub fn rect_iter_at<'a>(&'a self, rect: Rect,
                            start: Coord) -> Option<RectIter<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) && rect.contains_coord(start) {
            Some(RectIter {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: unsafe {
                    self.elems.as_ptr().offset(self.start_offset(start))
                },
                rect: rect,
                cur_coord: start
            })
        }
        else {
            None
        }
    }

    /// Mutable iterater over the entire Vec2D.
    pub fn iter_mut<'a>(&'a mut self) -> RectIterMut<'a, Elem> {
        let rect = self.size.rect();
        self.rect_iter_mut(rect).unwrap()
    }

    /// Create a mutable iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter_mut<'a>(&'a mut self, rect: Rect) -> Option<RectIterMut<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) {
            Some(RectIterMut {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: self.elems.as_mut_ptr(),
                rect: rect,
                cur_coord: rect.min_coord
            })
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIter<'a, Elem> {
    type Item = (Coord, &'a Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { & *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIterMut<'a, Elem> {
    type Item = (Coord, &'a mut Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { &mut *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl Rect {
    /// Create a new Rect defined by inclusive minimum and maximum
    /// coordinates. If min_coord is greater than max_coord on either
    /// axis then None is returned.
    pub fn new(min_coord: Coord, max_coord: Coord) -> Option<Rect> {
        if min_coord.x <= max_coord.x && min_coord.y <= max_coord.y {
            Some(Rect {
                min_coord: min_coord,
                max_coord: max_coord
            })
        }
        else {
            None
        }
    }

}

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

    #[test]
    fn test_coord() {
        let coord = Coord::new(1, 2);
        assert_eq!(coord.x, 1);
        assert_eq!(coord.y, 2);
    }

    #[test]
    fn test_coord_add() {
        let a = Coord::new(1, 2);
        let b = Coord::new(5, 9);
        assert_eq!(a + b, Coord::new(6, 11));
    }

    #[test]
    fn test_rect() {
        let rect = Rect::new(Coord::new(1, 2), Coord::new(5, 3)).unwrap();
        assert_eq!(rect.width(), 5);
        assert_eq!(rect.height(), 2);

        assert_eq!(rect.width(), rect.size().width);
        assert_eq!(rect.height(), rect.size().height);

        assert_eq!(rect.contains_coord(Coord::new(0, 0)), false);
        assert_eq!(rect.contains_coord(Coord::new(4, 3)), true);
    }

    #[test]
    fn test_bad_rect() {
        assert_eq!(Rect::new(Coord::new(2, 1), Coord::new(1, 1)).is_none(),
                   true);
        assert_eq!(Rect::new(Coord::new(1, 2), Coord::new(1, 1)).is_none(),
                   true);
    }

    #[test]
    fn test_size() {
        let size = Size::new(3, 2);
        assert_eq!(size.width, 3);
        assert_eq!(size.height, 2);

        assert_eq!(size.area(), 6);

        assert_eq!(size.contains_coord(Coord::new(1, 1)), true);
        assert_eq!(size.contains_coord(Coord::new(4, 1)), false);
        assert_eq!(size.contains_coord(Coord::new(1, 3)), false);

        let rect = size.rect();
        assert_eq!(rect.min_coord, Coord::new(0, 0));
        assert_eq!(rect.max_coord, Coord::new(2, 1));
    }

    #[test]
    fn test_rect_iter_mut() {
        let elems = vec![1, 2, 3, 4];
        let mut grid = Vec2D::from_vec(Size::new(2, 2), elems).unwrap();
        let rect = Rect::new(Coord::new(0, 0), Coord::new(1, 1)).unwrap();

        let mut actual_coords = Vec::new();
        for (coord, elem) in grid.rect_iter_mut(rect).unwrap() {
            *elem = -(*elem);
            actual_coords.push((coord.x, coord.y));
        }
        assert_eq!(actual_coords, [(0, 0), (1, 0), (0, 1), (1, 1)]);
        assert_eq!(grid.elems, [-1, -2, -3, -4]);
    }

    #[test]
    fn test_two_iterators() {
        let size = Size::new(2, 1);
        let v = Vec2D::from_vec(size, vec![0, 1]).unwrap();

        let iter1 = v.rect_iter(size.rect()).unwrap();
        let iter2 = v.rect_iter(size.rect()).unwrap();

        for ((coord1, elem1), (coord2, elem2)) in iter1.zip(iter2) {
            assert_eq!(coord1, coord2);
            assert_eq!(elem1, elem2);
        }
    }

    #[test]
    fn test_rect_iter_at() {
        let size = Size::new(1, 2);
        let v = Vec2D::from_vec(size, vec![0, 1]).unwrap();

        let start = Coord::new(0, 1);
        let mut iter = v.rect_iter_at(size.rect(), start).unwrap();
        let (coord, elem) = iter.next().unwrap();
        assert_eq!((coord, *elem), (start, 1));
        assert_eq!(iter.next().is_none(), true);
    }
}
<commit_msg>add missing rect_iter_mut_at function and fix rect_iter_mut to use it<commit_after>// Copyright 2015 Nicholas Bishop
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Vec2D is a very simple 2D container for storing rectangular data

#![deny(missing_docs)]

/// 2D coordinate
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Coord {
    /// X component
    pub x: usize,

    /// Y component
    pub y: usize
}

/// Rectangle defined by inclusive minimum and maximum coordinates
#[derive(Clone, Copy, Eq, Debug, PartialEq)]
pub struct Rect {
    /// Minimum coordinate (inclusive)
    min_coord: Coord,

    /// Maximum coordinate (inclusive)
    max_coord: Coord
}

/// Rectangle dimensions
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Size {
    /// Width of rectangle
    pub width: usize,
    /// Height of rectangle
    pub height: usize
}

/// Container for 2D data
#[derive(Clone, Debug)]
pub struct Vec2D<T> {
    elems: Vec<T>,
    size: Size
}

/// Iterator over a rectangle within a Vec2D
pub struct RectIter<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *const Elem,
    cur_coord: Coord,
    stride: isize
}

/// Mutable iterator over a rectangle within a Vec2D
pub struct RectIterMut<'a, Elem: 'a> {
    grid: std::marker::PhantomData<&'a mut Vec2D<Elem>>,

    rect: Rect,
    cur_elem: *mut Elem,
    cur_coord: Coord,
    stride: isize
}

impl Coord {
    /// Create a coordinate at (x, y)
    pub fn new(x: usize, y: usize) -> Coord {
        Coord {
            x: x,
            y: y
        }
    }
}

impl std::ops::Add for Coord {
    type Output = Coord;

    fn add(self, other: Coord) -> Coord {
        Coord::new(self.x + other.x,
                   self.y + other.y)
    }
}

impl Rect {
    /// Calculate rectangle width
    pub fn width(&self) -> usize {
        return self.max_coord.x - self.min_coord.x + 1;
    }

    /// Calculate rectangle height
    pub fn height(&self) -> usize {
        return self.max_coord.y - self.min_coord.y + 1;
    }

    /// Calculate rectangle size
    pub fn size(&self) -> Size {
        Size::new(self.width(), self.height())
    }

    /// Return true if the coordinate is between `min_coord` and
    /// `max_coord` (inclusive).
    pub fn contains_coord(&self, coord: Coord) -> bool {
        (coord.x >= self.min_coord.x && coord.x <= self.max_coord.x &&
         coord.y >= self.min_coord.y && coord.y <= self.max_coord.y)
    }
}

impl Size {
    /// Create a 2D size of (width, height)
    pub fn new(width: usize, height: usize) -> Size {
        Size {
            width: width,
            height: height
        }
    }

    /// width * height
    pub fn area(&self) -> usize {
        self.width * self.height
    }

    /// Return true if the coordinate fits within self's width and
    /// height, false otherwise.
    pub fn contains_coord(&self, coord: Coord) -> bool {
        coord.x < self.width && coord.y < self.height
    }

    /// Create a rectangle starting at (0, 0) with `self`'s size.
    pub fn rect(&self) -> Rect {
        Rect {
            min_coord: Coord::new(0, 0),
            max_coord: Coord::new(self.width - 1, self.height - 1)
        }
    }
}

impl<Elem: Copy> Vec2D<Elem> {
    /// Create a Vec2D with the given `size`. All elements are
    /// initialized as copies of the `example` element.
    pub fn from_example(size: Size, example: &Elem) -> Vec2D<Elem> {
        Vec2D {
            elems: vec![*example; size.area()],
            size: size
        }
    }

    /// Create a Vec2D with the given `size`. The contents are set to
    /// `src`. None is returned if the `size` does not match the
    /// length of `src`.
    pub fn from_vec(size: Size, src: Vec<Elem>) -> Option<Vec2D<Elem>> {
        if size.area() == src.len() {
            Some(Vec2D {
                elems: src,
                size: size
            })
        }
        else {
            None
        }
    }

    /// Shortcut for self.size.rect()
    pub fn rect(&self) -> Rect {
        self.size.rect()
    }

    fn stride(&self, rect: &Rect) -> isize {
        (self.size.width + 1 - rect.width()) as isize
    }

    /// Calculate pointer offset for `start` element.
    fn start_offset(&self, start: Coord) -> isize {
        debug_assert_eq!(self.size.contains_coord(start), true);
        (start.y * self.size.width + start.x) as isize
    }

    /// Iterator over the entire Vec2D.
    pub fn iter<'a>(&'a self) -> RectIter<'a, Elem> {
        self.rect_iter(self.size.rect()).unwrap()
    }

    /// Create an iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter<'a>(&'a self, rect: Rect) -> Option<RectIter<'a, Elem>> {
        self.rect_iter_at(rect, rect.min_coord)
    }

    /// Create an iterator over a rectangular region of the Vec2D with
    /// the `start` coord. None is returned if the given `rect` does
    /// not fit entirely within the Vec2D or if the `start` coord is
    /// not within `rect`.
    pub fn rect_iter_at<'a>(&'a self, rect: Rect,
                            start: Coord) -> Option<RectIter<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) && rect.contains_coord(start) {
            Some(RectIter {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: unsafe {
                    self.elems.as_ptr().offset(self.start_offset(start))
                },
                rect: rect,
                cur_coord: start
            })
        }
        else {
            None
        }
    }

    /// Mutable iterater over the entire Vec2D.
    pub fn iter_mut<'a>(&'a mut self) -> RectIterMut<'a, Elem> {
        let rect = self.size.rect();
        self.rect_iter_mut(rect).unwrap()
    }

    /// Create a mutable iterator over a rectangular region of the
    /// Vec2D. None is returned if the given `rect` does not fit
    /// entirely within the Vec2D.
    pub fn rect_iter_mut<'a>(&'a mut self, rect: Rect) -> Option<RectIterMut<'a, Elem>> {
        self.rect_iter_mut_at(rect, rect.min_coord)
    }

    /// Create a mutable iterator over a rectangular region of the
    /// Vec2D with the `start` coord. None is returned if the given
    /// `rect` does not fit entirely within the Vec2D or if the
    /// `start` coord is not within `rect`.
    pub fn rect_iter_mut_at<'a>(&'a mut self, rect: Rect,
                            start: Coord) -> Option<RectIterMut<'a, Elem>> {
        if self.size.contains_coord(rect.max_coord) && rect.contains_coord(start) {
            Some(RectIterMut {
                grid: std::marker::PhantomData,
                stride: self.stride(&rect),
                cur_elem: unsafe {
                    self.elems.as_mut_ptr().offset(self.start_offset(start))
                },
                rect: rect,
                cur_coord: start
            })
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIter<'a, Elem> {
    type Item = (Coord, &'a Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { & *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl<'a, Elem> Iterator for RectIterMut<'a, Elem> {
    type Item = (Coord, &'a mut Elem);

    fn next(&mut self) -> Option<Self::Item> {
        if self.cur_coord.y <= self.rect.max_coord.y {
            let result = (self.cur_coord, unsafe { &mut *self.cur_elem });

            self.cur_coord.x += 1;
            if self.cur_coord.x <= self.rect.max_coord.x {
                unsafe { self.cur_elem = self.cur_elem.offset(1); }
            }
            else {
                self.cur_coord.x = self.rect.min_coord.x;
                self.cur_coord.y += 1;
                unsafe { self.cur_elem = self.cur_elem.offset(self.stride); }
            }
            Some(result)
        }
        else {
            None
        }
    }
}

impl Rect {
    /// Create a new Rect defined by inclusive minimum and maximum
    /// coordinates. If min_coord is greater than max_coord on either
    /// axis then None is returned.
    pub fn new(min_coord: Coord, max_coord: Coord) -> Option<Rect> {
        if min_coord.x <= max_coord.x && min_coord.y <= max_coord.y {
            Some(Rect {
                min_coord: min_coord,
                max_coord: max_coord
            })
        }
        else {
            None
        }
    }

}

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

    #[test]
    fn test_coord() {
        let coord = Coord::new(1, 2);
        assert_eq!(coord.x, 1);
        assert_eq!(coord.y, 2);
    }

    #[test]
    fn test_coord_add() {
        let a = Coord::new(1, 2);
        let b = Coord::new(5, 9);
        assert_eq!(a + b, Coord::new(6, 11));
    }

    #[test]
    fn test_rect() {
        let rect = Rect::new(Coord::new(1, 2), Coord::new(5, 3)).unwrap();
        assert_eq!(rect.width(), 5);
        assert_eq!(rect.height(), 2);

        assert_eq!(rect.width(), rect.size().width);
        assert_eq!(rect.height(), rect.size().height);

        assert_eq!(rect.contains_coord(Coord::new(0, 0)), false);
        assert_eq!(rect.contains_coord(Coord::new(4, 3)), true);
    }

    #[test]
    fn test_bad_rect() {
        assert_eq!(Rect::new(Coord::new(2, 1), Coord::new(1, 1)).is_none(),
                   true);
        assert_eq!(Rect::new(Coord::new(1, 2), Coord::new(1, 1)).is_none(),
                   true);
    }

    #[test]
    fn test_size() {
        let size = Size::new(3, 2);
        assert_eq!(size.width, 3);
        assert_eq!(size.height, 2);

        assert_eq!(size.area(), 6);

        assert_eq!(size.contains_coord(Coord::new(1, 1)), true);
        assert_eq!(size.contains_coord(Coord::new(4, 1)), false);
        assert_eq!(size.contains_coord(Coord::new(1, 3)), false);

        let rect = size.rect();
        assert_eq!(rect.min_coord, Coord::new(0, 0));
        assert_eq!(rect.max_coord, Coord::new(2, 1));
    }

    #[test]
    fn test_rect_iter_mut() {
        let elems = vec![1, 2, 3, 4];
        let mut grid = Vec2D::from_vec(Size::new(2, 2), elems).unwrap();
        let rect = Rect::new(Coord::new(0, 0), Coord::new(1, 1)).unwrap();

        let mut actual_coords = Vec::new();
        for (coord, elem) in grid.rect_iter_mut(rect).unwrap() {
            *elem = -(*elem);
            actual_coords.push((coord.x, coord.y));
        }
        assert_eq!(actual_coords, [(0, 0), (1, 0), (0, 1), (1, 1)]);
        assert_eq!(grid.elems, [-1, -2, -3, -4]);
    }

    #[test]
    fn test_two_iterators() {
        let size = Size::new(2, 1);
        let v = Vec2D::from_vec(size, vec![0, 1]).unwrap();

        let iter1 = v.rect_iter(size.rect()).unwrap();
        let iter2 = v.rect_iter(size.rect()).unwrap();

        for ((coord1, elem1), (coord2, elem2)) in iter1.zip(iter2) {
            assert_eq!(coord1, coord2);
            assert_eq!(elem1, elem2);
        }
    }

    #[test]
    fn test_rect_iter_at() {
        let size = Size::new(1, 2);
        let v = Vec2D::from_vec(size, vec![0, 1]).unwrap();

        let start = Coord::new(0, 1);
        let mut iter = v.rect_iter_at(size.rect(), start).unwrap();
        let (coord, elem) = iter.next().unwrap();
        assert_eq!((coord, *elem), (start, 1));
        assert_eq!(iter.next().is_none(), true);
    }
}
<|endoftext|>


<commit_before>extern mod extra;

use extra::digest::Digest;
use extra::md5::Md5;
use extra::url::Url;
use std::cell::Cell;
use std::rt::io::io_error;
use std::rt::io::net::ip::SocketAddr;
use std::rt::io::net::tcp::TcpStream;
use std::str;

use message::*;

mod message;

pub struct PostgresConnection {
    priv stream: Cell<TcpStream>,
    priv next_stmt_id: Cell<int>
}

impl Drop for PostgresConnection {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.write_message(&Terminate);
        }
    }
}

impl PostgresConnection {
    pub fn connect(url: &str) -> PostgresConnection {
        let url: Url = FromStr::from_str(url).unwrap();

        let socket_url = fmt!("%s:%s", url.host,
                              url.port.get_ref().as_slice());
        let addr: SocketAddr = FromStr::from_str(socket_url).unwrap();
        let conn = PostgresConnection {
            stream: Cell::new(TcpStream::connect(addr).unwrap()),
            next_stmt_id: Cell::new(0)
        };

        let mut args = url.query.clone();
        args.push((~"user", url.user.get_ref().user.clone()));
        if !url.path.is_empty() {
            args.push((~"database", url.path.clone()));
        }
        conn.write_message(&StartupMessage(args.as_slice()));

        conn.handle_auth(&url);

        loop {
            match conn.read_message() {
                ParameterStatus(param, value) =>
                    info!("Parameter %s = %s", param, value),
                BackendKeyData(*) => (),
                ReadyForQuery(*) => break,
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }

        conn
    }

    fn write_message(&self, message: &FrontendMessage) {
        do self.stream.with_mut_ref |s| {
            s.write_message(message);
        }
    }

    fn read_message(&self) -> BackendMessage {
        do self.stream.with_mut_ref |s| {
            s.read_message()
        }
    }

    fn handle_auth(&self, url: &Url) {
        loop {
            match self.read_message() {
                AuthenticationOk => break,
                AuthenticationCleartextPassword => {
                    let pass = url.user.get_ref().pass.get_ref().as_slice();
                    self.write_message(&PasswordMessage(pass));
                }
                AuthenticationMD5Password(nonce) => {
                    let input = url.user.get_ref().pass.get_ref().as_slice() +
                            url.user.get_ref().user.as_slice();
                    let mut md5 = Md5::new();
                    md5.input_str(input);
                    let output = md5.result_str();
                    md5.reset();
                    md5.input_str(output);
                    md5.input(nonce);
                    let output = "md5" + md5.result_str();
                    self.write_message(&PasswordMessage(output.as_slice()));
                }
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }
    }

    pub fn prepare<'a>(&'a self, query: &str) -> PostgresStatement<'a> {
        let id = self.next_stmt_id.take();
        let stmt_name = format!("statement_{}", id);
        self.next_stmt_id.put_back(id + 1);

        let types = [];
        self.write_message(&Parse(stmt_name, query, types));
        self.write_message(&Sync);

        match self.read_message() {
            ParseComplete => (),
            resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
            resp => fail!("Bad response: %?", resp.to_str())
        }

        self.wait_for_ready();

        self.write_message(&Describe('S' as u8, stmt_name));
        self.write_message(&Sync);

        let num_params = match self.read_message() {
            ParameterDescription(ref types) => types.len(),
            resp => fail!("Bad response: %?", resp.to_str())
        };

        match self.read_message() {
            RowDescription(*) | NoData => (),
            resp => fail!("Bad response: %?", resp.to_str())
        }

        self.wait_for_ready();

        PostgresStatement {
            conn: self,
            name: stmt_name,
            num_params: num_params,
            next_portal_id: Cell::new(0)
        }
    }

    pub fn in_transaction<T, E: ToStr>(&self, blk: &fn(&PostgresConnection)
                                                       -> Result<T, E>)
                                       -> Result<T, E> {
        self.quick_query("BEGIN");

        // If this fails, Postgres will rollback when the connection closes
        let ret = blk(self);

        if ret.is_ok() {
            self.quick_query("COMMIT");
        } else {
            self.quick_query("ROLLBACK");
        }

        ret
    }

    fn quick_query(&self, query: &str) {
        self.write_message(&Query(query));

        loop {
            match self.read_message() {
                ReadyForQuery(*) => break,
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                _ => ()
            }
        }
    }

    fn wait_for_ready(&self) {
        loop {
            match self.read_message() {
                ReadyForQuery(*) => break,
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }
    }
}

pub struct PostgresStatement<'self> {
    priv conn: &'self PostgresConnection,
    priv name: ~str,
    priv num_params: uint,
    priv next_portal_id: Cell<uint>
}

#[unsafe_destructor]
impl<'self> Drop for PostgresStatement<'self> {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.conn.write_message(&Close('S' as u8, self.name.as_slice()));
            self.conn.write_message(&Sync);
            loop {
                match self.conn.read_message() {
                    ReadyForQuery(*) => break,
                    _ => ()
                }
            }
        }
    }
}

impl<'self> PostgresStatement<'self> {
    pub fn num_params(&self) -> uint {
        self.num_params
    }

    fn execute(&self, portal_name: &str, params: &[&ToSql]) {
        if self.num_params != params.len() {
            fail!("Expected %u params but got %u", self.num_params,
                  params.len());
        }

        let formats = [];
        let values: ~[Option<~[u8]>] = params.iter().map(|val| val.to_sql())
                .collect();
        let result_formats = [];

        self.conn.write_message(&Bind(portal_name, self.name.as_slice(),
                                      formats, values, result_formats));
        self.conn.write_message(&Execute(portal_name.as_slice(), 0));
        self.conn.write_message(&Sync);

        match self.conn.read_message() {
            BindComplete => (),
            resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
            resp => fail!("Bad response: %?", resp.to_str())
        }
    }

    pub fn update(&self, params: &[&ToSql]) -> uint {
        // The unnamed portal is automatically cleaned up at sync time
        self.execute("", params);

        let mut num = 0;
        loop {
            match self.conn.read_message() {
                CommandComplete(ret) => {
                    let s = ret.split_iter(' ').last().unwrap();
                    match FromStr::from_str(s) {
                        None => (),
                        Some(n) => num = n
                    }
                    break;
                }
                DataRow(*) => (),
                EmptyQueryResponse => break,
                NoticeResponse(*) => (),
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }
        self.conn.wait_for_ready();

        num
    }

    pub fn query<'a>(&'a self, params: &[&ToSql]) -> PostgresResult<'a> {
        let id = self.next_portal_id.take();
        let portal_name = format!("{:s}_portal_{}", self.name.as_slice(), id);
        self.next_portal_id.put_back(id + 1);

        self.execute(portal_name, params);

        let mut data = ~[];
        loop {
            match self.conn.read_message() {
                EmptyQueryResponse => break,
                DataRow(row) => data.push(row),
                CommandComplete(*) => break,
                NoticeResponse(*) => (),
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }

        PostgresResult {
            stmt: self,
            name: portal_name,
            data: data
        }
    }
}

pub struct PostgresResult<'self> {
    priv stmt: &'self PostgresStatement<'self>,
    priv name: ~str,
    priv data: ~[~[Option<~[u8]>]]
}

#[unsafe_destructor]
impl<'self> Drop for PostgresResult<'self> {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.stmt.conn.write_message(&Close('P' as u8,
                                                self.name.as_slice()));
            self.stmt.conn.write_message(&Sync);
            loop {
                match self.stmt.conn.read_message() {
                    ReadyForQuery(*) => break,
                    _ => ()
                }
            }
        }
    }
}

impl<'self> PostgresResult<'self> {
    pub fn iter<'a>(&'a self) -> PostgresResultIterator<'a> {
        PostgresResultIterator { result: self, next_row: 0 }
    }
}

pub struct PostgresResultIterator<'self> {
    priv result: &'self PostgresResult<'self>,
    priv next_row: uint
}

impl<'self> Iterator<PostgresRow<'self>> for PostgresResultIterator<'self> {
    fn next(&mut self) -> Option<PostgresRow<'self>> {
        if self.next_row == self.result.data.len() {
            return None;
        }

        let row = self.next_row;
        self.next_row += 1;
        Some(PostgresRow { result: self.result, row: row })
    }
}

pub struct PostgresRow<'self> {
    priv result: &'self PostgresResult<'self>,
    priv row: uint
}

impl<'self> Container for PostgresRow<'self> {
    fn len(&self) -> uint {
        self.result.data[self.row].len()
    }
}

impl<'self, T: FromSql> Index<uint, T> for PostgresRow<'self> {
    fn index(&self, idx: &uint) -> T {
        self.get(*idx)
    }
}

impl<'self> PostgresRow<'self> {
    pub fn get<T: FromSql>(&self, idx: uint) -> T {
        FromSql::from_sql(&self.result.data[self.row][idx])
    }
}

pub trait FromSql {
    fn from_sql(raw: &Option<~[u8]>) -> Self;
}

macro_rules! from_str_impl(
    ($t:ty) => (
        impl FromSql for Option<$t> {
            fn from_sql(raw: &Option<~[u8]>) -> Option<$t> {
                match *raw {
                    None => None,
                    Some(ref buf) => {
                        let s = str::from_bytes_slice(buf.as_slice());
                        Some(FromStr::from_str(s).unwrap())
                    }
                }
            }
        }
    )
)

macro_rules! from_option_impl(
    ($t:ty) => (
        impl FromSql for $t {
            fn from_sql(raw: &Option<~[u8]>) -> $t {
                FromSql::from_sql::<Option<$t>>(raw).unwrap()
            }
        }
    )
)

from_str_impl!(int)
from_option_impl!(int)
from_str_impl!(i8)
from_option_impl!(i8)
from_str_impl!(i16)
from_option_impl!(i16)
from_str_impl!(i32)
from_option_impl!(i32)
from_str_impl!(i64)
from_option_impl!(i64)
from_str_impl!(uint)
from_option_impl!(uint)
from_str_impl!(u8)
from_option_impl!(u8)
from_str_impl!(u16)
from_option_impl!(u16)
from_str_impl!(u32)
from_option_impl!(u32)
from_str_impl!(u64)
from_option_impl!(u64)
from_str_impl!(float)
from_option_impl!(float)
from_str_impl!(f32)
from_option_impl!(f32)
from_str_impl!(f64)
from_option_impl!(f64)

impl FromSql for Option<~str> {
    fn from_sql(raw: &Option<~[u8]>) -> Option<~str> {
        do raw.chain_ref |buf| {
            Some(str::from_bytes(buf.as_slice()))
        }
    }
}
from_option_impl!(~str)

pub trait ToSql {
    fn to_sql(&self) -> Option<~[u8]>;
}

macro_rules! to_str_impl(
    ($t:ty) => (
        impl ToSql for $t {
            fn to_sql(&self) -> Option<~[u8]> {
                Some(self.to_str().into_bytes())
            }
        }
    )
)

macro_rules! to_option_impl(
    ($t:ty) => (
        impl ToSql for Option<$t> {
            fn to_sql(&self) -> Option<~[u8]> {
                do self.chain |val| {
                    val.to_sql()
                }
            }
        }
    )
)

to_str_impl!(int)
to_option_impl!(int)
to_str_impl!(i8)
to_option_impl!(i8)
to_str_impl!(i16)
to_option_impl!(i16)
to_str_impl!(i32)
to_option_impl!(i32)
to_str_impl!(i64)
to_option_impl!(i64)
to_str_impl!(uint)
to_option_impl!(uint)
to_str_impl!(u8)
to_option_impl!(u8)
to_str_impl!(u16)
to_option_impl!(u16)
to_str_impl!(u32)
to_option_impl!(u32)
to_str_impl!(u64)
to_option_impl!(u64)
to_str_impl!(float)
to_option_impl!(float)
to_str_impl!(f32)
to_option_impl!(f32)
to_str_impl!(f64)
to_option_impl!(f64)

impl<'self> ToSql for &'self str {
    fn to_sql(&self) -> Option<~[u8]> {
        Some(self.as_bytes().to_owned())
    }
}

impl ToSql for Option<~str> {
    fn to_sql(&self) -> Option<~[u8]> {
        do self.chain_ref |val| {
            val.to_sql()
        }
    }
}

impl<'self> ToSql for Option<&'self str> {
    fn to_sql(&self) -> Option<~[u8]> {
        do self.chain |val| {
            val.to_sql()
        }
    }
}
<commit_msg>Closer-to-proper error handling for connect<commit_after>extern mod extra;

use extra::digest::Digest;
use extra::md5::Md5;
use extra::url::{UserInfo, Url};
use std::cell::Cell;
use std::rt::io::io_error;
use std::rt::io::net::ip::SocketAddr;
use std::rt::io::net::tcp::TcpStream;
use std::str;

use message::*;

mod message;

pub struct PostgresConnection {
    priv stream: Cell<TcpStream>,
    priv next_stmt_id: Cell<int>
}

impl Drop for PostgresConnection {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.write_message(&Terminate);
        }
    }
}

#[deriving(ToStr)]
pub struct PostgresError;

#[deriving(ToStr)]
pub enum PostgresConnectError {
    InvalidUrl,
    MissingUser,
    DbError(PostgresError),
    MissingPassword,
    UnsupportedAuthentication
}

impl PostgresConnection {
    pub fn connect(url: &str) -> PostgresConnection {
        match PostgresConnection::try_connect(url) {
            Ok(conn) => conn,
            Err(err) => fail!("Failed to connect: %s", err.to_str())
        }
    }

    pub fn try_connect(url: &str) -> Result<PostgresConnection,
                                            PostgresConnectError> {
        let Url {
            host,
            port,
            user,
            path,
            query: args,
            _
        }: Url = match FromStr::from_str(url) {
            Some(url) => url,
            None => return Err(InvalidUrl)
        };
        let user = match user {
            Some(user) => user,
            None => return Err(MissingUser)
        };
        let mut args = args;

        let socket_url = fmt!("%s:%s", host, port.unwrap_or_default(~"5432"));
        let addr: SocketAddr = match FromStr::from_str(socket_url) {
            Some(addr) => addr,
            None => return Err(InvalidUrl)
        };

        let conn = PostgresConnection {
            // Need to figure out what to do about unwrap here
            stream: Cell::new(TcpStream::connect(addr).unwrap()),
            next_stmt_id: Cell::new(0)
        };

        // we have to clone here since we need the user again for auth
        args.push((~"user", user.user.clone()));
        if !path.is_empty() {
            args.push((~"database", path));
        }
        conn.write_message(&StartupMessage(args.as_slice()));

        match conn.handle_auth(user) {
            Some(err) => return Err(err),
            None => ()
        }

        loop {
            match conn.read_message() {
                ParameterStatus(param, value) =>
                    info!("Parameter %s = %s", param, value),
                BackendKeyData(*) => (),
                ReadyForQuery(*) => break,
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }

        Ok(conn)
    }

    fn write_message(&self, message: &FrontendMessage) {
        do self.stream.with_mut_ref |s| {
            s.write_message(message);
        }
    }

    fn read_message(&self) -> BackendMessage {
        do self.stream.with_mut_ref |s| {
            s.read_message()
        }
    }

    fn handle_auth(&self, user: UserInfo) -> Option<PostgresConnectError> {
        match self.read_message() {
            AuthenticationOk => return None,
            AuthenticationCleartextPassword => {
                let pass = match user.pass {
                    Some(pass) => pass,
                    None => return Some(MissingPassword)
                };
                self.write_message(&PasswordMessage(pass));
            }
            AuthenticationMD5Password(nonce) => {
                let UserInfo { user, pass } = user;
                let pass = match pass {
                    Some(pass) => pass,
                    None => return Some(MissingPassword)
                };
                let input = pass + user;
                let mut md5 = Md5::new();
                md5.input_str(input);
                let output = md5.result_str();
                md5.reset();
                md5.input_str(output);
                md5.input(nonce);
                let output = "md5" + md5.result_str();
                self.write_message(&PasswordMessage(output.as_slice()));
            }
            resp => fail!("Bad response: %?", resp.to_str())
        }

        match self.read_message() {
            AuthenticationOk => None,
            ErrorResponse(*) => Some(DbError(PostgresError)),
            resp => fail!("Bad response: %?", resp.to_str())
        }
    }

    pub fn prepare<'a>(&'a self, query: &str) -> PostgresStatement<'a> {
        let id = self.next_stmt_id.take();
        let stmt_name = format!("statement_{}", id);
        self.next_stmt_id.put_back(id + 1);

        let types = [];
        self.write_message(&Parse(stmt_name, query, types));
        self.write_message(&Sync);

        match self.read_message() {
            ParseComplete => (),
            resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
            resp => fail!("Bad response: %?", resp.to_str())
        }

        self.wait_for_ready();

        self.write_message(&Describe('S' as u8, stmt_name));
        self.write_message(&Sync);

        let num_params = match self.read_message() {
            ParameterDescription(ref types) => types.len(),
            resp => fail!("Bad response: %?", resp.to_str())
        };

        match self.read_message() {
            RowDescription(*) | NoData => (),
            resp => fail!("Bad response: %?", resp.to_str())
        }

        self.wait_for_ready();

        PostgresStatement {
            conn: self,
            name: stmt_name,
            num_params: num_params,
            next_portal_id: Cell::new(0)
        }
    }

    pub fn in_transaction<T, E: ToStr>(&self, blk: &fn(&PostgresConnection)
                                                       -> Result<T, E>)
                                       -> Result<T, E> {
        self.quick_query("BEGIN");

        // If this fails, Postgres will rollback when the connection closes
        let ret = blk(self);

        if ret.is_ok() {
            self.quick_query("COMMIT");
        } else {
            self.quick_query("ROLLBACK");
        }

        ret
    }

    fn quick_query(&self, query: &str) {
        self.write_message(&Query(query));

        loop {
            match self.read_message() {
                ReadyForQuery(*) => break,
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                _ => ()
            }
        }
    }

    fn wait_for_ready(&self) {
        loop {
            match self.read_message() {
                ReadyForQuery(*) => break,
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }
    }
}

pub struct PostgresStatement<'self> {
    priv conn: &'self PostgresConnection,
    priv name: ~str,
    priv num_params: uint,
    priv next_portal_id: Cell<uint>
}

#[unsafe_destructor]
impl<'self> Drop for PostgresStatement<'self> {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.conn.write_message(&Close('S' as u8, self.name.as_slice()));
            self.conn.write_message(&Sync);
            loop {
                match self.conn.read_message() {
                    ReadyForQuery(*) => break,
                    _ => ()
                }
            }
        }
    }
}

impl<'self> PostgresStatement<'self> {
    pub fn num_params(&self) -> uint {
        self.num_params
    }

    fn execute(&self, portal_name: &str, params: &[&ToSql]) {
        if self.num_params != params.len() {
            fail!("Expected %u params but got %u", self.num_params,
                  params.len());
        }

        let formats = [];
        let values: ~[Option<~[u8]>] = params.iter().map(|val| val.to_sql())
                .collect();
        let result_formats = [];

        self.conn.write_message(&Bind(portal_name, self.name.as_slice(),
                                      formats, values, result_formats));
        self.conn.write_message(&Execute(portal_name.as_slice(), 0));
        self.conn.write_message(&Sync);

        match self.conn.read_message() {
            BindComplete => (),
            resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
            resp => fail!("Bad response: %?", resp.to_str())
        }
    }

    pub fn update(&self, params: &[&ToSql]) -> uint {
        // The unnamed portal is automatically cleaned up at sync time
        self.execute("", params);

        let mut num = 0;
        loop {
            match self.conn.read_message() {
                CommandComplete(ret) => {
                    let s = ret.split_iter(' ').last().unwrap();
                    match FromStr::from_str(s) {
                        None => (),
                        Some(n) => num = n
                    }
                    break;
                }
                DataRow(*) => (),
                EmptyQueryResponse => break,
                NoticeResponse(*) => (),
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }
        self.conn.wait_for_ready();

        num
    }

    pub fn query<'a>(&'a self, params: &[&ToSql]) -> PostgresResult<'a> {
        let id = self.next_portal_id.take();
        let portal_name = format!("{:s}_portal_{}", self.name.as_slice(), id);
        self.next_portal_id.put_back(id + 1);

        self.execute(portal_name, params);

        let mut data = ~[];
        loop {
            match self.conn.read_message() {
                EmptyQueryResponse => break,
                DataRow(row) => data.push(row),
                CommandComplete(*) => break,
                NoticeResponse(*) => (),
                resp @ ErrorResponse(*) => fail!("Error: %?", resp.to_str()),
                resp => fail!("Bad response: %?", resp.to_str())
            }
        }

        PostgresResult {
            stmt: self,
            name: portal_name,
            data: data
        }
    }
}

pub struct PostgresResult<'self> {
    priv stmt: &'self PostgresStatement<'self>,
    priv name: ~str,
    priv data: ~[~[Option<~[u8]>]]
}

#[unsafe_destructor]
impl<'self> Drop for PostgresResult<'self> {
    fn drop(&self) {
        do io_error::cond.trap(|_| {}).inside {
            self.stmt.conn.write_message(&Close('P' as u8,
                                                self.name.as_slice()));
            self.stmt.conn.write_message(&Sync);
            loop {
                match self.stmt.conn.read_message() {
                    ReadyForQuery(*) => break,
                    _ => ()
                }
            }
        }
    }
}

impl<'self> PostgresResult<'self> {
    pub fn iter<'a>(&'a self) -> PostgresResultIterator<'a> {
        PostgresResultIterator { result: self, next_row: 0 }
    }
}

pub struct PostgresResultIterator<'self> {
    priv result: &'self PostgresResult<'self>,
    priv next_row: uint
}

impl<'self> Iterator<PostgresRow<'self>> for PostgresResultIterator<'self> {
    fn next(&mut self) -> Option<PostgresRow<'self>> {
        if self.next_row == self.result.data.len() {
            return None;
        }

        let row = self.next_row;
        self.next_row += 1;
        Some(PostgresRow { result: self.result, row: row })
    }
}

pub struct PostgresRow<'self> {
    priv result: &'self PostgresResult<'self>,
    priv row: uint
}

impl<'self> Container for PostgresRow<'self> {
    fn len(&self) -> uint {
        self.result.data[self.row].len()
    }
}

impl<'self, T: FromSql> Index<uint, T> for PostgresRow<'self> {
    fn index(&self, idx: &uint) -> T {
        self.get(*idx)
    }
}

impl<'self> PostgresRow<'self> {
    pub fn get<T: FromSql>(&self, idx: uint) -> T {
        FromSql::from_sql(&self.result.data[self.row][idx])
    }
}

pub trait FromSql {
    fn from_sql(raw: &Option<~[u8]>) -> Self;
}

macro_rules! from_str_impl(
    ($t:ty) => (
        impl FromSql for Option<$t> {
            fn from_sql(raw: &Option<~[u8]>) -> Option<$t> {
                match *raw {
                    None => None,
                    Some(ref buf) => {
                        let s = str::from_bytes_slice(buf.as_slice());
                        Some(FromStr::from_str(s).unwrap())
                    }
                }
            }
        }
    )
)

macro_rules! from_option_impl(
    ($t:ty) => (
        impl FromSql for $t {
            fn from_sql(raw: &Option<~[u8]>) -> $t {
                FromSql::from_sql::<Option<$t>>(raw).unwrap()
            }
        }
    )
)

from_str_impl!(int)
from_option_impl!(int)
from_str_impl!(i8)
from_option_impl!(i8)
from_str_impl!(i16)
from_option_impl!(i16)
from_str_impl!(i32)
from_option_impl!(i32)
from_str_impl!(i64)
from_option_impl!(i64)
from_str_impl!(uint)
from_option_impl!(uint)
from_str_impl!(u8)
from_option_impl!(u8)
from_str_impl!(u16)
from_option_impl!(u16)
from_str_impl!(u32)
from_option_impl!(u32)
from_str_impl!(u64)
from_option_impl!(u64)
from_str_impl!(float)
from_option_impl!(float)
from_str_impl!(f32)
from_option_impl!(f32)
from_str_impl!(f64)
from_option_impl!(f64)

impl FromSql for Option<~str> {
    fn from_sql(raw: &Option<~[u8]>) -> Option<~str> {
        do raw.chain_ref |buf| {
            Some(str::from_bytes(buf.as_slice()))
        }
    }
}
from_option_impl!(~str)

pub trait ToSql {
    fn to_sql(&self) -> Option<~[u8]>;
}

macro_rules! to_str_impl(
    ($t:ty) => (
        impl ToSql for $t {
            fn to_sql(&self) -> Option<~[u8]> {
                Some(self.to_str().into_bytes())
            }
        }
    )
)

macro_rules! to_option_impl(
    ($t:ty) => (
        impl ToSql for Option<$t> {
            fn to_sql(&self) -> Option<~[u8]> {
                do self.chain |val| {
                    val.to_sql()
                }
            }
        }
    )
)

to_str_impl!(int)
to_option_impl!(int)
to_str_impl!(i8)
to_option_impl!(i8)
to_str_impl!(i16)
to_option_impl!(i16)
to_str_impl!(i32)
to_option_impl!(i32)
to_str_impl!(i64)
to_option_impl!(i64)
to_str_impl!(uint)
to_option_impl!(uint)
to_str_impl!(u8)
to_option_impl!(u8)
to_str_impl!(u16)
to_option_impl!(u16)
to_str_impl!(u32)
to_option_impl!(u32)
to_str_impl!(u64)
to_option_impl!(u64)
to_str_impl!(float)
to_option_impl!(float)
to_str_impl!(f32)
to_option_impl!(f32)
to_str_impl!(f64)
to_option_impl!(f64)

impl<'self> ToSql for &'self str {
    fn to_sql(&self) -> Option<~[u8]> {
        Some(self.as_bytes().to_owned())
    }
}

impl ToSql for Option<~str> {
    fn to_sql(&self) -> Option<~[u8]> {
        do self.chain_ref |val| {
            val.to_sql()
        }
    }
}

impl<'self> ToSql for Option<&'self str> {
    fn to_sql(&self) -> Option<~[u8]> {
        do self.chain |val| {
            val.to_sql()
        }
    }
}
<|endoftext|>