<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 AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN";
const BLDR_URL_ENVVAR: &str = "HAB_BLDR_URL";
const CTL_SECRET_ENVVAR: &str= "HAB_CTL_SECRET";
const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

// fn set_env_var(env_var: &str) -> Result<()> {
//   let config = config::load()?;
//   if henv::var(env_var).is_err() {
//     if let Some(config_val) = match env_var {
//       AUTH_TOKEN_ENVVAR => config.auth_token,
//       _ => println!("uh oh, this shouldn't happen"),
//     };
//     println!("config var is: {}", config_val);
//     env::set_var(env_var, config_val)
//   }
// }

fn set_env_var(env_var: &str, val: String) -> Result<()> {
    debug!("Setting {}={} via config file", env_var, &val);
    env::set_var(env_var, val);
    Ok(())
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;
    if henv::var(AUTH_TOKEN_ENVVAR).is_err() {
        if let Some(auth_token) = config.auth_token {
          set_env_var(AUTH_TOKEN_ENVVAR, auth_token)?;
        }
    }

    if henv::var(BLDR_URL_ENVVAR).is_err() {
        if let Some(bldr_url) = config.bldr_url {
            set_env_var(BLDR_URL_ENVVAR, bldr_url)?;
        }
    }

    if henv::var(CTL_SECRET_ENVVAR).is_err() {
        if let Some(ctl_secret) = config.ctl_secret {
            set_env_var(CTL_SECRET_ENVVAR, ctl_secret)?;
        }
    }

    if henv::var(ORIGIN_ENVVAR).is_err() {
        if let Some(default_origin) = config.origin {
            set_env_var(ORIGIN_ENVVAR, default_origin)?;
        }
    }

    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>Further abstract out set_env_var_from_config to reduce copy/pasted code and make code more obvious<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 AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN";
const BLDR_URL_ENVVAR: &str = "HAB_BLDR_URL";
const CTL_SECRET_ENVVAR: &str= "HAB_CTL_SECRET";
const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

fn set_env_var_from_config(env_var: &str, config_val: Option<String>) {
    if henv::var(env_var).is_err() {
        if let Some(val) = config_val {
            debug!("Setting {}={} via config file", env_var, val);
            env::set_var(env_var, val);
        }
    }
}

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    let config = config::load()?;

    set_env_var_from_config(AUTH_TOKEN_ENVVAR, config.auth_token);
    set_env_var_from_config(BLDR_URL_ENVVAR, config.bldr_url);
    set_env_var_from_config(CTL_SECRET_ENVVAR, config.ctl_secret);
    set_env_var_from_config(ORIGIN_ENVVAR, config.origin);

    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>use std::fmt;
// use std::iter;
use std::slice;
use std::ops;
use std::collections::HashMap;
use error::*;

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum Operator {
    Add,
    Sub,
    Mul,
    Div,
}

#[cfg_attr(rustfmt, rustfmt_skip)]
impl fmt::Display for Operator {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Operator::*;
        write!(f, "{}", match *self {
            Add => "+",
            Sub => "-",
            Mul => "*",
            Div => "/",
        })
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum Arity {
    Nullary,
    Unary,
    Binary,
    Ternary,
    Multiary,  // 2+ arguments
}

impl fmt::Display for Arity {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Arity::*;
        let text = match *self {
            Nullary => "0",
            Unary   => "1",
            Binary  => "2",
            Ternary => "3",
            Multiary => "at least 2",
        };

        write!(f, "{} args", text)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Atom {
    Op(Operator),
    // Bool(bool),
    Int(i64),
    // Float(f64),
    // Str(String),
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Atom::*;
        match *self {
            Op(ref t) => write!(f, "{}", t),
            // Bool(ref t) => write!(f, "{}", t),
            Int(ref t) => write!(f, "{}", t),
            // Float(ref t) => write!(f, "{}", t),
            // Str(ref t) => write!(f, "{}", t),
        }
    }
}

impl From<Operator> for Atom {
    fn from(x: Operator) -> Self {
        Atom::Op(x)
    }
}

impl From<i64> for Atom {
    fn from(x: i64) -> Self {
        Atom::Int(x)
    }
}

impl ops::Add<Atom> for Atom {
    type Output = Atom;
    fn add(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i + j),
            _ => panic!(),
        }
    }
}

impl ops::Sub<Atom> for Atom {
    type Output = Atom;
    fn sub(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i - j),
            _ => panic!(),
        }
    }
}

impl ops::Mul<Atom> for Atom {
    type Output = Atom;
    fn mul(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i * j),
            _ => panic!(),
        }
    }
}

impl ops::Div<Atom> for Atom {
    type Output = Atom;
    fn div(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i / j),
            _ => panic!(),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Value {
    Atom(Atom),
    List(List),
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Value::*;
        match *self {
            Atom(ref t) => write!(f, "{}", t),
            List(ref t) => write!(f, "{}", t),
        }
    }
}

impl<T> From<T> for Value where T: Into<Atom> {
    fn from(v: T) -> Self {
        Value::Atom(v.into())
    }
}

impl ops::Add<Value> for Value {
    type Output = Value;
    fn add(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i + j),
            _ => panic!(),
        }
    }
}

impl ops::Sub<Value> for Value {
    type Output = Value;
    fn sub(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i - j),
            _ => panic!(),
        }
    }
}

impl ops::Mul<Value> for Value {
    type Output = Value;
    fn mul(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i * j),
            _ => panic!(),
        }
    }
}

impl ops::Div<Value> for Value {
    type Output = Value;
    fn div(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i / j),
            _ => panic!(),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Node {
    Atom (Atom),
    List (List),
    Expr (Expr),
}

impl Node {
    fn eval(&self) -> Result<Value> {
        use self::Node::*;
        match *self {
            Atom(a) => Ok(Value::Atom(a)),
            List(ref l) => Ok(Value::List(l.clone())),
            Expr(ref e) => e.eval(),
        }
    }
}

impl<T> From<T> for Node where T: Into<Atom> {
    fn from(v: T) -> Self {
        Node::Atom(v.into())
    }
}

// impl<T> From<T> for Node where T: Into<List> {
//     fn from(l: T) -> Self {
//         Node::List(l.into())
//     }
// }
//
// impl<T> From<T> for Node where T: Into<Expr> {
//     fn from(e: T) -> Self {
//         Node::Expr(e.into())
//     }
// }

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Node::*;
        match *self {
            Atom(ref a) => write!(f, "{}", a),
            List(ref l) => write!(f, "{}", l),
            Expr(ref e) => write!(f, "{}", e),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct List {
    inner: Vec<Node>
}

impl List {
    pub fn new() -> Self {
        List { inner: Vec::new() }
    }
}

impl From<Vec<Node>> for List {
    fn from(v: Vec<Node>) -> Self {
        List { inner: v }
    }
}

impl fmt::Display for List {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let elements = self.inner.iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(" ");
        write!(f, "( {} )", elements)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Expr {
    op: Operator,
    args: Vec<Node>
}

impl Expr {
    pub fn new(op: Operator, args: Vec<Node>) -> Self {
        Expr { op: op, args: args }
    }

    pub fn op(&self) -> Operator {
        self.op
    }

    pub fn args(&self) -> slice::Iter<Node> {
        self.args.iter()
    }

    pub fn num_args(&self) -> usize {
        self.args.len()
    }

    pub fn expect_arity(&self, arity: Arity) -> bool {
        use self::Arity::*;
        let num_args = self.num_args();
        match arity {
            Nullary  => num_args == 0,
            Unary    => num_args == 1,
            Binary   => num_args == 2,
            Ternary  => num_args == 3,
            Multiary => num_args >= 2,
        }
    }

    pub fn eval(&self) -> Result<Value> {
        // use self::Node::*;
        // use self::Atom::*;
        use self::Operator::*;
        use self::Arity::*;

        // Arity = number of arguments accepted
        let mut arity = HashMap::new();

        arity.insert(Add, Multiary);
        arity.insert(Sub, Binary);
        arity.insert(Mul, Multiary);
        arity.insert(Div, Binary);

        let op = self.op();

        let expected_arity = match arity.get(&op) {
            Some(e) => e,
            None => return Err(format!("operator {} not found", op).into())
        };

        // Check if number of arguments match expected number
        if !self.expect_arity(*expected_arity) {
            return Err(format!(
                "{} operator expected {} arguments, but received {}",
                op, expected_arity, self.num_args()
            ).into())
        }

        let mut args = self.args().map(Node::eval);

        // Return the first expr that cannot eval correctly
        if let Some(err) = args.find(|x| x.is_err()) {
            err
        } else {
            let mut eval_args = args.map(Result::unwrap);
            match op {
                Add => Ok(eval_args.fold(0.into(), |acc, x| acc + x)),
                Sub => {
                    let a = eval_args.next().unwrap();
                    let b = eval_args.next().unwrap();
                    Ok(a - b)
                },
                Mul => Ok(eval_args.fold(1.into(), |acc, x| acc * x)),
                Div => {
                    let a = eval_args.next().unwrap();
                    let b = eval_args.next().unwrap();
                    Ok(a / b)
                }
            }
        }
    }
}

impl fmt::Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let elements = self.args.iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(" ");
        write!(f, "( {} {} )", self.op, elements)
    }
}
<commit_msg>Fix eval<commit_after>use std::fmt;
// use std::iter;
use std::slice;
use std::ops;
use std::collections::HashMap;
use error::*;

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum Operator {
    Add,
    Sub,
    Mul,
    Div,
}

#[cfg_attr(rustfmt, rustfmt_skip)]
impl fmt::Display for Operator {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Operator::*;
        write!(f, "{}", match *self {
            Add => "+",
            Sub => "-",
            Mul => "*",
            Div => "/",
        })
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum Arity {
    Nullary,
    Unary,
    Binary,
    Ternary,
    Multiary,  // 2+ arguments
}

impl fmt::Display for Arity {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Arity::*;
        let text = match *self {
            Nullary => "0",
            Unary   => "1",
            Binary  => "2",
            Ternary => "3",
            Multiary => "at least 2",
        };

        write!(f, "{} args", text)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Atom {
    Op(Operator),
    // Bool(bool),
    Int(i64),
    // Float(f64),
    // Str(String),
}

impl fmt::Display for Atom {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Atom::*;
        match *self {
            Op(ref t) => write!(f, "{}", t),
            // Bool(ref t) => write!(f, "{}", t),
            Int(ref t) => write!(f, "{}", t),
            // Float(ref t) => write!(f, "{}", t),
            // Str(ref t) => write!(f, "{}", t),
        }
    }
}

impl From<Operator> for Atom {
    fn from(x: Operator) -> Self {
        Atom::Op(x)
    }
}

impl From<i64> for Atom {
    fn from(x: i64) -> Self {
        Atom::Int(x)
    }
}

impl ops::Add<Atom> for Atom {
    type Output = Atom;
    fn add(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i + j),
            _ => panic!(),
        }
    }
}

impl ops::Sub<Atom> for Atom {
    type Output = Atom;
    fn sub(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i - j),
            _ => panic!(),
        }
    }
}

impl ops::Mul<Atom> for Atom {
    type Output = Atom;
    fn mul(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i * j),
            _ => panic!(),
        }
    }
}

impl ops::Div<Atom> for Atom {
    type Output = Atom;
    fn div(self, rhs: Atom) -> Self::Output {
        use self::Atom::*;
        match (self, rhs) {
            (Int(i), Int(j)) => Int(i / j),
            _ => panic!(),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Value {
    Atom(Atom),
    List(List),
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Value::*;
        match *self {
            Atom(ref t) => write!(f, "{}", t),
            List(ref t) => write!(f, "{}", t),
        }
    }
}

impl<T> From<T> for Value where T: Into<Atom> {
    fn from(v: T) -> Self {
        Value::Atom(v.into())
    }
}

impl ops::Add<Value> for Value {
    type Output = Value;
    fn add(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i + j),
            _ => panic!(),
        }
    }
}

impl ops::Sub<Value> for Value {
    type Output = Value;
    fn sub(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i - j),
            _ => panic!(),
        }
    }
}

impl ops::Mul<Value> for Value {
    type Output = Value;
    fn mul(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i * j),
            _ => panic!(),
        }
    }
}

impl ops::Div<Value> for Value {
    type Output = Value;
    fn div(self, rhs: Value) -> Self::Output {
        use self::Value::*;
        match (self, rhs) {
            (Atom(i), Atom(j)) => Atom(i / j),
            _ => panic!(),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Node {
    Atom (Atom),
    List (List),
    Expr (Expr),
}

impl Node {
    fn eval(&self) -> Result<Value> {
        use self::Node::*;
        match *self {
            Atom(a) => Ok(Value::Atom(a)),
            List(ref l) => Ok(Value::List(l.clone())),
            Expr(ref e) => e.eval(),
        }
    }
}

impl<T> From<T> for Node where T: Into<Atom> {
    fn from(v: T) -> Self {
        Node::Atom(v.into())
    }
}

// impl<T> From<T> for Node where T: Into<List> {
//     fn from(l: T) -> Self {
//         Node::List(l.into())
//     }
// }
//
// impl<T> From<T> for Node where T: Into<Expr> {
//     fn from(e: T) -> Self {
//         Node::Expr(e.into())
//     }
// }

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Node::*;
        match *self {
            Atom(ref a) => write!(f, "{}", a),
            List(ref l) => write!(f, "{}", l),
            Expr(ref e) => write!(f, "{}", e),
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct List {
    inner: Vec<Node>
}

impl List {
    pub fn new() -> Self {
        List { inner: Vec::new() }
    }
}

impl From<Vec<Node>> for List {
    fn from(v: Vec<Node>) -> Self {
        List { inner: v }
    }
}

impl fmt::Display for List {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let elements = self.inner.iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(" ");
        write!(f, "( {} )", elements)
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Expr {
    op: Operator,
    args: Vec<Node>
}

impl Expr {
    pub fn new(op: Operator, args: Vec<Node>) -> Self {
        Expr { op: op, args: args }
    }

    pub fn op(&self) -> Operator {
        self.op
    }

    pub fn args(&self) -> slice::Iter<Node> {
        self.args.iter()
    }

    pub fn num_args(&self) -> usize {
        self.args.len()
    }

    pub fn expect_arity(&self, arity: Arity) -> bool {
        use self::Arity::*;
        let num_args = self.num_args();
        match arity {
            Nullary  => num_args == 0,
            Unary    => num_args == 1,
            Binary   => num_args == 2,
            Ternary  => num_args == 3,
            Multiary => num_args >= 2,
        }
    }

    pub fn eval(&self) -> Result<Value> {
        // use self::Node::*;
        // use self::Atom::*;
        use self::Operator::*;
        use self::Arity::*;

        // Arity = number of arguments accepted
        let mut arity = HashMap::new();

        arity.insert(Add, Multiary);
        arity.insert(Sub, Binary);
        arity.insert(Mul, Multiary);
        arity.insert(Div, Binary);

        let op = self.op();

        let expected_arity = match arity.get(&op) {
            Some(e) => e,
            None => return Err(format!("operator {} not found", op).into())
        };

        // Check if number of arguments match expected number
        if !self.expect_arity(*expected_arity) {
            return Err(format!(
                "{} operator expected {} arguments, but received {}",
                op, expected_arity, self.num_args()
            ).into())
        }

        let evaled_args: Result<Vec<_>> = self.args().map(Node::eval).collect();

        // Return the first expr that cannot eval correctly
        match evaled_args {
            Err(e) => Err(e),
            Ok(v) => {
                let mut eval_args = v.into_iter();
                match op {
                    Add => Ok(eval_args.fold(0.into(), |acc, x| acc + x)),
                    Sub => {
                        let a = eval_args.next().unwrap();
                        let b = eval_args.next().unwrap();
                        Ok(a - b)
                    },
                    Mul => Ok(eval_args.fold(1.into(), |acc, x| acc * x)),
                    Div => {
                        let a = eval_args.next().unwrap();
                        let b = eval_args.next().unwrap();
                        Ok(a / b)
                    }
                }

            }
        }
    }
}

impl fmt::Display for Expr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let elements = self.args.iter()
            .map(ToString::to_string)
            .collect::<Vec<_>>()
            .join(" ");
        write!(f, "( {} {} )", self.op, elements)
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use env_logger::Env;
use log::{error, info};
use relay_compiler::{compiler::Compiler, config::Config};
use std::{path::PathBuf, sync::Arc};
use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(
    name = "Relay Compiler",
    about = "Compiler to produce Relay generated files."
)]
struct Opt {
    /// Compile and watch for changes
    #[structopt(long, short)]
    watch: bool,

    /// Path to the compiler config file
    config: PathBuf,
}

#[tokio::main]
async fn main() {
    env_logger::from_env(Env::default().default_filter_or("info")).init();

    let opt = Opt::from_args();

    let config = match Config::load(opt.config) {
        Ok(config) => config,
        Err(err) => {
            error!("{}", err);
            std::process::exit(1);
        }
    };

    let compiler = Compiler::new(config, Arc::new(common::NoopPerfLogger));

    if opt.watch {
        if let Err(err) = compiler.watch().await {
            error!("{}", err);
            std::process::exit(1);
        }
    } else {
        match compiler.compile().await {
            Ok(_compiler_state) => {
                info!("Done");
            }
            Err(_) => {
                std::process::exit(1);
            }
        }
    }
}
<commit_msg>log error when compiler fails (#3286)<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use env_logger::Env;
use log::{error, info};
use relay_compiler::{compiler::Compiler, config::Config};
use std::{path::PathBuf, sync::Arc};
use structopt::StructOpt;

#[derive(StructOpt)]
#[structopt(
    name = "Relay Compiler",
    about = "Compiler to produce Relay generated files."
)]
struct Opt {
    /// Compile and watch for changes
    #[structopt(long, short)]
    watch: bool,

    /// Path to the compiler config file
    config: PathBuf,
}

#[tokio::main]
async fn main() {
    env_logger::from_env(Env::default().default_filter_or("info")).init();

    let opt = Opt::from_args();

    let config = match Config::load(opt.config) {
        Ok(config) => config,
        Err(err) => {
            error!("{}", err);
            std::process::exit(1);
        }
    };

    let compiler = Compiler::new(config, Arc::new(common::NoopPerfLogger));

    if opt.watch {
        if let Err(err) = compiler.watch().await {
            error!("{}", err);
            std::process::exit(1);
        }
    } else {
        match compiler.compile().await {
            Ok(_compiler_state) => {
                info!("Done");
            }
            Err(err) => {
                error!("{}", err);
                std::process::exit(1);
            }
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![allow(unsafe_code)]

//! A drop-in replacement for string_cache, but backed by Gecko `nsAtom`s.

use gecko_bindings::bindings::Gecko_AddRefAtom;
use gecko_bindings::bindings::Gecko_Atomize;
use gecko_bindings::bindings::Gecko_Atomize16;
use gecko_bindings::bindings::Gecko_ReleaseAtom;
use gecko_bindings::structs::{nsAtom, nsAtom_AtomKind, nsDynamicAtom, nsStaticAtom};
use nsstring::{nsAString, nsStr};
use precomputed_hash::PrecomputedHash;
use std::{mem, slice, str};
use std::borrow::{Borrow, Cow};
use std::char::{self, DecodeUtf16};
use std::fmt::{self, Write};
use std::hash::{Hash, Hasher};
use std::iter::Cloned;
use std::ops::Deref;
use style_traits::SpecifiedValueInfo;

#[macro_use]
#[allow(improper_ctypes, non_camel_case_types, missing_docs)]
pub mod atom_macro {
    include!(concat!(env!("OUT_DIR"), "/gecko/atom_macro.rs"));
}

#[macro_use]
pub mod namespace;

pub use self::namespace::{Namespace, WeakNamespace};

macro_rules! local_name {
    ($s:tt) => {
        atom!($s)
    };
}

/// A strong reference to a Gecko atom.
#[derive(Eq, PartialEq)]
pub struct Atom(*mut WeakAtom);

/// An atom *without* a strong reference.
///
/// Only usable as `&'a WeakAtom`,
/// where `'a` is the lifetime of something that holds a strong reference to that atom.
pub struct WeakAtom(nsAtom);

/// A BorrowedAtom for Gecko is just a weak reference to a `nsAtom`, that
/// hasn't been bumped.
pub type BorrowedAtom<'a> = &'a WeakAtom;

impl Deref for Atom {
    type Target = WeakAtom;

    #[inline]
    fn deref(&self) -> &WeakAtom {
        unsafe { &*self.0 }
    }
}

impl PrecomputedHash for Atom {
    #[inline]
    fn precomputed_hash(&self) -> u32 {
        self.get_hash()
    }
}

impl Borrow<WeakAtom> for Atom {
    #[inline]
    fn borrow(&self) -> &WeakAtom {
        self
    }
}

impl Eq for WeakAtom {}
impl PartialEq for WeakAtom {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        let weak: *const WeakAtom = self;
        let other: *const WeakAtom = other;
        weak == other
    }
}

unsafe impl Send for Atom {}
unsafe impl Sync for Atom {}
unsafe impl Sync for WeakAtom {}

impl WeakAtom {
    /// Construct a `WeakAtom` from a raw `nsAtom`.
    #[inline]
    pub unsafe fn new<'a>(atom: *const nsAtom) -> &'a mut Self {
        &mut *(atom as *mut WeakAtom)
    }

    /// Clone this atom, bumping the refcount if the atom is not static.
    #[inline]
    pub fn clone(&self) -> Atom {
        unsafe { Atom::from_raw(self.as_ptr()) }
    }

    /// Get the atom hash.
    #[inline]
    pub fn get_hash(&self) -> u32 {
        self.0.mHash
    }

    /// Get the atom as a slice of utf-16 chars.
    #[inline]
    pub fn as_slice(&self) -> &[u16] {
        let string = if self.is_static() {
            let atom_ptr = self.as_ptr() as *const nsStaticAtom;
            let string_offset = unsafe { (*atom_ptr).mStringOffset };
            let string_offset = -(string_offset as isize);
            let u8_ptr = atom_ptr as *const u8;
            // It is safe to use offset() here because both addresses are within
            // the same struct, e.g. mozilla::detail::gGkAtoms.
            unsafe { u8_ptr.offset(string_offset) as *const u16 }
        } else {
            let atom_ptr = self.as_ptr() as *const nsDynamicAtom;
            // Dynamic atom chars are stored at the end of the object.
            unsafe { atom_ptr.offset(1) as *const u16 }
        };
        unsafe { slice::from_raw_parts(string, self.len() as usize) }
    }

    // NOTE: don't expose this, since it's slow, and easy to be misused.
    fn chars(&self) -> DecodeUtf16<Cloned<slice::Iter<u16>>> {
        char::decode_utf16(self.as_slice().iter().cloned())
    }

    /// Execute `cb` with the string that this atom represents.
    ///
    /// Find alternatives to this function when possible, please, since it's
    /// pretty slow.
    pub fn with_str<F, Output>(&self, cb: F) -> Output
    where
        F: FnOnce(&str) -> Output,
    {
        let mut buffer: [u8; 64] = unsafe { mem::uninitialized() };

        // The total string length in utf16 is going to be less than or equal
        // the slice length (each utf16 character is going to take at least one
        // and at most 2 items in the utf16 slice).
        //
        // Each of those characters will take at most four bytes in the utf8
        // one. Thus if the slice is less than 64 / 4 (16) we can guarantee that
        // we'll decode it in place.
        let owned_string;
        let len = self.len();
        let utf8_slice = if len <= 16 {
            let mut total_len = 0;

            for c in self.chars() {
                let c = c.unwrap_or(char::REPLACEMENT_CHARACTER);
                let utf8_len = c.encode_utf8(&mut buffer[total_len..]).len();
                total_len += utf8_len;
            }

            let slice = unsafe { str::from_utf8_unchecked(&buffer[..total_len]) };
            debug_assert_eq!(slice, String::from_utf16_lossy(self.as_slice()));
            slice
        } else {
            owned_string = String::from_utf16_lossy(self.as_slice());
            &*owned_string
        };

        cb(utf8_slice)
    }

    /// Returns whether this atom is static.
    #[inline]
    pub fn is_static(&self) -> bool {
        unsafe { (*self.as_ptr()).mKind() == nsAtom_AtomKind::Static as u32 }
    }

    /// Returns the length of the atom string.
    #[inline]
    pub fn len(&self) -> u32 {
        unsafe { (*self.as_ptr()).mLength() }
    }

    /// Returns whether this atom is the empty string.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the atom as a mutable pointer.
    #[inline]
    pub fn as_ptr(&self) -> *mut nsAtom {
        let const_ptr: *const nsAtom = &self.0;
        const_ptr as *mut nsAtom
    }

    /// Convert this atom to ASCII lower-case
    pub fn to_ascii_lowercase(&self) -> Atom {
        let slice = self.as_slice();
        match slice
            .iter()
            .position(|&char16| (b'A' as u16) <= char16 && char16 <= (b'Z' as u16))
        {
            None => self.clone(),
            Some(i) => {
                let mut buffer: [u16; 64] = unsafe { mem::uninitialized() };
                let mut vec;
                let mutable_slice = if let Some(buffer_prefix) = buffer.get_mut(..slice.len()) {
                    buffer_prefix.copy_from_slice(slice);
                    buffer_prefix
                } else {
                    vec = slice.to_vec();
                    &mut vec
                };
                for char16 in &mut mutable_slice[i..] {
                    if *char16 <= 0x7F {
                        *char16 = (*char16 as u8).to_ascii_lowercase() as u16
                    }
                }
                Atom::from(&*mutable_slice)
            },
        }
    }

    /// Return whether two atoms are ASCII-case-insensitive matches
    pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool {
        if self == other {
            return true;
        }

        let a = self.as_slice();
        let b = other.as_slice();
        a.len() == b.len() && a.iter().zip(b).all(|(&a16, &b16)| {
            if a16 <= 0x7F && b16 <= 0x7F {
                (a16 as u8).eq_ignore_ascii_case(&(b16 as u8))
            } else {
                a16 == b16
            }
        })
    }

    /// Return whether this atom is an ASCII-case-insensitive match for the given string
    pub fn eq_str_ignore_ascii_case(&self, other: &str) -> bool {
        self.chars()
            .map(|r| r.map(|c: char| c.to_ascii_lowercase()))
            .eq(other.chars().map(|c: char| Ok(c.to_ascii_lowercase())))
    }
}

impl fmt::Debug for WeakAtom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        write!(w, "Gecko WeakAtom({:p}, {})", self, self)
    }
}

impl fmt::Display for WeakAtom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        for c in self.chars() {
            w.write_char(c.unwrap_or(char::REPLACEMENT_CHARACTER))?
        }
        Ok(())
    }
}

impl Atom {
    /// Execute a callback with the atom represented by `ptr`.
    pub unsafe fn with<F, R>(ptr: *mut nsAtom, callback: F) -> R
    where
        F: FnOnce(&Atom) -> R,
    {
        let atom = Atom(WeakAtom::new(ptr));
        let ret = callback(&atom);
        mem::forget(atom);
        ret
    }

    /// Creates an atom from an static atom pointer without checking in release
    /// builds.
    ///
    /// Right now it's only used by the atom macro, and ideally it should keep
    /// that way, now we have sugar for is_static, creating atoms using
    /// Atom::from_raw should involve almost no overhead.
    #[inline]
    pub unsafe fn from_static(ptr: *mut nsStaticAtom) -> Self {
        let atom = Atom(ptr as *mut WeakAtom);
        debug_assert!(
            atom.is_static(),
            "Called from_static for a non-static atom!"
        );
        atom
    }

    /// Creates an atom from an atom pointer.
    #[inline(always)]
    pub unsafe fn from_raw(ptr: *mut nsAtom) -> Self {
        let atom = Atom(ptr as *mut WeakAtom);
        if !atom.is_static() {
            Gecko_AddRefAtom(ptr);
        }
        atom
    }

    /// Creates an atom from a dynamic atom pointer that has already had AddRef
    /// called on it.
    #[inline]
    pub unsafe fn from_addrefed(ptr: *mut nsAtom) -> Self {
        assert!(!ptr.is_null());
        Atom(WeakAtom::new(ptr))
    }

    /// Convert this atom into an addrefed nsAtom pointer.
    #[inline]
    pub fn into_addrefed(self) -> *mut nsAtom {
        let ptr = self.as_ptr();
        mem::forget(self);
        ptr
    }
}

impl Hash for Atom {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        state.write_u32(self.get_hash());
    }
}

impl Hash for WeakAtom {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        state.write_u32(self.get_hash());
    }
}

impl Clone for Atom {
    #[inline(always)]
    fn clone(&self) -> Atom {
        unsafe { Atom::from_raw(self.as_ptr()) }
    }
}

impl Drop for Atom {
    #[inline]
    fn drop(&mut self) {
        if !self.is_static() {
            unsafe {
                Gecko_ReleaseAtom(self.as_ptr());
            }
        }
    }
}

impl Default for Atom {
    #[inline]
    fn default() -> Self {
        atom!("")
    }
}

impl fmt::Debug for Atom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        write!(w, "Gecko Atom({:p}, {})", self.0, self)
    }
}

impl fmt::Display for Atom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        unsafe { (&*self.0).fmt(w) }
    }
}

impl<'a> From<&'a str> for Atom {
    #[inline]
    fn from(string: &str) -> Atom {
        debug_assert!(string.len() <= u32::max_value() as usize);
        unsafe {
            Atom(WeakAtom::new(Gecko_Atomize(
                string.as_ptr() as *const _,
                string.len() as u32,
            )))
        }
    }
}

impl<'a> From<&'a [u16]> for Atom {
    #[inline]
    fn from(slice: &[u16]) -> Atom {
        Atom::from(&*nsStr::from(slice))
    }
}

impl<'a> From<&'a nsAString> for Atom {
    #[inline]
    fn from(string: &nsAString) -> Atom {
        unsafe { Atom(WeakAtom::new(Gecko_Atomize16(string))) }
    }
}

impl<'a> From<Cow<'a, str>> for Atom {
    #[inline]
    fn from(string: Cow<'a, str>) -> Atom {
        Atom::from(&*string)
    }
}

impl From<String> for Atom {
    #[inline]
    fn from(string: String) -> Atom {
        Atom::from(&*string)
    }
}

malloc_size_of_is_0!(Atom);

impl SpecifiedValueInfo for Atom {}
<commit_msg>style: Remove unused BorrowedAtom.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![allow(unsafe_code)]

//! A drop-in replacement for string_cache, but backed by Gecko `nsAtom`s.

use gecko_bindings::bindings::Gecko_AddRefAtom;
use gecko_bindings::bindings::Gecko_Atomize;
use gecko_bindings::bindings::Gecko_Atomize16;
use gecko_bindings::bindings::Gecko_ReleaseAtom;
use gecko_bindings::structs::{nsAtom, nsAtom_AtomKind, nsDynamicAtom, nsStaticAtom};
use nsstring::{nsAString, nsStr};
use precomputed_hash::PrecomputedHash;
use std::{mem, slice, str};
use std::borrow::{Borrow, Cow};
use std::char::{self, DecodeUtf16};
use std::fmt::{self, Write};
use std::hash::{Hash, Hasher};
use std::iter::Cloned;
use std::ops::Deref;
use style_traits::SpecifiedValueInfo;

#[macro_use]
#[allow(improper_ctypes, non_camel_case_types, missing_docs)]
pub mod atom_macro {
    include!(concat!(env!("OUT_DIR"), "/gecko/atom_macro.rs"));
}

#[macro_use]
pub mod namespace;

pub use self::namespace::{Namespace, WeakNamespace};

macro_rules! local_name {
    ($s:tt) => {
        atom!($s)
    };
}

/// A strong reference to a Gecko atom.
#[derive(Eq, PartialEq)]
pub struct Atom(*mut WeakAtom);

/// An atom *without* a strong reference.
///
/// Only usable as `&'a WeakAtom`,
/// where `'a` is the lifetime of something that holds a strong reference to that atom.
pub struct WeakAtom(nsAtom);

impl Deref for Atom {
    type Target = WeakAtom;

    #[inline]
    fn deref(&self) -> &WeakAtom {
        unsafe { &*self.0 }
    }
}

impl PrecomputedHash for Atom {
    #[inline]
    fn precomputed_hash(&self) -> u32 {
        self.get_hash()
    }
}

impl Borrow<WeakAtom> for Atom {
    #[inline]
    fn borrow(&self) -> &WeakAtom {
        self
    }
}

impl Eq for WeakAtom {}
impl PartialEq for WeakAtom {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        let weak: *const WeakAtom = self;
        let other: *const WeakAtom = other;
        weak == other
    }
}

unsafe impl Send for Atom {}
unsafe impl Sync for Atom {}
unsafe impl Sync for WeakAtom {}

impl WeakAtom {
    /// Construct a `WeakAtom` from a raw `nsAtom`.
    #[inline]
    pub unsafe fn new<'a>(atom: *const nsAtom) -> &'a mut Self {
        &mut *(atom as *mut WeakAtom)
    }

    /// Clone this atom, bumping the refcount if the atom is not static.
    #[inline]
    pub fn clone(&self) -> Atom {
        unsafe { Atom::from_raw(self.as_ptr()) }
    }

    /// Get the atom hash.
    #[inline]
    pub fn get_hash(&self) -> u32 {
        self.0.mHash
    }

    /// Get the atom as a slice of utf-16 chars.
    #[inline]
    pub fn as_slice(&self) -> &[u16] {
        let string = if self.is_static() {
            let atom_ptr = self.as_ptr() as *const nsStaticAtom;
            let string_offset = unsafe { (*atom_ptr).mStringOffset };
            let string_offset = -(string_offset as isize);
            let u8_ptr = atom_ptr as *const u8;
            // It is safe to use offset() here because both addresses are within
            // the same struct, e.g. mozilla::detail::gGkAtoms.
            unsafe { u8_ptr.offset(string_offset) as *const u16 }
        } else {
            let atom_ptr = self.as_ptr() as *const nsDynamicAtom;
            // Dynamic atom chars are stored at the end of the object.
            unsafe { atom_ptr.offset(1) as *const u16 }
        };
        unsafe { slice::from_raw_parts(string, self.len() as usize) }
    }

    // NOTE: don't expose this, since it's slow, and easy to be misused.
    fn chars(&self) -> DecodeUtf16<Cloned<slice::Iter<u16>>> {
        char::decode_utf16(self.as_slice().iter().cloned())
    }

    /// Execute `cb` with the string that this atom represents.
    ///
    /// Find alternatives to this function when possible, please, since it's
    /// pretty slow.
    pub fn with_str<F, Output>(&self, cb: F) -> Output
    where
        F: FnOnce(&str) -> Output,
    {
        let mut buffer: [u8; 64] = unsafe { mem::uninitialized() };

        // The total string length in utf16 is going to be less than or equal
        // the slice length (each utf16 character is going to take at least one
        // and at most 2 items in the utf16 slice).
        //
        // Each of those characters will take at most four bytes in the utf8
        // one. Thus if the slice is less than 64 / 4 (16) we can guarantee that
        // we'll decode it in place.
        let owned_string;
        let len = self.len();
        let utf8_slice = if len <= 16 {
            let mut total_len = 0;

            for c in self.chars() {
                let c = c.unwrap_or(char::REPLACEMENT_CHARACTER);
                let utf8_len = c.encode_utf8(&mut buffer[total_len..]).len();
                total_len += utf8_len;
            }

            let slice = unsafe { str::from_utf8_unchecked(&buffer[..total_len]) };
            debug_assert_eq!(slice, String::from_utf16_lossy(self.as_slice()));
            slice
        } else {
            owned_string = String::from_utf16_lossy(self.as_slice());
            &*owned_string
        };

        cb(utf8_slice)
    }

    /// Returns whether this atom is static.
    #[inline]
    pub fn is_static(&self) -> bool {
        unsafe { (*self.as_ptr()).mKind() == nsAtom_AtomKind::Static as u32 }
    }

    /// Returns the length of the atom string.
    #[inline]
    pub fn len(&self) -> u32 {
        unsafe { (*self.as_ptr()).mLength() }
    }

    /// Returns whether this atom is the empty string.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the atom as a mutable pointer.
    #[inline]
    pub fn as_ptr(&self) -> *mut nsAtom {
        let const_ptr: *const nsAtom = &self.0;
        const_ptr as *mut nsAtom
    }

    /// Convert this atom to ASCII lower-case
    pub fn to_ascii_lowercase(&self) -> Atom {
        let slice = self.as_slice();
        match slice
            .iter()
            .position(|&char16| (b'A' as u16) <= char16 && char16 <= (b'Z' as u16))
        {
            None => self.clone(),
            Some(i) => {
                let mut buffer: [u16; 64] = unsafe { mem::uninitialized() };
                let mut vec;
                let mutable_slice = if let Some(buffer_prefix) = buffer.get_mut(..slice.len()) {
                    buffer_prefix.copy_from_slice(slice);
                    buffer_prefix
                } else {
                    vec = slice.to_vec();
                    &mut vec
                };
                for char16 in &mut mutable_slice[i..] {
                    if *char16 <= 0x7F {
                        *char16 = (*char16 as u8).to_ascii_lowercase() as u16
                    }
                }
                Atom::from(&*mutable_slice)
            },
        }
    }

    /// Return whether two atoms are ASCII-case-insensitive matches
    pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool {
        if self == other {
            return true;
        }

        let a = self.as_slice();
        let b = other.as_slice();
        a.len() == b.len() && a.iter().zip(b).all(|(&a16, &b16)| {
            if a16 <= 0x7F && b16 <= 0x7F {
                (a16 as u8).eq_ignore_ascii_case(&(b16 as u8))
            } else {
                a16 == b16
            }
        })
    }

    /// Return whether this atom is an ASCII-case-insensitive match for the given string
    pub fn eq_str_ignore_ascii_case(&self, other: &str) -> bool {
        self.chars()
            .map(|r| r.map(|c: char| c.to_ascii_lowercase()))
            .eq(other.chars().map(|c: char| Ok(c.to_ascii_lowercase())))
    }
}

impl fmt::Debug for WeakAtom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        write!(w, "Gecko WeakAtom({:p}, {})", self, self)
    }
}

impl fmt::Display for WeakAtom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        for c in self.chars() {
            w.write_char(c.unwrap_or(char::REPLACEMENT_CHARACTER))?
        }
        Ok(())
    }
}

impl Atom {
    /// Execute a callback with the atom represented by `ptr`.
    pub unsafe fn with<F, R>(ptr: *mut nsAtom, callback: F) -> R
    where
        F: FnOnce(&Atom) -> R,
    {
        let atom = Atom(WeakAtom::new(ptr));
        let ret = callback(&atom);
        mem::forget(atom);
        ret
    }

    /// Creates an atom from an static atom pointer without checking in release
    /// builds.
    ///
    /// Right now it's only used by the atom macro, and ideally it should keep
    /// that way, now we have sugar for is_static, creating atoms using
    /// Atom::from_raw should involve almost no overhead.
    #[inline]
    pub unsafe fn from_static(ptr: *mut nsStaticAtom) -> Self {
        let atom = Atom(ptr as *mut WeakAtom);
        debug_assert!(
            atom.is_static(),
            "Called from_static for a non-static atom!"
        );
        atom
    }

    /// Creates an atom from an atom pointer.
    #[inline(always)]
    pub unsafe fn from_raw(ptr: *mut nsAtom) -> Self {
        let atom = Atom(ptr as *mut WeakAtom);
        if !atom.is_static() {
            Gecko_AddRefAtom(ptr);
        }
        atom
    }

    /// Creates an atom from a dynamic atom pointer that has already had AddRef
    /// called on it.
    #[inline]
    pub unsafe fn from_addrefed(ptr: *mut nsAtom) -> Self {
        assert!(!ptr.is_null());
        Atom(WeakAtom::new(ptr))
    }

    /// Convert this atom into an addrefed nsAtom pointer.
    #[inline]
    pub fn into_addrefed(self) -> *mut nsAtom {
        let ptr = self.as_ptr();
        mem::forget(self);
        ptr
    }
}

impl Hash for Atom {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        state.write_u32(self.get_hash());
    }
}

impl Hash for WeakAtom {
    fn hash<H>(&self, state: &mut H)
    where
        H: Hasher,
    {
        state.write_u32(self.get_hash());
    }
}

impl Clone for Atom {
    #[inline(always)]
    fn clone(&self) -> Atom {
        unsafe { Atom::from_raw(self.as_ptr()) }
    }
}

impl Drop for Atom {
    #[inline]
    fn drop(&mut self) {
        if !self.is_static() {
            unsafe {
                Gecko_ReleaseAtom(self.as_ptr());
            }
        }
    }
}

impl Default for Atom {
    #[inline]
    fn default() -> Self {
        atom!("")
    }
}

impl fmt::Debug for Atom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        write!(w, "Gecko Atom({:p}, {})", self.0, self)
    }
}

impl fmt::Display for Atom {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        unsafe { (&*self.0).fmt(w) }
    }
}

impl<'a> From<&'a str> for Atom {
    #[inline]
    fn from(string: &str) -> Atom {
        debug_assert!(string.len() <= u32::max_value() as usize);
        unsafe {
            Atom(WeakAtom::new(Gecko_Atomize(
                string.as_ptr() as *const _,
                string.len() as u32,
            )))
        }
    }
}

impl<'a> From<&'a [u16]> for Atom {
    #[inline]
    fn from(slice: &[u16]) -> Atom {
        Atom::from(&*nsStr::from(slice))
    }
}

impl<'a> From<&'a nsAString> for Atom {
    #[inline]
    fn from(string: &nsAString) -> Atom {
        unsafe { Atom(WeakAtom::new(Gecko_Atomize16(string))) }
    }
}

impl<'a> From<Cow<'a, str>> for Atom {
    #[inline]
    fn from(string: Cow<'a, str>) -> Atom {
        Atom::from(&*string)
    }
}

impl From<String> for Atom {
    #[inline]
    fn from(string: String) -> Atom {
        Atom::from(&*string)
    }
}

malloc_size_of_is_0!(Atom);

impl SpecifiedValueInfo for Atom {}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use std::collections::BTreeSet;
use std::fmt::{self, Display};
use parse::NextPrev;

#[derive(Clone, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Range(pub char, pub char);

impl fmt::Display for Range {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({}, {})", self.0, self.1)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Set(BTreeSet<Range>);

impl fmt::Display for Set {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let Set(ref set) = *self;
        let len = BTreeSet::len(set);

        for (count, s) in set.iter().enumerate() {
            if count < len - 1 { try!(write!(f, "{}, ", s)) }
            else { return write!(f, "{}", s) }
        }

        Ok(())
    }
}

impl Set {
    pub fn new() -> Self { Set(BTreeSet::new()) }
    pub fn insert(&mut self, value: Range) {
        let mut ret = BTreeSet::new();
        // value is a complete subset of one of the other ranges.
        let mut subset = false;

        // Borrowing self blocks later operation. Add a new scope.
        {   let Set(ref set) = *self;

            let Range(mut min_val, mut max_val) = value;
            if min_val > max_val { panic!("First value cannot be greater than the second.") }

            // Loop over set adding old disjoint pieces and supersets back.
            // When partially overlapped, expand value to the union. At the
            // end, insert union after it has been fully expanded.
            for &Range(min, max) in &*set {
                // value overlaps at the beginning.
                if min_val < min && max_val >= min && max_val < max { max_val = max }
                // value overlaps at the end.
                else if min_val > min && min_val <= max && max_val > max { min_val = min }
                // value is entirely contained between min and max. Insert original
                // into new array because new is a subset.
                else if min_val > min && max_val < max {
                    ret.insert(Range(min, max));
                    subset = true;
                }
                // value is a superset to the current so don't add current.
                else if min_val <= min && max_val >= max {}
                // value is disjoint with current so add current.
                else { ret.insert(Range(min, max)); }
            }

            // Insert value only when it's not a subset.
            if !subset { ret.insert(Range(min_val, max_val)); }
        }

        *self = Set(ret);
    }
    pub fn remove(&mut self, value: Range) {
        let mut ret = BTreeSet::new();

        // Borrowing self blocks later modification. Make a new scope to contain it.
        {   let Set(ref set) = *self;

            let Range(min_val, max_val) = value;
            if min_val > max_val { panic!("First value cannot be greater than the second.") }

            // Loop over set inserting whatever doesn't intersect.
            for &Range(min, max) in &*set {
                // value overlaps at the beginning.
                if min_val <= min && max_val >= min && max_val < max { ret.insert(Range(max_val.next(), max)); }
                // value overlaps at the end.
                else if min_val > min && min_val <= max && max_val >= max { ret.insert(Range(min, min_val.prev())); }
                // value is entirely contained between min and max. Split set
                // into two pieces.
                else if min_val > min && max_val < max {
                    ret.insert(Range(min, min_val.prev()));
                    ret.insert(Range(max_val.next(), max));

                    // Current piece was a superset so value cannot be anywhere else.
                    break;
                // value is a superset to the current so don't add current.
                } else if min_val <= min && max_val >= max {}
                // value is disjoint with current so add current.
                else { ret.insert(Range(min, max)); }
            }
        }

        *self = Set(ret)
    }
    pub fn union(&mut self, value: Self) {
        for x in value.0 { self.insert(x) }
    }
}

#[cfg(test)]
mod test {
    use super::{Range, Set};

    fn generate(vec: Vec<(char, char)>) -> Set {
        let mut set = Set::new();

        for (a, b) in vec {
            set.insert(Range(a, b));
        }

        set
    }
    #[test]
    fn insert_partial_overlap() {
        let set1 = generate(vec![('3', '5'), ('4', '6')]);
        let set2 = generate(vec![('3', '5'), ('1', '4')]);

        let other1 = generate(vec![('3', '6')]);
        let other2 = generate(vec![('1', '5')]);

        assert_eq!(other1, set1);
        assert_eq!(other2, set2);
    }
    #[test]
    fn insert_subset() {
        let set   = generate(vec![('1', '5'), ('2', '3')]);
        let other = generate(vec![('1', '5')]);

        assert_eq!(other, set);
    }
    #[test]
    fn insert_superset() {
        let set   = generate(vec![('3', '5'), ('2', '6')]);
        let other = generate(vec![('2', '6')]);

        assert_eq!(other, set);
    }
    #[test]
    #[should_panic]
    fn panic_on_decreasing_order() {
        generate(vec![('5', '3')]);
    }
    #[test]
    fn set_union() {
        let mut set1 = generate(vec![('1', '4'), ('5', '6')]);
        let set2 =     generate(vec![('0', '5'), ('8', '9')]);

        set1.union(set2);

        let other = generate(vec![('0', '6'), ('8', '9')]);
        assert_eq!(set1, other);
    }
    #[test]
    fn remove_partial_overlap() {
        let mut set1 = generate(vec![('5', '9')]);
        set1.remove(Range('3', '6')); // beginnings different.

        let mut set2 = generate(vec![('5', '9')]);
        set2.remove(Range('5', '6')); // beginnings same.

        let mut set3 = generate(vec![('2', '6')]);
        set3.remove(Range('5', '9')); // end different.

        let mut set4 = generate(vec![('2', '6')]);
        set4.remove(Range('5', '6')); // end same.

        let other12 = generate(vec![('7', '9')]);
        let other34 = generate(vec![('2', '4')]);

        assert_eq!(set1, other12);
        assert_eq!(set2, other12);
        assert_eq!(set3, other34);
        assert_eq!(set4, other34);
    }
    #[test]
    fn remove_subset() {
        let mut set = generate(vec![('1', '9')]);
        set.remove(Range('3', '6'));

        let other = generate(vec![('1', '2'), ('7', '9')]);

        assert_eq!(set, other);
    }
    #[test]
    fn remove_superset() {
        let mut set = generate(vec![('5', '6')]);
        set.remove(Range('3', '8'));

        let other = generate(vec![]);
        assert_eq!(set, other);
    }
    #[test]
    fn remove_disjoint() {
        let mut set = generate(vec![('2', '3')]);
        set.remove(Range('6', '8'));

        let other = generate(vec![('2', '3')]);
        assert_eq!(set, other);
    }
}
<commit_msg>Correct subset insertion with same endpoints. Add tests.<commit_after>#![allow(dead_code)]

use std::collections::BTreeSet;
use std::fmt::{self, Display};
use parse::NextPrev;

#[derive(Clone, Debug, PartialOrd, Ord, PartialEq, Eq)]
pub struct Range(pub char, pub char);

impl fmt::Display for Range {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({}, {})", self.0, self.1)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Set(BTreeSet<Range>);

impl fmt::Display for Set {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let Set(ref set) = *self;
        let len = BTreeSet::len(set);

        for (count, s) in set.iter().enumerate() {
            if count < len - 1 { try!(write!(f, "{}, ", s)) }
            else { return write!(f, "{}", s) }
        }

        Ok(())
    }
}

impl Set {
    pub fn new() -> Self { Set(BTreeSet::new()) }
    pub fn insert(&mut self, value: Range) {
        let mut ret = BTreeSet::new();
        // value is a complete subset of one of the other ranges.
        let mut subset = false;

        // Borrowing self blocks later operation. Add a new scope.
        {   let Set(ref set) = *self;

            let Range(mut min_val, mut max_val) = value;
            if min_val > max_val { panic!("First value cannot be greater than the second.") }

            // Loop over set adding old disjoint pieces and supersets back.
            // When partially overlapped, expand value to the union. At the
            // end, insert union after it has been fully expanded.
            for &Range(min, max) in &*set {
                // value overlaps at the beginning.
                if min_val < min && max_val >= min && max_val < max { max_val = max }
                // value overlaps at the end.
                else if min_val > min && min_val <= max && max_val > max { min_val = min }
                // value is entirely contained between min and max. Insert original
                // into new array because new is a subset.
                else if min_val >= min && max_val <= max {
                    ret.insert(Range(min, max));
                    subset = true;
                }
                // value is a superset to the current so don't add current.
                else if min_val < min && max_val > max {}
                // value is disjoint with current so add current.
                else { ret.insert(Range(min, max)); }
            }

            // Insert value only when it's not a subset.
            if !subset { ret.insert(Range(min_val, max_val)); }
        }

        *self = Set(ret);
    }
    pub fn remove(&mut self, value: Range) {
        let mut ret = BTreeSet::new();

        // Borrowing self blocks later modification. Make a new scope to contain it.
        {   let Set(ref set) = *self;

            let Range(min_val, max_val) = value;
            if min_val > max_val { panic!("First value cannot be greater than the second.") }

            // Loop over set inserting whatever doesn't intersect.
            for &Range(min, max) in &*set {
                // value overlaps at the beginning.
                if min_val <= min && max_val >= min && max_val < max { ret.insert(Range(max_val.next(), max)); }
                // value overlaps at the end.
                else if min_val > min && min_val <= max && max_val >= max { ret.insert(Range(min, min_val.prev())); }
                // value is entirely contained between min and max. Split set
                // into two pieces.
                else if min_val > min && max_val < max {
                    ret.insert(Range(min, min_val.prev()));
                    ret.insert(Range(max_val.next(), max));

                    // Current piece was a superset so value cannot be anywhere else.
                    break;
                // value is a superset to the current so don't add current.
                } else if min_val <= min && max_val >= max {}
                // value is disjoint with current so add current.
                else { ret.insert(Range(min, max)); }
            }
        }

        *self = Set(ret)
    }
    pub fn union(&mut self, value: Self) {
        for x in value.0 { self.insert(x) }
    }
}

#[cfg(test)]
mod test {
    use super::{Range, Set};

    fn generate(vec: Vec<(char, char)>) -> Set {
        let mut set = Set::new();

        for (a, b) in vec {
            set.insert(Range(a, b));
        }

        set
    }
    #[test]
    fn insert_partial_overlap() {
        let set1 = generate(vec![('3', '5'), ('4', '6')]);
        let set2 = generate(vec![('3', '5'), ('1', '4')]);

        let other1 = generate(vec![('3', '6')]);
        let other2 = generate(vec![('1', '5')]);

        assert_eq!(other1, set1);
        assert_eq!(other2, set2);
    }
    #[test]
    fn insert_subset() {
        let set1 = generate(vec![('1', '5'), ('2', '3')]); // Complete overlap.
        let set2 = generate(vec![('1', '5'), ('1', '3')]); // Left is exact.
        let set3 = generate(vec![('1', '5'), ('2', '5')]); // Right is exact.
        let set4 = generate(vec![('1', '5'), ('1', '5')]); // Both are exact.

        let other = generate(vec![('1', '5')]);

        assert_eq!(set1, other);
        assert_eq!(set2, other);
        assert_eq!(set3, other);
        assert_eq!(set4, other);
    }
    #[test]
    fn insert_superset() {
        let set   = generate(vec![('3', '5'), ('2', '6')]);
        let other = generate(vec![('2', '6')]);

        assert_eq!(other, set);
    }
    #[test]
    #[should_panic]
    fn panic_on_decreasing_order() {
        generate(vec![('5', '3')]);
    }
    #[test]
    fn set_union() {
        let mut set1 = generate(vec![('1', '4'), ('5', '6')]);
        let set2 =     generate(vec![('0', '5'), ('8', '9')]);

        set1.union(set2);

        let other = generate(vec![('0', '6'), ('8', '9')]);
        assert_eq!(set1, other);
    }
    #[test]
    fn remove_partial_overlap() {
        let mut set1 = generate(vec![('5', '9')]);
        set1.remove(Range('3', '6')); // beginnings different.

        let mut set2 = generate(vec![('5', '9')]);
        set2.remove(Range('5', '6')); // beginnings same.

        let mut set3 = generate(vec![('2', '6')]);
        set3.remove(Range('5', '9')); // end different.

        let mut set4 = generate(vec![('2', '6')]);
        set4.remove(Range('5', '6')); // end same.

        let other12 = generate(vec![('7', '9')]);
        let other34 = generate(vec![('2', '4')]);

        assert_eq!(set1, other12);
        assert_eq!(set2, other12);
        assert_eq!(set3, other34);
        assert_eq!(set4, other34);
    }
    #[test]
    fn remove_subset() {
        let mut set = generate(vec![('1', '9')]);
        set.remove(Range('3', '6'));

        let other = generate(vec![('1', '2'), ('7', '9')]);

        assert_eq!(set, other);
    }
    #[test]
    fn remove_superset() {
        let mut set = generate(vec![('5', '6')]);
        set.remove(Range('3', '8'));

        let other = generate(vec![]);
        assert_eq!(set, other);
    }
    #[test]
    fn remove_disjoint() {
        let mut set = generate(vec![('2', '3')]);
        set.remove(Range('6', '8'));

        let other = generate(vec![('2', '3')]);
        assert_eq!(set, other);
    }
}
<|endoftext|>


<commit_before>use {Color, Result};
use point::{Format, ScanDirection};
use std::io::{Cursor, ErrorKind, Read, Write};

/// A raw point.
#[derive(Clone, Copy, Debug, Default)]
#[allow(missing_docs)]
pub struct Point {
    pub x: i32,
    pub y: i32,
    pub z: i32,
    pub intensity: u16,
    pub flags: u8,
    pub classification: u8,
    pub scan_angle_rank: i8,
    pub user_data: u8,
    pub point_source_id: u16,
    pub gps_time: Option<f64>,
    pub color: Option<Color>,
}

impl Point {
    /// Reads a raw point.
    ///
    /// If there are exactly zero bytes left in the `Read`, then this function returns `Ok(None)`.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::{Seek, SeekFrom};
    /// use std::fs::File;
    /// use las::raw::Point;
    /// let mut file = File::open("tests/data/autzen.las").unwrap();
    /// file.seek(SeekFrom::Start(1994)).unwrap();
    /// let point = Point::read_from(file, 1.into()).unwrap();
    /// ```
    pub fn read_from<R: Read>(mut read: R, format: Format) -> Result<Option<Point>> {
        use byteorder::{LittleEndian, ReadBytesExt};
        let byte = match read.read_u8() {
            Ok(byte) => byte,
            Err(err) => {
                match err.kind() {
                    ErrorKind::UnexpectedEof => return Ok(None),
                    _ => return Err(err.into()),
                }
            }
        };
        let mut next_three = [0; 3];
        read.read_exact(&mut next_three)?;
        let mut cursor = Cursor::new([byte, next_three[0], next_three[1], next_three[2]]);
        let x = cursor.read_i32::<LittleEndian>()?;
        let y = read.read_i32::<LittleEndian>()?;
        let z = read.read_i32::<LittleEndian>()?;
        let intensity = read.read_u16::<LittleEndian>()?;
        let flags = read.read_u8()?;
        let classification = read.read_u8()?;
        let scan_angle_rank = read.read_i8()?;
        let user_data = read.read_u8()?;
        let point_source_id = read.read_u16::<LittleEndian>()?;
        let gps_time = if format.has_gps_time() {
            Some(read.read_f64::<LittleEndian>()?)
        } else {
            None
        };
        let color = if format.has_color() {
            let red = read.read_u16::<LittleEndian>()?;
            let green = read.read_u16::<LittleEndian>()?;
            let blue = read.read_u16::<LittleEndian>()?;
            Some(Color::new(red, green, blue))
        } else {
            None
        };
        // TODO read extra bytes
        Ok(Some(Point {
            x: x,
            y: y,
            z: z,
            intensity: intensity,
            flags: flags,
            classification: classification,
            scan_angle_rank: scan_angle_rank,
            user_data: user_data,
            point_source_id: point_source_id,
            gps_time: gps_time,
            color: color,
        }))
    }

    /// Calculates the return number from the flag byte.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 1, ..Default::default() };
    /// assert_eq!(1, point.return_number());
    /// ```
    pub fn return_number(&self) -> u8 {
        self.flags & 7
    }

    /// Calculates the number of returns from the flag byte.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 8, ..Default::default() };
    /// assert_eq!(1, point.number_of_returns());
    /// ```
    pub fn number_of_returns(&self) -> u8 {
        (self.flags & 56) >> 3
    }

    /// Returns the scan direction as determined by the scan direction flag.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// use las::point::ScanDirection;
    /// let point = Point { flags: 64, ..Default::default() };
    /// assert_eq!(ScanDirection::Positive, point.scan_direction());
    /// ```
    pub fn scan_direction(&self) -> ScanDirection {
        if (self.flags & 64) == 64 {
            ScanDirection::Positive
        } else {
            ScanDirection::Negative
        }
    }

    /// Returns true if the flags indicate that this point is the edge of a flight line.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 128, ..Default::default() };
    /// assert!(point.edge_of_flight_line());
    /// ```
    pub fn edge_of_flight_line(&self) -> bool {
        (self.flags & 128) == 128
    }

    /// Returns true if this point is synthetic.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 32, ..Default::default() };
    /// assert!(point.synthetic());
    pub fn synthetic(&self) -> bool {
        (self.classification & 32) == 32
    }

    /// Returns true if this point is a key point.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 64, ..Default::default() };
    /// assert!(point.key_point());
    pub fn key_point(&self) -> bool {
        (self.classification & 64) == 64
    }

    /// Returns true if this point is withheld.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 128, ..Default::default() };
    /// assert!(point.withheld());
    pub fn withheld(&self) -> bool {
        (self.classification & 128) == 128
    }

    /// Writes a raw pont.
    ///
    /// # Examples
    ///
    /// `Write` implements `WriteRawPoint`.
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::raw::Point;
    /// let mut cursor = Cursor::new(Vec::new());
    /// let point = Point::default();
    /// point.write_to(cursor, 0.into()).unwrap();
    /// ```
    pub fn write_to<W: Write>(&self, mut write: W, format: Format) -> Result<()> {
        use byteorder::{LittleEndian, WriteBytesExt};
        use Error;

        write.write_i32::<LittleEndian>(self.x)?;
        write.write_i32::<LittleEndian>(self.y)?;
        write.write_i32::<LittleEndian>(self.z)?;
        write.write_u16::<LittleEndian>(self.intensity)?;
        write.write_u8(self.flags)?;
        write.write_u8(self.classification)?;
        write.write_i8(self.scan_angle_rank)?;
        write.write_u8(self.user_data)?;
        write.write_u16::<LittleEndian>(self.point_source_id)?;
        if format.has_gps_time() {
            if let Some(gps_time) = self.gps_time {
                write.write_f64::<LittleEndian>(gps_time)?;
            } else {
                return Err(Error::MissingGpsTime(format, *self));
            }
        }
        if format.has_color() {
            if let Some(color) = self.color {
                write.write_u16::<LittleEndian>(color.red)?;
                write.write_u16::<LittleEndian>(color.green)?;
                write.write_u16::<LittleEndian>(color.blue)?;
            } else {
                return Err(Error::MissingColor(format, *self));
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use point::{Format, ScanDirection};
    use std::io::Cursor;

    #[test]
    fn return_number() {
        assert_eq!(
            0,
            Point {
                flags: 0,
                ..Default::default()
            }.return_number()
        );
        assert_eq!(
            7,
            Point {
                flags: 7,
                ..Default::default()
            }.return_number()
        );
        assert_eq!(
            0,
            Point {
                flags: 8,
                ..Default::default()
            }.return_number()
        );
    }

    #[test]
    fn number_of_returns() {
        assert_eq!(
            0,
            Point {
                flags: 0,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            1,
            Point {
                flags: 8,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            7,
            Point {
                flags: 56,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            0,
            Point {
                flags: 64,
                ..Default::default()
            }.number_of_returns()
        );
    }

    #[test]
    fn scan_direction() {
        assert_eq!(
            ScanDirection::Negative,
            Point { ..Default::default() }.scan_direction()
        );
        assert_eq!(
            ScanDirection::Positive,
            Point {
                flags: 64,
                ..Default::default()
            }.scan_direction()
        );
    }

    #[test]
    fn edge_of_flight_line() {
        assert!(!Point { ..Default::default() }.edge_of_flight_line());
        assert!(
            Point {
                flags: 128,
                ..Default::default()
            }.edge_of_flight_line()
        );
    }

    #[test]
    fn write_without_gps_time() {
        let point = Point { ..Default::default() };
        let write = Cursor::new(Vec::new());
        assert!(point.write_to(write, Format::from(1)).is_err());
    }

    #[test]
    fn read_eof() {
        let cursor = Cursor::new(Vec::new());
        assert!(Point::read_from(cursor, 0.into()).unwrap().is_none());
    }

    #[test]
    fn read_one_byte() {
        let cursor = Cursor::new(vec![1]);
        assert!(Point::read_from(cursor, 0.into()).is_err());
    }
}
<commit_msg>Move some imports into a functino<commit_after>use {Color, Result};
use point::{Format, ScanDirection};
use std::io::{Read, Write};

/// A raw point.
#[derive(Clone, Copy, Debug, Default)]
#[allow(missing_docs)]
pub struct Point {
    pub x: i32,
    pub y: i32,
    pub z: i32,
    pub intensity: u16,
    pub flags: u8,
    pub classification: u8,
    pub scan_angle_rank: i8,
    pub user_data: u8,
    pub point_source_id: u16,
    pub gps_time: Option<f64>,
    pub color: Option<Color>,
}

impl Point {
    /// Reads a raw point.
    ///
    /// If there are exactly zero bytes left in the `Read`, then this function returns `Ok(None)`.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::{Seek, SeekFrom};
    /// use std::fs::File;
    /// use las::raw::Point;
    /// let mut file = File::open("tests/data/autzen.las").unwrap();
    /// file.seek(SeekFrom::Start(1994)).unwrap();
    /// let point = Point::read_from(file, 1.into()).unwrap();
    /// ```
    pub fn read_from<R: Read>(mut read: R, format: Format) -> Result<Option<Point>> {
        use byteorder::{LittleEndian, ReadBytesExt};
        use std::io::{Cursor, ErrorKind};

        let byte = match read.read_u8() {
            Ok(byte) => byte,
            Err(err) => {
                match err.kind() {
                    ErrorKind::UnexpectedEof => return Ok(None),
                    _ => return Err(err.into()),
                }
            }
        };
        let mut next_three = [0; 3];
        read.read_exact(&mut next_three)?;
        let mut cursor = Cursor::new([byte, next_three[0], next_three[1], next_three[2]]);
        let x = cursor.read_i32::<LittleEndian>()?;
        let y = read.read_i32::<LittleEndian>()?;
        let z = read.read_i32::<LittleEndian>()?;
        let intensity = read.read_u16::<LittleEndian>()?;
        let flags = read.read_u8()?;
        let classification = read.read_u8()?;
        let scan_angle_rank = read.read_i8()?;
        let user_data = read.read_u8()?;
        let point_source_id = read.read_u16::<LittleEndian>()?;
        let gps_time = if format.has_gps_time() {
            Some(read.read_f64::<LittleEndian>()?)
        } else {
            None
        };
        let color = if format.has_color() {
            let red = read.read_u16::<LittleEndian>()?;
            let green = read.read_u16::<LittleEndian>()?;
            let blue = read.read_u16::<LittleEndian>()?;
            Some(Color::new(red, green, blue))
        } else {
            None
        };
        // TODO read extra bytes
        Ok(Some(Point {
            x: x,
            y: y,
            z: z,
            intensity: intensity,
            flags: flags,
            classification: classification,
            scan_angle_rank: scan_angle_rank,
            user_data: user_data,
            point_source_id: point_source_id,
            gps_time: gps_time,
            color: color,
        }))
    }

    /// Calculates the return number from the flag byte.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 1, ..Default::default() };
    /// assert_eq!(1, point.return_number());
    /// ```
    pub fn return_number(&self) -> u8 {
        self.flags & 7
    }

    /// Calculates the number of returns from the flag byte.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 8, ..Default::default() };
    /// assert_eq!(1, point.number_of_returns());
    /// ```
    pub fn number_of_returns(&self) -> u8 {
        (self.flags & 56) >> 3
    }

    /// Returns the scan direction as determined by the scan direction flag.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// use las::point::ScanDirection;
    /// let point = Point { flags: 64, ..Default::default() };
    /// assert_eq!(ScanDirection::Positive, point.scan_direction());
    /// ```
    pub fn scan_direction(&self) -> ScanDirection {
        if (self.flags & 64) == 64 {
            ScanDirection::Positive
        } else {
            ScanDirection::Negative
        }
    }

    /// Returns true if the flags indicate that this point is the edge of a flight line.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { flags: 128, ..Default::default() };
    /// assert!(point.edge_of_flight_line());
    /// ```
    pub fn edge_of_flight_line(&self) -> bool {
        (self.flags & 128) == 128
    }

    /// Returns true if this point is synthetic.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 32, ..Default::default() };
    /// assert!(point.synthetic());
    pub fn synthetic(&self) -> bool {
        (self.classification & 32) == 32
    }

    /// Returns true if this point is a key point.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 64, ..Default::default() };
    /// assert!(point.key_point());
    pub fn key_point(&self) -> bool {
        (self.classification & 64) == 64
    }

    /// Returns true if this point is withheld.
    ///
    /// # Examples
    ///
    /// ```
    /// use las::raw::Point;
    /// let point = Point { classification: 128, ..Default::default() };
    /// assert!(point.withheld());
    pub fn withheld(&self) -> bool {
        (self.classification & 128) == 128
    }

    /// Writes a raw pont.
    ///
    /// # Examples
    ///
    /// `Write` implements `WriteRawPoint`.
    ///
    /// ```
    /// use std::io::Cursor;
    /// use las::raw::Point;
    /// let mut cursor = Cursor::new(Vec::new());
    /// let point = Point::default();
    /// point.write_to(cursor, 0.into()).unwrap();
    /// ```
    pub fn write_to<W: Write>(&self, mut write: W, format: Format) -> Result<()> {
        use byteorder::{LittleEndian, WriteBytesExt};
        use Error;

        write.write_i32::<LittleEndian>(self.x)?;
        write.write_i32::<LittleEndian>(self.y)?;
        write.write_i32::<LittleEndian>(self.z)?;
        write.write_u16::<LittleEndian>(self.intensity)?;
        write.write_u8(self.flags)?;
        write.write_u8(self.classification)?;
        write.write_i8(self.scan_angle_rank)?;
        write.write_u8(self.user_data)?;
        write.write_u16::<LittleEndian>(self.point_source_id)?;
        if format.has_gps_time() {
            if let Some(gps_time) = self.gps_time {
                write.write_f64::<LittleEndian>(gps_time)?;
            } else {
                return Err(Error::MissingGpsTime(format, *self));
            }
        }
        if format.has_color() {
            if let Some(color) = self.color {
                write.write_u16::<LittleEndian>(color.red)?;
                write.write_u16::<LittleEndian>(color.green)?;
                write.write_u16::<LittleEndian>(color.blue)?;
            } else {
                return Err(Error::MissingColor(format, *self));
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use point::{Format, ScanDirection};
    use std::io::Cursor;

    #[test]
    fn return_number() {
        assert_eq!(
            0,
            Point {
                flags: 0,
                ..Default::default()
            }.return_number()
        );
        assert_eq!(
            7,
            Point {
                flags: 7,
                ..Default::default()
            }.return_number()
        );
        assert_eq!(
            0,
            Point {
                flags: 8,
                ..Default::default()
            }.return_number()
        );
    }

    #[test]
    fn number_of_returns() {
        assert_eq!(
            0,
            Point {
                flags: 0,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            1,
            Point {
                flags: 8,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            7,
            Point {
                flags: 56,
                ..Default::default()
            }.number_of_returns()
        );
        assert_eq!(
            0,
            Point {
                flags: 64,
                ..Default::default()
            }.number_of_returns()
        );
    }

    #[test]
    fn scan_direction() {
        assert_eq!(
            ScanDirection::Negative,
            Point { ..Default::default() }.scan_direction()
        );
        assert_eq!(
            ScanDirection::Positive,
            Point {
                flags: 64,
                ..Default::default()
            }.scan_direction()
        );
    }

    #[test]
    fn edge_of_flight_line() {
        assert!(!Point { ..Default::default() }.edge_of_flight_line());
        assert!(
            Point {
                flags: 128,
                ..Default::default()
            }.edge_of_flight_line()
        );
    }

    #[test]
    fn write_without_gps_time() {
        let point = Point { ..Default::default() };
        let write = Cursor::new(Vec::new());
        assert!(point.write_to(write, Format::from(1)).is_err());
    }

    #[test]
    fn read_eof() {
        let cursor = Cursor::new(Vec::new());
        assert!(Point::read_from(cursor, 0.into()).unwrap().is_none());
    }

    #[test]
    fn read_one_byte() {
        let cursor = Cursor::new(vec![1]);
        assert!(Point::read_from(cursor, 0.into()).is_err());
    }
}
<|endoftext|>