<commit_before>/// Struct to represent an inheritable opacity property
/// <https://www.w3.org/TR/SVG/masking.html#OpacityProperty>

use ::cssparser::{Parser, ParserInput, Token};
use ::libc;

use std::str::FromStr;

use ::glib::translate::*;

use parsers::ParseError;
use error::*;

// Keep this in sync with rsvg-css.h:RsvgOpacityKind
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum OpacityKind {
    Inherit,
    Specified,
    ParseError
}

// Keep this in sync with rsvg-css.h:RsvgOpacitySpec
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct OpacitySpec {
    kind: OpacityKind,
    opacity: u8
}

// This is the Rust version of the above
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Opacity {
    Inherit,
    Specified (f64)
}

impl From<Result<Opacity, AttributeError>> for OpacitySpec {
    fn from (result: Result<Opacity, AttributeError>) -> OpacitySpec {
        match result {
            Ok (Opacity::Inherit) =>
                OpacitySpec {
                    kind: OpacityKind::Inherit,
                    opacity: 0
                },

            Ok (Opacity::Specified (val)) =>
                OpacitySpec {
                    kind: OpacityKind::Specified,
                    opacity: opacity_to_u8 (val)
                },

            _ =>
                OpacitySpec {
                    kind: OpacityKind::ParseError,
                    opacity: 0
                }
        }
    }
}

pub fn opacity_to_u8 (val: f64) -> u8 {
    (val * 255.0 + 0.5).floor () as u8
}

fn make_err () -> AttributeError {
    AttributeError::Parse (ParseError::new ("expected 'inherit' or number"))
}

impl FromStr for Opacity {
    type Err = AttributeError;

    fn from_str (s: &str) -> Result<Opacity, AttributeError> {
        let mut input = ParserInput::new (s);
        let mut parser = Parser::new (&mut input);

        let opacity = {
            let token = parser.next ()
                .map_err (|_| make_err ())?;

            match *token {
                Token::Ident (ref cow) => {
                    if cow.as_ref () == "inherit" {
                        Opacity::Inherit
                    } else {
                        return Err (make_err ());
                    }
                },

                Token::Number { value, .. } => {
                    if value < 0.0 {
                        Opacity::Specified (0.0)
                    } else if value > 1.0 {
                        Opacity::Specified (1.0)
                    } else {
                        Opacity::Specified (f64::from(value))
                    }
                },

                _ => return Err (make_err ())
            }
        };

        parser.expect_exhausted ().map_err (|_| make_err ())?;

        Ok (opacity)
    }
}

impl Opacity {
    pub fn from_opacity_spec (spec: &OpacitySpec) -> Result <Opacity, AttributeError> {
        match *spec {
            OpacitySpec { kind: OpacityKind::Inherit, .. }      => Ok (Opacity::Inherit),

            OpacitySpec { kind: OpacityKind::Specified, opacity } => Ok (Opacity::Specified (f64::from(opacity) / 255.0)),

            OpacitySpec { kind: OpacityKind::ParseError, .. } => Err (AttributeError::Parse (ParseError::new ("parse error")))
        }
    }
}

#[no_mangle]
pub extern fn rsvg_css_parse_opacity (string: *const libc::c_char) -> OpacitySpec {
    let s = unsafe { String::from_glib_none (string) };

    OpacitySpec::from (Opacity::from_str (&s))
}

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

    #[test]
    fn parses_inherit () {
        assert_eq! (Opacity::from_str ("inherit"), Ok (Opacity::Inherit));
    }

    #[test]
    fn parses_number () {
        assert_eq! (Opacity::from_str ("0"),   Ok (Opacity::Specified (0.0)));
        assert_eq! (Opacity::from_str ("1"),   Ok (Opacity::Specified (1.0)));
        assert_eq! (Opacity::from_str ("0.5"), Ok (Opacity::Specified (0.5)));
    }

    #[test]
    fn parses_out_of_range_number () {
        assert_eq! (Opacity::from_str ("-10"), Ok (Opacity::Specified (0.0)));
        assert_eq! (Opacity::from_str ("10"),  Ok (Opacity::Specified (1.0)));
    }

    #[test]
    fn errors_on_invalid_input () {
        assert! (is_parse_error (&Opacity::from_str ("")));
        assert! (is_parse_error (&Opacity::from_str ("foo")));
        assert! (is_parse_error (&Opacity::from_str ("-x")));
    }

    #[test]
    fn errors_on_extra_input () {
        assert! (is_parse_error (&Opacity::from_str ("inherit a million dollars")));
        assert! (is_parse_error (&Opacity::from_str ("0.0foo")));
    }

    fn parse (s: &str) -> OpacitySpec {
        rsvg_css_parse_opacity (s.to_glib_none ().0)
    }

    #[test]
    fn converts_result_to_opacity_spec () {
        assert_eq! (parse ("inherit"),
                    OpacitySpec { kind: OpacityKind::Inherit,
                                  opacity: 0 });

        assert_eq! (parse ("0"),
                    OpacitySpec { kind: OpacityKind::Specified,
                                  opacity: 0 });
        assert_eq! (parse ("1"),
                    OpacitySpec { kind: OpacityKind::Specified,
                                  opacity: 255 });

        assert_eq! (parse ("foo"),
                    OpacitySpec { kind: OpacityKind::ParseError,
                                  opacity: 0 });
    }

    fn test_roundtrip (s: &str) {
        let result = Opacity::from_str (s);
        let result2 = result.clone ();
        let spec = OpacitySpec::from (result2);

        if result.is_ok () {
            assert_eq! (Opacity::from_opacity_spec (&spec), result);
        } else {
            assert! (Opacity::from_opacity_spec (&spec).is_err ());
        }
    }

    #[test]
    fn roundtrips () {
        test_roundtrip ("inherit");
        test_roundtrip ("0");
        test_roundtrip ("1.0");
        test_roundtrip ("chilaquil");
    }
}
<commit_msg>rsvg_css_parse_opacity(): Remove an allocation<commit_after>/// Struct to represent an inheritable opacity property
/// <https://www.w3.org/TR/SVG/masking.html#OpacityProperty>

use ::cssparser::{Parser, ParserInput, Token};
use ::libc;

use std::ffi::CStr;
use std::str::FromStr;

use parsers::ParseError;
use error::*;

// Keep this in sync with rsvg-css.h:RsvgOpacityKind
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum OpacityKind {
    Inherit,
    Specified,
    ParseError
}

// Keep this in sync with rsvg-css.h:RsvgOpacitySpec
#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct OpacitySpec {
    kind: OpacityKind,
    opacity: u8
}

// This is the Rust version of the above
#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Opacity {
    Inherit,
    Specified (f64)
}

impl From<Result<Opacity, AttributeError>> for OpacitySpec {
    fn from (result: Result<Opacity, AttributeError>) -> OpacitySpec {
        match result {
            Ok (Opacity::Inherit) =>
                OpacitySpec {
                    kind: OpacityKind::Inherit,
                    opacity: 0
                },

            Ok (Opacity::Specified (val)) =>
                OpacitySpec {
                    kind: OpacityKind::Specified,
                    opacity: opacity_to_u8 (val)
                },

            _ =>
                OpacitySpec {
                    kind: OpacityKind::ParseError,
                    opacity: 0
                }
        }
    }
}

pub fn opacity_to_u8 (val: f64) -> u8 {
    (val * 255.0 + 0.5).floor () as u8
}

fn make_err () -> AttributeError {
    AttributeError::Parse (ParseError::new ("expected 'inherit' or number"))
}

impl FromStr for Opacity {
    type Err = AttributeError;

    fn from_str (s: &str) -> Result<Opacity, AttributeError> {
        let mut input = ParserInput::new (s);
        let mut parser = Parser::new (&mut input);

        let opacity = {
            let token = parser.next ()
                .map_err (|_| make_err ())?;

            match *token {
                Token::Ident (ref cow) => {
                    if cow.as_ref () == "inherit" {
                        Opacity::Inherit
                    } else {
                        return Err (make_err ());
                    }
                },

                Token::Number { value, .. } => {
                    if value < 0.0 {
                        Opacity::Specified (0.0)
                    } else if value > 1.0 {
                        Opacity::Specified (1.0)
                    } else {
                        Opacity::Specified (f64::from(value))
                    }
                },

                _ => return Err (make_err ())
            }
        };

        parser.expect_exhausted ().map_err (|_| make_err ())?;

        Ok (opacity)
    }
}

impl Opacity {
    pub fn from_opacity_spec (spec: &OpacitySpec) -> Result <Opacity, AttributeError> {
        match *spec {
            OpacitySpec { kind: OpacityKind::Inherit, .. }      => Ok (Opacity::Inherit),

            OpacitySpec { kind: OpacityKind::Specified, opacity } => Ok (Opacity::Specified (f64::from(opacity) / 255.0)),

            OpacitySpec { kind: OpacityKind::ParseError, .. } => Err (AttributeError::Parse (ParseError::new ("parse error")))
        }
    }
}

#[no_mangle]
pub extern fn rsvg_css_parse_opacity (string: *const libc::c_char) -> OpacitySpec {
    // we can unwrap because libxml2 already validated this for UTF-8
    let s = unsafe { CStr::from_ptr(string).to_str().unwrap() };

    OpacitySpec::from(Opacity::from_str(s))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;
    use glib::translate::*;

    #[test]
    fn parses_inherit () {
        assert_eq! (Opacity::from_str ("inherit"), Ok (Opacity::Inherit));
    }

    #[test]
    fn parses_number () {
        assert_eq! (Opacity::from_str ("0"),   Ok (Opacity::Specified (0.0)));
        assert_eq! (Opacity::from_str ("1"),   Ok (Opacity::Specified (1.0)));
        assert_eq! (Opacity::from_str ("0.5"), Ok (Opacity::Specified (0.5)));
    }

    #[test]
    fn parses_out_of_range_number () {
        assert_eq! (Opacity::from_str ("-10"), Ok (Opacity::Specified (0.0)));
        assert_eq! (Opacity::from_str ("10"),  Ok (Opacity::Specified (1.0)));
    }

    #[test]
    fn errors_on_invalid_input () {
        assert! (is_parse_error (&Opacity::from_str ("")));
        assert! (is_parse_error (&Opacity::from_str ("foo")));
        assert! (is_parse_error (&Opacity::from_str ("-x")));
    }

    #[test]
    fn errors_on_extra_input () {
        assert! (is_parse_error (&Opacity::from_str ("inherit a million dollars")));
        assert! (is_parse_error (&Opacity::from_str ("0.0foo")));
    }

    fn parse (s: &str) -> OpacitySpec {
        rsvg_css_parse_opacity (s.to_glib_none ().0)
    }

    #[test]
    fn converts_result_to_opacity_spec () {
        assert_eq! (parse ("inherit"),
                    OpacitySpec { kind: OpacityKind::Inherit,
                                  opacity: 0 });

        assert_eq! (parse ("0"),
                    OpacitySpec { kind: OpacityKind::Specified,
                                  opacity: 0 });
        assert_eq! (parse ("1"),
                    OpacitySpec { kind: OpacityKind::Specified,
                                  opacity: 255 });

        assert_eq! (parse ("foo"),
                    OpacitySpec { kind: OpacityKind::ParseError,
                                  opacity: 0 });
    }

    fn test_roundtrip (s: &str) {
        let result = Opacity::from_str (s);
        let result2 = result.clone ();
        let spec = OpacitySpec::from (result2);

        if result.is_ok () {
            assert_eq! (Opacity::from_opacity_spec (&spec), result);
        } else {
            assert! (Opacity::from_opacity_spec (&spec).is_err ());
        }
    }

    #[test]
    fn roundtrips () {
        test_roundtrip ("inherit");
        test_roundtrip ("0");
        test_roundtrip ("1.0");
        test_roundtrip ("chilaquil");
    }
}
<|endoftext|>


<commit_before>use crate::util::dirs;
use anyhow::{bail, Result};
use futures::StreamExt;
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use strum::{Display, EnumString, EnumVariantNames};

#[derive(Debug, Display, EnumString, EnumVariantNames, PartialEq, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum Kind {
    #[cfg(any(feature = "plugins-docker"))]
    Docker,
    #[cfg(any(feature = "plugins-binary"))]
    Binary,
    #[cfg(any(feature = "plugins-package"))]
    Package,
}

/// Split a plugin spec (ie. alias/name and, optionally, version) into alias and version
pub fn spec_to_alias_version(spec: &str) -> (String, String) {
    if spec.contains('@') {
        let parts: Vec<&str> = spec.split('@').collect();
        (parts[0].into(), parts[1].into())
    } else {
        (spec.into(), "latest".into())
    }
}

/// Resolve a plugin alias to a plugin name
///
/// If the provided string is a registered alias then returns the corresponding
/// plugin name, otherwise assumes the string is a plugin name
pub fn alias_to_name(alias: &str, aliases: &HashMap<String, String>) -> String {
    match aliases.get(alias) {
        Some(name) => name,
        None => alias,
    }
    .into()
}

/// Get the path to the manifest file of a plugin
pub fn manifest_path(name: &str) -> Result<PathBuf> {
    Ok(dirs::plugins(false)?.join(format!("{}.json", name)))
}

/// Get the manifest for an installed plugin
pub fn manifest(name: &str) -> Result<String> {
    // TODO: deserialize JSON into a manifest object
    match fs::read_to_string(manifest_path(name)?) {
        Ok(json) => Ok(json),
        Err(_) => bail!("Plugin '{}' is not installed", name),
    }
}

/// Test whether a plugin is installed
pub fn is_installed(name: &str) -> Result<bool> {
    Ok(manifest_path(name)?.exists())
}

/// List installed manifests
pub fn list() -> Result<Vec<String>> {
    // TODO: Return a Vec<Manifest>
    let dir = dirs::plugins(true)?;
    let mut plugins: Vec<String> = vec![];
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if let Some(ext) = path.extension() {
            if ext == "json" {
                let name = path.file_stem().unwrap().to_string_lossy().to_string();
                plugins.push(name)
            }
        }
    }
    Ok(plugins)
}

/// Add a plugin as a pulled Docker image
///
/// For this to succeed must be able to connect to the local
/// Docker server and be able to pull an image with corresponding
/// name.
#[cfg(any(feature = "plugins-docker"))]
pub async fn install_docker(name: &str, version: &str) -> Result<()> {
    let docker = bollard::Docker::connect_with_local_defaults()?;

    let image = if name.contains('/') {
        name.to_string()
    } else {
        format!("stencila/{}", name)
    };
    let image = format!("{}:{}", image, version);

    // Pull the image (by creating an image from it)
    let mut stream = docker.create_image(
        Some(bollard::image::CreateImageOptions {
            from_image: image.clone(),
            ..Default::default()
        }),
        None,
        None,
    );
    while let Some(item) = stream.next().await {
        match item {
            Ok(info) => {
                if let Some(error) = info.error {
                    bail!("{}", error)
                } else if let Some(status) = info.status {
                    // TODO display of status and progress could be improved and displayed
                    // per id (layer) as in docker CLI
                    // println!("{:?} {:?} {:?}", info.id, info.progress, info.progress_detail);
                    if let Some(progress) = info.progress {
                        tracing::info!("{} {}", status, progress)
                    } else {
                        tracing::info!("{}", status)
                    }
                }
            }
            Err(error) => match error {
                bollard::errors::Error::DockerResponseNotFoundError { .. } => {
                    bail!("Unable to find Docker image '{}'", image)
                }
                _ => bail!("{}", error),
            },
        }
    }

    // Create a container to obtain manifest from
    let container_name: String = rand::thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .take(24)
        .map(char::from)
        .collect();
    let response = docker
        .create_container(
            Some(bollard::container::CreateContainerOptions {
                name: &container_name,
            }),
            bollard::container::Config {
                image: Some(image.clone()),
                cmd: Some(vec!["manifest".to_string()]),
                ..Default::default()
            },
        )
        .await?;
    if response.warnings.len() > 0 {
        for warning in response.warnings {
            tracing::warn!("When creating container: {}", warning);
        }
    }

    // Start the container
    docker
        .start_container(
            &container_name,
            None::<bollard::container::StartContainerOptions<String>>,
        )
        .await?;

    // Capture the container output into `manifest`
    let mut stream = docker.logs::<String>(
        &container_name,
        Some(bollard::container::LogsOptions {
            follow: true,
            stderr: true,
            stdout: true,
            ..Default::default()
        }),
    );
    let mut stdout = vec![];
    let mut stderr = vec![];
    while let Some(output) = stream.next().await {
        if let Ok(output) = output {
            match output {
                bollard::container::LogOutput::StdOut { message } => {
                    stdout.extend_from_slice(&message)
                }
                bollard::container::LogOutput::StdErr { message } => {
                    stderr.extend_from_slice(&message)
                }
                _ => {}
            }
        }
    }

    if stderr.len() > 0 {
        tracing::warn!("{}", std::str::from_utf8(&stderr)?);
    }

    let manifest = if stdout.len() > 0 {
        match std::str::from_utf8(&stdout) {
            Ok(stdout) => stdout,
            Err(error) => bail!("Error converting stream to UTF8: {}", error),
        }
    } else {
        bail!(
            "No output from Docker container manifest command; is {} a Stencila plugin?",
            image
        )
    };

    // TODO: desrialize manifest and then serialize to file
    fs::write(manifest_path(name)?, manifest)?;

    Ok(())
}

/// Add a plugin as a downloaded binary
pub fn install_binary(name: &str, version: &str) -> Result<()> {
    // TODO
    bail!("Unable to add plugin '{}@{}' as binary", name, version)
}

/// Add a plugin as a programing language package
pub fn install_package(name: &str, version: &str) -> Result<()> {
    // TODO
    bail!(
        "Unable to add plugin '{}@{}' as programming language package",
        name,
        version
    )
}

/// Add a plugin
pub async fn add(plugin: &str, kinds: &Vec<Kind>, aliases: &HashMap<String, String>) -> Result<()> {
    let (alias, version) = spec_to_alias_version(plugin);
    let name = alias_to_name(&alias, aliases);

    if is_installed(&name)? {
        bail!(
            "Plugin '{}' is already installed, consider using `stencila plugins upgrade {}`",
            name,
            name
        );
    }

    for kind in kinds {
        let result = match kind {
            Kind::Docker => install_docker(&name, &version).await,
            Kind::Binary => install_binary(&name, &version),
            Kind::Package => install_package(&name, &version),
        };
        match result {
            // Success, so just return now
            Ok(_) => return Ok(()),
            // Error, keep trying other kinds, or if there is
            // only one kind, return that error.
            Err(error) => {
                if kinds.len() == 1 {
                    return Err(error);
                }
            }
        }
    }

    bail!(
        "Unable to install plugin '{}', tried kinds {:?}",
        plugin,
        kinds
    )
}

/// Add a list of plugins
pub async fn add_list(
    plugins: Vec<String>,
    kinds: Vec<Kind>,
    aliases: HashMap<String, String>,
) -> Result<()> {
    for plugin in plugins {
        match add(&plugin, &kinds, &aliases).await {
            Ok(_) => tracing::info!("Added plugin {}", plugin),
            Err(error) => tracing::error!("{}", error),
        }
    }
    Ok(())
}

/// Remove a plugin
pub fn remove(plugin: &str, aliases: &HashMap<String, String>) -> Result<()> {
    let (alias, ..) = spec_to_alias_version(plugin);
    let name = alias_to_name(&alias, aliases);

    let file = manifest_path(&name)?;
    if file.exists() {
        fs::remove_file(file)?;
    }

    Ok(())
}

/// Remove a list of plugins
pub fn remove_list(plugins: Vec<String>, aliases: HashMap<String, String>) -> Result<()> {
    for plugin in plugins {
        match remove(&plugin, &aliases) {
            Ok(_) => tracing::info!("Removed plugin {}", plugin),
            Err(error) => tracing::error!("{}", error),
        }
    }
    Ok(())
}

#[cfg(feature = "config")]
pub mod config {
    use super::*;
    use validator::Validate;

    #[derive(Debug, PartialEq, Deserialize, Serialize, Validate)]
    pub struct Config {
        #[serde(default = "default_kinds")]
        pub kinds: Vec<Kind>,

        #[serde(default = "default_aliases")]
        pub aliases: HashMap<String, String>,
    }

    /// Default configuration
    ///
    /// These values are used when `config.toml` does not
    /// contain any config for `upgrade`.
    impl Default for Config {
        fn default() -> Self {
            Config {
                kinds: default_kinds(),
                aliases: default_aliases(),
            }
        }
    }

    /// Get the default value for `kinds`
    pub fn default_kinds() -> Vec<Kind> {
        vec![Kind::Docker, Kind::Binary, Kind::Package]
    }

    /// Get the default value for `aliases`
    pub fn default_aliases() -> HashMap<String, String> {
        let mut aliases = HashMap::new();
        aliases.insert("javascript".to_string(), "jesta".to_string());
        aliases.insert("js".to_string(), "jesta".to_string());
        aliases
    }
}

#[cfg(feature = "cli")]
pub mod cli {
    use super::*;
    use structopt::StructOpt;

    #[derive(Debug, StructOpt)]
    #[structopt(
        about = "Manage plugins",
        setting = structopt::clap::AppSettings::DeriveDisplayOrder
    )]
    pub struct Args {
        #[structopt(subcommand)]
        pub action: Action,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(
        setting = structopt::clap::AppSettings::DeriveDisplayOrder
    )]
    pub enum Action {
        List,
        Show(Show),
        Add(Add),
        Remove(Remove),
        Upgrade(Upgrade),
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Show the manifest of a plugin")]
    pub struct Show {
        /// The name of the plugin to show
        #[structopt()]
        pub plugin: String,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Add one or more plugins")]
    pub struct Add {
        /// Attempt to add plugins as Docker image
        #[structopt(short, long)]
        pub docker: bool,

        /// Attempt to add plugins as binary
        #[structopt(short, long)]
        pub binary: bool,

        /// Attempt to add plugins as language package
        #[structopt(short, long)]
        pub package: bool,

        /// The names or aliases of plugins to add
        #[structopt(required = true, multiple = true)]
        pub plugins: Vec<String>,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Remove one or more plugins")]
    pub struct Remove {
        /// The names or aliases of plugins to remove
        #[structopt(required = true, multiple = true)]
        pub plugins: Vec<String>,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Upgrade plugins")]
    pub struct Upgrade {
        /// The names or aliases of plugins to upgrade
        /// (omit to upgrade all plugins)
        #[structopt(multiple = true)]
        pub plugins: Vec<String>,
    }

    pub async fn plugins(args: Args) -> Result<()> {
        let Args { action } = args;
        let config::Config { kinds, aliases } = crate::config::get()?.plugins;

        match action {
            Action::List => {
                let plugins = list()?;
                println!("{:?}", plugins);
                Ok(())
            }
            Action::Show(action) => {
                let Show { plugin } = action;

                let manifest = manifest(&plugin)?;
                println!("{}", manifest);
                Ok(())
            }
            Action::Add(action) => {
                let Add {
                    docker,
                    binary,
                    package,
                    plugins,
                } = action;

                let mut kinds_local = vec![];
                if docker {
                    kinds_local.push(Kind::Docker)
                }
                if binary {
                    kinds_local.push(Kind::Binary)
                }
                if package {
                    kinds_local.push(Kind::Package)
                }
                if kinds_local.len() == 0 {
                    kinds_local = kinds
                }

                add_list(plugins, kinds_local, aliases).await
            }
            Action::Remove(action) => {
                let Remove { plugins } = action;

                remove_list(plugins, aliases)
            }
            Action::Upgrade(action) => {
                let Upgrade { plugins } = action;

                let plugins = if plugins.len() == 0 { list()? } else { plugins };

                // Note: Currently, `upgrade` is just an alias for `add`
                // and does not warn user if plugin is not yet installed.
                add_list(plugins, kinds, aliases).await
            }
        }
    }
}
<commit_msg>feat(Plugins): Initial implementation for adding binary plugins<commit_after>use crate::util::dirs;
use anyhow::{anyhow, bail, Result};
use futures::StreamExt;
use rand::Rng;
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;
use std::thread;
use std::{collections::HashMap, process::Command};
use strum::{Display, EnumString, EnumVariantNames};

#[derive(Debug, Display, EnumString, EnumVariantNames, PartialEq, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum Kind {
    #[cfg(any(feature = "plugins-docker"))]
    Docker,
    #[cfg(any(feature = "plugins-binary"))]
    Binary,
    #[cfg(any(feature = "plugins-package"))]
    Package,
}

/// Split a plugin spec (ie. alias/name and, optionally, version) into alias and version
pub fn spec_to_alias_version(spec: &str) -> (String, String) {
    if spec.contains('@') {
        let parts: Vec<&str> = spec.split('@').collect();
        (parts[0].into(), parts[1].into())
    } else {
        (spec.into(), "latest".into())
    }
}

/// Resolve a plugin alias to a plugin name
///
/// If the provided string is a registered alias then returns the corresponding
/// plugin name, otherwise assumes the string is a plugin name
pub fn alias_to_name(alias: &str, aliases: &HashMap<String, String>) -> String {
    match aliases.get(alias) {
        Some(name) => name,
        None => alias,
    }
    .into()
}

/// Get the path to the manifest file of a plugin
pub fn manifest_path(name: &str) -> Result<PathBuf> {
    Ok(dirs::plugins(false)?.join(format!("{}.json", name)))
}

/// Get the manifest for an installed plugin
pub fn manifest(name: &str) -> Result<String> {
    // TODO: deserialize JSON into a manifest object
    match fs::read_to_string(manifest_path(name)?) {
        Ok(json) => Ok(json),
        Err(_) => bail!("Plugin '{}' is not installed", name),
    }
}

/// Test whether a plugin is installed
pub fn is_installed(name: &str) -> Result<bool> {
    Ok(manifest_path(name)?.exists())
}

/// List installed manifests
pub fn list() -> Result<Vec<String>> {
    // TODO: Return a Vec<Manifest>
    let dir = dirs::plugins(true)?;
    let mut plugins: Vec<String> = vec![];
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if let Some(ext) = path.extension() {
            if ext == "json" {
                let name = path.file_stem().unwrap().to_string_lossy().to_string();
                plugins.push(name)
            }
        }
    }
    Ok(plugins)
}

/// Add a plugin as a pulled Docker image
///
/// For this to succeed must be able to connect to the local
/// Docker server and be able to pull an image with corresponding
/// name.
#[cfg(any(feature = "plugins-docker"))]
pub async fn install_docker(name: &str, version: &str) -> Result<()> {
    let docker = bollard::Docker::connect_with_local_defaults()?;

    let image = if name.contains('/') {
        name.to_string()
    } else {
        format!("stencila/{}", name)
    };
    let image = format!("{}:{}", image, version);

    // Pull the image (by creating an image from it)
    let mut stream = docker.create_image(
        Some(bollard::image::CreateImageOptions {
            from_image: image.clone(),
            ..Default::default()
        }),
        None,
        None,
    );
    while let Some(item) = stream.next().await {
        match item {
            Ok(info) => {
                if let Some(error) = info.error {
                    bail!("{}", error)
                } else if let Some(status) = info.status {
                    // TODO display of status and progress could be improved and displayed
                    // per id (layer) as in docker CLI
                    // println!("{:?} {:?} {:?}", info.id, info.progress, info.progress_detail);
                    if let Some(progress) = info.progress {
                        tracing::info!("{} {}", status, progress)
                    } else {
                        tracing::info!("{}", status)
                    }
                }
            }
            Err(error) => match error {
                bollard::errors::Error::DockerResponseNotFoundError { .. } => {
                    bail!("Unable to find Docker image '{}'", image)
                }
                _ => bail!("{}", error),
            },
        }
    }

    // Create a container to obtain manifest from
    let container_name: String = rand::thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .take(24)
        .map(char::from)
        .collect();
    let response = docker
        .create_container(
            Some(bollard::container::CreateContainerOptions {
                name: &container_name,
            }),
            bollard::container::Config {
                image: Some(image.clone()),
                cmd: Some(vec!["manifest".to_string()]),
                ..Default::default()
            },
        )
        .await?;
    if response.warnings.len() > 0 {
        for warning in response.warnings {
            tracing::warn!("When creating container: {}", warning);
        }
    }

    // Start the container
    docker
        .start_container(
            &container_name,
            None::<bollard::container::StartContainerOptions<String>>,
        )
        .await?;

    // Capture the container output into `manifest`
    let mut stream = docker.logs::<String>(
        &container_name,
        Some(bollard::container::LogsOptions {
            follow: true,
            stderr: true,
            stdout: true,
            ..Default::default()
        }),
    );
    let mut stdout = vec![];
    let mut stderr = vec![];
    while let Some(output) = stream.next().await {
        if let Ok(output) = output {
            match output {
                bollard::container::LogOutput::StdOut { message } => {
                    stdout.extend_from_slice(&message)
                }
                bollard::container::LogOutput::StdErr { message } => {
                    stderr.extend_from_slice(&message)
                }
                _ => {}
            }
        }
    }

    if stderr.len() > 0 {
        tracing::warn!("{}", std::str::from_utf8(&stderr)?);
    }

    let manifest = if stdout.len() > 0 {
        match std::str::from_utf8(&stdout) {
            Ok(stdout) => stdout,
            Err(error) => bail!("Error converting stream to UTF8: {}", error),
        }
    } else {
        bail!(
            "No output from Docker container manifest command; is {} a Stencila plugin?",
            image
        )
    };

    // TODO: deserialize manifest and then serialize to file
    fs::write(manifest_path(name)?, manifest)?;

    Ok(())
}

/// Add a plugin as a downloaded binary
#[cfg(any(feature = "plugins-binary"))]
pub fn install_binary(name: &str, version: &str) -> Result<()> {
    let (owner, name) = if name.contains('/') {
        let parts: Vec<&str> = name.split('/').collect();
        (parts[0], parts[1])
    } else {
        ("stencila", name)
    };

    // Get the current version from the manifest (if any)
    let current_version = if let Ok(_manifest) = manifest(name) {
        // TODO extract version from manifest
        "0.1.0"
    } else {
        // Not yet installed, so artificially low version
        "0.0.0"
    };

    // Create the directory where the binary will be downloaded to
    let install_dir = dirs::plugins(false)?.join(name);
    fs::create_dir_all(&install_dir)?;
    let install_path = install_dir.join(name);

    let mut builder = self_update::backends::github::Update::configure();
    builder
        .repo_owner(owner)
        .repo_name(name)
        .bin_name(name)
        .current_version(current_version)
        .bin_install_path(&install_path)
        .show_output(true)
        .show_download_progress(true);
    if version != "latest" {
        builder.target_version_tag(format!("v{}", version).as_str());
    }

    // The download has to be done in another thread because it spawns
    // a new tokio runtime
    thread::spawn(move || -> Result<()> {
        if let Err(error) = builder.build()?.update() {
            match error {
                self_update::errors::Error::Network(message) => {
                    if message.contains("404") {
                        bail!("Could not find repository or corresponding release in repository")
                    } else {
                        bail!(message)
                    }
                }
                _ => bail!(error.to_string()),
            }
        } else {
            Ok(())
        }
    })
    .join()
    .map_err(|_| anyhow!("Error joining thread"))??;

    // Get and store manifest
    let manifest = Command::new(&install_path).arg("manifest").output()?.stdout;
    let manifest = std::str::from_utf8(&manifest)?;
    // TODO: deserialize manifest and then serialize to file
    fs::write(manifest_path(name)?, manifest)?;

    Ok(())
}

/// Add a plugin as a programing language package
#[cfg(any(feature = "plugins-package"))]
pub fn install_package(name: &str, version: &str) -> Result<()> {
    // TODO
    bail!(
        "Unable to add plugin '{}@{}' as programming language package",
        name,
        version
    )
}

/// Add a plugin
pub async fn add(plugin: &str, kinds: &Vec<Kind>, aliases: &HashMap<String, String>) -> Result<()> {
    let (alias, version) = spec_to_alias_version(plugin);
    let name = alias_to_name(&alias, aliases);

    if is_installed(&name)? {
        bail!(
            "Plugin '{}' is already installed, consider using `stencila plugins upgrade {}`",
            name,
            name
        );
    }

    for kind in kinds {
        let result = match kind {
            Kind::Docker => install_docker(&name, &version).await,
            Kind::Binary => install_binary(&name, &version),
            Kind::Package => install_package(&name, &version),
        };
        match result {
            // Success, so just return now
            Ok(_) => return Ok(()),
            // Error, keep trying other kinds, or if there is
            // only one kind, return that error.
            Err(error) => {
                if kinds.len() == 1 {
                    return Err(error);
                }
            }
        }
    }

    bail!(
        "Unable to install plugin '{}', tried kinds {:?}",
        plugin,
        kinds
    )
}

/// Add a list of plugins
pub async fn add_list(
    plugins: Vec<String>,
    kinds: Vec<Kind>,
    aliases: HashMap<String, String>,
) -> Result<()> {
    for plugin in plugins {
        match add(&plugin, &kinds, &aliases).await {
            Ok(_) => tracing::info!("Added plugin {}", plugin),
            Err(error) => tracing::error!("{}", error),
        }
    }
    Ok(())
}

/// Remove a plugin
pub fn remove(plugin: &str, aliases: &HashMap<String, String>) -> Result<()> {
    let (alias, ..) = spec_to_alias_version(plugin);
    let name = alias_to_name(&alias, aliases);

    let file = manifest_path(&name)?;
    if file.exists() {
        fs::remove_file(file)?;
    }

    Ok(())
}

/// Remove a list of plugins
pub fn remove_list(plugins: Vec<String>, aliases: HashMap<String, String>) -> Result<()> {
    for plugin in plugins {
        match remove(&plugin, &aliases) {
            Ok(_) => tracing::info!("Removed plugin {}", plugin),
            Err(error) => tracing::error!("{}", error),
        }
    }
    Ok(())
}

#[cfg(feature = "config")]
pub mod config {
    use super::*;
    use validator::Validate;

    #[derive(Debug, PartialEq, Deserialize, Serialize, Validate)]
    pub struct Config {
        #[serde(default = "default_kinds")]
        pub kinds: Vec<Kind>,

        #[serde(default = "default_aliases")]
        pub aliases: HashMap<String, String>,
    }

    /// Default configuration
    ///
    /// These values are used when `config.toml` does not
    /// contain any config for `upgrade`.
    impl Default for Config {
        fn default() -> Self {
            Config {
                kinds: default_kinds(),
                aliases: default_aliases(),
            }
        }
    }

    /// Get the default value for `kinds`
    pub fn default_kinds() -> Vec<Kind> {
        vec![Kind::Docker, Kind::Binary, Kind::Package]
    }

    /// Get the default value for `aliases`
    pub fn default_aliases() -> HashMap<String, String> {
        let mut aliases = HashMap::new();
        aliases.insert("javascript".to_string(), "jesta".to_string());
        aliases.insert("js".to_string(), "jesta".to_string());
        aliases
    }
}

#[cfg(feature = "cli")]
pub mod cli {
    use super::*;
    use structopt::StructOpt;

    #[derive(Debug, StructOpt)]
    #[structopt(
        about = "Manage plugins",
        setting = structopt::clap::AppSettings::DeriveDisplayOrder
    )]
    pub struct Args {
        #[structopt(subcommand)]
        pub action: Action,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(
        setting = structopt::clap::AppSettings::DeriveDisplayOrder
    )]
    pub enum Action {
        List,
        Show(Show),
        Add(Add),
        Remove(Remove),
        Upgrade(Upgrade),
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Show the manifest of a plugin")]
    pub struct Show {
        /// The name of the plugin to show
        #[structopt()]
        pub plugin: String,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Add one or more plugins")]
    pub struct Add {
        /// Attempt to add plugins as Docker image
        #[structopt(short, long)]
        pub docker: bool,

        /// Attempt to add plugins as binary
        #[structopt(short, long)]
        pub binary: bool,

        /// Attempt to add plugins as language package
        #[structopt(short, long)]
        pub package: bool,

        /// The names or aliases of plugins to add
        #[structopt(required = true, multiple = true)]
        pub plugins: Vec<String>,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Remove one or more plugins")]
    pub struct Remove {
        /// The names or aliases of plugins to remove
        #[structopt(required = true, multiple = true)]
        pub plugins: Vec<String>,
    }

    #[derive(Debug, StructOpt)]
    #[structopt(about = "Upgrade plugins")]
    pub struct Upgrade {
        /// The names or aliases of plugins to upgrade
        /// (omit to upgrade all plugins)
        #[structopt(multiple = true)]
        pub plugins: Vec<String>,
    }

    pub async fn plugins(args: Args) -> Result<()> {
        let Args { action } = args;
        let config::Config { kinds, aliases } = crate::config::get()?.plugins;

        match action {
            Action::List => {
                let plugins = list()?;
                println!("{:?}", plugins);
                Ok(())
            }
            Action::Show(action) => {
                let Show { plugin } = action;

                let manifest = manifest(&plugin)?;
                println!("{}", manifest);
                Ok(())
            }
            Action::Add(action) => {
                let Add {
                    docker,
                    binary,
                    package,
                    plugins,
                } = action;

                let mut kinds_local = vec![];
                if docker {
                    kinds_local.push(Kind::Docker)
                }
                if binary {
                    kinds_local.push(Kind::Binary)
                }
                if package {
                    kinds_local.push(Kind::Package)
                }
                if kinds_local.len() == 0 {
                    kinds_local = kinds
                }

                add_list(plugins, kinds_local, aliases).await
            }
            Action::Remove(action) => {
                let Remove { plugins } = action;

                remove_list(plugins, aliases)
            }
            Action::Upgrade(action) => {
                let Upgrade { plugins } = action;

                let plugins = if plugins.len() == 0 { list()? } else { plugins };

                // Note: Currently, `upgrade` is just an alias for `add`
                // and does not warn user if plugin is not yet installed.
                add_list(plugins, kinds, aliases).await
            }
        }
    }
}
<|endoftext|>


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

use std::{io, os};
use std::default::Default;
use std::vec::MoveItems;

use test::{black_box, Bencher, TestDesc, TestDescAndFn};
use test::{DynTestName, DynBenchFn, TDynBenchFn};

use html5ever::tokenizer::{TokenSink, Token, Tokenizer, TokenizerOpts};

struct Sink;

impl TokenSink for Sink {
    fn process_token(&mut self, token: Token) {
        // Don't use the token, but make sure we don't get
        // optimized out entirely.
        black_box(token);
    }
}

// This could almost be the TokenSink too, but it's not
// mut within run().
struct Bench {
    input: String,
    clone_only: bool,
    opts: TokenizerOpts,
}

impl Bench {
    fn new(name: &str, size: Option<uint>, clone_only: bool,
           opts: TokenizerOpts) -> Bench {
        let mut path = os::self_exe_path().expect("can't get exe path");
        path.push("../data/bench/");
        path.push(name);
        let mut file = io::File::open(&path).ok().expect("can't open file");
        let file_input = file.read_to_string().ok().expect("can't read file");

        let input = match size {
            None => file_input,
            Some(size) => {
                // Replicate the input in memory up to the desired size.
                let mut input = String::with_capacity(size);
                while input.len() < size {
                    input.push_str(file_input.as_slice());
                }
                input
            }
        };

        Bench {
            input: input,
            clone_only: clone_only,
            opts: opts,
        }
    }
}

impl TDynBenchFn for Bench {
    fn run(&self, bh: &mut Bencher) {
        bh.iter(|| {
            let input = self.input.clone();
            if self.clone_only {
                // Because the tokenizer consumes its buffers, we need
                // to clone inside iter().  We can benchmark this
                // separately and subtract it out.
                black_box(input);
            } else {
                let mut tok = Tokenizer::new(Sink, self.opts.clone());
                tok.feed(input);
                tok.end();
            }
        });
    }
}

fn make_bench(name: &str, size: Option<uint>, clone_only: bool,
              opts: TokenizerOpts) -> TestDescAndFn {
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName([
                "tokenize ".to_string(),
                name.to_string(),
                size.map_or("".to_string(), |s| format!(" size {:7u}", s)),
                (if clone_only { " (clone only)" } else { "" }).to_string(),
                (if opts.exact_errors { " (exact errors)" } else { "" }).to_string(),
            ].concat().to_string()),
            ignore: false,
            should_fail: false,
        },
        testfn: DynBenchFn(box Bench::new(name, size, clone_only, opts)),
    }
}

pub fn tests() -> MoveItems<TestDescAndFn> {
    let mut tests = vec!(make_bench("lipsum.html", Some(1024*1024), true, Default::default()));

    let mut opts_vec = vec!(Default::default());
    if os::getenv("BENCH_EXACT_ERRORS").is_some() {
        opts_vec.push(TokenizerOpts {
            exact_errors: true,
            .. Default::default()
        });
    }

    for opts in opts_vec.iter() {
        for &file in ["lipsum.html", "lipsum-zh.html", "strong.html"].iter() {
            for &sz in [1024, 1024*1024].iter() {
                tests.push(make_bench(file, Some(sz), false, opts.clone()));
            }
        }

        for &file in ["tiny-fragment.html", "small-fragment.html", "medium-fragment.html"].iter() {
            tests.push(make_bench(file, None, false, opts.clone()));
        }

        if os::getenv("BENCH_UNCOMMITTED").is_some() {
            // Not checked into the repo, so don't include by default.
            for &file in ["sina.com.cn.html", "wikipedia.html"].iter() {
                let name = format!("uncommitted/{:s}", file);
                tests.push(make_bench(name.as_slice(), None, false, opts.clone()));
            }
        }
    }

    tests.into_iter()
}
<commit_msg>Benchmark on small input chunks<commit_after>// Copyright 2014 The html5ever Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::{io, os, cmp};
use std::default::Default;
use std::vec::MoveItems;

use test::{black_box, Bencher, TestDesc, TestDescAndFn};
use test::{DynTestName, DynBenchFn, TDynBenchFn};

use html5ever::tokenizer::{TokenSink, Token, Tokenizer, TokenizerOpts};

struct Sink;

impl TokenSink for Sink {
    fn process_token(&mut self, token: Token) {
        // Don't use the token, but make sure we don't get
        // optimized out entirely.
        black_box(token);
    }
}

// This could almost be the TokenSink too, but it's not
// mut within run().
struct Bench {
    input: Vec<String>,
    clone_only: bool,
    opts: TokenizerOpts,
}

impl Bench {
    fn new(name: &str, size: Option<uint>, clone_only: bool,
           opts: TokenizerOpts) -> Bench {
        let mut path = os::self_exe_path().expect("can't get exe path");
        path.push("../data/bench/");
        path.push(name);
        let mut file = io::File::open(&path).ok().expect("can't open file");

        // Read the file and treat it as an infinitely repeating sequence of characters.
        let file_input = file.read_to_string().ok().expect("can't read file");
        let size = size.unwrap_or(file_input.len());
        let mut stream = file_input.as_slice().chars().cycle();

        // Break the input into chunks of 1024 chars (= a few kB).
        // This simulates reading from the network.
        let mut input = vec![];
        let mut total = 0u;
        while total < size {
            // The by_ref() call is important, otherwise we get wrong results!
            // See rust-lang/rust#18045.
            let sz = cmp::min(1024, size - total);
            input.push(stream.by_ref().take(sz).collect());
            total += sz;
        }

        Bench {
            input: input,
            clone_only: clone_only,
            opts: opts,
        }
    }
}

impl TDynBenchFn for Bench {
    fn run(&self, bh: &mut Bencher) {
        bh.iter(|| {
            let input = self.input.clone();
            if self.clone_only {
                // Because the tokenizer consumes its buffers, we need
                // to clone inside iter().  We can benchmark this
                // separately and subtract it out.
                //
                // See rust-lang/rust#18043.
                black_box(input);
            } else {
                let mut tok = Tokenizer::new(Sink, self.opts.clone());
                for buf in input.into_iter() {
                    tok.feed(buf);
                }
                tok.end();
            }
        });
    }
}

fn make_bench(name: &str, size: Option<uint>, clone_only: bool,
              opts: TokenizerOpts) -> TestDescAndFn {
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName([
                "tokenize ".to_string(),
                name.to_string(),
                size.map_or("".to_string(), |s| format!(" size {:7u}", s)),
                (if clone_only { " (clone only)" } else { "" }).to_string(),
                (if opts.exact_errors { " (exact errors)" } else { "" }).to_string(),
            ].concat().to_string()),
            ignore: false,
            should_fail: false,
        },
        testfn: DynBenchFn(box Bench::new(name, size, clone_only, opts)),
    }
}

pub fn tests() -> MoveItems<TestDescAndFn> {
    let mut tests = vec!(make_bench("lipsum.html", Some(1024*1024), true, Default::default()));

    let mut opts_vec = vec!(Default::default());
    if os::getenv("BENCH_EXACT_ERRORS").is_some() {
        opts_vec.push(TokenizerOpts {
            exact_errors: true,
            .. Default::default()
        });
    }

    for opts in opts_vec.iter() {
        for &file in ["lipsum.html", "lipsum-zh.html", "strong.html"].iter() {
            for &sz in [1024, 1024*1024].iter() {
                tests.push(make_bench(file, Some(sz), false, opts.clone()));
            }
        }

        for &file in ["tiny-fragment.html", "small-fragment.html", "medium-fragment.html"].iter() {
            tests.push(make_bench(file, None, false, opts.clone()));
        }

        if os::getenv("BENCH_UNCOMMITTED").is_some() {
            // Not checked into the repo, so don't include by default.
            for &file in ["sina.com.cn.html", "wikipedia.html"].iter() {
                let name = format!("uncommitted/{:s}", file);
                tests.push(make_bench(name.as_slice(), None, false, opts.clone()));
            }
        }
    }

    tests.into_iter()
}
<|endoftext|>


<commit_before>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<commit_msg>impl `Step` for `Page`<commit_after>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
#[cfg(feature = "step_trait")]
use core::iter::Step;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

#[cfg(feature = "step_trait")]
impl<S: PageSize> Step for Page<S> {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        Step::steps_between(&start.start_address, &end.start_address)
            .map(|steps| steps / S::SIZE as usize)
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::forward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::backward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<|endoftext|>


<commit_before>//! Various extension methods to ast Expr Nodes, which are hard to code-generate.

use crate::{
    SyntaxToken, SyntaxElement, SmolStr,
    ast::{self, AstNode, AstChildren, children, child_opt},
    SyntaxKind::*,
    T
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ElseBranch<'a> {
    Block(&'a ast::Block),
    IfExpr(&'a ast::IfExpr),
}

impl ast::IfExpr {
    pub fn then_branch(&self) -> Option<&ast::Block> {
        self.blocks().nth(0)
    }
    pub fn else_branch(&self) -> Option<ElseBranch> {
        let res = match self.blocks().nth(1) {
            Some(block) => ElseBranch::Block(block),
            None => {
                let elif: &ast::IfExpr = child_opt(self)?;
                ElseBranch::IfExpr(elif)
            }
        };
        Some(res)
    }

    fn blocks(&self) -> AstChildren<ast::Block> {
        children(self)
    }
}

impl ast::RefExpr {
    pub fn is_mut(&self) -> bool {
        self.syntax().children_with_tokens().any(|n| n.kind() == T![mut])
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum PrefixOp {
    /// The `*` operator for dereferencing
    Deref,
    /// The `!` operator for logical inversion
    Not,
    /// The `-` operator for negation
    Neg,
}

impl ast::PrefixExpr {
    pub fn op_kind(&self) -> Option<PrefixOp> {
        match self.op_token()?.kind() {
            T![*] => Some(PrefixOp::Deref),
            T![!] => Some(PrefixOp::Not),
            T![-] => Some(PrefixOp::Neg),
            _ => None,
        }
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.syntax().first_child_or_token()?.as_token()
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinOp {
    /// The `||` operator for boolean OR
    BooleanOr,
    /// The `&&` operator for boolean AND
    BooleanAnd,
    /// The `==` operator for equality testing
    EqualityTest,
    /// The `!=` operator for equality testing
    NegatedEqualityTest,
    /// The `<=` operator for lesser-equal testing
    LesserEqualTest,
    /// The `>=` operator for greater-equal testing
    GreaterEqualTest,
    /// The `<` operator for comparison
    LesserTest,
    /// The `>` operator for comparison
    GreaterTest,
    /// The `+` operator for addition
    Addition,
    /// The `*` operator for multiplication
    Multiplication,
    /// The `-` operator for subtraction
    Subtraction,
    /// The `/` operator for division
    Division,
    /// The `%` operator for remainder after division
    Remainder,
    /// The `<<` operator for left shift
    LeftShift,
    /// The `>>` operator for right shift
    RightShift,
    /// The `^` operator for bitwise XOR
    BitwiseXor,
    /// The `|` operator for bitwise OR
    BitwiseOr,
    /// The `&` operator for bitwise AND
    BitwiseAnd,
    /// The `..` operator for right-open ranges
    RangeRightOpen,
    /// The `..=` operator for right-closed ranges
    RangeRightClosed,
    /// The `=` operator for assignment
    Assignment,
    /// The `+=` operator for assignment after addition
    AddAssign,
    /// The `/=` operator for assignment after division
    DivAssign,
    /// The `*=` operator for assignment after multiplication
    MulAssign,
    /// The `%=` operator for assignment after remainders
    RemAssign,
    /// The `>>=` operator for assignment after shifting right
    ShrAssign,
    /// The `<<=` operator for assignment after shifting left
    ShlAssign,
    /// The `-=` operator for assignment after subtraction
    SubAssign,
    /// The `|=` operator for assignment after bitwise OR
    BitOrAssign,
    /// The `&=` operator for assignment after bitwise AND
    BitAndAssign,
    /// The `^=` operator for assignment after bitwise XOR
    BitXorAssign,
}

impl ast::BinExpr {
    fn op_details(&self) -> Option<(SyntaxToken, BinOp)> {
        self.syntax().children_with_tokens().filter_map(|it| it.as_token()).find_map(|c| {
            match c.kind() {
                T![||] => Some((c, BinOp::BooleanOr)),
                T![&&] => Some((c, BinOp::BooleanAnd)),
                T![==] => Some((c, BinOp::EqualityTest)),
                T![!=] => Some((c, BinOp::NegatedEqualityTest)),
                T![<=] => Some((c, BinOp::LesserEqualTest)),
                T![>=] => Some((c, BinOp::GreaterEqualTest)),
                T![<] => Some((c, BinOp::LesserTest)),
                T![>] => Some((c, BinOp::GreaterTest)),
                T![+] => Some((c, BinOp::Addition)),
                T![*] => Some((c, BinOp::Multiplication)),
                T![-] => Some((c, BinOp::Subtraction)),
                T![/] => Some((c, BinOp::Division)),
                T![%] => Some((c, BinOp::Remainder)),
                T![<<] => Some((c, BinOp::LeftShift)),
                T![>>] => Some((c, BinOp::RightShift)),
                T![^] => Some((c, BinOp::BitwiseXor)),
                T![|] => Some((c, BinOp::BitwiseOr)),
                T![&] => Some((c, BinOp::BitwiseAnd)),
                T![..] => Some((c, BinOp::RangeRightOpen)),
                T![..=] => Some((c, BinOp::RangeRightClosed)),
                T![=] => Some((c, BinOp::Assignment)),
                T![+=] => Some((c, BinOp::AddAssign)),
                T![/=] => Some((c, BinOp::DivAssign)),
                T![*=] => Some((c, BinOp::MulAssign)),
                T![%=] => Some((c, BinOp::RemAssign)),
                T![>>=] => Some((c, BinOp::ShrAssign)),
                T![<<=] => Some((c, BinOp::ShlAssign)),
                T![-=] => Some((c, BinOp::SubAssign)),
                T![|=] => Some((c, BinOp::BitOrAssign)),
                T![&=] => Some((c, BinOp::BitAndAssign)),
                T![^=] => Some((c, BinOp::BitXorAssign)),
                _ => None,
            }
        })
    }

    pub fn op_kind(&self) -> Option<BinOp> {
        self.op_details().map(|t| t.1)
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.op_details().map(|t| t.0)
    }

    pub fn lhs(&self) -> Option<&ast::Expr> {
        children(self).nth(0)
    }

    pub fn rhs(&self) -> Option<&ast::Expr> {
        children(self).nth(1)
    }

    pub fn sub_exprs(&self) -> (Option<&ast::Expr>, Option<&ast::Expr>) {
        let mut children = children(self);
        let first = children.next();
        let second = children.next();
        (first, second)
    }
}

pub enum ArrayExprKind<'a> {
    Repeat { initializer: Option<&'a ast::Expr>, repeat: Option<&'a ast::Expr> },
    ElementList(AstChildren<'a, ast::Expr>),
}

impl ast::ArrayExpr {
    pub fn kind(&self) -> ArrayExprKind {
        if self.is_repeat() {
            ArrayExprKind::Repeat {
                initializer: children(self).nth(0),
                repeat: children(self).nth(1),
            }
        } else {
            ArrayExprKind::ElementList(children(self))
        }
    }

    fn is_repeat(&self) -> bool {
        self.syntax().children_with_tokens().any(|it| it.kind() == T![;])
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LiteralKind {
    String,
    ByteString,
    Char,
    Byte,
    IntNumber { suffix: Option<SmolStr> },
    FloatNumber { suffix: Option<SmolStr> },
    Bool,
}

impl ast::Literal {
    pub fn token(&self) -> SyntaxToken {
        match self.syntax().first_child_or_token().unwrap() {
            SyntaxElement::Token(token) => token,
            _ => unreachable!(),
        }
    }

    pub fn kind(&self) -> LiteralKind {
        match self.token().kind() {
            INT_NUMBER => {
                let allowed_suffix_list = [
                    "isize", "i128", "i64", "i32", "i16", "i8", "usize", "u128", "u64", "u32",
                    "u16", "u8",
                ];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::IntNumber { suffix }
            }
            FLOAT_NUMBER => {
                let allowed_suffix_list = ["f64", "f32"];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::FloatNumber { suffix: suffix }
            }
            STRING | RAW_STRING => LiteralKind::String,
            T![true] | T![false] => LiteralKind::Bool,
            BYTE_STRING | RAW_BYTE_STRING => LiteralKind::ByteString,
            CHAR => LiteralKind::Char,
            BYTE => LiteralKind::Byte,
            _ => unreachable!(),
        }
    }
}

impl ast::NamedField {
    pub fn parent_struct_lit(&self) -> &ast::StructLit {
        self.syntax().ancestors().find_map(ast::StructLit::cast).unwrap()
    }
}
<commit_msg>Skip attrs in `Literal::token`<commit_after>//! Various extension methods to ast Expr Nodes, which are hard to code-generate.

use crate::{
    SyntaxToken, SyntaxElement, SmolStr,
    ast::{self, AstNode, AstChildren, children, child_opt},
    SyntaxKind::*,
    T
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ElseBranch<'a> {
    Block(&'a ast::Block),
    IfExpr(&'a ast::IfExpr),
}

impl ast::IfExpr {
    pub fn then_branch(&self) -> Option<&ast::Block> {
        self.blocks().nth(0)
    }
    pub fn else_branch(&self) -> Option<ElseBranch> {
        let res = match self.blocks().nth(1) {
            Some(block) => ElseBranch::Block(block),
            None => {
                let elif: &ast::IfExpr = child_opt(self)?;
                ElseBranch::IfExpr(elif)
            }
        };
        Some(res)
    }

    fn blocks(&self) -> AstChildren<ast::Block> {
        children(self)
    }
}

impl ast::RefExpr {
    pub fn is_mut(&self) -> bool {
        self.syntax().children_with_tokens().any(|n| n.kind() == T![mut])
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum PrefixOp {
    /// The `*` operator for dereferencing
    Deref,
    /// The `!` operator for logical inversion
    Not,
    /// The `-` operator for negation
    Neg,
}

impl ast::PrefixExpr {
    pub fn op_kind(&self) -> Option<PrefixOp> {
        match self.op_token()?.kind() {
            T![*] => Some(PrefixOp::Deref),
            T![!] => Some(PrefixOp::Not),
            T![-] => Some(PrefixOp::Neg),
            _ => None,
        }
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.syntax().first_child_or_token()?.as_token()
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum BinOp {
    /// The `||` operator for boolean OR
    BooleanOr,
    /// The `&&` operator for boolean AND
    BooleanAnd,
    /// The `==` operator for equality testing
    EqualityTest,
    /// The `!=` operator for equality testing
    NegatedEqualityTest,
    /// The `<=` operator for lesser-equal testing
    LesserEqualTest,
    /// The `>=` operator for greater-equal testing
    GreaterEqualTest,
    /// The `<` operator for comparison
    LesserTest,
    /// The `>` operator for comparison
    GreaterTest,
    /// The `+` operator for addition
    Addition,
    /// The `*` operator for multiplication
    Multiplication,
    /// The `-` operator for subtraction
    Subtraction,
    /// The `/` operator for division
    Division,
    /// The `%` operator for remainder after division
    Remainder,
    /// The `<<` operator for left shift
    LeftShift,
    /// The `>>` operator for right shift
    RightShift,
    /// The `^` operator for bitwise XOR
    BitwiseXor,
    /// The `|` operator for bitwise OR
    BitwiseOr,
    /// The `&` operator for bitwise AND
    BitwiseAnd,
    /// The `..` operator for right-open ranges
    RangeRightOpen,
    /// The `..=` operator for right-closed ranges
    RangeRightClosed,
    /// The `=` operator for assignment
    Assignment,
    /// The `+=` operator for assignment after addition
    AddAssign,
    /// The `/=` operator for assignment after division
    DivAssign,
    /// The `*=` operator for assignment after multiplication
    MulAssign,
    /// The `%=` operator for assignment after remainders
    RemAssign,
    /// The `>>=` operator for assignment after shifting right
    ShrAssign,
    /// The `<<=` operator for assignment after shifting left
    ShlAssign,
    /// The `-=` operator for assignment after subtraction
    SubAssign,
    /// The `|=` operator for assignment after bitwise OR
    BitOrAssign,
    /// The `&=` operator for assignment after bitwise AND
    BitAndAssign,
    /// The `^=` operator for assignment after bitwise XOR
    BitXorAssign,
}

impl ast::BinExpr {
    fn op_details(&self) -> Option<(SyntaxToken, BinOp)> {
        self.syntax().children_with_tokens().filter_map(|it| it.as_token()).find_map(|c| {
            match c.kind() {
                T![||] => Some((c, BinOp::BooleanOr)),
                T![&&] => Some((c, BinOp::BooleanAnd)),
                T![==] => Some((c, BinOp::EqualityTest)),
                T![!=] => Some((c, BinOp::NegatedEqualityTest)),
                T![<=] => Some((c, BinOp::LesserEqualTest)),
                T![>=] => Some((c, BinOp::GreaterEqualTest)),
                T![<] => Some((c, BinOp::LesserTest)),
                T![>] => Some((c, BinOp::GreaterTest)),
                T![+] => Some((c, BinOp::Addition)),
                T![*] => Some((c, BinOp::Multiplication)),
                T![-] => Some((c, BinOp::Subtraction)),
                T![/] => Some((c, BinOp::Division)),
                T![%] => Some((c, BinOp::Remainder)),
                T![<<] => Some((c, BinOp::LeftShift)),
                T![>>] => Some((c, BinOp::RightShift)),
                T![^] => Some((c, BinOp::BitwiseXor)),
                T![|] => Some((c, BinOp::BitwiseOr)),
                T![&] => Some((c, BinOp::BitwiseAnd)),
                T![..] => Some((c, BinOp::RangeRightOpen)),
                T![..=] => Some((c, BinOp::RangeRightClosed)),
                T![=] => Some((c, BinOp::Assignment)),
                T![+=] => Some((c, BinOp::AddAssign)),
                T![/=] => Some((c, BinOp::DivAssign)),
                T![*=] => Some((c, BinOp::MulAssign)),
                T![%=] => Some((c, BinOp::RemAssign)),
                T![>>=] => Some((c, BinOp::ShrAssign)),
                T![<<=] => Some((c, BinOp::ShlAssign)),
                T![-=] => Some((c, BinOp::SubAssign)),
                T![|=] => Some((c, BinOp::BitOrAssign)),
                T![&=] => Some((c, BinOp::BitAndAssign)),
                T![^=] => Some((c, BinOp::BitXorAssign)),
                _ => None,
            }
        })
    }

    pub fn op_kind(&self) -> Option<BinOp> {
        self.op_details().map(|t| t.1)
    }

    pub fn op_token(&self) -> Option<SyntaxToken> {
        self.op_details().map(|t| t.0)
    }

    pub fn lhs(&self) -> Option<&ast::Expr> {
        children(self).nth(0)
    }

    pub fn rhs(&self) -> Option<&ast::Expr> {
        children(self).nth(1)
    }

    pub fn sub_exprs(&self) -> (Option<&ast::Expr>, Option<&ast::Expr>) {
        let mut children = children(self);
        let first = children.next();
        let second = children.next();
        (first, second)
    }
}

pub enum ArrayExprKind<'a> {
    Repeat { initializer: Option<&'a ast::Expr>, repeat: Option<&'a ast::Expr> },
    ElementList(AstChildren<'a, ast::Expr>),
}

impl ast::ArrayExpr {
    pub fn kind(&self) -> ArrayExprKind {
        if self.is_repeat() {
            ArrayExprKind::Repeat {
                initializer: children(self).nth(0),
                repeat: children(self).nth(1),
            }
        } else {
            ArrayExprKind::ElementList(children(self))
        }
    }

    fn is_repeat(&self) -> bool {
        self.syntax().children_with_tokens().any(|it| it.kind() == T![;])
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum LiteralKind {
    String,
    ByteString,
    Char,
    Byte,
    IntNumber { suffix: Option<SmolStr> },
    FloatNumber { suffix: Option<SmolStr> },
    Bool,
}

impl ast::Literal {
    pub fn token(&self) -> SyntaxToken {
        let elem = self.syntax()
            .children_with_tokens()
            .find(|e| e.kind() != ATTR && !e.kind().is_trivia());
        match elem {
            Some(SyntaxElement::Token(token)) => token,
            _ => unreachable!(),
        }
    }

    pub fn kind(&self) -> LiteralKind {
        match self.token().kind() {
            INT_NUMBER => {
                let allowed_suffix_list = [
                    "isize", "i128", "i64", "i32", "i16", "i8", "usize", "u128", "u64", "u32",
                    "u16", "u8",
                ];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::IntNumber { suffix }
            }
            FLOAT_NUMBER => {
                let allowed_suffix_list = ["f64", "f32"];
                let text = self.token().text().to_string();
                let suffix = allowed_suffix_list
                    .iter()
                    .find(|&s| text.ends_with(s))
                    .map(|&suf| SmolStr::new(suf));
                LiteralKind::FloatNumber { suffix: suffix }
            }
            STRING | RAW_STRING => LiteralKind::String,
            T![true] | T![false] => LiteralKind::Bool,
            BYTE_STRING | RAW_BYTE_STRING => LiteralKind::ByteString,
            CHAR => LiteralKind::Char,
            BYTE => LiteralKind::Byte,
            _ => unreachable!(),
        }
    }
}

#[test]
fn test_literal_with_attr() {
    let parse = ast::SourceFile::parse(r#"const _: &str = { #[attr] "Hello" };"#);
    let lit = parse
        .tree
        .syntax()
        .descendants()
        .find_map(ast::Literal::cast)
        .unwrap();
    assert_eq!(lit.token().text(), r#""Hello""#);
}

impl ast::NamedField {
    pub fn parent_struct_lit(&self) -> &ast::StructLit {
        self.syntax().ancestors().find_map(ast::StructLit::cast).unwrap()
    }
}
<|endoftext|>