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

#![allow(dead_code)] // runtime init functions not used during testing

use os::windows::prelude::*;
use sys::windows::os::current_exe;
use sys::c;
use ffi::OsString;
use fmt;
use vec;
use core::iter;
use slice;
use path::PathBuf;

pub unsafe fn init(_argc: isize, _argv: *const *const u8) { }

pub unsafe fn cleanup() { }

pub fn args() -> Args {
    unsafe {
        let lp_cmd_line = c::GetCommandLineW();
        let parsed_args_list = parse_lp_cmd_line(
            lp_cmd_line as *const u16,
            || current_exe().map(PathBuf::into_os_string).unwrap_or_else(|_| OsString::new()));

        Args { parsed_args_list: parsed_args_list }
    }
}

/// Implements the Windows command-line argument parsing algorithm, described at
/// <https://docs.microsoft.com/en-us/previous-versions//17w5ykft(v=vs.85)>.
///
/// Windows includes a function to do this in shell32.dll,
/// but linking with that DLL causes the process to be registered as a GUI application.
/// GUI applications add a bunch of overhead, even if no windows are drawn. See
/// <https://randomascii.wordpress.com/2018/12/03/a-not-called-function-can-cause-a-5x-slowdown/>.
unsafe fn parse_lp_cmd_line<F: Fn() -> OsString>(lp_cmd_line: *const u16, exe_name: F)
                                                 -> vec::IntoIter<OsString> {
    const BACKSLASH: u16 = '\\' as u16;
    const QUOTE: u16 = '"' as u16;
    const TAB: u16 = '\t' as u16;
    const SPACE: u16 = ' ' as u16;
    let mut in_quotes = false;
    let mut was_in_quotes = false;
    let mut backslash_count: usize = 0;
    let mut ret_val = Vec::new();
    let mut cur = Vec::new();
    if lp_cmd_line.is_null() || *lp_cmd_line == 0 {
        ret_val.push(exe_name());
        return ret_val.into_iter();
    }
    let mut i = 0;
    // The executable name at the beginning is special.
    match *lp_cmd_line {
        // The executable name ends at the next quote mark,
        // no matter what.
        QUOTE => {
            loop {
                i += 1;
                if *lp_cmd_line.offset(i) == 0 {
                    ret_val.push(OsString::from_wide(
                        slice::from_raw_parts(lp_cmd_line.offset(1), i as usize - 1)
                    ));
                    return ret_val.into_iter();
                }
                if *lp_cmd_line.offset(i) == QUOTE {
                    break;
                }
            }
            ret_val.push(OsString::from_wide(
                slice::from_raw_parts(lp_cmd_line.offset(1), i as usize - 1)
            ));
            i += 1;
        }
        // Implement quirk: when they say whitespace here,
        // they include the entire ASCII control plane:
        // "However, if lpCmdLine starts with any amount of whitespace, CommandLineToArgvW
        // will consider the first argument to be an empty string. Excess whitespace at the
        // end of lpCmdLine is ignored."
        0...SPACE => {
            ret_val.push(OsString::new());
            i += 1;
        },
        // The executable name ends at the next whitespace,
        // no matter what.
        _ => {
            loop {
                i += 1;
                if *lp_cmd_line.offset(i) == 0 {
                    ret_val.push(OsString::from_wide(
                        slice::from_raw_parts(lp_cmd_line, i as usize)
                    ));
                    return ret_val.into_iter();
                }
                if let 0...SPACE = *lp_cmd_line.offset(i) {
                    break;
                }
            }
            ret_val.push(OsString::from_wide(
                slice::from_raw_parts(lp_cmd_line, i as usize)
            ));
            i += 1;
        }
    }
    loop {
        let c = *lp_cmd_line.offset(i);
        match c {
            // backslash
            BACKSLASH => {
                backslash_count += 1;
                was_in_quotes = false;
            },
            QUOTE if backslash_count % 2 == 0 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count / 2));
                backslash_count = 0;
                if was_in_quotes {
                    cur.push('"' as u16);
                    was_in_quotes = false;
                } else {
                    was_in_quotes = in_quotes;
                    in_quotes = !in_quotes;
                }
            }
            QUOTE if backslash_count % 2 != 0 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count / 2));
                backslash_count = 0;
                was_in_quotes = false;
                cur.push(b'"' as u16);
            }
            SPACE | TAB if !in_quotes => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                if !cur.is_empty() || was_in_quotes {
                    ret_val.push(OsString::from_wide(&cur[..]));
                    cur.truncate(0);
                }
                backslash_count = 0;
                was_in_quotes = false;
            }
            0x00 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                // include empty quoted strings at the end of the arguments list
                if !cur.is_empty() || was_in_quotes || in_quotes {
                    ret_val.push(OsString::from_wide(&cur[..]));
                }
                break;
            }
            _ => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                backslash_count = 0;
                was_in_quotes = false;
                cur.push(c);
            }
        }
        i += 1;
    }
    ret_val.into_iter()
}

pub struct Args {
    parsed_args_list: vec::IntoIter<OsString>,
}

pub struct ArgsInnerDebug<'a> {
    args: &'a Args,
}

impl<'a> fmt::Debug for ArgsInnerDebug<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("[")?;
        let mut first = true;
        for i in self.args.parsed_args_list.clone() {
            if !first {
                f.write_str(", ")?;
            }
            first = false;

            fmt::Debug::fmt(&i, f)?;
        }
        f.write_str("]")?;
        Ok(())
    }
}

impl Args {
    pub fn inner_debug(&self) -> ArgsInnerDebug {
        ArgsInnerDebug {
            args: self
        }
    }
}

impl Iterator for Args {
    type Item = OsString;
    fn next(&mut self) -> Option<OsString> { self.parsed_args_list.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.parsed_args_list.size_hint() }
}

impl DoubleEndedIterator for Args {
    fn next_back(&mut self) -> Option<OsString> { self.parsed_args_list.next_back() }
}

impl ExactSizeIterator for Args {
    fn len(&self) -> usize { self.parsed_args_list.len() }
}

#[cfg(test)]
mod tests {
    use sys::windows::args::*;
    use ffi::OsString;

    fn chk(string: &str, parts: &[&str]) {
        let mut wide: Vec<u16> = OsString::from(string).encode_wide().collect();
        wide.push(0);
        let parsed = unsafe {
            parse_lp_cmd_line(wide.as_ptr() as *const u16, || OsString::from("TEST.EXE"))
        };
        let expected: Vec<OsString> = parts.iter().map(|k| OsString::from(k)).collect();
        assert_eq!(parsed.as_slice(), expected.as_slice());
    }

    #[test]
    fn empty() {
        chk("", &["TEST.EXE"]);
        chk("\0", &["TEST.EXE"]);
    }

    #[test]
    fn single_words() {
        chk("EXE one_word", &["EXE", "one_word"]);
        chk("EXE a", &["EXE", "a"]);
        chk("EXE 😅", &["EXE", "😅"]);
        chk("EXE 😅🤦", &["EXE", "😅🤦"]);
    }

    #[test]
    fn official_examples() {
        chk(r#"EXE "abc" d e"#, &["EXE", "abc", "d", "e"]);
        chk(r#"EXE a\\\b d"e f"g h"#, &["EXE", r#"a\\\b"#, "de fg", "h"]);
        chk(r#"EXE a\\\"b c d"#, &["EXE", r#"a\"b"#, "c", "d"]);
        chk(r#"EXE a\\\\"b c" d e"#, &["EXE", r#"a\\b c"#, "d", "e"]);
    }

    #[test]
    fn whitespace_behavior() {
        chk(r#" test"#, &["", "test"]);
        chk(r#"  test"#, &["", "test"]);
        chk(r#" test test2"#, &["", "test", "test2"]);
        chk(r#" test  test2"#, &["", "test", "test2"]);
        chk(r#"test test2 "#, &["test", "test2"]);
        chk(r#"test  test2 "#, &["test", "test2"]);
        chk(r#"test "#, &["test"]);
    }

    #[test]
    fn genius_quotes() {
        chk(r#"EXE "" """#, &["EXE", "", ""]);
        chk(r#"EXE "" """"#, &["EXE", "", "\""]);
        chk(
            r#"EXE "this is """all""" in the same argument""#,
            &["EXE", "this is \"all\" in the same argument"]
        );
        chk(r#"EXE "a"""#, &["EXE", "a\""]);
        chk(r#"EXE "a"" a"#, &["EXE", "a\"", "a"]);
    }
}
<commit_msg>Fix nits<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(dead_code)] // runtime init functions not used during testing

use os::windows::prelude::*;
use sys::windows::os::current_exe;
use sys::c;
use ffi::OsString;
use fmt;
use vec;
use core::iter;
use slice;
use path::PathBuf;

pub unsafe fn init(_argc: isize, _argv: *const *const u8) { }

pub unsafe fn cleanup() { }

pub fn args() -> Args {
    unsafe {
        let lp_cmd_line = c::GetCommandLineW();
        let parsed_args_list = parse_lp_cmd_line(
            lp_cmd_line as *const u16,
            || current_exe().map(PathBuf::into_os_string).unwrap_or_else(|_| OsString::new()));

        Args { parsed_args_list: parsed_args_list }
    }
}

/// Implements the Windows command-line argument parsing algorithm.
///
/// Microsoft's documentation for the Windows CLI argument format can be found at
/// <https://docs.microsoft.com/en-us/previous-versions//17w5ykft(v=vs.85)>.
///
/// Windows includes a function to do this in shell32.dll,
/// but linking with that DLL causes the process to be registered as a GUI application.
/// GUI applications add a bunch of overhead, even if no windows are drawn. See
/// <https://randomascii.wordpress.com/2018/12/03/a-not-called-function-can-cause-a-5x-slowdown/>.
///
/// This function was tested for equivalence to the shell32.dll implementation in
/// Windows 10 Pro v1803, using an exhaustive test suite available at
/// <https://gist.github.com/notriddle/dde431930c392e428055b2dc22e638f5> or
/// <https://paste.gg/p/anonymous/47d6ed5f5bd549168b1c69c799825223>.
unsafe fn parse_lp_cmd_line<F: Fn() -> OsString>(lp_cmd_line: *const u16, exe_name: F)
                                                 -> vec::IntoIter<OsString> {
    const BACKSLASH: u16 = '\\' as u16;
    const QUOTE: u16 = '"' as u16;
    const TAB: u16 = '\t' as u16;
    const SPACE: u16 = ' ' as u16;
    let mut in_quotes = false;
    let mut was_in_quotes = false;
    let mut backslash_count: usize = 0;
    let mut ret_val = Vec::new();
    let mut cur = Vec::new();
    if lp_cmd_line.is_null() || *lp_cmd_line == 0 {
        ret_val.push(exe_name());
        return ret_val.into_iter();
    }
    let mut i = 0;
    // The executable name at the beginning is special.
    match *lp_cmd_line {
        // The executable name ends at the next quote mark,
        // no matter what.
        QUOTE => {
            loop {
                i += 1;
                if *lp_cmd_line.offset(i) == 0 {
                    ret_val.push(OsString::from_wide(
                        slice::from_raw_parts(lp_cmd_line.offset(1), i as usize - 1)
                    ));
                    return ret_val.into_iter();
                }
                if *lp_cmd_line.offset(i) == QUOTE {
                    break;
                }
            }
            ret_val.push(OsString::from_wide(
                slice::from_raw_parts(lp_cmd_line.offset(1), i as usize - 1)
            ));
            i += 1;
        }
        // Implement quirk: when they say whitespace here,
        // they include the entire ASCII control plane:
        // "However, if lpCmdLine starts with any amount of whitespace, CommandLineToArgvW
        // will consider the first argument to be an empty string. Excess whitespace at the
        // end of lpCmdLine is ignored."
        0...SPACE => {
            ret_val.push(OsString::new());
            i += 1;
        },
        // The executable name ends at the next whitespace,
        // no matter what.
        _ => {
            loop {
                i += 1;
                if *lp_cmd_line.offset(i) == 0 {
                    ret_val.push(OsString::from_wide(
                        slice::from_raw_parts(lp_cmd_line, i as usize)
                    ));
                    return ret_val.into_iter();
                }
                if let 0...SPACE = *lp_cmd_line.offset(i) {
                    break;
                }
            }
            ret_val.push(OsString::from_wide(
                slice::from_raw_parts(lp_cmd_line, i as usize)
            ));
            i += 1;
        }
    }
    loop {
        let c = *lp_cmd_line.offset(i);
        match c {
            // backslash
            BACKSLASH => {
                backslash_count += 1;
                was_in_quotes = false;
            },
            QUOTE if backslash_count % 2 == 0 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count / 2));
                backslash_count = 0;
                if was_in_quotes {
                    cur.push('"' as u16);
                    was_in_quotes = false;
                } else {
                    was_in_quotes = in_quotes;
                    in_quotes = !in_quotes;
                }
            }
            QUOTE if backslash_count % 2 != 0 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count / 2));
                backslash_count = 0;
                was_in_quotes = false;
                cur.push(b'"' as u16);
            }
            SPACE | TAB if !in_quotes => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                if !cur.is_empty() || was_in_quotes {
                    ret_val.push(OsString::from_wide(&cur[..]));
                    cur.truncate(0);
                }
                backslash_count = 0;
                was_in_quotes = false;
            }
            0x00 => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                // include empty quoted strings at the end of the arguments list
                if !cur.is_empty() || was_in_quotes || in_quotes {
                    ret_val.push(OsString::from_wide(&cur[..]));
                }
                break;
            }
            _ => {
                cur.extend(iter::repeat(b'\\' as u16).take(backslash_count));
                backslash_count = 0;
                was_in_quotes = false;
                cur.push(c);
            }
        }
        i += 1;
    }
    ret_val.into_iter()
}

pub struct Args {
    parsed_args_list: vec::IntoIter<OsString>,
}

pub struct ArgsInnerDebug<'a> {
    args: &'a Args,
}

impl<'a> fmt::Debug for ArgsInnerDebug<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str("[")?;
        let mut first = true;
        for i in self.args.parsed_args_list.as_slice() {
            if !first {
                f.write_str(", ")?;
            }
            first = false;

            fmt::Debug::fmt(i, f)?;
        }
        f.write_str("]")?;
        Ok(())
    }
}

impl Args {
    pub fn inner_debug(&self) -> ArgsInnerDebug {
        ArgsInnerDebug {
            args: self
        }
    }
}

impl Iterator for Args {
    type Item = OsString;
    fn next(&mut self) -> Option<OsString> { self.parsed_args_list.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.parsed_args_list.size_hint() }
}

impl DoubleEndedIterator for Args {
    fn next_back(&mut self) -> Option<OsString> { self.parsed_args_list.next_back() }
}

impl ExactSizeIterator for Args {
    fn len(&self) -> usize { self.parsed_args_list.len() }
}

#[cfg(test)]
mod tests {
    use sys::windows::args::*;
    use ffi::OsString;

    fn chk(string: &str, parts: &[&str]) {
        let mut wide: Vec<u16> = OsString::from(string).encode_wide().collect();
        wide.push(0);
        let parsed = unsafe {
            parse_lp_cmd_line(wide.as_ptr() as *const u16, || OsString::from("TEST.EXE"))
        };
        let expected: Vec<OsString> = parts.iter().map(|k| OsString::from(k)).collect();
        assert_eq!(parsed.as_slice(), expected.as_slice());
    }

    #[test]
    fn empty() {
        chk("", &["TEST.EXE"]);
        chk("\0", &["TEST.EXE"]);
    }

    #[test]
    fn single_words() {
        chk("EXE one_word", &["EXE", "one_word"]);
        chk("EXE a", &["EXE", "a"]);
        chk("EXE 😅", &["EXE", "😅"]);
        chk("EXE 😅🤦", &["EXE", "😅🤦"]);
    }

    #[test]
    fn official_examples() {
        chk(r#"EXE "abc" d e"#, &["EXE", "abc", "d", "e"]);
        chk(r#"EXE a\\\b d"e f"g h"#, &["EXE", r#"a\\\b"#, "de fg", "h"]);
        chk(r#"EXE a\\\"b c d"#, &["EXE", r#"a\"b"#, "c", "d"]);
        chk(r#"EXE a\\\\"b c" d e"#, &["EXE", r#"a\\b c"#, "d", "e"]);
    }

    #[test]
    fn whitespace_behavior() {
        chk(r#" test"#, &["", "test"]);
        chk(r#"  test"#, &["", "test"]);
        chk(r#" test test2"#, &["", "test", "test2"]);
        chk(r#" test  test2"#, &["", "test", "test2"]);
        chk(r#"test test2 "#, &["test", "test2"]);
        chk(r#"test  test2 "#, &["test", "test2"]);
        chk(r#"test "#, &["test"]);
    }

    #[test]
    fn genius_quotes() {
        chk(r#"EXE "" """#, &["EXE", "", ""]);
        chk(r#"EXE "" """"#, &["EXE", "", "\""]);
        chk(
            r#"EXE "this is """all""" in the same argument""#,
            &["EXE", "this is \"all\" in the same argument"]
        );
        chk(r#"EXE "a"""#, &["EXE", "a\""]);
        chk(r#"EXE "a"" a"#, &["EXE", "a\"", "a"]);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fs::{File, OpenOptions, remove_file};
use std::io::{Read, Write};

use rustc::ty::layout::Size;

use crate::stacked_borrows::Tag;
use crate::*;

pub struct FileHandle {
    file: File,
    flag: i32,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        // The first two bits of the flag correspond to the access mode of the file in linux. This
        // is done this way because `O_RDONLY` is zero in several platforms.
        let access_mode = flag & 0b11;

        if access_mode == this.eval_libc_i32("O_RDONLY")? {
            options.read(true);
        } else if access_mode == this.eval_libc_i32("O_WRONLY")? {
            options.write(true);
        } else if access_mode == this.eval_libc_i32("O_RDWR")? {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // This flag is a noop for now because `std` already sets it.
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag);
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file, flag });
            fh.low
        });

        this.consume_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`fcntl` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor
        if cmd == this.eval_libc_i32("F_GETFD")? {
            this.get_handle_and(fd, |handle| Ok(handle.flag))
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`close` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        this.remove_handle_and(fd, |handle, this| {
            this.consume_result(handle.file.sync_all().map(|_| 0i32))
        })
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`read` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Reading zero bytes should not change `buf`
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf_scalar = this.read_scalar(buf_op)?.not_undef()?;

        // Remove the file handle to avoid borrowing issues
        this.remove_handle_and(fd, |mut handle, this| {
            // Don't use `?` to avoid returning before reinserting the handle
            let bytes = this.force_ptr(buf_scalar).and_then(|buf| {
                this.memory_mut()
                    .get_mut(buf.alloc_id)?
                    .get_bytes_mut(tcx, buf, Size::from_bytes(count))
                    .map(|buffer| handle.file.read(buffer))
            });
            // Reinsert the file handle
            this.machine.file_handler.handles.insert(fd, handle);
            this.consume_result(bytes?.map(|bytes| bytes as i64))
        })
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`write` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Writing zero bytes should not change `buf`
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;

        this.remove_handle_and(fd, |mut handle, this| {
            let bytes = this.memory().get(buf.alloc_id).and_then(|alloc| {
                alloc
                    .get_bytes(tcx, buf, Size::from_bytes(count))
                    .map(|bytes| handle.file.write(bytes).map(|bytes| bytes as i64))
            });
            this.machine.file_handler.handles.insert(fd, handle);
            this.consume_result(bytes?)
        })
    }

    fn unlink( &mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`write` not available when isolation is enabled")
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        let result = remove_file(path).map(|_| 0);

        this.consume_result(result)
    }

    /// Helper function that gets a `FileHandle` immutable reference and allows to manipulate it
    /// using the `f` closure.
    ///
    /// If the `fd` file descriptor does not correspond to a file, this functions returns `Ok(-1)`
    /// and sets `Evaluator::last_error` to `libc::EBADF` (invalid file descriptor).
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn get_handle_and<F, T: From<i32>>(&mut self, fd: i32, f: F) -> InterpResult<'tcx, T>
    where
        F: Fn(&FileHandle) -> InterpResult<'tcx, T>,
    {
        let this = self.eval_context_mut();
        if let Some(handle) = this.machine.file_handler.handles.get(&fd) {
            f(handle)
        } else {
            let ebadf = this.eval_libc("EBADF")?;
            this.set_last_error(ebadf)?;
            Ok((-1).into())
        }
    }

    /// Helper function that removes a `FileHandle` and allows to manipulate it using the `f`
    /// closure. This function is quite useful when you need to modify a `FileHandle` but you need
    /// to modify `MiriEvalContext` at the same time, so you can modify the handle and reinsert it
    /// using `f`.
    ///
    /// If the `fd` file descriptor does not correspond to a file, this functions returns `Ok(-1)`
    /// and sets `Evaluator::last_error` to `libc::EBADF` (invalid file descriptor).
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn remove_handle_and<F, T: From<i32>>(&mut self, fd: i32, mut f: F) -> InterpResult<'tcx, T>
    where
        F: FnMut(FileHandle, &mut MiriEvalContext<'mir, 'tcx>) -> InterpResult<'tcx, T>,
    {
        let this = self.eval_context_mut();
        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            f(handle, this)
        } else {
            let ebadf = this.eval_libc("EBADF")?;
            this.set_last_error(ebadf)?;
            Ok((-1).into())
        }
    }

    /// Helper function that consumes an `std::io::Result<T>` and returns an
    /// `InterpResult<'tcx,T>::Ok` instead. It is expected that the result can be converted to an
    /// OS error using `std::io::Error::raw_os_error`.
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn consume_result<T: From<i32>>(
        &mut self,
        result: std::io::Result<T>,
    ) -> InterpResult<'tcx, T> {
        match result {
            Ok(ok) => Ok(ok),
            Err(e) => {
                self.eval_context_mut().consume_io_error(e)?;
                Ok((-1).into())
            }
        }
    }
}
<commit_msg>Correct fcntl behavior<commit_after>use std::collections::HashMap;
use std::fs::{File, OpenOptions, remove_file};
use std::io::{Read, Write};

use rustc::ty::layout::Size;

use crate::stacked_borrows::Tag;
use crate::*;

pub struct FileHandle {
    file: File,
}

pub struct FileHandler {
    handles: HashMap<i32, FileHandle>,
    low: i32,
}

impl Default for FileHandler {
    fn default() -> Self {
        FileHandler {
            handles: Default::default(),
            // 0, 1 and 2 are reserved for stdin, stdout and stderr
            low: 3,
        }
    }
}

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    fn open(
        &mut self,
        path_op: OpTy<'tcx, Tag>,
        flag_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`open` not available when isolation is enabled")
        }

        let flag = this.read_scalar(flag_op)?.to_i32()?;

        let mut options = OpenOptions::new();

        // The first two bits of the flag correspond to the access mode of the file in linux. This
        // is done this way because `O_RDONLY` is zero in several platforms.
        let access_mode = flag & 0b11;

        if access_mode == this.eval_libc_i32("O_RDONLY")? {
            options.read(true);
        } else if access_mode == this.eval_libc_i32("O_WRONLY")? {
            options.write(true);
        } else if access_mode == this.eval_libc_i32("O_RDWR")? {
            options.read(true).write(true);
        } else {
            throw_unsup_format!("Unsupported access mode {:#x}", access_mode);
        }
        // We need to check that there aren't unsupported options in `flag`. For this we try to
        // reproduce the content of `flag` in the `mirror` variable using only the supported
        // options.
        let mut mirror = access_mode;

        let o_append = this.eval_libc_i32("O_APPEND")?;
        if flag & o_append != 0 {
            options.append(true);
            mirror |= o_append;
        }
        let o_trunc = this.eval_libc_i32("O_TRUNC")?;
        if flag & o_trunc != 0 {
            options.truncate(true);
            mirror |= o_trunc;
        }
        let o_creat = this.eval_libc_i32("O_CREAT")?;
        if flag & o_creat != 0 {
            options.create(true);
            mirror |= o_creat;
        }
        let o_cloexec = this.eval_libc_i32("O_CLOEXEC")?;
        if flag & o_cloexec != 0 {
            // We do not need to do anything for this flag because `std` already sets it.
            // (Technically we do not support *not* setting this flag, but we ignore that.)
            mirror |= o_cloexec;
        }
        // If `flag` is not equal to `mirror`, there is an unsupported option enabled in `flag`,
        // then we throw an error.
        if flag != mirror {
            throw_unsup_format!("unsupported flags {:#x}", flag & !mirror);
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        let fd = options.open(path).map(|file| {
            let mut fh = &mut this.machine.file_handler;
            fh.low += 1;
            fh.handles.insert(fh.low, FileHandle { file });
            fh.low
        });

        this.consume_result(fd)
    }

    fn fcntl(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        cmd_op: OpTy<'tcx, Tag>,
        _arg1_op: Option<OpTy<'tcx, Tag>>,
    ) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`fcntl` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let cmd = this.read_scalar(cmd_op)?.to_i32()?;
        // We only support getting the flags for a descriptor
        if cmd == this.eval_libc_i32("F_GETFD")? {
            // Currently this is the only flag that `F_GETFD` returns. It is OK to just return the
            // `FD_CLOEXEC` value without checking if the flag is set for the file because `std`
            // always sets this flag when opening a file. However we still need to check that the
            // file itself is open.
            this.get_handle_and(fd, |_| Ok(0))?;
            this.eval_libc_i32("FD_CLOEXEC")
        } else {
            throw_unsup_format!("The {:#x} command is not supported for `fcntl`)", cmd);
        }
    }

    fn close(&mut self, fd_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`close` not available when isolation is enabled")
        }

        let fd = this.read_scalar(fd_op)?.to_i32()?;

        this.remove_handle_and(fd, |handle, this| {
            this.consume_result(handle.file.sync_all().map(|_| 0i32))
        })
    }

    fn read(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`read` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Reading zero bytes should not change `buf`
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf_scalar = this.read_scalar(buf_op)?.not_undef()?;

        // Remove the file handle to avoid borrowing issues
        this.remove_handle_and(fd, |mut handle, this| {
            // Don't use `?` to avoid returning before reinserting the handle
            let bytes = this.force_ptr(buf_scalar).and_then(|buf| {
                this.memory_mut()
                    .get_mut(buf.alloc_id)?
                    .get_bytes_mut(tcx, buf, Size::from_bytes(count))
                    .map(|buffer| handle.file.read(buffer))
            });
            // Reinsert the file handle
            this.machine.file_handler.handles.insert(fd, handle);
            this.consume_result(bytes?.map(|bytes| bytes as i64))
        })
    }

    fn write(
        &mut self,
        fd_op: OpTy<'tcx, Tag>,
        buf_op: OpTy<'tcx, Tag>,
        count_op: OpTy<'tcx, Tag>,
    ) -> InterpResult<'tcx, i64> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`write` not available when isolation is enabled")
        }

        let tcx = &{ this.tcx.tcx };

        let count = this.read_scalar(count_op)?.to_usize(&*this.tcx)?;
        // Writing zero bytes should not change `buf`
        if count == 0 {
            return Ok(0);
        }
        let fd = this.read_scalar(fd_op)?.to_i32()?;
        let buf = this.force_ptr(this.read_scalar(buf_op)?.not_undef()?)?;

        this.remove_handle_and(fd, |mut handle, this| {
            let bytes = this.memory().get(buf.alloc_id).and_then(|alloc| {
                alloc
                    .get_bytes(tcx, buf, Size::from_bytes(count))
                    .map(|bytes| handle.file.write(bytes).map(|bytes| bytes as i64))
            });
            this.machine.file_handler.handles.insert(fd, handle);
            this.consume_result(bytes?)
        })
    }

    fn unlink( &mut self, path_op: OpTy<'tcx, Tag>) -> InterpResult<'tcx, i32> {
        let this = self.eval_context_mut();

        if !this.machine.communicate {
            throw_unsup_format!("`write` not available when isolation is enabled")
        }

        let path_bytes = this
            .memory()
            .read_c_str(this.read_scalar(path_op)?.not_undef()?)?;
        let path = std::str::from_utf8(path_bytes)
            .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", path_bytes))?;

        let result = remove_file(path).map(|_| 0);

        this.consume_result(result)
    }

    /// Helper function that gets a `FileHandle` immutable reference and allows to manipulate it
    /// using the `f` closure.
    ///
    /// If the `fd` file descriptor does not correspond to a file, this functions returns `Ok(-1)`
    /// and sets `Evaluator::last_error` to `libc::EBADF` (invalid file descriptor).
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn get_handle_and<F, T: From<i32>>(&mut self, fd: i32, f: F) -> InterpResult<'tcx, T>
    where
        F: Fn(&FileHandle) -> InterpResult<'tcx, T>,
    {
        let this = self.eval_context_mut();
        if let Some(handle) = this.machine.file_handler.handles.get(&fd) {
            f(handle)
        } else {
            let ebadf = this.eval_libc("EBADF")?;
            this.set_last_error(ebadf)?;
            Ok((-1).into())
        }
    }

    /// Helper function that removes a `FileHandle` and allows to manipulate it using the `f`
    /// closure. This function is quite useful when you need to modify a `FileHandle` but you need
    /// to modify `MiriEvalContext` at the same time, so you can modify the handle and reinsert it
    /// using `f`.
    ///
    /// If the `fd` file descriptor does not correspond to a file, this functions returns `Ok(-1)`
    /// and sets `Evaluator::last_error` to `libc::EBADF` (invalid file descriptor).
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn remove_handle_and<F, T: From<i32>>(&mut self, fd: i32, mut f: F) -> InterpResult<'tcx, T>
    where
        F: FnMut(FileHandle, &mut MiriEvalContext<'mir, 'tcx>) -> InterpResult<'tcx, T>,
    {
        let this = self.eval_context_mut();
        if let Some(handle) = this.machine.file_handler.handles.remove(&fd) {
            f(handle, this)
        } else {
            let ebadf = this.eval_libc("EBADF")?;
            this.set_last_error(ebadf)?;
            Ok((-1).into())
        }
    }

    /// Helper function that consumes an `std::io::Result<T>` and returns an
    /// `InterpResult<'tcx,T>::Ok` instead. It is expected that the result can be converted to an
    /// OS error using `std::io::Error::raw_os_error`.
    ///
    /// This function uses `T: From<i32>` instead of `i32` directly because some IO related
    /// functions return different integer types (like `read`, that returns an `i64`)
    fn consume_result<T: From<i32>>(
        &mut self,
        result: std::io::Result<T>,
    ) -> InterpResult<'tcx, T> {
        match result {
            Ok(ok) => Ok(ok),
            Err(e) => {
                self.eval_context_mut().consume_io_error(e)?;
                Ok((-1).into())
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, WriteBatch, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' | b'x' => {
            // Sequence of two byte document ids
            // d = directory
            // x = deletion list

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        b's' => {
            // Statistic
            // An i64 number that can be incremented or decremented
            let mut value = match existing_val {
                Some(existing_val) => BigEndian::read_i64(existing_val),
                None => 0
            };

            for op in operands {
                value += BigEndian::read_i64(op);
            }

            let mut buf = [0; 8];
            BigEndian::write_i64(&mut buf, value);
            buf.iter().cloned().collect()
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
    doc_key_mapping: RwLock<BTreeMap<Vec<u8>, DocRef>>,
    deleted_docs: RwLock<Vec<DocRef>>,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Start write batch
        let mut write_batch = WriteBatch::default();

        // Insert contents
        let mut token_count: i64 = 0;
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                token_count += 1;
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                write_batch.merge(&key, &doc_id_bytes);
            }
        }

        // Increment total docs
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_docs".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, 1);
        write_batch.merge(&key, &inc_bytes);

        // Increment total tokens
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_tokens".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, token_count);
        write_batch.merge(&key, &inc_bytes);

        // Write document data
        self.db.write(write_batch);

        // Update doc_key_mapping
        let mut write_batch = WriteBatch::default();
        let previous_doc_ref = self.doc_key_mapping.write().unwrap().insert(doc.key.as_bytes().iter().cloned().collect(), doc_ref);

        let mut key = Vec::with_capacity(1 + doc.key.len());
        key.push(b'k');
        for byte in doc.key.as_bytes() {
            key.push(*byte);
        }
        let mut doc_ref_bytes = [0; 6];
        BigEndian::write_u32(&mut doc_ref_bytes, doc_ref.chunk());
        BigEndian::write_u16(&mut doc_ref_bytes[4..], doc_ref.ord());
        write_batch.put(&key, &doc_ref_bytes);

        // If there was a document there previously, mark it as deleted
        if let Some(previous_doc_ref) = previous_doc_ref {
            self.deleted_docs.write().unwrap().push(previous_doc_ref);

            let mut key = Vec::with_capacity(5);
            key.push(b'x');
            for byte in previous_doc_ref.chunk().to_string().as_bytes() {
                key.push(*byte);
            }

            let mut previous_doc_id_bytes = [0; 2];
            BigEndian::write_u16(&mut previous_doc_id_bytes, previous_doc_ref.ord());
            write_batch.merge(&key, &previous_doc_id_bytes);
        }

        // Write document data
        self.db.write(write_batch);
    }
}

#[cfg(test)]
mod tests {
    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    fn clean_test_indices() {
        use std::fs::remove_dir_all;

        remove_dir_all("test_indices");
    }

    #[test]
    fn test_create() {
        clean_test_indices();

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        clean_test_indices();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "another_test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        clean_test_indices();

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<commit_msg>Make each test only remove its own directory<commit_after>#![feature(integer_atomics)]

extern crate abra;
extern crate rocksdb;
extern crate rustc_serialize;
#[macro_use]
extern crate maplit;
extern crate byteorder;

use std::sync::{Arc, RwLock};
use std::sync::atomic::{AtomicU32, Ordering};
use std::collections::BTreeMap;

use rocksdb::{DB, WriteBatch, Writable, Options, MergeOperands};
use abra::{Term, Document};
use abra::schema::{Schema, FieldType, FieldRef, AddFieldError};
use rustc_serialize::{json, Encodable};
use byteorder::{ByteOrder, BigEndian};

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct TermRef(u32);

impl TermRef {
    pub fn ord(&self) -> u32 {
        self.0
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
struct DocRef(u32, u16);

impl DocRef {
    pub fn chunk(&self) -> u32 {
        self.0
    }

    pub fn ord(&self) -> u16 {
        self.1
    }
}

fn merge_keys(key: &[u8], existing_val: Option<&[u8]>, operands: &mut MergeOperands) -> Vec<u8> {
    match key[0] {
        b'd' | b'x' => {
            // Sequence of two byte document ids
            // d = directory
            // x = deletion list

            // Allocate vec for new Value
            let new_size = match existing_val {
                Some(existing_val) => existing_val.len(),
                None => 0,
            } + operands.size_hint().0 * 2;

            let mut new_val = Vec::with_capacity(new_size);

            // Push existing value
            existing_val.map(|v| {
                for b in v {
                    new_val.push(*b);
                }
            });

            // Append new entries
            for op in operands {
                for b in op {
                    new_val.push(*b);
                }
            }

            new_val
        }
        b's' => {
            // Statistic
            // An i64 number that can be incremented or decremented
            let mut value = match existing_val {
                Some(existing_val) => BigEndian::read_i64(existing_val),
                None => 0
            };

            for op in operands {
                value += BigEndian::read_i64(op);
            }

            let mut buf = [0; 8];
            BigEndian::write_i64(&mut buf, value);
            buf.iter().cloned().collect()
        }
        _ => {
            // Unrecognised key, fallback to emulating a put operation (by taking the last value)
            operands.last().unwrap().iter().cloned().collect()
        }
    }
}

pub struct RocksDBIndexStore {
    schema: Arc<Schema>,
    db: DB,
    next_term_ref: AtomicU32,
    term_dictionary: RwLock<BTreeMap<Vec<u8>, TermRef>>,
    next_chunk: AtomicU32,
    doc_key_mapping: RwLock<BTreeMap<Vec<u8>, DocRef>>,
    deleted_docs: RwLock<Vec<DocRef>>,
}

impl RocksDBIndexStore {
    pub fn create(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        opts.create_if_missing(true);
        let db = try!(DB::open(&opts, path));

        // Schema
        let schema = Schema::new();
        db.put(b".schema", json::encode(&schema).unwrap().as_bytes());

        // Next term ref
        db.put(b".next_term_ref", b"1");

        // Next chunk
        db.put(b".next_chunk", b"1");

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(1),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(1),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn open(path: &str) -> Result<RocksDBIndexStore, String> {
        let mut opts = Options::default();
        opts.add_merge_operator("merge operator", merge_keys);
        let db = try!(DB::open(&opts, path));

        let schema = match db.get(b".schema") {
            Ok(Some(schema)) => {
                let schema = schema.to_utf8().unwrap().to_string();
                json::decode(&schema).unwrap()
            }
            Ok(None) => Schema::new(),  // TODO: error
            Err(_) => Schema::new(),  // TODO: error
        };

        let next_term_ref = match db.get(b".next_term_ref") {
            Ok(Some(next_term_ref)) => {
                next_term_ref.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        let next_chunk = match db.get(b".next_chunk") {
            Ok(Some(next_chunk)) => {
                next_chunk.to_utf8().unwrap().parse::<u32>().unwrap()
            }
            Ok(None) => 1,  // TODO: error
            Err(_) => 1,  // TODO: error
        };

        Ok(RocksDBIndexStore {
            schema: Arc::new(schema),
            db: db,
            next_term_ref: AtomicU32::new(next_term_ref),
            term_dictionary: RwLock::new(BTreeMap::new()),
            next_chunk: AtomicU32::new(next_chunk),
            doc_key_mapping: RwLock::new(BTreeMap::new()),
            deleted_docs: RwLock::new(Vec::new()),
        })
    }

    pub fn add_field(&mut self, name: String, field_type: FieldType) -> Result<FieldRef, AddFieldError> {
        let mut schema_copy = (*self.schema).clone();
        let field_ref = try!(schema_copy.add_field(name, field_type));
        self.schema = Arc::new(schema_copy);

        self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());

        Ok(field_ref)
    }

    pub fn remove_field(&mut self, field_ref: &FieldRef) -> bool {
        let mut schema_copy = (*self.schema).clone();
        let field_removed = schema_copy.remove_field(field_ref);

        if field_removed {
            self.schema = Arc::new(schema_copy);
            self.db.put(b".schema", json::encode(&self.schema).unwrap().as_bytes());
        }

        field_removed
    }

    fn get_or_create_term(&mut self, term: &Term) -> TermRef {
        let term_bytes = term.to_bytes();

        if let Some(term_ref) = self.term_dictionary.read().unwrap().get(&term_bytes) {
            return *term_ref;
        }

        // Term doesn't exist in the term dictionary

        // Increment next_term_ref
        let next_term_ref = self.next_term_ref.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_term_ref", (next_term_ref + 1).to_string().as_bytes());

        // Create term ref
        let term_ref = TermRef(next_term_ref);

        // Get exclusive lock to term dictionary
        let mut term_dictionary = self.term_dictionary.write().unwrap();

        // It's possible that another thread has written the term to the dictionary
        // since we checked earlier. If this is the case, We should forget about
        // writing our TermRef and use the one that has been inserted already.
        if let Some(term_ref) = term_dictionary.get(&term_bytes) {
            return *term_ref;
        }

        // Write it to the on-disk term dictionary
        let mut key = Vec::with_capacity(1 + term_bytes.len());
        key.push(b't');
        for byte in term_bytes.iter() {
            key.push(*byte);
        }
        self.db.put(&key, next_term_ref.to_string().as_bytes());

        // Write it to the term dictionary
        term_dictionary.insert(term_bytes, term_ref);

        term_ref
    }

    pub fn insert_or_update_document(&mut self, doc: Document) {
        // Allocate a new chunk for the document

        // Chunk merges are very slow so we should avoid doing them at runtime
        // which is why each new document is created in a fresh chunk.
        // Later on, a background process will come and merge any small chunks
        // together.

        // For best performance, documents should be inserted in batches.

        // Increment next_chunk
        let next_chunk = self.next_chunk.fetch_add(1, Ordering::SeqCst);
        self.db.put(b".next_chunk", (next_chunk + 1).to_string().as_bytes());

        // Create doc ref
        let doc_ref = DocRef(next_chunk, 0);

        // Start write batch
        let mut write_batch = WriteBatch::default();

        // Insert contents
        let mut token_count: i64 = 0;
        for (field_ref, tokens) in doc.fields.iter() {
            for token in tokens.iter() {
                token_count += 1;
                let term_ref = self.get_or_create_term(&token.term);

                let mut key = Vec::with_capacity(20);
                key.push(b'd');
                for byte in field_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in term_ref.ord().to_string().as_bytes() {
                    key.push(*byte);
                }
                key.push(b'/');
                for byte in doc_ref.chunk().to_string().as_bytes() {
                    key.push(*byte);
                }

                let mut doc_id_bytes = [0; 2];
                BigEndian::write_u16(&mut doc_id_bytes, doc_ref.ord());
                write_batch.merge(&key, &doc_id_bytes);
            }
        }

        // Increment total docs
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_docs".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, 1);
        write_batch.merge(&key, &inc_bytes);

        // Increment total tokens
        let mut key = Vec::with_capacity(5);
        for byte in b"stotal_tokens".iter() {
            key.push(*byte);
        }
        key.push(b'/');
        for byte in doc_ref.chunk().to_string().as_bytes() {
            key.push(*byte);
        }
        let mut inc_bytes = [0; 8];
        BigEndian::write_i64(&mut inc_bytes, token_count);
        write_batch.merge(&key, &inc_bytes);

        // Write document data
        self.db.write(write_batch);

        // Update doc_key_mapping
        let mut write_batch = WriteBatch::default();
        let previous_doc_ref = self.doc_key_mapping.write().unwrap().insert(doc.key.as_bytes().iter().cloned().collect(), doc_ref);

        let mut key = Vec::with_capacity(1 + doc.key.len());
        key.push(b'k');
        for byte in doc.key.as_bytes() {
            key.push(*byte);
        }
        let mut doc_ref_bytes = [0; 6];
        BigEndian::write_u32(&mut doc_ref_bytes, doc_ref.chunk());
        BigEndian::write_u16(&mut doc_ref_bytes[4..], doc_ref.ord());
        write_batch.put(&key, &doc_ref_bytes);

        // If there was a document there previously, mark it as deleted
        if let Some(previous_doc_ref) = previous_doc_ref {
            self.deleted_docs.write().unwrap().push(previous_doc_ref);

            let mut key = Vec::with_capacity(5);
            key.push(b'x');
            for byte in previous_doc_ref.chunk().to_string().as_bytes() {
                key.push(*byte);
            }

            let mut previous_doc_id_bytes = [0; 2];
            BigEndian::write_u16(&mut previous_doc_id_bytes, previous_doc_ref.ord());
            write_batch.merge(&key, &previous_doc_id_bytes);
        }

        // Write document data
        self.db.write(write_batch);
    }
}

#[cfg(test)]
mod tests {
    use std::fs::remove_dir_all;

    use rocksdb::{DB, Options, IteratorMode};
    use abra::{Term, Token, Document};
    use abra::schema::{Schema, FieldType, FieldRef};

    use super::RocksDBIndexStore;

    #[test]
    fn test_create() {
        remove_dir_all("test_indices/test_create");

        let store = RocksDBIndexStore::create("test_indices/test_create");
        assert!(store.is_ok());
    }

    #[test]
    fn test_open() {
        remove_dir_all("test_indices/test_open");

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_err());

        // Create DB
        let mut opts = Options::default();
        opts.create_if_missing(true);
        DB::open(&opts, "test_indices/test_open").unwrap();

        let store = RocksDBIndexStore::open("test_indices/test_open");
        assert!(store.is_ok());
    }

    fn make_test_store(path: &str) -> RocksDBIndexStore {
        let mut store = RocksDBIndexStore::create(path).unwrap();
        let mut title_field = store.add_field("title".to_string(), FieldType::Text).unwrap();
        let mut body_field = store.add_field("body".to_string(), FieldType::Text).unwrap();

        store.insert_or_update_document(Document {
            key: "test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("hello".to_string()), position: 1 },
                    Token { term: Term::String("world".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store.insert_or_update_document(Document {
            key: "another_test_doc".to_string(),
            fields: hashmap! {
                title_field => vec![
                    Token { term: Term::String("howdy".to_string()), position: 1 },
                    Token { term: Term::String("partner".to_string()), position: 2 },
                ],
                body_field => vec![
                    Token { term: Term::String("lorem".to_string()), position: 1 },
                    Token { term: Term::String("ipsum".to_string()), position: 2 },
                    Token { term: Term::String("dolar".to_string()), position: 3 },
                ],
            }
        });

        store
    }

    pub fn print_keys(db: &DB) {
        fn bytes_to_string(bytes: &Box<[u8]>) -> String {
            use std::char;

            let mut string = String::new();

            for byte in bytes.iter() {
                if *byte < 128 {
                    // ASCII character
                    string.push(char::from_u32(*byte as u32).unwrap());
                } else {
                    string.push('?');
                }
            }

            string
        }

        for (key, value) in db.iterator(IteratorMode::Start) {
            println!("{} = {:?}", bytes_to_string(&key), value);
        }
    }

    #[test]
    fn test() {
        remove_dir_all("test_indices/test");

        let store = make_test_store("test_indices/test");

        print_keys(&store.db);

    }
}
<|endoftext|>


<commit_before>pub use binary::BinaryTrait;
use binary::{
    async_trait::async_trait,
    binary_clone_box,
    eyre::{bail, Result},
    semver_versions_matching, Binary,
};

pub struct PodmanBinary;

#[async_trait]
impl BinaryTrait for PodmanBinary {
    fn spec(&self) -> Binary {
        Binary::new("podman", &[], &[])
    }

    binary_clone_box!();

    async fn versions(&self, _os: &str) -> Result<Vec<String>> {
        self.versions_github_releases("containers", "podman")
            .await
            .map(|versions| semver_versions_matching(versions, ">=3"))
    }

    async fn install_version(&self, version: &str, os: &str, _arch: &str) -> Result<()> {
        let url = format!(
            "https://github.com/containers/podman/releases/download/v{version}/podman-remote-",
            version = version
        ) + match os {
            "linux" => "static.tar.gz",
            "macos" => "release-darwin.zip",
            "windows" => "release-windows.zip",
            _ => bail!(
                "Installation of `podman` for operating system `{}` is not supported",
                os
            ),
        };
        let archive = self.download(&url, None, None).await?;

        let dest = self.dir(Some(version.into()), true)?;
        self.extract(&archive, 0, &dest)?;
        self.executables(&dest, &["podman"])?;

        Ok(())
    }
}
<commit_msg>fix(Podman binary): Include version in downloaded file names<commit_after>pub use binary::BinaryTrait;
use binary::{
    async_trait::async_trait,
    binary_clone_box,
    eyre::{bail, Result},
    semver_versions_matching, Binary,
};

pub struct PodmanBinary;

#[async_trait]
impl BinaryTrait for PodmanBinary {
    fn spec(&self) -> Binary {
        Binary::new("podman", &[], &[])
    }

    binary_clone_box!();

    async fn versions(&self, _os: &str) -> Result<Vec<String>> {
        self.versions_github_releases("containers", "podman")
            .await
            .map(|versions| semver_versions_matching(versions, ">=3"))
    }

    async fn install_version(&self, version: &str, os: &str, _arch: &str) -> Result<()> {
        let suffix = match os {
            "linux" => "static.tar.gz",
            "macos" => "release-darwin.zip",
            "windows" => "release-windows.zip",
            _ => bail!(
                "Installation of `podman` for operating system `{}` is not supported",
                os
            ),
        };
        let url = format!(
            "https://github.com/containers/podman/releases/download/v{version}/podman-remote-",
            version = version
        ) + suffix;
        let filename = ["podman-remote-", version, "-", suffix].concat();
        let archive = self.download(&url, Some(filename), None).await?;

        let dest = self.dir(Some(version.into()), true)?;
        self.extract(&archive, 0, &dest)?;
        self.executables(&dest, &["podman"])?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::ffi::OsString;
use std::path::Path;

use binary::{
    async_trait::async_trait, binary_clone_box, eyre::Result, tracing, Binary, BinaryTrait,
};
use binary_python::PythonBinary;

/// A `BinaryTrait` for `poetry`
pub struct PoetryBinary;

#[async_trait]
impl BinaryTrait for PoetryBinary {
    fn spec(&self) -> Binary {
        Binary::new("poetry", &[], &[])
    }

    binary_clone_box!();

    async fn versions(&self, _os: &str) -> Result<Vec<String>> {
        self.versions_github_releases("python-poetry", "poetry")
            .await
            // "installer does not support Poetry releases < 0.12.0"
            .map(|versions| self.semver_versions_matching(versions, ">=0.12"))
    }

    async fn install_version(
        &self,
        version: &str,
        dest: &Path,
        _os: &str,
        _arch: &str,
    ) -> Result<()> {
        let script = self
            .download(
                "https://install.python-poetry.org",
                Some("install-poetry.py".to_string()),
                None,
            )
            .await?;

        tracing::info!(
            "Running `install-poetry.py` to install Poetry `{}` to `{}`",
            version,
            dest.display()
        );
        let mut python = PythonBinary {}
            .require(Some(">=3.6".to_string()), true)
            .await?;
        python.env_list(&[
            ("POETRY_HOME", dest.into()),
            ("POETRY_VERSION", OsString::from(version)),
            ("POETRY_ACCEPT", OsString::from("yes")),
        ]);
        python.run(&[&script.display().to_string()]).await?;

        Ok(())
    }
}
<commit_msg>fix(Poetry binary): Requires at least Python 3.7<commit_after>use std::ffi::OsString;
use std::path::Path;

use binary::{
    async_trait::async_trait, binary_clone_box, eyre::Result, tracing, Binary, BinaryTrait,
};
use binary_python::PythonBinary;

/// A `BinaryTrait` for `poetry`
pub struct PoetryBinary;

#[async_trait]
impl BinaryTrait for PoetryBinary {
    fn spec(&self) -> Binary {
        Binary::new("poetry", &[], &[])
    }

    binary_clone_box!();

    async fn versions(&self, _os: &str) -> Result<Vec<String>> {
        self.versions_github_releases("python-poetry", "poetry")
            .await
            // "installer does not support Poetry releases < 0.12.0"
            .map(|versions| self.semver_versions_matching(versions, ">=0.12"))
    }

    async fn install_version(
        &self,
        version: &str,
        dest: &Path,
        _os: &str,
        _arch: &str,
    ) -> Result<()> {
        let script = self
            .download(
                "https://install.python-poetry.org",
                Some("install-poetry.py".to_string()),
                None,
            )
            .await?;

        tracing::info!(
            "Running `install-poetry.py` to install Poetry `{}` to `{}`",
            version,
            dest.display()
        );
        let mut python = PythonBinary {}
            .require(Some(">=3.7".to_string()), true)
            .await?;
        python.env_list(&[
            ("POETRY_HOME", dest.into()),
            ("POETRY_VERSION", OsString::from(version)),
            ("POETRY_ACCEPT", OsString::from("yes")),
        ]);
        python.run(&[&script.display().to_string()]).await?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Deployments interface

use self::super::{Github, Result};
use rustc_serialize::json;
use rep::{Deployment, DeploymentReq, DeploymentStatus, DeploymentStatusReq};

/// Interface for repository deployements
pub struct Deployments<'a> {
    github: &'a Github<'a>,
    owner: String,
    repo: String,
}

/// INterface for deployment statuses
pub struct DeploymentStatuses<'a> {
    github: &'a Github<'a>,
    owner: String,
    repo: String,
    id: u64,
}

impl<'a> DeploymentStatuses<'a> {
    /// creates a new deployment status
    pub fn new<O, R>(github: &'a Github<'a>, owner: O, repo: R, id: u64) -> DeploymentStatuses<'a>
        where O: Into<String>,
              R: Into<String>
    {
        DeploymentStatuses {
            github: github,
            owner: owner.into(),
            repo: repo.into(),
            id: id,
        }
    }

    fn path(&self, more: &str) -> String {
        format!("/repos/{}/{}/deployments/{}/statuses{}",
                self.owner,
                self.repo,
                self.id,
                more)
    }

    /// lists all statuses associated with a deployment
    pub fn list(&self) -> Result<Vec<DeploymentStatus>> {
        let body = try!(self.github.get(&self.path("")));
        Ok(json::decode::<Vec<DeploymentStatus>>(&body).unwrap())
    }

    /// creates a new deployment status. For convenience, a DeploymentStatusReq.builder
    /// interface is required for building up a request
    pub fn create(&self, status: &DeploymentStatusReq) -> Result<DeploymentStatus> {
        let data = json::encode::<DeploymentStatusReq>(&status).unwrap();
        let body = try!(self.github.post(&self.path(""), &data.as_bytes()));
        Ok(json::decode::<DeploymentStatus>(&body).unwrap())
    }
}

impl<'a> Deployments<'a> {
    /// Create a new deployments instance
    pub fn new<O, R>(github: &'a Github<'a>, owner: O, repo: R) -> Deployments<'a>
        where O: Into<String>,
              R: Into<String>
    {
        Deployments {
            github: github,
            owner: owner.into(),
            repo: repo.into(),
        }
    }

    fn path(&self, more: &str) -> String {
        format!("/repos/{}/{}/deployments{}", self.owner, self.repo, more)
    }

    /// lists all deployments for a repository
    pub fn list(&self) -> Result<String> {
        let body = try!(self.github.get(&self.path("")));
        Ok(body)
    }

    /// creates a new deployment for this repository
    pub fn create(&self, dep: &DeploymentReq) -> Result<String> {
        let data = json::encode(&dep).unwrap();
        let body = try!(self.github.post(&self.path(""), data.as_bytes()));
        Ok(body)
    }

    /// get a reference to the statuses api for a give deployment
    pub fn statuses(&self, id: u64) -> DeploymentStatuses {
        DeploymentStatuses::new(self.github, self.owner.as_ref(), self.repo.as_ref(), id)
    }
}
<commit_msg>more warnings<commit_after>//! Deployments interface

use self::super::{Github, Result};
use rustc_serialize::json;
use rep::{DeploymentReq, DeploymentStatus, DeploymentStatusReq};

/// Interface for repository deployements
pub struct Deployments<'a> {
    github: &'a Github<'a>,
    owner: String,
    repo: String,
}

/// INterface for deployment statuses
pub struct DeploymentStatuses<'a> {
    github: &'a Github<'a>,
    owner: String,
    repo: String,
    id: u64,
}

impl<'a> DeploymentStatuses<'a> {
    /// creates a new deployment status
    pub fn new<O, R>(github: &'a Github<'a>, owner: O, repo: R, id: u64) -> DeploymentStatuses<'a>
        where O: Into<String>,
              R: Into<String>
    {
        DeploymentStatuses {
            github: github,
            owner: owner.into(),
            repo: repo.into(),
            id: id,
        }
    }

    fn path(&self, more: &str) -> String {
        format!("/repos/{}/{}/deployments/{}/statuses{}",
                self.owner,
                self.repo,
                self.id,
                more)
    }

    /// lists all statuses associated with a deployment
    pub fn list(&self) -> Result<Vec<DeploymentStatus>> {
        let body = try!(self.github.get(&self.path("")));
        Ok(json::decode::<Vec<DeploymentStatus>>(&body).unwrap())
    }

    /// creates a new deployment status. For convenience, a DeploymentStatusReq.builder
    /// interface is required for building up a request
    pub fn create(&self, status: &DeploymentStatusReq) -> Result<DeploymentStatus> {
        let data = json::encode::<DeploymentStatusReq>(&status).unwrap();
        let body = try!(self.github.post(&self.path(""), &data.as_bytes()));
        Ok(json::decode::<DeploymentStatus>(&body).unwrap())
    }
}

impl<'a> Deployments<'a> {
    /// Create a new deployments instance
    pub fn new<O, R>(github: &'a Github<'a>, owner: O, repo: R) -> Deployments<'a>
        where O: Into<String>,
              R: Into<String>
    {
        Deployments {
            github: github,
            owner: owner.into(),
            repo: repo.into(),
        }
    }

    fn path(&self, more: &str) -> String {
        format!("/repos/{}/{}/deployments{}", self.owner, self.repo, more)
    }

    /// lists all deployments for a repository
    pub fn list(&self) -> Result<String> {
        let body = try!(self.github.get(&self.path("")));
        Ok(body)
    }

    /// creates a new deployment for this repository
    pub fn create(&self, dep: &DeploymentReq) -> Result<String> {
        let data = json::encode(&dep).unwrap();
        let body = try!(self.github.post(&self.path(""), data.as_bytes()));
        Ok(body)
    }

    /// get a reference to the statuses api for a give deployment
    pub fn statuses(&self, id: u64) -> DeploymentStatuses {
        DeploymentStatuses::new(self.github, self.owner.as_ref(), self.repo.as_ref(), id)
    }
}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

use std::rc::*;

use super::super::tree::*;

use super::component::*;
use super::subscriptionmanager::*;

///
/// Stores a registration of a consumer
///
#[derive(Clone)]
struct ConsumerRegistration {
    address: TreeAddress,
    extent: TreeExtent
}

///
/// Consumer for data written by an immediate publisher
///
struct ImmediateConsumer {
    ///
    /// Where subscriptions can be registered for this consumer
    ///
    subscriptions: Rc<SubscriptionManager<ConsumerRegistration>>
}

impl Consumer for ImmediateConsumer {
    ///
    /// Calls a function whenever a particular section of the tree has changed
    ///
    fn subscribe(&mut self, address: TreeAddress, extent: TreeExtent, callback: ConsumerCallback) {
        // Need to persuade rust that it can call the FnMut (assign parameter to a mutable variable)
        let mut also_callback = callback;

        self.subscriptions.add_subscription(ConsumerRegistration { address: address.clone(), extent: extent }, Box::new(move |change| {
            // The change we get from the subscription will have an address relative to the root of the tree
            // Make the subscription change relative to the address that was subscribed to
            let maybe_relative_change = change.relative_to(&address);
            if let Some(relative_change) = maybe_relative_change {
                also_callback(&relative_change);
            }
        }));
    }
}

///
/// Publisher that immediately sends changes to its consumers. Can be used as a consumer factory.
///
pub struct ImmediatePublisher {
    ///
    /// Subscriptions for this publisher
    ///
    subscriptions: Rc<SubscriptionManager<ConsumerRegistration>>
}

impl ImmediatePublisher {
    ///
    /// Creates a new immediate publisher
    ///
    pub fn new() -> Box<ImmediatePublisher> {
        Box::new(ImmediatePublisher { subscriptions: Rc::new(SubscriptionManager::new()) })
    }

    ///
    /// Creates a consumer that will receive notifications from this publisher
    ///
    pub fn create_consumer(&self) -> ConsumerRef {
        Box::new(ImmediateConsumer { subscriptions: self.subscriptions.clone() })
    }
}

impl Publisher for ImmediatePublisher {
    ///
    /// Publishes a change to the consumers of this component
    ///
    fn publish(&mut self, change: TreeChange) {
        self.subscriptions.call_subscriptions(&|registration| {
            change.applies_to(&registration.address, &registration.extent).unwrap_or(false)
        }, &change);
    }
}

#[cfg(test)]
mod immediate_publisher_tests {
    use std::cell::*;
    use std::rc::*;

    use super::super::super::component::*;
    use super::*;

    #[test]
    fn can_consume_published_item() {
        let mut publisher   = ImmediatePublisher::new();
        let mut consumer    = publisher.create_consumer();

        let our_count       = Rc::new(Cell::new(0));
        let their_count     = our_count.clone();

        let mut consumer_tree = "empty".to_tree_node();

        // Subscribe to the second child of the root
        // Ie, tree should look like this:
        //
        // root
        //   +- Anything
        //   +- ConsumerTree
        //        +- add(value)
        consumer.subscribe(1.to_tree_address(), TreeExtent::SubTree, Box::new(move |change| {
            // Update the tree
            consumer_tree = change.apply(&consumer_tree);

            // Tree can have an 'add' node that specifies how much to add to the count for this change
            let old_val     = their_count.get();
            let tree_value  = consumer_tree.get_child_ref_at("add").map(|val| { val.get_value().to_int(0) }).unwrap_or(0);

            their_count.set(old_val + tree_value);
        }));

        // count is initially 0
        assert!(our_count.get() == 0);

        // Publish an add of 1 to set the count to 1
        let just_add = ("add", 1).to_tree_node();

        // Change the child of the second child of the root to the 'just_add' tree
        let modify_child_of_subscribed_tree = TreeChange::new(&(0, 1).to_tree_address(), TreeChangeType::Child, Some(&just_add));
        assert!(!modify_child_of_subscribed_tree.relative_to(&1.to_tree_address()).is_none());
        publisher.publish(modify_child_of_subscribed_tree);

        assert!(our_count.get() == 1);

        // Publish an add of 1 to set the count to 1
        let whole_tree = tree!("root", "some_other_tree", tree!("consumer_target", ("add", 2)));

        // Replace the entire tree with the tree above
        let modify_entire_tree = TreeChange::new(&TreeAddress::Here, TreeChangeType::Child, Some(&whole_tree));
        publisher.publish(modify_entire_tree);

        assert!(our_count.get() == 3);
    }
}
<commit_msg>Should fix the immediate_publisher tests<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

use std::rc::*;

use super::super::tree::*;

use super::component::*;
use super::subscriptionmanager::*;

///
/// Stores a registration of a consumer
///
#[derive(Clone)]
struct ConsumerRegistration {
    address: TreeAddress,
    extent: TreeExtent
}

///
/// Consumer for data written by an immediate publisher
///
struct ImmediateConsumer {
    ///
    /// Where subscriptions can be registered for this consumer
    ///
    subscriptions: Rc<SubscriptionManager<ConsumerRegistration>>
}

impl Consumer for ImmediateConsumer {
    ///
    /// Calls a function whenever a particular section of the tree has changed
    ///
    fn subscribe(&mut self, address: TreeAddress, extent: TreeExtent, callback: ConsumerCallback) {
        // Need to persuade rust that it can call the FnMut (assign parameter to a mutable variable)
        let mut also_callback = callback;

        self.subscriptions.add_subscription(ConsumerRegistration { address: address.clone(), extent: extent }, Box::new(move |change| {
            // The change we get from the subscription will have an address relative to the root of the tree
            // Make the subscription change relative to the address that was subscribed to
            let maybe_relative_change = change.relative_to(&address);
            if let Some(relative_change) = maybe_relative_change {
                also_callback(&relative_change);
            }
        }));
    }
}

///
/// Publisher that immediately sends changes to its consumers. Can be used as a consumer factory.
///
pub struct ImmediatePublisher {
    ///
    /// Subscriptions for this publisher
    ///
    subscriptions: Rc<SubscriptionManager<ConsumerRegistration>>
}

impl ImmediatePublisher {
    ///
    /// Creates a new immediate publisher
    ///
    pub fn new() -> Box<ImmediatePublisher> {
        Box::new(ImmediatePublisher { subscriptions: Rc::new(SubscriptionManager::new()) })
    }

    ///
    /// Creates a consumer that will receive notifications from this publisher
    ///
    pub fn create_consumer(&self) -> ConsumerRef {
        Box::new(ImmediateConsumer { subscriptions: self.subscriptions.clone() })
    }
}

impl Publisher for ImmediatePublisher {
    ///
    /// Publishes a change to the consumers of this component
    ///
    fn publish(&mut self, change: TreeChange) {
        self.subscriptions.call_subscriptions(&|registration| {
            change.applies_to(&registration.address, &registration.extent).unwrap_or(false)
        }, &change);
    }
}

#[cfg(test)]
mod immediate_publisher_tests {
    use std::cell::*;
    use std::rc::*;

    use super::super::super::component::*;
    use super::*;

    #[test]
    fn can_consume_published_item() {
        let mut publisher   = ImmediatePublisher::new();
        let mut consumer    = publisher.create_consumer();

        let our_count       = Rc::new(Cell::new(0));
        let their_count     = our_count.clone();

        let mut consumer_tree = "empty".to_tree_node();

        // Subscribe to the second child of the root
        // Ie, tree should look like this:
        //
        // root
        //   +- Anything
        //   +- ConsumerTree
        //        +- add(value)
        consumer.subscribe(1.to_tree_address(), TreeExtent::SubTree, Box::new(move |change| {
            // Update the tree
            consumer_tree = change.apply(&consumer_tree);

            // Tree can have an 'add' node that specifies how much to add to the count for this change
            let old_val     = their_count.get();
            let tree_value  = consumer_tree.get_child_ref_at("add").map(|val| { val.get_value().to_int(0) }).unwrap_or(0);

            their_count.set(old_val + tree_value);
        }));

        // count is initially 0
        assert!(our_count.get() == 0);

        // Publish an add of 1 to set the count to 1
        let just_add = ("add", 1).to_tree_node();

        // Change the child of the second child of the root to the 'just_add' tree
        let modify_child_of_subscribed_tree = TreeChange::new(&(1, 0).to_tree_address(), &just_add);
        assert!(!modify_child_of_subscribed_tree.relative_to(&1.to_tree_address()).is_none());
        publisher.publish(modify_child_of_subscribed_tree);

        assert!(our_count.get() == 1);

        // Publish an add of 1 to set the count to 1
        let whole_tree = tree!("root", "some_other_tree", tree!("consumer_target", ("add", 2)));

        // Replace the entire tree with the tree above
        let modify_entire_tree = TreeChange::new(&TreeAddress::Here, &whole_tree);
        publisher.publish(modify_entire_tree);

        assert!(our_count.get() == 3);
    }
}
<|endoftext|>


<commit_before><commit_msg>Removed outdated and ignored test<commit_after><|endoftext|>