<commit_before>//! Utilities for manipulating the primary selection
//!
//! The primary selection is an additional protocol modeled after the data device to represent
//! and additional selection (copy/paste), a concept taken from the X11 Server.
//! This primary selection is a shortcut to the common clipboard selection,
//! where text just needs to be selected in order to allow copying it elsewhere
//! The de facto way to perform this action is the middle mouse button, although it is not limited to this one.
//!
//! This module provides the freestanding [`set_primary_focus`] function:
//!   This function sets the data device focus for a given seat; you'd typically call it
//!   whenever the keyboard focus changes, to follow it (for example in the focus hook of your keyboards).
//!
//! The module also provides an additional mechanism allowing your compositor to see and interact with
//! the contents of the primary selection:
//!
//! - the freestanding function [`set_primary_selection`]
//!   allows you to set the contents of the selection for your clients
//! - the `PrimarySelectionHandle` gives you the option to inspect new selections
//!   by overriding [`PrimarySelectionHandler::new_selection].
//!
//! ## Initialization
//!
//! To initialize this implementation, create the [`PrimarySelectionState`], store it inside your `State` struct
//! and implement the [`PrimarySelectionHandler`], as shown in this example:
//!
//! ```
//! # extern crate wayland_server;
//! # #[macro_use] extern crate smithay;
//! use smithay::delegate_primary_selection;
//! use smithay::wayland::primary_selection::{PrimarySelectionState, PrimarySelectionHandler};
//! # use smithay::input::{Seat, SeatHandler, SeatState, pointer::CursorImageStatus};
//! # use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
//!
//! # struct State { primary_selection_state: PrimarySelectionState }
//! # let mut display = wayland_server::Display::<State>::new().unwrap();
//! // Create the primary_selection state
//! let primary_selection_state = PrimarySelectionState::new::<State, _>(
//!     &display.handle(),
//!     None // We don't add a logger in this example
//! );
//!
//! // insert the PrimarySelectionState into your state
//! // ..
//!
//! // implement the necessary traits
//! # impl SeatHandler for State {
//! #     type KeyboardFocus = WlSurface;
//! #     type PointerFocus = WlSurface;
//! #     fn seat_state(&mut self) -> &mut SeatState<Self> { unimplemented!() }
//! #     fn focus_changed(&mut self, seat: &Seat<Self>, focused: Option<&WlSurface>) { unimplemented!() }
//! #     fn cursor_image(&mut self, seat: &Seat<Self>, image: CursorImageStatus) { unimplemented!() }
//! # }
//! impl PrimarySelectionHandler for State {
//!     fn primary_selection_state(&self) -> &PrimarySelectionState { &self.primary_selection_state }
//!     // ... override default implementations here to customize handling ...
//! }
//! delegate_primary_selection!(State);
//!
//! // You're now ready to go!
//! ```

use std::cell::RefCell;

use io_lifetimes::OwnedFd;
use wayland_protocols::wp::primary_selection::zv1::server::{
    zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1 as PrimaryDeviceManager,
    zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1 as PrimarySource,
};
use wayland_server::{backend::GlobalId, Client, DisplayHandle, GlobalDispatch};

use crate::input::{Seat, SeatHandler};

mod device;
mod seat_data;
mod source;

pub use device::PrimaryDeviceUserData;
pub use source::{with_source_metadata, PrimarySourceUserData, SourceMetadata};

use seat_data::{SeatData, Selection};

/// Events that are generated by interactions of the clients with the data device
pub trait PrimarySelectionHandler: Sized {
    /// [PrimarySelectionState] getter
    fn primary_selection_state(&self) -> &PrimarySelectionState;

    /// A client has set the selection
    #[allow(unused_variables)]
    fn new_selection(&mut self, source: Option<PrimarySource>) {}

    /// A client requested to read the server-set selection
    ///
    /// * `mime_type` - the requested mime type
    /// * `fd` - the fd to write into
    #[allow(unused_variables)]
    fn send_selection(&mut self, mime_type: String, fd: OwnedFd) {}
}

/// State of data device
#[derive(Debug)]
pub struct PrimarySelectionState {
    log: slog::Logger,
    manager_global: GlobalId,
}

impl PrimarySelectionState {
    /// Regiseter new [ZwpPrimarySelectionDeviceManagerV1] global
    pub fn new<D, L>(display: &DisplayHandle, logger: L) -> Self
    where
        L: Into<Option<::slog::Logger>>,
        D: GlobalDispatch<PrimaryDeviceManager, ()> + 'static,
        D: PrimarySelectionHandler,
    {
        let log = crate::slog_or_fallback(logger).new(slog::o!("smithay_module" => "primary_selection_mgr"));

        let manager_global = display.create_global::<D, PrimaryDeviceManager, _>(1, ());

        Self { log, manager_global }
    }

    /// [ZwpPrimarySelectionDeviceManagerV1] GlobalId getter
    pub fn global(&self) -> GlobalId {
        self.manager_global.clone()
    }
}

/// Set the primary selection focus to a certain client for a given seat
pub fn set_primary_focus<D>(dh: &DisplayHandle, seat: &Seat<D>, client: Option<Client>)
where
    D: SeatHandler + PrimarySelectionHandler + 'static,
{
    seat.user_data()
        .insert_if_missing(|| RefCell::new(SeatData::new()));
    let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
    seat_data.borrow_mut().set_focus::<D>(dh, client);
}

/// Set a compositor-provided primary selection for this seat
///
/// You need to provide the available mime types for this selection.
///
/// Whenever a client requests to read the selection, your callback will
/// receive a [`PrimarySelectionHandler::send_selection`] event.
pub fn set_primary_selection<D>(dh: &DisplayHandle, seat: &Seat<D>, mime_types: Vec<String>)
where
    D: SeatHandler + PrimarySelectionHandler + 'static,
{
    seat.user_data()
        .insert_if_missing(|| RefCell::new(SeatData::new()));
    let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
    seat_data
        .borrow_mut()
        .set_selection::<D>(dh, Selection::Compositor(SourceMetadata { mime_types }));
}

mod handlers {
    use std::cell::RefCell;

    use slog::error;
    use wayland_protocols::wp::primary_selection::zv1::server::{
        zwp_primary_selection_device_manager_v1::{
            self as primary_device_manager, ZwpPrimarySelectionDeviceManagerV1 as PrimaryDeviceManager,
        },
        zwp_primary_selection_device_v1::ZwpPrimarySelectionDeviceV1 as PrimaryDevice,
        zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1 as PrimarySource,
    };
    use wayland_server::{Dispatch, DisplayHandle, GlobalDispatch};

    use crate::input::{Seat, SeatHandler};

    use super::{device::PrimaryDeviceUserData, seat_data::SeatData, source::PrimarySourceUserData};
    use super::{PrimarySelectionHandler, PrimarySelectionState};

    impl<D> GlobalDispatch<PrimaryDeviceManager, (), D> for PrimarySelectionState
    where
        D: GlobalDispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimarySource, PrimarySourceUserData>,
        D: Dispatch<PrimaryDevice, PrimaryDeviceUserData>,
        D: PrimarySelectionHandler,
        D: 'static,
    {
        fn bind(
            _state: &mut D,
            _handle: &DisplayHandle,
            _client: &wayland_server::Client,
            resource: wayland_server::New<PrimaryDeviceManager>,
            _global_data: &(),
            data_init: &mut wayland_server::DataInit<'_, D>,
        ) {
            data_init.init(resource, ());
        }
    }

    impl<D> Dispatch<PrimaryDeviceManager, (), D> for PrimarySelectionState
    where
        D: Dispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimarySource, PrimarySourceUserData>,
        D: Dispatch<PrimaryDevice, PrimaryDeviceUserData>,
        D: PrimarySelectionHandler,
        D: SeatHandler,
        D: 'static,
    {
        fn request(
            state: &mut D,
            _client: &wayland_server::Client,
            _resource: &PrimaryDeviceManager,
            request: primary_device_manager::Request,
            _data: &(),
            _dhandle: &DisplayHandle,
            data_init: &mut wayland_server::DataInit<'_, D>,
        ) {
            let primary_selection_state = state.primary_selection_state();

            match request {
                primary_device_manager::Request::CreateSource { id } => {
                    data_init.init(id, PrimarySourceUserData::new());
                }
                primary_device_manager::Request::GetDevice { id, seat: wl_seat } => {
                    match Seat::<D>::from_resource(&wl_seat) {
                        Some(seat) => {
                            seat.user_data()
                                .insert_if_missing(|| RefCell::new(SeatData::new()));

                            let device = data_init.init(id, PrimaryDeviceUserData { wl_seat });

                            let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
                            seat_data.borrow_mut().add_device(device);
                        }
                        None => {
                            error!(
                                &primary_selection_state.log,
                                "Unmanaged seat given to a primary selection device."
                            );
                        }
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

#[allow(missing_docs)] // TODO
#[macro_export]
macro_rules! delegate_primary_selection {
    ($(@<$( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+>)? $ty: ty) => {
        $crate::reexports::wayland_server::delegate_global_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1: ()
        ] => $crate::wayland::primary_selection::PrimarySelectionState);

        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1: ()
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_v1::ZwpPrimarySelectionDeviceV1: $crate::wayland::primary_selection::PrimaryDeviceUserData
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1: $crate::wayland::primary_selection::PrimarySourceUserData
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
    };
}
<commit_msg>fix: add arm for primary_device_manager::Request::Destroy<commit_after>//! Utilities for manipulating the primary selection
//!
//! The primary selection is an additional protocol modeled after the data device to represent
//! and additional selection (copy/paste), a concept taken from the X11 Server.
//! This primary selection is a shortcut to the common clipboard selection,
//! where text just needs to be selected in order to allow copying it elsewhere
//! The de facto way to perform this action is the middle mouse button, although it is not limited to this one.
//!
//! This module provides the freestanding [`set_primary_focus`] function:
//!   This function sets the data device focus for a given seat; you'd typically call it
//!   whenever the keyboard focus changes, to follow it (for example in the focus hook of your keyboards).
//!
//! The module also provides an additional mechanism allowing your compositor to see and interact with
//! the contents of the primary selection:
//!
//! - the freestanding function [`set_primary_selection`]
//!   allows you to set the contents of the selection for your clients
//! - the `PrimarySelectionHandle` gives you the option to inspect new selections
//!   by overriding [`PrimarySelectionHandler::new_selection].
//!
//! ## Initialization
//!
//! To initialize this implementation, create the [`PrimarySelectionState`], store it inside your `State` struct
//! and implement the [`PrimarySelectionHandler`], as shown in this example:
//!
//! ```
//! # extern crate wayland_server;
//! # #[macro_use] extern crate smithay;
//! use smithay::delegate_primary_selection;
//! use smithay::wayland::primary_selection::{PrimarySelectionState, PrimarySelectionHandler};
//! # use smithay::input::{Seat, SeatHandler, SeatState, pointer::CursorImageStatus};
//! # use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
//!
//! # struct State { primary_selection_state: PrimarySelectionState }
//! # let mut display = wayland_server::Display::<State>::new().unwrap();
//! // Create the primary_selection state
//! let primary_selection_state = PrimarySelectionState::new::<State, _>(
//!     &display.handle(),
//!     None // We don't add a logger in this example
//! );
//!
//! // insert the PrimarySelectionState into your state
//! // ..
//!
//! // implement the necessary traits
//! # impl SeatHandler for State {
//! #     type KeyboardFocus = WlSurface;
//! #     type PointerFocus = WlSurface;
//! #     fn seat_state(&mut self) -> &mut SeatState<Self> { unimplemented!() }
//! #     fn focus_changed(&mut self, seat: &Seat<Self>, focused: Option<&WlSurface>) { unimplemented!() }
//! #     fn cursor_image(&mut self, seat: &Seat<Self>, image: CursorImageStatus) { unimplemented!() }
//! # }
//! impl PrimarySelectionHandler for State {
//!     fn primary_selection_state(&self) -> &PrimarySelectionState { &self.primary_selection_state }
//!     // ... override default implementations here to customize handling ...
//! }
//! delegate_primary_selection!(State);
//!
//! // You're now ready to go!
//! ```

use std::cell::RefCell;

use io_lifetimes::OwnedFd;
use wayland_protocols::wp::primary_selection::zv1::server::{
    zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1 as PrimaryDeviceManager,
    zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1 as PrimarySource,
};
use wayland_server::{backend::GlobalId, Client, DisplayHandle, GlobalDispatch};

use crate::input::{Seat, SeatHandler};

mod device;
mod seat_data;
mod source;

pub use device::PrimaryDeviceUserData;
pub use source::{with_source_metadata, PrimarySourceUserData, SourceMetadata};

use seat_data::{SeatData, Selection};

/// Events that are generated by interactions of the clients with the data device
pub trait PrimarySelectionHandler: Sized {
    /// [PrimarySelectionState] getter
    fn primary_selection_state(&self) -> &PrimarySelectionState;

    /// A client has set the selection
    #[allow(unused_variables)]
    fn new_selection(&mut self, source: Option<PrimarySource>) {}

    /// A client requested to read the server-set selection
    ///
    /// * `mime_type` - the requested mime type
    /// * `fd` - the fd to write into
    #[allow(unused_variables)]
    fn send_selection(&mut self, mime_type: String, fd: OwnedFd) {}
}

/// State of data device
#[derive(Debug)]
pub struct PrimarySelectionState {
    log: slog::Logger,
    manager_global: GlobalId,
}

impl PrimarySelectionState {
    /// Regiseter new [ZwpPrimarySelectionDeviceManagerV1] global
    pub fn new<D, L>(display: &DisplayHandle, logger: L) -> Self
    where
        L: Into<Option<::slog::Logger>>,
        D: GlobalDispatch<PrimaryDeviceManager, ()> + 'static,
        D: PrimarySelectionHandler,
    {
        let log = crate::slog_or_fallback(logger).new(slog::o!("smithay_module" => "primary_selection_mgr"));

        let manager_global = display.create_global::<D, PrimaryDeviceManager, _>(1, ());

        Self { log, manager_global }
    }

    /// [ZwpPrimarySelectionDeviceManagerV1] GlobalId getter
    pub fn global(&self) -> GlobalId {
        self.manager_global.clone()
    }
}

/// Set the primary selection focus to a certain client for a given seat
pub fn set_primary_focus<D>(dh: &DisplayHandle, seat: &Seat<D>, client: Option<Client>)
where
    D: SeatHandler + PrimarySelectionHandler + 'static,
{
    seat.user_data()
        .insert_if_missing(|| RefCell::new(SeatData::new()));
    let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
    seat_data.borrow_mut().set_focus::<D>(dh, client);
}

/// Set a compositor-provided primary selection for this seat
///
/// You need to provide the available mime types for this selection.
///
/// Whenever a client requests to read the selection, your callback will
/// receive a [`PrimarySelectionHandler::send_selection`] event.
pub fn set_primary_selection<D>(dh: &DisplayHandle, seat: &Seat<D>, mime_types: Vec<String>)
where
    D: SeatHandler + PrimarySelectionHandler + 'static,
{
    seat.user_data()
        .insert_if_missing(|| RefCell::new(SeatData::new()));
    let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
    seat_data
        .borrow_mut()
        .set_selection::<D>(dh, Selection::Compositor(SourceMetadata { mime_types }));
}

mod handlers {
    use std::cell::RefCell;

    use slog::error;
    use wayland_protocols::wp::primary_selection::zv1::server::{
        zwp_primary_selection_device_manager_v1::{
            self as primary_device_manager, ZwpPrimarySelectionDeviceManagerV1 as PrimaryDeviceManager,
        },
        zwp_primary_selection_device_v1::ZwpPrimarySelectionDeviceV1 as PrimaryDevice,
        zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1 as PrimarySource,
    };
    use wayland_server::{Dispatch, DisplayHandle, GlobalDispatch};

    use crate::input::{Seat, SeatHandler};

    use super::{device::PrimaryDeviceUserData, seat_data::SeatData, source::PrimarySourceUserData};
    use super::{PrimarySelectionHandler, PrimarySelectionState};

    impl<D> GlobalDispatch<PrimaryDeviceManager, (), D> for PrimarySelectionState
    where
        D: GlobalDispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimarySource, PrimarySourceUserData>,
        D: Dispatch<PrimaryDevice, PrimaryDeviceUserData>,
        D: PrimarySelectionHandler,
        D: 'static,
    {
        fn bind(
            _state: &mut D,
            _handle: &DisplayHandle,
            _client: &wayland_server::Client,
            resource: wayland_server::New<PrimaryDeviceManager>,
            _global_data: &(),
            data_init: &mut wayland_server::DataInit<'_, D>,
        ) {
            data_init.init(resource, ());
        }
    }

    impl<D> Dispatch<PrimaryDeviceManager, (), D> for PrimarySelectionState
    where
        D: Dispatch<PrimaryDeviceManager, ()>,
        D: Dispatch<PrimarySource, PrimarySourceUserData>,
        D: Dispatch<PrimaryDevice, PrimaryDeviceUserData>,
        D: PrimarySelectionHandler,
        D: SeatHandler,
        D: 'static,
    {
        fn request(
            state: &mut D,
            _client: &wayland_server::Client,
            _resource: &PrimaryDeviceManager,
            request: primary_device_manager::Request,
            _data: &(),
            _dhandle: &DisplayHandle,
            data_init: &mut wayland_server::DataInit<'_, D>,
        ) {
            let primary_selection_state = state.primary_selection_state();

            match request {
                primary_device_manager::Request::CreateSource { id } => {
                    data_init.init(id, PrimarySourceUserData::new());
                }
                primary_device_manager::Request::GetDevice { id, seat: wl_seat } => {
                    match Seat::<D>::from_resource(&wl_seat) {
                        Some(seat) => {
                            seat.user_data()
                                .insert_if_missing(|| RefCell::new(SeatData::new()));

                            let device = data_init.init(id, PrimaryDeviceUserData { wl_seat });

                            let seat_data = seat.user_data().get::<RefCell<SeatData>>().unwrap();
                            seat_data.borrow_mut().add_device(device);
                        }
                        None => {
                            error!(
                                &primary_selection_state.log,
                                "Unmanaged seat given to a primary selection device."
                            );
                        }
                    }
                }
                primary_device_manager::Request::Destroy => {}
                _ => unreachable!(),
            }
        }
    }
}

#[allow(missing_docs)] // TODO
#[macro_export]
macro_rules! delegate_primary_selection {
    ($(@<$( $lt:tt $( : $clt:tt $(+ $dlt:tt )* )? ),+>)? $ty: ty) => {
        $crate::reexports::wayland_server::delegate_global_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1: ()
        ] => $crate::wayland::primary_selection::PrimarySelectionState);

        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_manager_v1::ZwpPrimarySelectionDeviceManagerV1: ()
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_device_v1::ZwpPrimarySelectionDeviceV1: $crate::wayland::primary_selection::PrimaryDeviceUserData
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
        $crate::reexports::wayland_server::delegate_dispatch!($(@< $( $lt $( : $clt $(+ $dlt )* )? ),+ >)? $ty: [
            $crate::reexports::wayland_protocols::wp::primary_selection::zv1::server::zwp_primary_selection_source_v1::ZwpPrimarySelectionSourceV1: $crate::wayland::primary_selection::PrimarySourceUserData
        ] => $crate::wayland::primary_selection::PrimarySelectionState);
    };
}
<|endoftext|>


<commit_before>//! Support for the wlroots Gamma Control Protocol

use wayland_sys::server::wl_display as wl_server_display;
use wlroots_sys::{wl_display, wlr_gtk_primary_selection_device_manager, wlr_gtk_primary_selection_device_manager_create,
    wlr_gtk_primary_selection_device_manager_destroy};

#[derive(Debug)]
/// Manager that implements GTK primary selection
pub struct Manager {
    manager: *mut wlr_gtk_primary_selection_device_manager
}

impl Manager {
    pub(crate) unsafe fn new(display: *mut wl_server_display) -> Option<Self> {
        let manager_raw = wlr_gtk_primary_selection_device_manager_create(display as *mut wl_display);

        if !manager_raw.is_null() {
            Some(Manager { manager: manager_raw })
        } else {
            None
        }
    }

}

impl Drop for Manager {
    fn drop(&mut self) {
        unsafe { wlr_gtk_primary_selection_device_manager_destroy(self.manager) }
    }
}
<commit_msg>fix docs<commit_after>//! Support for the GTK Primary Selection Protocol

use wayland_sys::server::wl_display as wl_server_display;
use wlroots_sys::{wl_display, wlr_gtk_primary_selection_device_manager, wlr_gtk_primary_selection_device_manager_create,
    wlr_gtk_primary_selection_device_manager_destroy};

#[derive(Debug)]
/// Manager that implements GTK primary selection
pub struct Manager {
    manager: *mut wlr_gtk_primary_selection_device_manager
}

impl Manager {
    pub(crate) unsafe fn new(display: *mut wl_server_display) -> Option<Self> {
        let manager_raw = wlr_gtk_primary_selection_device_manager_create(display as *mut wl_display);

        if !manager_raw.is_null() {
            Some(Manager { manager: manager_raw })
        } else {
            None
        }
    }

}

impl Drop for Manager {
    fn drop(&mut self) {
        unsafe { wlr_gtk_primary_selection_device_manager_destroy(self.manager) }
    }
}
<|endoftext|>


<commit_before>//! # env
//!
//! Sets up the env vars before running the tasks.
//!

pub(crate) mod crateinfo;
mod gitinfo;

#[cfg(test)]
#[path = "./mod_test.rs"]
mod mod_test;

use command;
use indexmap::IndexMap;
use rust_info;
use rust_info::types::{RustChannel, RustInfo};
use std::env;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use types::{
    Config, CrateInfo, EnvInfo, EnvValue, EnvValueInfo, GitInfo, PackageInfo, Step, Workspace,
};

fn evaluate_env_value(env_value: &EnvValueInfo) -> String {
    match command::run_script_get_output(&env_value.script, None, true) {
        Ok(output) => {
            let exit_code = output.0;
            let stdout = output.1;

            command::validate_exit_code(exit_code);

            if exit_code == 0 {
                let mut lines: Vec<&str> = stdout.split("\n").collect();
                lines.retain(|&line| line.len() > 0);

                if lines.len() > 0 {
                    let line = lines[lines.len() - 1].to_string();

                    let line_str = str::replace(&line, "\r", "");

                    line_str.to_string()
                } else {
                    "".to_string()
                }
            } else {
                "".to_string()
            }
        }
        _ => "".to_string(),
    }
}

fn expand_value(value: &str) -> String {
    let mut value_string = value.to_string();
    match value_string.find("${") {
        Some(_) => {
            for (existing_key, existing_value) in env::vars() {
                let mut key_pattern = "${".to_string();
                key_pattern.push_str(&existing_key);
                key_pattern.push_str("}");

                value_string = str::replace(&value_string, &key_pattern, &existing_value);
            }

            value_string
        }
        None => value_string,
    }
}

fn evaluate_and_set_env(key: &str, value: &str) {
    let env_value = expand_value(&value);

    env::set_var(&key, &env_value);
}

fn set_env_for_info(key: &str, env_value: &EnvValueInfo) {
    let value = evaluate_env_value(&env_value);

    evaluate_and_set_env(&key, &value);
}

/// Updates the env based on the provided data
pub(crate) fn set_env(env: IndexMap<String, EnvValue>) {
    debug!("Setting Up Env.");

    for (key, env_value) in &env {
        debug!("Setting env: {} = {:#?}", &key, &env_value);

        match *env_value {
            EnvValue::Value(ref value) => evaluate_and_set_env(&key, value),
            EnvValue::Info(ref info) => set_env_for_info(&key, info),
        };
    }
}

/// Updates the env for the current execution based on the descriptor.
fn initialize_env(config: &Config) {
    info!("Setting Up Env.");

    set_env(config.env.clone());
}

fn setup_env_for_crate() -> CrateInfo {
    let crate_info = crateinfo::load();
    let crate_info_clone = crate_info.clone();

    let package_info = crate_info.package.unwrap_or(PackageInfo::new());

    if package_info.name.is_some() {
        let crate_name = package_info.name.unwrap();
        env::set_var("CARGO_MAKE_CRATE_NAME", &crate_name);

        let crate_fs_name = str::replace(&crate_name, "-", "_");
        env::set_var("CARGO_MAKE_CRATE_FS_NAME", &crate_fs_name);
    }

    if package_info.version.is_some() {
        env::set_var("CARGO_MAKE_CRATE_VERSION", &package_info.version.unwrap());
    }

    if package_info.description.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_DESCRIPTION",
            &package_info.description.unwrap(),
        );
    }

    if package_info.license.is_some() {
        env::set_var("CARGO_MAKE_CRATE_LICENSE", &package_info.license.unwrap());
    }

    if package_info.documentation.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_DOCUMENTATION",
            &package_info.documentation.unwrap(),
        );
    }

    if package_info.homepage.is_some() {
        env::set_var("CARGO_MAKE_CRATE_HOMEPAGE", &package_info.homepage.unwrap());
    }

    if package_info.repository.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_REPOSITORY",
            &package_info.repository.unwrap(),
        );
    }

    let has_dependencies = match crate_info.dependencies {
        Some(ref dependencies) => dependencies.len() > 0,
        None => crate_info.workspace.is_some(),
    };

    let has_dependencies_var_value = if has_dependencies { "TRUE" } else { "FALSE" };
    env::set_var(
        "CARGO_MAKE_CRATE_HAS_DEPENDENCIES",
        has_dependencies_var_value,
    );

    let is_workspace_var_value = if crate_info.workspace.is_none() {
        "FALSE"
    } else {
        "TRUE"
    };
    env::set_var("CARGO_MAKE_CRATE_IS_WORKSPACE", is_workspace_var_value);

    let workspace = crate_info.workspace.unwrap_or(Workspace::new());
    let members = workspace.members.unwrap_or(vec![]);
    let members_string = members.join(",");
    env::set_var("CARGO_MAKE_CRATE_WORKSPACE_MEMBERS", &members_string);

    // check if Cargo.lock file exists in working directory
    let lock_file = Path::new("Cargo.lock");
    let lock_file_exists_var_value = if lock_file.exists() { "TRUE" } else { "FALSE" };
    env::set_var(
        "CARGO_MAKE_CRATE_LOCK_FILE_EXISTS",
        lock_file_exists_var_value,
    );

    crate_info_clone
}

fn setup_env_for_git_repo() -> GitInfo {
    let git_info = gitinfo::load();
    let git_info_clone = git_info.clone();

    if git_info.branch.is_some() {
        env::set_var("CARGO_MAKE_GIT_BRANCH", &git_info.branch.unwrap());
    }

    if git_info.user_name.is_some() {
        env::set_var("CARGO_MAKE_GIT_USER_NAME", &git_info.user_name.unwrap());
    }

    if git_info.user_email.is_some() {
        env::set_var("CARGO_MAKE_GIT_USER_EMAIL", &git_info.user_email.unwrap());
    }

    git_info_clone
}

fn setup_env_for_rust() -> RustInfo {
    let rustinfo = rust_info::get();
    let rust_info_clone = rustinfo.clone();

    if rustinfo.version.is_some() {
        env::set_var("CARGO_MAKE_RUST_VERSION", &rustinfo.version.unwrap());
    }

    if rustinfo.channel.is_some() {
        let channel_option = rustinfo.channel.unwrap();

        let channel = match channel_option {
            RustChannel::Stable => "stable",
            RustChannel::Beta => "beta",
            RustChannel::Nightly => "nightly",
        };

        env::set_var("CARGO_MAKE_RUST_CHANNEL", channel.to_string());
    }

    env::set_var(
        "CARGO_MAKE_RUST_TARGET_ARCH",
        &rustinfo.target_arch.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_ENV",
        &rustinfo.target_env.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_OS",
        &rustinfo.target_os.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_POINTER_WIDTH",
        &rustinfo
            .target_pointer_width
            .unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_VENDOR",
        &rustinfo.target_vendor.unwrap_or("unknown".to_string()),
    );

    rust_info_clone
}

/// Sets up the env before the tasks execution.
pub(crate) fn setup_env(config: &Config, task: &str) -> EnvInfo {
    initialize_env(config);

    env::set_var("CARGO_MAKE", "true");
    env::set_var("CARGO_MAKE_TASK", &task);

    // load crate info
    let crate_info = setup_env_for_crate();

    // load git info
    let git_info = setup_env_for_git_repo();

    // load rust info
    let rust_info = setup_env_for_rust();

    EnvInfo {
        rust_info,
        crate_info,
        git_info,
    }
}

pub(crate) fn setup_cwd(cwd: Option<&str>) {
    let directory = cwd.unwrap_or(".");

    debug!("Changing working directory to: {}", &directory);

    let mut directory_path_buf = PathBuf::from(&directory);
    if !cfg!(windows) {
        directory_path_buf = directory_path_buf
            .canonicalize()
            .unwrap_or(directory_path_buf);
    }
    let directory_path = directory_path_buf.as_path();

    match env::set_current_dir(&directory_path) {
        Err(error) => error!(
            "Unable to set current working directory to: {} {:#?}",
            &directory, error
        ),
        _ => {
            env::set_var("CARGO_MAKE_WORKING_DIRECTORY", directory_path);

            debug!("Working directory changed to: {}", &directory);
        }
    }
}

pub(crate) fn get_env(key: &str, default_value: &str) -> String {
    match env::var(key) {
        Ok(value) => value.to_string(),
        _ => default_value.to_string(),
    }
}

pub(crate) fn get_env_as_bool(key: &str, default_value: bool) -> bool {
    let default_str = if default_value { "true" } else { "false" };

    let mut value = get_env(key, default_str);
    value = value.to_lowercase();

    if value == "true" || value == "yes" || value == "1" {
        true
    } else {
        false
    }
}

pub(crate) fn parse_env_file(env_file: Option<String>) -> Option<Vec<String>> {
    match env_file {
        Some(file_name) => {
            let file_path = if file_name.starts_with(".") {
                let base_path = get_env("CARGO_MAKE_WORKING_DIRECTORY", ".");
                Path::new(&base_path).join(file_name)
            } else {
                Path::new(&file_name).to_path_buf()
            };

            if file_path.exists() {
                debug!("Opening env file: {:#?}", &file_path);
                let mut file = match File::open(&file_path) {
                    Ok(value) => value,
                    Err(error) => panic!(
                        "Unable to open env file: {} error: {}",
                        file_path.to_str().unwrap_or(""),
                        error
                    ),
                };

                let mut env_content = String::new();
                file.read_to_string(&mut env_content).unwrap();

                let mut env: Vec<String> = vec![];

                let lines: Vec<&str> = env_content.split('\n').collect();

                for mut line in lines {
                    line = line.trim();

                    if !line.starts_with("#") {
                        env.push(line.to_string());
                    }
                }

                Some(env)
            } else {
                None
            }
        }
        None => None,
    }
}

pub(crate) fn get_cargo_make_home() -> Option<PathBuf> {
    match env::var("CARGO_MAKE_HOME") {
        Ok(directory) => Some(PathBuf::from(directory)),
        _ => match env::home_dir() {
            Some(directory) => Some(directory.join(".cargo-make")),
            None => None,
        },
    }
}

fn get_project_root_for_path(directory: &PathBuf) -> Option<String> {
    let file_path = Path::new(directory).join("Cargo.toml");

    if file_path.exists() {
        match directory.to_str() {
            Some(directory_string) => Some(directory_string.to_string()),
            _ => None,
        }
    } else {
        match directory.parent() {
            Some(parent_directory) => {
                let parent_directory_path = parent_directory.to_path_buf();
                get_project_root_for_path(&parent_directory_path)
            }
            None => None,
        }
    }
}

pub(crate) fn get_project_root() -> Option<String> {
    match env::current_dir() {
        Ok(directory) => get_project_root_for_path(&directory),
        _ => None,
    }
}

pub(crate) fn expand_env(step: &Step) -> Step {
    //clone data before modify
    let mut config = step.config.clone();

    //update command by replacing any env vars
    match config.command {
        Some(value) => {
            config.command = Some(expand_value(&value));
        }
        None => {}
    };

    //update args by replacing any env vars
    match config.args {
        Some(ref mut args) => {
            for index in 0..args.len() {
                args[index] = expand_value(&args[index]);
            }
        }
        None => {}
    };

    Step {
        name: step.name.clone(),
        config,
    }
}
<commit_msg>fix workspace coverage not being generated #97<commit_after>//! # env
//!
//! Sets up the env vars before running the tasks.
//!

pub(crate) mod crateinfo;
mod gitinfo;

#[cfg(test)]
#[path = "./mod_test.rs"]
mod mod_test;

use command;
use indexmap::IndexMap;
use rust_info;
use rust_info::types::{RustChannel, RustInfo};
use std::env;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};
use types::{
    Config, CrateInfo, EnvInfo, EnvValue, EnvValueInfo, GitInfo, PackageInfo, Step, Workspace,
};

fn evaluate_env_value(env_value: &EnvValueInfo) -> String {
    match command::run_script_get_output(&env_value.script, None, true) {
        Ok(output) => {
            let exit_code = output.0;
            let stdout = output.1;

            command::validate_exit_code(exit_code);

            if exit_code == 0 {
                let mut lines: Vec<&str> = stdout.split("\n").collect();
                lines.retain(|&line| line.len() > 0);

                if lines.len() > 0 {
                    let line = lines[lines.len() - 1].to_string();

                    let line_str = str::replace(&line, "\r", "");

                    line_str.to_string()
                } else {
                    "".to_string()
                }
            } else {
                "".to_string()
            }
        }
        _ => "".to_string(),
    }
}

fn expand_value(value: &str) -> String {
    let mut value_string = value.to_string();
    match value_string.find("${") {
        Some(_) => {
            for (existing_key, existing_value) in env::vars() {
                let mut key_pattern = "${".to_string();
                key_pattern.push_str(&existing_key);
                key_pattern.push_str("}");

                value_string = str::replace(&value_string, &key_pattern, &existing_value);
            }

            value_string
        }
        None => value_string,
    }
}

fn evaluate_and_set_env(key: &str, value: &str) {
    let env_value = expand_value(&value);

    env::set_var(&key, &env_value);
}

fn set_env_for_info(key: &str, env_value: &EnvValueInfo) {
    let value = evaluate_env_value(&env_value);

    evaluate_and_set_env(&key, &value);
}

/// Updates the env based on the provided data
pub(crate) fn set_env(env: IndexMap<String, EnvValue>) {
    debug!("Setting Up Env.");

    for (key, env_value) in &env {
        debug!("Setting env: {} = {:#?}", &key, &env_value);

        match *env_value {
            EnvValue::Value(ref value) => evaluate_and_set_env(&key, value),
            EnvValue::Info(ref info) => set_env_for_info(&key, info),
        };
    }
}

/// Updates the env for the current execution based on the descriptor.
fn initialize_env(config: &Config) {
    info!("Setting Up Env.");

    set_env(config.env.clone());
}

fn setup_env_for_crate() -> CrateInfo {
    let crate_info = crateinfo::load();
    let crate_info_clone = crate_info.clone();

    let package_info = crate_info.package.unwrap_or(PackageInfo::new());

    if package_info.name.is_some() {
        let crate_name = package_info.name.unwrap();
        env::set_var("CARGO_MAKE_CRATE_NAME", &crate_name);

        let crate_fs_name = str::replace(&crate_name, "-", "_");
        env::set_var("CARGO_MAKE_CRATE_FS_NAME", &crate_fs_name);
    }

    if package_info.version.is_some() {
        env::set_var("CARGO_MAKE_CRATE_VERSION", &package_info.version.unwrap());
    }

    if package_info.description.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_DESCRIPTION",
            &package_info.description.unwrap(),
        );
    }

    if package_info.license.is_some() {
        env::set_var("CARGO_MAKE_CRATE_LICENSE", &package_info.license.unwrap());
    }

    if package_info.documentation.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_DOCUMENTATION",
            &package_info.documentation.unwrap(),
        );
    }

    if package_info.homepage.is_some() {
        env::set_var("CARGO_MAKE_CRATE_HOMEPAGE", &package_info.homepage.unwrap());
    }

    if package_info.repository.is_some() {
        env::set_var(
            "CARGO_MAKE_CRATE_REPOSITORY",
            &package_info.repository.unwrap(),
        );
    }

    let has_dependencies = match crate_info.dependencies {
        Some(ref dependencies) => dependencies.len() > 0,
        None => crate_info.workspace.is_some(),
    };

    let has_dependencies_var_value = if has_dependencies { "TRUE" } else { "FALSE" };
    env::set_var(
        "CARGO_MAKE_CRATE_HAS_DEPENDENCIES",
        has_dependencies_var_value,
    );

    let is_workspace_var_value = if crate_info.workspace.is_none() {
        "FALSE"
    } else {
        "TRUE"
    };
    env::set_var("CARGO_MAKE_CRATE_IS_WORKSPACE", is_workspace_var_value);

    let workspace = crate_info.workspace.unwrap_or(Workspace::new());
    let members = workspace.members.unwrap_or(vec![]);
    let members_string = members.join(",");
    env::set_var("CARGO_MAKE_CRATE_WORKSPACE_MEMBERS", &members_string);

    // check if Cargo.lock file exists in working directory
    let lock_file = Path::new("Cargo.lock");
    let lock_file_exists_var_value = if lock_file.exists() { "TRUE" } else { "FALSE" };
    env::set_var(
        "CARGO_MAKE_CRATE_LOCK_FILE_EXISTS",
        lock_file_exists_var_value,
    );

    crate_info_clone
}

fn setup_env_for_git_repo() -> GitInfo {
    let git_info = gitinfo::load();
    let git_info_clone = git_info.clone();

    if git_info.branch.is_some() {
        env::set_var("CARGO_MAKE_GIT_BRANCH", &git_info.branch.unwrap());
    }

    if git_info.user_name.is_some() {
        env::set_var("CARGO_MAKE_GIT_USER_NAME", &git_info.user_name.unwrap());
    }

    if git_info.user_email.is_some() {
        env::set_var("CARGO_MAKE_GIT_USER_EMAIL", &git_info.user_email.unwrap());
    }

    git_info_clone
}

fn setup_env_for_rust() -> RustInfo {
    let rustinfo = rust_info::get();
    let rust_info_clone = rustinfo.clone();

    if rustinfo.version.is_some() {
        env::set_var("CARGO_MAKE_RUST_VERSION", &rustinfo.version.unwrap());
    }

    if rustinfo.channel.is_some() {
        let channel_option = rustinfo.channel.unwrap();

        let channel = match channel_option {
            RustChannel::Stable => "stable",
            RustChannel::Beta => "beta",
            RustChannel::Nightly => "nightly",
        };

        env::set_var("CARGO_MAKE_RUST_CHANNEL", channel.to_string());
    }

    env::set_var(
        "CARGO_MAKE_RUST_TARGET_ARCH",
        &rustinfo.target_arch.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_ENV",
        &rustinfo.target_env.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_OS",
        &rustinfo.target_os.unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_POINTER_WIDTH",
        &rustinfo
            .target_pointer_width
            .unwrap_or("unknown".to_string()),
    );
    env::set_var(
        "CARGO_MAKE_RUST_TARGET_VENDOR",
        &rustinfo.target_vendor.unwrap_or("unknown".to_string()),
    );

    rust_info_clone
}

/// Sets up the env before the tasks execution.
pub(crate) fn setup_env(config: &Config, task: &str) -> EnvInfo {
    env::set_var("CARGO_MAKE", "true");
    env::set_var("CARGO_MAKE_TASK", &task);

    // load crate info
    let crate_info = setup_env_for_crate();

    // load git info
    let git_info = setup_env_for_git_repo();

    // load rust info
    let rust_info = setup_env_for_rust();

    // load env vars
    initialize_env(config);

    EnvInfo {
        rust_info,
        crate_info,
        git_info,
    }
}

pub(crate) fn setup_cwd(cwd: Option<&str>) {
    let directory = cwd.unwrap_or(".");

    debug!("Changing working directory to: {}", &directory);

    let mut directory_path_buf = PathBuf::from(&directory);
    if !cfg!(windows) {
        directory_path_buf = directory_path_buf
            .canonicalize()
            .unwrap_or(directory_path_buf);
    }
    let directory_path = directory_path_buf.as_path();

    match env::set_current_dir(&directory_path) {
        Err(error) => error!(
            "Unable to set current working directory to: {} {:#?}",
            &directory, error
        ),
        _ => {
            env::set_var("CARGO_MAKE_WORKING_DIRECTORY", directory_path);

            debug!("Working directory changed to: {}", &directory);
        }
    }
}

pub(crate) fn get_env(key: &str, default_value: &str) -> String {
    match env::var(key) {
        Ok(value) => value.to_string(),
        _ => default_value.to_string(),
    }
}

pub(crate) fn get_env_as_bool(key: &str, default_value: bool) -> bool {
    let default_str = if default_value { "true" } else { "false" };

    let mut value = get_env(key, default_str);
    value = value.to_lowercase();

    if value == "true" || value == "yes" || value == "1" {
        true
    } else {
        false
    }
}

pub(crate) fn parse_env_file(env_file: Option<String>) -> Option<Vec<String>> {
    match env_file {
        Some(file_name) => {
            let file_path = if file_name.starts_with(".") {
                let base_path = get_env("CARGO_MAKE_WORKING_DIRECTORY", ".");
                Path::new(&base_path).join(file_name)
            } else {
                Path::new(&file_name).to_path_buf()
            };

            if file_path.exists() {
                debug!("Opening env file: {:#?}", &file_path);
                let mut file = match File::open(&file_path) {
                    Ok(value) => value,
                    Err(error) => panic!(
                        "Unable to open env file: {} error: {}",
                        file_path.to_str().unwrap_or(""),
                        error
                    ),
                };

                let mut env_content = String::new();
                file.read_to_string(&mut env_content).unwrap();

                let mut env: Vec<String> = vec![];

                let lines: Vec<&str> = env_content.split('\n').collect();

                for mut line in lines {
                    line = line.trim();

                    if !line.starts_with("#") {
                        env.push(line.to_string());
                    }
                }

                Some(env)
            } else {
                None
            }
        }
        None => None,
    }
}

pub(crate) fn get_cargo_make_home() -> Option<PathBuf> {
    match env::var("CARGO_MAKE_HOME") {
        Ok(directory) => Some(PathBuf::from(directory)),
        _ => match env::home_dir() {
            Some(directory) => Some(directory.join(".cargo-make")),
            None => None,
        },
    }
}

fn get_project_root_for_path(directory: &PathBuf) -> Option<String> {
    let file_path = Path::new(directory).join("Cargo.toml");

    if file_path.exists() {
        match directory.to_str() {
            Some(directory_string) => Some(directory_string.to_string()),
            _ => None,
        }
    } else {
        match directory.parent() {
            Some(parent_directory) => {
                let parent_directory_path = parent_directory.to_path_buf();
                get_project_root_for_path(&parent_directory_path)
            }
            None => None,
        }
    }
}

pub(crate) fn get_project_root() -> Option<String> {
    match env::current_dir() {
        Ok(directory) => get_project_root_for_path(&directory),
        _ => None,
    }
}

pub(crate) fn expand_env(step: &Step) -> Step {
    //clone data before modify
    let mut config = step.config.clone();

    //update command by replacing any env vars
    match config.command {
        Some(value) => {
            config.command = Some(expand_value(&value));
        }
        None => {}
    };

    //update args by replacing any env vars
    match config.args {
        Some(ref mut args) => {
            for index in 0..args.len() {
                args[index] = expand_value(&args[index]);
            }
        }
        None => {}
    };

    Step {
        name: step.name.clone(),
        config,
    }
}
<|endoftext|>


<commit_before>use std::iter::Rev;
use core::slice::Items;
use std::collections::hashmap::{HashMap, HashSet, Entries};

pub const NS_XMLNS_PREFIX: &'static str = "xmlns";
pub const NS_XMLNS_URI: &'static str    = "http://www.w3.org/2000/xmlns/";
pub const NS_XML_PREFIX: &'static str   = "xml";
pub const NS_XML_URI: &'static str      = "http://www.w3.org/XML/1998/namespace";
pub const NS_EMPTY_URI: &'static str    = "";

/// Denotes something which contains namespace URI mappings.
///
/// A URI mapping is a pair of type `(Option<&str>, &str)`, where the first item
/// is namespace prefix (`None` meaning default prefix) and the second item
/// is a URI mapped to the prefix.
pub trait NamespaceIterable<'a, I: Iterator<(Option<&'a str>, &'a str)>> {
    fn uri_mappings(&'a self) -> I;
}

/// Namespace is a map from prefixes to namespace URIs.
///
/// `None` prefix means no prefix (i.e. default namespace).
#[deriving(PartialEq, Clone)]
pub struct Namespace(pub HashMap<Option<String>, String>);

impl Namespace {
    /// Returns an empty namespace.
    #[inline]
    pub fn empty() -> Namespace { Namespace(HashMap::with_capacity(2)) }

    /// Checks whether this namespace is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        let Namespace(ref hm) = *self;
        hm.is_empty()
    }

    /// Checks whether this namespace is essentially empty, that is, it does not contain
    /// anything but default mappings.
    pub fn is_essentially_empty(&self) -> bool {
        let Namespace(ref hm) = *self;
        for (k, v) in hm.iter() {
            match (k.as_ref().map(|k| k.as_slice()), v.as_slice()) {
                (None, u)    if u == NS_EMPTY_URI                         => {},
                (Some(p), u) if p == NS_XMLNS_PREFIX && u == NS_XMLNS_URI => {},
                (Some(p), u) if p == NS_XML_PREFIX   && u == NS_XML_URI   => {},
                _ => return false
            }
        }
        true
    }

    /// Puts a mapping into this namespace.
    ///
    /// This method does not override already existing mapping.
    ///
    /// Returns a boolean flag indicating whether the map already contained
    /// the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace);
    /// * `uri`    --- namespace URI.
    ///
    /// # Return value
    /// `true` if `prefix` has been inserted successfully; `false` if the `prefix`
    /// was already present in the namespace.
    pub fn put(&mut self, prefix: Option<String>, uri: String) -> bool {
        match *self {
            Namespace(ref mut hm) => hm.insert(prefix, uri)
        }
    }

    /// Queries the namespace for the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace).
    ///
    /// # Return value
    /// Namespace URI corresponding to the given prefix, if it is present.
    pub fn get<'a>(&'a self, prefix: &Option<String>) -> Option<&'a str> {
        match *self {
            Namespace(ref hm) => hm.find(prefix).map(|s| s.as_slice())
        }
    }
}

/// An iterator over mappings from prefixes to URIs in a namespace.
pub struct NamespaceMappings<'a> {
    entries: Entries<'a, Option<String>, String>
}

impl<'a> Iterator<(Option<&'a str>, &'a str)> for NamespaceMappings<'a> {
    fn next(&mut self) -> Option<(Option<&'a str>, &'a str)> {
        self.entries.next().map(|(prefix, uri)| {
            (prefix.as_ref().map(|p| p.as_slice()), uri.as_slice())
        })
    }
}

impl<'a> NamespaceIterable<'a, NamespaceMappings<'a>> for Namespace {
    fn uri_mappings(&'a self) -> NamespaceMappings<'a> {
        let Namespace(ref hm) = *self;
        NamespaceMappings { entries: hm.iter() }
    }
}

/// Namespace stack is a sequence of namespaces.
///
/// Namespace stack is used to represent cumulative namespace consisting of
/// combined namespaces from nested elements.
#[deriving(Clone, PartialEq)]
pub struct NamespaceStack(pub Vec<Namespace>);

impl NamespaceStack {
    /// Returns an empty namespace stack.
    #[inline]
    pub fn empty() -> NamespaceStack { NamespaceStack(Vec::with_capacity(2)) }

    /// Returns a namespace stack with default items in it.
    ///
    /// Default items are the following:
    ///
    /// * `xml` → `http://www.w3.org/XML/1998/namespace`;
    /// * `xmlns` → `http://www.w3.org/2000/xmlns/`.
    #[inline]
    pub fn default() -> NamespaceStack {
        let mut nst = NamespaceStack::empty();
        nst.push_empty();
        // xml namespace
        nst.put(Some(NS_XML_PREFIX.to_string()), NS_XML_URI.to_string());
        // xmlns namespace
        nst.put(Some(NS_XMLNS_PREFIX.to_string()), NS_XMLNS_URI.to_string());
        // empty namespace
        nst.put(None, NS_EMPTY_URI.to_string());
        nst
    }

    /// Adds an empty namespace to the top of this stack.
    #[inline]
    pub fn push_empty(&mut self) {
        let NamespaceStack(ref mut nst) = *self;
        nst.push(Namespace::empty());
    }

    /// Removes a namespace at the top of the stack.
    ///
    /// Fails if the stack is empty.
    #[inline]
    pub fn pop(&mut self) -> Namespace {
        let NamespaceStack(ref mut nst) = *self;
        nst.pop().unwrap()
    }

    /// Returns a namespace at the top of the stack, leaving the stack intact.
    ///
    /// Fails if the stack is empty.
    #[inline]
    pub fn peek<'a>(&'a mut self) -> &'a mut Namespace {
        let NamespaceStack(ref mut nst) = *self;
        nst.last_mut().unwrap()
    }

    /// Puts a mapping into the topmost namespace in this stack.
    ///
    /// This method does not override a mapping in the topmost namespace if it is
    /// already present, however, it does not depend on other namespaces in the stack,
    /// so it is possible to put a mapping which is present in lower namespaces.
    ///
    /// Returns a boolean flag indicating whether the topmost namespace
    /// already contained the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace);
    /// * `uri`    --- namespace URI.
    ///
    /// # Return value
    /// `true` if `prefix` has been inserted successfully; `false` if the `prefix`
    /// was already present in the namespace.
    #[inline]
    pub fn put(&mut self, prefix: Option<String>, uri: String) -> bool {
        let NamespaceStack(ref mut nst) = *self;
        nst.last_mut().unwrap().put(prefix, uri)
    }

    /// Performs a search for the given prefix in the whole stack.
    ///
    /// This method walks the stack from top to bottom, querying each namespace
    /// in order for the given prefix. If none of the namespaces contains the prefix,
    /// `None` is returned.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace)
    #[inline]
    pub fn get<'a>(&'a self, prefix: &Option<String>) -> Option<&'a str> {
        let NamespaceStack(ref nst) = *self;
        for ns in nst.iter().rev() {
            match ns.get(prefix) {
                None => {},
                r => return r,
            }
        }
        None
    }

    /// Combines this stack of namespaces into a single namespace.
    ///
    /// Namespaces are combined in left-to-right order, that is, rightmost namespace
    /// elements take priority over leftmost ones.
    pub fn squash(&self) -> Namespace {
        let NamespaceStack(ref nstack) = *self;
        let mut result = HashMap::new();
        for &Namespace(ref ns) in nstack.iter() {
            result.extend(ns.iter().map(|(k, v)| (k.clone(), v.to_string())));
        }
        Namespace(result)
    }
}

/// An iterator over mappings from prefixes to URIs in a namespace stack.
pub struct NamespaceStackMappings<'a> {
    namespaces: Rev<Items<'a, Namespace>>,
    current_namespace: Option<NamespaceMappings<'a>>,
    used_keys: HashSet<Option<&'a str>>
}

impl<'a> NamespaceStackMappings<'a> {
    fn to_next_namespace(&mut self) -> bool {
        self.current_namespace = self.namespaces.next().map(|ns| ns.uri_mappings());
        self.current_namespace.is_some()
    }
}

impl<'a> Iterator<(Option<&'a str>, &'a str)> for NamespaceStackMappings<'a> {
    fn next(&mut self) -> Option<(Option<&'a str>, &'a str)> {
        // If there is no current namespace and no next namespace, we're finished
        if self.current_namespace.is_none() && !self.to_next_namespace() {
            return None;
        }
        let next_item = self.current_namespace.as_mut().unwrap().next();

        match next_item {
            // There is an element in the current namespace
            Some((k, v)) => if self.used_keys.contains(&k) {
                // If the current key is used, go to the next one
                self.next()
            } else {
                // Otherwise insert the current key to the set of used keys and
                // return the mapping
                self.used_keys.insert(k);
                Some((k, v))
            },
            // Current namespace is exhausted
            None => if self.to_next_namespace() {
                // If there is next namespace, continue from it
                self.next()
            } else {
                // No next namespace, exiting
                None
            }
        }
    }
}

impl<'a> NamespaceIterable<'a, NamespaceStackMappings<'a>> for NamespaceStack {
    fn uri_mappings(&'a self) -> NamespaceStackMappings<'a> {
        let NamespaceStack(ref nst) = *self;
        NamespaceStackMappings {
            namespaces: nst.iter().rev(),
            current_namespace: None,
            used_keys: HashSet::new()
        }
    }
}
<commit_msg>Switched to tuple indexing in namespaces<commit_after>use std::iter::Rev;
use core::slice::Items;
use std::collections::hashmap::{HashMap, HashSet, Entries};

pub const NS_XMLNS_PREFIX: &'static str = "xmlns";
pub const NS_XMLNS_URI: &'static str    = "http://www.w3.org/2000/xmlns/";
pub const NS_XML_PREFIX: &'static str   = "xml";
pub const NS_XML_URI: &'static str      = "http://www.w3.org/XML/1998/namespace";
pub const NS_EMPTY_URI: &'static str    = "";

/// Denotes something which contains namespace URI mappings.
///
/// A URI mapping is a pair of type `(Option<&str>, &str)`, where the first item
/// is namespace prefix (`None` meaning default prefix) and the second item
/// is a URI mapped to the prefix.
pub trait NamespaceIterable<'a, I: Iterator<(Option<&'a str>, &'a str)>> {
    fn uri_mappings(&'a self) -> I;
}

/// Namespace is a map from prefixes to namespace URIs.
///
/// `None` prefix means no prefix (i.e. default namespace).
#[deriving(PartialEq, Clone)]
pub struct Namespace(pub HashMap<Option<String>, String>);

impl Namespace {
    /// Returns an empty namespace.
    #[inline]
    pub fn empty() -> Namespace { Namespace(HashMap::with_capacity(2)) }

    /// Checks whether this namespace is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Checks whether this namespace is essentially empty, that is, it does not contain
    /// anything but default mappings.
    pub fn is_essentially_empty(&self) -> bool {
        for (k, v) in self.0.iter() {
            match (k.as_ref().map(|k| k.as_slice()), v.as_slice()) {
                (None, u)    if u == NS_EMPTY_URI                         => {},
                (Some(p), u) if p == NS_XMLNS_PREFIX && u == NS_XMLNS_URI => {},
                (Some(p), u) if p == NS_XML_PREFIX   && u == NS_XML_URI   => {},
                _ => return false
            }
        }
        true
    }

    /// Puts a mapping into this namespace.
    ///
    /// This method does not override already existing mapping.
    ///
    /// Returns a boolean flag indicating whether the map already contained
    /// the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace);
    /// * `uri`    --- namespace URI.
    ///
    /// # Return value
    /// `true` if `prefix` has been inserted successfully; `false` if the `prefix`
    /// was already present in the namespace.
    pub fn put(&mut self, prefix: Option<String>, uri: String) -> bool {
        self.0.insert(prefix, uri)
    }

    /// Queries the namespace for the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace).
    ///
    /// # Return value
    /// Namespace URI corresponding to the given prefix, if it is present.
    pub fn get<'a>(&'a self, prefix: &Option<String>) -> Option<&'a str> {
        self.0.find(prefix).map(|s| s.as_slice())
    }
}

/// An iterator over mappings from prefixes to URIs in a namespace.
pub struct NamespaceMappings<'a> {
    entries: Entries<'a, Option<String>, String>
}

impl<'a> Iterator<(Option<&'a str>, &'a str)> for NamespaceMappings<'a> {
    fn next(&mut self) -> Option<(Option<&'a str>, &'a str)> {
        self.entries.next().map(|(prefix, uri)| {
            (prefix.as_ref().map(|p| p.as_slice()), uri.as_slice())
        })
    }
}

impl<'a> NamespaceIterable<'a, NamespaceMappings<'a>> for Namespace {
    fn uri_mappings(&'a self) -> NamespaceMappings<'a> {
        NamespaceMappings { entries: self.0.iter() }
    }
}

/// Namespace stack is a sequence of namespaces.
///
/// Namespace stack is used to represent cumulative namespace consisting of
/// combined namespaces from nested elements.
#[deriving(Clone, PartialEq)]
pub struct NamespaceStack(pub Vec<Namespace>);

impl NamespaceStack {
    /// Returns an empty namespace stack.
    #[inline]
    pub fn empty() -> NamespaceStack { NamespaceStack(Vec::with_capacity(2)) }

    /// Returns a namespace stack with default items in it.
    ///
    /// Default items are the following:
    ///
    /// * `xml` → `http://www.w3.org/XML/1998/namespace`;
    /// * `xmlns` → `http://www.w3.org/2000/xmlns/`.
    #[inline]
    pub fn default() -> NamespaceStack {
        let mut nst = NamespaceStack::empty();
        nst.push_empty();
        // xml namespace
        nst.put(Some(NS_XML_PREFIX.to_string()), NS_XML_URI.to_string());
        // xmlns namespace
        nst.put(Some(NS_XMLNS_PREFIX.to_string()), NS_XMLNS_URI.to_string());
        // empty namespace
        nst.put(None, NS_EMPTY_URI.to_string());
        nst
    }

    /// Adds an empty namespace to the top of this stack.
    #[inline]
    pub fn push_empty(&mut self) {
        self.0.push(Namespace::empty());
    }

    /// Removes a namespace at the top of the stack.
    ///
    /// Fails if the stack is empty.
    #[inline]
    pub fn pop(&mut self) -> Namespace {
        self.0.pop().unwrap()
    }

    /// Returns a namespace at the top of the stack, leaving the stack intact.
    ///
    /// Fails if the stack is empty.
    #[inline]
    pub fn peek<'a>(&'a mut self) -> &'a mut Namespace {
        self.0.last_mut().unwrap()
    }

    /// Puts a mapping into the topmost namespace in this stack.
    ///
    /// This method does not override a mapping in the topmost namespace if it is
    /// already present, however, it does not depend on other namespaces in the stack,
    /// so it is possible to put a mapping which is present in lower namespaces.
    ///
    /// Returns a boolean flag indicating whether the topmost namespace
    /// already contained the given prefix.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace);
    /// * `uri`    --- namespace URI.
    ///
    /// # Return value
    /// `true` if `prefix` has been inserted successfully; `false` if the `prefix`
    /// was already present in the namespace.
    #[inline]
    pub fn put(&mut self, prefix: Option<String>, uri: String) -> bool {
        self.0.last_mut().unwrap().put(prefix, uri)
    }

    /// Performs a search for the given prefix in the whole stack.
    ///
    /// This method walks the stack from top to bottom, querying each namespace
    /// in order for the given prefix. If none of the namespaces contains the prefix,
    /// `None` is returned.
    ///
    /// # Parameters
    /// * `prefix` --- namespace prefix (`None` means default namespace)
    #[inline]
    pub fn get<'a>(&'a self, prefix: &Option<String>) -> Option<&'a str> {
        for ns in self.0.iter().rev() {
            match ns.get(prefix) {
                None => {},
                r => return r,
            }
        }
        None
    }

    /// Combines this stack of namespaces into a single namespace.
    ///
    /// Namespaces are combined in left-to-right order, that is, rightmost namespace
    /// elements take priority over leftmost ones.
    pub fn squash(&self) -> Namespace {
        let mut result = HashMap::new();
        for ns in self.0.iter() {
            result.extend(ns.0.iter().map(|(k, v)| (k.clone(), v.to_string())));
        }
        Namespace(result)
    }
}

/// An iterator over mappings from prefixes to URIs in a namespace stack.
pub struct NamespaceStackMappings<'a> {
    namespaces: Rev<Items<'a, Namespace>>,
    current_namespace: Option<NamespaceMappings<'a>>,
    used_keys: HashSet<Option<&'a str>>
}

impl<'a> NamespaceStackMappings<'a> {
    fn to_next_namespace(&mut self) -> bool {
        self.current_namespace = self.namespaces.next().map(|ns| ns.uri_mappings());
        self.current_namespace.is_some()
    }
}

impl<'a> Iterator<(Option<&'a str>, &'a str)> for NamespaceStackMappings<'a> {
    fn next(&mut self) -> Option<(Option<&'a str>, &'a str)> {
        // If there is no current namespace and no next namespace, we're finished
        if self.current_namespace.is_none() && !self.to_next_namespace() {
            return None;
        }
        let next_item = self.current_namespace.as_mut().unwrap().next();

        match next_item {
            // There is an element in the current namespace
            Some((k, v)) => if self.used_keys.contains(&k) {
                // If the current key is used, go to the next one
                self.next()
            } else {
                // Otherwise insert the current key to the set of used keys and
                // return the mapping
                self.used_keys.insert(k);
                Some((k, v))
            },
            // Current namespace is exhausted
            None => if self.to_next_namespace() {
                // If there is next namespace, continue from it
                self.next()
            } else {
                // No next namespace, exiting
                None
            }
        }
    }
}

impl<'a> NamespaceIterable<'a, NamespaceStackMappings<'a>> for NamespaceStack {
    fn uri_mappings(&'a self) -> NamespaceStackMappings<'a> {
        NamespaceStackMappings {
            namespaces: self.0.iter().rev(),
            current_namespace: None,
            used_keys: HashSet::new()
        }
    }
}
<|endoftext|>


<commit_before>//! Executes compiled JMESPath bytecode instructions.

extern crate rustc_serialize;

use self::rustc_serialize::json::Json;

use ast::{Ast, Comparator, KeyValuePair};

#[derive(Clone, PartialEq, Debug)]
pub enum Opcode {
    // Pops N operands from stack and calls a built-in function by name.
    Call(u8, String),
    // Pops two elements from the stack and pushes true/false.
    Cmp(Comparator),
    // Duplicates the top of the stack
    Dup,
    // Pops TOS and iterates over each element if it is an array. If it
    // is not an array, then jumps to the given address.
    Each(usize),
    // Pops TOS and pushes the <operand> key value of TOS onto the stack.
    // If TOS is not an "object", null is pushed onto the stack.
    Field(String),
    // Pops TOS and pushes the <operand> index value of TOS onto the stack.
    // If TOS is not an "array", null is pushed onto the stack.
    Index(usize),
    // Pops TOS and pushes the <operand> index from the end of the length of
    // of TOS onto the stack. If TOS is not an "array", null is pushed onto
    /// the stack.
    NegativeIndex(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "array". Pushes T1 onto T2 and pushes T2 back onto the stack.
    InjectIdx(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "object". Adds T1 at the given operand key onto T2 and pushes T2
    // back onto the stack.
    InjectKey(String),
    // Unconditional jump to an address.
    Jump(usize),
    // Pops the TOS and jumps if 0 or false
    Jumpz(usize),
    // Pops the call stack
    Ret(usize),
    // Pops the TOS
    Pop,
    // Pushes a value onto the stack
    Push(Json),
    // Pops TOS and pushes true/false if the value is truthy.
    Truthy,
    // Pops TOS and pushes the type of the popped value.
    Type
}

pub type Program = Vec<Opcode>;

/// JMESPath VM
pub struct Vm<'a> {
    /// The compiled program.
    opcodes: &'a Program,
    /// The opcode index.
    index: usize,
    /// The VM stack consisting of JSON typed values.
    stack: Vec<Json>
}

impl<'a> Vm<'a> {
    /// Creates a new VM using the given program and data.
    pub fn new(opcodes: &'a Vec<Opcode>, data: Json) -> Vm<'a> {
        Vm {
            opcodes: opcodes,
            stack: vec![data],
            index: 0
        }
    }

    pub fn run(&mut self) -> Result<Json, String> {
        macro_rules! tos {
            () => {{
                self.stack.pop().unwrap_or(Json::Null)
            }};
        }
        while self.index < self.opcodes.len() {
            match &self.opcodes[self.index] {
                &Opcode::Push(ref j) => {
                    self.stack.push(j.clone());
                },
                &Opcode::Pop => {
                    self.stack.pop();
                },
                &Opcode::Dup => {
                    if self.stack.is_empty() {
                        self.stack.push(Json::Null);
                    } else {
                        let tos = self.stack.pop().unwrap();
                        self.stack.push(tos.clone());
                        self.stack.push(tos);
                    }
                },
                &Opcode::Field(ref f) => {
                    match tos!().find(f) {
                        Some(j) => self.stack.push(j.clone()),
                        _ => self.stack.push(Json::Null)
                    };
                },
                &Opcode::Index(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|_| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::NegativeIndex(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|i| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::Truthy => {
                    let tos = tos!();
                    self.stack.push(Json::Boolean(match tos {
                        Json::Boolean(b) if b == false => false,
                        Json::String(ref s) if s.len() == 0 => false,
                        Json::Array(ref a) if a.is_empty() => false,
                        Json::Object(ref o) if o.is_empty() => false,
                        Json::Null => false,
                        _ => true
                    }));
                },
                &Opcode::Type => {
                    let tos = tos!();
                    self.stack.push(match tos {
                        Json::Boolean(_) => Json::String("boolean".to_string()),
                        Json::String(_) => Json::String("string".to_string()),
                        Json::I64(_) | Json::U64(_) | Json::F64(_) => {
                            Json::String("number".to_string())
                        },
                        Json::Array(_) => Json::String("array".to_string()),
                        Json::Object(_) => Json::String("object".to_string()),
                        Json::Null => Json::String("null".to_string()),
                    });
                },
                &Opcode::Jump(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    self.index = address;
                    continue;
                },
                &Opcode::Jumpz(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    let tos = tos!();
                    self.index = match tos {
                        Json::Boolean(b) if b == false => address,
                        Json::I64(d) if d == 0 => address,
                        Json::F64(f) if f == 0f64 => address,
                        Json::U64(u) if u == 0 => address,
                        _ => self.index + 1
                    };
                    continue;
                },
                &Opcode::Cmp(ref cmp) => {
                    let (a, b) = (tos!(), tos!());
                    self.stack.push(Json::Boolean(match cmp {
                        &Comparator::Lt
                        | &Comparator::Lte
                        | &Comparator::Gt
                        | &Comparator::Gte => Vm::cmp_numbers(cmp, a, b),
                        &Comparator::Eq => a == b,
                        &Comparator::Ne => a != b,
                    }));
                },
                _ => panic!("Not implemented yet!")
            }
            self.index += 1;
        }

        Ok(self.stack.pop().unwrap_or(Json::Null))
    }

    /// Compares two JSON values, ensuring that each is a number.
    ///
    /// If the provided values are not numbers, then false is returned.
    /// Otherwise, a comparison is made between the two values.
    fn cmp_numbers(cmp: &Comparator, a: Json, b: Json) -> bool {
        if !a.is_number() || !b.is_number() {
            false
        } else {
            let (af, bf) = (a.as_f64().unwrap(), b.as_f64().unwrap());
            match cmp {
                &Comparator::Lt => af < bf,
                &Comparator::Lte => af <= bf,
                &Comparator::Gt => af > bf,
                &Comparator::Gte => af >= bf,
                _ => false
            }
        }
    }
}

#[cfg(test)]
mod test {
    extern crate rustc_serialize;
    use super::{Opcode, Vm};
    use ast::{Comparator};
    use self::rustc_serialize::json::Json;

    #[test] fn truthy_test() {
        let opcodes = vec![Opcode::Truthy];
        let tests = vec![("\"foo\"", true), ("\"\"", false),
                         ("[1]", true), ("[]", false),
                         ("{\"foo\": 1}", true), ("{}", false),
                         ("true", true), ("false", false),
                         ("null", false), ("1", true)];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, Json::from_str(js).unwrap());
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }

    #[test] fn pushes_and_pops() {
        let opcodes = vec![Opcode::Pop, Opcode::Push(Json::String("foo".to_string()))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::String("foo".to_string()), vm.run().unwrap());
    }

    #[test] fn dups_tos() {
        let opcodes = vec![Opcode::Dup, Opcode::Pop];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::I64(1), vm.run().unwrap());
    }

    #[test] fn dups_null() {
        let opcodes = vec![Opcode::Pop, Opcode::Dup];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_fields() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"a\": 1}").unwrap());
        assert_eq!(Json::U64(1), vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_missing_field() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"b\": 1}").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_field_from_non_object() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::Index(0), Json::U64(0)),
                         (Opcode::Index(1), Json::U64(1)),
                         (Opcode::Index(2), Json::Null),
                         (Opcode::Index(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn extracts_negative_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::NegativeIndex(0), Json::U64(1)),
                         (Opcode::NegativeIndex(1), Json::U64(0)),
                         (Opcode::NegativeIndex(2), Json::Null),
                         (Opcode::NegativeIndex(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn pushes_type_on_stack() {
        let js = Json::from_str("[0, 1]").unwrap();
        let opcodes = vec![Opcode::Type];
        let tests = vec![(Json::String("a".to_string()), Json::String("string".to_string())),
                         (Json::Array(vec![]), Json::String("array".to_string())),
                         (Json::from_str("{}").unwrap(), Json::String("object".to_string())),
                         (Json::I64(1), Json::String("number".to_string())),
                         (Json::Null, Json::String("null".to_string())),
                         (Json::Boolean(true), Json::String("boolean".to_string()))];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, js);
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn can_jump() {
        let opcodes = vec![Opcode::Jump(2),
                           Opcode::Jump(3),
                           Opcode::Jump(1),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn checks_if_jump_is_valid() {
        let opcodes = vec![Opcode::Jump(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn checks_if_jumpz_is_valid() {
        let opcodes = vec![Opcode::Jumpz(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn can_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Index(0),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(false));
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn can_not_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Jump(100),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(true));
        assert_eq!(Err("Invalid jump address 100".to_string()), vm.run());
    }

    #[test] fn compares_conditionally() {
        let tests = vec![("[0, 1]", Comparator::Lt, false),
                         ("[0, 1]", Comparator::Lte, false),
                         ("[0, 1]", Comparator::Gt, true),
                         ("[0, 1]", Comparator::Gte, true),
                         ("[0, 1]", Comparator::Eq, false),
                         ("[0, \"a\"]", Comparator::Ne, true),
                         ("[0, \"a\"]", Comparator::Eq, false),
                         ("[\"a\", \"a\"]", Comparator::Eq, true),
                         ("[\"a\", \"b\"]", Comparator::Eq, false),
                         ("[\"a\", \"b\"]", Comparator::Ne, true)];
        for (js, cmp, result) in tests {
            let parsed = Json::from_str(js).unwrap();
            let opcodes = vec![Opcode::Push(parsed[0].clone()),
                               Opcode::Push(parsed[1].clone()),
                               Opcode::Cmp(cmp)];
            let mut vm = Vm::new(&opcodes, Json::Null);
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }
}
<commit_msg>Adding inject and insert opcodes<commit_after>//! Executes compiled JMESPath bytecode instructions.

extern crate rustc_serialize;

use self::rustc_serialize::json::Json;

use ast::{Ast, Comparator, KeyValuePair};

#[derive(Clone, PartialEq, Debug)]
pub enum Opcode {
    // Pops N operands from stack and calls a built-in function by name.
    Call(u8, String),
    // Pops two elements from the stack and pushes true/false.
    Cmp(Comparator),
    // Duplicates the top of the stack
    Dup,
    // Pops TOS and iterates over each element if it is an array. If it
    // is not an array, then jumps to the given address.
    Each(usize),
    // Pops TOS and pushes the <operand> key value of TOS onto the stack.
    // If TOS is not an "object", null is pushed onto the stack.
    Field(String),
    // Pops TOS and pushes the <operand> index value of TOS onto the stack.
    // If TOS is not an "array", null is pushed onto the stack.
    Index(usize),
    // Pops TOS and pushes the <operand> index from the end of the length of
    // of TOS onto the stack. If TOS is not an "array", null is pushed onto
    /// the stack.
    NegativeIndex(usize),
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "array". Pushes T1 onto T2 and pushes T2 back onto the stack.
    Insert,
    // Pops two elements, T1 and T2 where T1 is a value and T2 is an
    // "object". Adds T1 at the given operand key onto T2 and pushes T2
    // back onto the stack.
    InjectKey(String),
    // Unconditional jump to an address.
    Jump(usize),
    // Pops the TOS and jumps if 0 or false
    Jumpz(usize),
    // Pops the call stack
    Ret(usize),
    // Pops the TOS
    Pop,
    // Pushes a value onto the stack
    Push(Json),
    // Pops TOS and pushes true/false if the value is truthy.
    Truthy,
    // Pops TOS and pushes the type of the popped value.
    Type
}

pub type Program = Vec<Opcode>;

/// JMESPath VM
pub struct Vm<'a> {
    /// The compiled program.
    opcodes: &'a Program,
    /// The opcode index.
    index: usize,
    /// The VM stack consisting of JSON typed values.
    stack: Vec<Json>
}

impl<'a> Vm<'a> {
    /// Creates a new VM using the given program and data.
    pub fn new(opcodes: &'a Vec<Opcode>, data: Json) -> Vm<'a> {
        Vm {
            opcodes: opcodes,
            stack: vec![data],
            index: 0
        }
    }

    pub fn run(&mut self) -> Result<Json, String> {
        macro_rules! tos {
            () => {{
                self.stack.pop().unwrap_or(Json::Null)
            }};
        }
        while self.index < self.opcodes.len() {
            match &self.opcodes[self.index] {
                &Opcode::Push(ref j) => {
                    self.stack.push(j.clone());
                },
                &Opcode::Pop => {
                    self.stack.pop();
                },
                &Opcode::Dup => {
                    if self.stack.is_empty() {
                        self.stack.push(Json::Null);
                    } else {
                        let tos = self.stack.pop().unwrap();
                        self.stack.push(tos.clone());
                        self.stack.push(tos);
                    }
                },
                &Opcode::Field(ref f) => {
                    match tos!().find(f) {
                        Some(j) => self.stack.push(j.clone()),
                        _ => self.stack.push(Json::Null)
                    };
                },
                &Opcode::Index(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|_| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::NegativeIndex(i) => {
                    let tos = tos!();
                    self.stack.push(tos.as_array()
                        .and_then(|a| a.len().checked_sub(i + 1))
                        .and_then(|i| Some(tos[i].clone()))
                        .unwrap_or(Json::Null));
                },
                &Opcode::Truthy => {
                    let tos = tos!();
                    self.stack.push(Json::Boolean(match tos {
                        Json::Boolean(b) if b == false => false,
                        Json::String(ref s) if s.len() == 0 => false,
                        Json::Array(ref a) if a.is_empty() => false,
                        Json::Object(ref o) if o.is_empty() => false,
                        Json::Null => false,
                        _ => true
                    }));
                },
                &Opcode::Type => {
                    let tos = tos!();
                    self.stack.push(match tos {
                        Json::Boolean(_) => Json::String("boolean".to_string()),
                        Json::String(_) => Json::String("string".to_string()),
                        Json::I64(_) | Json::U64(_) | Json::F64(_) => {
                            Json::String("number".to_string())
                        },
                        Json::Array(_) => Json::String("array".to_string()),
                        Json::Object(_) => Json::String("object".to_string()),
                        Json::Null => Json::String("null".to_string()),
                    });
                },
                &Opcode::Jump(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    self.index = address;
                    continue;
                },
                &Opcode::Jumpz(address) => {
                    if self.opcodes.get(address).is_none() {
                        return Err(format!("Invalid jump address {}", address));
                    }
                    let tos = tos!();
                    self.index = match tos {
                        Json::Boolean(b) if b == false => address,
                        Json::I64(d) if d == 0 => address,
                        Json::F64(f) if f == 0f64 => address,
                        Json::U64(u) if u == 0 => address,
                        _ => self.index + 1
                    };
                    continue;
                },
                &Opcode::Cmp(ref cmp) => {
                    let (a, b) = (tos!(), tos!());
                    self.stack.push(Json::Boolean(match cmp {
                        &Comparator::Lt
                        | &Comparator::Lte
                        | &Comparator::Gt
                        | &Comparator::Gte => Vm::cmp_numbers(cmp, a, b),
                        &Comparator::Eq => a == b,
                        &Comparator::Ne => a != b,
                    }));
                },
                &Opcode::InjectKey(ref k) => {
                    let (new_val, mut coll) = (tos!(), tos!());
                    self.stack.push(coll.as_object_mut()
                        .map(|obj| obj.insert(k.clone(), new_val))
                        .map(|_| coll)
                        .unwrap_or(Json::Null));
                },
                &Opcode::Insert => {
                    let (new_val, mut coll) = (tos!(), tos!());
                    self.stack.push(coll.as_array_mut()
                        .map(|arr| arr.push(new_val))
                        .map(|_| coll)
                        .unwrap_or(Json::Null));
                },
                _ => panic!("Not implemented yet!")
            }
            self.index += 1;
        }

        Ok(self.stack.pop().unwrap_or(Json::Null))
    }

    /// Compares two JSON values, ensuring that each is a number.
    ///
    /// If the provided values are not numbers, then false is returned.
    /// Otherwise, a comparison is made between the two values.
    fn cmp_numbers(cmp: &Comparator, a: Json, b: Json) -> bool {
        if !a.is_number() || !b.is_number() {
            false
        } else {
            let (af, bf) = (a.as_f64().unwrap(), b.as_f64().unwrap());
            match cmp {
                &Comparator::Lt => af < bf,
                &Comparator::Lte => af <= bf,
                &Comparator::Gt => af > bf,
                &Comparator::Gte => af >= bf,
                _ => false
            }
        }
    }
}

#[cfg(test)]
mod test {
    extern crate rustc_serialize;
    use super::{Opcode, Vm};
    use ast::{Comparator};
    use self::rustc_serialize::json::Json;

    #[test] fn truthy_test() {
        let opcodes = vec![Opcode::Truthy];
        let tests = vec![("\"foo\"", true), ("\"\"", false),
                         ("[1]", true), ("[]", false),
                         ("{\"foo\": 1}", true), ("{}", false),
                         ("true", true), ("false", false),
                         ("null", false), ("1", true)];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, Json::from_str(js).unwrap());
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }

    #[test] fn pushes_and_pops() {
        let opcodes = vec![Opcode::Pop, Opcode::Push(Json::String("foo".to_string()))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::String("foo".to_string()), vm.run().unwrap());
    }

    #[test] fn dups_tos() {
        let opcodes = vec![Opcode::Dup, Opcode::Pop];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::I64(1), vm.run().unwrap());
    }

    #[test] fn dups_null() {
        let opcodes = vec![Opcode::Pop, Opcode::Dup];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_fields() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"a\": 1}").unwrap());
        assert_eq!(Json::U64(1), vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_missing_field() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{\"b\": 1}").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn pushes_null_when_extracting_field_from_non_object() {
        let opcodes = vec![Opcode::Field("a".to_string())];
        let mut vm = Vm::new(&opcodes, Json::I64(1));
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn extracts_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::Index(0), Json::U64(0)),
                         (Opcode::Index(1), Json::U64(1)),
                         (Opcode::Index(2), Json::Null),
                         (Opcode::Index(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn extracts_negative_indices() {
        let js = Json::from_str("[0, 1]").unwrap();
        let tests = vec![(Opcode::NegativeIndex(0), Json::U64(1)),
                         (Opcode::NegativeIndex(1), Json::U64(0)),
                         (Opcode::NegativeIndex(2), Json::Null),
                         (Opcode::NegativeIndex(3), Json::Null)];
        for (op, result) in tests {
            let opcodes = vec![op];
            let mut vm = Vm::new(&opcodes, js.clone());
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn pushes_type_on_stack() {
        let js = Json::from_str("[0, 1]").unwrap();
        let opcodes = vec![Opcode::Type];
        let tests = vec![(Json::String("a".to_string()), Json::String("string".to_string())),
                         (Json::Array(vec![]), Json::String("array".to_string())),
                         (Json::from_str("{}").unwrap(), Json::String("object".to_string())),
                         (Json::I64(1), Json::String("number".to_string())),
                         (Json::Null, Json::String("null".to_string())),
                         (Json::Boolean(true), Json::String("boolean".to_string()))];
        for (js, result) in tests {
            let mut vm = Vm::new(&opcodes, js);
            assert_eq!(result, vm.run().unwrap());
        }
    }

    #[test] fn can_jump() {
        let opcodes = vec![Opcode::Jump(2),
                           Opcode::Jump(3),
                           Opcode::Jump(1),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn checks_if_jump_is_valid() {
        let opcodes = vec![Opcode::Jump(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn checks_if_jumpz_is_valid() {
        let opcodes = vec![Opcode::Jumpz(2)];
        let mut vm = Vm::new(&opcodes, Json::Null);
        assert_eq!(Err("Invalid jump address 2".to_string()), vm.run());
    }

    #[test] fn can_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Index(0),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(false));
        assert_eq!(Json::Boolean(true), vm.run().unwrap());
    }

    #[test] fn can_not_jump_conditionally() {
        let opcodes = vec![Opcode::Jumpz(2),
                           Opcode::Jump(100),
                           Opcode::Push(Json::Boolean(true))];
        let mut vm = Vm::new(&opcodes, Json::Boolean(true));
        assert_eq!(Err("Invalid jump address 100".to_string()), vm.run());
    }

    #[test] fn compares_conditionally() {
        let tests = vec![("[0, 1]", Comparator::Lt, false),
                         ("[0, 1]", Comparator::Lte, false),
                         ("[0, 1]", Comparator::Gt, true),
                         ("[0, 1]", Comparator::Gte, true),
                         ("[0, 1]", Comparator::Eq, false),
                         ("[0, \"a\"]", Comparator::Ne, true),
                         ("[0, \"a\"]", Comparator::Eq, false),
                         ("[\"a\", \"a\"]", Comparator::Eq, true),
                         ("[\"a\", \"b\"]", Comparator::Eq, false),
                         ("[\"a\", \"b\"]", Comparator::Ne, true)];
        for (js, cmp, result) in tests {
            let parsed = Json::from_str(js).unwrap();
            let opcodes = vec![Opcode::Push(parsed[0].clone()),
                               Opcode::Push(parsed[1].clone()),
                               Opcode::Cmp(cmp)];
            let mut vm = Vm::new(&opcodes, Json::Null);
            assert_eq!(Json::Boolean(result), vm.run().unwrap());
        }
    }

    #[test] fn injects_key_into_valid_object() {
        let opcodes = vec![Opcode::Push(Json::String("bar".to_string())),
                           Opcode::InjectKey("foo".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("{}").unwrap());
        let result = vm.run().unwrap();
        assert_eq!("{\"foo\":\"bar\"}", format!("{}", result));
    }

    #[test] fn pushes_false_when_injecting_key_onto_non_object() {
        let opcodes = vec![Opcode::Push(Json::String("bar".to_string())),
                           Opcode::InjectKey("foo".to_string())];
        let mut vm = Vm::new(&opcodes, Json::from_str("[]").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }

    #[test] fn inserts_into_valid_array() {
        let opcodes = vec![Opcode::Push(Json::String("bar".to_string())),
                           Opcode::Insert,
                           Opcode::Push(Json::String("baz".to_string())),
                           Opcode::Insert];
        let mut vm = Vm::new(&opcodes, Json::from_str("[]").unwrap());
        let result = vm.run().unwrap();
        assert_eq!("[\"bar\",\"baz\"]", format!("{}", result));
    }

    #[test] fn pushes_false_when_inserting_onto_non_array() {
        let opcodes = vec![Opcode::Push(Json::String("bar".to_string())),
                           Opcode::Insert];
        let mut vm = Vm::new(&opcodes, Json::from_str("{}").unwrap());
        assert_eq!(Json::Null, vm.run().unwrap());
    }
}
<|endoftext|>