<commit_before>// NOTE converted from assets/kate-syntax/rust.xml

use gfx::Color;

#[derive(Copy, Clone, PartialEq, Debug)]
pub struct Style {
    pub color: Color,
    pub bold: bool
}

// extensions=r".*\.rs" mimetype="text/x-rust"
pub struct Rust<'a, I> {
    lines: I,
    depth: usize,
    current_line: &'a str,
    output: Vec<(usize, Vec<(usize, Style)>)>
}

impl<'a, I> Rust<'a, I> where I: Iterator<Item=&'a str> {
    pub fn run(mut lines: I) -> (usize, Vec<(usize, Vec<(usize, Style)>)>) {
        let mut hl = Rust {
            current_line: lines.next().unwrap_or(""),
            lines: lines,
            depth: 0,
            output: vec![(0, vec![])]
        };
        hl.normal();
        (hl.depth, hl.output)
    }

    fn advance(&mut self, len: usize, style: Style) {
        self.current_line = &self.current_line[len..];

        let last = &mut self.output.last_mut().unwrap().1;
        if let Some(&mut (ref mut prev_len, prev_style)) = last.last_mut() {
            if style == prev_style {
                *prev_len += len;
                return;
            }
        }
        last.push((len, style))
    }

    fn next_line(&mut self) -> bool {
        if let Some(line) = self.lines.next() {
            self.current_line = line;
            self.output.push((self.depth, vec![]));
            true
        } else {
            self.current_line = "";
            false
        }
    }
}

macro_rules! rust_ident { () => ("[a-zA-Z_][a-zA-Z_0-9]*") }
macro_rules! rust_int_suf { () => ("([iu](8|16|32|64)?)?") }

macro_rules! re {
    ($($r:expr),+) => (regex!(concat!("^(?:", $($r,)* ")")))
}
macro_rules! kw {
    ($($r:expr),+) => (regex!(concat!("^(?:" $(,$r,)"|"* r")\b")))
}

macro_rules! goto {
    ($s:expr, {}) => (());
    ($s:expr, return) => ({ $s.depth -= 1; return;});
    ($s:expr, $cx:ident) => ({ $s.depth += 1; $s.$cx() })
}

#[cfg(not(test))]
macro_rules! cases {
    ($s:expr, $ds:ident:) => { {
        let ch_len = {
            let mut idx = $s.current_line.char_indices().map(|(i, _)| i);
            idx.next();
            idx.next().unwrap_or($s.current_line.len())
        };
        $s.advance(ch_len, styles::$ds)
    } };
    ($s:expr, $ds:ident: $re:expr; $($rest:tt)*) => { cases!($s, $ds: $re, $ds => {}; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr => $goto:tt; $($rest:tt)*) => { cases!($s, $ds: $re, $ds => $goto; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr, $style:ident; $($rest:tt)*) => { cases!($s, $ds: $re, $style => {}; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr, $style:ident => $goto:tt; $($rest:tt)*) => {
        if let Some((0, len)) = $re.find($s.current_line) {
            $s.advance(len, styles::$style);
            goto!($s, $goto);
        } else {
            cases!($s, $ds: $($rest)*)
        }
    }
}

#[cfg(test)]
macro_rules! cases { ($($rest:tt)*) => {()} }

macro_rules! context {
    ($name:ident, $ds:ident: $($rest:tt)*) => { context!{$name, $ds => {}: $($rest)*} };
    ($name:ident, $ds:ident => $line_end_cx:tt: $($rest:tt)*) => {
        fn $name(&mut self) {
            loop {
                while !self.current_line.is_empty() {
                    cases!(self, $ds: $($rest)*);
                }
                if self.next_line() {
                    goto!(self, $line_end_cx);
                } else {
                    return;
                }
            }
        }
    }
}

macro_rules! normal_like {
    ($($x:tt)+) => {context!{$($x)*
        re!(r"\s+");
        kw!("fn"), Keyword => function;
        kw!("type"), Keyword => type_;

        kw!("abstract",
            "alignof",
            "become",
            "do",
            "final",
            "offsetof",
            "override",
            "priv",
            "pure",
            "sizeof",
            "typeof",
            "unsized",
            "yield"), Keyword;
         kw!("as",
            "box",
            "break",
            "const",
            "continue",
            "crate",
            "else",
            "enum",
            "extern",
            "for",
            "if",
            "impl",
            "in",
            "let",
            "loop",
            "match",
            "mod",
            "move",
            "mut",
            "pub",
            "ref",
            "return",
            "static",
            "struct",
            "super",
            "trait",
            "unsafe",
            "use",
            "virtual",
            "where",
            "while"), Keyword;
        kw!("bool",
            "int",
            "isize",
            "uint",
            "usize",
            "i8",
            "i16",
            "i32",
            "i64",
            "u8",
            "u16",
            "u32",
            "u64",
            "f32",
            "f64",
            "float",
            "char",
            "str",
            "Option",
            "Result",
            "Self",
            "Box",
            "Vec",
            "String"), Type;
        kw!("AsSlice",
            "CharExt",
            "Clone",
            "Copy",
            "Debug",
            "Decodable",
            "Default",
            "Display",
            "DoubleEndedIterator",
            "Drop",
            "Encodable",
            "Eq",
            "Default",
            "Extend",
            "Fn",
            "FnMut",
            "FnOnce",
            "FromPrimitive",
            "Hash",
            "Iterator",
            "IteratorExt",
            "MutPtrExt",
            "Ord",
            "PartialEq",
            "PartialOrd",
            "PtrExt",
            "Rand",
            "Send",
            "Sized",
            "SliceConcatExt",
            "SliceExt",
            "Str",
            "StrExt",
            "Sync",
            "ToString"), Trait;
        kw!("c_float",
            "c_double",
            "c_void",
            "FILE",
            "fpos_t",
            "DIR",
            "dirent",
            "c_char",
            "c_schar",
            "c_uchar",
            "c_short",
            "c_ushort",
            "c_int",
            "c_uint",
            "c_long",
            "c_ulong",
            "size_t",
            "ptrdiff_t",
            "clock_t",
            "time_t",
            "c_longlong",
            "c_ulonglong",
            "intptr_t",
            "uintptr_t",
            "off_t",
            "dev_t",
            "ino_t",
            "pid_t",
            "mode_t",
            "ssize_t"), CType;
        kw!("self"), SelfKw;
        kw!("true",
            "false",
            "Some",
            "None",
            "Ok",
            "Err",
            "Success",
            "Failure",
            "Cons",
            "Nil"), Constant;
        kw!("EXIT_FAILURE",
            "EXIT_SUCCESS",
            "RAND_MAX",
            "EOF",
            "SEEK_SET",
            "SEEK_CUR",
            "SEEK_END",
            "_IOFBF",
            "_IONBF",
            "_IOLBF",
            "BUFSIZ",
            "FOPEN_MAX",
            "FILENAME_MAX",
            "L_tmpnam",
            "TMP_MAX",
            "O_RDONLY",
            "O_WRONLY",
            "O_RDWR",
            "O_APPEND",
            "O_CREAT",
            "O_EXCL",
            "O_TRUNC",
            "S_IFIFO",
            "S_IFCHR",
            "S_IFBLK",
            "S_IFDIR",
            "S_IFREG",
            "S_IFMT",
            "S_IEXEC",
            "S_IWRITE",
            "S_IREAD",
            "S_IRWXU",
            "S_IXUSR",
            "S_IWUSR",
            "S_IRUSR",
            "F_OK",
            "R_OK",
            "W_OK",
            "X_OK",
            "STDIN_FILENO",
            "STDOUT_FILENO",
            "STDERR_FILENO"), CConstant;

        // Match special comments for region markers
        re!("//BEGIN"), RegionMarker => region_marker;
        re!("//END"), RegionMarker => region_marker;

        // Match comments
        re!("//"), Comment => comment;
        re!(r"/\*"), Comment => comment2;
        re!(r"0x[0-9a-fA-F_]+", rust_int_suf!()), Hex;
        re!(r"0o[0-7_]+", rust_int_suf!()), Octal;
        re!(r"0b[0-1_]+", rust_int_suf!()), Binary;
        re!(r"[0-9][0-9_]*\.[0-9_]*([eE][+-]?[0-9_]+)?(f32|f64|f)?"), Float;
        re!(r"[0-9][0-9_]*", rust_int_suf!()), Decimal;
        re!(r"#\["), Attribute => attribute;
        re!(r"#!\["), Attribute => attribute;
        re!(rust_ident!(), "::"), Scope;
        re!(rust_ident!(), "!"), Macro;
        //re!("'", rust_ident!(), "(?!')"), Lifetime;
        re!(r"[{}\[\]]"), Symbol;
        re!("r\""), String => raw_string;
        re!("r##\""), String => raw_hashed2;
        re!("r#\""), String => raw_hashed1;
        re!("\""), String => string;
        re!(r"'"), Character => character;
        re!(rust_ident!());
    }}
}

impl<'a, I> Rust<'a, I> where I: Iterator<Item=&'a str> {
    normal_like!{normal, NormalText:}
    normal_like!{attribute, Attribute:
        re!(r"\]"), Attribute => return;
    }
    context!{function, Definition:
        re!(r"\s+");
        re!(r"\("), NormalText => return;
        re!(r">"), NormalText => return;
    }
    context!{type_, Definition:
        re!(r"\s+");
        re!(r"[=<;]"), NormalText => return;
    }
    // Rustc allows strings to extend over multiple lines, and the
    // only thing a backshash at end-of-line does is remove the whitespace.
    context!{string, String:
        re!(r"\\"), CharEscape => char_escape;
        re!("\""), String => return;
    }
    context!{raw_string, String:
        re!("\""), String => return;
    }
    // These rules are't complete: they won't match r###"abc"###
    context!{raw_hashed1, String:
        re!("\"#"), String => return;
    }
    context!{raw_hashed2, String:
        re!("\"##"), String => return;
    }
    context!{character, Character => return:
        re!(r"\\"), CharEscape => char_escape;
        re!(r"'"), Character => return;
    }
    context!{char_escape, CharEscape => return:
        re!("[nrt'\"]"), CharEscape => return;
        re!(r"x[0-9a-fA-F]{2}"), CharEscape => return;
        re!(r"u\{[0-9a-fA-F]{1,6}\}"), CharEscape => return;
        re!(r"u[0-9a-fA-F]{4}"), CharEscape => return;
        re!(r"U[0-9a-fA-F]{8}"), CharEscape => return;
        re!(r"\."), Error => return;
    }
    context!{region_marker, RegionMarker => return:}
    context!{comment, Comment => return:
        re!(r"\s+");
    }
    context!{comment2, Comment:
        re!(r"\s+");
        re!(r"\*/"), Comment => return;
    }
}

#[allow(non_upper_case_globals)]
pub mod base_styles {
    use super::Style;

    const DEFAULT: Style = Style { color: [1.0, 1.0, 1.0, 1.0], bold: false };

    macro_rules! styles {
        () => {};
        ($name:ident => ($r:expr,$g:expr,$b:expr) $($rest:tt)*) => {
            styles!($name => color=[$r as f32 / 255.0, $g as f32 / 255.0, $b as f32 / 255.0, 1.0] $($rest)*);
        };
        ($name:ident => $($extra:ident=$val:expr),*; $($rest:tt)*) => {
            pub const $name: Style = Style { $($extra: $val,)* ..DEFAULT };
            styles!($($rest)*);
        }
    }

    styles!{
        Normal => (239,240,241);
        Keyword => (239,240,241), bold=true;
        Function => (246,116,0);
        DataType => (41,128,185);
        DecVal => (246,116,0);
        BaseN => (246,116,0);
        Float => (246,116,0);
        Char => (61,174,233);
        String => (218,68,83);
        Comment => (189,195,199);
        RegionMarker => (41,128,185);
        Others => (39,174,96);
        Error => (218,68,83);
    }
}

#[allow(non_upper_case_globals)]
pub mod styles {
    use super::{Style, base_styles};

    macro_rules! styles {
        () => {};
        ($name:ident => $base:ident; $($rest:tt)*) => {
            pub const $name: Style = base_styles::$base;
            styles!($($rest)*);
        };
        ($name:ident => $base:ident, $($extra:ident=$val:expr),*; $($rest:tt)*) => {
            pub const $name: Style = Style { $($extra: $val,)* ..base_styles::$base };
            styles!($($rest)*);
        }
    }

    styles!{
        NormalText => Normal;
        Keyword => Keyword;
        SelfKw => Function, bold=true;
        Type => DataType, bold=true;
        Trait => DataType, bold=true;
        CType => DataType;
        Constant => Keyword;
        CConstant => Normal;
        Definition => Function;
        Comment => Comment;
        Scope => Others, bold=true;
        Decimal => DecVal;
        Hex => BaseN;
        Octal => BaseN;
        Binary => BaseN;
        Float => Float;
        String => String;
        CharEscape => Char;
        Character => Char;
        Macro => Others;
        Attribute => Others;
        Lifetime => Char, bold=true;
        RegionMarker => RegionMarker;
        Error => Error;

        Symbol => Normal;
    }
}
<commit_msg>highlight: partial fix for Rust lifetime syntax.<commit_after>// NOTE converted from assets/kate-syntax/rust.xml

use gfx::Color;

#[derive(Copy, Clone, PartialEq, Debug)]
pub struct Style {
    pub color: Color,
    pub bold: bool
}

// extensions=r".*\.rs" mimetype="text/x-rust"
pub struct Rust<'a, I> {
    lines: I,
    depth: usize,
    current_line: &'a str,
    output: Vec<(usize, Vec<(usize, Style)>)>
}

impl<'a, I> Rust<'a, I> where I: Iterator<Item=&'a str> {
    pub fn run(mut lines: I) -> (usize, Vec<(usize, Vec<(usize, Style)>)>) {
        let mut hl = Rust {
            current_line: lines.next().unwrap_or(""),
            lines: lines,
            depth: 0,
            output: vec![(0, vec![])]
        };
        hl.normal();
        (hl.depth, hl.output)
    }

    fn advance(&mut self, len: usize, style: Style) {
        self.current_line = &self.current_line[len..];

        let last = &mut self.output.last_mut().unwrap().1;
        if let Some(&mut (ref mut prev_len, prev_style)) = last.last_mut() {
            if style == prev_style {
                *prev_len += len;
                return;
            }
        }
        last.push((len, style))
    }

    fn next_line(&mut self) -> bool {
        if let Some(line) = self.lines.next() {
            self.current_line = line;
            self.output.push((self.depth, vec![]));
            true
        } else {
            self.current_line = "";
            false
        }
    }
}

macro_rules! rust_ident { () => ("[a-zA-Z_][a-zA-Z_0-9]*") }
macro_rules! rust_int_suf { () => ("([iu](8|16|32|64)?)?") }

macro_rules! re {
    ($($r:expr),+) => (regex!(concat!("^(?:", $($r,)* ")")))
}
macro_rules! kw {
    ($($r:expr),+) => (regex!(concat!("^(?:" $(,$r,)"|"* r")\b")))
}

macro_rules! goto {
    ($s:expr, {}) => (());
    ($s:expr, return) => ({ $s.depth -= 1; return;});
    ($s:expr, $cx:ident) => ({ $s.depth += 1; $s.$cx() })
}

#[cfg(not(test))]
macro_rules! cases {
    ($s:expr, $ds:ident:) => { {
        let ch_len = {
            let mut idx = $s.current_line.char_indices().map(|(i, _)| i);
            idx.next();
            idx.next().unwrap_or($s.current_line.len())
        };
        $s.advance(ch_len, styles::$ds)
    } };
    ($s:expr, $ds:ident: $re:expr; $($rest:tt)*) => { cases!($s, $ds: $re, $ds => {}; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr => $goto:tt; $($rest:tt)*) => { cases!($s, $ds: $re, $ds => $goto; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr, $style:ident; $($rest:tt)*) => { cases!($s, $ds: $re, $style => {}; $($rest)*) };
    ($s:expr, $ds:ident: $re:expr, $style:ident => $goto:tt; $($rest:tt)*) => {
        if let Some((0, len)) = $re.find($s.current_line) {
            $s.advance(len, styles::$style);
            goto!($s, $goto);
        } else {
            cases!($s, $ds: $($rest)*)
        }
    }
}

#[cfg(test)]
macro_rules! cases { ($($rest:tt)*) => {()} }

macro_rules! context {
    ($name:ident, $ds:ident: $($rest:tt)*) => { context!{$name, $ds => {}: $($rest)*} };
    ($name:ident, $ds:ident => $line_end_cx:tt: $($rest:tt)*) => {
        fn $name(&mut self) {
            loop {
                while !self.current_line.is_empty() {
                    cases!(self, $ds: $($rest)*);
                }
                if self.next_line() {
                    goto!(self, $line_end_cx);
                } else {
                    return;
                }
            }
        }
    }
}

macro_rules! normal_like {
    ($($x:tt)+) => {context!{$($x)*
        re!(r"\s+");
        kw!("fn"), Keyword => function;
        kw!("type"), Keyword => type_;

        kw!("abstract",
            "alignof",
            "become",
            "do",
            "final",
            "offsetof",
            "override",
            "priv",
            "pure",
            "sizeof",
            "typeof",
            "unsized",
            "yield"), Keyword;
         kw!("as",
            "box",
            "break",
            "const",
            "continue",
            "crate",
            "else",
            "enum",
            "extern",
            "for",
            "if",
            "impl",
            "in",
            "let",
            "loop",
            "match",
            "mod",
            "move",
            "mut",
            "pub",
            "ref",
            "return",
            "static",
            "struct",
            "super",
            "trait",
            "unsafe",
            "use",
            "virtual",
            "where",
            "while"), Keyword;
        kw!("bool",
            "int",
            "isize",
            "uint",
            "usize",
            "i8",
            "i16",
            "i32",
            "i64",
            "u8",
            "u16",
            "u32",
            "u64",
            "f32",
            "f64",
            "float",
            "char",
            "str",
            "Option",
            "Result",
            "Self",
            "Box",
            "Vec",
            "String"), Type;
        kw!("AsSlice",
            "CharExt",
            "Clone",
            "Copy",
            "Debug",
            "Decodable",
            "Default",
            "Display",
            "DoubleEndedIterator",
            "Drop",
            "Encodable",
            "Eq",
            "Default",
            "Extend",
            "Fn",
            "FnMut",
            "FnOnce",
            "FromPrimitive",
            "Hash",
            "Iterator",
            "IteratorExt",
            "MutPtrExt",
            "Ord",
            "PartialEq",
            "PartialOrd",
            "PtrExt",
            "Rand",
            "Send",
            "Sized",
            "SliceConcatExt",
            "SliceExt",
            "Str",
            "StrExt",
            "Sync",
            "ToString"), Trait;
        kw!("c_float",
            "c_double",
            "c_void",
            "FILE",
            "fpos_t",
            "DIR",
            "dirent",
            "c_char",
            "c_schar",
            "c_uchar",
            "c_short",
            "c_ushort",
            "c_int",
            "c_uint",
            "c_long",
            "c_ulong",
            "size_t",
            "ptrdiff_t",
            "clock_t",
            "time_t",
            "c_longlong",
            "c_ulonglong",
            "intptr_t",
            "uintptr_t",
            "off_t",
            "dev_t",
            "ino_t",
            "pid_t",
            "mode_t",
            "ssize_t"), CType;
        kw!("self"), SelfKw;
        kw!("true",
            "false",
            "Some",
            "None",
            "Ok",
            "Err",
            "Success",
            "Failure",
            "Cons",
            "Nil"), Constant;
        kw!("EXIT_FAILURE",
            "EXIT_SUCCESS",
            "RAND_MAX",
            "EOF",
            "SEEK_SET",
            "SEEK_CUR",
            "SEEK_END",
            "_IOFBF",
            "_IONBF",
            "_IOLBF",
            "BUFSIZ",
            "FOPEN_MAX",
            "FILENAME_MAX",
            "L_tmpnam",
            "TMP_MAX",
            "O_RDONLY",
            "O_WRONLY",
            "O_RDWR",
            "O_APPEND",
            "O_CREAT",
            "O_EXCL",
            "O_TRUNC",
            "S_IFIFO",
            "S_IFCHR",
            "S_IFBLK",
            "S_IFDIR",
            "S_IFREG",
            "S_IFMT",
            "S_IEXEC",
            "S_IWRITE",
            "S_IREAD",
            "S_IRWXU",
            "S_IXUSR",
            "S_IWUSR",
            "S_IRUSR",
            "F_OK",
            "R_OK",
            "W_OK",
            "X_OK",
            "STDIN_FILENO",
            "STDOUT_FILENO",
            "STDERR_FILENO"), CConstant;

        // Match special comments for region markers
        re!("//BEGIN"), RegionMarker => region_marker;
        re!("//END"), RegionMarker => region_marker;

        // Match comments
        re!("//"), Comment => comment;
        re!(r"/\*"), Comment => comment2;
        re!(r"0x[0-9a-fA-F_]+", rust_int_suf!()), Hex;
        re!(r"0o[0-7_]+", rust_int_suf!()), Octal;
        re!(r"0b[0-1_]+", rust_int_suf!()), Binary;
        re!(r"[0-9][0-9_]*\.[0-9_]*([eE][+-]?[0-9_]+)?(f32|f64|f)?"), Float;
        re!(r"[0-9][0-9_]*", rust_int_suf!()), Decimal;
        re!(r"#\["), Attribute => attribute;
        re!(r"#!\["), Attribute => attribute;
        re!(rust_ident!(), "::"), Scope;
        re!(rust_ident!(), "!"), Macro;
        re!("'", rust_ident!()/*, "(?!')"*/), Lifetime;
        re!(r"[{}\[\]]"), Symbol;
        re!("r\""), String => raw_string;
        re!("r##\""), String => raw_hashed2;
        re!("r#\""), String => raw_hashed1;
        re!("\""), String => string;
        re!(r"'"), Character => character;
        re!(rust_ident!());
    }}
}

impl<'a, I> Rust<'a, I> where I: Iterator<Item=&'a str> {
    normal_like!{normal, NormalText:}
    normal_like!{attribute, Attribute:
        re!(r"\]"), Attribute => return;
    }
    context!{function, Definition:
        re!(r"\s+");
        re!(r"\("), NormalText => return;
        re!(r">"), NormalText => return;
    }
    context!{type_, Definition:
        re!(r"\s+");
        re!(r"[=<;]"), NormalText => return;
    }
    // Rustc allows strings to extend over multiple lines, and the
    // only thing a backshash at end-of-line does is remove the whitespace.
    context!{string, String:
        re!(r"\\"), CharEscape => char_escape;
        re!("\""), String => return;
    }
    context!{raw_string, String:
        re!("\""), String => return;
    }
    // These rules are't complete: they won't match r###"abc"###
    context!{raw_hashed1, String:
        re!("\"#"), String => return;
    }
    context!{raw_hashed2, String:
        re!("\"##"), String => return;
    }
    context!{character, Character => return:
        re!(r"\\"), CharEscape => char_escape;
        re!(r"'"), Character => return;
    }
    context!{char_escape, CharEscape => return:
        re!("[nrt'\"]"), CharEscape => return;
        re!(r"x[0-9a-fA-F]{2}"), CharEscape => return;
        re!(r"u\{[0-9a-fA-F]{1,6}\}"), CharEscape => return;
        re!(r"u[0-9a-fA-F]{4}"), CharEscape => return;
        re!(r"U[0-9a-fA-F]{8}"), CharEscape => return;
        re!(r"\."), Error => return;
    }
    context!{region_marker, RegionMarker => return:}
    context!{comment, Comment => return:
        re!(r"\s+");
    }
    context!{comment2, Comment:
        re!(r"\s+");
        re!(r"\*/"), Comment => return;
    }
}

#[allow(non_upper_case_globals)]
pub mod base_styles {
    use super::Style;

    const DEFAULT: Style = Style { color: [1.0, 1.0, 1.0, 1.0], bold: false };

    macro_rules! styles {
        () => {};
        ($name:ident => ($r:expr,$g:expr,$b:expr) $($rest:tt)*) => {
            styles!($name => color=[$r as f32 / 255.0, $g as f32 / 255.0, $b as f32 / 255.0, 1.0] $($rest)*);
        };
        ($name:ident => $($extra:ident=$val:expr),*; $($rest:tt)*) => {
            pub const $name: Style = Style { $($extra: $val,)* ..DEFAULT };
            styles!($($rest)*);
        }
    }

    styles!{
        Normal => (239,240,241);
        Keyword => (239,240,241), bold=true;
        Function => (246,116,0);
        DataType => (41,128,185);
        DecVal => (246,116,0);
        BaseN => (246,116,0);
        Float => (246,116,0);
        Char => (61,174,233);
        String => (218,68,83);
        Comment => (189,195,199);
        RegionMarker => (41,128,185);
        Others => (39,174,96);
        Error => (218,68,83);
    }
}

#[allow(non_upper_case_globals)]
pub mod styles {
    use super::{Style, base_styles};

    macro_rules! styles {
        () => {};
        ($name:ident => $base:ident; $($rest:tt)*) => {
            pub const $name: Style = base_styles::$base;
            styles!($($rest)*);
        };
        ($name:ident => $base:ident, $($extra:ident=$val:expr),*; $($rest:tt)*) => {
            pub const $name: Style = Style { $($extra: $val,)* ..base_styles::$base };
            styles!($($rest)*);
        }
    }

    styles!{
        NormalText => Normal;
        Keyword => Keyword;
        SelfKw => Function, bold=true;
        Type => DataType, bold=true;
        Trait => DataType, bold=true;
        CType => DataType;
        Constant => Keyword;
        CConstant => Normal;
        Definition => Function;
        Comment => Comment;
        Scope => Others, bold=true;
        Decimal => DecVal;
        Hex => BaseN;
        Octal => BaseN;
        Binary => BaseN;
        Float => Float;
        String => String;
        CharEscape => Char;
        Character => Char;
        Macro => Others;
        Attribute => Others;
        Lifetime => Char, bold=true;
        RegionMarker => RegionMarker;
        Error => Error;

        Symbol => Normal;
    }
}
<|endoftext|>


<commit_before>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::*;
use client::AsyncClient;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

pub const LEVEL_ELEMENTS_MULTIPLIER: usize = 10;
pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 10;

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_M: usize = 24;
const LEVEL_1: usize = LEVEL_M * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_PAGE_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<LevelTree>>;
pub type Ptr = NodeCellRef;
pub type Key = EntryKey;

with_levels! {
    LM, LEVEL_M;
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    // L4, LEVEL_4; // See https://github.com/rust-lang/rust/issues/58164
}

pub struct LSMTree {
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        debug!("Initializing LSM-tree...");
        let (trees, max_sizes) = init_lsm_level_trees(neb_client);
        debug!("Initialized LSM-tree");
        Arc::new(LSMTree { trees, max_sizes })
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) {
        key_with_id(&mut key, id);
        self.trees[0].insert_into(&key)
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> bool {
        key_with_id(&mut key, id);
        self.trees
            .iter()
            .map(|tree| tree.mark_key_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors: Vec<Box<Cursor>> = vec![];
        for tree in &self.trees {
            cursors.push(tree.seek_for(&key, ordering));
        }
        return LSMTreeCursor::new(cursors);
    }

    pub fn check_and_merge(&self) {
        for i in 0..self.trees.len() - 1 {
            debug!("Checking tree merge {}", i);
            let lower = &*self.trees[i];
            let upper = &*self.trees[i + 1];
            if lower.count() > self.max_sizes[i] {
                lower.merge_to(upper);
            }
        }
    }

    pub fn start_sentinel(this: &Arc<Self>) {
        let this = this.clone();
        thread::Builder::new().name("LSM-Tree Sentinel".to_string()).spawn(move || {
            loop {
                this.check_and_merge();
                thread::sleep(Duration::from_millis(500));
            }
        });
    }

    pub fn level_sizes(&self) -> Vec<usize> {
        self.trees.iter().map(|t| t.count()).collect()
    }

    pub fn len(&self) -> usize {
        self.trees.iter().map(|tree| tree.count()).sum::<usize>()
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>) -> Self {
        LSMTreeCursor {
            level_cursors: cursors,
        }
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        if let Some(prev_key) = self.current().map(|k| k.to_owned()) {
            self.level_cursors
                .iter_mut()
                .filter(|c| c.current().is_some())
                .min_by(|a, b| a.current().unwrap().cmp(b.current().unwrap()))
                .map(|c| c.next());
            let dedupe_current = if let Some(current) = self.current() {
                if current <= &prev_key {
                    None
                } else {
                    Some(true)
                }
            } else {
                Some(false)
            };
            if let Some(has_next) = dedupe_current {
                has_next
            } else {
                self.next()
            }
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    fn dump_trees(lsm_tree: &LSMTree, name: &str) {
        for i in 0..lsm_tree.trees.len() {
            lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
        }
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5700");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 3 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(btree::page_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("331776".to_string())
            .parse::<u64>()
            .unwrap();

        let tree_clone = tree.clone();
        debug!("Testing LSM-tree");
        thread::spawn(move || loop {
            thread::sleep(Duration::from_secs(10));
            let tree_len = tree_clone.len();
            debug!(
                "LSM-Tree now have {}/{} elements, total {:.2}%",
                tree_len,
                num,
                tree_len as f32 / num as f32 * 100.0
            );
        });

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_1_before_insertion");
        for _ in 0..20 {
            tree.check_and_merge();
        }
        dump_trees(&*tree, "stage_1_after_insertion");

        debug!("Start point search validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });

        LSMTree::start_sentinel(&tree);

        (num..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_2_after_insertion");
        thread::sleep(Duration::new(30, 0));
        dump_trees(&*tree, "stage_2_waited_insertion");

        (0..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });
    }
}
<commit_msg>randomize<commit_after>use super::btree::{BPlusTree, RTCursor as BPlusTreeCursor};
use super::*;
use client::AsyncClient;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use itertools::Itertools;
use parking_lot::RwLock;
use ram::segs::MAX_SEGMENT_SIZE;
use std::collections::BTreeSet;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use std::{mem, ptr};

pub const LEVEL_ELEMENTS_MULTIPLIER: usize = 10;
pub const LEVEL_PAGE_DIFF_MULTIPLIER: usize = 10;

const LEVEL_M_MAX_ELEMENTS_COUNT: usize = LEVEL_M * LEVEL_M * LEVEL_M;
const LEVEL_M: usize = 24;
const LEVEL_1: usize = LEVEL_M * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_2: usize = LEVEL_1 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_3: usize = LEVEL_2 * LEVEL_PAGE_DIFF_MULTIPLIER;
const LEVEL_4: usize = LEVEL_3 * LEVEL_PAGE_DIFF_MULTIPLIER;

type LevelTrees = Vec<Box<LevelTree>>;
pub type Ptr = NodeCellRef;
pub type Key = EntryKey;

with_levels! {
    LM, LEVEL_M;
    L1, LEVEL_1;
    L2, LEVEL_2;
    L3, LEVEL_3;
    // L4, LEVEL_4; // See https://github.com/rust-lang/rust/issues/58164
}

pub struct LSMTree {
    trees: LevelTrees,
    // use Vec here for convenience
    max_sizes: Vec<usize>,
}

unsafe impl Send for LSMTree {}
unsafe impl Sync for LSMTree {}

impl LSMTree {
    pub fn new(neb_client: &Arc<AsyncClient>) -> Arc<Self> {
        debug!("Initializing LSM-tree...");
        let (trees, max_sizes) = init_lsm_level_trees(neb_client);
        debug!("Initialized LSM-tree");
        Arc::new(LSMTree { trees, max_sizes })
    }

    pub fn insert(&self, mut key: EntryKey, id: &Id) {
        key_with_id(&mut key, id);
        self.trees[0].insert_into(&key)
    }

    pub fn remove(&self, mut key: EntryKey, id: &Id) -> bool {
        key_with_id(&mut key, id);
        self.trees
            .iter()
            .map(|tree| tree.mark_key_deleted(&key))
            .collect_vec() // collect here to prevent short circuit
            .into_iter()
            .any(|d| d)
    }

    pub fn seek(&self, mut key: EntryKey, ordering: Ordering) -> LSMTreeCursor {
        match ordering {
            Ordering::Forward => key_with_id(&mut key, &Id::unit_id()),
            Ordering::Backward => key_with_id(&mut key, &Id::new(::std::u64::MAX, ::std::u64::MAX)),
        };
        let mut cursors: Vec<Box<Cursor>> = vec![];
        for tree in &self.trees {
            cursors.push(tree.seek_for(&key, ordering));
        }
        return LSMTreeCursor::new(cursors);
    }

    pub fn check_and_merge(&self) {
        for i in 0..self.trees.len() - 1 {
            debug!("Checking tree merge {}", i);
            let lower = &*self.trees[i];
            let upper = &*self.trees[i + 1];
            if lower.count() > self.max_sizes[i] {
                lower.merge_to(upper);
            }
        }
    }

    pub fn start_sentinel(this: &Arc<Self>) {
        let this = this.clone();
        thread::Builder::new().name("LSM-Tree Sentinel".to_string()).spawn(move || {
            loop {
                this.check_and_merge();
                thread::sleep(Duration::from_millis(500));
            }
        });
    }

    pub fn level_sizes(&self) -> Vec<usize> {
        self.trees.iter().map(|t| t.count()).collect()
    }

    pub fn len(&self) -> usize {
        self.trees.iter().map(|tree| tree.count()).sum::<usize>()
    }
}

pub struct LSMTreeCursor {
    level_cursors: Vec<Box<Cursor>>,
}

impl LSMTreeCursor {
    fn new(cursors: Vec<Box<Cursor>>) -> Self {
        LSMTreeCursor {
            level_cursors: cursors,
        }
    }
}

impl Cursor for LSMTreeCursor {
    fn next(&mut self) -> bool {
        if let Some(prev_key) = self.current().map(|k| k.to_owned()) {
            self.level_cursors
                .iter_mut()
                .filter(|c| c.current().is_some())
                .min_by(|a, b| a.current().unwrap().cmp(b.current().unwrap()))
                .map(|c| c.next());
            let dedupe_current = if let Some(current) = self.current() {
                if current <= &prev_key {
                    None
                } else {
                    Some(true)
                }
            } else {
                Some(false)
            };
            if let Some(has_next) = dedupe_current {
                has_next
            } else {
                self.next()
            }
        } else {
            false
        }
    }

    fn current(&self) -> Option<&EntryKey> {
        self.level_cursors.iter().filter_map(|c| c.current()).min()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use byteorder::BigEndian;
    use byteorder::WriteBytesExt;
    use client;
    use futures::prelude::*;
    use rand::distributions::Uniform;
    use rand::thread_rng;
    use rand::Rng;
    use rayon::prelude::*;
    use server::NebServer;
    use server::ServerOptions;
    use std::env;
    use std::io::Cursor as StdCursor;

    fn u64_to_slice(n: u64) -> [u8; 8] {
        let mut key_slice = [0u8; 8];
        {
            let mut cursor = StdCursor::new(&mut key_slice[..]);
            cursor.write_u64::<BigEndian>(n);
        };
        key_slice
    }

    fn dump_trees(lsm_tree: &LSMTree, name: &str) {
        for i in 0..lsm_tree.trees.len() {
            lsm_tree.trees[i].dump(&format!("{}_lsm_{}_dump.json", name, i));
        }
    }

    #[test]
    pub fn insertions() {
        env_logger::init();
        let server_group = "sstable_index_init";
        let server_addr = String::from("127.0.0.1:5700");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 1,
                memory_size: 3 * 1024 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
            },
            &server_addr,
            &server_group,
        );
        let client = Arc::new(
            client::AsyncClient::new(&server.rpc, &vec![server_addr], server_group).unwrap(),
        );
        client.new_schema_with_id(btree::page_schema()).wait();
        let tree = LSMTree::new(&client);
        let num = env::var("LSM_TREE_TEST_ITEMS")
            // this value cannot do anything useful to the test
            // must arrange a long-term test to cover every levels
            .unwrap_or("331776".to_string())
            .parse::<u64>()
            .unwrap();

        let tree_clone = tree.clone();
        debug!("Testing LSM-tree");
        thread::spawn(move || loop {
            thread::sleep(Duration::from_secs(10));
            let tree_len = tree_clone.len();
            debug!(
                "LSM-Tree now have {}/{} elements, total {:.2}%",
                tree_len,
                num,
                tree_len as f32 / num as f32 * 100.0
            );
        });

        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_1_before_insertion");
        for _ in 0..20 {
            tree.check_and_merge();
        }
        dump_trees(&*tree, "stage_1_after_insertion");

        debug!("Start point search validations");
        (0..num).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });

        LSMTree::start_sentinel(&tree);

        let mut test_data = (num..num * 2).collect_vec();
        thread_rng().shuffle(test_data.as_mut_slice());
        test_data.par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let key = SmallVec::from_slice(&key_slice);
            tree.insert(key, &id);
        });

        dump_trees(&*tree, "stage_2_after_insertion");
        thread::sleep(Duration::new(30, 0));
        dump_trees(&*tree, "stage_2_waited_insertion");

        (0..num * 2).collect::<Vec<_>>().par_iter().for_each(|i| {
            let i = *i;
            let id = Id::new(0, i);
            let key_slice = u64_to_slice(i);
            let mut key = SmallVec::from_slice(&key_slice);
            debug!("checking: {}", i);
            let mut cursor = tree.seek(key.clone(), Ordering::Forward);
            key_with_id(&mut key, &id);
            assert_eq!(cursor.current(), Some(&key), "{}", i);
        });
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::process::Command;
use std::time::{Duration, Instant};

fn main() {
    let mut args = env::args_os().skip(1).collect::<Vec<_>>();
    let rustc = env::var_os("RUSTC_REAL").unwrap();

    if let Some(pos) = args.iter().position(|arg| arg == "--wrap-rustc-with") {
        // Strip out the flag and its argument, and run rustc under the wrapper
        // program named by the argument.
        args.remove(pos);
        let wrapper = args.remove(pos);
        let wrapper = wrapper.to_str().unwrap();

        raise_priority();

        match wrapper {
            "perf-stat" => {
                let mut cmd = Command::new("perf");
                let has_perf = cmd.output().is_ok();
                assert!(has_perf);
                cmd.arg("stat")
                    .arg("-x;")
                    .arg("-e")
                    .arg("instructions:u,cycles:u,task-clock,cpu-clock,faults")
                    .arg("--log-fd")
                    .arg("1")
                    .arg(&rustc)
                    .args(&args);

                let start = Instant::now();
                assert!(cmd.status().expect("failed to spawn").success());
                let dur = start.elapsed();
                print_memory();
                print_time(dur);
            }

            "time-passes" => {
                let mut cmd = Command::new(&rustc);
                cmd.arg("-Ztime-passes")
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "perf-record" => {
                let mut cmd = Command::new("perf");
                let has_perf = cmd.output().is_ok();
                assert!(has_perf);
                cmd.arg("record")
                    .arg("--call-graph=dwarf")  // njn: best?
                    .arg("--output=perf")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "cachegrind" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);

                // With --cache-sim=no and --branch-sim=no, Cachegrind just
                // collects instruction counts.
                cmd.arg("--tool=cachegrind")
                    .arg("--cache-sim=no")
                    .arg("--branch-sim=no")
                    .arg("--cachegrind-out-file=cgout")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "callgrind" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);

                // With --cache-sim=no and --branch-sim=no, Callgrind just
                // collects instruction counts.
                cmd.arg("--tool=callgrind")
                    .arg("--cache-sim=no")
                    .arg("--branch-sim=no")
                    .arg("--callgrind-out-file=clgout")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "dhat" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);
                cmd.arg("--tool=exp-dhat")
                    .arg("--show-top-n=500")
                    .arg("--num-callers=4")
                    .arg("--sort-by=tot-blocks-allocd")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "massif" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);
                cmd.arg("--tool=massif")
                    .arg("--heap-admin=0")
                    .arg("--depth=15")
                    .arg("--threshold=0.2")
                    .arg("--massif-out-file=msout")
                    .arg("--alloc-fn=__rdl_alloc")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "eprintln" => {
                let mut cmd = Command::new(&rustc);
                cmd.args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            _ => {
                panic!("unknown wrapper: {}", wrapper);
            }
        }

    } else {
        let mut cmd = Command::new(&rustc);
        cmd.args(&args);
        exec(&mut cmd);
    }
}

#[cfg(unix)]
fn exec(cmd: &mut Command) -> ! {
    use std::os::unix::prelude::*;
    let error = cmd.exec();
    panic!("failed to exec: {}", error);
}

#[cfg(unix)]
fn raise_priority() {
    extern crate libc;

    unsafe {
        // Try to reduce jitter in wall time by increasing our priority to the
        // maximum
        for i in (1..21).rev() {
            let r = libc::setpriority(libc::PRIO_PROCESS as _, libc::getpid() as libc::id_t, -i);
            if r == 0 {
                break;
            }
        }
    }
}

#[cfg(unix)]
fn print_memory() {
    extern crate libc;

    use std::mem;

    unsafe {
        let mut usage = mem::zeroed();
        let r = libc::getrusage(libc::RUSAGE_CHILDREN, &mut usage);
        if r == 0 {
            // for explanation of all the semicolons, see `print_time` below
            println!("{};;max-rss;3;100.00", usage.ru_maxrss);
        }
    }
}

fn print_time(dur: Duration) {
    // Format output the same as `perf stat` in CSV mode, explained at
    // http://man7.org/linux/man-pages/man1/perf-stat.1.html#CSV_FORMAT
    //
    // tl;dr; it's:
    //
    //      $value ; $unit ; $name ; $runtime ; $pct
    println!(
        "{}.{:09};;wall-time;4;100.00",
        dur.as_secs(),
        dur.subsec_nanos()
    );
}

#[cfg(windows)]
fn raise_priority() {}

#[cfg(windows)]
fn print_memory() {}
<commit_msg>Tweak the `perf-record` invocation.<commit_after>use std::env;
use std::process::Command;
use std::time::{Duration, Instant};

fn main() {
    let mut args = env::args_os().skip(1).collect::<Vec<_>>();
    let rustc = env::var_os("RUSTC_REAL").unwrap();

    if let Some(pos) = args.iter().position(|arg| arg == "--wrap-rustc-with") {
        // Strip out the flag and its argument, and run rustc under the wrapper
        // program named by the argument.
        args.remove(pos);
        let wrapper = args.remove(pos);
        let wrapper = wrapper.to_str().unwrap();

        raise_priority();

        match wrapper {
            "perf-stat" => {
                let mut cmd = Command::new("perf");
                let has_perf = cmd.output().is_ok();
                assert!(has_perf);
                cmd.arg("stat")
                    .arg("-x;")
                    .arg("-e")
                    .arg("instructions:u,cycles:u,task-clock,cpu-clock,faults")
                    .arg("--log-fd")
                    .arg("1")
                    .arg(&rustc)
                    .args(&args);

                let start = Instant::now();
                assert!(cmd.status().expect("failed to spawn").success());
                let dur = start.elapsed();
                print_memory();
                print_time(dur);
            }

            "time-passes" => {
                let mut cmd = Command::new(&rustc);
                cmd.arg("-Ztime-passes")
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "perf-record" => {
                let mut cmd = Command::new("perf");
                let has_perf = cmd.output().is_ok();
                assert!(has_perf);
                cmd.arg("record")
                    .arg("--call-graph=dwarf")
                    .arg("--output=perf")
                    .arg("--freq=99")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "cachegrind" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);

                // With --cache-sim=no and --branch-sim=no, Cachegrind just
                // collects instruction counts.
                cmd.arg("--tool=cachegrind")
                    .arg("--cache-sim=no")
                    .arg("--branch-sim=no")
                    .arg("--cachegrind-out-file=cgout")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "callgrind" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);

                // With --cache-sim=no and --branch-sim=no, Callgrind just
                // collects instruction counts.
                cmd.arg("--tool=callgrind")
                    .arg("--cache-sim=no")
                    .arg("--branch-sim=no")
                    .arg("--callgrind-out-file=clgout")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "dhat" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);
                cmd.arg("--tool=exp-dhat")
                    .arg("--show-top-n=500")
                    .arg("--num-callers=4")
                    .arg("--sort-by=tot-blocks-allocd")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "massif" => {
                let mut cmd = Command::new("valgrind");
                let has_valgrind = cmd.output().is_ok();
                assert!(has_valgrind);
                cmd.arg("--tool=massif")
                    .arg("--heap-admin=0")
                    .arg("--depth=15")
                    .arg("--threshold=0.2")
                    .arg("--massif-out-file=msout")
                    .arg("--alloc-fn=__rdl_alloc")
                    .arg(&rustc)
                    .args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            "eprintln" => {
                let mut cmd = Command::new(&rustc);
                cmd.args(&args);

                assert!(cmd.status().expect("failed to spawn").success());
            }

            _ => {
                panic!("unknown wrapper: {}", wrapper);
            }
        }

    } else {
        let mut cmd = Command::new(&rustc);
        cmd.args(&args);
        exec(&mut cmd);
    }
}

#[cfg(unix)]
fn exec(cmd: &mut Command) -> ! {
    use std::os::unix::prelude::*;
    let error = cmd.exec();
    panic!("failed to exec: {}", error);
}

#[cfg(unix)]
fn raise_priority() {
    extern crate libc;

    unsafe {
        // Try to reduce jitter in wall time by increasing our priority to the
        // maximum
        for i in (1..21).rev() {
            let r = libc::setpriority(libc::PRIO_PROCESS as _, libc::getpid() as libc::id_t, -i);
            if r == 0 {
                break;
            }
        }
    }
}

#[cfg(unix)]
fn print_memory() {
    extern crate libc;

    use std::mem;

    unsafe {
        let mut usage = mem::zeroed();
        let r = libc::getrusage(libc::RUSAGE_CHILDREN, &mut usage);
        if r == 0 {
            // for explanation of all the semicolons, see `print_time` below
            println!("{};;max-rss;3;100.00", usage.ru_maxrss);
        }
    }
}

fn print_time(dur: Duration) {
    // Format output the same as `perf stat` in CSV mode, explained at
    // http://man7.org/linux/man-pages/man1/perf-stat.1.html#CSV_FORMAT
    //
    // tl;dr; it's:
    //
    //      $value ; $unit ; $name ; $runtime ; $pct
    println!(
        "{}.{:09};;wall-time;4;100.00",
        dur.as_secs(),
        dur.subsec_nanos()
    );
}

#[cfg(windows)]
fn raise_priority() {}

#[cfg(windows)]
fn print_memory() {}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;

use clap::{App, Arg};
use std::fs::File;
use std::io::{BufReader, Read};

extern crate script_extractor;
use script_extractor::*;

fn main() {
    let args = App::new("script-extractor")
                   .version(&crate_version!())
                   .about("Parse movie scripts into a structured format")
                   .after_help("Input has to be in the format generated by \
                                poppler's 'pdftohtml -xml'.\n\
                                Reads from stdin if no file or '-' is specified.")
                   .arg(Arg::with_name("input-file")
                            .help("input file in poppler's xml format")
                            .index(1)
                            .validator(check_file_exists))
                   .arg(Arg::with_name("pages")
                            .help("Specify a single page or range of pages to extract")
                            .short("p")
                            .long("pages")
                            .takes_value(true)
                            .validator(|v| if extract_range(&v).is_some() {
                                               Ok(())
                                           } else {
                                               Err(format!("Invalid page range '{}'", v))
                                           }))
                   .get_matches();

    let mut input: Box<Read> = if let Some(input_file) = args.value_of("input-file") {
        if input_file == "-" {
            Box::new(BufReader::new(std::io::stdin()))
        } else {
            let file_reader = File::open(input_file).expect("Cannot open input-file");
            Box::new(BufReader::new(file_reader))
        }
    } else {
        Box::new(BufReader::new(std::io::stdin()))
    };

    let mut script = parse::parse_script(&mut input);

    // filter by pages if requested
    if let Some(range_string) = args.value_of("pages") {
        let range = extract_range(range_string).unwrap_or((0,u32::max_value()));
        script = filter_script(script, range);
    }

    serialize::xml::format_script(&script, &mut std::io::stdout()).unwrap();
}

/// Check if the given file exists. Also considers "-" as a valid file.
fn check_file_exists(file_name: String) -> Result<(), String> {
    if file_name == "-" {
        Ok(())
    } else if let Ok(metadata) = std::fs::metadata(&file_name) {
        if metadata.is_file() && !metadata.permissions().readonly() {
            Ok(())
        } else {
            Err(format!("Cannot read file '{}'", file_name))
        }
    } else {
        Err(format!("File '{}' not found", file_name))
    }
}
<commit_msg>Add cli flags for selection between xml and json output<commit_after>#[macro_use]
extern crate clap;

use clap::{App, Arg, ArgGroup};
use std::fs::File;
use std::io::{BufReader, Read};

extern crate script_extractor;
use script_extractor::*;

fn main() {
    let args = App::new("script-extractor")
                   .version(&crate_version!())
                   .about("Parse movie scripts into a structured format")
                   .after_help("Input has to be in the format generated by \
                                poppler's 'pdftohtml -xml'.\n\
                                Reads from stdin if no file or '-' is specified.")
                   .arg(Arg::with_name("input-file")
                            .help("input file in poppler's xml format")
                            .index(1)
                            .validator(check_file_exists))
                   .arg(Arg::with_name("xml")
                            .long("xml")
                            .help("Output script in xml format"))
                   .arg(Arg::with_name("json")
                            .long("json")
                            .help("Output script in json format"))
                   .arg_group(ArgGroup::with_name("output-format")
                                       .add_all(&["xml", "json"])
                                       .required(true))
                   .arg(Arg::with_name("pages")
                            .help("Specify a single page or range of pages to extract")
                            .short("p")
                            .long("pages")
                            .takes_value(true)
                            .validator(|v| if extract_range(&v).is_some() {
                                               Ok(())
                                           } else {
                                               Err(format!("Invalid page range '{}'", v))
                                           }))
                   .get_matches();

    let mut input: Box<Read> = if let Some(input_file) = args.value_of("input-file") {
        if input_file == "-" {
            Box::new(BufReader::new(std::io::stdin()))
        } else {
            let file_reader = File::open(input_file).expect("Cannot open input-file");
            Box::new(BufReader::new(file_reader))
        }
    } else {
        Box::new(BufReader::new(std::io::stdin()))
    };

    let mut script = parse::parse_script(&mut input);

    // filter by pages if requested
    if let Some(range_string) = args.value_of("pages") {
        let range = extract_range(range_string).unwrap_or((0,u32::max_value()));
        script = filter_script(script, range);
    }

    if args.is_present("xml") {
        serialize::xml::format_script(&script, &mut std::io::stdout()).unwrap();
    } else if args.is_present("json") {
        serialize::json::format_script(&script, &mut std::io::stdout()).unwrap();
    }
}

/// Check if the given file exists. Also considers "-" as a valid file.
fn check_file_exists(file_name: String) -> Result<(), String> {
    if file_name == "-" {
        Ok(())
    } else if let Ok(metadata) = std::fs::metadata(&file_name) {
        if metadata.is_file() && !metadata.permissions().readonly() {
            Ok(())
        } else {
            Err(format!("Cannot read file '{}'", file_name))
        }
    } else {
        Err(format!("File '{}' not found", file_name))
    }
}
<|endoftext|>


<commit_before>extern crate onmf;
use onmf::helpers::{ToImage, Normalize};
use onmf::testimage_generator;
use onmf::testimage_generator::{magnify};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate rand;
use rand::{StdRng, SeedableRng};

fn main() {
    let mag_factor = 10;

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let steps = 6;
    let per_step = 1000;
    let nsamples = steps * per_step;
    let nobserved = 10 * 10;
    let nhidden = 20;

    let mut data = DMat::<f64>::new_zeros(nsamples, nobserved);

    // write on testimage into each row of data
    for (step, i, factor) in testimage_generator::testimages::<f64, _>(per_step, &mut rng) {
        let row = step * i;
        for (col, val) in factor.as_vec().iter().enumerate() {
            data[(row, col)] = val.clone();
        }
    }

    let mut nmf = onmf::OrthogonalNMF::<f64>::init_randomly(
        nhidden, nobserved, nsamples);

    for _ in 0..200 { nmf.iterate(&data) }

    // take a snapshot
    // read testimage out of each row of nmf.hidden
    for irow in 0..nhidden {
        let mut column = Vec::<f64>::new();
        for icol in 0..nobserved {
            column.push(nmf.hidden[(irow, icol)]);
        }
        let image = DMat::<f64>::from_col_vec(10, 10, &column[..]);
        // println!("image {}", irow);
        // println!("{:?}", image);
        magnify(image.normalize(), mag_factor).save_to_png(&format!("orthogonal-nmf-hidden-snapshot-{}.png", irow)[..]).unwrap();
    }
}
<commit_msg>ortho nmf: detect even more latent variables to avoid mixtures<commit_after>extern crate onmf;
use onmf::helpers::{ToImage, Normalize};
use onmf::testimage_generator;
use onmf::testimage_generator::{magnify};

extern crate nalgebra;
use self::nalgebra::{DMat};

extern crate rand;
use rand::{StdRng, SeedableRng};

fn main() {
    let mag_factor = 10;

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let steps = 6;
    let per_step = 1000;
    let nsamples = steps * per_step;
    let nobserved = 10 * 10;
    let nhidden = 20 + steps * 2;

    let mut data = DMat::<f64>::new_zeros(nsamples, nobserved);

    // write on testimage into each row of data
    for (step, i, factor) in testimage_generator::testimages::<f64, _>(per_step, &mut rng) {
        let row = step * i;
        for (col, val) in factor.as_vec().iter().enumerate() {
            data[(row, col)] = val.clone();
        }
    }

    let mut nmf = onmf::OrthogonalNMF::<f64>::init_randomly(
        nhidden, nobserved, nsamples);

    for _ in 0..200 { nmf.iterate(&data) }

    // take a snapshot
    // read testimage out of each row of nmf.hidden
    for irow in 0..nhidden {
        let mut column = Vec::<f64>::new();
        for icol in 0..nobserved {
            column.push(nmf.hidden[(irow, icol)]);
        }
        let image = DMat::<f64>::from_col_vec(10, 10, &column[..]);
        // println!("image {}", irow);
        // println!("{:?}", image);
        magnify(image.normalize(), mag_factor).save_to_png(&format!("orthogonal-nmf-hidden-snapshot-{}.png", irow)[..]).unwrap();
    }
}
<|endoftext|>


<commit_before>use rustc::hir::intravisit as visit;
use rustc::hir::*;
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::middle::expr_use_visitor::*;
use rustc::middle::mem_categorization::{cmt_, Categorization};
use rustc::ty::layout::LayoutOf;
use rustc::ty::{self, Ty};
use rustc::util::nodemap::HirIdSet;
use rustc::{declare_tool_lint, impl_lint_pass};
use syntax::source_map::Span;

use crate::utils::span_lint;

#[derive(Copy, Clone)]
pub struct BoxedLocal {
    pub too_large_for_stack: u64,
}

declare_clippy_lint! {
    /// **What it does:** Checks for usage of `Box<T>` where an unboxed `T` would
    /// work fine.
    ///
    /// **Why is this bad?** This is an unnecessary allocation, and bad for
    /// performance. It is only necessary to allocate if you wish to move the box
    /// into something.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// fn main() {
    ///     let x = Box::new(1);
    ///     foo(*x);
    ///     println!("{}", *x);
    /// }
    /// ```
    pub BOXED_LOCAL,
    perf,
    "using `Box<T>` where unnecessary"
}

fn is_non_trait_box(ty: Ty<'_>) -> bool {
    ty.is_box() && !ty.boxed_ty().is_trait()
}

struct EscapeDelegate<'a, 'tcx> {
    cx: &'a LateContext<'a, 'tcx>,
    set: HirIdSet,
    too_large_for_stack: u64,
}

impl_lint_pass!(BoxedLocal => [BOXED_LOCAL]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for BoxedLocal {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: visit::FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        hir_id: HirId,
    ) {
        // If the method is an impl for a trait, don't warn.
        let parent_id = cx.tcx.hir().get_parent_item(hir_id);
        let parent_node = cx.tcx.hir().find(parent_id);

        if let Some(Node::Item(item)) = parent_node {
            if let ItemKind::Impl(_, _, _, _, Some(..), _, _) = item.node {
                return;
            }
        }

        let mut v = EscapeDelegate {
            cx,
            set: HirIdSet::default(),
            too_large_for_stack: self.too_large_for_stack,
        };

        let fn_def_id = cx.tcx.hir().local_def_id(hir_id);
        let region_scope_tree = &cx.tcx.region_scope_tree(fn_def_id);
        ExprUseVisitor::new(
            &mut v,
            cx.tcx,
            fn_def_id,
            cx.param_env,
            region_scope_tree,
            cx.tables,
            None,
        )
        .consume_body(body);

        for node in v.set {
            span_lint(
                cx,
                BOXED_LOCAL,
                cx.tcx.hir().span(node),
                "local variable doesn't need to be boxed here",
            );
        }
    }
}

impl<'a, 'tcx> Delegate<'tcx> for EscapeDelegate<'a, 'tcx> {
    fn consume(&mut self, _: HirId, _: Span, cmt: &cmt_<'tcx>, mode: ConsumeMode) {
        if let Categorization::Local(lid) = cmt.cat {
            if let Move(DirectRefMove) | Move(CaptureMove) = mode {
                // moved out or in. clearly can't be localized
                self.set.remove(&lid);
            }
        }
    }
    fn matched_pat(&mut self, _: &Pat, _: &cmt_<'tcx>, _: MatchMode) {}
    fn consume_pat(&mut self, consume_pat: &Pat, cmt: &cmt_<'tcx>, _: ConsumeMode) {
        let map = &self.cx.tcx.hir();
        if map.is_argument(consume_pat.hir_id) {
            // Skip closure arguments
            if let Some(Node::Expr(..)) = map.find(map.get_parent_node(consume_pat.hir_id)) {
                return;
            }
            if is_non_trait_box(cmt.ty) && !self.is_large_box(cmt.ty) {
                self.set.insert(consume_pat.hir_id);
            }
            return;
        }
        if let Categorization::Rvalue(..) = cmt.cat {
            if let Some(Node::Stmt(st)) = map.find(map.get_parent_node(cmt.hir_id)) {
                if let StmtKind::Local(ref loc) = st.node {
                    if let Some(ref ex) = loc.init {
                        if let ExprKind::Box(..) = ex.node {
                            if is_non_trait_box(cmt.ty) && !self.is_large_box(cmt.ty) {
                                // let x = box (...)
                                self.set.insert(consume_pat.hir_id);
                            }
                            // TODO Box::new
                            // TODO vec![]
                            // TODO "foo".to_owned() and friends
                        }
                    }
                }
            }
        }
        if let Categorization::Local(lid) = cmt.cat {
            if self.set.contains(&lid) {
                // let y = x where x is known
                // remove x, insert y
                self.set.insert(consume_pat.hir_id);
                self.set.remove(&lid);
            }
        }
    }
    fn borrow(
        &mut self,
        _: HirId,
        _: Span,
        cmt: &cmt_<'tcx>,
        _: ty::Region<'_>,
        _: ty::BorrowKind,
        loan_cause: LoanCause,
    ) {
        if let Categorization::Local(lid) = cmt.cat {
            match loan_cause {
                // `x.foo()`
                // Used without autoderef-ing (i.e., `x.clone()`).
                LoanCause::AutoRef |

                // `&x`
                // `foo(&x)` where no extra autoref-ing is happening.
                LoanCause::AddrOf |

                // `match x` can move.
                LoanCause::MatchDiscriminant => {
                    self.set.remove(&lid);
                }

                // Do nothing for matches, etc. These can't escape.
                _ => {}
            }
        }
    }
    fn decl_without_init(&mut self, _: HirId, _: Span) {}
    fn mutate(&mut self, _: HirId, _: Span, _: &cmt_<'tcx>, _: MutateMode) {}
}

impl<'a, 'tcx> EscapeDelegate<'a, 'tcx> {
    fn is_large_box(&self, ty: Ty<'tcx>) -> bool {
        // Large types need to be boxed to avoid stack overflows.
        if ty.is_box() {
            self.cx.layout_of(ty.boxed_ty()).ok().map_or(0, |l| l.size.bytes()) > self.too_large_for_stack
        } else {
            false
        }
    }
}
<commit_msg>Auto merge of #4306 - mikerite:fix-breakage-20190730, r=flip1995<commit_after>use rustc::hir::intravisit as visit;
use rustc::hir::{self, *};
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::middle::expr_use_visitor::*;
use rustc::middle::mem_categorization::{cmt_, Categorization};
use rustc::ty::layout::LayoutOf;
use rustc::ty::{self, Ty};
use rustc::util::nodemap::HirIdSet;
use rustc::{declare_tool_lint, impl_lint_pass};
use syntax::source_map::Span;

use crate::utils::span_lint;

#[derive(Copy, Clone)]
pub struct BoxedLocal {
    pub too_large_for_stack: u64,
}

declare_clippy_lint! {
    /// **What it does:** Checks for usage of `Box<T>` where an unboxed `T` would
    /// work fine.
    ///
    /// **Why is this bad?** This is an unnecessary allocation, and bad for
    /// performance. It is only necessary to allocate if you wish to move the box
    /// into something.
    ///
    /// **Known problems:** None.
    ///
    /// **Example:**
    /// ```rust
    /// fn main() {
    ///     let x = Box::new(1);
    ///     foo(*x);
    ///     println!("{}", *x);
    /// }
    /// ```
    pub BOXED_LOCAL,
    perf,
    "using `Box<T>` where unnecessary"
}

fn is_non_trait_box(ty: Ty<'_>) -> bool {
    ty.is_box() && !ty.boxed_ty().is_trait()
}

struct EscapeDelegate<'a, 'tcx> {
    cx: &'a LateContext<'a, 'tcx>,
    set: HirIdSet,
    too_large_for_stack: u64,
}

impl_lint_pass!(BoxedLocal => [BOXED_LOCAL]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for BoxedLocal {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        _: visit::FnKind<'tcx>,
        _: &'tcx FnDecl,
        body: &'tcx Body,
        _: Span,
        hir_id: HirId,
    ) {
        // If the method is an impl for a trait, don't warn.
        let parent_id = cx.tcx.hir().get_parent_item(hir_id);
        let parent_node = cx.tcx.hir().find(parent_id);

        if let Some(Node::Item(item)) = parent_node {
            if let ItemKind::Impl(_, _, _, _, Some(..), _, _) = item.node {
                return;
            }
        }

        let mut v = EscapeDelegate {
            cx,
            set: HirIdSet::default(),
            too_large_for_stack: self.too_large_for_stack,
        };

        let fn_def_id = cx.tcx.hir().local_def_id(hir_id);
        let region_scope_tree = &cx.tcx.region_scope_tree(fn_def_id);
        ExprUseVisitor::new(
            &mut v,
            cx.tcx,
            fn_def_id,
            cx.param_env,
            region_scope_tree,
            cx.tables,
            None,
        )
        .consume_body(body);

        for node in v.set {
            span_lint(
                cx,
                BOXED_LOCAL,
                cx.tcx.hir().span(node),
                "local variable doesn't need to be boxed here",
            );
        }
    }
}

// TODO: Replace with Map::is_argument(..) when it's fixed
fn is_argument(map: &hir::map::Map<'_>, id: HirId) -> bool {
    match map.find(id) {
        Some(Node::Binding(_)) => (),
        _ => return false,
    }

    match map.find(map.get_parent_node(id)) {
        Some(Node::Arg(_)) => true,
        _ => false,
    }
}

impl<'a, 'tcx> Delegate<'tcx> for EscapeDelegate<'a, 'tcx> {
    fn consume(&mut self, _: HirId, _: Span, cmt: &cmt_<'tcx>, mode: ConsumeMode) {
        if let Categorization::Local(lid) = cmt.cat {
            if let Move(DirectRefMove) | Move(CaptureMove) = mode {
                // moved out or in. clearly can't be localized
                self.set.remove(&lid);
            }
        }
    }
    fn matched_pat(&mut self, _: &Pat, _: &cmt_<'tcx>, _: MatchMode) {}
    fn consume_pat(&mut self, consume_pat: &Pat, cmt: &cmt_<'tcx>, _: ConsumeMode) {
        let map = &self.cx.tcx.hir();
        if is_argument(map, consume_pat.hir_id) {
            // Skip closure arguments
            let parent_id = map.get_parent_node(consume_pat.hir_id);
            if let Some(Node::Expr(..)) = map.find(map.get_parent_node(parent_id)) {
                return;
            }

            if is_non_trait_box(cmt.ty) && !self.is_large_box(cmt.ty) {
                self.set.insert(consume_pat.hir_id);
            }
            return;
        }
        if let Categorization::Rvalue(..) = cmt.cat {
            if let Some(Node::Stmt(st)) = map.find(map.get_parent_node(cmt.hir_id)) {
                if let StmtKind::Local(ref loc) = st.node {
                    if let Some(ref ex) = loc.init {
                        if let ExprKind::Box(..) = ex.node {
                            if is_non_trait_box(cmt.ty) && !self.is_large_box(cmt.ty) {
                                // let x = box (...)
                                self.set.insert(consume_pat.hir_id);
                            }
                            // TODO Box::new
                            // TODO vec![]
                            // TODO "foo".to_owned() and friends
                        }
                    }
                }
            }
        }
        if let Categorization::Local(lid) = cmt.cat {
            if self.set.contains(&lid) {
                // let y = x where x is known
                // remove x, insert y
                self.set.insert(consume_pat.hir_id);
                self.set.remove(&lid);
            }
        }
    }
    fn borrow(
        &mut self,
        _: HirId,
        _: Span,
        cmt: &cmt_<'tcx>,
        _: ty::Region<'_>,
        _: ty::BorrowKind,
        loan_cause: LoanCause,
    ) {
        if let Categorization::Local(lid) = cmt.cat {
            match loan_cause {
                // `x.foo()`
                // Used without autoderef-ing (i.e., `x.clone()`).
                LoanCause::AutoRef |

                // `&x`
                // `foo(&x)` where no extra autoref-ing is happening.
                LoanCause::AddrOf |

                // `match x` can move.
                LoanCause::MatchDiscriminant => {
                    self.set.remove(&lid);
                }

                // Do nothing for matches, etc. These can't escape.
                _ => {}
            }
        }
    }
    fn decl_without_init(&mut self, _: HirId, _: Span) {}
    fn mutate(&mut self, _: HirId, _: Span, _: &cmt_<'tcx>, _: MutateMode) {}
}

impl<'a, 'tcx> EscapeDelegate<'a, 'tcx> {
    fn is_large_box(&self, ty: Ty<'tcx>) -> bool {
        // Large types need to be boxed to avoid stack overflows.
        if ty.is_box() {
            self.cx.layout_of(ty.boxed_ty()).ok().map_or(0, |l| l.size.bytes()) > self.too_large_for_stack
        } else {
            false
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    MultiStore,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let line = line.trim_end();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let value = serde_json::from_str(&line)?;
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => {
                        let value = vec![value];
                        serde_json::to_string(&value)?
                    },
                    Some(v) => match v {
                        Value::Json(e) => {
                            let mut current: serde_json::Value = serde_json::from_str(&e)?;
                            let current = current.as_array_mut().unwrap();
                            current.push(value);
                            serde_json::to_string(&current)?
                        },
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    println!("Commiting count: {}", count);
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    // let mut builder = Rkv::environment_builder::<Lmdb>();
    // builder.set_max_dbs(20);
    // builder.set_make_dir_if_needed(true);
    // builder.set_map_size(200000000000);
    // println!("Created BUILDER");
    // let arc = manager
    //     .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
    //     .with_context(|| "Failed to create arc")?;
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let parsed = serde_json::from_str(&raw)?;
                                to_update.push(parsed);
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::warn!("No data found for key {}", &buf);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {}", &id).into());
                        },
                    }

                    serde_json::to_writer(&mut writer, &data)?;
                    writeln!(&mut writer)?;
                    buf.clear();
                }
            },
        }
    }

    Ok(())
}
<commit_msg>Skip unneeded parsing<commit_after>use std::{
    collections::HashMap,
    error::Error,
    fs::File,
    io::{
        BufRead,
        BufReader,
    },
    path::{
        Path,
        PathBuf,
    },
    str,
    thread,
};

use serde_query::{
    DeserializeQuery,
    Query,
};

use anyhow::{
    Context,
    Result,
};

// use crossbeam_channel::{
//     unbounded,
//     Sender,
// };

use serde_json::Deserializer;

use rkv::{
    backend::{
        BackendEnvironmentBuilder,
        Lmdb,
        LmdbDatabase,
        LmdbEnvironment,
    },
    Manager,
    Rkv,
    SingleStore,
    StoreOptions,
    Value,
};

#[derive(Debug, DeserializeQuery)]
struct DocId {
    #[query(".id")]
    id: String,
}

pub struct Spec<'a> {
    path: &'a Path,
    allow_missing: bool,
    commit_size: usize,
    threads: usize,
}

impl<'a> Spec<'a> {
    pub fn new(path: &Path) -> Spec {
        Spec {
            path,
            allow_missing: false,
            commit_size: 1_000_000usize,
            threads: 4,
        }
    }

    pub fn set_allow_missing(&mut self, allow_missing: bool) -> () {
        self.allow_missing = allow_missing;
    }

    pub fn set_commit_size(&mut self, commit_size: usize) -> () {
        self.commit_size = commit_size;
    }
}

pub fn index(spec: &Spec, data_type: &str, filename: &Path) -> Result<(), Box<dyn Error>> {
    let mut reader = rnc_utils::buf_reader(&filename)?;

    let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    let mut builder = Rkv::environment_builder::<Lmdb>();
    builder.set_max_dbs(20);
    builder.set_make_dir_if_needed(true);
    builder.set_map_size(200000000000);
    let arc = manager
        .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
        .with_context(|| "Failed to create arc")?;
    let env = arc.read().unwrap();
    let store = env
        .open_single(data_type, StoreOptions::create())
        .with_context(|| format!("Failed to open store {:?}", data_type))?;

    let mut writer = env.write().with_context(|| "Failed to create writer")?;
    let mut buf = String::new();
    let mut count = 0usize;
    loop {
        match reader.read_line(&mut buf)? {
            0 => break,
            _ => {
                let line = buf.replace("\\\\", "\\");
                let value = line.trim_end().to_owned();
                let id: DocId = serde_json::from_str::<Query<DocId>>(&line)?.into();
                let existing = store.get(&writer, &id.id)?;
                let to_add = match existing {
                    None => value,
                    Some(v) => match v {
                        Value::Json(e) => format!("{}{}", value, e),
                        _ => panic!("Invalid existing data, should never happen"),
                    },
                };

                store.put(&mut writer, &id.id, &Value::Json(&to_add))?;
                buf.clear();
                count += 1;
                if count % spec.commit_size == 0 {
                    println!("Commiting count: {}", count);
                    writer
                        .commit()
                        .with_context(|| format!("Failed to write data from {:?}", filename))?;
                    writer = env.write().with_context(|| "Could not create for new chunk")?;
                }
            },
        }
    }
    writer.commit().with_context(|| format!("Failed to write data from {:?}", filename))?;

    Ok(())
}

// fn path_as_column_name(path: &Path) -> String {
//     path.file_stem().unwrap().to_str().unwrap().to_string().replace("-", "_")
// }

// fn send_file_lines(path: &Path, sender: Sender<(String, DocId, String)>) {
//     let file = File::open(path).unwrap();
//     let data_type = path_as_column_name(&path);
//     let mut reader = BufReader::new(file);
//     let mut buf = String::new();
//     loop {
//         match reader.read_line(&mut buf).unwrap() {
//             0 => break,
//             _ => {
//                 let line = buf.replace("\\\\", "\\");
//                 let id: DocId =
// serde_json::from_str::<Query<DocId>>(&line).unwrap().into();
// sender.send((data_type.to_string(), id, line.to_string())).unwrap();
// buf.clear();             },
//         }
//     }
// }

// pub fn index_files(spec: &Spec, filename: &Path) -> Result<(), Box<dyn Error>> {
//     let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
//     let arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>)?;
//     let env = arc.read();

//     // let stores = HashMap::<String, Multi>
//     let reader = rnc_utils::buf_reader(&filename)?;
//     let (json_sender, json_reciever) = unbounded();
//     for line in reader.lines() {
//         let line = line?;
//         let filename = line.trim_end();
//         let path = PathBuf::from(filename);
//         let sender = json_sender.clone();
//         let store = env
//             .open_multi(&data_type, StoreOptions::create())
//             .with_context(|| format!("Failed to open store {:?}", data_type))?;
//         stores.insert(name, store);
//         thread::spawn(move || send_file_lines(&path, sender));
//     }

//     let writer = env.write().with_context(|| "Failed to create writer")?;

//     for (data_type, id, data) in json_reciever {
//         let store = stores.get(&data_type)?;
//         store
//             .put(&mut writer, id.id, &Value::Json(data))
//             .with_context(|| format!("Failed to write data for {:}/{:?}", &data_type,
// &id))?;     }

//     Ok(())
// }

pub fn lookup(spec: &Spec, key_file: &Path, output: &Path) -> Result<(), Box<dyn Error>> {
    let mut writer = rnc_utils::buf_writer(&output)
        .with_context(|| format!("Could not open {:?} for writing", &output))?;
    let mut keys = rnc_utils::buf_reader(&key_file)
        .with_context(|| format!("Could not open key file {:?}", &key_file))?;

    // let mut manager = Manager::<LmdbEnvironment>::singleton().write()?;
    // let mut builder = Rkv::environment_builder::<Lmdb>();
    // builder.set_max_dbs(20);
    // builder.set_make_dir_if_needed(true);
    // builder.set_map_size(200000000000);
    // println!("Created BUILDER");
    // let arc = manager
    //     .get_or_create_from_builder(spec.path, builder, Rkv::from_builder::<Lmdb>)
    //     .with_context(|| "Failed to create arc")?;
    // let env = arc.read().unwrap();

    let mut manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
    let created_arc = manager.get_or_create(spec.path, Rkv::new::<Lmdb>).unwrap();
    let env = created_arc.read().unwrap();

    let names = env.get_dbs()?.iter().filter_map(|db| db.to_owned()).collect::<Vec<String>>();

    let mut stores: Vec<(String, SingleStore<LmdbDatabase>)> = Vec::new();
    for name in names {
        let store = env
            .open_single(name.as_str(), StoreOptions::default())
            .with_context(|| format!("Failed to open store {:?}", &name))?;
        stores.push((name.to_owned(), store));
    }

    let reader = env.read().with_context(|| "Failed to create writer")?;

    let mut buf = String::new();
    loop {
        match keys.read_line(&mut buf)? {
            0 => break,
            _ => {
                let mut data = HashMap::new();
                let id = buf.trim_end().to_string();
                data.insert("id", serde_json::Value::String(id.to_string()));
                let mut seen = false;
                for (name, store) in &stores {
                    let default = serde_json::Value::Array(Vec::new());
                    let to_update = data.entry(&name).or_insert(default).as_array_mut().unwrap();
                    let current = store
                        .get(&reader, &id)
                        .with_context(|| format!("Failed to get results for {:?}", &id))?;

                    match current {
                        None => (),
                        Some(v) => match v {
                            Value::Json(raw) => {
                                seen = true;
                                let values = Deserializer::from_str(raw).into_iter::<serde_json::Value>();
                                for value in values {
                                    to_update.push(value?);
                                }
                            },
                            _ => panic!("Invalid lookup data"),
                        },
                    }

                    match (seen, spec.allow_missing) {
                        (true, _) => (),
                        (false, true) => {
                            log::warn!("No data found for key {}", &buf);
                            continue;
                        },
                        (false, false) => {
                            return Err(format!("No data found for key {}", &id).into());
                        },
                    }

                    serde_json::to_writer(&mut writer, &data)?;
                    writeln!(&mut writer)?;
                    buf.clear();
                }
            },
        }
    }

    Ok(())
}
<|endoftext|>