<commit_before>//! A data structure for tracking source positions in language implementations, inspired by the
//! [CodeMap type in rustc's libsyntax](https://github.com/rust-lang/rust/blob/master/src/libsyntax/codemap.rs).
//!
//! The `CodeMap` tracks all source files and maps positions within them to linear indexes as if all
//! source files were concatenated. This allows a source position to be represented by a small
//! 32-bit `Pos` indexing into the `CodeMap`, under the assumption that the total amount of parsed
//! source code will not exceed 4GiB. The `CodeMap` can look up the source file, line, and column
//! of a `Pos` or `Span`, as well as provide source code snippets for error reporting.
//!
//! # Example
//! ```
//! use codemap::CodeMap;
//! let mut codemap = CodeMap::new();
//! let file = codemap.add_file("test.rs".to_string(), "fn test(){\n    println!(\"Hello\");\n}\n".to_string());
//! let string_literal_span = file.span.subspan(24, 31);
//!
//! let location = codemap.look_up_span(string_literal_span);
//! assert_eq!(location.file.name(), "test.rs");
//! assert_eq!(location.begin.line, 1);
//! assert_eq!(location.begin.column, 13);
//! assert_eq!(location.end.line, 1);
//! assert_eq!(location.end.column, 20);
//! ```

use std::cmp::{self, Ordering};
use std::ops::{Add, Sub, Deref};
use std::fmt;
use std::sync::Arc;
use std::hash::{Hash, Hasher};

/// A small, `Copy`, value representing a position in a `CodeMap`'s file.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Pos(u32);

impl Add<u64> for Pos {
    type Output = Pos;
    fn add(self, other: u64) -> Pos {
        Pos(self.0 + other as u32)
    }
}

impl Sub<Pos> for Pos {
    type Output = u64;
    fn sub(self, other: Pos) -> u64 {
        (self.0 - other.0) as u64
    }
}

/// A range of text within a CodeMap.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub struct Span {
    /// The position in the codemap representing the first byte of the span.
    low: Pos,

    /// The position after the last byte of the span.
    high: Pos,
}

impl Span {
    /// Makes a span from offsets relative to the start of this span.
    ///
    /// # Panics
    ///   * If `end < begin`
    ///   * If `end` is beyond the length of the span
    pub fn subspan(&self, begin: u64, end: u64) -> Span {
        assert!(end >= begin);
        assert!(self.low + end <= self.high);
        Span {
            low: self.low + begin,
            high: self.low + end,
        }
    }

    /// Checks if a span is contained within this span.
    pub fn contains(&self, other: Span) -> bool {
        self.low <= other.low && self.high >= other.high
    }

    /// The position in the codemap representing the first byte of the span.
    pub fn low(&self) -> Pos {
        self.low
    }

    /// The position after the last byte of the span.
    pub fn high(&self) -> Pos {
        self.high
    }

    /// The length in bytes of the text of the span
    pub fn len(&self) -> u64 {
        self.high - self.low
    }

    /// Create a span that encloses both `self` and `other`.
    pub fn merge(&self, other: Span) -> Span {
        Span {
            low: cmp::min(self.low, other.low),
            high: cmp::max(self.high, other.high),
        }
    }
}

/// Associate a Span with a value of arbitrary type (e.g. an AST node).
#[derive(Clone, PartialEq, Eq, Hash, Debug, Copy)]
pub struct Spanned<T> {
    pub node: T,
    pub span: Span,
}

impl<T> Deref for Spanned<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.node
    }
}

/// A data structure recording source code files for position lookup.
pub struct CodeMap {
    files: Vec<Arc<File>>,
    _allow_priv: ()
}

impl CodeMap {
    /// Creates an empty `CodeMap`.
    pub fn new() -> CodeMap {
        CodeMap {
            files: vec![],
            _allow_priv: ()
        }
    }

    /// Adds a file with the given name and contents.
    ///
    /// Use the returned `File` and its `.span` property to create `Spans`
    /// representing substrings of the file.
    pub fn add_file(&mut self, name: String, source: String) -> Arc<File> {
        let low = self.end_pos() + 1;
        let high = low + source.len() as u64;
        let mut lines = vec![low];
        lines.extend(source.match_indices('\n').map(|(p, _)| { low + (p + 1) as u64 }));

        let file = Arc::new(File {
            span: Span { low: low, high: high },
            name: name,
            source: source,
            lines: lines,
        });

        self.files.push(file.clone());
        file
    }

    fn end_pos(&self) -> Pos {
        self.files.last().map(|x| x.span.high).unwrap_or(Pos(0))
    }

    /// Looks up the `File` that contains the specified position.
    pub fn find_file(&self, pos: Pos) -> &Arc<File> {
        self.files.binary_search_by(|file| {
            if file.span.high < pos {
                Ordering::Less
            } else if file.span.low > pos {
                Ordering::Greater
            } else {
                Ordering::Equal
            }
        }).ok().map(|i| &self.files[i]).expect("Mapping unknown source location")
    }

    /// Gets the file, line, and column represented by a `Pos`.
    pub fn look_up_pos(&self, pos: Pos) -> Loc {
        let file = self.find_file(pos);
        let position = file.find_line_col(pos);
        Loc { file: file.clone(), position }
    }

    /// Gets the file and its line and column ranges represented by a `Span`.
    pub fn look_up_span(&self, span: Span) -> SpanLoc {
        let file = self.find_file(span.low);
        let begin = file.find_line_col(span.low);
        let end = file.find_line_col(span.high);
        SpanLoc { file: file.clone(), begin, end }
    }
}

/// A `CodeMap`'s record of a source file.
pub struct File {
    /// The span representing the entire file.
    pub span: Span,

    /// The filename as it would be displayed in an error message.
    name: String,

    /// Contents of the file.
    source: String,

    /// Byte positions of line beginnings.
    lines: Vec<Pos>,
}

impl File {
    /// Gets the name of the file
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Gets the line number of a Pos.
    ///
    /// The lines are 0-indexed (first line is numbered 0)
    ///
    /// # Panics
    ///
    ///  * If `pos` is not within this file's span
    pub fn find_line(&self, pos: Pos) -> usize {
        assert!(pos >= self.span.low);
        assert!(pos <= self.span.high);
        (match self.lines.binary_search(&pos) {
            Ok(i) => i,
            Err(i) => i - 1,
        })
    }

    /// Gets the line and column of a Pos.
    ///
    /// # Panics
    ///
    /// * If `pos` is not with this file's span
    /// * If `pos` points to a byte in the middle of a UTF-8 character
    pub fn find_line_col(&self, pos: Pos) -> LineCol {
        let line = self.find_line(pos);
        let line_span = self.line_span(line);
        let byte_col = pos - line_span.low;
        let col = self.source_slice(line_span)[..byte_col as usize].chars().count();

        LineCol{ line: line, column: col }
    }

    /// Gets the full source text of the file
    pub fn source(&self) -> &str {
        &self.source
    }

    /// Gets the source text of a Span.
    ///
    /// # Panics
    ///
    ///   * If `span` is not entirely within this file.
    pub fn source_slice(&self, span: Span) -> &str {
        assert!(self.span.contains(span));
        &self.source[((span.low - self.span.low) as usize)..((span.high - self.span.low) as usize)]
    }

    /// Gets the span representing a line by line number.
    ///
    /// The line number is 0-indexed (first line is numbered 0). The returned span includes the
    /// line terminator.
    ///
    /// # Panics
    ///
    ///  * If the line number is out of range
    pub fn line_span(&self, line: usize) -> Span {
        assert!(line < self.lines.len());
        Span {
            low: self.lines[line],
            high: *self.lines.get(line + 1).unwrap_or(&self.span.high)
        }
    }

    /// Gets the source text of a line.
    ///
    /// The string returned does not include the terminating \r or \n characters.
    ///
    /// # Panics
    ///
    ///  * If the line number is out of range
    pub fn source_line(&self, line: usize) -> &str {
        self.source_slice(self.line_span(line)).trim_right_matches(&['\n', '\r'][..])
    }

    /// Gets the number of lines in the file
    pub fn num_lines(&self) -> usize {
        self.lines.len()
    }
}

impl fmt::Debug for File {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "File({:?})", self.name)
    }
}

impl PartialEq for File {
    /// Compares by identity
    fn eq(&self, other: &File) -> bool {
        self as *const _ == other as *const _
    }
}

impl Eq for File {}

impl Hash for File {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.span.hash(hasher);
    }
}

/// A line and column.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub struct LineCol {
    /// The line number within the file (0-indexed).
    pub line: usize,

    /// The column within the line (0-indexed).
    pub column: usize,
}

/// A file, and a line and column within it.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct Loc {
    pub file: Arc<File>,
    pub position: LineCol,
}

impl fmt::Display for Loc {
    /// Formats the location as `filename:line:column`, with a 1-indexed
    /// line and column.
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}:{}:{}", self.file.name, self.position.line+1, self.position.column+1)
    }
}

/// A file, and a line and column range within it.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct SpanLoc {
    pub file: Arc<File>,
    pub begin: LineCol,
    pub end: LineCol,
}

impl fmt::Display for SpanLoc {
    /// Formats the span as `filename:start_line:start_column: end_line:end_column`,
    /// or if the span is zero-length, `filename:line:column`, with a 1-indexed line and column.
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        if self.begin == self.end {
            write!(f, "{}:{}:{}", self.file.name, self.begin.line+1, self.begin.column+1)
        } else {
            write!(f, "{}:{}:{}: {}:{}", self.file.name, self.begin.line+1, self.begin.column+1, self.end.line+1, self.end.column+1)
        }
    }
}

#[test]
fn test_codemap() {
    let mut codemap = CodeMap::new();
    let f1 = codemap.add_file("test1.rs".to_string(), "abcd\nefghij\nqwerty".to_string());
    let f2 = codemap.add_file("test2.rs".to_string(), "foo\nbar".to_string());

    assert_eq!(codemap.find_file(f1.span.low()).name(), "test1.rs");
    assert_eq!(codemap.find_file(f1.span.high()).name(), "test1.rs");
    assert_eq!(codemap.find_file(f2.span.low()).name(), "test2.rs");
    assert_eq!(codemap.find_file(f2.span.high()).name(), "test2.rs");

    let x = f1.span.subspan(5, 10);
    let f = codemap.find_file(x.low);
    assert_eq!(f.name, "test1.rs");
    assert_eq!(f.find_line_col(f.span.low()), LineCol { line: 0, column: 0 });
    assert_eq!(f.find_line_col(f.span.low() + 4), LineCol { line: 0, column: 4 });
    assert_eq!(f.find_line_col(f.span.low() + 5), LineCol { line: 1, column: 0 });
    assert_eq!(f.find_line_col(f.span.low() + 16), LineCol { line: 2, column: 4 });

    let x = f2.span.subspan(4, 7);
    assert_eq!(codemap.find_file(x.low()).name(), "test2.rs");
    assert_eq!(codemap.find_file(x.high()).name(), "test2.rs");
}

#[test]
fn test_issue2() {
    let mut codemap = CodeMap::new();
    let content = "a \nxyz\r\n";
    let file = codemap.add_file("<test>".to_owned(), content.to_owned());

    let span = file.span.subspan(2, 3);
    assert_eq!(codemap.look_up_span(span), SpanLoc {
        file: file.clone(),
        begin: LineCol { line: 0, column: 2 },
        end: LineCol { line: 1, column: 0 }
    });

    assert_eq!(file.source_line(0), "a ");
    assert_eq!(file.source_line(1), "xyz");
    assert_eq!(file.source_line(2), "");
}

#[test]
fn test_multibyte() {
    let mut codemap = CodeMap::new();
    let content = "65°00′N 18°00′W 汉语\n🔬";
    let file = codemap.add_file("<test>".to_owned(), content.to_owned());

    assert_eq!(codemap.look_up_pos(file.span.low() + 21), Loc { file: file.clone(), position: LineCol { line: 0, column: 15 } });
    assert_eq!(codemap.look_up_pos(file.span.low() + 28), Loc { file: file.clone(), position: LineCol { line: 0, column: 18 } });
    assert_eq!(codemap.look_up_pos(file.span.low() + 33), Loc { file: file.clone(), position: LineCol { line: 1, column: 1 } });
}
<commit_msg>impl Default for CodeMap<commit_after>//! A data structure for tracking source positions in language implementations, inspired by the
//! [CodeMap type in rustc's libsyntax](https://github.com/rust-lang/rust/blob/master/src/libsyntax/codemap.rs).
//!
//! The `CodeMap` tracks all source files and maps positions within them to linear indexes as if all
//! source files were concatenated. This allows a source position to be represented by a small
//! 32-bit `Pos` indexing into the `CodeMap`, under the assumption that the total amount of parsed
//! source code will not exceed 4GiB. The `CodeMap` can look up the source file, line, and column
//! of a `Pos` or `Span`, as well as provide source code snippets for error reporting.
//!
//! # Example
//! ```
//! use codemap::CodeMap;
//! let mut codemap = CodeMap::new();
//! let file = codemap.add_file("test.rs".to_string(), "fn test(){\n    println!(\"Hello\");\n}\n".to_string());
//! let string_literal_span = file.span.subspan(24, 31);
//!
//! let location = codemap.look_up_span(string_literal_span);
//! assert_eq!(location.file.name(), "test.rs");
//! assert_eq!(location.begin.line, 1);
//! assert_eq!(location.begin.column, 13);
//! assert_eq!(location.end.line, 1);
//! assert_eq!(location.end.column, 20);
//! ```

use std::cmp::{self, Ordering};
use std::ops::{Add, Sub, Deref};
use std::fmt;
use std::sync::Arc;
use std::hash::{Hash, Hasher};

/// A small, `Copy`, value representing a position in a `CodeMap`'s file.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Debug)]
pub struct Pos(u32);

impl Add<u64> for Pos {
    type Output = Pos;
    fn add(self, other: u64) -> Pos {
        Pos(self.0 + other as u32)
    }
}

impl Sub<Pos> for Pos {
    type Output = u64;
    fn sub(self, other: Pos) -> u64 {
        (self.0 - other.0) as u64
    }
}

/// A range of text within a CodeMap.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub struct Span {
    /// The position in the codemap representing the first byte of the span.
    low: Pos,

    /// The position after the last byte of the span.
    high: Pos,
}

impl Span {
    /// Makes a span from offsets relative to the start of this span.
    ///
    /// # Panics
    ///   * If `end < begin`
    ///   * If `end` is beyond the length of the span
    pub fn subspan(&self, begin: u64, end: u64) -> Span {
        assert!(end >= begin);
        assert!(self.low + end <= self.high);
        Span {
            low: self.low + begin,
            high: self.low + end,
        }
    }

    /// Checks if a span is contained within this span.
    pub fn contains(&self, other: Span) -> bool {
        self.low <= other.low && self.high >= other.high
    }

    /// The position in the codemap representing the first byte of the span.
    pub fn low(&self) -> Pos {
        self.low
    }

    /// The position after the last byte of the span.
    pub fn high(&self) -> Pos {
        self.high
    }

    /// The length in bytes of the text of the span
    pub fn len(&self) -> u64 {
        self.high - self.low
    }

    /// Create a span that encloses both `self` and `other`.
    pub fn merge(&self, other: Span) -> Span {
        Span {
            low: cmp::min(self.low, other.low),
            high: cmp::max(self.high, other.high),
        }
    }
}

/// Associate a Span with a value of arbitrary type (e.g. an AST node).
#[derive(Clone, PartialEq, Eq, Hash, Debug, Copy)]
pub struct Spanned<T> {
    pub node: T,
    pub span: Span,
}

impl<T> Deref for Spanned<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.node
    }
}

/// A data structure recording source code files for position lookup.
#[derive(Default)]
pub struct CodeMap {
    files: Vec<Arc<File>>,
    _allow_priv: ()
}

impl CodeMap {
    /// Creates an empty `CodeMap`.
    pub fn new() -> CodeMap {
        Default::default()
    }

    /// Adds a file with the given name and contents.
    ///
    /// Use the returned `File` and its `.span` property to create `Spans`
    /// representing substrings of the file.
    pub fn add_file(&mut self, name: String, source: String) -> Arc<File> {
        let low = self.end_pos() + 1;
        let high = low + source.len() as u64;
        let mut lines = vec![low];
        lines.extend(source.match_indices('\n').map(|(p, _)| { low + (p + 1) as u64 }));

        let file = Arc::new(File {
            span: Span { low: low, high: high },
            name: name,
            source: source,
            lines: lines,
        });

        self.files.push(file.clone());
        file
    }

    fn end_pos(&self) -> Pos {
        self.files.last().map(|x| x.span.high).unwrap_or(Pos(0))
    }

    /// Looks up the `File` that contains the specified position.
    pub fn find_file(&self, pos: Pos) -> &Arc<File> {
        self.files.binary_search_by(|file| {
            if file.span.high < pos {
                Ordering::Less
            } else if file.span.low > pos {
                Ordering::Greater
            } else {
                Ordering::Equal
            }
        }).ok().map(|i| &self.files[i]).expect("Mapping unknown source location")
    }

    /// Gets the file, line, and column represented by a `Pos`.
    pub fn look_up_pos(&self, pos: Pos) -> Loc {
        let file = self.find_file(pos);
        let position = file.find_line_col(pos);
        Loc { file: file.clone(), position }
    }

    /// Gets the file and its line and column ranges represented by a `Span`.
    pub fn look_up_span(&self, span: Span) -> SpanLoc {
        let file = self.find_file(span.low);
        let begin = file.find_line_col(span.low);
        let end = file.find_line_col(span.high);
        SpanLoc { file: file.clone(), begin, end }
    }
}

/// A `CodeMap`'s record of a source file.
pub struct File {
    /// The span representing the entire file.
    pub span: Span,

    /// The filename as it would be displayed in an error message.
    name: String,

    /// Contents of the file.
    source: String,

    /// Byte positions of line beginnings.
    lines: Vec<Pos>,
}

impl File {
    /// Gets the name of the file
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Gets the line number of a Pos.
    ///
    /// The lines are 0-indexed (first line is numbered 0)
    ///
    /// # Panics
    ///
    ///  * If `pos` is not within this file's span
    pub fn find_line(&self, pos: Pos) -> usize {
        assert!(pos >= self.span.low);
        assert!(pos <= self.span.high);
        (match self.lines.binary_search(&pos) {
            Ok(i) => i,
            Err(i) => i - 1,
        })
    }

    /// Gets the line and column of a Pos.
    ///
    /// # Panics
    ///
    /// * If `pos` is not with this file's span
    /// * If `pos` points to a byte in the middle of a UTF-8 character
    pub fn find_line_col(&self, pos: Pos) -> LineCol {
        let line = self.find_line(pos);
        let line_span = self.line_span(line);
        let byte_col = pos - line_span.low;
        let col = self.source_slice(line_span)[..byte_col as usize].chars().count();

        LineCol{ line: line, column: col }
    }

    /// Gets the full source text of the file
    pub fn source(&self) -> &str {
        &self.source
    }

    /// Gets the source text of a Span.
    ///
    /// # Panics
    ///
    ///   * If `span` is not entirely within this file.
    pub fn source_slice(&self, span: Span) -> &str {
        assert!(self.span.contains(span));
        &self.source[((span.low - self.span.low) as usize)..((span.high - self.span.low) as usize)]
    }

    /// Gets the span representing a line by line number.
    ///
    /// The line number is 0-indexed (first line is numbered 0). The returned span includes the
    /// line terminator.
    ///
    /// # Panics
    ///
    ///  * If the line number is out of range
    pub fn line_span(&self, line: usize) -> Span {
        assert!(line < self.lines.len());
        Span {
            low: self.lines[line],
            high: *self.lines.get(line + 1).unwrap_or(&self.span.high)
        }
    }

    /// Gets the source text of a line.
    ///
    /// The string returned does not include the terminating \r or \n characters.
    ///
    /// # Panics
    ///
    ///  * If the line number is out of range
    pub fn source_line(&self, line: usize) -> &str {
        self.source_slice(self.line_span(line)).trim_right_matches(&['\n', '\r'][..])
    }

    /// Gets the number of lines in the file
    pub fn num_lines(&self) -> usize {
        self.lines.len()
    }
}

impl fmt::Debug for File {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "File({:?})", self.name)
    }
}

impl PartialEq for File {
    /// Compares by identity
    fn eq(&self, other: &File) -> bool {
        self as *const _ == other as *const _
    }
}

impl Eq for File {}

impl Hash for File {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.span.hash(hasher);
    }
}

/// A line and column.
#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
pub struct LineCol {
    /// The line number within the file (0-indexed).
    pub line: usize,

    /// The column within the line (0-indexed).
    pub column: usize,
}

/// A file, and a line and column within it.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct Loc {
    pub file: Arc<File>,
    pub position: LineCol,
}

impl fmt::Display for Loc {
    /// Formats the location as `filename:line:column`, with a 1-indexed
    /// line and column.
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}:{}:{}", self.file.name, self.position.line+1, self.position.column+1)
    }
}

/// A file, and a line and column range within it.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct SpanLoc {
    pub file: Arc<File>,
    pub begin: LineCol,
    pub end: LineCol,
}

impl fmt::Display for SpanLoc {
    /// Formats the span as `filename:start_line:start_column: end_line:end_column`,
    /// or if the span is zero-length, `filename:line:column`, with a 1-indexed line and column.
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        if self.begin == self.end {
            write!(f, "{}:{}:{}", self.file.name, self.begin.line+1, self.begin.column+1)
        } else {
            write!(f, "{}:{}:{}: {}:{}", self.file.name, self.begin.line+1, self.begin.column+1, self.end.line+1, self.end.column+1)
        }
    }
}

#[test]
fn test_codemap() {
    let mut codemap = CodeMap::new();
    let f1 = codemap.add_file("test1.rs".to_string(), "abcd\nefghij\nqwerty".to_string());
    let f2 = codemap.add_file("test2.rs".to_string(), "foo\nbar".to_string());

    assert_eq!(codemap.find_file(f1.span.low()).name(), "test1.rs");
    assert_eq!(codemap.find_file(f1.span.high()).name(), "test1.rs");
    assert_eq!(codemap.find_file(f2.span.low()).name(), "test2.rs");
    assert_eq!(codemap.find_file(f2.span.high()).name(), "test2.rs");

    let x = f1.span.subspan(5, 10);
    let f = codemap.find_file(x.low);
    assert_eq!(f.name, "test1.rs");
    assert_eq!(f.find_line_col(f.span.low()), LineCol { line: 0, column: 0 });
    assert_eq!(f.find_line_col(f.span.low() + 4), LineCol { line: 0, column: 4 });
    assert_eq!(f.find_line_col(f.span.low() + 5), LineCol { line: 1, column: 0 });
    assert_eq!(f.find_line_col(f.span.low() + 16), LineCol { line: 2, column: 4 });

    let x = f2.span.subspan(4, 7);
    assert_eq!(codemap.find_file(x.low()).name(), "test2.rs");
    assert_eq!(codemap.find_file(x.high()).name(), "test2.rs");
}

#[test]
fn test_issue2() {
    let mut codemap = CodeMap::new();
    let content = "a \nxyz\r\n";
    let file = codemap.add_file("<test>".to_owned(), content.to_owned());

    let span = file.span.subspan(2, 3);
    assert_eq!(codemap.look_up_span(span), SpanLoc {
        file: file.clone(),
        begin: LineCol { line: 0, column: 2 },
        end: LineCol { line: 1, column: 0 }
    });

    assert_eq!(file.source_line(0), "a ");
    assert_eq!(file.source_line(1), "xyz");
    assert_eq!(file.source_line(2), "");
}

#[test]
fn test_multibyte() {
    let mut codemap = CodeMap::new();
    let content = "65°00′N 18°00′W 汉语\n🔬";
    let file = codemap.add_file("<test>".to_owned(), content.to_owned());

    assert_eq!(codemap.look_up_pos(file.span.low() + 21), Loc { file: file.clone(), position: LineCol { line: 0, column: 15 } });
    assert_eq!(codemap.look_up_pos(file.span.low() + 28), Loc { file: file.clone(), position: LineCol { line: 0, column: 18 } });
    assert_eq!(codemap.look_up_pos(file.span.low() + 33), Loc { file: file.clone(), position: LineCol { line: 1, column: 1 } });
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2015 Guillaume Gomez
//

#![crate_name = "sysinfo"]
#![crate_type = "lib"]
#![crate_type = "rlib"]
#![doc = include_str!("../README.md")]
#![allow(unknown_lints)]
#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
#![allow(clippy::upper_case_acronyms)]
#![allow(renamed_and_removed_lints)]

#[cfg(feature = "debug")]
#[doc(hidden)]
#[allow(unused)]
macro_rules! sysinfo_debug {
    ($($x:tt)*) => {{
        eprintln!($($x)*);
    }}
}

#[cfg(not(feature = "debug"))]
#[doc(hidden)]
#[allow(unused)]
macro_rules! sysinfo_debug {
    ($($x:tt)*) => {{}};
}

cfg_if::cfg_if! {
    if #[cfg(any(target_os = "macos", target_os = "ios"))] {
        mod apple;
        use apple as sys;
        extern crate core_foundation_sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else if #[cfg(windows)] {
        mod windows;
        use windows as sys;
        extern crate winapi;
        extern crate ntapi;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else if #[cfg(any(target_os = "linux", target_os = "android"))] {
        mod linux;
        use linux as sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else {
        mod unknown;
        use unknown as sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 0;
    }
}

pub use common::{
    AsU32, DiskType, DiskUsage, Gid, LoadAvg, NetworksIter, Pid, RefreshKind, Signal, Uid, User,
};
pub use sys::{Component, Disk, NetworkData, Networks, Process, ProcessStatus, Processor, System};
pub use traits::{
    ComponentExt, DiskExt, NetworkExt, NetworksExt, ProcessExt, ProcessorExt, SystemExt, UserExt,
};

#[cfg(feature = "c-interface")]
pub use c_interface::*;
pub use utils::get_current_pid;

#[cfg(feature = "c-interface")]
mod c_interface;
mod common;
mod debug;
mod system;
mod traits;
mod utils;

/// This function is only used on linux targets, on the other platforms it does nothing and returns
/// `false`.
///
/// On linux, to improve performance, we keep a `/proc` file open for each process we index with
/// a maximum number of files open equivalent to half of the system limit.
///
/// The problem is that some users might need all the available file descriptors so we need to
/// allow them to change this limit.
///
/// Note that if you set a limit bigger than the system limit, the system limit will be set.
///
/// Returns `true` if the new value has been set.
///
/// ```no_run
/// use sysinfo::{System, SystemExt, set_open_files_limit};
///
/// // We call the function before any call to the processes update.
/// if !set_open_files_limit(10) {
///     // It'll always return false on non-linux targets.
///     eprintln!("failed to update the open files limit...");
/// }
/// let s = System::new_all();
/// ```
pub fn set_open_files_limit(mut _new_limit: isize) -> bool {
    #[cfg(any(target_os = "linux", target_os = "android"))]
    {
        if _new_limit < 0 {
            _new_limit = 0;
        }
        let max = sys::system::get_max_nb_fds();
        if _new_limit > max {
            _new_limit = max;
        }
        if let Ok(ref mut x) = unsafe { sys::system::REMAINING_FILES.lock() } {
            // If files are already open, to be sure that the number won't be bigger when those
            // files are closed, we subtract the current number of opened files to the new limit.
            let diff = max - **x;
            **x = _new_limit - diff;
            true
        } else {
            false
        }
    }
    #[cfg(all(not(target_os = "linux"), not(target_os = "android")))]
    {
        false
    }
}

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

    #[test]
    fn check_process_memory_usage() {
        let mut s = System::new();
        s.refresh_all();

        if System::IS_SUPPORTED {
            // No process should have 0 as memory usage.
            #[cfg(not(feature = "apple-sandbox"))]
            assert!(!s.processes().iter().all(|(_, proc_)| proc_.memory() == 0));
        } else {
            // There should be no process, but if there is one, its memory usage should be 0.
            assert!(s.processes().iter().all(|(_, proc_)| proc_.memory() == 0));
        }
    }

    #[test]
    fn check_memory_usage() {
        let mut s = System::new();

        assert_eq!(s.total_memory(), 0);
        assert_eq!(s.free_memory(), 0);
        assert_eq!(s.available_memory(), 0);
        assert_eq!(s.used_memory(), 0);
        assert_eq!(s.total_swap(), 0);
        assert_eq!(s.free_swap(), 0);
        assert_eq!(s.used_swap(), 0);

        s.refresh_memory();
        if System::IS_SUPPORTED {
            assert!(s.total_memory() > 0);
            assert!(s.used_memory() > 0);
            if s.total_swap() > 0 {
                // I think it's pretty safe to assume that there is still some swap left...
                assert!(s.free_swap() > 0);
            }
        } else {
            assert_eq!(s.total_memory(), 0);
            assert_eq!(s.used_memory(), 0);
            assert_eq!(s.total_swap(), 0);
            assert_eq!(s.free_swap(), 0);
        }
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn check_cpu_usage() {
        let mut s = System::new();

        s.refresh_all();
        // All CPU usage will start at zero until the second refresh
        assert_eq!(
            s.processes()
                .iter()
                .all(|(_, proc_)| proc_.cpu_usage() == 0.0),
            true
        );

        // Wait a bit to update CPU usage values
        std::thread::sleep(std::time::Duration::from_millis(100));
        s.refresh_all();
        assert_eq!(
            s.processes()
                .iter()
                .all(|(_, proc_)| proc_.cpu_usage() >= 0.0
                    && proc_.cpu_usage() <= (s.processors().len() as f32) * 100.0),
            true
        );
    }

    #[test]
    fn check_users() {
        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_users_list();
        assert!(s.users().len() >= MIN_USERS);

        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_all();
        assert!(s.users().is_empty());

        let s = System::new_all();
        assert!(s.users().len() >= MIN_USERS);
    }

    #[test]
    fn check_uid_gid() {
        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_users_list();
        let users = s.users();
        assert!(users.len() >= MIN_USERS);

        if System::IS_SUPPORTED {
            #[cfg(not(target_os = "windows"))]
            {
                let user = users
                    .iter()
                    .find(|u| u.name() == "root")
                    .expect("no root user");
                assert_eq!(*user.uid(), 0);
                assert_eq!(*user.gid(), 0);
                if let Some(user) = users.iter().find(|u| *u.gid() > 0) {
                    assert!(*user.uid() > 0);
                    assert!(*user.gid() > 0);
                }
            }
            #[cfg(not(target_os = "windows"))]
            {
                assert!(users.iter().filter(|u| *u.uid() > 0).count() > 0);
            }
        }
    }

    #[test]
    fn check_system_info() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let s = System::new();
            assert!(!s.name().expect("Failed to get system name").is_empty());

            assert!(!s
                .kernel_version()
                .expect("Failed to get kernel version")
                .is_empty());

            assert!(!s.os_version().expect("Failed to get os version").is_empty());

            assert!(!s
                .long_os_version()
                .expect("Failed to get long OS version")
                .is_empty());
        }
    }

    #[test]
    fn check_host_name() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let s = System::new();
            assert!(!s.host_name().expect("Failed to get host name").is_empty());
        }
    }

    #[test]
    fn check_refresh_process_return_value() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let _pid = get_current_pid().expect("Failed to get current PID");

            #[cfg(not(feature = "apple-sandbox"))]
            {
                let mut s = System::new();
                // First check what happens in case the process isn't already in our process list.
                assert!(s.refresh_process(_pid));
                // Then check that it still returns true if the process is already in our process list.
                assert!(s.refresh_process(_pid));
            }
        }
    }

    #[test]
    fn ensure_is_supported_is_set_correctly() {
        if MIN_USERS > 0 {
            assert!(System::IS_SUPPORTED);
        } else {
            assert!(!System::IS_SUPPORTED);
        }
    }

    #[test]
    fn check_processors_number() {
        let s = System::new();

        if System::IS_SUPPORTED {
            assert!(!s.processors().is_empty());
            // In case we are running inside a VM, it's possible to not have a physical core, only
            // logical ones, which is why we don't test `physical_cores_count > 0`.
            let physical_cores_count = s
                .physical_core_count()
                .expect("failed to get number of physical cores");
            assert!(physical_cores_count <= s.processors().len());
        } else {
            assert!(s.processors().is_empty());
            assert_eq!(s.physical_core_count(), None);
        }
    }
}

// Used to check that System is Send and Sync.
#[cfg(doctest)]
/// ```
/// fn is_send<T: Send>() {}
/// is_send::<sysinfo::System>();
/// ```
///
/// ```
/// fn is_sync<T: Sync>() {}
/// is_sync::<sysinfo::System>();
/// ```
pub mod check_if_system_is_send_and_sync {}
<commit_msg>Cargo fix + clippy suggestions (#559)<commit_after>//
// Sysinfo
//
// Copyright (c) 2015 Guillaume Gomez
//

#![crate_name = "sysinfo"]
#![crate_type = "lib"]
#![crate_type = "rlib"]
#![doc = include_str!("../README.md")]
#![allow(unknown_lints)]
#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
#![allow(clippy::upper_case_acronyms)]
#![allow(renamed_and_removed_lints)]

#[cfg(feature = "debug")]
#[doc(hidden)]
#[allow(unused)]
macro_rules! sysinfo_debug {
    ($($x:tt)*) => {{
        eprintln!($($x)*);
    }}
}

#[cfg(not(feature = "debug"))]
#[doc(hidden)]
#[allow(unused)]
macro_rules! sysinfo_debug {
    ($($x:tt)*) => {{}};
}

cfg_if::cfg_if! {
    if #[cfg(any(target_os = "macos", target_os = "ios"))] {
        mod apple;
        use apple as sys;
        extern crate core_foundation_sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else if #[cfg(windows)] {
        mod windows;
        use windows as sys;
        extern crate winapi;
        extern crate ntapi;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else if #[cfg(any(target_os = "linux", target_os = "android"))] {
        mod linux;
        use linux as sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 1;
    } else {
        mod unknown;
        use unknown as sys;

        #[cfg(test)]
        pub(crate) const MIN_USERS: usize = 0;
    }
}

pub use common::{
    AsU32, DiskType, DiskUsage, Gid, LoadAvg, NetworksIter, Pid, RefreshKind, Signal, Uid, User,
};
pub use sys::{Component, Disk, NetworkData, Networks, Process, ProcessStatus, Processor, System};
pub use traits::{
    ComponentExt, DiskExt, NetworkExt, NetworksExt, ProcessExt, ProcessorExt, SystemExt, UserExt,
};

#[cfg(feature = "c-interface")]
pub use c_interface::*;
pub use utils::get_current_pid;

#[cfg(feature = "c-interface")]
mod c_interface;
mod common;
mod debug;
mod system;
mod traits;
mod utils;

/// This function is only used on linux targets, on the other platforms it does nothing and returns
/// `false`.
///
/// On linux, to improve performance, we keep a `/proc` file open for each process we index with
/// a maximum number of files open equivalent to half of the system limit.
///
/// The problem is that some users might need all the available file descriptors so we need to
/// allow them to change this limit.
///
/// Note that if you set a limit bigger than the system limit, the system limit will be set.
///
/// Returns `true` if the new value has been set.
///
/// ```no_run
/// use sysinfo::{System, SystemExt, set_open_files_limit};
///
/// // We call the function before any call to the processes update.
/// if !set_open_files_limit(10) {
///     // It'll always return false on non-linux targets.
///     eprintln!("failed to update the open files limit...");
/// }
/// let s = System::new_all();
/// ```
pub fn set_open_files_limit(mut _new_limit: isize) -> bool {
    #[cfg(any(target_os = "linux", target_os = "android"))]
    {
        if _new_limit < 0 {
            _new_limit = 0;
        }
        let max = sys::system::get_max_nb_fds();
        if _new_limit > max {
            _new_limit = max;
        }
        if let Ok(ref mut x) = unsafe { sys::system::REMAINING_FILES.lock() } {
            // If files are already open, to be sure that the number won't be bigger when those
            // files are closed, we subtract the current number of opened files to the new limit.
            let diff = max - **x;
            **x = _new_limit - diff;
            true
        } else {
            false
        }
    }
    #[cfg(all(not(target_os = "linux"), not(target_os = "android")))]
    {
        false
    }
}

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

    #[test]
    fn check_process_memory_usage() {
        let mut s = System::new();
        s.refresh_all();

        if System::IS_SUPPORTED {
            // No process should have 0 as memory usage.
            #[cfg(not(feature = "apple-sandbox"))]
            assert!(!s.processes().iter().all(|(_, proc_)| proc_.memory() == 0));
        } else {
            // There should be no process, but if there is one, its memory usage should be 0.
            assert!(s.processes().iter().all(|(_, proc_)| proc_.memory() == 0));
        }
    }

    #[test]
    fn check_memory_usage() {
        let mut s = System::new();

        assert_eq!(s.total_memory(), 0);
        assert_eq!(s.free_memory(), 0);
        assert_eq!(s.available_memory(), 0);
        assert_eq!(s.used_memory(), 0);
        assert_eq!(s.total_swap(), 0);
        assert_eq!(s.free_swap(), 0);
        assert_eq!(s.used_swap(), 0);

        s.refresh_memory();
        if System::IS_SUPPORTED {
            assert!(s.total_memory() > 0);
            assert!(s.used_memory() > 0);
            if s.total_swap() > 0 {
                // I think it's pretty safe to assume that there is still some swap left...
                assert!(s.free_swap() > 0);
            }
        } else {
            assert_eq!(s.total_memory(), 0);
            assert_eq!(s.used_memory(), 0);
            assert_eq!(s.total_swap(), 0);
            assert_eq!(s.free_swap(), 0);
        }
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn check_cpu_usage() {
        let mut s = System::new();

        s.refresh_all();
        // All CPU usage will start at zero until the second refresh
        assert!(s
            .processes()
            .iter()
            .all(|(_, proc_)| proc_.cpu_usage() == 0.0));

        // Wait a bit to update CPU usage values
        std::thread::sleep(std::time::Duration::from_millis(100));
        s.refresh_all();
        assert!(s
            .processes()
            .iter()
            .all(|(_, proc_)| proc_.cpu_usage() >= 0.0
                && proc_.cpu_usage() <= (s.processors().len() as f32) * 100.0));
    }

    #[test]
    fn check_users() {
        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_users_list();
        assert!(s.users().len() >= MIN_USERS);

        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_all();
        assert!(s.users().is_empty());

        let s = System::new_all();
        assert!(s.users().len() >= MIN_USERS);
    }

    #[test]
    fn check_uid_gid() {
        let mut s = System::new();
        assert!(s.users().is_empty());
        s.refresh_users_list();
        let users = s.users();
        assert!(users.len() >= MIN_USERS);

        if System::IS_SUPPORTED {
            #[cfg(not(target_os = "windows"))]
            {
                let user = users
                    .iter()
                    .find(|u| u.name() == "root")
                    .expect("no root user");
                assert_eq!(*user.uid(), 0);
                assert_eq!(*user.gid(), 0);
                if let Some(user) = users.iter().find(|u| *u.gid() > 0) {
                    assert!(*user.uid() > 0);
                    assert!(*user.gid() > 0);
                }
            }
            #[cfg(not(target_os = "windows"))]
            {
                assert!(users.iter().filter(|u| *u.uid() > 0).count() > 0);
            }
        }
    }

    #[test]
    fn check_system_info() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let s = System::new();
            assert!(!s.name().expect("Failed to get system name").is_empty());

            assert!(!s
                .kernel_version()
                .expect("Failed to get kernel version")
                .is_empty());

            assert!(!s.os_version().expect("Failed to get os version").is_empty());

            assert!(!s
                .long_os_version()
                .expect("Failed to get long OS version")
                .is_empty());
        }
    }

    #[test]
    fn check_host_name() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let s = System::new();
            assert!(!s.host_name().expect("Failed to get host name").is_empty());
        }
    }

    #[test]
    fn check_refresh_process_return_value() {
        // We don't want to test on unsupported systems.
        if System::IS_SUPPORTED {
            let _pid = get_current_pid().expect("Failed to get current PID");

            #[cfg(not(feature = "apple-sandbox"))]
            {
                let mut s = System::new();
                // First check what happens in case the process isn't already in our process list.
                assert!(s.refresh_process(_pid));
                // Then check that it still returns true if the process is already in our process list.
                assert!(s.refresh_process(_pid));
            }
        }
    }

    #[test]
    fn ensure_is_supported_is_set_correctly() {
        if MIN_USERS > 0 {
            assert!(System::IS_SUPPORTED);
        } else {
            assert!(!System::IS_SUPPORTED);
        }
    }

    #[test]
    fn check_processors_number() {
        let s = System::new();

        if System::IS_SUPPORTED {
            assert!(!s.processors().is_empty());
            // In case we are running inside a VM, it's possible to not have a physical core, only
            // logical ones, which is why we don't test `physical_cores_count > 0`.
            let physical_cores_count = s
                .physical_core_count()
                .expect("failed to get number of physical cores");
            assert!(physical_cores_count <= s.processors().len());
        } else {
            assert!(s.processors().is_empty());
            assert_eq!(s.physical_core_count(), None);
        }
    }
}

// Used to check that System is Send and Sync.
#[cfg(doctest)]
/// ```
/// fn is_send<T: Send>() {}
/// is_send::<sysinfo::System>();
/// ```
///
/// ```
/// fn is_sync<T: Sync>() {}
/// is_sync::<sysinfo::System>();
/// ```
pub mod check_if_system_is_send_and_sync {}
<|endoftext|>


<commit_before>extern crate bit_vec;
extern crate byteorder;
extern crate crc;
extern crate libc;
extern crate libz_sys;
extern crate num_cpus;
extern crate scoped_pool;

use scoped_pool::Pool;
use std::collections::{HashMap, HashSet};
use std::fs::{File, copy};
use std::io::{BufWriter, Write, stderr, stdout};
use std::path::{Path, PathBuf};

pub mod deflate {
    pub mod deflate;
    pub mod stream;
}
pub mod png;

#[derive(Clone,Debug)]
/// Options controlling the output of the `optimize` function
pub struct Options {
    /// Whether the input file should be backed up before writing the output
    pub backup: bool,
    /// Path to write the output file to
    pub out_file: PathBuf,
    /// Used only in CLI interface
    pub out_dir: Option<PathBuf>,
    /// Write to stdout instead of a file
    pub stdout: bool,
    /// Attempt to fix errors when decoding the input file
    pub fix_errors: bool,
    /// Don't actually write any output, just calculate the best results
    pub pretend: bool,
    /// Used only in CLI interface
    pub recursive: bool,
    /// Overwrite existing output files
    pub clobber: bool,
    /// Create new output files if they don't exist
    pub create: bool,
    /// Write to output even if there was no improvement in compression
    pub force: bool,
    /// Ensure the output file has the same permissions as the input file
    pub preserve_attrs: bool,
    /// How verbose the console logging should be (`None` for quiet, `Some(0)` for normal, `Some(1)` for verbose)
    pub verbosity: Option<u8>,
    /// Which filters to try on the file (0-5)
    pub filter: HashSet<u8>,
    /// Whether to change the interlacing type of the file
    /// `None` will not change the current interlacing type
    /// `Some(x)` will change the file to interlacing mode `x`
    pub interlace: Option<u8>,
    /// Which zlib compression levels to try on the file (1-9)
    pub compression: HashSet<u8>,
    /// Which zlib memory levels to try on the file (1-9)
    pub memory: HashSet<u8>,
    /// Which zlib compression strategies to try on the file (0-3)
    pub strategies: HashSet<u8>,
    /// Window size to use when compressing the file, as `2^window` bytes
    /// Doesn't affect compression but may affect speed and memory usage
    /// 15 is recommended default, 8-15 are valid values
    pub window: u8,
    /// Whether to attempt bit depth reduction
    pub bit_depth_reduction: bool,
    /// Whether to attempt color type reduction
    pub color_type_reduction: bool,
    /// Whether to attempt palette reduction
    pub palette_reduction: bool,
    /// Whether to perform IDAT recoding
    /// If any type of reduction is performed, IDAT recoding will be performed
    /// regardless of this setting
    pub idat_recoding: bool,
    /// Which headers to strip from the PNG file, if any
    pub strip: png::Headers,
    /// Whether to use heuristics to pick the best filter and compression
    /// Intended for use with `-o 1` from the CLI interface
    pub use_heuristics: bool,
}

/// Perform optimization on the input file using the options provided
pub fn optimize(filepath: &Path, opts: &Options) -> Result<(), String> {
    // Decode PNG from file
    if opts.verbosity.is_some() {
        writeln!(&mut stderr(), "Processing: {}", filepath.to_str().unwrap()).ok();
    }
    let in_file = Path::new(filepath);
    let mut png = match png::PngData::new(&in_file) {
        Ok(x) => x,
        Err(x) => return Err(x),
    };

    // Print png info
    let idat_original_size = png.idat_data.len();
    let file_original_size = filepath.metadata().unwrap().len() as usize;
    if opts.verbosity.is_some() {
        writeln!(&mut stderr(),
                 "    {}x{} pixels, PNG format",
                 png.ihdr_data.width,
                 png.ihdr_data.height)
            .ok();
        if let Some(palette) = png.palette.clone() {
            writeln!(&mut stderr(),
                     "    {} bits/pixel, {} colors in palette",
                     png.ihdr_data.bit_depth,
                     palette.len() / 3)
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    {}x{} bits/pixel, {:?}",
                     png.channels_per_pixel(),
                     png.ihdr_data.bit_depth,
                     png.ihdr_data.color_type)
                .ok();
        }
        writeln!(&mut stderr(),
                 "    IDAT size = {} bytes",
                 idat_original_size)
            .ok();
        writeln!(&mut stderr(),
                 "    File size = {} bytes",
                 file_original_size)
            .ok();
    }

    let mut filter = opts.filter.clone();
    let compression = opts.compression.clone();
    let memory = opts.memory.clone();
    let mut strategies = opts.strategies.clone();

    if opts.use_heuristics {
        // Heuristically determine which set of options to use
        if png.ihdr_data.bit_depth.as_u8() >= 8 &&
           png.ihdr_data.color_type != png::ColorType::Indexed {
            if filter.is_empty() {
                filter.insert(5);
            }
            if strategies.is_empty() {
                strategies.insert(1);
            }
        } else {
            if filter.is_empty() {
                filter.insert(0);
            }
            if strategies.is_empty() {
                strategies.insert(0);
            }
        }
    }

    let mut something_changed = false;

    if opts.bit_depth_reduction {
        if png.reduce_bit_depth() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if opts.color_type_reduction {
        if png.reduce_color_type() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if opts.palette_reduction {
        if png.reduce_palette() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if something_changed && opts.verbosity.is_some() {
        report_reduction(&png);
    }

    if let Some(interlacing) = opts.interlace {
        if png.change_interlacing(interlacing) {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        }
    }

    if opts.idat_recoding || something_changed {
        // Use 1 thread on single-core, otherwise use threads = 1.5x CPU cores
        let num_cpus = num_cpus::get();
        let thread_count = num_cpus + (num_cpus >> 1);
        let pool = Pool::new(thread_count);
        // Go through selected permutations and determine the best
        let mut best: Option<(u8, u8, u8, u8, Vec<u8>)> = None;
        let combinations = filter.len() * compression.len() * memory.len() * strategies.len();
        let mut results: Vec<(u8, u8, u8, u8, Vec<u8>)> = Vec::with_capacity(combinations);
        let mut filters: HashMap<u8, Vec<u8>> = HashMap::with_capacity(filter.len());
        if opts.verbosity.is_some() {
            writeln!(&mut stderr(), "Trying: {} combinations", combinations).ok();
        }

        for f in &filter {
            let filtered = png.filter_image(*f);
            filters.insert(*f, filtered.clone());
            for zc in &compression {
                for zm in &memory {
                    for zs in &strategies {
                        results.push((*f, *zc, *zm, *zs, Vec::new()));
                    }
                }
            }
        }
        pool.scoped(|scope| {
            for trial in &mut results {
                let filtered = filters.get(&trial.0).unwrap();
                scope.execute(move || {
                    let new_idat = deflate::deflate::deflate(filtered,
                                                             trial.1,
                                                             trial.2,
                                                             trial.3,
                                                             opts.window)
                                       .unwrap();

                    if opts.verbosity == Some(1) {
                        writeln!(&mut stderr(),
                                 "    zc = {}  zm = {}  zs = {}  f = {}        {} bytes",
                                 trial.1,
                                 trial.2,
                                 trial.3,
                                 trial.0,
                                 new_idat.len())
                            .ok();
                    }

                    trial.4 = new_idat;
                });
            }
        });

        for result in &results {
            if (best.is_some() && result.4.len() < best.as_ref().map(|x| x.4.len()).unwrap()) ||
               (best.is_none() &&
                (result.4.len() < png.idat_data.len() ||
                 (opts.interlace.is_some() && opts.interlace != Some(png.ihdr_data.interlaced)) ||
                 opts.force)) {
                best = Some(result.clone());
            }
        }

        if let Some(better) = best {
            png.idat_data = better.4.clone();
            if opts.verbosity.is_some() {
                writeln!(&mut stderr(), "Found better combination:").ok();
                writeln!(&mut stderr(),
                         "    zc = {}  zm = {}  zs = {}  f = {}        {} bytes",
                         better.1,
                         better.2,
                         better.3,
                         better.0,
                         png.idat_data.len())
                    .ok();
            }
        }
    }

    match opts.strip.clone() {
        // Strip headers
        png::Headers::None => (),
        png::Headers::Some(hdrs) => {
            for hdr in &hdrs {
                png.aux_headers.remove(hdr);
            }
        }
        png::Headers::Safe => {
            const PRESERVED_HEADERS: [&'static str; 9] = ["cHRM", "gAMA", "iCCP", "sBIT", "sRGB",
                                                          "bKGD", "hIST", "pHYs", "sPLT"];
            let mut preserved = HashMap::new();
            for (hdr, contents) in &png.aux_headers {
                if PRESERVED_HEADERS.contains(&hdr.as_ref()) {
                    preserved.insert(hdr.clone(), contents.clone());
                }
            }
            png.aux_headers = preserved;
        }
        png::Headers::All => {
            png.aux_headers = HashMap::new();
        }
    }

    let output_data = png.output();
    if file_original_size <= output_data.len() && !opts.force && opts.interlace.is_none() {
        writeln!(&mut stderr(), "File already optimized").ok();
        return Ok(());
    }

    if opts.pretend {
        writeln!(&mut stderr(), "Running in pretend mode, no output").ok();
    } else {
        if opts.backup {
            match copy(in_file,
                       in_file.with_extension(format!("bak.{}",
                                                      in_file.extension()
                                                             .unwrap()
                                                             .to_str()
                                                             .unwrap()))) {
                Ok(x) => x,
                Err(_) => {
                    return Err(format!("Unable to write to backup file at {}",
                                       opts.out_file.display()))
                }
            };
        }

        if opts.stdout {
            let mut buffer = BufWriter::new(stdout());
            match buffer.write_all(&output_data) {
                Ok(_) => (),
                Err(_) => return Err("Unable to write to stdout".to_owned()),
            }
        } else {
            let out_file = match File::create(&opts.out_file) {
                Ok(x) => x,
                Err(_) => {
                    return Err(format!("Unable to write to file {}", opts.out_file.display()))
                }
            };
            let mut buffer = BufWriter::new(out_file);
            match buffer.write_all(&output_data) {
                Ok(_) => {
                    if opts.verbosity.is_some() {
                        writeln!(&mut stderr(), "Output: {}", opts.out_file.display()).ok();
                    }
                }
                Err(_) => {
                    return Err(format!("Unable to write to file {}", opts.out_file.display()))
                }
            }
        }
    }
    if opts.verbosity.is_some() {
        if idat_original_size >= png.idat_data.len() {
            writeln!(&mut stderr(),
                     "    IDAT size = {} bytes ({} bytes decrease)",
                     png.idat_data.len(),
                     idat_original_size - png.idat_data.len())
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    IDAT size = {} bytes ({} bytes increase)",
                     png.idat_data.len(),
                     png.idat_data.len() - idat_original_size)
                .ok();
        }
        if file_original_size >= output_data.len() {
            writeln!(&mut stderr(),
                     "    file size = {} bytes ({} bytes = {:.2}% decrease)",
                     output_data.len(),
                     file_original_size - output_data.len(),
                     (file_original_size - output_data.len()) as f64 / file_original_size as f64 *
                     100f64)
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    file size = {} bytes ({} bytes = {:.2}% increase)",
                     output_data.len(),
                     output_data.len() - file_original_size,
                     (output_data.len() - file_original_size) as f64 / file_original_size as f64 *
                     100f64)
                .ok();
        }
    }
    Ok(())
}

fn report_reduction(png: &png::PngData) {
    if let Some(palette) = png.palette.clone() {
        writeln!(&mut stderr(),
                 "Reducing image to {} bits/pixel, {} colors in palette",
                 png.ihdr_data.bit_depth,
                 palette.len() / 3)
            .ok();
    } else {
        writeln!(&mut stderr(),
                 "Reducing image to {}x{} bits/pixel, {}",
                 png.channels_per_pixel(),
                 png.ihdr_data.bit_depth,
                 png.ihdr_data.color_type)
            .ok();
    }
}
<commit_msg>Don't store intermediate copies of compressed image data unless it is current best.<commit_after>extern crate bit_vec;
extern crate byteorder;
extern crate crc;
extern crate libc;
extern crate libz_sys;
extern crate num_cpus;
extern crate scoped_pool;

use scoped_pool::Pool;
use std::collections::{HashMap, HashSet};
use std::fs::{File, copy};
use std::io::{BufWriter, Write, stderr, stdout};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

pub mod deflate {
    pub mod deflate;
    pub mod stream;
}
pub mod png;

#[derive(Clone,Debug)]
/// Options controlling the output of the `optimize` function
pub struct Options {
    /// Whether the input file should be backed up before writing the output
    pub backup: bool,
    /// Path to write the output file to
    pub out_file: PathBuf,
    /// Used only in CLI interface
    pub out_dir: Option<PathBuf>,
    /// Write to stdout instead of a file
    pub stdout: bool,
    /// Attempt to fix errors when decoding the input file
    pub fix_errors: bool,
    /// Don't actually write any output, just calculate the best results
    pub pretend: bool,
    /// Used only in CLI interface
    pub recursive: bool,
    /// Overwrite existing output files
    pub clobber: bool,
    /// Create new output files if they don't exist
    pub create: bool,
    /// Write to output even if there was no improvement in compression
    pub force: bool,
    /// Ensure the output file has the same permissions as the input file
    pub preserve_attrs: bool,
    /// How verbose the console logging should be (`None` for quiet, `Some(0)` for normal, `Some(1)` for verbose)
    pub verbosity: Option<u8>,
    /// Which filters to try on the file (0-5)
    pub filter: HashSet<u8>,
    /// Whether to change the interlacing type of the file
    /// `None` will not change the current interlacing type
    /// `Some(x)` will change the file to interlacing mode `x`
    pub interlace: Option<u8>,
    /// Which zlib compression levels to try on the file (1-9)
    pub compression: HashSet<u8>,
    /// Which zlib memory levels to try on the file (1-9)
    pub memory: HashSet<u8>,
    /// Which zlib compression strategies to try on the file (0-3)
    pub strategies: HashSet<u8>,
    /// Window size to use when compressing the file, as `2^window` bytes
    /// Doesn't affect compression but may affect speed and memory usage
    /// 15 is recommended default, 8-15 are valid values
    pub window: u8,
    /// Whether to attempt bit depth reduction
    pub bit_depth_reduction: bool,
    /// Whether to attempt color type reduction
    pub color_type_reduction: bool,
    /// Whether to attempt palette reduction
    pub palette_reduction: bool,
    /// Whether to perform IDAT recoding
    /// If any type of reduction is performed, IDAT recoding will be performed
    /// regardless of this setting
    pub idat_recoding: bool,
    /// Which headers to strip from the PNG file, if any
    pub strip: png::Headers,
    /// Whether to use heuristics to pick the best filter and compression
    /// Intended for use with `-o 1` from the CLI interface
    pub use_heuristics: bool,
}

/// Perform optimization on the input file using the options provided
pub fn optimize(filepath: &Path, opts: &Options) -> Result<(), String> {
    type TrialWithData = (u8, u8, u8, u8, Vec<u8>);

    // Decode PNG from file
    if opts.verbosity.is_some() {
        writeln!(&mut stderr(), "Processing: {}", filepath.to_str().unwrap()).ok();
    }
    let in_file = Path::new(filepath);
    let mut png = match png::PngData::new(&in_file) {
        Ok(x) => x,
        Err(x) => return Err(x),
    };

    // Print png info
    let idat_original_size = png.idat_data.len();
    let file_original_size = filepath.metadata().unwrap().len() as usize;
    if opts.verbosity.is_some() {
        writeln!(&mut stderr(),
                 "    {}x{} pixels, PNG format",
                 png.ihdr_data.width,
                 png.ihdr_data.height)
            .ok();
        if let Some(palette) = png.palette.clone() {
            writeln!(&mut stderr(),
                     "    {} bits/pixel, {} colors in palette",
                     png.ihdr_data.bit_depth,
                     palette.len() / 3)
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    {}x{} bits/pixel, {:?}",
                     png.channels_per_pixel(),
                     png.ihdr_data.bit_depth,
                     png.ihdr_data.color_type)
                .ok();
        }
        writeln!(&mut stderr(),
                 "    IDAT size = {} bytes",
                 idat_original_size)
            .ok();
        writeln!(&mut stderr(),
                 "    File size = {} bytes",
                 file_original_size)
            .ok();
    }

    let mut filter = opts.filter.clone();
    let compression = opts.compression.clone();
    let memory = opts.memory.clone();
    let mut strategies = opts.strategies.clone();

    if opts.use_heuristics {
        // Heuristically determine which set of options to use
        if png.ihdr_data.bit_depth.as_u8() >= 8 &&
           png.ihdr_data.color_type != png::ColorType::Indexed {
            if filter.is_empty() {
                filter.insert(5);
            }
            if strategies.is_empty() {
                strategies.insert(1);
            }
        } else {
            if filter.is_empty() {
                filter.insert(0);
            }
            if strategies.is_empty() {
                strategies.insert(0);
            }
        }
    }

    let mut something_changed = false;

    if opts.bit_depth_reduction {
        if png.reduce_bit_depth() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if opts.color_type_reduction {
        if png.reduce_color_type() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if opts.palette_reduction {
        if png.reduce_palette() {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        };
    }

    if something_changed && opts.verbosity.is_some() {
        report_reduction(&png);
    }

    if let Some(interlacing) = opts.interlace {
        if png.change_interlacing(interlacing) {
            something_changed = true;
            if opts.verbosity == Some(1) {
                report_reduction(&png);
            }
        }
    }

    if opts.idat_recoding || something_changed {
        // Use 1 thread on single-core, otherwise use threads = 1.5x CPU cores
        let num_cpus = num_cpus::get();
        let thread_count = num_cpus + (num_cpus >> 1);
        let pool = Pool::new(thread_count);
        // Go through selected permutations and determine the best
        let best: Arc<Mutex<Option<TrialWithData>>> = Arc::new(Mutex::new(None));
        let combinations = filter.len() * compression.len() * memory.len() * strategies.len();
        let mut results: Vec<(u8, u8, u8, u8)> = Vec::with_capacity(combinations);
        let mut filters: HashMap<u8, Vec<u8>> = HashMap::with_capacity(filter.len());
        if opts.verbosity.is_some() {
            writeln!(&mut stderr(), "Trying: {} combinations", combinations).ok();
        }

        for f in &filter {
            let filtered = png.filter_image(*f);
            filters.insert(*f, filtered.clone());
            for zc in &compression {
                for zm in &memory {
                    for zs in &strategies {
                        results.push((*f, *zc, *zm, *zs));
                    }
                }
            }
        }
        pool.scoped(|scope| {
            let original_len = png.idat_data.len();
            let interlacing_changed = opts.interlace.is_some() &&
                                      opts.interlace != Some(png.ihdr_data.interlaced);
            for trial in &results {
                let filtered = filters.get(&trial.0).unwrap();
                let best = best.clone();
                scope.execute(move || {
                    let new_idat = deflate::deflate::deflate(filtered,
                                                             trial.1,
                                                             trial.2,
                                                             trial.3,
                                                             opts.window)
                                       .unwrap();

                    if opts.verbosity == Some(1) {
                        writeln!(&mut stderr(),
                                 "    zc = {}  zm = {}  zs = {}  f = {}        {} bytes",
                                 trial.1,
                                 trial.2,
                                 trial.3,
                                 trial.0,
                                 new_idat.len())
                            .ok();
                    }

                    let mut best = best.lock().unwrap();
                    if (best.is_some() &&
                        new_idat.len() < best.as_ref().map(|x| x.4.len()).unwrap()) ||
                       (best.is_none() &&
                        (new_idat.len() < original_len || interlacing_changed || opts.force)) {
                        *best = Some((trial.0, trial.1, trial.2, trial.3, new_idat));
                    }
                });
            }
        });

        let mut final_best = best.lock().unwrap();
        if let Some(better) = final_best.take() {
            png.idat_data = better.4.clone();
            if opts.verbosity.is_some() {
                writeln!(&mut stderr(), "Found better combination:").ok();
                writeln!(&mut stderr(),
                         "    zc = {}  zm = {}  zs = {}  f = {}        {} bytes",
                         better.1,
                         better.2,
                         better.3,
                         better.0,
                         png.idat_data.len())
                    .ok();
            }
        }
    }

    match opts.strip.clone() {
        // Strip headers
        png::Headers::None => (),
        png::Headers::Some(hdrs) => {
            for hdr in &hdrs {
                png.aux_headers.remove(hdr);
            }
        }
        png::Headers::Safe => {
            const PRESERVED_HEADERS: [&'static str; 9] = ["cHRM", "gAMA", "iCCP", "sBIT", "sRGB",
                                                          "bKGD", "hIST", "pHYs", "sPLT"];
            let mut preserved = HashMap::new();
            for (hdr, contents) in &png.aux_headers {
                if PRESERVED_HEADERS.contains(&hdr.as_ref()) {
                    preserved.insert(hdr.clone(), contents.clone());
                }
            }
            png.aux_headers = preserved;
        }
        png::Headers::All => {
            png.aux_headers = HashMap::new();
        }
    }

    let output_data = png.output();
    if file_original_size <= output_data.len() && !opts.force && opts.interlace.is_none() {
        writeln!(&mut stderr(), "File already optimized").ok();
        return Ok(());
    }

    if opts.pretend {
        writeln!(&mut stderr(), "Running in pretend mode, no output").ok();
    } else {
        if opts.backup {
            match copy(in_file,
                       in_file.with_extension(format!("bak.{}",
                                                      in_file.extension()
                                                             .unwrap()
                                                             .to_str()
                                                             .unwrap()))) {
                Ok(x) => x,
                Err(_) => {
                    return Err(format!("Unable to write to backup file at {}",
                                       opts.out_file.display()))
                }
            };
        }

        if opts.stdout {
            let mut buffer = BufWriter::new(stdout());
            match buffer.write_all(&output_data) {
                Ok(_) => (),
                Err(_) => return Err("Unable to write to stdout".to_owned()),
            }
        } else {
            let out_file = match File::create(&opts.out_file) {
                Ok(x) => x,
                Err(_) => {
                    return Err(format!("Unable to write to file {}", opts.out_file.display()))
                }
            };
            let mut buffer = BufWriter::new(out_file);
            match buffer.write_all(&output_data) {
                Ok(_) => {
                    if opts.verbosity.is_some() {
                        writeln!(&mut stderr(), "Output: {}", opts.out_file.display()).ok();
                    }
                }
                Err(_) => {
                    return Err(format!("Unable to write to file {}", opts.out_file.display()))
                }
            }
        }
    }
    if opts.verbosity.is_some() {
        if idat_original_size >= png.idat_data.len() {
            writeln!(&mut stderr(),
                     "    IDAT size = {} bytes ({} bytes decrease)",
                     png.idat_data.len(),
                     idat_original_size - png.idat_data.len())
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    IDAT size = {} bytes ({} bytes increase)",
                     png.idat_data.len(),
                     png.idat_data.len() - idat_original_size)
                .ok();
        }
        if file_original_size >= output_data.len() {
            writeln!(&mut stderr(),
                     "    file size = {} bytes ({} bytes = {:.2}% decrease)",
                     output_data.len(),
                     file_original_size - output_data.len(),
                     (file_original_size - output_data.len()) as f64 / file_original_size as f64 *
                     100f64)
                .ok();
        } else {
            writeln!(&mut stderr(),
                     "    file size = {} bytes ({} bytes = {:.2}% increase)",
                     output_data.len(),
                     output_data.len() - file_original_size,
                     (output_data.len() - file_original_size) as f64 / file_original_size as f64 *
                     100f64)
                .ok();
        }
    }
    Ok(())
}

fn report_reduction(png: &png::PngData) {
    if let Some(palette) = png.palette.clone() {
        writeln!(&mut stderr(),
                 "Reducing image to {} bits/pixel, {} colors in palette",
                 png.ihdr_data.bit_depth,
                 palette.len() / 3)
            .ok();
    } else {
        writeln!(&mut stderr(),
                 "Reducing image to {}x{} bits/pixel, {}",
                 png.channels_per_pixel(),
                 png.ihdr_data.bit_depth,
                 png.ihdr_data.color_type)
            .ok();
    }
}
<|endoftext|>


<commit_before>import option.none;
import option.some;
import util.orb;

type vbuf = rustrt.vbuf;

type operator2[T,U,V] = fn(&T, &U) -> V;

native "rust" mod rustrt {
    type vbuf;

    fn vec_buf[T](vec[T] v, uint offset) -> vbuf;

    fn vec_len[T](vec[T] v) -> uint;
    /**
     * Sometimes we modify the vec internal data via vec_buf and need to
     * update the vec's fill length accordingly.
     */
    fn vec_len_set[T](vec[T] v, uint n);

    /**
     * The T in vec_alloc[T, U] is the type of the vec to allocate.  The
     * U is the type of an element in the vec.  So to allocate a vec[U] we
     * want to invoke this as vec_alloc[vec[U], U].
     */
    fn vec_alloc[T, U](uint n_elts) -> vec[U];
    fn vec_alloc_mut[T, U](uint n_elts) -> vec[mutable U];

    fn refcount[T](vec[T] v) -> uint;

    fn vec_print_debug_info[T](vec[T] v);

    fn vec_from_vbuf[T](vbuf v, uint n_elts) -> vec[T];
}

fn alloc[T](uint n_elts) -> vec[T] {
    ret rustrt.vec_alloc[vec[T], T](n_elts);
}

fn alloc_mut[T](uint n_elts) -> vec[mutable T] {
    ret rustrt.vec_alloc_mut[vec[mutable T], T](n_elts);
}

fn refcount[T](vec[mutable? T] v) -> uint {
    auto r = rustrt.refcount[T](v);
    if (r == dbg.const_refcount) {
        ret r;
    } else {
        // -1 because calling this function incremented the refcount.
        ret  r - 1u;
    }
}

unsafe fn vec_from_vbuf[T](vbuf v, uint n_elts) -> vec[T] {
    ret rustrt.vec_from_vbuf[T](v, n_elts);
}

type init_op[T] = fn(uint i) -> T;

fn init_fn[T](&init_op[T] op, uint n_elts) -> vec[T] {
    let vec[T] v = alloc[T](n_elts);
    let uint i = 0u;
    while (i < n_elts) {
        v += vec(op(i));
        i += 1u;
    }
    ret v;
}

fn init_fn_mut[T](&init_op[T] op, uint n_elts) -> vec[mutable T] {
    let vec[mutable T] v = alloc_mut[T](n_elts);
    let uint i = 0u;
    while (i < n_elts) {
        v += vec(mutable op(i));
        i += 1u;
    }
    ret v;
}

fn init_elt[T](&T t, uint n_elts) -> vec[T] {
    /**
     * FIXME (issue #81): should be:
     *
     * fn elt_op[T](&T x, uint i) -> T { ret x; }
     * let init_op[T] inner = bind elt_op[T](t, _);
     * ret init_fn[T](inner, n_elts);
     */
    let vec[T] v = alloc[T](n_elts);
    let uint i = n_elts;
    while (i > 0u) {
        i -= 1u;
        v += vec(t);
    }
    ret v;
}

fn init_elt_mut[T](&T t, uint n_elts) -> vec[mutable T] {
    let vec[mutable T] v = alloc_mut[T](n_elts);
    let uint i = n_elts;
    while (i > 0u) {
        i -= 1u;
        v += vec(mutable t);
    }
    ret v;
}

fn buf[T](vec[mutable? T] v) -> vbuf {
    ret rustrt.vec_buf[T](v, 0u);
}

fn len[T](vec[mutable? T] v) -> uint {
    ret rustrt.vec_len[T](v);
}

fn len_set[T](vec[mutable? T] v, uint n) {
    rustrt.vec_len_set[T](v, n);
}

fn buf_off[T](vec[mutable? T] v, uint offset) -> vbuf {
    check (offset < len[T](v));
    ret rustrt.vec_buf[T](v, offset);
}

fn print_debug_info[T](vec[mutable? T] v) {
    rustrt.vec_print_debug_info[T](v);
}

// Returns the last element of v.
fn last[T](vec[mutable? T] v) -> option.t[T] {
    auto l = len[T](v);
    if (l == 0u) {
        ret none[T];
    }
    ret some[T](v.(l - 1u));
}

// Returns elements from [start..end) from v.
fn slice[T](vec[mutable? T] v, uint start, uint end) -> vec[T] {
    check (start <= end);
    check (end <= len[T](v));
    auto result = alloc[T](end - start);
    let uint i = start;
    while (i < end) {
        result += vec(v.(i));
        i += 1u;
    }
    ret result;
}

fn shift[T](&mutable vec[mutable? T] v) -> T {
    auto ln = len[T](v);
    check(ln > 0u);
    auto e = v.(0);
    v = slice[T](v, 1u, ln);
    ret e;
}

fn pop[T](&mutable vec[mutable? T] v) -> T {
    auto ln = len[T](v);
    check(ln > 0u);
    ln -= 1u;
    auto e = v.(ln);
    v = slice[T](v, 0u, ln);
    ret e;
}

fn push[T](&mutable vec[mutable? T] v, &T t) {
    v += vec(t);
}

fn unshift[T](&mutable vec[mutable? T] v, &T t) {
    auto res = alloc[T](len[T](v) + 1u);
    res += vec(t);
    res += v;
    v = res;
}

fn grow[T](&mutable vec[mutable? T] v, int n, &T initval) {
    let int i = n;
    while (i > 0) {
        i -= 1;
        v += vec(initval);
    }
}

fn map[T, U](&option.operator[T,U] f, &vec[mutable? T] v) -> vec[U] {
    let vec[U] u = alloc[U](len[T](v));
    for (T ve in v) {
        u += vec(f(ve));
    }
    ret u;
}

fn map2[T,U,V](&operator2[T,U,V] f, &vec[mutable? T] v0, &vec[mutable? U] v1)
        -> vec[V] {
    auto v0_len = len[T](v0);
    if (v0_len != len[U](v1)) {
        fail;
    }

    let vec[V] u = alloc[V](v0_len);
    auto i = 0u;
    while (i < v0_len) {
        u += vec(f(v0.(i), v1.(i)));
        i += 1u;
    }

    ret u;
}

fn find[T](fn (&T) -> bool f, &vec[mutable? T] v) -> option.t[T] {
    for (T elt in v) {
        if (f(elt)) {
            ret some[T](elt);
        }
    }

    ret none[T];
}

fn foldl[T, U](fn (&U, &T) -> U p, &U z, &vec[T] v) -> U {
    auto sz = len[T](v);

    if (sz == 0u) {
        ret z;
    }
    else {
        auto rest = slice[T](v, 1u, sz);

        ret (p(foldl[T,U](p, z, rest), v.(0)));
    }
}

fn unzip[T, U](&vec[tup(T, U)] v) -> tup(vec[T], vec[U]) {
    auto sz = len[tup(T, U)](v);

    if (sz == 0u) {
        ret tup(alloc[T](0u), alloc[U](0u));
    }
    else {
        auto rest = slice[tup(T, U)](v, 1u, sz);
        auto tl   = unzip[T, U](rest);
        auto a    = vec(v.(0)._0);
        auto b    = vec(v.(0)._1);
        ret tup(a + tl._0, b + tl._1);
    }
}

fn or(&vec[bool] v) -> bool {
    auto f = orb;
    be _vec.foldl[bool, bool](f, false, v);
}

fn clone[T](&vec[T] v) -> vec[T] {
    ret slice[T](v, 0u, len[T](v));
}

fn plus_option[T](&vec[T] v, &option.t[T] o) -> () {
    alt (o) {
        case (none[T]) {}
        case (some[T](?x)) { v += vec(x); }
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<commit_msg>stdlib: Add some vector functions: empty, empty_mut, grow_set, and freeze. Change _vec.grow to use uint arguments.<commit_after>import option.none;
import option.some;
import util.orb;

type vbuf = rustrt.vbuf;

type operator2[T,U,V] = fn(&T, &U) -> V;

native "rust" mod rustrt {
    type vbuf;

    fn vec_buf[T](vec[T] v, uint offset) -> vbuf;

    fn vec_len[T](vec[T] v) -> uint;
    /**
     * Sometimes we modify the vec internal data via vec_buf and need to
     * update the vec's fill length accordingly.
     */
    fn vec_len_set[T](vec[T] v, uint n);

    /**
     * The T in vec_alloc[T, U] is the type of the vec to allocate.  The
     * U is the type of an element in the vec.  So to allocate a vec[U] we
     * want to invoke this as vec_alloc[vec[U], U].
     */
    fn vec_alloc[T, U](uint n_elts) -> vec[U];
    fn vec_alloc_mut[T, U](uint n_elts) -> vec[mutable U];

    fn refcount[T](vec[T] v) -> uint;

    fn vec_print_debug_info[T](vec[T] v);

    fn vec_from_vbuf[T](vbuf v, uint n_elts) -> vec[T];
}

fn alloc[T](uint n_elts) -> vec[T] {
    ret rustrt.vec_alloc[vec[T], T](n_elts);
}

fn alloc_mut[T](uint n_elts) -> vec[mutable T] {
    ret rustrt.vec_alloc_mut[vec[mutable T], T](n_elts);
}

fn refcount[T](vec[mutable? T] v) -> uint {
    auto r = rustrt.refcount[T](v);
    if (r == dbg.const_refcount) {
        ret r;
    } else {
        // -1 because calling this function incremented the refcount.
        ret  r - 1u;
    }
}

unsafe fn vec_from_vbuf[T](vbuf v, uint n_elts) -> vec[T] {
    ret rustrt.vec_from_vbuf[T](v, n_elts);
}

// FIXME: Remove me; this is a botch to get around rustboot's bad typechecker.
fn empty[T]() -> vec[T] {
    ret alloc[T](0u);
}

// FIXME: Remove me; this is a botch to get around rustboot's bad typechecker.
fn empty_mut[T]() -> vec[mutable T] {
    ret alloc_mut[T](0u);
}

type init_op[T] = fn(uint i) -> T;

fn init_fn[T](&init_op[T] op, uint n_elts) -> vec[T] {
    let vec[T] v = alloc[T](n_elts);
    let uint i = 0u;
    while (i < n_elts) {
        v += vec(op(i));
        i += 1u;
    }
    ret v;
}

fn init_fn_mut[T](&init_op[T] op, uint n_elts) -> vec[mutable T] {
    let vec[mutable T] v = alloc_mut[T](n_elts);
    let uint i = 0u;
    while (i < n_elts) {
        v += vec(mutable op(i));
        i += 1u;
    }
    ret v;
}

fn init_elt[T](&T t, uint n_elts) -> vec[T] {
    /**
     * FIXME (issue #81): should be:
     *
     * fn elt_op[T](&T x, uint i) -> T { ret x; }
     * let init_op[T] inner = bind elt_op[T](t, _);
     * ret init_fn[T](inner, n_elts);
     */
    let vec[T] v = alloc[T](n_elts);
    let uint i = n_elts;
    while (i > 0u) {
        i -= 1u;
        v += vec(t);
    }
    ret v;
}

fn init_elt_mut[T](&T t, uint n_elts) -> vec[mutable T] {
    let vec[mutable T] v = alloc_mut[T](n_elts);
    let uint i = n_elts;
    while (i > 0u) {
        i -= 1u;
        v += vec(mutable t);
    }
    ret v;
}

fn buf[T](vec[mutable? T] v) -> vbuf {
    ret rustrt.vec_buf[T](v, 0u);
}

fn len[T](vec[mutable? T] v) -> uint {
    ret rustrt.vec_len[T](v);
}

fn len_set[T](vec[mutable? T] v, uint n) {
    rustrt.vec_len_set[T](v, n);
}

fn buf_off[T](vec[mutable? T] v, uint offset) -> vbuf {
    check (offset < len[T](v));
    ret rustrt.vec_buf[T](v, offset);
}

fn print_debug_info[T](vec[mutable? T] v) {
    rustrt.vec_print_debug_info[T](v);
}

// Returns the last element of v.
fn last[T](vec[mutable? T] v) -> option.t[T] {
    auto l = len[T](v);
    if (l == 0u) {
        ret none[T];
    }
    ret some[T](v.(l - 1u));
}

// Returns elements from [start..end) from v.
fn slice[T](vec[mutable? T] v, uint start, uint end) -> vec[T] {
    check (start <= end);
    check (end <= len[T](v));
    auto result = alloc[T](end - start);
    let uint i = start;
    while (i < end) {
        result += vec(v.(i));
        i += 1u;
    }
    ret result;
}

fn shift[T](&mutable vec[mutable? T] v) -> T {
    auto ln = len[T](v);
    check(ln > 0u);
    auto e = v.(0);
    v = slice[T](v, 1u, ln);
    ret e;
}

fn pop[T](&mutable vec[mutable? T] v) -> T {
    auto ln = len[T](v);
    check(ln > 0u);
    ln -= 1u;
    auto e = v.(ln);
    v = slice[T](v, 0u, ln);
    ret e;
}

fn push[T](&mutable vec[mutable? T] v, &T t) {
    v += vec(t);
}

fn unshift[T](&mutable vec[mutable? T] v, &T t) {
    auto res = alloc[T](len[T](v) + 1u);
    res += vec(t);
    res += v;
    v = res;
}

fn grow[T](&mutable vec[mutable? T] v, uint n, &T initval) {
    let uint i = n;
    while (i > 0u) {
        i -= 1u;
        v += vec(initval);
    }
}

fn grow_set[T](&mutable vec[mutable T] v, uint index, &T initval, &T val) {
    auto length = _vec.len[mutable T](v);
    if (index >= length) {
        grow[mutable T](v, index - length + 1u, initval);
    }
    v.(index) = val;
}

fn map[T, U](&option.operator[T,U] f, &vec[mutable? T] v) -> vec[U] {
    let vec[U] u = alloc[U](len[T](v));
    for (T ve in v) {
        u += vec(f(ve));
    }
    ret u;
}

fn map2[T,U,V](&operator2[T,U,V] f, &vec[mutable? T] v0, &vec[mutable? U] v1)
        -> vec[V] {
    auto v0_len = len[T](v0);
    if (v0_len != len[U](v1)) {
        fail;
    }

    let vec[V] u = alloc[V](v0_len);
    auto i = 0u;
    while (i < v0_len) {
        u += vec(f(v0.(i), v1.(i)));
        i += 1u;
    }

    ret u;
}

fn find[T](fn (&T) -> bool f, &vec[mutable? T] v) -> option.t[T] {
    for (T elt in v) {
        if (f(elt)) {
            ret some[T](elt);
        }
    }

    ret none[T];
}

fn foldl[T, U](fn (&U, &T) -> U p, &U z, &vec[T] v) -> U {
    auto sz = len[T](v);

    if (sz == 0u) {
        ret z;
    }
    else {
        auto rest = slice[T](v, 1u, sz);

        ret (p(foldl[T,U](p, z, rest), v.(0)));
    }
}

fn unzip[T, U](&vec[tup(T, U)] v) -> tup(vec[T], vec[U]) {
    auto sz = len[tup(T, U)](v);

    if (sz == 0u) {
        ret tup(alloc[T](0u), alloc[U](0u));
    }
    else {
        auto rest = slice[tup(T, U)](v, 1u, sz);
        auto tl   = unzip[T, U](rest);
        auto a    = vec(v.(0)._0);
        auto b    = vec(v.(0)._1);
        ret tup(a + tl._0, b + tl._1);
    }
}

fn or(&vec[bool] v) -> bool {
    auto f = orb;
    be _vec.foldl[bool, bool](f, false, v);
}

fn clone[T](&vec[T] v) -> vec[T] {
    ret slice[T](v, 0u, len[T](v));
}

fn plus_option[T](&vec[T] v, &option.t[T] o) -> () {
    alt (o) {
        case (none[T]) {}
        case (some[T](?x)) { v += vec(x); }
    }
}

// TODO: Remove in favor of built-in "freeze" operation when it's implemented.
fn freeze[T](vec[mutable T] v) -> vec[T] {
    let vec[T] result = vec();
    for (T elem in v) {
        result += vec(elem);
    }
    ret result;
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<|endoftext|>