<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN"
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    if henv::var(ORIGIN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }
    if henv::var(AUTH_TOKEN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(auth_token) = config.auth_token {
            debug!("Setting default origin {} via CLI config", &auth_token);
            env::set_var("HAB_AUTH_TOKEN", auth_token;
        }
    }

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<commit_msg>fix typos<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    if henv::var(ORIGIN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }
    if henv::var(AUTH_TOKEN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(auth_token) = config.auth_token {
            debug!("Setting default origin {} via CLI config", &auth_token);
            env::set_var("HAB_AUTH_TOKEN", auth_token);
        }
    }

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<|endoftext|>


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

use stb_image::bindgen::*;

use std::libc;
use std::libc::{c_void, c_int};
use std::vec::raw::from_buf_raw;

pub struct Image<T> {
    width   : uint,
    height  : uint,
    depth   : uint,
    data    : ~[T],
}

pub fn new_image<T>(width: uint, height: uint, depth: uint, data: ~[T]) -> Image<T> {
    Image::<T> {
        width   : width,
        height  : height,
        depth   : depth,
        data    : data,
    }
}

pub enum LoadResult {
    Error(~str),
    ImageU8(Image<u8>),
    ImageF32(Image<f32>),
}

impl LoadResult {
    pub fn from_result(res: Result<LoadResult,~Any>)-> LoadResult {
        match res {
            Ok(res) => res,
            Err(e)  => Error(e.to_str()),
        }
    }
}

pub fn load(path: ~str) -> LoadResult {
    let force_depth = 0;
    load_with_depth(path, force_depth, false)
}

fn load_internal<T>(buf : *T, w : c_int, h : c_int, d : c_int) -> Image<T> {
    unsafe {
        // FIXME: Shouldn't copy; instead we should use a sendable resource. They
        // aren't particularly safe yet though.
        let data = from_buf_raw(buf, (w * h * d) as uint);
        libc::free(buf as *mut c_void);
        Image::<T>{
            width   : w as uint,
            height  : h as uint,
            depth   : d as uint,
            data    : data}
    }
}

pub fn load_with_depth(path: ~str, force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        path.to_c_str().with_ref(|bytes| {
            if !convert_hdr && stbi_is_hdr(bytes)!=0   {
                let buffer = stbi_loadf(bytes,
                                        &mut width,
                                        &mut height,
                                        &mut depth,
                                        force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_loadf failed")
                } else {
                    ImageF32(load_internal(buffer, width, height, depth))
                }
            } else {
                let buffer = stbi_load(bytes,
                                       &mut width,
                                       &mut height,
                                       &mut depth,
                                       force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_load failed")
                } else {
                    ImageU8(load_internal(buffer, width, height, depth))
                }
            }
        })
    }
}

pub fn load_from_memory(buffer: &[u8]) -> LoadResult {
    let force_depth = 0;
    load_from_memory_with_depth(buffer, force_depth, false)
}

pub fn load_from_memory_with_depth(buffer: &[u8], force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        if !convert_hdr && stbi_is_hdr_from_memory(buffer.as_ptr(), buffer.len() as c_int) != 0 {
            let buffer = stbi_loadf_from_memory(buffer.as_ptr(),
                                                buffer.len() as c_int,
                                                &mut width,
                                                &mut height,
                                                &mut depth,
                                                force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_loadf_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageF32(load_internal(buffer, width, height, actual_depth))
            }
        } else {
            let buffer = stbi_load_from_memory(buffer.as_ptr(),
                                               buffer.len() as c_int,
                                               &mut width,
                                               &mut height,
                                               &mut depth,
                                               force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_load_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageU8(load_internal(buffer, width, height, actual_depth))
            }
        }
    }
}
<commit_msg>Updated to latest Rust: use std::any::Any<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use stb_image::bindgen::*;

use std::any::Any;
use std::libc;
use std::libc::{c_void, c_int};
use std::vec::raw::from_buf_raw;

pub struct Image<T> {
    width   : uint,
    height  : uint,
    depth   : uint,
    data    : ~[T],
}

pub fn new_image<T>(width: uint, height: uint, depth: uint, data: ~[T]) -> Image<T> {
    Image::<T> {
        width   : width,
        height  : height,
        depth   : depth,
        data    : data,
    }
}

pub enum LoadResult {
    Error(~str),
    ImageU8(Image<u8>),
    ImageF32(Image<f32>),
}

impl LoadResult {
    pub fn from_result(res: Result<LoadResult,~Any>)-> LoadResult {
        match res {
            Ok(res) => res,
            Err(e)  => Error(e.to_str()),
        }
    }
}

pub fn load(path: ~str) -> LoadResult {
    let force_depth = 0;
    load_with_depth(path, force_depth, false)
}

fn load_internal<T>(buf : *T, w : c_int, h : c_int, d : c_int) -> Image<T> {
    unsafe {
        // FIXME: Shouldn't copy; instead we should use a sendable resource. They
        // aren't particularly safe yet though.
        let data = from_buf_raw(buf, (w * h * d) as uint);
        libc::free(buf as *mut c_void);
        Image::<T>{
            width   : w as uint,
            height  : h as uint,
            depth   : d as uint,
            data    : data}
    }
}

pub fn load_with_depth(path: ~str, force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        path.to_c_str().with_ref(|bytes| {
            if !convert_hdr && stbi_is_hdr(bytes)!=0   {
                let buffer = stbi_loadf(bytes,
                                        &mut width,
                                        &mut height,
                                        &mut depth,
                                        force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_loadf failed")
                } else {
                    ImageF32(load_internal(buffer, width, height, depth))
                }
            } else {
                let buffer = stbi_load(bytes,
                                       &mut width,
                                       &mut height,
                                       &mut depth,
                                       force_depth as c_int);
                if buffer.is_null() {
                    Error(~"stbi_load failed")
                } else {
                    ImageU8(load_internal(buffer, width, height, depth))
                }
            }
        })
    }
}

pub fn load_from_memory(buffer: &[u8]) -> LoadResult {
    let force_depth = 0;
    load_from_memory_with_depth(buffer, force_depth, false)
}

pub fn load_from_memory_with_depth(buffer: &[u8], force_depth: uint, convert_hdr:bool) -> LoadResult {
    unsafe {
        let mut width = 0 as c_int;
        let mut height = 0 as c_int;
        let mut depth = 0 as c_int;
        if !convert_hdr && stbi_is_hdr_from_memory(buffer.as_ptr(), buffer.len() as c_int) != 0 {
            let buffer = stbi_loadf_from_memory(buffer.as_ptr(),
                                                buffer.len() as c_int,
                                                &mut width,
                                                &mut height,
                                                &mut depth,
                                                force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_loadf_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageF32(load_internal(buffer, width, height, actual_depth))
            }
        } else {
            let buffer = stbi_load_from_memory(buffer.as_ptr(),
                                               buffer.len() as c_int,
                                               &mut width,
                                               &mut height,
                                               &mut depth,
                                               force_depth as c_int);
            if buffer.is_null() {
                Error(~"stbi_load_from_memory failed")
            } else {
                let actual_depth = if force_depth != 0 { force_depth as c_int } else { depth };
                ImageU8(load_internal(buffer, width, height, actual_depth))
            }
        }
    }
}
<|endoftext|>


<commit_before>#![allow(clippy::new_without_default)]

//! Interprets the terminal events we care about (keyboard input).
//!
//! Keyboard events are read from `stdin`. We look for byte strings of key combinations that we
//! care about. E.g. Alt-arrow keys, C-w etc.
//!
//! NOTE: Make sure to either set `stdin` to non-blocking mode or enable non-canonical input. For
//! the latter, see:
//!
//! - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
//! - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
//!
//! Also make sure to set `VMIN` and `VTIME` to zero in non-canonical mode.

#[cfg(test)]
mod tests;

use std::char;
use std::collections::VecDeque;
use std::os::unix::io::RawFd;
use std::pin::Pin;
use std::task::{Context, Poll};

use tokio::io::unix::AsyncFd;
use tokio::stream::Stream;

use term_input_macros::byte_seq_parser;

////////////////////////////////////////////////////////////////////////////////////////////////////
// Public types
////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Key {
    AltArrow(Arrow),
    AltChar(char),
    Arrow(Arrow),
    Backspace,
    Char(char),
    Ctrl(char),
    CtrlArrow(Arrow),
    Del,
    End,
    Esc,
    Home,
    PageDown,
    PageUp,
    ShiftDown,
    ShiftUp,
    Tab,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arrow {
    Left,
    Right,
    Up,
    Down,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Event {
    /// A single key input.
    Key(Key),

    /// Usually a paste.
    String(String),

    // These are not generated anymore. I needed these for a feature, but never really got around
    // implementing it.
    // FocusGained,
    // FocusLost,
    /// An unknown sequence of bytes (probably for a key combination that we don't care about).
    Unknown(Vec<u8>),
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Byte sequences of key presses we want to capture
////////////////////////////////////////////////////////////////////////////////////////////////////

byte_seq_parser! {
    parse_key_bytes -> Key, // Function name. Generation function type is
                            // fn(&[u8]) -> Option<(Key, usize)>.

    [27, 91, 49, 59, 51, 66] => Key::AltArrow(Arrow::Down),
    [27, 91, 49, 59, 51, 68] => Key::AltArrow(Arrow::Left),
    [27, 91, 49, 59, 51, 67] => Key::AltArrow(Arrow::Right),
    [27, 91, 49, 59, 51, 65] => Key::AltArrow(Arrow::Up),
    [27, 91, 66] => Key::Arrow(Arrow::Down),
    [27, 79, 66] => Key::Arrow(Arrow::Down),
    [27, 91, 68] => Key::Arrow(Arrow::Left),
    [27, 79, 68] => Key::Arrow(Arrow::Left),
    [27, 91, 67] => Key::Arrow(Arrow::Right),
    [27, 79, 67] => Key::Arrow(Arrow::Right),
    [27, 91, 65] => Key::Arrow(Arrow::Up),
    [27, 79, 65] => Key::Arrow(Arrow::Up),
    [27, 91, 49, 59, 53, 66] => Key::CtrlArrow(Arrow::Down),
    [27, 91, 49, 59, 53, 68] => Key::CtrlArrow(Arrow::Left),
    [27, 91, 49, 59, 53, 67] => Key::CtrlArrow(Arrow::Right),
    [27, 91, 49, 59, 53, 65] => Key::CtrlArrow(Arrow::Up),
    [27, 91, 51, 126] => Key::Del,
    [27, 91, 54, 126] => Key::PageDown,
    [27, 91, 53, 126] => Key::PageUp,
    [27, 91, 49, 59, 50, 65] => Key::ShiftUp,
    [27, 91, 49, 59, 50, 66] => Key::ShiftDown,
    [27, 91, 72] => Key::Home,
    [27, 91, 70] => Key::End,
    [27, 79, 72] => Key::Home,
    [27, 79, 70] => Key::End,
    [27, 91, 52, 126] => Key::End,
    [9] => Key::Tab,
    [127] => Key::Backspace,
    [1] => Key::Ctrl('a'),
    [5] => Key::Ctrl('e'),
    [23] => Key::Ctrl('w'),
    [11] => Key::Ctrl('k'),
    [4] => Key::Ctrl('d'),
    [3] => Key::Ctrl('c'),
    [17] => Key::Ctrl('q'),
    [16] => Key::Ctrl('p'),
    [14] => Key::Ctrl('n'),
    [21] => Key::Ctrl('u'),
    [24] => Key::Ctrl('x'),
}

// static XTERM_FOCUS_GAINED: [u8; 3] = [27, 91, 73];
// static XTERM_FOCUS_LOST: [u8; 3] = [27, 91, 79];

////////////////////////////////////////////////////////////////////////////////////////////////////

pub struct Input {
    /// Queue of events waiting to be polled.
    evs: VecDeque<Event>,

    /// Used when reading from stdin.
    buf: Vec<u8>,

    stdin: AsyncFd<RawFd>,
}

impl Input {
    /// Create a new input handler. Requires a `tokio` reactor to be running. Make sure to either
    /// set `stdin` to non-blocking mode or enable non-canonical input, otherwise the `Stream`
    /// implementation will not work. If you are using this with `termbox` it does the latter so no
    /// need to do anything. See module documentation for links on how to enable non-canonical
    /// input.
    pub fn new() -> Input {
        Input {
            evs: VecDeque::new(),
            buf: Vec::with_capacity(100),
            stdin: AsyncFd::new(libc::STDIN_FILENO).unwrap(),
        }
    }

    fn parse_buffer(&mut self) {
        let mut buf_slice: &[u8] = &self.buf;

        while !buf_slice.is_empty() {
            // Special treatment for 127 (backspace, 0x1B) and 13 ('\r', 0xD)
            let fst = buf_slice[0];
            let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
                parse_key_comb
            } else {
                parse_chars
            };

            match parse_fn(buf_slice) {
                Some((ev, used)) => {
                    buf_slice = &buf_slice[used..];
                    self.evs.push_back(ev);
                }
                None => {
                    self.evs.push_back(Event::Unknown(buf_slice.to_owned()));
                    break;
                }
            }
        }

        self.buf.clear();
    }
}

impl Stream for Input {
    type Item = std::io::Result<Event>;

    fn poll_next(mut self: Pin<&mut Input>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let self_: &mut Input = &mut *self;

        'main: loop {
            // Try to parse any bytes in the input buffer from the last poll
            self_.parse_buffer();

            // Yield pending events
            if let Some(ev) = self_.evs.pop_front() {
                return Poll::Ready(Some(Ok(ev)));
            }

            // Otherwise read stdin and loop if successful
            let mut poll_ret = self_.stdin.poll_read_ready(cx);
            loop {
                match poll_ret {
                    Poll::Ready(Ok(mut ready)) => {
                        // We will read stdin until it's empty, clear the readiness state now
                        // (before reading) to avoid a potential race that would occur if stdin
                        // becomes available again between the last read and clear_ready.
                        ready.clear_ready();

                        // Read stdin until it fails with EWOULDBLOCK or returns 0
                        let read_ret = read_stdin(&mut self_.buf);

                        // Poll again so that we will be notified when stdin becomes ready again
                        poll_ret = self_.stdin.poll_read_ready(cx);

                        match read_ret {
                            Ok(()) => {}
                            Err(err) => {
                                // NOTE: `poll_ret` is ignored here but I think that's OK?
                                // TODO: Contents of `buf` is ignored here until `poll_next` is
                                // called again.
                                let err =
                                    std::io::Error::from(err.as_errno().expect("Weird nix error"));
                                return Poll::Ready(Some(Err(err)));
                            }
                        }

                        if self_.buf.is_empty() {
                            continue;
                        } else {
                            continue 'main;
                        }
                    }
                    Poll::Ready(Err(err)) => {
                        return Poll::Ready(Some(Err(err)));
                    }
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub(crate) fn parse_single_event(buf: &[u8]) -> Event {
    let fst = buf[0];
    let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
        parse_key_comb
    } else {
        parse_chars
    };

    let (ev, used) = match parse_fn(buf) {
        Some((ev, used)) => (ev, used),
        None => (Event::Unknown(buf.to_owned()), buf.len()),
    };

    assert_eq!(buf.len(), used);

    ev
}

fn parse_chars(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    // Use a fast path for the common case: Single utf-8 character.
    // (TODO: What about other encodings?)

    utf8_char_len(buf[0]).map(|char_len| {
        if char_len as usize == buf.len() {
            // Fast path: single character
            let ev = Event::Key(Key::Char(get_utf8_char(buf, char_len)));
            (ev, char_len as usize)
        } else {
            // Probably a paste: allocate a string and collect chars
            let mut string = String::with_capacity(1000);
            let mut start_idx = 0;
            loop {
                if start_idx == buf.len() {
                    break;
                }
                match utf8_char_len(buf[start_idx]) {
                    Some(char_len) => {
                        string.push(get_utf8_char(&buf[start_idx..], char_len));
                        start_idx += char_len as usize;
                    }
                    None => {
                        break;
                    }
                }
            }
            let ev = Event::String(string);
            (ev, start_idx)
        }
    })
}

fn parse_key_comb(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    if let Some((key, used)) = parse_key_bytes(buf) {
        return Some((Event::Key(key), used));
    }

    if buf[0] == 27 {
        // 0x1B, ESC
        // 27 not followed by anything is an actual ESC
        if buf.len() == 1 {
            return Some((Event::Key(Key::Esc), 1));
        }
        // otherwise it's probably alt + key
        else {
            debug_assert!(buf.len() >= 2);
            return utf8_char_len(buf[1]).map(|char_len| {
                let ev = Event::Key(Key::AltChar(get_utf8_char(&buf[1..], char_len)));
                (ev, char_len as usize + 1)
            });
        }
    }

    None
}

fn utf8_char_len(byte: u8) -> Option<u8> {
    if byte >> 7 == 0b0 {
        Some(1)
    } else if byte >> 5 == 0b110 {
        Some(2)
    } else if byte >> 4 == 0b1110 {
        Some(3)
    } else if byte >> 3 == 0b11110 {
        Some(4)
    } else {
        None
    }
}

fn get_utf8_char(buf: &[u8], len: u8) -> char {
    let codepoint: u32 = {
        if len == 1 {
            u32::from(buf[0] & 0b0111_1111)
        } else if len == 2 {
            ((u32::from(buf[0] & 0b0001_1111)) << 6) + (u32::from(buf[1] & 0b0011_1111))
        } else if len == 3 {
            ((u32::from(buf[0] & 0b0000_1111)) << 12)
                + ((u32::from(buf[1] & 0b0011_1111)) << 6)
                + (u32::from(buf[2] & 0b0011_1111))
        } else {
            debug_assert!(len == 4);
            ((u32::from(buf[0] & 0b0000_0111)) << 18)
                + ((u32::from(buf[1] & 0b0011_1111)) << 12)
                + ((u32::from(buf[2] & 0b0011_1111)) << 6)
                + (u32::from(buf[3] & 0b0011_1111))
        }
    };

    char::from_u32(codepoint).unwrap()
}

/// Read `stdin` until `read` fails with `EWOULDBLOCK` (happens in canonical mode, when `stdin` is
/// set to non-blocking mode) or returns 0 (happens in non-canonical mode when `VMIN` and `VTIME`
/// are 0). If you are using `term_input` with `termbox` then you don't need to set `stdin` to
/// non-blocking mode as `termbox` enables non-canonical mode.
///
/// See also:
///
/// - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
/// - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
pub fn read_stdin(buf: &mut Vec<u8>) -> Result<(), nix::Error> {
    loop {
        let old_len = buf.len();
        buf.reserve(100);
        unsafe {
            buf.set_len(old_len + 100);
        }

        match nix::unistd::read(libc::STDIN_FILENO, &mut buf[old_len..]) {
            Ok(n_read) => {
                unsafe { buf.set_len(old_len + n_read) };
                if n_read == 0 {
                    // We're in non-canonical mode, or stdin is closed. We can't distinguish the
                    // two here but I think it's fine to return OK when stdin is closed.
                    return Ok(());
                }
            }
            Err(err) => {
                unsafe { buf.set_len(old_len) };
                match err {
                    nix::Error::Sys(nix::errno::EWOULDBLOCK) => {
                        return Ok(());
                    }
                    _ => {
                        return Err(err);
                    }
                }
            }
        }
    }
}
<commit_msg>term_input: More typos<commit_after>#![allow(clippy::new_without_default)]

//! Interprets the terminal events we care about (keyboard input).
//!
//! Keyboard events are read from `stdin`. We look for byte strings of key combinations that we
//! care about. E.g. Alt-arrow keys, C-w etc.
//!
//! NOTE: Make sure to either set `stdin` to non-blocking mode or enable non-canonical input. For
//! the latter, see:
//!
//! - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
//! - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
//!
//! Also make sure to set `VMIN` and `VTIME` to zero in non-canonical mode.

#[cfg(test)]
mod tests;

use std::char;
use std::collections::VecDeque;
use std::os::unix::io::RawFd;
use std::pin::Pin;
use std::task::{Context, Poll};

use tokio::io::unix::AsyncFd;
use tokio::stream::Stream;

use term_input_macros::byte_seq_parser;

////////////////////////////////////////////////////////////////////////////////////////////////////
// Public types
////////////////////////////////////////////////////////////////////////////////////////////////////

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Key {
    AltArrow(Arrow),
    AltChar(char),
    Arrow(Arrow),
    Backspace,
    Char(char),
    Ctrl(char),
    CtrlArrow(Arrow),
    Del,
    End,
    Esc,
    Home,
    PageDown,
    PageUp,
    ShiftDown,
    ShiftUp,
    Tab,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Arrow {
    Left,
    Right,
    Up,
    Down,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Event {
    /// A single key input.
    Key(Key),

    /// Usually a paste.
    String(String),

    // These are not generated anymore. I needed these for a feature, but never really got around
    // implementing it.
    // FocusGained,
    // FocusLost,
    /// An unknown sequence of bytes (probably for a key combination that we don't care about).
    Unknown(Vec<u8>),
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Byte sequences of key presses we want to capture
////////////////////////////////////////////////////////////////////////////////////////////////////

byte_seq_parser! {
    parse_key_bytes -> Key, // Function name. Generation function type is
                            // fn(&[u8]) -> Option<(Key, usize)>.

    [27, 91, 49, 59, 51, 66] => Key::AltArrow(Arrow::Down),
    [27, 91, 49, 59, 51, 68] => Key::AltArrow(Arrow::Left),
    [27, 91, 49, 59, 51, 67] => Key::AltArrow(Arrow::Right),
    [27, 91, 49, 59, 51, 65] => Key::AltArrow(Arrow::Up),
    [27, 91, 66] => Key::Arrow(Arrow::Down),
    [27, 79, 66] => Key::Arrow(Arrow::Down),
    [27, 91, 68] => Key::Arrow(Arrow::Left),
    [27, 79, 68] => Key::Arrow(Arrow::Left),
    [27, 91, 67] => Key::Arrow(Arrow::Right),
    [27, 79, 67] => Key::Arrow(Arrow::Right),
    [27, 91, 65] => Key::Arrow(Arrow::Up),
    [27, 79, 65] => Key::Arrow(Arrow::Up),
    [27, 91, 49, 59, 53, 66] => Key::CtrlArrow(Arrow::Down),
    [27, 91, 49, 59, 53, 68] => Key::CtrlArrow(Arrow::Left),
    [27, 91, 49, 59, 53, 67] => Key::CtrlArrow(Arrow::Right),
    [27, 91, 49, 59, 53, 65] => Key::CtrlArrow(Arrow::Up),
    [27, 91, 51, 126] => Key::Del,
    [27, 91, 54, 126] => Key::PageDown,
    [27, 91, 53, 126] => Key::PageUp,
    [27, 91, 49, 59, 50, 65] => Key::ShiftUp,
    [27, 91, 49, 59, 50, 66] => Key::ShiftDown,
    [27, 91, 72] => Key::Home,
    [27, 91, 70] => Key::End,
    [27, 79, 72] => Key::Home,
    [27, 79, 70] => Key::End,
    [27, 91, 52, 126] => Key::End,
    [9] => Key::Tab,
    [127] => Key::Backspace,
    [1] => Key::Ctrl('a'),
    [5] => Key::Ctrl('e'),
    [23] => Key::Ctrl('w'),
    [11] => Key::Ctrl('k'),
    [4] => Key::Ctrl('d'),
    [3] => Key::Ctrl('c'),
    [17] => Key::Ctrl('q'),
    [16] => Key::Ctrl('p'),
    [14] => Key::Ctrl('n'),
    [21] => Key::Ctrl('u'),
    [24] => Key::Ctrl('x'),
}

// static XTERM_FOCUS_GAINED: [u8; 3] = [27, 91, 73];
// static XTERM_FOCUS_LOST: [u8; 3] = [27, 91, 79];

////////////////////////////////////////////////////////////////////////////////////////////////////

pub struct Input {
    /// Queue of events waiting to be polled.
    evs: VecDeque<Event>,

    /// Used when reading from stdin.
    buf: Vec<u8>,

    stdin: AsyncFd<RawFd>,
}

impl Input {
    /// Create a new input handler. Requires a `tokio` reactor to be running. Make sure to either
    /// set `stdin` to non-blocking mode or enable non-canonical input, otherwise the `Stream`
    /// implementation will not work. If you are using this with `termbox` it does the latter so no
    /// need to do anything. See module documentation for links on how to enable non-canonical
    /// input.
    pub fn new() -> Input {
        Input {
            evs: VecDeque::new(),
            buf: Vec::with_capacity(100),
            stdin: AsyncFd::new(libc::STDIN_FILENO).unwrap(),
        }
    }

    fn parse_buffer(&mut self) {
        let mut buf_slice: &[u8] = &self.buf;

        while !buf_slice.is_empty() {
            // Special treatment for 127 (backspace, 0x1B) and 13 ('\r', 0xD)
            let fst = buf_slice[0];
            let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
                parse_key_comb
            } else {
                parse_chars
            };

            match parse_fn(buf_slice) {
                Some((ev, used)) => {
                    buf_slice = &buf_slice[used..];
                    self.evs.push_back(ev);
                }
                None => {
                    self.evs.push_back(Event::Unknown(buf_slice.to_owned()));
                    break;
                }
            }
        }

        self.buf.clear();
    }
}

impl Stream for Input {
    type Item = std::io::Result<Event>;

    fn poll_next(mut self: Pin<&mut Input>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let self_: &mut Input = &mut *self;

        'main: loop {
            // Try to parse any bytes in the input buffer from the last poll
            self_.parse_buffer();

            // Yield pending events
            if let Some(ev) = self_.evs.pop_front() {
                return Poll::Ready(Some(Ok(ev)));
            }

            // Otherwise read stdin and loop if successful
            let mut poll_ret = self_.stdin.poll_read_ready(cx);
            loop {
                match poll_ret {
                    Poll::Ready(Ok(mut ready)) => {
                        // We will read stdin until it's empty, clear the readiness state now
                        // (before reading) to avoid a potential race that would occur if stdin
                        // becomes available again between the last read and clear_ready.
                        ready.clear_ready();

                        // Read stdin until it fails with EWOULDBLOCK or returns 0
                        let read_ret = read_stdin(&mut self_.buf);

                        // Poll again so that we will be notified when stdin becomes ready again
                        poll_ret = self_.stdin.poll_read_ready(cx);

                        match read_ret {
                            Ok(()) => {}
                            Err(err) => {
                                // NOTE: `poll_ret` is ignored here but I think that's OK?
                                // TODO: Contents of `buf` is ignored here until `poll_next` is
                                // called again.
                                let err =
                                    std::io::Error::from(err.as_errno().expect("Weird nix error"));
                                return Poll::Ready(Some(Err(err)));
                            }
                        }

                        if self_.buf.is_empty() {
                            continue;
                        } else {
                            continue 'main;
                        }
                    }
                    Poll::Ready(Err(err)) => {
                        return Poll::Ready(Some(Err(err)));
                    }
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
pub(crate) fn parse_single_event(buf: &[u8]) -> Event {
    let fst = buf[0];
    let parse_fn = if (fst < 32 && fst != 13) || fst == 127 {
        parse_key_comb
    } else {
        parse_chars
    };

    let (ev, used) = match parse_fn(buf) {
        Some((ev, used)) => (ev, used),
        None => (Event::Unknown(buf.to_owned()), buf.len()),
    };

    assert_eq!(buf.len(), used);

    ev
}

fn parse_chars(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    // Use a fast path for the common case: Single utf-8 character.
    // (TODO: What about other encodings?)

    utf8_char_len(buf[0]).map(|char_len| {
        if char_len as usize == buf.len() {
            // Fast path: single character
            let ev = Event::Key(Key::Char(get_utf8_char(buf, char_len)));
            (ev, char_len as usize)
        } else {
            // Probably a paste: allocate a string and collect chars
            let mut string = String::with_capacity(1000);
            let mut start_idx = 0;
            loop {
                if start_idx == buf.len() {
                    break;
                }
                match utf8_char_len(buf[start_idx]) {
                    Some(char_len) => {
                        string.push(get_utf8_char(&buf[start_idx..], char_len));
                        start_idx += char_len as usize;
                    }
                    None => {
                        break;
                    }
                }
            }
            let ev = Event::String(string);
            (ev, start_idx)
        }
    })
}

fn parse_key_comb(buf: &[u8]) -> Option<(Event, usize)> {
    debug_assert!(!buf.is_empty());

    if let Some((key, used)) = parse_key_bytes(buf) {
        return Some((Event::Key(key), used));
    }

    if buf[0] == 27 {
        // 0x1B, ESC
        // 27 not followed by anything is an actual ESC
        if buf.len() == 1 {
            return Some((Event::Key(Key::Esc), 1));
        }
        // otherwise it's probably alt + key
        else {
            debug_assert!(buf.len() >= 2);
            return utf8_char_len(buf[1]).map(|char_len| {
                let ev = Event::Key(Key::AltChar(get_utf8_char(&buf[1..], char_len)));
                (ev, char_len as usize + 1)
            });
        }
    }

    None
}

fn utf8_char_len(byte: u8) -> Option<u8> {
    if byte >> 7 == 0b0 {
        Some(1)
    } else if byte >> 5 == 0b110 {
        Some(2)
    } else if byte >> 4 == 0b1110 {
        Some(3)
    } else if byte >> 3 == 0b11110 {
        Some(4)
    } else {
        None
    }
}

fn get_utf8_char(buf: &[u8], len: u8) -> char {
    let codepoint: u32 = {
        if len == 1 {
            u32::from(buf[0] & 0b0111_1111)
        } else if len == 2 {
            ((u32::from(buf[0] & 0b0001_1111)) << 6) + (u32::from(buf[1] & 0b0011_1111))
        } else if len == 3 {
            ((u32::from(buf[0] & 0b0000_1111)) << 12)
                + ((u32::from(buf[1] & 0b0011_1111)) << 6)
                + (u32::from(buf[2] & 0b0011_1111))
        } else {
            debug_assert!(len == 4);
            ((u32::from(buf[0] & 0b0000_0111)) << 18)
                + ((u32::from(buf[1] & 0b0011_1111)) << 12)
                + ((u32::from(buf[2] & 0b0011_1111)) << 6)
                + (u32::from(buf[3] & 0b0011_1111))
        }
    };

    char::from_u32(codepoint).unwrap()
}

/// Read `stdin` until `read` fails with `EWOULDBLOCK` (happens in non-canonical mode, when `stdin`
/// is set to non-blocking mode) or returns 0 (happens in non-canonical mode when `VMIN` and
/// `VTIME` are 0). If you are using `term_input` with `termbox` then you don't need to set `stdin`
/// to non-blocking mode as `termbox` enables non-canonical mode.
///
/// See also:
///
/// - https://www.gnu.org/software/libc/manual/html_node/Canonical-or-Not.html
/// - https://www.gnu.org/software/libc/manual/html_node/Noncanonical-Input.html
pub fn read_stdin(buf: &mut Vec<u8>) -> Result<(), nix::Error> {
    loop {
        let old_len = buf.len();
        buf.reserve(100);
        unsafe {
            buf.set_len(old_len + 100);
        }

        match nix::unistd::read(libc::STDIN_FILENO, &mut buf[old_len..]) {
            Ok(n_read) => {
                unsafe { buf.set_len(old_len + n_read) };
                if n_read == 0 {
                    // We're in non-canonical mode, or stdin is closed. We can't distinguish the
                    // two here but I think it's fine to return OK when stdin is closed.
                    return Ok(());
                }
            }
            Err(err) => {
                unsafe { buf.set_len(old_len) };
                match err {
                    nix::Error::Sys(nix::errno::EWOULDBLOCK) => {
                        return Ok(());
                    }
                    _ => {
                        return Err(err);
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate elf;
extern crate getopts;

use getopts::Options;
use std::cmp;
use std::env;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::Write;
use std::mem;
use std::path::Path;
use std::slice;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct LoadInfo {
    version: u32,
    total_size: u32,
    entry_offset: u32,
    rel_data_offset: u32,
    rel_data_size: u32,
    text_offset: u32,
    text_size: u32,
    got_offset: u32,
    got_size: u32,
    data_offset: u32,
    data_size: u32,
    bss_mem_offset: u32,
    bss_size: u32,
    min_stack_len: u32,
    min_app_heap_len: u32,
    min_kernel_heap_len: u32,
    package_name_offset: u32,
    package_name_size: u32,
    checksum: u32,
}

impl fmt::Display for LoadInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "
            version: {:>8} {:>#10X}
         total_size: {:>8} {:>#10X}
       entry_offset: {:>8} {:>#10X}
    rel_data_offset: {:>8} {:>#10X}
      rel_data_size: {:>8} {:>#10X}
        text_offset: {:>8} {:>#10X}
          text_size: {:>8} {:>#10X}
         got_offset: {:>8} {:>#10X}
           got_size: {:>8} {:>#10X}
        data_offset: {:>8} {:>#10X}
          data_size: {:>8} {:>#10X}
     bss_mem_offset: {:>8} {:>#10X}
           bss_size: {:>8} {:>#10X}
      min_stack_len: {:>8} {:>#10X}
   min_app_heap_len: {:>8} {:>#10X}
min_kernel_heap_len: {:>8} {:>#10X}
package_name_offset: {:>8} {:>#10X}
  package_name_size: {:>8} {:>#10X}
           checksum: {:>8} {:>#10X}
",
        self.version, self.version,
        self.total_size, self.total_size,
        self.entry_offset, self.entry_offset,
        self.rel_data_offset, self.rel_data_offset,
        self.rel_data_size, self.rel_data_size,
        self.text_offset, self.text_offset,
        self.text_size, self.text_size,
        self.got_offset, self.got_offset,
        self.got_size, self.got_size,
        self.data_offset, self.data_offset,
        self.data_size, self.data_size,
        self.bss_mem_offset, self.bss_mem_offset,
        self.bss_size, self.bss_size,
        self.min_stack_len, self.min_stack_len,
        self.min_app_heap_len, self.min_app_heap_len,
        self.min_kernel_heap_len, self.min_kernel_heap_len,
        self.package_name_offset, self.package_name_offset,
        self.package_name_size, self.package_name_size,
        self.checksum, self.checksum,
        )
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "", "set output file name", "OUTFILE");
    opts.optopt("n", "", "set package name", "PACKAGE_NAME");
    opts.optflag("v", "verbose", "be verbose");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    let output = matches.opt_str("o");
    let package_name = matches.opt_str("n");
    let verbose = matches.opt_present("v");
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let path = Path::new(&input);
    let file = match elf::File::open_path(&path) {
        Ok(f) => f,
        Err(e) => panic!("Error: {:?}", e),
    };

    match output {
            None => {
                let mut out = io::stdout();
                do_work(&file, &mut out, package_name, verbose)
            }
            Some(name) => {
                match File::create(Path::new(&name)) {
                    Ok(mut f) => do_work(&file, &mut f, package_name, verbose),
                    Err(e) => panic!("Error: {:?}", e),
                }
            }
        }
        .expect("Failed to write output");
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [-o OUTFILE] FILE", program);
    print!("{}", opts.usage(&brief));
}

fn get_section<'a>(input: &'a elf::File, name: &str) -> elf::Section {
    match input.get_section(name) {
        Some(section) => {
            elf::Section {
                data: section.data.clone(),
                shdr: section.shdr.clone(),
            }
        }
        None => {
            elf::Section {
                data: Vec::new(),
                shdr: elf::types::SectionHeader {
                    name: String::from(name),
                    shtype: elf::types::SHT_NULL,
                    flags: elf::types::SHF_NONE,
                    addr: 0,
                    offset: 0,
                    size: 0,
                    link: 0,
                    info: 0,
                    addralign: 0,
                    entsize: 0,
                },
            }
        }
    }
}

unsafe fn as_byte_slice<'a, T: Copy>(input: &'a T) -> &'a [u8] {
    slice::from_raw_parts(input as *const T as *const u8, mem::size_of::<T>())
}

fn do_work(input: &elf::File,
           output: &mut Write,
           package_name: Option<String>,
           verbose: bool)
           -> io::Result<()> {
    let package_name = package_name.unwrap_or(String::new());
    let (rel_data_size, rel_data) = match input.sections
        .iter()
        .find(|section| section.shdr.name == ".rel.data".as_ref()) {
        Some(section) => (section.shdr.size, section.data.as_ref()),
        None => (0 as u64, &[] as &[u8]),

    };
    let text = get_section(input, ".text");
    let got = get_section(input, ".got");
    let data = get_section(input, ".data");
    let bss = get_section(input, ".bss");

    // For these, we only care about the length
    let stack_len = get_section(input, ".stack").data.len() as u32;
    let app_heap_len = get_section(input, ".app_heap").data.len() as u32;
    let kernel_heap_len = get_section(input, ".kernel_heap").data.len() as u32;

    let mut total_size = (mem::size_of::<LoadInfo>() + rel_data.len() + text.data.len() +
                          got.data.len() +
                          data.data.len() + package_name.len()) as u32;

    let pad = if total_size.count_ones() > 1 {
        let power2len = cmp::max(1 << (32 - total_size.leading_zeros()), 512);
        power2len - total_size
    } else {
        0
    };
    total_size = total_size + pad;

    let rel_data_offset = mem::size_of::<LoadInfo>() as u32;
    let text_offset = rel_data_offset + (rel_data_size as u32);
    let text_size = text.shdr.size as u32;
    let entry_offset = (input.ehdr.entry ^ 0x80000000) as u32 + text_offset;
    let got_offset = text_offset + text_size;
    let got_size = got.shdr.size as u32;
    let data_offset = got_offset + got_size;
    let data_size = data.shdr.size as u32;
    let package_name_offset = data_offset + data_size;
    let package_name_size = package_name.len() as u32;

    let load_info_version = 1;

    let load_info = LoadInfo {
        version: load_info_version,
        total_size: total_size,
        entry_offset: entry_offset,
        rel_data_offset: rel_data_offset,
        rel_data_size: rel_data_size as u32,
        text_offset: text_offset,
        text_size: text_size,
        got_offset: got_offset,
        got_size: got_size,
        data_offset: data_offset,
        data_size: data_size,
        bss_mem_offset: bss.shdr.addr as u32,
        bss_size: bss.shdr.size as u32,
        min_stack_len: stack_len,
        min_app_heap_len: app_heap_len,
        min_kernel_heap_len: kernel_heap_len,
        package_name_offset: package_name_offset,
        package_name_size: package_name_size,
        checksum: load_info_version ^ total_size ^ entry_offset ^ rel_data_offset ^
                  rel_data_size as u32 ^ text_offset ^ text_size ^ got_offset ^
                  got_size ^
                  data_offset ^ data_size ^ bss.shdr.addr as u32 ^
                  bss.shdr.size as u32 ^
                  stack_len ^ app_heap_len ^
                  kernel_heap_len ^ package_name_offset ^ package_name_size,
    };

    if verbose {
        print!("{}", load_info);
    }

    try!(output.write_all(unsafe { as_byte_slice(&load_info) }));
    try!(output.write_all(rel_data.as_ref()));
    try!(output.write_all(text.data.as_ref()));
    try!(output.write_all(got.data.as_ref()));
    try!(output.write_all(data.data.as_ref()));
    try!(output.write_all(package_name.as_ref()));

    let mut pad = pad as usize;
    let zero_buf = [0u8; 512];
    while pad > 0 {
        let amount_to_write = cmp::min(zero_buf.len(), pad);
        pad -= try!(output.write(&zero_buf[..amount_to_write]));
    }

    Ok(())
}
<commit_msg>rust fmt<commit_after>extern crate elf;
extern crate getopts;

use getopts::Options;
use std::cmp;
use std::env;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::Write;
use std::mem;
use std::path::Path;
use std::slice;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
struct LoadInfo {
    version: u32,
    total_size: u32,
    entry_offset: u32,
    rel_data_offset: u32,
    rel_data_size: u32,
    text_offset: u32,
    text_size: u32,
    got_offset: u32,
    got_size: u32,
    data_offset: u32,
    data_size: u32,
    bss_mem_offset: u32,
    bss_size: u32,
    min_stack_len: u32,
    min_app_heap_len: u32,
    min_kernel_heap_len: u32,
    package_name_offset: u32,
    package_name_size: u32,
    checksum: u32,
}

impl fmt::Display for LoadInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "
            version: {:>8} {:>#10X}
         total_size: {:>8} {:>#10X}
       entry_offset: {:>8} {:>#10X}
    rel_data_offset: {:>8} {:>#10X}
      rel_data_size: {:>8} {:>#10X}
        text_offset: {:>8} {:>#10X}
          text_size: {:>8} {:>#10X}
         got_offset: {:>8} {:>#10X}
           got_size: {:>8} {:>#10X}
        data_offset: {:>8} {:>#10X}
          data_size: {:>8} {:>#10X}
     bss_mem_offset: {:>8} {:>#10X}
           bss_size: {:>8} {:>#10X}
      min_stack_len: {:>8} {:>#10X}
   min_app_heap_len: {:>8} {:>#10X}
min_kernel_heap_len: {:>8} {:>#10X}
package_name_offset: {:>8} {:>#10X}
  package_name_size: {:>8} {:>#10X}
           checksum: {:>8} {:>#10X}
",
        self.version, self.version,
        self.total_size, self.total_size,
        self.entry_offset, self.entry_offset,
        self.rel_data_offset, self.rel_data_offset,
        self.rel_data_size, self.rel_data_size,
        self.text_offset, self.text_offset,
        self.text_size, self.text_size,
        self.got_offset, self.got_offset,
        self.got_size, self.got_size,
        self.data_offset, self.data_offset,
        self.data_size, self.data_size,
        self.bss_mem_offset, self.bss_mem_offset,
        self.bss_size, self.bss_size,
        self.min_stack_len, self.min_stack_len,
        self.min_app_heap_len, self.min_app_heap_len,
        self.min_kernel_heap_len, self.min_kernel_heap_len,
        self.package_name_offset, self.package_name_offset,
        self.package_name_size, self.package_name_size,
        self.checksum, self.checksum,
        )
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("o", "", "set output file name", "OUTFILE");
    opts.optopt("n", "", "set package name", "PACKAGE_NAME");
    opts.optflag("v", "verbose", "be verbose");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    let output = matches.opt_str("o");
    let package_name = matches.opt_str("n");
    let verbose = matches.opt_present("v");
    let input = if !matches.free.is_empty() {
        matches.free[0].clone()
    } else {
        print_usage(&program, opts);
        return;
    };

    let path = Path::new(&input);
    let file = match elf::File::open_path(&path) {
        Ok(f) => f,
        Err(e) => panic!("Error: {:?}", e),
    };

    match output {
            None => {
                let mut out = io::stdout();
                do_work(&file, &mut out, package_name, verbose)
            }
            Some(name) => {
                match File::create(Path::new(&name)) {
                    Ok(mut f) => do_work(&file, &mut f, package_name, verbose),
                    Err(e) => panic!("Error: {:?}", e),
                }
            }
        }
        .expect("Failed to write output");
}

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [-o OUTFILE] FILE", program);
    print!("{}", opts.usage(&brief));
}

fn get_section<'a>(input: &'a elf::File, name: &str) -> elf::Section {
    match input.get_section(name) {
        Some(section) => {
            elf::Section {
                data: section.data.clone(),
                shdr: section.shdr.clone(),
            }
        }
        None => {
            elf::Section {
                data: Vec::new(),
                shdr: elf::types::SectionHeader {
                    name: String::from(name),
                    shtype: elf::types::SHT_NULL,
                    flags: elf::types::SHF_NONE,
                    addr: 0,
                    offset: 0,
                    size: 0,
                    link: 0,
                    info: 0,
                    addralign: 0,
                    entsize: 0,
                },
            }
        }
    }
}

unsafe fn as_byte_slice<'a, T: Copy>(input: &'a T) -> &'a [u8] {
    slice::from_raw_parts(input as *const T as *const u8, mem::size_of::<T>())
}

fn do_work(input: &elf::File,
           output: &mut Write,
           package_name: Option<String>,
           verbose: bool)
           -> io::Result<()> {
    let package_name = package_name.unwrap_or(String::new());
    let (rel_data_size, rel_data) =
        match input
                  .sections
                  .iter()
                  .find(|section| section.shdr.name == ".rel.data".as_ref()) {
            Some(section) => (section.shdr.size, section.data.as_ref()),
            None => (0 as u64, &[] as &[u8]),

        };
    let text = get_section(input, ".text");
    let got = get_section(input, ".got");
    let data = get_section(input, ".data");
    let bss = get_section(input, ".bss");

    // For these, we only care about the length
    let stack_len = get_section(input, ".stack").data.len() as u32;
    let app_heap_len = get_section(input, ".app_heap").data.len() as u32;
    let kernel_heap_len = get_section(input, ".kernel_heap").data.len() as u32;

    let mut total_size = (mem::size_of::<LoadInfo>() + rel_data.len() + text.data.len() +
                          got.data.len() +
                          data.data.len() + package_name.len()) as u32;

    let pad = if total_size.count_ones() > 1 {
        let power2len = cmp::max(1 << (32 - total_size.leading_zeros()), 512);
        power2len - total_size
    } else {
        0
    };
    total_size = total_size + pad;

    let rel_data_offset = mem::size_of::<LoadInfo>() as u32;
    let text_offset = rel_data_offset + (rel_data_size as u32);
    let text_size = text.shdr.size as u32;
    let entry_offset = (input.ehdr.entry ^ 0x80000000) as u32 + text_offset;
    let got_offset = text_offset + text_size;
    let got_size = got.shdr.size as u32;
    let data_offset = got_offset + got_size;
    let data_size = data.shdr.size as u32;
    let package_name_offset = data_offset + data_size;
    let package_name_size = package_name.len() as u32;

    let load_info_version = 1;

    let load_info = LoadInfo {
        version: load_info_version,
        total_size: total_size,
        entry_offset: entry_offset,
        rel_data_offset: rel_data_offset,
        rel_data_size: rel_data_size as u32,
        text_offset: text_offset,
        text_size: text_size,
        got_offset: got_offset,
        got_size: got_size,
        data_offset: data_offset,
        data_size: data_size,
        bss_mem_offset: bss.shdr.addr as u32,
        bss_size: bss.shdr.size as u32,
        min_stack_len: stack_len,
        min_app_heap_len: app_heap_len,
        min_kernel_heap_len: kernel_heap_len,
        package_name_offset: package_name_offset,
        package_name_size: package_name_size,
        checksum: load_info_version ^ total_size ^ entry_offset ^ rel_data_offset ^
                  rel_data_size as u32 ^ text_offset ^ text_size ^ got_offset ^
                  got_size ^
                  data_offset ^ data_size ^ bss.shdr.addr as u32 ^
                  bss.shdr.size as u32 ^
                  stack_len ^ app_heap_len ^
                  kernel_heap_len ^ package_name_offset ^ package_name_size,
    };

    if verbose {
        print!("{}", load_info);
    }

    try!(output.write_all(unsafe { as_byte_slice(&load_info) }));
    try!(output.write_all(rel_data.as_ref()));
    try!(output.write_all(text.data.as_ref()));
    try!(output.write_all(got.data.as_ref()));
    try!(output.write_all(data.data.as_ref()));
    try!(output.write_all(package_name.as_ref()));

    let mut pad = pad as usize;
    let zero_buf = [0u8; 512];
    while pad > 0 {
        let amount_to_write = cmp::min(zero_buf.len(), pad);
        pad -= try!(output.write(&zero_buf[..amount_to_write]));
    }

    Ok(())
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;

use croaring::Bitmap;
use roaring::RoaringBitmap;
use test::{black_box, Bencher};

#[bench]
fn perf_comp_create_croaring(b: &mut Bencher) {
    b.iter(|| {
        let bitmap = Bitmap::create();

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_rust_roaring(b: &mut Bencher) {
    b.iter(|| {
        let bitmap: RoaringBitmap = RoaringBitmap::new();

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_and_add_one_croaring(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap = Bitmap::create();
        bitmap.add(black_box(1));

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_and_add_one_rust_roaring(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap: RoaringBitmap = RoaringBitmap::new();
        bitmap.insert(black_box(1));

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_add_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();

    b.iter(|| {
        bitmap.add(black_box(1));
        bitmap.add(black_box(10));
        bitmap.add(black_box(100));
        bitmap.add(black_box(1000));
        bitmap.add(black_box(10000));
        bitmap.add(black_box(100000));
        bitmap.add(black_box(1000000));
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_add_many_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();

    b.iter(|| {
        let int_slice = &[1, 10, 100, 1000, 10_000, 100_000, 1_000_000];
        bitmap.add_many(int_slice);
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_add_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();

    b.iter(|| {
        bitmap.insert(black_box(1));
        bitmap.insert(black_box(10));
        bitmap.insert(black_box(100));
        bitmap.insert(black_box(1000));
        bitmap.insert(black_box(10000));
        bitmap.insert(black_box(100000));
        bitmap.insert(black_box(1000000));
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_remove_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();
    bitmap.add(1);
    bitmap.add(10);
    bitmap.add(100);
    bitmap.add(1000);
    bitmap.add(10000);
    bitmap.add(100000);
    bitmap.add(1000000);

    b.iter(|| {
        bitmap.remove(black_box(1000000));
    });
}

#[bench]
fn perf_comp_remove_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();
    bitmap.insert(1);
    bitmap.insert(10);
    bitmap.insert(100);
    bitmap.insert(1000);
    bitmap.insert(10000);
    bitmap.insert(100000);
    bitmap.insert(1000000);

    b.iter(|| {
        bitmap.remove(black_box(1000000));
    });
}

#[bench]
fn perf_comp_contains_true_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();
    bitmap.add(1);

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_true_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();
    bitmap.insert(1);

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_false_croaring(b: &mut Bencher) {
    let bitmap = Bitmap::create();

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_false_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = RoaringBitmap::new();

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_cardinality_100000_croaring(b: &mut Bencher) {
    let bitmap: Bitmap = (1..100000).collect();

    b.iter(|| {
        black_box(bitmap.cardinality());
    });
}

#[bench]
fn perf_comp_cardinality_100000_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = (1..100000).collect();

    b.iter(|| {
        black_box(bitmap.len());
    });
}

#[bench]
fn perf_comp_and_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.and(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_and_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 & &bitmap2);
}

#[bench]
fn perf_comp_and_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.and_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_and_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.intersect_with(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_or_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.or(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_or_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 | &bitmap2);
}

#[bench]
fn perf_comp_or_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.or_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_or_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.union_with(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_xor_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.xor(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_xor_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 ^ &bitmap2);
}

#[bench]
fn perf_comp_xor_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.xor_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_xor_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.symmetric_difference_with(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_iter_croaring(b: &mut Bencher) {
    let bitmap: Bitmap = (1..10000).collect();

    b.iter(|| {
        let mut sum: u32 = 0;

        for (_, element) in bitmap.iter().enumerate() {
            sum += element;
        }

        assert_eq!(sum, 49995000);
    });
}

#[bench]
fn perf_comp_iter_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = (1..10000).collect();

    b.iter(|| {
        let mut sum: u32 = 0;

        for (_, element) in bitmap.iter().enumerate() {
            sum += element;
        }

        assert_eq!(sum, 49995000);
    });
}
<commit_msg>Fix benchmarks<commit_after>#![feature(test)]

extern crate test;

use croaring::Bitmap;
use roaring::RoaringBitmap;
use test::{black_box, Bencher};

#[bench]
fn perf_comp_create_croaring(b: &mut Bencher) {
    b.iter(|| {
        let bitmap = Bitmap::create();

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_rust_roaring(b: &mut Bencher) {
    b.iter(|| {
        let bitmap: RoaringBitmap = RoaringBitmap::new();

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_and_add_one_croaring(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap = Bitmap::create();
        bitmap.add(black_box(1));

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_create_and_add_one_rust_roaring(b: &mut Bencher) {
    b.iter(|| {
        let mut bitmap: RoaringBitmap = RoaringBitmap::new();
        bitmap.insert(black_box(1));

        black_box(bitmap);
    });
}

#[bench]
fn perf_comp_add_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();

    b.iter(|| {
        bitmap.add(black_box(1));
        bitmap.add(black_box(10));
        bitmap.add(black_box(100));
        bitmap.add(black_box(1000));
        bitmap.add(black_box(10000));
        bitmap.add(black_box(100000));
        bitmap.add(black_box(1000000));
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_add_many_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();

    b.iter(|| {
        let int_slice = &[1, 10, 100, 1000, 10_000, 100_000, 1_000_000];
        bitmap.add_many(int_slice);
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_add_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();

    b.iter(|| {
        bitmap.insert(black_box(1));
        bitmap.insert(black_box(10));
        bitmap.insert(black_box(100));
        bitmap.insert(black_box(1000));
        bitmap.insert(black_box(10000));
        bitmap.insert(black_box(100000));
        bitmap.insert(black_box(1000000));
    });
    black_box(bitmap);
}

#[bench]
fn perf_comp_remove_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();
    bitmap.add(1);
    bitmap.add(10);
    bitmap.add(100);
    bitmap.add(1000);
    bitmap.add(10000);
    bitmap.add(100000);
    bitmap.add(1000000);

    b.iter(|| {
        bitmap.remove(black_box(1000000));
    });
}

#[bench]
fn perf_comp_remove_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();
    bitmap.insert(1);
    bitmap.insert(10);
    bitmap.insert(100);
    bitmap.insert(1000);
    bitmap.insert(10000);
    bitmap.insert(100000);
    bitmap.insert(1000000);

    b.iter(|| {
        bitmap.remove(black_box(1000000));
    });
}

#[bench]
fn perf_comp_contains_true_croaring(b: &mut Bencher) {
    let mut bitmap = Bitmap::create();
    bitmap.add(1);

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_true_rust_roaring(b: &mut Bencher) {
    let mut bitmap: RoaringBitmap = RoaringBitmap::new();
    bitmap.insert(1);

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_false_croaring(b: &mut Bencher) {
    let bitmap = Bitmap::create();

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_contains_false_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = RoaringBitmap::new();

    b.iter(|| {
        bitmap.contains(black_box(1));
    });
}

#[bench]
fn perf_comp_cardinality_100000_croaring(b: &mut Bencher) {
    let bitmap: Bitmap = (1..100000).collect();

    b.iter(|| {
        black_box(bitmap.cardinality());
    });
}

#[bench]
fn perf_comp_cardinality_100000_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = (1..100000).collect();

    b.iter(|| {
        black_box(bitmap.len());
    });
}

#[bench]
fn perf_comp_and_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.and(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_and_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 & &bitmap2);
}

#[bench]
fn perf_comp_and_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.and_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_and_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1 &= black_box(&bitmap2);
    });
}

#[bench]
fn perf_comp_or_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.or(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_or_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 | &bitmap2);
}

#[bench]
fn perf_comp_or_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.or_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_or_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1 |= &bitmap2;
    });
}

#[bench]
fn perf_comp_xor_new_croaring(b: &mut Bencher) {
    let bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.xor(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_xor_new_rust_roaring(b: &mut Bencher) {
    let bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| &bitmap1 ^ &bitmap2);
}

#[bench]
fn perf_comp_xor_inplace_croaring(b: &mut Bencher) {
    let mut bitmap1: Bitmap = (1..100).collect();
    let bitmap2: Bitmap = (100..200).collect();

    b.iter(|| {
        bitmap1.xor_inplace(black_box(&bitmap2));
    });
}

#[bench]
fn perf_comp_xor_inplace_rust_roaring(b: &mut Bencher) {
    let mut bitmap1: RoaringBitmap = (1..100).collect();
    let bitmap2: RoaringBitmap = (100..200).collect();

    b.iter(|| {
        bitmap1 ^= black_box(&bitmap2);
    });
}

#[bench]
fn perf_comp_iter_croaring(b: &mut Bencher) {
    let bitmap: Bitmap = (1..10000).collect();

    b.iter(|| {
        let mut sum: u32 = 0;

        for (_, element) in bitmap.iter().enumerate() {
            sum += element;
        }

        assert_eq!(sum, 49995000);
    });
}

#[bench]
fn perf_comp_iter_rust_roaring(b: &mut Bencher) {
    let bitmap: RoaringBitmap = (1..10000).collect();

    b.iter(|| {
        let mut sum: u32 = 0;

        for (_, element) in bitmap.iter().enumerate() {
            sum += element;
        }

        assert_eq!(sum, 49995000);
    });
}
<|endoftext|>