<commit_before>// Copyright (c) 2017 Jason White
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::io;
use std::io::Write as IoWrite;
use std::fmt;
use std::fmt::Write as FmtWrite;
use std::time::Duration;
use std::path::PathBuf;
use std::process;
use std::collections::BTreeMap;
use std::ffi::OsString;

use tempfile::{NamedTempFile, TempPath};

use node::{Error, Task};
use util::{NeverAlwaysAuto, Counter};
use retry;

/// A task that executes a single command. A command is simply a process to be
/// spawned.
#[derive(Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq, Hash, Clone)]
pub struct Command {
    /// Process and arguments to spawn.
    args: Vec<String>,

    /// Optional working directory to spawn the process in. If `None`, uses the
    /// working directory of the parent process (i.e., the build process).
    cwd: Option<PathBuf>,

    /// Optional environment variables.
    env: Option<BTreeMap<String, String>>,

    /// Response file creation.
    ///
    /// If `Always`, creates a temporary response file with all the command line
    /// arguments and passes that as the first command line argument instead.
    /// This is useful for very long command lines that exceed operating system
    /// limits.
    #[serde(default)]
    response_file: NeverAlwaysAuto,

    /// Redirect standard output to a file instead.
    stdout: Option<PathBuf>,

    /// String to display when executing the task. If `None`, the command
    /// arguments are displayed in full instead.
    display: Option<String>,

    /// How much time to give the command to execute. If `None`, there is no
    /// time limit.
    timeout: Option<Duration>,

    /// Retry settings.
    #[serde(default)]
    retry: retry::Retry,
}

impl Command {
    #[cfg(test)]
    pub fn new(args: Vec<String>) -> Command {
        Command {
            args: args,
            cwd: None,
            env: None,
            stdout: None,
            display: None,
            timeout: None,
            retry: retry::Retry::new(),
        }
    }
}

impl Command {
    // Sets the working directory for the command.
    #[allow(dead_code)]
    pub fn cwd(mut self, path: PathBuf) -> Command {
        self.cwd = Some(path);
        self
    }

    // Sets the stdout file for the command.
    #[allow(dead_code)]
    pub fn stdout(mut self, path: PathBuf) -> Command {
        self.stdout = Some(path);
        self
    }

    // Sets the display string for the command.
    #[allow(dead_code)]
    pub fn display(mut self, display: String) -> Command {
        self.display = Some(display);
        self
    }

    // Sets the timeout for the command.
    #[allow(dead_code)]
    pub fn timeout(mut self, timeout: Duration) -> Command {
        self.timeout = Some(timeout);
        self
    }

    // Sets the retry configuration.
    #[allow(dead_code)]
    pub fn retry(mut self, retry: retry::Retry) -> Command {
        self.retry = retry;
        self
    }

    fn execute_impl(&self, log: &mut io::Write) -> Result<(), Error> {
        //writeln!(log, "Executing `{}`", self)?;

        assert!(self.args.len() > 0);

        // TODO:
        //  1. Spawn the process
        //  2. Capture stdout/stderr appropriately.
        //  4. Add implicit dependency detection framework and refactor all of
        //     the above to make it work.
        //  5. Implement timeouts.

        let mut cmd = process::Command::new(&self.args[0]);
        cmd.stdin(process::Stdio::null());

        // Generate a response file if necessary.
        let generate_response_file = match self.response_file {
            NeverAlwaysAuto::Never => false,
            NeverAlwaysAuto::Always => true,
            NeverAlwaysAuto::Auto => args_too_large(&self.args),
        };

        // The temporary response file needs to outlive the spawned process, so
        // it needs to be bound to a variable even if it is never used.
        let _rsp = if generate_response_file {
            let temp = response_file(&self.args[1..])?;

            let mut arg = OsString::new();
            arg.push("@");
            arg.push(&temp);
            cmd.arg(&arg);

            Some(temp)
        } else {
            cmd.args(&self.args[1..]);
            None
        };

        if let Some(ref cwd) = self.cwd {
            cmd.current_dir(cwd);
        }

        if let Some(ref env) = self.env {
            cmd.envs(env);
        }

        let output = cmd.output()?;

        // TODO: Interleave stdout and stderr.
        log.write(&output.stdout)?;
        log.write(&output.stderr)?;

        if output.status.success() {
            Ok(())
        } else {
            match output.status.code() {
                Some(code) => {
                    Err(io::Error::new(io::ErrorKind::Other,
                                       format!("Process exited with error code {}",
                                               code)))
                }
                None => {
                    Err(io::Error::new(
                            io::ErrorKind::Other, "Process terminated by signal"))
                }
            }
        }
    }
}

impl fmt::Display for Command {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref display) = self.display {
            write!(f, "{}", display)
        } else {
            let mut args = self.args.iter();

            if let Some(arg) = args.next() {
                write!(f, "{}", Arg::new(arg))?;
            }

            for arg in args {
                write!(f, " {}", Arg::new(arg))?;
            }

            Ok(())
        }
    }
}

impl fmt::Debug for Command {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        write!(f, "\"")?;

        let mut args = self.args.iter();

        if let Some(arg) = args.next() {
            write!(f, "{}", Arg::new(arg))?;
        }

        for arg in args {
            write!(f, " {}", Arg::new(arg))?;
        }

        write!(f, "\"")?;

        Ok(())
    }
}

impl Task for Command {
    fn execute(&self, log: &mut io::Write) -> Result<(), Error> {
        self.retry
            .call(|| self.execute_impl(log), retry::progress_dummy)
    }
}

/// Helper type for formatting command line arguments.
struct Arg<'a> {
    arg: &'a str,
}

impl<'a> Arg<'a> {
    pub fn new(arg: &'a str) -> Arg<'a> {
        Arg { arg: arg }
    }

    /// Quotes the argument such that it is safe to pass to the shell.
    #[cfg(windows)]
    pub fn quote(&self, writer: &mut fmt::Write) -> fmt::Result {
        let quote = self.arg.chars().any(|c| c == ' ' || c == '\t') ||
                    self.arg.is_empty();

        if quote {
            writer.write_char('"')?;
        }

        let mut backslashes: usize = 0;

        for x in self.arg.chars() {
            if x == '\\' {
                backslashes += 1;
            } else {
                // Dump backslashes if we hit a quotation mark.
                if x == '"' {
                    // We need 2n+1 backslashes to escape a quote.
                    for _ in 0..(backslashes+1) {
                        writer.write_char('\\')?;
                    }
                }

                backslashes = 0;
            }

            writer.write_char(x)?;
        }

        if quote {
            // Escape any trailing backslashes.
            for _ in 0..backslashes {
                writer.write_char('\\')?;
            }

            writer.write_char('"')?;
        }

        Ok(())
    }
}

impl<'a> fmt::Display for Arg<'a> {
    /// Converts an argument such that it is safe to append to a command line
    /// string.
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.quote(f)
    }
}

/// Writes the response file to a stream.
fn write_response_file<S, I>(args: I, writer: &mut io::Write) -> io::Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let mut iter = args.into_iter();

    // Write UTF-8 BOM. Some tools require this to properly decode it as UTF-8
    // instead of ASCII.
    writer.write(b"\xEF\xBB\xBF")?;

    if let Some(arg) = iter.next() {
        write!(writer, "{}", Arg::new(arg.as_ref())).unwrap();
    }

    for arg in iter {
        write!(writer, " {}", Arg::new(arg.as_ref())).unwrap();
    }

    // Some programs require a trailing new line (notably LIB.exe and LINK.exe).
    writer.write(b"\n")?;

    Ok(())
}

/// Generates a temporary response file for the given command line arguments.
fn response_file<S, I>(args: I) -> io::Result<TempPath>
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let tempfile = NamedTempFile::new()?;

    {
        let mut writer = io::BufWriter::new(&tempfile);
        write_response_file(args, &mut writer)?;

        // Explicitly flush to catch any errors.
        writer.flush()?;
    }

    Ok(tempfile.into_temp_path())
}

/// Checks if the given command line arguments are too large and should instead
/// go into a response file. The entire list of arguments, including the program
/// name should be passed to this function.
#[cfg(windows)]
fn args_too_large<S, I>(args: I) -> bool
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    // The maximum length is 32768 characters, including the NULL terminator.
    let mut counter = Counter::new();

    let mut iter = args.into_iter();

    if let Some(arg) = iter.next() {
        write!(counter, "{}", Arg::new(arg.as_ref())).unwrap();
    }

    for arg in iter {
        write!(counter, " {}", Arg::new(arg.as_ref())).unwrap();
    }

    // Final NULL terminator.
    counter += 1;

    counter.count() > 32768
}

#[cfg(unix)]
fn args_too_large<S, I>(args: I) -> bool
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let mut size: usize = 0;

    for arg in args.into_iter() {
        // +1 for the NULL terminator.
        size += arg.len() + 1;
    }

    // +1 for the final NULL terminator.
    size += 1;

    // Can't be larger than 128 kB.
    size > 0x20000
}

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

    #[test]
    fn test_command_display() {
        assert_eq!(format!("{}", ));
    }
}
<commit_msg>Add test for argument quoting<commit_after>// Copyright (c) 2017 Jason White
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::io;
use std::io::Write as IoWrite;
use std::fmt;
use std::fmt::Write as FmtWrite;
use std::time::Duration;
use std::path::PathBuf;
use std::process;
use std::collections::BTreeMap;
use std::ffi::OsString;

use tempfile::{NamedTempFile, TempPath};

use node::{Error, Task};
use util::{NeverAlwaysAuto, Counter};
use retry;

/// A task that executes a single command. A command is simply a process to be
/// spawned.
#[derive(Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq, Hash, Clone)]
pub struct Command {
    /// Process and arguments to spawn.
    args: Vec<String>,

    /// Optional working directory to spawn the process in. If `None`, uses the
    /// working directory of the parent process (i.e., the build process).
    cwd: Option<PathBuf>,

    /// Optional environment variables.
    env: Option<BTreeMap<String, String>>,

    /// Response file creation.
    ///
    /// If `Always`, creates a temporary response file with all the command line
    /// arguments and passes that as the first command line argument instead.
    /// This is useful for very long command lines that exceed operating system
    /// limits.
    #[serde(default)]
    response_file: NeverAlwaysAuto,

    /// Redirect standard output to a file instead.
    stdout: Option<PathBuf>,

    /// String to display when executing the task. If `None`, the command
    /// arguments are displayed in full instead.
    display: Option<String>,

    /// How much time to give the command to execute. If `None`, there is no
    /// time limit.
    timeout: Option<Duration>,

    /// Retry settings.
    #[serde(default)]
    retry: retry::Retry,
}

impl Command {
    #[cfg(test)]
    pub fn new(args: Vec<String>) -> Command {
        Command {
            args: args,
            cwd: None,
            env: None,
            stdout: None,
            display: None,
            timeout: None,
            retry: retry::Retry::new(),
        }
    }
}

impl Command {
    // Sets the working directory for the command.
    #[allow(dead_code)]
    pub fn cwd(mut self, path: PathBuf) -> Command {
        self.cwd = Some(path);
        self
    }

    // Sets the stdout file for the command.
    #[allow(dead_code)]
    pub fn stdout(mut self, path: PathBuf) -> Command {
        self.stdout = Some(path);
        self
    }

    // Sets the display string for the command.
    #[allow(dead_code)]
    pub fn display(mut self, display: String) -> Command {
        self.display = Some(display);
        self
    }

    // Sets the timeout for the command.
    #[allow(dead_code)]
    pub fn timeout(mut self, timeout: Duration) -> Command {
        self.timeout = Some(timeout);
        self
    }

    // Sets the retry configuration.
    #[allow(dead_code)]
    pub fn retry(mut self, retry: retry::Retry) -> Command {
        self.retry = retry;
        self
    }

    fn execute_impl(&self, log: &mut io::Write) -> Result<(), Error> {
        //writeln!(log, "Executing `{}`", self)?;

        assert!(self.args.len() > 0);

        // TODO:
        //  1. Spawn the process
        //  2. Capture stdout/stderr appropriately.
        //  4. Add implicit dependency detection framework and refactor all of
        //     the above to make it work.
        //  5. Implement timeouts.

        let mut cmd = process::Command::new(&self.args[0]);
        cmd.stdin(process::Stdio::null());

        // Generate a response file if necessary.
        let generate_response_file = match self.response_file {
            NeverAlwaysAuto::Never => false,
            NeverAlwaysAuto::Always => true,
            NeverAlwaysAuto::Auto => args_too_large(&self.args),
        };

        // The temporary response file needs to outlive the spawned process, so
        // it needs to be bound to a variable even if it is never used.
        let _rsp = if generate_response_file {
            let temp = response_file(&self.args[1..])?;

            let mut arg = OsString::new();
            arg.push("@");
            arg.push(&temp);
            cmd.arg(&arg);

            Some(temp)
        } else {
            cmd.args(&self.args[1..]);
            None
        };

        if let Some(ref cwd) = self.cwd {
            cmd.current_dir(cwd);
        }

        if let Some(ref env) = self.env {
            cmd.envs(env);
        }

        let output = cmd.output()?;

        // TODO: Interleave stdout and stderr.
        log.write(&output.stdout)?;
        log.write(&output.stderr)?;

        if output.status.success() {
            Ok(())
        } else {
            match output.status.code() {
                Some(code) => {
                    Err(io::Error::new(io::ErrorKind::Other,
                                       format!("Process exited with error code {}",
                                               code)))
                }
                None => {
                    Err(io::Error::new(
                            io::ErrorKind::Other, "Process terminated by signal"))
                }
            }
        }
    }
}

impl fmt::Display for Command {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(ref display) = self.display {
            write!(f, "{}", display)
        } else {
            let mut args = self.args.iter();

            if let Some(arg) = args.next() {
                write!(f, "{}", Arg::new(arg))?;
            }

            for arg in args {
                write!(f, " {}", Arg::new(arg))?;
            }

            Ok(())
        }
    }
}

impl fmt::Debug for Command {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        write!(f, "\"")?;

        let mut args = self.args.iter();

        if let Some(arg) = args.next() {
            write!(f, "{}", Arg::new(arg))?;
        }

        for arg in args {
            write!(f, " {}", Arg::new(arg))?;
        }

        write!(f, "\"")?;

        Ok(())
    }
}

impl Task for Command {
    fn execute(&self, log: &mut io::Write) -> Result<(), Error> {
        self.retry
            .call(|| self.execute_impl(log), retry::progress_dummy)
    }
}

/// Helper type for formatting command line arguments.
struct Arg<'a> {
    arg: &'a str,
}

impl<'a> Arg<'a> {
    pub fn new(arg: &'a str) -> Arg<'a> {
        Arg { arg: arg }
    }

    /// Quotes the argument such that it is safe to pass to the shell.
    #[cfg(windows)]
    pub fn quote(&self, writer: &mut fmt::Write) -> fmt::Result {
        let quote = self.arg.chars().any(|c| c == ' ' || c == '\t') ||
                    self.arg.is_empty();

        if quote {
            writer.write_char('"')?;
        }

        let mut backslashes: usize = 0;

        for x in self.arg.chars() {
            if x == '\\' {
                backslashes += 1;
            } else {
                // Dump backslashes if we hit a quotation mark.
                if x == '"' {
                    // We need 2n+1 backslashes to escape a quote.
                    for _ in 0..(backslashes+1) {
                        writer.write_char('\\')?;
                    }
                }

                backslashes = 0;
            }

            writer.write_char(x)?;
        }

        if quote {
            // Escape any trailing backslashes.
            for _ in 0..backslashes {
                writer.write_char('\\')?;
            }

            writer.write_char('"')?;
        }

        Ok(())
    }
}

impl<'a> fmt::Display for Arg<'a> {
    /// Converts an argument such that it is safe to append to a command line
    /// string.
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.quote(f)
    }
}

/// Writes the response file to a stream.
fn write_response_file<S, I>(args: I, writer: &mut io::Write) -> io::Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let mut iter = args.into_iter();

    // Write UTF-8 BOM. Some tools require this to properly decode it as UTF-8
    // instead of ASCII.
    writer.write(b"\xEF\xBB\xBF")?;

    if let Some(arg) = iter.next() {
        write!(writer, "{}", Arg::new(arg.as_ref())).unwrap();
    }

    for arg in iter {
        write!(writer, " {}", Arg::new(arg.as_ref())).unwrap();
    }

    // Some programs require a trailing new line (notably LIB.exe and LINK.exe).
    writer.write(b"\n")?;

    Ok(())
}

/// Generates a temporary response file for the given command line arguments.
fn response_file<S, I>(args: I) -> io::Result<TempPath>
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let tempfile = NamedTempFile::new()?;

    {
        let mut writer = io::BufWriter::new(&tempfile);
        write_response_file(args, &mut writer)?;

        // Explicitly flush to catch any errors.
        writer.flush()?;
    }

    Ok(tempfile.into_temp_path())
}

/// Checks if the given command line arguments are too large and should instead
/// go into a response file. The entire list of arguments, including the program
/// name should be passed to this function.
#[cfg(windows)]
fn args_too_large<S, I>(args: I) -> bool
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    // The maximum length is 32768 characters, including the NULL terminator.
    let mut counter = Counter::new();

    let mut iter = args.into_iter();

    if let Some(arg) = iter.next() {
        write!(counter, "{}", Arg::new(arg.as_ref())).unwrap();
    }

    for arg in iter {
        write!(counter, " {}", Arg::new(arg.as_ref())).unwrap();
    }

    // Final NULL terminator.
    counter += 1;

    counter.count() > 32768
}

#[cfg(unix)]
fn args_too_large<S, I>(args: I) -> bool
    where I: IntoIterator<Item = S>,
          S: AsRef<str>
{
    let mut size: usize = 0;

    for arg in args.into_iter() {
        // +1 for the NULL terminator.
        size += arg.len() + 1;
    }

    // +1 for the final NULL terminator.
    size += 1;

    // Can't be larger than 128 kB.
    size > 0x20000
}

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

    #[test]
    #[cfg(windows)]
    fn test_arg_display() {
        assert_eq!(format!("{}", Arg::new("foo bar")), "\"foo bar\"");
        assert_eq!(format!("{}", Arg::new("foo\tbar")), "\"foo\tbar\"");
        assert_eq!(format!("{}", Arg::new("foobar")), "foobar");
        assert_eq!(format!("{}", Arg::new("\"foo bar\"")), "\"\\\"foo bar\\\"\"");
        assert_eq!(format!("{}", Arg::new(r"C:\foo\bar")), r"C:\foo\bar");
        assert_eq!(format!("{}", Arg::new(r"\\foo\bar")), r"\\foo\bar");
    }
}
<|endoftext|>


<commit_before>use mime::Mime;

header! {
    #[doc="`Content-Type` header, defined in"]
    #[doc="[RFC7231](http://tools.ietf.org/html/rfc7231#section-3.1.1.5)"]
    #[doc=""]
    #[doc="The `Content-Type` header field indicates the media type of the"]
    #[doc="associated representation: either the representation enclosed in the"]
    #[doc="message payload or the selected representation, as determined by the"]
    #[doc="message semantics.  The indicated media type defines both the data"]
    #[doc="format and how that data is intended to be processed by a recipient,"]
    #[doc="within the scope of the received message semantics, after any content"]
    #[doc="codings indicated by Content-Encoding are decoded."]
    #[doc=""]
    #[doc="# ABNF"]
    #[doc="```plain"]
    #[doc="Content-Type = media-type"]
    #[doc="```"]
    #[doc=""]
    #[doc="# Example values"]
    #[doc="* `text/html; charset=ISO-8859-4`"]
    (ContentType, "Content-Type") => [Mime]

    test_content_type {
        test_header!(
            test1,
            // FIXME: Should be b"text/html; charset=ISO-8859-4" but mime crate lowercases
            // the whole value so parsing and formatting the value gives a different result
            vec![b"text/html; charset=iso-8859-4"],
            Some(HeaderField(Mime(
                TopLevel::Text,
                SubLevel::Html,
                vec![(Attr::Charset, Value::Ext("iso-8859-4".to_owned()))]))));
    }
}

impl ContentType {
    /// A constructor  to easily create a `Content-Type: application/json; charset=utf-8` header.
    #[inline]
    pub fn json() -> ContentType {
        ContentType(mime!(Application/Json; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/plain; charset=utf-8` header.
    #[inline]
    pub fn plaintext() -> ContentType {
        ContentType(mime!(Text/Plain; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/html; charset=utf-8` header.
    #[inline]
    pub fn html() -> ContentType {
        ContentType(mime!(Text/Html; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: image/jpeg` header.
    #[inline]
    pub fn jpeg() -> ContentType {
        ContentType(mime!(Image/Jpeg))
    }

    /// A constructor  to easily create a `Content-Type: image/png` header.
    #[inline]
    pub fn png() -> ContentType {
        ContentType(mime!(Image/Png))
    }
}
bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });
<commit_msg>feat(header): add ContentType::form_url_encoded() constructor<commit_after>use mime::Mime;

header! {
    #[doc="`Content-Type` header, defined in"]
    #[doc="[RFC7231](http://tools.ietf.org/html/rfc7231#section-3.1.1.5)"]
    #[doc=""]
    #[doc="The `Content-Type` header field indicates the media type of the"]
    #[doc="associated representation: either the representation enclosed in the"]
    #[doc="message payload or the selected representation, as determined by the"]
    #[doc="message semantics.  The indicated media type defines both the data"]
    #[doc="format and how that data is intended to be processed by a recipient,"]
    #[doc="within the scope of the received message semantics, after any content"]
    #[doc="codings indicated by Content-Encoding are decoded."]
    #[doc=""]
    #[doc="# ABNF"]
    #[doc="```plain"]
    #[doc="Content-Type = media-type"]
    #[doc="```"]
    #[doc=""]
    #[doc="# Example values"]
    #[doc="* `text/html; charset=ISO-8859-4`"]
    (ContentType, "Content-Type") => [Mime]

    test_content_type {
        test_header!(
            test1,
            // FIXME: Should be b"text/html; charset=ISO-8859-4" but mime crate lowercases
            // the whole value so parsing and formatting the value gives a different result
            vec![b"text/html; charset=iso-8859-4"],
            Some(HeaderField(Mime(
                TopLevel::Text,
                SubLevel::Html,
                vec![(Attr::Charset, Value::Ext("iso-8859-4".to_owned()))]))));
    }
}

impl ContentType {
    /// A constructor  to easily create a `Content-Type: application/json; charset=utf-8` header.
    #[inline]
    pub fn json() -> ContentType {
        ContentType(mime!(Application/Json; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/plain; charset=utf-8` header.
    #[inline]
    pub fn plaintext() -> ContentType {
        ContentType(mime!(Text/Plain; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: text/html; charset=utf-8` header.
    #[inline]
    pub fn html() -> ContentType {
        ContentType(mime!(Text/Html; Charset=Utf8))
    }

    /// A constructor  to easily create a `Content-Type: application/www-form-url-encoded` header.
    #[inline]
    pub fn form_url_encoded() -> ContentType {
        ContentType(mime!(Application/WwwFormUrlEncoded))
    }
    /// A constructor  to easily create a `Content-Type: image/jpeg` header.
    #[inline]
    pub fn jpeg() -> ContentType {
        ContentType(mime!(Image/Jpeg))
    }

    /// A constructor  to easily create a `Content-Type: image/png` header.
    #[inline]
    pub fn png() -> ContentType {
        ContentType(mime!(Image/Png))
    }
}

bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });
<|endoftext|>


<commit_before>use header::{Header, HeaderFormat};
use std::fmt::{self, Show};
use header::shared::util::from_one_raw_str;
use mime::Mime;

/// The `Content-Type` header.
///
/// Used to describe the MIME type of message body. Can be used with both
/// requests and responses.
#[derive(Clone, PartialEq, Show)]
pub struct ContentType(pub Mime);

deref!(ContentType => Mime);

impl Header for ContentType {
    fn header_name(_: Option<ContentType>) -> &'static str {
        "Content-Type"
    }

    fn parse_header(raw: &[Vec<u8>]) -> Option<ContentType> {
        from_one_raw_str(raw).map(|mime| ContentType(mime))
    }
}

impl HeaderFormat for ContentType {
    fn fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let ContentType(ref value) = *self;
        value.fmt(fmt)
    }
}

bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });

<commit_msg>fixes invalid format of mime types<commit_after>use header::{Header, HeaderFormat};
use std::fmt::{self, String};
use header::shared::util::from_one_raw_str;
use mime::Mime;

/// The `Content-Type` header.
///
/// Used to describe the MIME type of message body. Can be used with both
/// requests and responses.
#[derive(Clone, PartialEq, Show)]
pub struct ContentType(pub Mime);

deref!(ContentType => Mime);

impl Header for ContentType {
    fn header_name(_: Option<ContentType>) -> &'static str {
        "Content-Type"
    }

    fn parse_header(raw: &[Vec<u8>]) -> Option<ContentType> {
        from_one_raw_str(raw).map(|mime| ContentType(mime))
    }
}

impl HeaderFormat for ContentType {
    fn fmt_header(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let ContentType(ref value) = *self;
        value.fmt(fmt)
    }
}

bench_header!(bench, ContentType, { vec![b"application/json; charset=utf-8".to_vec()] });

<|endoftext|>


<commit_before>extern crate ggez;

use ggez::*;
use ggez::graphics::{Canvas, Color, DrawMode, DrawParam, Point2};

struct MainState {
    canvas: Canvas,
}

impl MainState {
    fn new(ctx: &mut Context) -> GameResult<MainState> {
        let canvas = Canvas::with_window_size(ctx)?;
        Ok(MainState { canvas })
    }
}

impl event::EventHandler for MainState {
    fn update(&mut self, _ctx: &mut Context) -> GameResult<()> {
        Ok(())
    }

    fn draw(&mut self, ctx: &mut Context) -> GameResult<()> {
        // first lets render to our canvas
        graphics::set_canvas(ctx, Some(&self.canvas));
        graphics::set_background_color(ctx, Color::new(0.1, 0.2, 0.3, 1.0));
        graphics::clear(ctx);
        graphics::circle(ctx, DrawMode::Fill, Point2::new(400.0, 300.0), 100.0, 2.0)?;
        graphics::present(ctx);

        // now lets render our scene once in the top right and in the bottom
        // right
        graphics::set_canvas(ctx, None);
        graphics::set_background_color(ctx, Color::new(0.0, 0.0, 0.0, 1.0));
        graphics::clear(ctx);
        graphics::draw_ex(
            ctx,
            &self.canvas,
            DrawParam {
                dest: Point2::new(200.0, 150.0),
                scale: Point2::new(0.5, 0.5),
                ..Default::default()
            },
        )?;
        graphics::draw_ex(
            ctx,
            &self.canvas,
            DrawParam {
                dest: Point2::new(600.0, 450.0),
                scale: Point2::new(0.5, 0.5),
                ..Default::default()
            },
        )?;
        graphics::present(ctx);

        Ok(())
    }
}

pub fn main() {
    let c = conf::Conf::new();
    let ctx = &mut Context::load_from_conf("super_simple", "ggez", c).unwrap();
    let state = &mut MainState::new(ctx).unwrap();
    event::run(ctx, state).unwrap();
}
<commit_msg>shouldn't do a present after the off-screen render<commit_after>extern crate ggez;

use ggez::*;
use ggez::graphics::{Canvas, Color, DrawMode, DrawParam, Point2};

struct MainState {
    canvas: Canvas,
}

impl MainState {
    fn new(ctx: &mut Context) -> GameResult<MainState> {
        let canvas = Canvas::with_window_size(ctx)?;
        Ok(MainState { canvas })
    }
}

impl event::EventHandler for MainState {
    fn update(&mut self, _ctx: &mut Context) -> GameResult<()> {
        Ok(())
    }

    fn draw(&mut self, ctx: &mut Context) -> GameResult<()> {
        // first lets render to our canvas
        graphics::set_canvas(ctx, Some(&self.canvas));
        graphics::set_background_color(ctx, Color::new(0.1, 0.2, 0.3, 1.0));
        graphics::clear(ctx);
        graphics::circle(ctx, DrawMode::Fill, Point2::new(400.0, 300.0), 100.0, 2.0)?;

        // now lets render our scene once in the top right and in the bottom
        // right
        graphics::set_canvas(ctx, None);
        graphics::set_background_color(ctx, Color::new(0.0, 0.0, 0.0, 1.0));
        graphics::clear(ctx);
        graphics::draw_ex(
            ctx,
            &self.canvas,
            DrawParam {
                dest: Point2::new(200.0, 150.0),
                scale: Point2::new(0.5, 0.5),
                ..Default::default()
            },
        )?;
        graphics::draw_ex(
            ctx,
            &self.canvas,
            DrawParam {
                dest: Point2::new(600.0, 450.0),
                scale: Point2::new(0.5, 0.5),
                ..Default::default()
            },
        )?;
        graphics::present(ctx);

        Ok(())
    }
}

pub fn main() {
    let c = conf::Conf::new();
    let ctx = &mut Context::load_from_conf("super_simple", "ggez", c).unwrap();
    let state = &mut MainState::new(ctx).unwrap();
    event::run(ctx, state).unwrap();
}
<|endoftext|>


<commit_before>#![deny(warnings)]
use std::ascii::*;
/// Converts a `String` to camelCase `String`
///
/// #Examples
/// ```0
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "fooBar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```1
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FOO_BAR".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```2
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "Foo Bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```3
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "foo_bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```4
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "Foo bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```5
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "foo-bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```6
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FooBar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```7
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FooBar3".to_string();
///     let expected_string: String = "fooBar3".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
pub fn to_camel_case(non_camelized_string: String) -> String {
    let mut new_word: bool = false;
    let mut last_char: char = ' ';
    non_camelized_string
        .chars()
        .fold("".to_string(), |mut result, character|
            if character == '-' || character == '_' || character == ' ' {
                new_word = true;
                result
            } else if new_word || (
                (last_char.is_lowercase() && character.is_uppercase()) &&
                (last_char != ' ')
                ){
                new_word = false;
                result.push(character.to_ascii_uppercase());
                result
            } else {
                last_char = character;
                result.push(character.to_ascii_lowercase());
                result
            }
        )
}

/// Determines if a `String` is camelCase bool``
///
/// #Examples
/// ```0
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```1
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```2
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo-bar-string-that-is-really-really-long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```3
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "FooBarIsAReallyReallyLongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```4
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "fooBarIsAReallyReally3LongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```5
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "fooBarIsAReallyReallyLongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```6
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "FOO_BAR_STRING_THAT_IS_REALLY_REALLY_LONG".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```7
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo_bar_string_that_is_really_really_long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```8
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo bar string that is really really long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```9
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo Bar Is A Really Really Long String".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
/// ```
pub fn is_camel_case(test_string: String) -> bool {
    to_camel_case(test_string.clone()) == test_string
}

#[cfg(all(feature = "unstable", test))]
mod tests {
    extern crate test;
    use self::test::Bencher;

    #[bench]
    fn bench_camel0(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "Foo bar".to_string();
            super::to_camel_case(test_string)
        });
    }

    #[bench]
    fn bench_camel1(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "foo_bar".to_string();
            super::to_camel_case(test_string)
        }
        );
    }

    #[bench]
    fn bench_camel2(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "fooBar".to_string();
            super::to_camel_case(test_string)
        });
    }

    #[bench]
    fn bench_is_camel(b: &mut Bencher) {
        b.iter(|| {
            let test_string: String  = "Foo bar".to_string();
            super::is_camel_case(test_string)
        }
        );
    }
}
<commit_msg>Minor patch to camel case for cases where numbers are in the string<commit_after>#![deny(warnings)]
use std::ascii::*;
/// Converts a `String` to camelCase `String`
///
/// #Examples
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "fooBar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FOO_BAR".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "Foo Bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "foo_bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "Foo bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "foo-bar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FooBar".to_string();
///     let expected_string: String = "fooBar".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
/// ```
///     use inflector::cases::camelcase::to_camel_case;
///     let mock_string: String = "FooBar3".to_string();
///     let expected_string: String = "fooBar3".to_string();
///     let asserted_string: String = to_camel_case(mock_string);
///     assert!(asserted_string == expected_string);
///
/// ```
pub fn to_camel_case(non_camelized_string: String) -> String {
    let mut new_word: bool = false;
    let mut last_char: char = ' ';
    non_camelized_string
        .chars()
        .fold("".to_string(), |mut result, character|
            if character == '-' || character == '_' || character == ' ' {
                new_word = true;
                result
            } else if character.is_numeric() {
                new_word = true;
                result.push(character);
                result
            } else if new_word || (
                (last_char.is_lowercase() && character.is_uppercase()) &&
                (last_char != ' ')
                ){
                new_word = false;
                result.push(character.to_ascii_uppercase());
                result
            } else {
                last_char = character;
                result.push(character.to_ascii_lowercase());
                result
            }
        )
}

/// Determines if a `String` is camelCase bool``
///
/// #Examples
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo-bar-string-that-is-really-really-long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "FooBarIsAReallyReallyLongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "fooBarIsAReallyReally3LongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "fooBarIsAReallyReallyLongString".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == true);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "FOO_BAR_STRING_THAT_IS_REALLY_REALLY_LONG".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "foo_bar_string_that_is_really_really_long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo bar string that is really really long".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
///
///
/// ```
/// ```
///     use inflector::cases::camelcase::is_camel_case;
///     let mock_string: String = "Foo Bar Is A Really Really Long String".to_string();
///     let asserted_bool: bool = is_camel_case(mock_string);
///     assert!(asserted_bool == false);
/// ```
pub fn is_camel_case(test_string: String) -> bool {
    to_camel_case(test_string.clone()) == test_string
}

#[cfg(all(feature = "unstable", test))]
mod tests {
    extern crate test;
    use self::test::Bencher;

    #[bench]
    fn bench_camel0(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "Foo bar".to_string();
            super::to_camel_case(test_string)
        });
    }

    #[bench]
    fn bench_camel1(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "foo_bar".to_string();
            super::to_camel_case(test_string)
        }
        );
    }

    #[bench]
    fn bench_camel2(b: &mut Bencher) {
        b.iter(|| {
            let test_string = "fooBar".to_string();
            super::to_camel_case(test_string)
        });
    }

    #[bench]
    fn bench_is_camel(b: &mut Bencher) {
        b.iter(|| {
            let test_string: String  = "Foo bar".to_string();
            super::is_camel_case(test_string)
        }
        );
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

mod memory_usage;
#[cfg(feature = "cpu_profiler")]
mod google_cpu_profiler;

use std::{
    cell::RefCell,
    collections::HashSet,
    io::{stderr, Write},
    iter::repeat,
    mem,
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use itertools::Itertools;
use once_cell::sync::Lazy;

pub use crate::memory_usage::{Bytes, MemoryUsage};

// We use jemalloc mainly to get heap usage statistics, actual performance
// difference is not measures.
#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }
        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }
        let allowed = &stack.filter_data.allowed;
        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
            return Profiler { desc: None };
        }

        stack.starts.push(Instant::now());
        Profiler { desc: Some(desc.to_string()) }
    })
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // Filtering syntax
    // env RA_PROFILE=*             // dump everything
    // env RA_PROFILE=foo|bar|baz   // enabled only selected entries
    // env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
    pub fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter_data: Default::default() }
    }
}

#[derive(Default, Clone)]
struct FilterData {
    depth: usize,
    version: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        // Convert to millis for comparison to avoid problems with rounding
                        // (otherwise we could print `0ms` despite user's `>0` filter when
                        // `duration` is just a few nanos).
                        if duration.as_millis() > longer_than.as_millis() {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than, None);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(
    lvl: usize,
    msgs: &[Message],
    out: &mut impl Write,
    longer_than: Duration,
    total: Option<Duration>,
) {
    if msgs.is_empty() {
        return;
    }
    // The index of the first element that will be included in the slice when we recurse.
    let mut next_start = 0;
    let indent = repeat("    ").take(lvl).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    let mut accounted_for = Duration::default();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        accounted_for += duration;
        if duration.as_millis() > longer_than.as_millis() {
            writeln!(out, "{}{:5}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[next_start..i], out, longer_than, Some(duration));
        } else {
            short.push((msg, duration))
        }

        next_start = i + 1;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{}{:5}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }

    if let Some(total) = total {
        if let Some(unaccounted) = total.checked_sub(accounted_for) {
            let unaccounted_millis = unaccounted.as_millis();
            if unaccounted_millis > longer_than.as_millis() && unaccounted_millis > 0 {
                writeln!(out, "{}{:5}ms - ???", indent, unaccounted_millis)
                    .expect("printing profiling info to stdout");
            }
        }
    }
}

/// Prints backtrace to stderr, useful for debugging.
pub fn print_backtrace() {
    let bt = backtrace::Backtrace::new();
    eprintln!("{:?}", bt);
}

thread_local!(static IN_SCOPE: RefCell<bool> = RefCell::new(false));

/// Allows to check if the current code is withing some dynamic scope, can be
/// useful during debugging to figure out why a function is called.
pub struct Scope {
    prev: bool,
}

impl Scope {
    pub fn enter() -> Scope {
        let prev = IN_SCOPE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), true));
        Scope { prev }
    }
    pub fn is_active() -> bool {
        IN_SCOPE.with(|slot| *slot.borrow())
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        IN_SCOPE.with(|slot| *slot.borrow_mut() = self.prev);
    }
}

/// A wrapper around google_cpu_profiler.
///
/// Usage:
/// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
/// 2. Build with `cpu_profiler` feature.
/// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
/// 4. Install pprof for visualization (https://github.com/google/pprof).
/// 5. Use something like `pprof -svg target/release/ra_cli ./out.profile` to see the results.
///
/// For example, here's how I run profiling on NixOS:
///
/// ```bash
/// $ nix-shell -p gperftools --run \
///     'cargo run --release -p ra_cli -- parse < ~/projects/rustbench/parser.rs > /dev/null'
/// ```
#[derive(Debug)]
pub struct CpuProfiler {
    _private: (),
}

pub fn cpu_profiler() -> CpuProfiler {
    #[cfg(feature = "cpu_profiler")]
    {
        google_cpu_profiler::start("./out.profile".as_ref())
    }

    #[cfg(not(feature = "cpu_profiler"))]
    {
        eprintln!("cpu_profiler feature is disabled")
    }

    CpuProfiler { _private: () }
}

impl Drop for CpuProfiler {
    fn drop(&mut self) {
        #[cfg(feature = "cpu_profiler")]
        {
            google_cpu_profiler::stop()
        }
    }
}

pub fn memory_usage() -> MemoryUsage {
    MemoryUsage::current()
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }

    #[test]
    fn test_longer_than() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_nanos(3), message: "bar".to_owned() },
            Message { level: 1, duration: Duration::from_nanos(2), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(1), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
        // The calls to `bar` are so short that they'll be rounded to 0ms and should get collapsed
        // when printing.
        assert_eq!(
            std::str::from_utf8(&result).unwrap(),
            "    1ms - foo\n        0ms - bar (2 calls)\n"
        );
    }

    #[test]
    fn test_unaccounted_for_topmost() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_millis(2), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(5), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    5ms - foo",
                "        2ms - bar",
                "        3ms - ???",
                // Dummy comment to improve formatting
            ]
        );
    }

    #[test]
    fn test_unaccounted_for_multiple_levels() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 2, duration: Duration::from_millis(3), message: "baz".to_owned() },
            Message { level: 1, duration: Duration::from_millis(5), message: "bar".to_owned() },
            Message { level: 2, duration: Duration::from_millis(2), message: "baz".to_owned() },
            Message { level: 1, duration: Duration::from_millis(4), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(9), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), None);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    9ms - foo",
                "        5ms - bar",
                "            3ms - baz",
                "            2ms - ???",
                "        4ms - bar",
                "            2ms - baz",
                "            2ms - ???",
            ]
        );
    }
}
<commit_msg>Add print_time helper<commit_after>//! FIXME: write short doc here

mod memory_usage;
#[cfg(feature = "cpu_profiler")]
mod google_cpu_profiler;

use std::{
    cell::RefCell,
    collections::HashSet,
    io::{stderr, Write},
    iter::repeat,
    mem,
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use itertools::Itertools;
use once_cell::sync::Lazy;

pub use crate::memory_usage::{Bytes, MemoryUsage};

// We use jemalloc mainly to get heap usage statistics, actual performance
// difference is not measures.
#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }
        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }
        let allowed = &stack.filter_data.allowed;
        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
            return Profiler { desc: None };
        }

        stack.starts.push(Instant::now());
        Profiler { desc: Some(desc.to_string()) }
    })
}

pub fn print_time(desc: &str) -> impl Drop + '_ {
    struct Guard<'a> {
        desc: &'a str,
        start: Instant,
    }

    impl Drop for Guard<'_> {
        fn drop(&mut self) {
            eprintln!("{}: {:?}", self.desc, self.start.elapsed())
        }
    }

    Guard { desc, start: Instant::now() }
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // Filtering syntax
    // env RA_PROFILE=*             // dump everything
    // env RA_PROFILE=foo|bar|baz   // enabled only selected entries
    // env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
    pub fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter_data: Default::default() }
    }
}

#[derive(Default, Clone)]
struct FilterData {
    depth: usize,
    version: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        // Convert to millis for comparison to avoid problems with rounding
                        // (otherwise we could print `0ms` despite user's `>0` filter when
                        // `duration` is just a few nanos).
                        if duration.as_millis() > longer_than.as_millis() {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than, None);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(
    lvl: usize,
    msgs: &[Message],
    out: &mut impl Write,
    longer_than: Duration,
    total: Option<Duration>,
) {
    if msgs.is_empty() {
        return;
    }
    // The index of the first element that will be included in the slice when we recurse.
    let mut next_start = 0;
    let indent = repeat("    ").take(lvl).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    let mut accounted_for = Duration::default();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        accounted_for += duration;
        if duration.as_millis() > longer_than.as_millis() {
            writeln!(out, "{}{:5}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[next_start..i], out, longer_than, Some(duration));
        } else {
            short.push((msg, duration))
        }

        next_start = i + 1;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{}{:5}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }

    if let Some(total) = total {
        if let Some(unaccounted) = total.checked_sub(accounted_for) {
            let unaccounted_millis = unaccounted.as_millis();
            if unaccounted_millis > longer_than.as_millis() && unaccounted_millis > 0 {
                writeln!(out, "{}{:5}ms - ???", indent, unaccounted_millis)
                    .expect("printing profiling info to stdout");
            }
        }
    }
}

/// Prints backtrace to stderr, useful for debugging.
pub fn print_backtrace() {
    let bt = backtrace::Backtrace::new();
    eprintln!("{:?}", bt);
}

thread_local!(static IN_SCOPE: RefCell<bool> = RefCell::new(false));

/// Allows to check if the current code is withing some dynamic scope, can be
/// useful during debugging to figure out why a function is called.
pub struct Scope {
    prev: bool,
}

impl Scope {
    pub fn enter() -> Scope {
        let prev = IN_SCOPE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), true));
        Scope { prev }
    }
    pub fn is_active() -> bool {
        IN_SCOPE.with(|slot| *slot.borrow())
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        IN_SCOPE.with(|slot| *slot.borrow_mut() = self.prev);
    }
}

/// A wrapper around google_cpu_profiler.
///
/// Usage:
/// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
/// 2. Build with `cpu_profiler` feature.
/// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
/// 4. Install pprof for visualization (https://github.com/google/pprof).
/// 5. Use something like `pprof -svg target/release/ra_cli ./out.profile` to see the results.
///
/// For example, here's how I run profiling on NixOS:
///
/// ```bash
/// $ nix-shell -p gperftools --run \
///     'cargo run --release -p ra_cli -- parse < ~/projects/rustbench/parser.rs > /dev/null'
/// ```
#[derive(Debug)]
pub struct CpuProfiler {
    _private: (),
}

pub fn cpu_profiler() -> CpuProfiler {
    #[cfg(feature = "cpu_profiler")]
    {
        google_cpu_profiler::start("./out.profile".as_ref())
    }

    #[cfg(not(feature = "cpu_profiler"))]
    {
        eprintln!("cpu_profiler feature is disabled")
    }

    CpuProfiler { _private: () }
}

impl Drop for CpuProfiler {
    fn drop(&mut self) {
        #[cfg(feature = "cpu_profiler")]
        {
            google_cpu_profiler::stop()
        }
    }
}

pub fn memory_usage() -> MemoryUsage {
    MemoryUsage::current()
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }

    #[test]
    fn test_longer_than() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_nanos(3), message: "bar".to_owned() },
            Message { level: 1, duration: Duration::from_nanos(2), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(1), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
        // The calls to `bar` are so short that they'll be rounded to 0ms and should get collapsed
        // when printing.
        assert_eq!(
            std::str::from_utf8(&result).unwrap(),
            "    1ms - foo\n        0ms - bar (2 calls)\n"
        );
    }

    #[test]
    fn test_unaccounted_for_topmost() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_millis(2), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(5), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    5ms - foo",
                "        2ms - bar",
                "        3ms - ???",
                // Dummy comment to improve formatting
            ]
        );
    }

    #[test]
    fn test_unaccounted_for_multiple_levels() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 2, duration: Duration::from_millis(3), message: "baz".to_owned() },
            Message { level: 1, duration: Duration::from_millis(5), message: "bar".to_owned() },
            Message { level: 2, duration: Duration::from_millis(2), message: "baz".to_owned() },
            Message { level: 1, duration: Duration::from_millis(4), message: "bar".to_owned() },
            Message { level: 0, duration: Duration::from_millis(9), message: "foo".to_owned() },
        ];
        print(0, &msgs, &mut result, Duration::from_millis(0), None);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    9ms - foo",
                "        5ms - bar",
                "            3ms - baz",
                "            2ms - ???",
                "        4ms - bar",
                "            2ms - baz",
                "            2ms - ???",
            ]
        );
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

mod memory_usage;
#[cfg(feature = "cpu_profiler")]
mod google_cpu_profiler;

use std::{
    cell::RefCell,
    collections::HashSet,
    io::{stderr, Write},
    iter::repeat,
    mem,
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use itertools::Itertools;
use once_cell::sync::Lazy;

pub use crate::memory_usage::{Bytes, MemoryUsage};

// We use jemalloc mainly to get heap usage statistics, actual performance
// difference is not measures.
#[cfg(feature = "jemalloc")]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }
        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }
        let allowed = &stack.filter_data.allowed;
        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
            return Profiler { desc: None };
        }

        stack.starts.push(Instant::now());
        Profiler { desc: Some(desc.to_string()) }
    })
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // Filtering syntax
    // env RA_PROFILE=*             // dump everything
    // env RA_PROFILE=foo|bar|baz   // enabled only selected entries
    // env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
    pub fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter_data: Default::default() }
    }
}

#[derive(Default, Clone)]
struct FilterData {
    depth: usize,
    version: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        if duration >= longer_than {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(lvl: usize, msgs: &[Message], out: &mut impl Write, longer_than: Duration) {
    let mut last = 0;
    let indent = repeat("    ").take(lvl + 1).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        if duration >= longer_than {
            writeln!(out, "{} {:6}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[last..i], out, longer_than);
        } else {
            short.push((msg, duration))
        }

        last = i;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{} {:6}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }
}

/// Prints backtrace to stderr, useful for debugging.
pub fn print_backtrace() {
    let bt = backtrace::Backtrace::new();
    eprintln!("{:?}", bt);
}

thread_local!(static IN_SCOPE: RefCell<bool> = RefCell::new(false));

/// Allows to check if the current code is withing some dynamic scope, can be
/// useful during debugging to figure out why a function is called.
pub struct Scope {
    prev: bool,
}

impl Scope {
    pub fn enter() -> Scope {
        let prev = IN_SCOPE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), true));
        Scope { prev }
    }
    pub fn is_active() -> bool {
        IN_SCOPE.with(|slot| *slot.borrow())
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        IN_SCOPE.with(|slot| *slot.borrow_mut() = self.prev);
    }
}

/// A wrapper around google_cpu_profiler.
///
/// Usage:
/// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
/// 2. Build with `cpu_profiler` feature.
/// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
/// 4. Install pprof for visualization (https://github.com/google/pprof).
/// 5. Use something like `pprof -svg target/release/ra_cli ./out.profile` to see the results.
///
/// For example, here's how I run profiling on NixOS:
///
/// ```bash
/// $ nix-shell -p gperftools --run \
///     'cargo run --release -p ra_cli -- parse < ~/projects/rustbench/parser.rs > /dev/null'
/// ```
#[derive(Debug)]
pub struct CpuProfiler {
    _private: (),
}

pub fn cpu_profiler() -> CpuProfiler {
    #[cfg(feature = "cpu_profiler")]
    {
        google_cpu_profiler::start("./out.profile".as_ref())
    }

    #[cfg(not(feature = "cpu_profiler"))]
    {
        eprintln!("cpu_profiler feature is disabled")
    }

    CpuProfiler { _private: () }
}

impl Drop for CpuProfiler {
    fn drop(&mut self) {
        #[cfg(feature = "cpu_profiler")]
        {
            google_cpu_profiler::stop()
        }
    }
}

pub fn memory_usage() -> MemoryUsage {
    MemoryUsage::current()
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }
}
<commit_msg>show unaccounted for time in profiling<commit_after>//! FIXME: write short doc here

mod memory_usage;
#[cfg(feature = "cpu_profiler")]
mod google_cpu_profiler;

use std::{
    cell::RefCell,
    collections::HashSet,
    io::{stderr, Write},
    iter::repeat,
    mem,
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use itertools::Itertools;
use once_cell::sync::Lazy;

pub use crate::memory_usage::{Bytes, MemoryUsage};

// We use jemalloc mainly to get heap usage statistics, actual performance
// difference is not measures.
#[cfg(feature = "jemalloc")]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }
        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }
        let allowed = &stack.filter_data.allowed;
        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
            return Profiler { desc: None };
        }

        stack.starts.push(Instant::now());
        Profiler { desc: Some(desc.to_string()) }
    })
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // Filtering syntax
    // env RA_PROFILE=*             // dump everything
    // env RA_PROFILE=foo|bar|baz   // enabled only selected entries
    // env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
    pub fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter_data: Default::default() }
    }
}

#[derive(Default, Clone)]
struct FilterData {
    depth: usize,
    version: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        if duration >= longer_than {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than, None);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(
    lvl: usize,
    msgs: &[Message],
    out: &mut impl Write,
    longer_than: Duration,
    total: Option<Duration>,
) {
    let mut last = 0;
    let indent = repeat("    ").take(lvl + 1).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    let mut accounted_for = Duration::default();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        accounted_for += duration;
        if duration >= longer_than {
            writeln!(out, "{} {:6}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[last..i], out, longer_than, Some(duration));
        } else {
            short.push((msg, duration))
        }

        last = i;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{} {:6}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }

    if let Some(total) = total {
        if let Some(unaccounted) = total.checked_sub(accounted_for) {
            if unaccounted >= longer_than && last > 0 {
                writeln!(out, "{} {:6}ms - ???", indent, unaccounted.as_millis())
                    .expect("printing profiling info to stdout");
            }
        }
    }
}

/// Prints backtrace to stderr, useful for debugging.
pub fn print_backtrace() {
    let bt = backtrace::Backtrace::new();
    eprintln!("{:?}", bt);
}

thread_local!(static IN_SCOPE: RefCell<bool> = RefCell::new(false));

/// Allows to check if the current code is withing some dynamic scope, can be
/// useful during debugging to figure out why a function is called.
pub struct Scope {
    prev: bool,
}

impl Scope {
    pub fn enter() -> Scope {
        let prev = IN_SCOPE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), true));
        Scope { prev }
    }
    pub fn is_active() -> bool {
        IN_SCOPE.with(|slot| *slot.borrow())
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        IN_SCOPE.with(|slot| *slot.borrow_mut() = self.prev);
    }
}

/// A wrapper around google_cpu_profiler.
///
/// Usage:
/// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
/// 2. Build with `cpu_profiler` feature.
/// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
/// 4. Install pprof for visualization (https://github.com/google/pprof).
/// 5. Use something like `pprof -svg target/release/ra_cli ./out.profile` to see the results.
///
/// For example, here's how I run profiling on NixOS:
///
/// ```bash
/// $ nix-shell -p gperftools --run \
///     'cargo run --release -p ra_cli -- parse < ~/projects/rustbench/parser.rs > /dev/null'
/// ```
#[derive(Debug)]
pub struct CpuProfiler {
    _private: (),
}

pub fn cpu_profiler() -> CpuProfiler {
    #[cfg(feature = "cpu_profiler")]
    {
        google_cpu_profiler::start("./out.profile".as_ref())
    }

    #[cfg(not(feature = "cpu_profiler"))]
    {
        eprintln!("cpu_profiler feature is disabled")
    }

    CpuProfiler { _private: () }
}

impl Drop for CpuProfiler {
    fn drop(&mut self) {
        #[cfg(feature = "cpu_profiler")]
        {
            google_cpu_profiler::stop()
        }
    }
}

pub fn memory_usage() -> MemoryUsage {
    MemoryUsage::current()
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }
}
<|endoftext|>