<commit_before>//! Request logging middleware
use std::env;
use std::fmt;
use std::fmt::{Display, Formatter};

use libc;
use time;
use regex::Regex;

use httprequest::HttpRequest;
use httpresponse::HttpResponse;
use middlewares::{Middleware, Started, Finished};

/// `Middleware` for logging request and response info to the terminal.
///
/// ## Usage
///
/// Create `Logger` middlewares with the specified `format`.
/// Default `Logger` could be created with `default` method, it uses the default format:
///
/// ```ignore
///  %a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T
/// ```
/// ```rust
/// extern crate actix_web;
/// use actix_web::Application;
/// use actix_web::middlewares::Logger;
///
/// fn main() {
///     let app = Application::default("/")
///         .middleware(Logger::default())
///         .middleware(Logger::new("%a %{User-Agent}i"))
///         .finish();
/// }
/// ```
///
/// ## Format
///
/// `%%`  The percent sign
///
/// `%a`  Remote IP-address (IP-address of proxy if using reverse proxy)
///
/// `%t`  Time when the request was started to process
///
/// `%P`  The process ID of the child that serviced the request
///
/// `%r`  First line of request
///
/// `%s`  Response status code
///
/// `%b`  Size of response in bytes, including HTTP headers
///
/// `%T` Time taken to serve the request, in seconds with floating fraction in .06f format
///
/// `%D`  Time taken to serve the request, in milliseconds
///
/// `%{FOO}i`  request.headers['FOO']
///
/// `%{FOO}o`  response.headers['FOO']
///
/// `%{FOO}e`  os.environ['FOO']
///
pub struct Logger {
    format: Format,
}

impl Logger {
    /// Create `Logger` middleware with the specified `format`.
    pub fn new(format: &str) -> Logger {
        Logger { format: Format::new(format) }
    }
}

impl Default for Logger {
    /// Create `Logger` middleware with format:
    ///
    /// ```ignore
    /// %a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T
    /// ```
    fn default() -> Logger {
        Logger { format: Format::default() }
    }
}

struct StartTime(time::Tm);

impl Logger {

    fn log(&self, req: &mut HttpRequest, resp: &HttpResponse) {
        let entry_time = req.extensions().get::<StartTime>().unwrap().0;

        let render = |fmt: &mut Formatter| {
            for unit in &self.format.0 {
                unit.render(fmt, req, resp, entry_time)?;
            }
            Ok(())
        };
        info!("{}", FormatDisplay(&render));
    }
}

impl Middleware for Logger {

    fn start(&self, req: &mut HttpRequest) -> Started {
        req.extensions().insert(StartTime(time::now()));
        Started::Done
    }

    fn finish(&self, req: &mut HttpRequest, resp: &HttpResponse) -> Finished {
        self.log(req, resp);
        Finished::Done
    }
}

/// A formatting style for the `Logger`, consisting of multiple
/// `FormatText`s concatenated into one line.
#[derive(Clone)]
#[doc(hidden)]
struct Format(Vec<FormatText>);

impl Default for Format {
    /// Return the default formatting style for the `Logger`:
    fn default() -> Format {
        Format::new(r#"%a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T"#)
    }
}

impl Format {
    /// Create a `Format` from a format string.
    ///
    /// Returns `None` if the format string syntax is incorrect.
    pub fn new(s: &str) -> Format {
        trace!("Access log format: {}", s);
        let fmt = Regex::new(r"%(\{([A-Za-z0-9\-_]+)\}([ioe])|[atPrsbTD]?)").unwrap();

        let mut idx = 0;
        let mut results = Vec::new();
        for cap in fmt.captures_iter(s) {
            let m = cap.get(0).unwrap();
            let pos = m.start();
            if idx != pos {
                results.push(FormatText::Str(s[idx..pos].to_owned()));
            }
            idx = m.end();

            if let Some(key) = cap.get(2) {
                results.push(
                    match cap.get(3).unwrap().as_str() {
                        "i" => FormatText::RequestHeader(key.as_str().to_owned()),
                        "o" => FormatText::ResponseHeader(key.as_str().to_owned()),
                        "e" => FormatText::EnvironHeader(key.as_str().to_owned()),
                        _ => unreachable!(),
                    })
            } else {
                let m = cap.get(1).unwrap();
                results.push(
                    match m.as_str() {
                        "%" => FormatText::Percent,
                        "a" => FormatText::RemoteAddr,
                        "t" => FormatText::RequestTime,
                        "P" => FormatText::Pid,
                        "r" => FormatText::RequestLine,
                        "s" => FormatText::ResponseStatus,
                        "b" => FormatText::ResponseSize,
                        "T" => FormatText::Time,
                        "D" => FormatText::TimeMillis,
                        _ => FormatText::Str(m.as_str().to_owned()),
                    }
                );
            }
        }
        if idx != s.len() {
            results.push(FormatText::Str(s[idx..].to_owned()));
        }

        Format(results)
    }
}

/// A string of text to be logged. This is either one of the data
/// fields supported by the `Logger`, or a custom `String`.
#[doc(hidden)]
#[derive(Debug, Clone)]
pub enum FormatText {
    Str(String),
    Pid,
    Percent,
    RequestLine,
    RequestTime,
    ResponseStatus,
    ResponseSize,
    Time,
    TimeMillis,
    RemoteAddr,
    RequestHeader(String),
    ResponseHeader(String),
    EnvironHeader(String),
}

impl FormatText {

    fn render(&self, fmt: &mut Formatter,
              req: &HttpRequest,
              resp: &HttpResponse,
              entry_time: time::Tm) -> Result<(), fmt::Error>
    {
        match *self {
            FormatText::Str(ref string) => fmt.write_str(string),
            FormatText::Percent => "%".fmt(fmt),
            FormatText::RequestLine => {
                if req.query_string().is_empty() {
                    fmt.write_fmt(format_args!(
                        "{} {} {:?}",
                        req.method(), req.path(), req.version()))
                } else {
                    fmt.write_fmt(format_args!(
                        "{} {}?{} {:?}",
                        req.method(), req.path(), req.query_string(), req.version()))
                }
            },
            FormatText::ResponseStatus => resp.status().as_u16().fmt(fmt),
            FormatText::ResponseSize => resp.response_size().fmt(fmt),
            FormatText::Pid => unsafe{libc::getpid().fmt(fmt)},
            FormatText::Time => {
                let response_time = time::now() - entry_time;
                let response_time = (response_time.num_seconds() * 1000) as f64 +
                    (response_time.num_nanoseconds().unwrap_or(0) as f64) / 1000000000.0;

                fmt.write_fmt(format_args!("{:.6}", response_time))
            },
            FormatText::TimeMillis => {
                let response_time = time::now() - entry_time;
                let response_time_ms = (response_time.num_seconds() * 1000) as f64 +
                    (response_time.num_nanoseconds().unwrap_or(0) as f64) / 1000000.0;

                fmt.write_fmt(format_args!("{:.6}", response_time_ms))
            },
            FormatText::RemoteAddr => {
                if let Some(addr) = req.remote() {
                    addr.fmt(fmt)
                } else {
                    "-".fmt(fmt)
                }
            }
            FormatText::RequestTime => {
                entry_time.strftime("[%d/%b/%Y:%H:%M:%S %z]")
                    .unwrap()
                    .fmt(fmt)
            }
            FormatText::RequestHeader(ref name) => {
                let s = if let Some(val) = req.headers().get(name) {
                    if let Ok(s) = val.to_str() { s } else { "-" }
                } else {
                    "-"
                };
                fmt.write_fmt(format_args!("{}", s))
            }
            FormatText::ResponseHeader(ref name) => {
                let s = if let Some(val) = resp.headers().get(name) {
                    if let Ok(s) = val.to_str() { s } else { "-" }
                } else {
                    "-"
                };
                fmt.write_fmt(format_args!("{}", s))
            }
            FormatText::EnvironHeader(ref name) => {
                if let Ok(val) = env::var(name) {
                    fmt.write_fmt(format_args!("{}", val))
                } else {
                    "-".fmt(fmt)
                }
            }
        }
    }
}

pub(crate) struct FormatDisplay<'a>(
    &'a Fn(&mut Formatter) -> Result<(), fmt::Error>);

impl<'a> fmt::Display for FormatDisplay<'a> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        (self.0)(fmt)
    }
}

#[cfg(test)]
mod tests {
    use Body;
    use super::*;
    use time;
    use http::{Method, Version, StatusCode};
    use http::header::{self, HeaderMap};

    #[test]
    fn test_default_logger() {
        let format = Format::default();

        let mut headers = HeaderMap::new();
        headers.insert(header::USER_AGENT, header::HeaderValue::from_static("ACTIX-WEB"));
        let req = HttpRequest::new(
            Method::GET, "/".to_owned(), Version::HTTP_11, headers, String::new());
        let resp = HttpResponse::builder(StatusCode::OK)
            .force_close().body(Body::Empty).unwrap();
        let entry_time = time::now();

        let render = |fmt: &mut Formatter| {
            for unit in format.0.iter() {
                unit.render(fmt, &req, &resp, entry_time)?;
            }
            Ok(())
        };
        let s = format!("{}", FormatDisplay(&render));
        assert!(s.contains("GET / HTTP/1.1"));
        assert!(s.contains("200 0"));
        assert!(s.contains("ACTIX-WEB"));

        let req = HttpRequest::new(
            Method::GET, "/?test".to_owned(), Version::HTTP_11, HeaderMap::new(), String::new());
        let resp = HttpResponse::builder(StatusCode::OK)
            .force_close().body(Body::Empty).unwrap();
        let entry_time = time::now();

        let render = |fmt: &mut Formatter| {
            for unit in format.0.iter() {
                unit.render(fmt, &req, &resp, entry_time)?;
            }
            Ok(())
        };
        let s = format!("{}", FormatDisplay(&render));
        assert!(s.contains("GET /?test HTTP/1.1"));
    }
}
<commit_msg>more logger tests<commit_after>//! Request logging middleware
use std::env;
use std::fmt;
use std::fmt::{Display, Formatter};

use libc;
use time;
use regex::Regex;

use httprequest::HttpRequest;
use httpresponse::HttpResponse;
use middlewares::{Middleware, Started, Finished};

/// `Middleware` for logging request and response info to the terminal.
///
/// ## Usage
///
/// Create `Logger` middlewares with the specified `format`.
/// Default `Logger` could be created with `default` method, it uses the default format:
///
/// ```ignore
///  %a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T
/// ```
/// ```rust
/// extern crate actix_web;
/// use actix_web::Application;
/// use actix_web::middlewares::Logger;
///
/// fn main() {
///     let app = Application::default("/")
///         .middleware(Logger::default())
///         .middleware(Logger::new("%a %{User-Agent}i"))
///         .finish();
/// }
/// ```
///
/// ## Format
///
/// `%%`  The percent sign
///
/// `%a`  Remote IP-address (IP-address of proxy if using reverse proxy)
///
/// `%t`  Time when the request was started to process
///
/// `%P`  The process ID of the child that serviced the request
///
/// `%r`  First line of request
///
/// `%s`  Response status code
///
/// `%b`  Size of response in bytes, including HTTP headers
///
/// `%T` Time taken to serve the request, in seconds with floating fraction in .06f format
///
/// `%D`  Time taken to serve the request, in milliseconds
///
/// `%{FOO}i`  request.headers['FOO']
///
/// `%{FOO}o`  response.headers['FOO']
///
/// `%{FOO}e`  os.environ['FOO']
///
pub struct Logger {
    format: Format,
}

impl Logger {
    /// Create `Logger` middleware with the specified `format`.
    pub fn new(format: &str) -> Logger {
        Logger { format: Format::new(format) }
    }
}

impl Default for Logger {
    /// Create `Logger` middleware with format:
    ///
    /// ```ignore
    /// %a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T
    /// ```
    fn default() -> Logger {
        Logger { format: Format::default() }
    }
}

struct StartTime(time::Tm);

impl Logger {

    fn log(&self, req: &mut HttpRequest, resp: &HttpResponse) {
        let entry_time = req.extensions().get::<StartTime>().unwrap().0;

        let render = |fmt: &mut Formatter| {
            for unit in &self.format.0 {
                unit.render(fmt, req, resp, entry_time)?;
            }
            Ok(())
        };
        info!("{}", FormatDisplay(&render));
    }
}

impl Middleware for Logger {

    fn start(&self, req: &mut HttpRequest) -> Started {
        req.extensions().insert(StartTime(time::now()));
        Started::Done
    }

    fn finish(&self, req: &mut HttpRequest, resp: &HttpResponse) -> Finished {
        self.log(req, resp);
        Finished::Done
    }
}

/// A formatting style for the `Logger`, consisting of multiple
/// `FormatText`s concatenated into one line.
#[derive(Clone)]
#[doc(hidden)]
struct Format(Vec<FormatText>);

impl Default for Format {
    /// Return the default formatting style for the `Logger`:
    fn default() -> Format {
        Format::new(r#"%a %t "%r" %s %b "%{Referrer}i" "%{User-Agent}i" %T"#)
    }
}

impl Format {
    /// Create a `Format` from a format string.
    ///
    /// Returns `None` if the format string syntax is incorrect.
    pub fn new(s: &str) -> Format {
        trace!("Access log format: {}", s);
        let fmt = Regex::new(r"%(\{([A-Za-z0-9\-_]+)\}([ioe])|[atPrsbTD]?)").unwrap();

        let mut idx = 0;
        let mut results = Vec::new();
        for cap in fmt.captures_iter(s) {
            let m = cap.get(0).unwrap();
            let pos = m.start();
            if idx != pos {
                results.push(FormatText::Str(s[idx..pos].to_owned()));
            }
            idx = m.end();

            if let Some(key) = cap.get(2) {
                results.push(
                    match cap.get(3).unwrap().as_str() {
                        "i" => FormatText::RequestHeader(key.as_str().to_owned()),
                        "o" => FormatText::ResponseHeader(key.as_str().to_owned()),
                        "e" => FormatText::EnvironHeader(key.as_str().to_owned()),
                        _ => unreachable!(),
                    })
            } else {
                let m = cap.get(1).unwrap();
                results.push(
                    match m.as_str() {
                        "%" => FormatText::Percent,
                        "a" => FormatText::RemoteAddr,
                        "t" => FormatText::RequestTime,
                        "P" => FormatText::Pid,
                        "r" => FormatText::RequestLine,
                        "s" => FormatText::ResponseStatus,
                        "b" => FormatText::ResponseSize,
                        "T" => FormatText::Time,
                        "D" => FormatText::TimeMillis,
                        _ => FormatText::Str(m.as_str().to_owned()),
                    }
                );
            }
        }
        if idx != s.len() {
            results.push(FormatText::Str(s[idx..].to_owned()));
        }

        Format(results)
    }
}

/// A string of text to be logged. This is either one of the data
/// fields supported by the `Logger`, or a custom `String`.
#[doc(hidden)]
#[derive(Debug, Clone)]
pub enum FormatText {
    Str(String),
    Pid,
    Percent,
    RequestLine,
    RequestTime,
    ResponseStatus,
    ResponseSize,
    Time,
    TimeMillis,
    RemoteAddr,
    RequestHeader(String),
    ResponseHeader(String),
    EnvironHeader(String),
}

impl FormatText {

    fn render(&self, fmt: &mut Formatter,
              req: &HttpRequest,
              resp: &HttpResponse,
              entry_time: time::Tm) -> Result<(), fmt::Error>
    {
        match *self {
            FormatText::Str(ref string) => fmt.write_str(string),
            FormatText::Percent => "%".fmt(fmt),
            FormatText::RequestLine => {
                if req.query_string().is_empty() {
                    fmt.write_fmt(format_args!(
                        "{} {} {:?}",
                        req.method(), req.path(), req.version()))
                } else {
                    fmt.write_fmt(format_args!(
                        "{} {}?{} {:?}",
                        req.method(), req.path(), req.query_string(), req.version()))
                }
            },
            FormatText::ResponseStatus => resp.status().as_u16().fmt(fmt),
            FormatText::ResponseSize => resp.response_size().fmt(fmt),
            FormatText::Pid => unsafe{libc::getpid().fmt(fmt)},
            FormatText::Time => {
                let response_time = time::now() - entry_time;
                let response_time = (response_time.num_seconds() * 1000) as f64 +
                    (response_time.num_nanoseconds().unwrap_or(0) as f64) / 1000000000.0;

                fmt.write_fmt(format_args!("{:.6}", response_time))
            },
            FormatText::TimeMillis => {
                let response_time = time::now() - entry_time;
                let response_time_ms = (response_time.num_seconds() * 1000) as f64 +
                    (response_time.num_nanoseconds().unwrap_or(0) as f64) / 1000000.0;

                fmt.write_fmt(format_args!("{:.6}", response_time_ms))
            },
            FormatText::RemoteAddr => {
                if let Some(addr) = req.remote() {
                    addr.fmt(fmt)
                } else {
                    "-".fmt(fmt)
                }
            }
            FormatText::RequestTime => {
                entry_time.strftime("[%d/%b/%Y:%H:%M:%S %z]")
                    .unwrap()
                    .fmt(fmt)
            }
            FormatText::RequestHeader(ref name) => {
                let s = if let Some(val) = req.headers().get(name) {
                    if let Ok(s) = val.to_str() { s } else { "-" }
                } else {
                    "-"
                };
                fmt.write_fmt(format_args!("{}", s))
            }
            FormatText::ResponseHeader(ref name) => {
                let s = if let Some(val) = resp.headers().get(name) {
                    if let Ok(s) = val.to_str() { s } else { "-" }
                } else {
                    "-"
                };
                fmt.write_fmt(format_args!("{}", s))
            }
            FormatText::EnvironHeader(ref name) => {
                if let Ok(val) = env::var(name) {
                    fmt.write_fmt(format_args!("{}", val))
                } else {
                    "-".fmt(fmt)
                }
            }
        }
    }
}

pub(crate) struct FormatDisplay<'a>(
    &'a Fn(&mut Formatter) -> Result<(), fmt::Error>);

impl<'a> fmt::Display for FormatDisplay<'a> {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), fmt::Error> {
        (self.0)(fmt)
    }
}

#[cfg(test)]
mod tests {
    use Body;
    use super::*;
    use time;
    use http::{Method, Version, StatusCode};
    use http::header::{self, HeaderMap};

    #[test]
    fn test_logger() {
        let logger = Logger::new("%% %{User-Agent}i %{X-Test}o %{HOME}e %D test");

        let mut headers = HeaderMap::new();
        headers.insert(header::USER_AGENT, header::HeaderValue::from_static("ACTIX-WEB"));
        let mut req = HttpRequest::new(
            Method::GET, "/".to_owned(), Version::HTTP_11, headers, String::new());
        let resp = HttpResponse::builder(StatusCode::OK)
            .header("X-Test", "ttt")
            .force_close().body(Body::Empty).unwrap();

        match logger.start(&mut req) {
            Started::Done => (),
            _ => panic!(),
        }
        match logger.finish(&mut req, &resp) {
            Finished::Done => (),
            _ => panic!(),
        }
    }

    #[test]
    fn test_default_format() {
        let format = Format::default();

        let mut headers = HeaderMap::new();
        headers.insert(header::USER_AGENT, header::HeaderValue::from_static("ACTIX-WEB"));
        let req = HttpRequest::new(
            Method::GET, "/".to_owned(), Version::HTTP_11, headers, String::new());
        let resp = HttpResponse::builder(StatusCode::OK)
            .force_close().body(Body::Empty).unwrap();
        let entry_time = time::now();

        let render = |fmt: &mut Formatter| {
            for unit in format.0.iter() {
                unit.render(fmt, &req, &resp, entry_time)?;
            }
            Ok(())
        };
        let s = format!("{}", FormatDisplay(&render));
        assert!(s.contains("GET / HTTP/1.1"));
        assert!(s.contains("200 0"));
        assert!(s.contains("ACTIX-WEB"));

        let req = HttpRequest::new(
            Method::GET, "/".to_owned(), Version::HTTP_11, HeaderMap::new(), "test".to_owned());
        let resp = HttpResponse::builder(StatusCode::OK)
            .force_close().body(Body::Empty).unwrap();
        let entry_time = time::now();

        let render = |fmt: &mut Formatter| {
            for unit in format.0.iter() {
                unit.render(fmt, &req, &resp, entry_time)?;
            }
            Ok(())
        };
        let s = format!("{}", FormatDisplay(&render));
        assert!(s.contains("GET /?test HTTP/1.1"));
    }
}
<|endoftext|>


<commit_before>use std;
import int;

tag tree { nil; node(@tree, @tree, int); }

fn item_check(t: @tree) -> int {
    alt *t {
      nil. { ret 0; }
      node(left, right, item) {
        ret item + item_check(left) - item_check(right);
      }
    }
}

fn bottom_up_tree(item: int, depth: int) -> @tree {
    if depth > 0 {
        ret @node(bottom_up_tree(2 * item - 1, depth - 1),
                  bottom_up_tree(2 * item, depth - 1), item);
    } else { ret @nil; }
}

fn main() {
    let n = 8;
    let min_depth = 4;
    let max_depth;
    if min_depth + 2 > n {
        max_depth = min_depth + 2;
    } else { max_depth = n; }
    let stretch_depth = max_depth + 1;
    let stretch_tree = bottom_up_tree(0, stretch_depth);
    #debug("stretch tree of depth %d\t check: %d", stretch_depth,
           item_check(stretch_tree));
    let long_lived_tree = bottom_up_tree(0, max_depth);
    let depth = min_depth;
    while depth <= max_depth {
        let iterations = int::pow(2, max_depth - depth + min_depth as uint);
        let chk = 0;
        let i = 1;
        while i <= iterations {
            let temp_tree = bottom_up_tree(i, depth);
            chk += item_check(temp_tree);
            temp_tree = bottom_up_tree(-i, depth);
            chk += item_check(temp_tree);
            i += 1;
        }
        #debug("%d\t trees of depth %d\t check: %d", iterations * 2, depth,
               chk);
        depth += 2;
    }
    #debug("long lived trees of depth %d\t check: %d", max_depth,
           item_check(long_lived_tree));
}
<commit_msg>bench: Update shootout-binarytrees to use unique boxes<commit_after>use std;
import int;

tag tree { nil; node(~tree, ~tree, int); }

fn item_check(t: ~tree) -> int {
    alt *t {
      nil. { ret 0; }
      node(left, right, item) {
        ret item + item_check(left) - item_check(right);
      }
    }
}

fn bottom_up_tree(item: int, depth: int) -> ~tree {
    if depth > 0 {
        ret ~node(bottom_up_tree(2 * item - 1, depth - 1),
                  bottom_up_tree(2 * item, depth - 1), item);
    } else { ret ~nil; }
}

fn main() {
    let n = 8;
    let min_depth = 4;
    let max_depth;
    if min_depth + 2 > n {
        max_depth = min_depth + 2;
    } else { max_depth = n; }
    let stretch_depth = max_depth + 1;
    let stretch_tree = bottom_up_tree(0, stretch_depth);
    #debug("stretch tree of depth %d\t check: %d", stretch_depth,
           item_check(stretch_tree));
    let long_lived_tree = bottom_up_tree(0, max_depth);
    let depth = min_depth;
    while depth <= max_depth {
        let iterations = int::pow(2, max_depth - depth + min_depth as uint);
        let chk = 0;
        let i = 1;
        while i <= iterations {
            let temp_tree = bottom_up_tree(i, depth);
            chk += item_check(temp_tree);
            temp_tree = bottom_up_tree(-i, depth);
            chk += item_check(temp_tree);
            i += 1;
        }
        #debug("%d\t trees of depth %d\t check: %d", iterations * 2, depth,
               chk);
        depth += 2;
    }
    #debug("long lived trees of depth %d\t check: %d", max_depth,
           item_check(long_lived_tree));
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (c) 2018 Dragoș Tiselice
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! A `mod` containing error data structures.

use std::cmp;
use std::error;
use std::fmt;
use std::mem;

use RuleType;
use position::Position;
use span::Span;

/// A `struct` defining errors.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Error<R> {
    /// Variant of the error
    pub variant: ErrorVariant<R>,
    pub location: InputLocation,
    line: String,
    continued_line: Option<String>,
    start: (usize, usize),
    end: Option<(usize, usize)>
}

/// An `enum` describing `Error` variants.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum ErrorVariant<R> {
    /// Generated parsing error with expected and unexpected `Rule`s
    ParsingError {
        /// Positive attempts
        positives: Vec<R>,
        /// Negative attempts
        negatives: Vec<R>
    },
    /// Custom error with a message
    CustomError {
        /// Short explanation
        message: String
    }
}

/// An `enum` describing where the `Error` has occurred.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum InputLocation {
    /// `Error` was created by `Error::new_from_pos`
    Pos(usize),
    /// `Error` was created by `Error::new_from_span`
    Span((usize, usize))
}

impl<R: RuleType> Error<R> {
    pub fn new_from_pos(variant: ErrorVariant<R>, pos: Position) -> Error<R> {
        Error {
            variant,
            location: InputLocation::Pos(pos.pos()),
            line: pos.line_of().to_owned(),
            continued_line: None,
            start: pos.line_col(),
            end: None
        }
    }

    pub fn new_from_span(variant: ErrorVariant<R>, span: Span) -> Error<R> {
        let continued_line = if span.start_pos().line_col().0 != span.end_pos().line_col().0 {
            Some(span.end_pos().line_of().to_owned())
        } else {
            None
        };

        Error {
            variant,
            location: InputLocation::Span((span.start(), span.end())),
            line: span.start_pos().line_of().to_owned(),
            continued_line,
            start: span.start_pos().line_col(),
            end: Some(span.end_pos().line_col())
        }
    }

    /// Renames all `Rule`s from a `ParsingError` variant returning a `CustomErrorPos`. It does
    /// nothing when called on `CustomErrorPos` and `CustomErrorSpan` variants.
    ///
    /// Useful in order to rename verbose rules or have detailed per-`Rule` formatting.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::error::{Error, ErrorVariant};
    /// # use pest::Position;
    /// # #[allow(non_camel_case_types)]
    /// # #[allow(dead_code)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// # enum Rule {
    /// #     open_paren,
    /// #     closed_paren
    /// # }
    /// # let input = "";
    /// # let pos = Position::from_start(input);
    /// Error::new_from_pos(
    ///     ErrorVariant::ParsingError {
    ///         positives: vec![Rule::open_paren],
    ///         negatives: vec![Rule::closed_paren]
    ///     },
    ///     pos
    /// ).renamed_rules(|rule| {
    ///     match *rule {
    ///         Rule::open_paren => "(".to_owned(),
    ///         Rule::closed_paren => "closed paren".to_owned()
    ///     }
    /// });
    /// ```
    pub fn renamed_rules<F>(mut self, f: F) -> Error<R>
    where
        F: FnMut(&R) -> String
    {
        let variant = match self.variant {
            ErrorVariant::ParsingError {
                positives,
                negatives
            } => {
                let message = Error::parsing_error_message(&positives, &negatives, f);
                ErrorVariant::CustomError { message }
            }
            variant => variant
        };

        self.variant = variant;

        self
    }

    fn spacing(&self) -> String {
        let line = if let Some((line, _)) = self.end {
            cmp::max(self.start.0, line)
        } else {
            self.start.0
        };

        let line_str_len = format!("{}", line).len();

        let mut spacing = String::new();
        for _ in 0..line_str_len {
            spacing.push(' ');
        }

        spacing
    }

    fn underline(&self) -> String {
        let mut underline = String::new();

        let mut start = self.start.1;
        let end = if let Some((_, mut end)) = self.end {
            let inverted_cols = start > end;
            if inverted_cols {
                mem::swap(&mut start, &mut end);
                start -= 1;
                end += 1;
            }

            Some(end)
        } else {
            None
        };
        let offset = start - 1;

        for _ in 0..offset {
            underline.push(' ');
        }

        if let Some(end) = end {
            if end - start > 1 {
                underline.push('^');
                for _ in 2..(end - start) {
                    underline.push('-');
                }
                underline.push('^');
            } else {
                underline.push('^');
            }
        } else {
            underline.push_str("^---")
        }

        underline
    }

    fn message(&self) -> String {
        match self.variant {
            ErrorVariant::ParsingError {
                ref positives,
                ref negatives
            } => Error::parsing_error_message(positives, negatives, |r| format!("{:?}", r)),
            ErrorVariant::CustomError { ref message } => {
                message.clone()
            }
        }
    }

    fn parsing_error_message<F>(positives: &[R], negatives: &[R], mut f: F) -> String
    where
        F: FnMut(&R) -> String
    {
        match (negatives.is_empty(), positives.is_empty()) {
            (false, false) => format!(
                "unexpected {}; expected {}",
                Error::enumerate(negatives, &mut f),
                Error::enumerate(positives, &mut f)
            ),
            (false, true) => format!("unexpected {}", Error::enumerate(negatives, &mut f)),
            (true, false) => format!("expected {}", Error::enumerate(positives, &mut f)),
            (true, true) => "unknown parsing error".to_owned()
        }
    }

    fn enumerate<F>(rules: &[R], f: &mut F) -> String
    where
        F: FnMut(&R) -> String
    {
        match rules.len() {
            1 => f(&rules[0]),
            2 => format!("{} or {}", f(&rules[0]), f(&rules[1])),
            l => {
                let separated = rules
                    .iter()
                    .take(l - 1)
                    .map(|r| f(r))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("{}, or {}", separated, f(&rules[l - 1]))
            }
        }
    }

    pub(crate) fn format(&self) -> String {
        let spacing = self.spacing();

        let mut result = format!("{}--> {}:{}\n", spacing, self.start.0, self.start.1);
        result.push_str(&format!("{} |\n", spacing));

        let underlined = self.underline();

        if let Some(end) = self.end {
            if let Some(ref continued_line) = self.continued_line {
                result.push_str(&format!("{:1$} | ", self.start.0, spacing.len()));
                result.push_str(&format!("{}\n", self.line));

                if end.0 - self.start.0 > 1 {
                    result.push_str(&format!("{} | ...\n", spacing));
                }

                result.push_str(&format!("{:1$} | ", end.0, spacing.len()));
                result.push_str(&format!("{}\n", continued_line));
                result.push_str(&format!("{} | {}\n", spacing, underlined));

                result.push_str(&format!("{} |\n", spacing));
                result.push_str(&format!("{} = {}", spacing, self.message()));

                return result;
            }
        }

        result.push_str(&format!("{} | ", self.start.0));
        result.push_str(&format!("{}\n", self.line));
        result.push_str(&format!("{} | {}\n", spacing, underlined));

        result.push_str(&format!("{} |\n", spacing));
        result.push_str(&format!("{} = {}", spacing, self.message()));

        result
    }
}

impl<R: RuleType> fmt::Display for Error<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.format())
    }
}

impl<'i, R: RuleType> error::Error for Error<R> {
    fn description(&self) -> &str {
        match self.variant {
            ErrorVariant::ParsingError { .. } => "parsing error",
            ErrorVariant::CustomError { ref message } => message
        }
    }
}

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

    #[test]
    fn display_parsing_error_mixed() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2, 3],
                negatives: vec![4, 5, 6]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 4, 5, or 6; expected 1, 2, or 3",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_positives() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2],
                negatives: vec![]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = expected 1 or 2",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_negatives() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![],
                negatives: vec![4, 5, 6]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 4, 5, or 6",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_unknown() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![],
                negatives: vec![]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unknown parsing error",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_pos() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_two_lines() {
        let input = b"ab\ncd\nefgh";
        let start = unsafe { position::new(input, 4) };
        let end = unsafe { position::new(input, 9) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "3 | efgh",
                "  |  ^^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_three_lines() {
        let input = b"ab\ncd\nefgh";
        let start = unsafe { position::new(input, 1) };
        let end = unsafe { position::new(input, 9) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 1:2",
                "  |",
                "1 | ab",
                "  | ...",
                "3 | efgh",
                "  |  ^^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_two_lines_inverted_cols() {
        let input = b"abcdef\ngh";
        let start = unsafe { position::new(input, 5) };
        let end = unsafe { position::new(input, 8) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 1:6",
                "  |",
                "1 | abcdef",
                "2 | gh",
                "  | ^----^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn mapped_parsing_error() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2, 3],
                negatives: vec![4, 5, 6]
            },
            pos
        ).renamed_rules(|n| format!("{}", n + 1));

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 5, 6, or 7; expected 2, 3, or 4",
            ].join("\n")
        );
    }
}
<commit_msg>Improved control flow.<commit_after>// pest. The Elegant Parser
// Copyright (c) 2018 Dragoș Tiselice
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! A `mod` containing error data structures.

use std::cmp;
use std::error;
use std::fmt;
use std::mem;

use RuleType;
use position::Position;
use span::Span;

/// A `struct` defining errors.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Error<R> {
    /// Variant of the error
    pub variant: ErrorVariant<R>,
    pub location: InputLocation,
    line: String,
    continued_line: Option<String>,
    start: (usize, usize),
    end: Option<(usize, usize)>
}

/// An `enum` describing `Error` variants.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum ErrorVariant<R> {
    /// Generated parsing error with expected and unexpected `Rule`s
    ParsingError {
        /// Positive attempts
        positives: Vec<R>,
        /// Negative attempts
        negatives: Vec<R>
    },
    /// Custom error with a message
    CustomError {
        /// Short explanation
        message: String
    }
}

/// An `enum` describing where the `Error` has occurred.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum InputLocation {
    /// `Error` was created by `Error::new_from_pos`
    Pos(usize),
    /// `Error` was created by `Error::new_from_span`
    Span((usize, usize))
}

impl<R: RuleType> Error<R> {
    pub fn new_from_pos(variant: ErrorVariant<R>, pos: Position) -> Error<R> {
        Error {
            variant,
            location: InputLocation::Pos(pos.pos()),
            line: pos.line_of().to_owned(),
            continued_line: None,
            start: pos.line_col(),
            end: None
        }
    }

    pub fn new_from_span(variant: ErrorVariant<R>, span: Span) -> Error<R> {
        let continued_line = if span.start_pos().line_col().0 != span.end_pos().line_col().0 {
            Some(span.end_pos().line_of().to_owned())
        } else {
            None
        };

        Error {
            variant,
            location: InputLocation::Span((span.start(), span.end())),
            line: span.start_pos().line_of().to_owned(),
            continued_line,
            start: span.start_pos().line_col(),
            end: Some(span.end_pos().line_col())
        }
    }

    /// Renames all `Rule`s from a `ParsingError` variant returning a `CustomErrorPos`. It does
    /// nothing when called on `CustomErrorPos` and `CustomErrorSpan` variants.
    ///
    /// Useful in order to rename verbose rules or have detailed per-`Rule` formatting.
    ///
    /// # Examples
    ///
    /// ```
    /// # use pest::error::{Error, ErrorVariant};
    /// # use pest::Position;
    /// # #[allow(non_camel_case_types)]
    /// # #[allow(dead_code)]
    /// # #[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
    /// # enum Rule {
    /// #     open_paren,
    /// #     closed_paren
    /// # }
    /// # let input = "";
    /// # let pos = Position::from_start(input);
    /// Error::new_from_pos(
    ///     ErrorVariant::ParsingError {
    ///         positives: vec![Rule::open_paren],
    ///         negatives: vec![Rule::closed_paren]
    ///     },
    ///     pos
    /// ).renamed_rules(|rule| {
    ///     match *rule {
    ///         Rule::open_paren => "(".to_owned(),
    ///         Rule::closed_paren => "closed paren".to_owned()
    ///     }
    /// });
    /// ```
    pub fn renamed_rules<F>(mut self, f: F) -> Error<R>
    where
        F: FnMut(&R) -> String
    {
        let variant = match self.variant {
            ErrorVariant::ParsingError {
                positives,
                negatives
            } => {
                let message = Error::parsing_error_message(&positives, &negatives, f);
                ErrorVariant::CustomError { message }
            }
            variant => variant
        };

        self.variant = variant;

        self
    }

    fn spacing(&self) -> String {
        let line = if let Some((line, _)) = self.end {
            cmp::max(self.start.0, line)
        } else {
            self.start.0
        };

        let line_str_len = format!("{}", line).len();

        let mut spacing = String::new();
        for _ in 0..line_str_len {
            spacing.push(' ');
        }

        spacing
    }

    fn underline(&self) -> String {
        let mut underline = String::new();

        let mut start = self.start.1;
        let end = if let Some((_, mut end)) = self.end {
            let inverted_cols = start > end;
            if inverted_cols {
                mem::swap(&mut start, &mut end);
                start -= 1;
                end += 1;
            }

            Some(end)
        } else {
            None
        };
        let offset = start - 1;

        for _ in 0..offset {
            underline.push(' ');
        }

        if let Some(end) = end {
            if end - start > 1 {
                underline.push('^');
                for _ in 2..(end - start) {
                    underline.push('-');
                }
                underline.push('^');
            } else {
                underline.push('^');
            }
        } else {
            underline.push_str("^---")
        }

        underline
    }

    fn message(&self) -> String {
        match self.variant {
            ErrorVariant::ParsingError {
                ref positives,
                ref negatives
            } => Error::parsing_error_message(positives, negatives, |r| format!("{:?}", r)),
            ErrorVariant::CustomError { ref message } => {
                message.clone()
            }
        }
    }

    fn parsing_error_message<F>(positives: &[R], negatives: &[R], mut f: F) -> String
    where
        F: FnMut(&R) -> String
    {
        match (negatives.is_empty(), positives.is_empty()) {
            (false, false) => format!(
                "unexpected {}; expected {}",
                Error::enumerate(negatives, &mut f),
                Error::enumerate(positives, &mut f)
            ),
            (false, true) => format!("unexpected {}", Error::enumerate(negatives, &mut f)),
            (true, false) => format!("expected {}", Error::enumerate(positives, &mut f)),
            (true, true) => "unknown parsing error".to_owned()
        }
    }

    fn enumerate<F>(rules: &[R], f: &mut F) -> String
    where
        F: FnMut(&R) -> String
    {
        match rules.len() {
            1 => f(&rules[0]),
            2 => format!("{} or {}", f(&rules[0]), f(&rules[1])),
            l => {
                let separated = rules
                    .iter()
                    .take(l - 1)
                    .map(|r| f(r))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("{}, or {}", separated, f(&rules[l - 1]))
            }
        }
    }

    pub(crate) fn format(&self) -> String {
        let spacing = self.spacing();

        let mut result = format!("{}--> {}:{}\n", spacing, self.start.0, self.start.1);
        result.push_str(&format!("{} |\n", spacing));

        let underlined = self.underline();

        if let (Some(end), Some(ref continued_line)) = (self.end, &self.continued_line) {
            result.push_str(&format!("{:1$} | ", self.start.0, spacing.len()));
            result.push_str(&format!("{}\n", self.line));

            if end.0 - self.start.0 > 1 {
                result.push_str(&format!("{} | ...\n", spacing));
            }

            result.push_str(&format!("{:1$} | ", end.0, spacing.len()));
            result.push_str(&format!("{}\n", continued_line));
            result.push_str(&format!("{} | {}\n", spacing, underlined));

            result.push_str(&format!("{} |\n", spacing));
            result.push_str(&format!("{} = {}", spacing, self.message()));
        } else {
            result.push_str(&format!("{} | ", self.start.0));
            result.push_str(&format!("{}\n", self.line));
            result.push_str(&format!("{} | {}\n", spacing, underlined));

            result.push_str(&format!("{} |\n", spacing));
            result.push_str(&format!("{} = {}", spacing, self.message()));
        }

        result
    }
}

impl<R: RuleType> fmt::Display for Error<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.format())
    }
}

impl<'i, R: RuleType> error::Error for Error<R> {
    fn description(&self) -> &str {
        match self.variant {
            ErrorVariant::ParsingError { .. } => "parsing error",
            ErrorVariant::CustomError { ref message } => message
        }
    }
}

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

    #[test]
    fn display_parsing_error_mixed() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2, 3],
                negatives: vec![4, 5, 6]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 4, 5, or 6; expected 1, 2, or 3",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_positives() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2],
                negatives: vec![]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = expected 1 or 2",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_negatives() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![],
                negatives: vec![4, 5, 6]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 4, 5, or 6",
            ].join("\n")
        );
    }

    #[test]
    fn display_parsing_error_unknown() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![],
                negatives: vec![]
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unknown parsing error",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_pos() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            pos
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_two_lines() {
        let input = b"ab\ncd\nefgh";
        let start = unsafe { position::new(input, 4) };
        let end = unsafe { position::new(input, 9) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "3 | efgh",
                "  |  ^^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_three_lines() {
        let input = b"ab\ncd\nefgh";
        let start = unsafe { position::new(input, 1) };
        let end = unsafe { position::new(input, 9) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 1:2",
                "  |",
                "1 | ab",
                "  | ...",
                "3 | efgh",
                "  |  ^^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn display_custom_span_two_lines_inverted_cols() {
        let input = b"abcdef\ngh";
        let start = unsafe { position::new(input, 5) };
        let end = unsafe { position::new(input, 8) };
        let error: Error<u32> = Error::new_from_span(
            ErrorVariant::CustomError {
                message: "error: big one".to_owned()
            },
            start.span(&end)
        );

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 1:6",
                "  |",
                "1 | abcdef",
                "2 | gh",
                "  | ^----^",
                "  |",
                "  = error: big one",
            ].join("\n")
        );
    }

    #[test]
    fn mapped_parsing_error() {
        let input = b"ab\ncd\nef";
        let pos = unsafe { position::new(input, 4) };
        let error: Error<u32> = Error::new_from_pos(
            ErrorVariant::ParsingError {
                positives: vec![1, 2, 3],
                negatives: vec![4, 5, 6]
            },
            pos
        ).renamed_rules(|n| format!("{}", n + 1));

        assert_eq!(
            format!("{}", error),
            vec![
                " --> 2:2",
                "  |",
                "2 | cd",
                "  |  ^---",
                "  |",
                "  = unexpected 5, 6, or 7; expected 2, 3, or 4",
            ].join("\n")
        );
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
mod tests {
    use crate::data_provider::OnDiskDataProvider;
    use crate::errors::Result;
    use crate::iterator::{ParallelIterator, PointLocation, PointQuery};
    use crate::octree::{build_octree, Octree};
    use crate::{AttributeData, NumberOfPoints, PointsBatch};
    use cgmath::{EuclideanSpace, Point3, Vector3};
    use collision::{Aabb, Aabb3};
    use tempdir::TempDir;

    const NUM_POINTS: usize = 100_001;

    impl<T> NumberOfPoints for std::vec::IntoIter<T> {
        fn num_points(&self) -> usize {
            self.len()
        }
    }

    fn build_big_test_octree() -> Octree {
        let mut batch = PointsBatch {
            //ECEF parking lot porter dr
            position: vec![Vector3::new(-2_699_182.0, -4_294_938.0, 3_853_373.0); 2 * NUM_POINTS],
            attributes: vec![(
                "color".to_string(),
                AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); 2 * NUM_POINTS]),
            )]
            .into_iter()
            .collect(),
        };

        batch.position[3] = Vector3::new(-2_702_846.0, -4_291_151.0, 3_855_012.0); // ECEF STANFORD

        let p = Point3::new(6_400_000.0, 6_400_000.0, 6_400_000.0);
        let bounding_box = Aabb3::new(-1.0 * p, p);

        let pool = scoped_pool::Pool::new(10);
        let tmp_dir = TempDir::new("octree").unwrap();

        build_octree(&pool, &tmp_dir, 1.0, bounding_box, vec![batch].into_iter());
        Octree::from_data_provider(Box::new(OnDiskDataProvider {
            directory: tmp_dir.into_path(),
        }))
        .unwrap()
    }

    fn build_test_octree() -> Octree {
        let mut batch = PointsBatch {
            position: vec![Vector3::new(0.0, 0.0, 0.0); NUM_POINTS],
            attributes: vec![(
                "color".to_string(),
                AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); NUM_POINTS]),
            )]
            .into_iter()
            .collect(),
        };

        batch.position[NUM_POINTS - 1] = Vector3::new(-200., -40., 30.);

        let bounding_box = Aabb3::zero().grow(Point3::from_vec(batch.position[100_000]));

        let pool = scoped_pool::Pool::new(10);
        let tmp_dir = TempDir::new("octree").unwrap();

        build_octree(&pool, &tmp_dir, 1.0, bounding_box, vec![batch].into_iter());
        Octree::from_data_provider(Box::new(OnDiskDataProvider {
            directory: tmp_dir.into_path(),
        }))
        .unwrap()
    }

    #[test]
    fn test_batch_iterator() {
        let batch_size = 5000;
        // define function
        let mut callback_count: usize = 0;
        let max_num_points = 13_000; // 2*batch size + 3000
        let mut delivered_points = 0;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len();
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                println!("Callback: Max Points reached!");
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };

        // octree and iterator
        let octree = build_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };
        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator = ParallelIterator::new(
            octree_slice,
            &location,
            batch_size,
            std::cmp::max(1, num_cpus::get() - 1),
            4,
        );

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect_err("Test error");

        //  requiring and returning less points than the total
        assert!(3 <= callback_count);
        assert!(3 * batch_size <= delivered_points);
        assert!(callback_count >= 3);
        assert!((3 + 1) * batch_size >= delivered_points);
        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
    }

    #[test]
    fn test_batch_iterator_more_points() {
        let batch_size = NUM_POINTS / 2;
        // define function

        // requiring more points than in the octree
        let mut callback_count = 0;
        let mut delivered_points = 0;
        let max_num_points = NUM_POINTS + 30_000;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len();
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                println!("Callback: Max Points reached!");
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };
        // octree and iterator
        let octree = build_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };

        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator =
            ParallelIterator::new(octree_slice, &location, batch_size, 2, 2);

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect("Test OK");
        assert!(delivered_points == NUM_POINTS); // only delivers all the points in the octree
    }

    #[test]
    //#[ignore]
    fn test_batch_iterator_big_octree() {
        let batch_size = 5000;
        // define function
        let mut callback_count: u64 = 0;
        let max_num_points = 13_001; // 2*batch size + 3000
        let mut delivered_points: u64 = 0;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len() as u64;
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };

        // octree and iterator
        let octree = build_big_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };
        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator = ParallelIterator::new(
            octree_slice,
            &location,
            batch_size,
            std::cmp::max(1, num_cpus::get() - 1),
            4,
        );

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect_err("Test error");

        assert!(callback_count >= 3);
        assert!((3 + 1) * batch_size as u64 >= delivered_points);
        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
        assert!(delivered_points as i32 - max_num_points as i32 <= batch_size as i32);
    }
}
<commit_msg>Fix num points for IntoIter (#366)<commit_after>#[cfg(test)]
mod tests {
    use crate::data_provider::OnDiskDataProvider;
    use crate::errors::Result;
    use crate::iterator::{ParallelIterator, PointLocation, PointQuery};
    use crate::octree::{build_octree, Octree};
    use crate::{AttributeData, NumberOfPoints, PointsBatch};
    use cgmath::{EuclideanSpace, Point3, Vector3};
    use collision::{Aabb, Aabb3};
    use tempdir::TempDir;

    const NUM_POINTS: usize = 100_001;

    impl NumberOfPoints for std::vec::IntoIter<PointsBatch> {
        fn num_points(&self) -> usize {
            self.clone().map(|b| b.position.len()).sum()
        }
    }

    fn build_big_test_octree() -> Octree {
        let mut batch = PointsBatch {
            //ECEF parking lot porter dr
            position: vec![Vector3::new(-2_699_182.0, -4_294_938.0, 3_853_373.0); 2 * NUM_POINTS],
            attributes: vec![(
                "color".to_string(),
                AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); 2 * NUM_POINTS]),
            )]
            .into_iter()
            .collect(),
        };

        batch.position[3] = Vector3::new(-2_702_846.0, -4_291_151.0, 3_855_012.0); // ECEF STANFORD

        let p = Point3::new(6_400_000.0, 6_400_000.0, 6_400_000.0);
        let bounding_box = Aabb3::new(-1.0 * p, p);

        let pool = scoped_pool::Pool::new(10);
        let tmp_dir = TempDir::new("octree").unwrap();

        build_octree(&pool, &tmp_dir, 1.0, bounding_box, vec![batch].into_iter());
        Octree::from_data_provider(Box::new(OnDiskDataProvider {
            directory: tmp_dir.into_path(),
        }))
        .unwrap()
    }

    fn build_test_octree() -> Octree {
        let mut batch = PointsBatch {
            position: vec![Vector3::new(0.0, 0.0, 0.0); NUM_POINTS],
            attributes: vec![(
                "color".to_string(),
                AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); NUM_POINTS]),
            )]
            .into_iter()
            .collect(),
        };

        batch.position[NUM_POINTS - 1] = Vector3::new(-200., -40., 30.);

        let bounding_box = Aabb3::zero().grow(Point3::from_vec(batch.position[100_000]));

        let pool = scoped_pool::Pool::new(10);
        let tmp_dir = TempDir::new("octree").unwrap();

        build_octree(&pool, &tmp_dir, 1.0, bounding_box, vec![batch].into_iter());
        Octree::from_data_provider(Box::new(OnDiskDataProvider {
            directory: tmp_dir.into_path(),
        }))
        .unwrap()
    }

    #[test]
    fn test_batch_iterator() {
        let batch_size = 5000;
        // define function
        let mut callback_count: usize = 0;
        let max_num_points = 13_000; // 2*batch size + 3000
        let mut delivered_points = 0;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len();
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                println!("Callback: Max Points reached!");
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };

        // octree and iterator
        let octree = build_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };
        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator = ParallelIterator::new(
            octree_slice,
            &location,
            batch_size,
            std::cmp::max(1, num_cpus::get() - 1),
            4,
        );

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect_err("Test error");

        //  requiring and returning less points than the total
        assert!(3 <= callback_count);
        assert!(3 * batch_size <= delivered_points);
        assert!(callback_count >= 3);
        assert!((3 + 1) * batch_size >= delivered_points);
        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
    }

    #[test]
    fn test_batch_iterator_more_points() {
        let batch_size = NUM_POINTS / 2;
        // define function

        // requiring more points than in the octree
        let mut callback_count = 0;
        let mut delivered_points = 0;
        let max_num_points = NUM_POINTS + 30_000;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len();
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                println!("Callback: Max Points reached!");
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };
        // octree and iterator
        let octree = build_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };

        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator =
            ParallelIterator::new(octree_slice, &location, batch_size, 2, 2);

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect("Test OK");
        assert!(delivered_points == NUM_POINTS); // only delivers all the points in the octree
    }

    #[test]
    //#[ignore]
    fn test_batch_iterator_big_octree() {
        let batch_size = 5000;
        // define function
        let mut callback_count: u64 = 0;
        let max_num_points = 13_001; // 2*batch size + 3000
        let mut delivered_points: u64 = 0;
        println!(
            "batch_size= {} ,  num_points= {}",
            batch_size, max_num_points
        );
        let callback_func = |points_batch: PointsBatch| -> Result<()> {
            callback_count += 1;
            delivered_points += points_batch.position.len() as u64;
            println!(
                "Callback_count {:}, delivered points {:}",
                callback_count, delivered_points
            );
            if delivered_points >= max_num_points {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::Interrupted,
                    format!("Maximum number of {} points reached.", max_num_points),
                )
                .into());
            }
            Ok(())
        };

        // octree and iterator
        let octree = build_big_test_octree();
        let location = PointQuery {
            location: PointLocation::AllPoints(),
            global_from_local: None,
        };
        let octree_slice: &[Octree] = std::slice::from_ref(&octree);
        let mut parallel_iterator = ParallelIterator::new(
            octree_slice,
            &location,
            batch_size,
            std::cmp::max(1, num_cpus::get() - 1),
            4,
        );

        let _err_stop = parallel_iterator
            .try_for_each_batch(callback_func)
            .expect_err("Test error");

        assert!(callback_count >= 3);
        assert!((3 + 1) * batch_size as u64 >= delivered_points);
        assert!(delivered_points as i32 - max_num_points as i32 >= 0);
        assert!(delivered_points as i32 - max_num_points as i32 <= batch_size as i32);
    }
}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate regex;
extern crate futures;
extern crate httpbis;
extern crate tokio_core;
#[macro_use]
extern crate log;
extern crate env_logger;

mod test_misc;

use futures::future;
use futures::future::Future;

use httpbis::Client;
use test_misc::*;

#[test]
fn smoke() {
    env_logger::init().ok();

    let server = ServerTest::new();

    let client: Client =
        Client::new_plain("::1", server.port, Default::default()).expect("client");

    let mut futures = Vec::new();
    for _ in 0..10 {
        futures.push(client.start_get("/blocks/200000/5", "localhost").collect());
    }

    let r = future::join_all(futures).wait().expect("wait");
    for rr in r {
        assert_eq!(200000 * 5, rr.body.len());
    }
}
<commit_msg>Two more useless smoke tests<commit_after>extern crate bytes;
extern crate regex;
extern crate futures;
extern crate httpbis;
extern crate tokio_core;
#[macro_use]
extern crate log;
extern crate env_logger;

mod test_misc;

use futures::future;
use futures::future::Future;
use futures::stream::Stream;

use httpbis::Client;
use test_misc::*;

#[test]
fn smoke() {
    env_logger::init().ok();

    let server = ServerTest::new();

    let client: Client =
        Client::new_plain("::1", server.port, Default::default()).expect("client");

    let mut futures = Vec::new();
    for _ in 0..10 {
        futures.push(client.start_get("/blocks/200000/5", "localhost").collect());
    }

    let r = future::join_all(futures).wait().expect("wait");
    for rr in r {
        assert_eq!(200000 * 5, rr.body.len());
    }
}

#[test]
fn parallel_large() {
    env_logger::init().ok();

    let server = ServerTest::new();

    let client: Client =
        Client::new_plain("::1", server.port, Default::default()).expect("client");

    let mut futures = Vec::new();
    for _ in 0..50 {
        futures.push(client.start_get("/blocks/100000/5", "localhost").collect());
    }

    let r = future::join_all(futures).wait().expect("wait");
    for rr in r {
        assert_eq!(100000 * 5, rr.body.len());
    }
}

#[test]
fn seq_long() {
    env_logger::init().ok();

    let server = ServerTest::new();

    let client: Client =
        Client::new_plain("::1", server.port, Default::default()).expect("client");

    let (headers, parts) = client.start_get("/blocks/100000/100", "localhost").0
        .wait().expect("get");

    assert_eq!(200, headers.status());

    let mut sum_len = 0;
    for b in parts.filter_data().wait() {
        sum_len += b.unwrap().len();
    }

    assert_eq!(100000 * 100, sum_len);
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs::{self, File};
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::process;

use regex::Regex;

use super::{Error, Result};

const CONTROLLER_CPU: &str = "cpu";

const CONTROLLER_CPUSET: &str = "cpuset";
const CPUSET_CPUS: &str = "cpuset.cpus";
const CPUSET_MEMS: &str = "cpuset.mems";

const CONTROLLER_PIDS: &str = "pids";

// The list of cgroup controllers we're interested in.
const CONTROLLERS: [&'static str; 3] = [CONTROLLER_CPU, CONTROLLER_CPUSET, CONTROLLER_PIDS];
const PROC_MOUNTS: &str = "/proc/mounts";

pub struct Cgroup {
    tasks_files: Vec<PathBuf>,
}

// It's called writeln_special because we have to use this rather convoluted way of writing
// to special cgroup files, to avoid getting errors. It would be nice to know why that happens :-s
fn writeln_special<T, V>(file_path: &T, value: V) -> Result<()>
where
    T: AsRef<Path>,
    V: ::std::fmt::Display,
{
    fs::write(file_path, format!("{}\n", value))
        .map_err(|e| Error::Write(PathBuf::from(file_path.as_ref()), e))
}

fn readln_special<T: AsRef<Path>>(file_path: &T) -> Result<String> {
    let mut line = fs::read_to_string(file_path)
        .map_err(|e| Error::ReadToString(PathBuf::from(file_path.as_ref()), e))?;

    // Remove the newline character at the end (if any).
    line.pop();

    Ok(line)
}

// If we call inherit_from_parent_aux(.../A/B/C, file, condition), the following will happen:
// 1) If .../A/B/C/file does not exist, or if .../A/B/file does not exist, return an error.
// 2) If .../A/B/file is not empty, write the first line of .../A/B/file into .../A/B/C/file
// and return.
// 3) If ../A/B/file exists but it is empty, call inherit_from_parent_aux(.../A/B, file, false).
// 4) If .../A/B/file is no longer empty, write the first line of .../A/B/file into
// .../A/B/C/file, and return.
// 5) Otherwise, return an error.

// How is this helpful? When creating cgroup folders for the jailer Firecracker instance, the jailer
// will create a hierarchy that looks like <cgroup_base>/firecracker/<id>. Depending on each
// particular cgroup controller, <cgroup_base> contains a number of configuration files. These are
// not actually present on a disk; they are special files exposed by the controller, and they
// usually contain a single line with some configuration value(s). When the "firecracker" and <id>
// subfolders are created, configuration files with the same name appear automatically in the new
// folders, but their contents are not always automatically populated. Moreover,
// if <cgroup_base>/firecracker/some_file is empty, then we cannot have a non-empty file with
// at <cgroup_base>/firecracker/<id>/some_file. The inherit_from_parent function (which is based
// on the following helper function) helps with propagating the values.

// There is also a potential race condition mentioned below. Here is what it refers to: let's say we
// start multiple jailer processes, and one of them calls
// inherit_from_parent_aux(/A/firecracker/id1, file, true), and hits case number 3) from the list
// above, thus recursively calling inherit_from_parent_aux(/A/firecracker, file, false). It's
// entirely possible there was another process in the exact same situations, and that process
// gets to write something to /A/firecracker/file first. In this case, the recursive call made by
// the first process to inherit_from_parent_aux(/A/firecracker, file, false) may fail when writing
// to /A/firecracker/file, but we can still continue, because step 4) only cares about the file
// no longer being empty, regardless of who actually got to populated its contents.

fn inherit_from_parent_aux(
    path: &mut PathBuf,
    file_name: &str,
    retry_one_level_up: bool,
) -> Result<()> {
    // The function with_file_name() replaces the last component of a path with the given name.
    let parent_file = path.with_file_name(file_name);

    let mut line = readln_special(&parent_file)?;
    if line.is_empty() {
        if retry_one_level_up {
            // We have to borrow "parent" from "parent_file" as opposed to "path", because then
            // we wouldn't be able to mutably borrow path at the end of this function (at least not
            // according to how the Rust borrow checker operates right now :-s)
            let parent = parent_file
                .parent()
                .ok_or_else(|| Error::MissingParent(parent_file.clone()))?;

            // Trying to avoid the race condition described above. We don't care about the result,
            // because we check once more if line.is_empty() after the end of this block.
            let _ = inherit_from_parent_aux(&mut parent.to_path_buf(), file_name, false);
            line = readln_special(&parent_file)?;
        }

        if line.is_empty() {
            return Err(Error::CgroupInheritFromParent(
                path.to_path_buf(),
                file_name.to_string(),
            ));
        }
    }

    path.push(file_name);
    writeln_special(&path, &line)?;
    path.pop();

    Ok(())
}

// The path reference is &mut here because we do a push to get the destination file name. However,
// a pop follows shortly after (see fn inherit_from_parent_aux), reverting to the original value.
fn inherit_from_parent(path: &mut PathBuf, file_name: &str) -> Result<()> {
    inherit_from_parent_aux(path, file_name, true)
}

impl Cgroup {
    pub fn new(id: &str, numa_node: u32, exec_file_name: &OsStr) -> Result<Self> {
        let f =
            File::open(PROC_MOUNTS).map_err(|e| Error::FileOpen(PathBuf::from(PROC_MOUNTS), e))?;

        let mut found_controllers: HashMap<&'static str, PathBuf> =
            HashMap::with_capacity(CONTROLLERS.len());

        // Regex courtesy of Filippo.
        let re = Regex::new(
            r"^cgroup[[:space:]](?P<dir>.*)[[:space:]]cgroup[[:space:]](?P<options>.*)[[:space:]]0[[:space:]]0$",
        ).map_err(Error::RegEx)?;

        for l in BufReader::new(f).lines() {
            let l = l.map_err(|e| Error::ReadLine(PathBuf::from(PROC_MOUNTS), e))?;
            if let Some(capture) = re.captures(&l) {
                // We could do the search in a more efficient manner but eh.
                let v: Vec<&str> = capture["options"].split(',').collect();

                for controller in CONTROLLERS.into_iter() {
                    if v.contains(controller) {
                        if let Some(_) =
                            found_controllers.insert(controller, PathBuf::from(&capture["dir"]))
                        {
                            return Err(Error::CgroupLineNotUnique(
                                PROC_MOUNTS.to_string(),
                                controller.to_string(),
                            ));
                        }
                    }
                }
            }
        }

        let keys_len = found_controllers.keys().len();

        if keys_len < CONTROLLERS.len() {
            // We return an error about the first one we didn't find.
            for controller in CONTROLLERS.into_iter() {
                if !found_controllers.contains_key(controller) {
                    return Err(Error::CgroupLineNotFound(
                        PROC_MOUNTS.to_string(),
                        controller.to_string(),
                    ));
                }
            }
        }

        // This is just a sanity check.
        assert_eq!(keys_len, CONTROLLERS.len());

        // We now both create the cgroup subfolders, and fill the tasks_files vector.
        let mut tasks_files = Vec::with_capacity(keys_len);

        for (controller, mut path_buf) in found_controllers.drain() {
            path_buf.push(exec_file_name);
            path_buf.push(id);

            fs::create_dir_all(&path_buf).map_err(|e| Error::CreateDir(path_buf.clone(), e))?;

            // For now, the jailer is only populating configuration values for the cpuset
            // controller, related to the cpu cores we are allowed to run on, and the numa node we
            // want to restrict to. The jailer only creates the folder hierarchy for other cgroups,
            // and the customer has to provide any desired configuration explicitly (if any).

            if controller == CONTROLLER_CPUSET {
                inherit_from_parent(&mut path_buf, CPUSET_CPUS)?;

                // TODO: this does make an unnecessary write, as we change the value of the
                // "cpuset.mems" file again at the end of the for block. Maybe fix this sometime.
                inherit_from_parent(&mut path_buf, CPUSET_MEMS)?;

                // Enforce NUMA node restriction.
                path_buf.push(CPUSET_MEMS);
                writeln_special(&path_buf, numa_node)?;
                path_buf.pop();
            }

            // And now add "tasks" to get the path of the corresponding tasks file.
            path_buf.push("tasks");
            if !tasks_files.contains(&path_buf) {
                tasks_files.push(path_buf);
            }
        }

        Ok(Cgroup { tasks_files })
    }

    // This writes the pid of the current process to each tasks file. These are special files that,
    // when written to, will assign the process associated with the pid to the respective cgroup.
    pub fn attach_pid(&self) -> Result<()> {
        let pid = process::id();
        for tasks_file in &self.tasks_files {
            writeln_special(tasks_file, pid)?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::{tempdir, tempdir_in, NamedTempFile};
    use super::*;
    use std::io::Write;
    use std::path::PathBuf;

    #[test]
    fn test_inherit_from_parent() {
        // 1. If parent file does not exist, return an error.

        // This is /A/B/ .
        let dir = tempdir().expect("Cannot create temporary directory.");
        // This is /A/B/C .
        let dir2 = tempdir_in(dir.path()).expect("Cannot create temporary directory.");
        let path2 = PathBuf::from(dir2.path());
        let result = inherit_from_parent(&mut PathBuf::from(&path2), "inexistent");
        assert!(result.is_err());
        assert!(format!("{:?}", result).contains("ReadToString"));

        // 2. If parent file exists and is empty, will go one level up, and return error because
        // the grandparent file does not exist.
        let mut named_file = NamedTempFile::new_in(dir.path()).expect("Cannot create named file.");
        let result = inherit_from_parent(
            &mut PathBuf::from(path2.clone()),
            named_file.path().file_name().unwrap().to_str().unwrap(),
        );
        assert!(result.is_err());
        assert!(format!("{:?}", result).contains("CgroupInheritFromParent"));

        let child_file = dir2
            .path()
            .join(named_file.path().file_name().unwrap().to_str().unwrap());

        // 3. If parent file exists and is not empty, will return ok and child file will have its
        // contents.
        let some_line = "Parent line";
        writeln!(named_file, "{}", some_line).expect("Cannot write to file.");
        let result = inherit_from_parent(
            &mut PathBuf::from(path2),
            named_file.path().file_name().unwrap().to_str().unwrap(),
        );
        assert!(result.is_ok());
        let res = readln_special(&child_file).expect("Cannot read from file.");
        assert!(res == some_line);
    }
}
<commit_msg>jailer: enforce correct numa cpu assignment...<commit_after>use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs::{self, File};
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
use std::process;

use regex::Regex;

use super::{Error, Result};

const CONTROLLER_CPU: &str = "cpu";

const CONTROLLER_CPUSET: &str = "cpuset";
const CPUSET_CPUS: &str = "cpuset.cpus";
const CPUSET_MEMS: &str = "cpuset.mems";

const CONTROLLER_PIDS: &str = "pids";

// The list of cgroup controllers we're interested in.
const CONTROLLERS: [&'static str; 3] = [CONTROLLER_CPU, CONTROLLER_CPUSET, CONTROLLER_PIDS];
const PROC_MOUNTS: &str = "/proc/mounts";
const NODE_TO_CPULIST: &str = "/sys/devices/system/node/node";

pub struct Cgroup {
    tasks_files: Vec<PathBuf>,
}

// It's called writeln_special because we have to use this rather convoluted way of writing
// to special cgroup files, to avoid getting errors. It would be nice to know why that happens :-s
fn writeln_special<T, V>(file_path: &T, value: V) -> Result<()>
where
    T: AsRef<Path>,
    V: ::std::fmt::Display,
{
    fs::write(file_path, format!("{}\n", value))
        .map_err(|e| Error::Write(PathBuf::from(file_path.as_ref()), e))
}

fn readln_special<T: AsRef<Path>>(file_path: &T) -> Result<String> {
    let mut line = fs::read_to_string(file_path)
        .map_err(|e| Error::ReadToString(PathBuf::from(file_path.as_ref()), e))?;

    // Remove the newline character at the end (if any).
    line.pop();

    Ok(line)
}

// If we call inherit_from_parent_aux(.../A/B/C, file, condition), the following will happen:
// 1) If .../A/B/C/file does not exist, or if .../A/B/file does not exist, return an error.
// 2) If .../A/B/file is not empty, write the first line of .../A/B/file into .../A/B/C/file
// and return.
// 3) If ../A/B/file exists but it is empty, call inherit_from_parent_aux(.../A/B, file, false).
// 4) If .../A/B/file is no longer empty, write the first line of .../A/B/file into
// .../A/B/C/file, and return.
// 5) Otherwise, return an error.

// How is this helpful? When creating cgroup folders for the jailer Firecracker instance, the jailer
// will create a hierarchy that looks like <cgroup_base>/firecracker/<id>. Depending on each
// particular cgroup controller, <cgroup_base> contains a number of configuration files. These are
// not actually present on a disk; they are special files exposed by the controller, and they
// usually contain a single line with some configuration value(s). When the "firecracker" and <id>
// subfolders are created, configuration files with the same name appear automatically in the new
// folders, but their contents are not always automatically populated. Moreover,
// if <cgroup_base>/firecracker/some_file is empty, then we cannot have a non-empty file with
// at <cgroup_base>/firecracker/<id>/some_file. The inherit_from_parent function (which is based
// on the following helper function) helps with propagating the values.

// There is also a potential race condition mentioned below. Here is what it refers to: let's say we
// start multiple jailer processes, and one of them calls
// inherit_from_parent_aux(/A/firecracker/id1, file, true), and hits case number 3) from the list
// above, thus recursively calling inherit_from_parent_aux(/A/firecracker, file, false). It's
// entirely possible there was another process in the exact same situations, and that process
// gets to write something to /A/firecracker/file first. In this case, the recursive call made by
// the first process to inherit_from_parent_aux(/A/firecracker, file, false) may fail when writing
// to /A/firecracker/file, but we can still continue, because step 4) only cares about the file
// no longer being empty, regardless of who actually got to populated its contents.

fn inherit_from_parent_aux(
    path: &mut PathBuf,
    file_name: &str,
    retry_one_level_up: bool,
) -> Result<()> {
    // The function with_file_name() replaces the last component of a path with the given name.
    let parent_file = path.with_file_name(file_name);

    let mut line = readln_special(&parent_file)?;
    if line.is_empty() {
        if retry_one_level_up {
            // We have to borrow "parent" from "parent_file" as opposed to "path", because then
            // we wouldn't be able to mutably borrow path at the end of this function (at least not
            // according to how the Rust borrow checker operates right now :-s)
            let parent = parent_file
                .parent()
                .ok_or_else(|| Error::MissingParent(parent_file.clone()))?;

            // Trying to avoid the race condition described above. We don't care about the result,
            // because we check once more if line.is_empty() after the end of this block.
            let _ = inherit_from_parent_aux(&mut parent.to_path_buf(), file_name, false);
            line = readln_special(&parent_file)?;
        }

        if line.is_empty() {
            return Err(Error::CgroupInheritFromParent(
                path.to_path_buf(),
                file_name.to_string(),
            ));
        }
    }

    path.push(file_name);
    writeln_special(&path, &line)?;
    path.pop();

    Ok(())
}

// The path reference is &mut here because we do a push to get the destination file name. However,
// a pop follows shortly after (see fn inherit_from_parent_aux), reverting to the original value.
fn inherit_from_parent(path: &mut PathBuf, file_name: &str) -> Result<()> {
    inherit_from_parent_aux(path, file_name, true)
}

impl Cgroup {
    pub fn new(id: &str, numa_node: u32, exec_file_name: &OsStr) -> Result<Self> {
        let f =
            File::open(PROC_MOUNTS).map_err(|e| Error::FileOpen(PathBuf::from(PROC_MOUNTS), e))?;

        let mut found_controllers: HashMap<&'static str, PathBuf> =
            HashMap::with_capacity(CONTROLLERS.len());

        // Regex courtesy of Filippo.
        let re = Regex::new(
            r"^cgroup[[:space:]](?P<dir>.*)[[:space:]]cgroup[[:space:]](?P<options>.*)[[:space:]]0[[:space:]]0$",
        ).map_err(Error::RegEx)?;
        for l in BufReader::new(f).lines() {
            let l = l.map_err(|e| Error::ReadLine(PathBuf::from(PROC_MOUNTS), e))?;
            if let Some(capture) = re.captures(&l) {
                // We could do the search in a more efficient manner but eh.
                let v: Vec<&str> = capture["options"].split(',').collect();

                for controller in CONTROLLERS.into_iter() {
                    if v.contains(controller) {
                        if let Some(_) =
                            found_controllers.insert(controller, PathBuf::from(&capture["dir"]))
                        {
                            return Err(Error::CgroupLineNotUnique(
                                PROC_MOUNTS.to_string(),
                                controller.to_string(),
                            ));
                        }
                    }
                }
            }
        }

        let keys_len = found_controllers.keys().len();

        if keys_len < CONTROLLERS.len() {
            // We return an error about the first one we didn't find.
            for controller in CONTROLLERS.into_iter() {
                if !found_controllers.contains_key(controller) {
                    return Err(Error::CgroupLineNotFound(
                        PROC_MOUNTS.to_string(),
                        controller.to_string(),
                    ));
                }
            }
        }

        // This is just a sanity check.
        assert_eq!(keys_len, CONTROLLERS.len());

        // We now both create the cgroup subfolders, and fill the tasks_files vector.
        let mut tasks_files = Vec::with_capacity(keys_len);

        for (controller, mut path_buf) in found_controllers.drain() {
            path_buf.push(exec_file_name);
            path_buf.push(id);

            fs::create_dir_all(&path_buf).map_err(|e| Error::CreateDir(path_buf.clone(), e))?;

            // For now, the jailer is only populating configuration values for the cpuset
            // controller, related to the cpu cores we are allowed to run on, and the numa node we
            // want to restrict to. The jailer only creates the folder hierarchy for other cgroups,
            // and the customer has to provide any desired configuration explicitly (if any).

            if controller == CONTROLLER_CPUSET {
                inherit_from_parent(&mut path_buf, CPUSET_CPUS)?;

                // TODO: this does make an unnecessary write, as we change the value of the
                // "cpuset.mems" file again at the end of the for block. Maybe fix this sometime.
                inherit_from_parent(&mut path_buf, CPUSET_MEMS)?;

                // Enforce NUMA node restriction.
                // The cpuset subsystem assigns individual CPUs and memory nodes to cgroups.
                // CPUSET_MEMS specifies the memory nodes that tasks in this cgroup are permitted to
                // access.
                path_buf.push(CPUSET_MEMS);
                writeln_special(&path_buf, numa_node)?;
                path_buf.pop();
                // Similar to how numactl library does, we are copying the contents of
                // /sys/devices/system/node/nodeX/cpulist to the cpuset.cpus file for ensuring
                // correct numa cpu assignment.
                let line = readln_special(&PathBuf::from(format!(
                    "{}{}/cpulist",
                    NODE_TO_CPULIST, numa_node
                )))?;
                path_buf.push(CPUSET_CPUS);
                writeln_special(&path_buf, line)?;
                path_buf.pop();
            }

            // And now add "tasks" to get the path of the corresponding tasks file.
            path_buf.push("tasks");
            if !tasks_files.contains(&path_buf) {
                tasks_files.push(path_buf);
            }
        }

        Ok(Cgroup { tasks_files })
    }

    // This writes the pid of the current process to each tasks file. These are special files that,
    // when written to, will assign the process associated with the pid to the respective cgroup.
    pub fn attach_pid(&self) -> Result<()> {
        let pid = process::id();
        for tasks_file in &self.tasks_files {
            writeln_special(tasks_file, pid)?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    extern crate tempfile;

    use self::tempfile::{tempdir, tempdir_in, NamedTempFile};
    use super::*;
    use std::io::Write;
    use std::path::PathBuf;

    #[test]
    fn test_inherit_from_parent() {
        // 1. If parent file does not exist, return an error.

        // This is /A/B/ .
        let dir = tempdir().expect("Cannot create temporary directory.");
        // This is /A/B/C .
        let dir2 = tempdir_in(dir.path()).expect("Cannot create temporary directory.");
        let path2 = PathBuf::from(dir2.path());
        let result = inherit_from_parent(&mut PathBuf::from(&path2), "inexistent");
        assert!(result.is_err());
        assert!(format!("{:?}", result).contains("ReadToString"));

        // 2. If parent file exists and is empty, will go one level up, and return error because
        // the grandparent file does not exist.
        let mut named_file = NamedTempFile::new_in(dir.path()).expect("Cannot create named file.");
        let result = inherit_from_parent(
            &mut PathBuf::from(path2.clone()),
            named_file.path().file_name().unwrap().to_str().unwrap(),
        );
        assert!(result.is_err());
        assert!(format!("{:?}", result).contains("CgroupInheritFromParent"));

        let child_file = dir2
            .path()
            .join(named_file.path().file_name().unwrap().to_str().unwrap());

        // 3. If parent file exists and is not empty, will return ok and child file will have its
        // contents.
        let some_line = "Parent line";
        writeln!(named_file, "{}", some_line).expect("Cannot write to file.");
        let result = inherit_from_parent(
            &mut PathBuf::from(path2),
            named_file.path().file_name().unwrap().to_str().unwrap(),
        );
        assert!(result.is_ok());
        let res = readln_special(&child_file).expect("Cannot read from file.");
        assert!(res == some_line);
    }
}
<|endoftext|>


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

struct Clam<'self> {
    chowder: &'self int
}

trait get_chowder {
    fn get_chowder() -> &self/int;
}

impl get_chowder for Clam {
    fn get_chowder() -> &self/int { return self.chowder; }
}

pub fn main() {
    let clam = Clam { chowder: &3 };
    log(debug, *clam.get_chowder());
    clam.get_chowder();
}

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

struct Clam<'self> {
    chowder: &'self int
}

trait get_chowder<'self> {
    fn get_chowder() -> &'self int;
}

impl<'self> get_chowder<'self> for Clam<'self> {
    fn get_chowder() -> &'self int { return self.chowder; }
}

pub fn main() {
    let clam = Clam { chowder: &3 };
    log(debug, *clam.get_chowder());
    clam.get_chowder();
}

<|endoftext|>