<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use command::studio::enter::ARTIFACT_PATH_ENVVAR;
use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use hcore::os::process;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;
    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }
    if args[0] == OsString::from("rm") {
        return rm_container(&docker_cmd);
    }
    return enter(&docker_cmd, &args);
}

fn enter(docker_cmd: &Path, args: &Vec<OsString>) -> Result<()> {
    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd)
        && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![
        String::from("DEBUG"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_ORIGIN"),
        String::from("HAB_ORIGIN_KEYS"),
        String::from("HAB_STUDIO_BACKLINE_PKG"),
        String::from("HAB_STUDIO_NOSTUDIORC"),
        String::from("HAB_STUDIO_SUP"),
        String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
        String::from("http_proxy"),
        String::from("https_proxy"),
        String::from("RUST_LOG"),
    ];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    if container_exists(&docker_cmd) {
        return start_container(&docker_cmd);
    }
    run_container(&docker_cmd, args.to_vec(), volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images")
        .arg(&image_identifier(docker_cmd))
        .arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd
        .spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image, result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
        && (stderr.as_str().contains("Mounts denied")
            || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn container_exists(docker_cmd: &Path) -> bool {
    let name = path_to_name().expect("Unable to load name from path");
    let cmd_args: Vec<OsString> = vec!["container".into(), "inspect".into(), name.into()];
    let output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");
    if output.status.success() {
        return true;
    }
    return false;
}

#[cfg(not(target_os = "windows"))]
fn path_to_name() -> Result<String> {
    let mut cwd = env::current_dir()?;
    cwd = cwd.strip_prefix("/")?.to_path_buf();
    let pathstr = cwd.to_str().expect("Path to be parseable");
    return Ok(pathstr.replace("/", "--"));
}

#[cfg(target_os = "windows")]
fn path_to_name() -> Result<String> {
    let cwd = env::current_dir()?;
    let pathstr = cwd.to_str().expect("Path to be parseable");
    let stripped_path: String = pathstr.chars().skip(3).take(pathstr.len() - 3).collect();
    return Ok(stripped_path.replace("\\", "--"));
}

fn rm_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["rm".into()];
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn run_container<I, J, S, T>(
    docker_cmd: &Path,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    let name = path_to_name()?;
    cmd_args.push("--name".into());
    cmd_args.push(name.into());
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR, opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn start_container(docker_cmd: &Path) -> Result<()> {
    let mut cmd_args: Vec<OsString> = vec!["start".into()];
    cmd_args.push("--interactive".into());
    let name = path_to_name()?;
    cmd_args.push(name.into());
    Ok(process::become_command(docker_cmd.to_path_buf(), cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<commit_msg>Revert "Make docker studio behave like chroot studio"<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use command::studio::enter::ARTIFACT_PATH_ENVVAR;
use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use hcore::os::process;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;

    if is_image_present(&docker_cmd) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd)?;
    }

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            mnt_prefix,
            CACHE_KEY_PATH
        ),
    ];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!(
            "{}:{}/{}",
            cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH
        ));
    }
    if !is_serving_windows_containers(&docker_cmd)
        && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![
        String::from("DEBUG"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_ORIGIN"),
        String::from("HAB_ORIGIN_KEYS"),
        String::from("HAB_STUDIO_BACKLINE_PKG"),
        String::from("HAB_STUDIO_NOSTUDIORC"),
        String::from("HAB_STUDIO_SUP"),
        String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
        String::from("http_proxy"),
        String::from("https_proxy"),
        String::from("RUST_LOG"),
    ];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images")
        .arg(&image_identifier(docker_cmd))
        .arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path) -> Result<()> {
    let image = image_identifier(docker_cmd);
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
        .arg(&image)
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd
        .spawn()
        .expect("Docker command failed to spawn")
        .wait_with_output()
        .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", &image);
    } else {
        debug!(
            "Pulling Docker image '{}' failed with exit code: {:?}",
            &image, result.status
        );

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(
                image_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).to_string(),
            ));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(docker_cmd).into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd)
        .args(&cmd_args)
        .output()
        .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
        && (stderr.as_str().contains("Mounts denied")
            || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(
    docker_cmd: PathBuf,
    args: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    J: IntoIterator<Item = T>,
    S: AsRef<OsStr>,
    T: AsRef<str>,
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    if !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            cmd_args.push("--tty".into());
            cmd_args.push("--interactive".into());
        }
    }
    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts.split(" ")
                .map(|v| v.into())
                // Ensure we're not passing something like `--tty` again here.
                .filter(|v| !cmd_args.contains(v))
                .collect::<Vec<_>>();
        if !opts.is_empty() {
            debug!(
                "Adding extra Docker options from {} = {:?}",
                DOCKER_OPTS_ENVVAR, opts
            );
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }
    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd, cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<|endoftext|>


<commit_before>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![no_std]
#![no_main]

use core::panic::PanicInfo;
use oak_linux_boot_params::{BootParams, CCBlobSevInfo, CCSetupData, SetupDataType};
use sev_guest::{
    cpuid::CpuidPage,
    msr::{get_sev_status, SevStatus},
    secrets::SecretsPage,
};
use x86_64::instructions::{hlt, interrupts::int3};

mod asm;
mod ghcb;
mod serial;

#[no_mangle]
pub extern "C" fn rust64_start(_: u64, boot_params: &BootParams) -> ! {
    serial::init_logging();
    log::info!("Hello World!");

    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ENABLED) {
        log::info!("SEV enabled");
    }
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        log::info!("SEV-ES enabled");
    }
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        log::info!("SNP active");

        // Walk the null-terminated setup_data list until we find the CC blob.
        let mut setup_data = boot_params.hdr.setup_data;

        // Safety: there's a lot of dereferences of raw pointers here; unfortunately necessary as we
        // need to access the C data structures set up by the bootloader. It is safe as long as the
        // bootloader populated the zero page correctly.
        while !setup_data.is_null() {
            let type_ = unsafe { (*setup_data).type_ };
            if type_ == SetupDataType::CCBlob {
                break;
            }
            setup_data = unsafe { (*setup_data).next };
        }

        if !setup_data.is_null() {
            let cc_setup_data = unsafe { &*(setup_data as *const CCSetupData) };
            let cc_blob = unsafe { &*(cc_setup_data.cc_blob_address as *const CCBlobSevInfo) };
            let magic = cc_blob.magic;
            if magic != oak_linux_boot_params::CC_BLOB_SEV_INFO_MAGIC {
                panic!("CCBlobSevInfo data structure has invalid magic: {}", magic);
            }
            let cpuid = unsafe { &*(cc_blob.cpuid_phys as *const CpuidPage) };
            let secrets = unsafe { &*(cc_blob.secrets_phys as *const SecretsPage) };

            log::info!("CPUID page: {:?}", cpuid);
            log::info!("Secrets page: {:?}", secrets);
        } else {
            log::warn!("CCSetupData data structure not found in setup_data");
        }
    }
    panic!();
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        let mut ghcb_protocol = crate::ghcb::init_ghcb(sev_status.contains(SevStatus::SNP_ACTIVE));
        // Use the i8042 device to shut down the machine. Assumes the VMM exposes the device.
        // This is safe as both qemu and crosvm expose the i8042 device by default.
        let _ = ghcb_protocol.io_write_u8(0x64, 0xFE);
    } else {
        // Trigger a breakpoint exception. As we don't have a #BP handler, this will triple fault
        // and terminate the program.
        int3();
    }

    // But if we're still here for some reason, just go in an infinite halt loop.
    loop {
        hlt();
    }
}
<commit_msg>Add a crude #GP handler to sev_snp_hello_world.<commit_after>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![no_std]
#![no_main]
#![feature(abi_x86_interrupt)]

use core::{arch::asm, ops::Deref, panic::PanicInfo};
use lazy_static::lazy_static;
use oak_linux_boot_params::{BootParams, CCBlobSevInfo, CCSetupData, SetupDataType};
use sev_guest::{
    cpuid::CpuidPage,
    msr::{get_sev_status, SevStatus},
    secrets::SecretsPage,
};
use x86_64::{
    instructions::{hlt, interrupts::int3},
    structures::idt::{InterruptDescriptorTable, InterruptStackFrame},
};

mod asm;
mod ghcb;
mod serial;

lazy_static! {
    static ref IDT: InterruptDescriptorTable = {
        let mut idt = InterruptDescriptorTable::new();
        idt.general_protection_fault.set_handler_fn(gp_handler);
        idt
    };
}

/// Crude implementation of a #GP handler that skips any `rdmsr` instructions that trigger #GP.
///
/// Note that this doesn't properly emulate reading the MSR! It doesn't change any register values;
/// the instruction is just skipped.
extern "x86-interrupt" fn gp_handler(mut frame: InterruptStackFrame, _: u64) {
    unsafe {
        // RDMSR can cause a #GP if you try to read a non-existing MSR
        if *frame.deref().instruction_pointer.as_ptr::<u16>() == 0x320Fu16 {
            // just skip the instruction
            frame.as_mut().update(|val| {
                val.instruction_pointer += 2u64;
            });
        } else {
            // not a RDMSR, trigger a double fault
            asm!("int 8", options(noreturn));
        }
    }
}

#[no_mangle]
pub extern "C" fn rust64_start(_: u64, boot_params: &BootParams) -> ! {
    IDT.load();
    serial::init_logging();
    log::info!("Hello World!");

    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ENABLED) {
        log::info!("SEV enabled");
    }
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        log::info!("SEV-ES enabled");
    }
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        log::info!("SNP active");

        // Walk the null-terminated setup_data list until we find the CC blob.
        let mut setup_data = boot_params.hdr.setup_data;

        // Safety: there's a lot of dereferences of raw pointers here; unfortunately necessary as we
        // need to access the C data structures set up by the bootloader. It is safe as long as the
        // bootloader populated the zero page correctly.
        while !setup_data.is_null() {
            let type_ = unsafe { (*setup_data).type_ };
            if type_ == SetupDataType::CCBlob {
                break;
            }
            setup_data = unsafe { (*setup_data).next };
        }

        if !setup_data.is_null() {
            let cc_setup_data = unsafe { &*(setup_data as *const CCSetupData) };
            let cc_blob = unsafe { &*(cc_setup_data.cc_blob_address as *const CCBlobSevInfo) };
            let magic = cc_blob.magic;
            if magic != oak_linux_boot_params::CC_BLOB_SEV_INFO_MAGIC {
                panic!("CCBlobSevInfo data structure has invalid magic: {}", magic);
            }
            let cpuid = unsafe { &*(cc_blob.cpuid_phys as *const CpuidPage) };
            let secrets = unsafe { &*(cc_blob.secrets_phys as *const SecretsPage) };

            log::info!("CPUID page: {:?}", cpuid);
            log::info!("Secrets page: {:?}", secrets);
        } else {
            log::warn!("CCSetupData data structure not found in setup_data");
        }
    }
    panic!();
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        let mut ghcb_protocol = crate::ghcb::init_ghcb(sev_status.contains(SevStatus::SNP_ACTIVE));
        // Use the i8042 device to shut down the machine. Assumes the VMM exposes the device.
        // This is safe as both qemu and crosvm expose the i8042 device by default.
        let _ = ghcb_protocol.io_write_u8(0x64, 0xFE);
    } else {
        // Trigger a breakpoint exception. As we don't have a #BP handler, this will triple fault
        // and terminate the program.
        int3();
    }

    // But if we're still here for some reason, just go in an infinite halt loop.
    loop {
        hlt();
    }
}
<|endoftext|>


<commit_before>use rustc_ast::ast::{Lit, LitFloatType, LitIntType, LitKind};
use std::iter;

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Radix {
    Binary,
    Octal,
    Decimal,
    Hexadecimal,
}

impl Radix {
    /// Returns a reasonable digit group size for this radix.
    #[must_use]
    fn suggest_grouping(self) -> usize {
        match self {
            Self::Binary | Self::Hexadecimal => 4,
            Self::Octal | Self::Decimal => 3,
        }
    }
}

/// A helper method to format numeric literals with digit grouping.
/// `lit` must be a valid numeric literal without suffix.
pub fn format(lit: &str, type_suffix: Option<&str>, float: bool) -> String {
    NumericLiteral::new(lit, type_suffix, float).format()
}

#[derive(Debug)]
pub struct NumericLiteral<'a> {
    /// Which radix the literal was represented in.
    pub radix: Radix,
    /// The radix prefix, if present.
    pub prefix: Option<&'a str>,

    /// The integer part of the number.
    pub integer: &'a str,
    /// The fraction part of the number.
    pub fraction: Option<&'a str>,
    /// The exponent separator (b'e' or b'E') including preceding underscore if present
    /// and the exponent part.
    pub exponent: Option<(&'a str, &'a str)>,

    /// The type suffix, including preceding underscore if present.
    pub suffix: Option<&'a str>,
}

impl<'a> NumericLiteral<'a> {
    pub fn from_lit(src: &'a str, lit: &Lit) -> Option<NumericLiteral<'a>> {
        NumericLiteral::from_lit_kind(src, &lit.kind)
    }

    pub fn from_lit_kind(src: &'a str, lit_kind: &LitKind) -> Option<NumericLiteral<'a>> {
        if lit_kind.is_numeric() && src.chars().next().map_or(false, |c| c.is_digit(10)) {
            let (unsuffixed, suffix) = split_suffix(src, lit_kind);
            let float = matches!(lit_kind, LitKind::Float(..));
            Some(NumericLiteral::new(unsuffixed, suffix, float))
        } else {
            None
        }
    }

    #[must_use]
    pub fn new(lit: &'a str, suffix: Option<&'a str>, float: bool) -> Self {
        // Determine delimiter for radix prefix, if present, and radix.
        let radix = if lit.starts_with("0x") {
            Radix::Hexadecimal
        } else if lit.starts_with("0b") {
            Radix::Binary
        } else if lit.starts_with("0o") {
            Radix::Octal
        } else {
            Radix::Decimal
        };

        // Grab part of the literal after prefix, if present.
        let (prefix, mut sans_prefix) = if radix == Radix::Decimal {
            (None, lit)
        } else {
            let (p, s) = lit.split_at(2);
            (Some(p), s)
        };

        if suffix.is_some() && sans_prefix.ends_with('_') {
            // The '_' before the suffix isn't part of the digits
            sans_prefix = &sans_prefix[..sans_prefix.len() - 1];
        }

        let (integer, fraction, exponent) = Self::split_digit_parts(sans_prefix, float);

        Self {
            radix,
            prefix,
            integer,
            fraction,
            exponent,
            suffix,
        }
    }

    pub fn is_decimal(&self) -> bool {
        self.radix == Radix::Decimal
    }

    pub fn split_digit_parts(digits: &str, float: bool) -> (&str, Option<&str>, Option<(&str, &str)>) {
        let mut integer = digits;
        let mut fraction = None;
        let mut exponent = None;

        if float {
            for (i, c) in digits.char_indices() {
                match c {
                    '.' => {
                        integer = &digits[..i];
                        fraction = Some(&digits[i + 1..]);
                    },
                    'e' | 'E' => {
                        let exp_start = if digits[..i].ends_with('_') { i - 1 } else { i };

                        if integer.len() > exp_start {
                            integer = &digits[..exp_start];
                        } else {
                            fraction = Some(&digits[integer.len() + 1..exp_start]);
                        };
                        exponent = Some((&digits[exp_start..=i], &digits[i + 1..]));
                        break;
                    },
                    _ => {},
                }
            }
        }

        (integer, fraction, exponent)
    }

    /// Returns literal formatted in a sensible way.
    pub fn format(&self) -> String {
        let mut output = String::new();

        if let Some(prefix) = self.prefix {
            output.push_str(prefix);
        }

        let group_size = self.radix.suggest_grouping();

        Self::group_digits(
            &mut output,
            self.integer,
            group_size,
            true,
            self.radix == Radix::Hexadecimal,
        );

        if let Some(fraction) = self.fraction {
            output.push('.');
            Self::group_digits(&mut output, fraction, group_size, false, false);
        }

        if let Some((separator, exponent)) = self.exponent {
            output.push_str(separator);
            Self::group_digits(&mut output, exponent, group_size, true, false);
        }

        if let Some(suffix) = self.suffix {
            if output.ends_with('.') {
                output.push('0');
            }
            output.push('_');
            output.push_str(suffix);
        }

        output
    }

    pub fn group_digits(output: &mut String, input: &str, group_size: usize, partial_group_first: bool, pad: bool) {
        debug_assert!(group_size > 0);

        let mut digits = input.chars().filter(|&c| c != '_');

        // The exponent may have a sign, output it early, otherwise it will be
        // treated as a digit
        if digits.clone().next() == Some('-') {
            let _ = digits.next();
            output.push('-');
        }

        let first_group_size;

        if partial_group_first {
            first_group_size = (digits.clone().count() - 1) % group_size + 1;
            if pad {
                for _ in 0..group_size - first_group_size {
                    output.push('0');
                }
            }
        } else {
            first_group_size = group_size;
        }

        for _ in 0..first_group_size {
            if let Some(digit) = digits.next() {
                output.push(digit);
            }
        }

        for (c, i) in iter::zip(digits, (0..group_size).cycle()) {
            if i == 0 {
                output.push('_');
            }
            output.push(c);
        }
    }
}

fn split_suffix<'a>(src: &'a str, lit_kind: &LitKind) -> (&'a str, Option<&'a str>) {
    debug_assert!(lit_kind.is_numeric());
    lit_suffix_length(lit_kind).map_or((src, None), |suffix_length| {
        let (unsuffixed, suffix) = src.split_at(src.len() - suffix_length);
        (unsuffixed, Some(suffix))
    })
}

fn lit_suffix_length(lit_kind: &LitKind) -> Option<usize> {
    debug_assert!(lit_kind.is_numeric());
    let suffix = match lit_kind {
        LitKind::Int(_, int_lit_kind) => match int_lit_kind {
            LitIntType::Signed(int_ty) => Some(int_ty.name_str()),
            LitIntType::Unsigned(uint_ty) => Some(uint_ty.name_str()),
            LitIntType::Unsuffixed => None,
        },
        LitKind::Float(_, float_lit_kind) => match float_lit_kind {
            LitFloatType::Suffixed(float_ty) => Some(float_ty.name_str()),
            LitFloatType::Unsuffixed => None,
        },
        _ => None,
    };

    suffix.map(str::len)
}
<commit_msg>Drop exponent suggestion when 0<commit_after>use rustc_ast::ast::{Lit, LitFloatType, LitIntType, LitKind};
use std::iter;

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Radix {
    Binary,
    Octal,
    Decimal,
    Hexadecimal,
}

impl Radix {
    /// Returns a reasonable digit group size for this radix.
    #[must_use]
    fn suggest_grouping(self) -> usize {
        match self {
            Self::Binary | Self::Hexadecimal => 4,
            Self::Octal | Self::Decimal => 3,
        }
    }
}

/// A helper method to format numeric literals with digit grouping.
/// `lit` must be a valid numeric literal without suffix.
pub fn format(lit: &str, type_suffix: Option<&str>, float: bool) -> String {
    NumericLiteral::new(lit, type_suffix, float).format()
}

#[derive(Debug)]
pub struct NumericLiteral<'a> {
    /// Which radix the literal was represented in.
    pub radix: Radix,
    /// The radix prefix, if present.
    pub prefix: Option<&'a str>,

    /// The integer part of the number.
    pub integer: &'a str,
    /// The fraction part of the number.
    pub fraction: Option<&'a str>,
    /// The exponent separator (b'e' or b'E') including preceding underscore if present
    /// and the exponent part.
    pub exponent: Option<(&'a str, &'a str)>,

    /// The type suffix, including preceding underscore if present.
    pub suffix: Option<&'a str>,
}

impl<'a> NumericLiteral<'a> {
    pub fn from_lit(src: &'a str, lit: &Lit) -> Option<NumericLiteral<'a>> {
        NumericLiteral::from_lit_kind(src, &lit.kind)
    }

    pub fn from_lit_kind(src: &'a str, lit_kind: &LitKind) -> Option<NumericLiteral<'a>> {
        if lit_kind.is_numeric() && src.chars().next().map_or(false, |c| c.is_digit(10)) {
            let (unsuffixed, suffix) = split_suffix(src, lit_kind);
            let float = matches!(lit_kind, LitKind::Float(..));
            Some(NumericLiteral::new(unsuffixed, suffix, float))
        } else {
            None
        }
    }

    #[must_use]
    pub fn new(lit: &'a str, suffix: Option<&'a str>, float: bool) -> Self {
        // Determine delimiter for radix prefix, if present, and radix.
        let radix = if lit.starts_with("0x") {
            Radix::Hexadecimal
        } else if lit.starts_with("0b") {
            Radix::Binary
        } else if lit.starts_with("0o") {
            Radix::Octal
        } else {
            Radix::Decimal
        };

        // Grab part of the literal after prefix, if present.
        let (prefix, mut sans_prefix) = if radix == Radix::Decimal {
            (None, lit)
        } else {
            let (p, s) = lit.split_at(2);
            (Some(p), s)
        };

        if suffix.is_some() && sans_prefix.ends_with('_') {
            // The '_' before the suffix isn't part of the digits
            sans_prefix = &sans_prefix[..sans_prefix.len() - 1];
        }

        let (integer, fraction, exponent) = Self::split_digit_parts(sans_prefix, float);

        Self {
            radix,
            prefix,
            integer,
            fraction,
            exponent,
            suffix,
        }
    }

    pub fn is_decimal(&self) -> bool {
        self.radix == Radix::Decimal
    }

    pub fn split_digit_parts(digits: &str, float: bool) -> (&str, Option<&str>, Option<(&str, &str)>) {
        let mut integer = digits;
        let mut fraction = None;
        let mut exponent = None;

        if float {
            for (i, c) in digits.char_indices() {
                match c {
                    '.' => {
                        integer = &digits[..i];
                        fraction = Some(&digits[i + 1..]);
                    },
                    'e' | 'E' => {
                        let exp_start = if digits[..i].ends_with('_') { i - 1 } else { i };

                        if integer.len() > exp_start {
                            integer = &digits[..exp_start];
                        } else {
                            fraction = Some(&digits[integer.len() + 1..exp_start]);
                        };
                        exponent = Some((&digits[exp_start..=i], &digits[i + 1..]));
                        break;
                    },
                    _ => {},
                }
            }
        }

        (integer, fraction, exponent)
    }

    /// Returns literal formatted in a sensible way.
    pub fn format(&self) -> String {
        let mut output = String::new();

        if let Some(prefix) = self.prefix {
            output.push_str(prefix);
        }

        let group_size = self.radix.suggest_grouping();

        Self::group_digits(
            &mut output,
            self.integer,
            group_size,
            true,
            self.radix == Radix::Hexadecimal,
        );

        if let Some(fraction) = self.fraction {
            output.push('.');
            Self::group_digits(&mut output, fraction, group_size, false, false);
        }

        if let Some((separator, exponent)) = self.exponent {
            if exponent != "0" {
                output.push_str(separator);
                Self::group_digits(&mut output, exponent, group_size, true, false);
            }
        }

        if let Some(suffix) = self.suffix {
            if output.ends_with('.') {
                output.push('0');
            }
            output.push('_');
            output.push_str(suffix);
        }

        output
    }

    pub fn group_digits(output: &mut String, input: &str, group_size: usize, partial_group_first: bool, pad: bool) {
        debug_assert!(group_size > 0);

        let mut digits = input.chars().filter(|&c| c != '_');

        // The exponent may have a sign, output it early, otherwise it will be
        // treated as a digit
        if digits.clone().next() == Some('-') {
            let _ = digits.next();
            output.push('-');
        }

        let first_group_size;

        if partial_group_first {
            first_group_size = (digits.clone().count() - 1) % group_size + 1;
            if pad {
                for _ in 0..group_size - first_group_size {
                    output.push('0');
                }
            }
        } else {
            first_group_size = group_size;
        }

        for _ in 0..first_group_size {
            if let Some(digit) = digits.next() {
                output.push(digit);
            }
        }

        for (c, i) in iter::zip(digits, (0..group_size).cycle()) {
            if i == 0 {
                output.push('_');
            }
            output.push(c);
        }
    }
}

fn split_suffix<'a>(src: &'a str, lit_kind: &LitKind) -> (&'a str, Option<&'a str>) {
    debug_assert!(lit_kind.is_numeric());
    lit_suffix_length(lit_kind).map_or((src, None), |suffix_length| {
        let (unsuffixed, suffix) = src.split_at(src.len() - suffix_length);
        (unsuffixed, Some(suffix))
    })
}

fn lit_suffix_length(lit_kind: &LitKind) -> Option<usize> {
    debug_assert!(lit_kind.is_numeric());
    let suffix = match lit_kind {
        LitKind::Int(_, int_lit_kind) => match int_lit_kind {
            LitIntType::Signed(int_ty) => Some(int_ty.name_str()),
            LitIntType::Unsigned(uint_ty) => Some(uint_ty.name_str()),
            LitIntType::Unsuffixed => None,
        },
        LitKind::Float(_, float_lit_kind) => match float_lit_kind {
            LitFloatType::Suffixed(float_ty) => Some(float_ty.name_str()),
            LitFloatType::Unsuffixed => None,
        },
        _ => None,
    };

    suffix.map(str::len)
}
<|endoftext|>


<commit_before><commit_msg>green<commit_after><|endoftext|>


<commit_before><commit_msg>added from_hex, to_hex for big<commit_after><|endoftext|>


<commit_before>use parser::Identifier;
use operations::{Operation, Error};
use operations::item_type::{ItemType, FuncArgType};
use operations::scope::{ScopeStack, ScopeItem, TypeId};

pub fn define_u8(scope: &mut ScopeStack) -> TypeId {
    // Taking advantage of the scope system to simulate modules
    // This will be replaced with something better in:
    // https://github.com/brain-lang/brain/issues/37
    scope.push_scope();

    let unit_type = scope.primitives().unit();

    let u8_type = scope.declare_type(
        Identifier::from("u8"),
        ItemType::Primitive(1)
    );
    scope.register_primitive("u8", u8_type);

    scope.declare_builtin_function(
        // Special method for converting from literal
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::convert::From<{unsigned integer}>"),
        ItemType::Function {
            // This takes a single literal of the type specific within
            // the curly braces {} in the name
            args: vec![FuncArgType::Array {item: u8_type, size: None}],
            // Return type signifies which type we are declaring supports integer literals
            return_type: u8_type,
        },
        move |scope, args, target| {
            let value = args[0].numeric_literal_value();

            // 8 is the size of this numeric type in bytes
            // We use >= here because 0 is reserved for zero so we
            // have (2^size - 1) numbers available
            if value >= (1 << 8) {
                Err(Error::OverflowingLiteral {
                    typ: scope.get_type(u8_type).clone()
                })
            }
            else {
                Ok(vec![
                    Operation::Increment {
                        target: target.position(),
                        amount: value as usize,
                    }
                ])
            }
        }
    );

    scope.declare_builtin_function(
        // Special method for displaying this primitive (used from print/println)
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::fmt::Display::print"),
        ItemType::Function {
            args: vec![FuncArgType::Arg(u8_type)],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::TypedBlock {memory, ..} => memory,
                _ => unreachable!(),
            };

            Ok(vec![
                Operation::Write {
                    target: mem,
                }
            ])
        }
    );

    scope.declare_builtin_function(
        // Special method for displaying this primitive (used from print/println)
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::fmt::Display::print"),
        ItemType::Function {
            args: vec![FuncArgType::Array {item: u8_type, size: None}],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::Array {memory, ..} => memory,
                _ => unreachable!(),
            };

            Ok(vec![
                Operation::Write {
                    target: mem,
                }
            ])
        }
    );

    u8_type
}

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

    use operations::scope::{ScopeType};

    #[test]
    fn defines_primitive() {
        let mut scope = ScopeStack::new();
        define_u8(&mut scope);

        let u8_type_id = match **scope.lookup_type(&Identifier::from("u8")).first().unwrap() {
            ScopeType::Type(id) => id,
        };
        assert_eq!(scope.primitives().u8(), u8_type_id);
    }
}
<commit_msg>increment() and decrement() functions for u8<commit_after>use parser::Identifier;
use operations::{Operation, Error};
use operations::item_type::{ItemType, FuncArgType};
use operations::scope::{ScopeStack, ScopeItem, TypeId};

pub fn define_u8(scope: &mut ScopeStack) -> TypeId {
    // Taking advantage of the scope system to simulate modules
    // This will be replaced with something better in:
    // https://github.com/brain-lang/brain/issues/37
    scope.push_scope();

    let unit_type = scope.primitives().unit();

    let u8_type = scope.declare_type(
        Identifier::from("u8"),
        ItemType::Primitive(1)
    );
    scope.register_primitive("u8", u8_type);

    scope.declare_builtin_function(
        // Special method for converting from literal
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::convert::From<{unsigned integer}>"),
        ItemType::Function {
            // This takes a single literal of the type specific within
            // the curly braces {} in the name
            args: vec![FuncArgType::Array {item: u8_type, size: None}],
            // Return type signifies which type we are declaring supports integer literals
            return_type: u8_type,
        },
        move |scope, args, target| {
            let value = args[0].numeric_literal_value();

            // 8 is the size of this numeric type in bytes
            // We use >= here because 0 is reserved for zero so we
            // have (2^size - 1) numbers available
            if value >= (1 << 8) {
                Err(Error::OverflowingLiteral {
                    typ: scope.get_type(u8_type).clone()
                })
            }
            else {
                Ok(vec![
                    Operation::Increment {
                        target: target.position(),
                        amount: value as usize,
                    }
                ])
            }
        }
    );

    scope.declare_builtin_function(
        // Special method for displaying this primitive (used from print/println)
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::fmt::Display::print"),
        ItemType::Function {
            args: vec![FuncArgType::Arg(u8_type)],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::TypedBlock {memory, ..} => memory,
                _ => unreachable!(),
            };

            Ok(vec![
                Operation::Write {
                    target: mem,
                }
            ])
        }
    );

    scope.declare_builtin_function(
        Identifier::from("increment"),
        ItemType::Function {
            args: vec![FuncArgType::Arg(u8_type)],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::TypedBlock {memory, ..} => memory,
                _ => unreachable!(),
            };

            debug_assert!(mem.size() == 1);

            Ok(vec![
                Operation::Increment {
                    target: mem.position(),
                    amount: 1,
                }
            ])
        }
    );

    scope.declare_builtin_function(
        Identifier::from("decrement"),
        ItemType::Function {
            args: vec![FuncArgType::Arg(u8_type)],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::TypedBlock {memory, ..} => memory,
                _ => unreachable!(),
            };

            debug_assert!(mem.size() == 1);

            Ok(vec![
                Operation::Decrement {
                    target: mem.position(),
                    amount: 1,
                }
            ])
        }
    );

    // Method on [u8; _]
    scope.declare_builtin_function(
        // Special method for displaying this primitive (used from print/println)
        // This name is such that it could never be called directly
        // from the language itself
        Identifier::from("std::fmt::Display::print"),
        ItemType::Function {
            args: vec![FuncArgType::Array {item: u8_type, size: None}],
            return_type: unit_type,
        },
        move |scope, args, target| {
            let mem = match args[0] {
                ScopeItem::Array {memory, ..} => memory,
                _ => unreachable!(),
            };

            Ok(vec![
                Operation::Write {
                    target: mem,
                }
            ])
        }
    );

    u8_type
}

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

    use operations::scope::{ScopeType};

    #[test]
    fn defines_primitive() {
        let mut scope = ScopeStack::new();
        define_u8(&mut scope);

        let u8_type_id = match **scope.lookup_type(&Identifier::from("u8")).first().unwrap() {
            ScopeType::Type(id) => id,
        };
        assert_eq!(scope.primitives().u8(), u8_type_id);
    }
}
<|endoftext|>


<commit_before>/// Define a type that supports parsing and printing a multi-character symbol
/// as if it were a token.
///
/// # Usage
///
/// As a convention, it is recommended that this macro be invoked within a
/// module called `punct` or `punctuation` and that the resulting parser be invoked
/// with a `punct::` or `punctuation::` prefix.
///
/// ```edition2018
/// mod punct {
///     syn::custom_punctuation!(LeftRightArrow, <=>);
/// }
/// ```
///
/// The generated syntax tree node supports the following operations just like
/// any built-in punctuation token.
///
/// - [Peeking] — `input.peek(punct::LeftRightArrow)`
///
/// - [Parsing] — `input.parse::<punct::LeftRightArrow>()?`
///
/// - [Printing] — `quote!( ... #left_right_arrow ... )`
///
/// - Construction from a [`Span`] — `let left_right_arrow = punct::LeftRightArrow(sp)`
///
/// - Construction from multiple [`Span`] — `let left_right_arrow = punct::LeftRightArrow([sp, sp, sp])`
///
/// - Field access to its spans — `let spans = left_right_arrow.spans`
///
/// [Peeking]: parse/struct.ParseBuffer.html#method.peek
/// [Parsing]: parse/struct.ParseBuffer.html#method.parse
/// [Printing]: https://docs.rs/quote/0.6/quote/trait.ToTokens.html
/// [`Span`]: struct.Span.html
///
/// # Example
///
/// ```edition2018
/// use syn::{Expr, ExprParen, parenthesized};
/// use syn::punctuated::Punctuated;
/// use syn::parse::{Parse, ParseStream, Result};
/// use syn::token::Paren;
///
/// mod punct {
///     syn::custom_punctuation!(PathSeparator, </>);
/// }
///
/// // (expr) </> (expr) </> (expr) ...
/// struct PathSegments {
///     segments: Punctuated<Expr, punct::PathSeparator>,
/// }
///
/// impl Parse for PathSegments {
///     fn parse(input: ParseStream) -> Result<Self> {
///         let mut segments = Punctuated::new();
///
///         let la = input.lookahead1();
///         if la.peek(Paren) {
///             let content;
///             let paren_token = parenthesized!(content in input);
///             let expr = Box::new(content.parse()?);
///             segments.push_value(Expr::Paren(ExprParen { attrs: vec![], paren_token, expr }));
///         } else {
///             return Err(la.error());
///         }
///
///         while input.peek(punct::PathSeparator) {
///             segments.push_punct(input.parse()?);
///             let content;
///             let paren_token = parenthesized!(content in input);
///             let expr = Box::new(content.parse()?);
///             segments.push_value(Expr::Paren(ExprParen { attrs: vec![], paren_token, expr }));
///         }
///
///         Ok(PathSegments { segments })
///     }
/// }
///
/// let _: PathSegments = syn::parse_str(r#"("five") </> ("hundred")"#).unwrap();
/// ```
///
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        pub struct $ident {
            pub spans: custom_punctuation_repr!($($tt)+),
        }

        #[doc(hidden)]
        #[allow(non_snake_case)]
        pub fn $ident<__S: $crate::export::IntoSpans<custom_punctuation_repr!($($tt)+)>>(
            spans: __S,
        ) -> $ident {
            let _validate_len = 0 $(+ custom_punctuation_len!(strict, $tt))*;
            $ident {
                spans: $crate::export::IntoSpans::into_spans(spans)
            }
        }

        impl $crate::export::Default for $ident {
            fn default() -> Self {
                $ident($crate::export::Span::call_site())
            }
        }

        impl_parse_for_custom_punctuation!($ident, $($tt)+);
        impl_to_tokens_for_custom_punctuation!($ident, $($tt)+);
        impl_clone_for_custom_punctuation!($ident, $($tt)+);
        impl_extra_traits_for_custom_punctuation!($ident, $($tt)+);
    };
}

// Not public API.
#[cfg(feature = "parsing")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_parse_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::token::CustomToken for $ident {
            fn peek(cursor: $crate::buffer::Cursor) -> bool {
                $crate::token::parsing::peek_punct(cursor, stringify_punct!($($tt)+))
            }

            fn display() -> &'static $crate::export::str {
                custom_punctuation_concat!("`", stringify_punct!($($tt)+), "`")
            }
        }

        impl $crate::parse::Parse for $ident {
            fn parse(input: $crate::parse::ParseStream) -> $crate::parse::Result<$ident> {
                let spans: custom_punctuation_repr!($($tt)+) =
                    $crate::token::parsing::punct(input, stringify_punct!($($tt)+))?;
                Ok($ident(spans))
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "parsing"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_parse_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "printing")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_to_tokens_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::ToTokens for $ident {
            fn to_tokens(&self, tokens: &mut $crate::export::TokenStream2) {
                $crate::token::printing::punct(stringify_punct!($($tt)+), &self.spans, tokens)
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "printing"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_to_tokens_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "clone-impls")]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_clone_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::Copy for $ident {}

        impl $crate::export::Clone for $ident {
            fn clone(&self) -> Self {
                *self
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "clone-impls"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_clone_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "extra-traits")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_extra_traits_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::Debug for $ident {
            fn fmt(&self, f: &mut $crate::export::Formatter) -> $crate::export::fmt::Result {
                $crate::export::Formatter::write_str(f, stringify_punct!($($tt)+))
            }
        }

        impl $crate::export::Eq for $ident {}

        impl $crate::export::PartialEq for $ident {
            fn eq(&self, _other: &Self) -> $crate::export::bool {
                true
            }
        }

        impl $crate::export::Hash for $ident {
            fn hash<__H: $crate::export::Hasher>(&self, _state: &mut __H) {}
        }
    };
}

// Not public API.
#[cfg(not(feature = "extra-traits"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_extra_traits_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation_repr {
    ($($tt:tt)+) => {
        [$crate::export::Span; 0 $(+ custom_punctuation_len!(lenient, $tt))+]
    };
}

// Not public API.
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation_len {
    ($mode:ident, +)     => { 1 };
    ($mode:ident, +=)    => { 2 };
    ($mode:ident, &)     => { 1 };
    ($mode:ident, &&)    => { 2 };
    ($mode:ident, &=)    => { 2 };
    ($mode:ident, @)     => { 1 };
    ($mode:ident, !)     => { 1 };
    ($mode:ident, ^)     => { 1 };
    ($mode:ident, ^=)    => { 2 };
    ($mode:ident, :)     => { 1 };
    ($mode:ident, ::)    => { 2 };
    ($mode:ident, ,)     => { 1 };
    ($mode:ident, /)     => { 1 };
    ($mode:ident, /=)    => { 2 };
    ($mode:ident, .)     => { 1 };
    ($mode:ident, ..)    => { 2 };
    ($mode:ident, ...)   => { 3 };
    ($mode:ident, ..=)   => { 3 };
    ($mode:ident, =)     => { 1 };
    ($mode:ident, ==)    => { 2 };
    ($mode:ident, >=)    => { 2 };
    ($mode:ident, >)     => { 1 };
    ($mode:ident, <=)    => { 2 };
    ($mode:ident, <)     => { 1 };
    ($mode:ident, *=)    => { 2 };
    ($mode:ident, !=)    => { 2 };
    ($mode:ident, |)     => { 1 };
    ($mode:ident, |=)    => { 2 };
    ($mode:ident, ||)    => { 2 };
    ($mode:ident, #)     => { 1 };
    ($mode:ident, ?)     => { 1 };
    ($mode:ident, ->)    => { 2 };
    ($mode:ident, <-)    => { 2 };
    ($mode:ident, %)     => { 1 };
    ($mode:ident, %=)    => { 2 };
    ($mode:ident, =>)    => { 2 };
    ($mode:ident, ;)     => { 1 };
    ($mode:ident, <<)    => { 2 };
    ($mode:ident, <<=)   => { 3 };
    ($mode:ident, >>)    => { 2 };
    ($mode:ident, >>=)   => { 3 };
    ($mode:ident, *)     => { 1 };
    ($mode:ident, -)     => { 1 };
    ($mode:ident, -=)    => { 2 };
    ($mode:ident, ~)     => { 1 };
    (lenient, $tt:tt)    => { 0 };
    (strict, $tt:tt)     => {{ custom_punctuation_unexpected!($tt); 0 }};
}

// Not public API.
#[doc(hidden)]
#[macro_export]
macro_rules! custom_punctuation_unexpected {
    () => {};
}

// Not public API.
#[doc(hidden)]
#[macro_export]
macro_rules! stringify_punct {
    ($($tt:tt)+) => {
        concat!($(stringify!($tt)),+)
    };
}

// Not public API.
// Without this, local_inner_macros breaks when looking for concat!
#[doc(hidden)]
#[macro_export]
macro_rules! custom_punctuation_concat {
    ($($tt:tt)*) => {
        concat!($($tt)*)
    };
}
<commit_msg>Touch up custom_punctuation documentation<commit_after>/// Define a type that supports parsing and printing a multi-character symbol
/// as if it were a punctuation token.
///
/// # Usage
///
/// ```edition2018
/// syn::custom_punctuation!(LeftRightArrow, <=>);
/// ```
///
/// The generated syntax tree node supports the following operations just like
/// any built-in punctuation token.
///
/// - [Peeking] — `input.peek(LeftRightArrow)`
///
/// - [Parsing] — `input.parse::<LeftRightArrow>()?`
///
/// - [Printing] — `quote!( ... #lrarrow ... )`
///
/// - Construction from a [`Span`] — `let lrarrow = LeftRightArrow(sp)`
///
/// - Construction from multiple [`Span`] — `let lrarrow = LeftRightArrow([sp, sp, sp])`
///
/// - Field access to its spans — `let spans = lrarrow.spans`
///
/// [Peeking]: parse/struct.ParseBuffer.html#method.peek
/// [Parsing]: parse/struct.ParseBuffer.html#method.parse
/// [Printing]: https://docs.rs/quote/0.6/quote/trait.ToTokens.html
/// [`Span`]: struct.Span.html
///
/// # Example
///
/// ```edition2018
/// use syn::parse::{Parse, ParseStream, Result};
/// use syn::punctuated::Punctuated;
/// use syn::token::Paren;
/// use syn::{parenthesized, Expr, ExprParen};
///
/// syn::custom_punctuation!(PathSeparator, </>);
///
/// // (expr) </> (expr) </> (expr) ...
/// struct PathSegments {
///     segments: Punctuated<Expr, PathSeparator>,
/// }
///
/// impl Parse for PathSegments {
///     fn parse(input: ParseStream) -> Result<Self> {
///         let mut segments = Punctuated::new();
///
///         let la = input.lookahead1();
///         if la.peek(Paren) {
///             let content;
///             let paren_token = parenthesized!(content in input);
///             let expr = Box::new(content.parse()?);
///             segments.push_value(Expr::Paren(ExprParen { attrs: vec![], paren_token, expr }));
///         } else {
///             return Err(la.error());
///         }
///
///         while input.peek(PathSeparator) {
///             segments.push_punct(input.parse()?);
///             let content;
///             let paren_token = parenthesized!(content in input);
///             let expr = Box::new(content.parse()?);
///             segments.push_value(Expr::Paren(ExprParen { attrs: vec![], paren_token, expr }));
///         }
///
///         Ok(PathSegments { segments })
///     }
/// }
///
/// fn main() {
///     let input = r#" ("five") </> ("hundred") "#;
///     let _: PathSegments = syn::parse_str(input).unwrap();
/// }
/// ```
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        pub struct $ident {
            pub spans: custom_punctuation_repr!($($tt)+),
        }

        #[doc(hidden)]
        #[allow(non_snake_case)]
        pub fn $ident<__S: $crate::export::IntoSpans<custom_punctuation_repr!($($tt)+)>>(
            spans: __S,
        ) -> $ident {
            let _validate_len = 0 $(+ custom_punctuation_len!(strict, $tt))*;
            $ident {
                spans: $crate::export::IntoSpans::into_spans(spans)
            }
        }

        impl $crate::export::Default for $ident {
            fn default() -> Self {
                $ident($crate::export::Span::call_site())
            }
        }

        impl_parse_for_custom_punctuation!($ident, $($tt)+);
        impl_to_tokens_for_custom_punctuation!($ident, $($tt)+);
        impl_clone_for_custom_punctuation!($ident, $($tt)+);
        impl_extra_traits_for_custom_punctuation!($ident, $($tt)+);
    };
}

// Not public API.
#[cfg(feature = "parsing")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_parse_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::token::CustomToken for $ident {
            fn peek(cursor: $crate::buffer::Cursor) -> bool {
                $crate::token::parsing::peek_punct(cursor, stringify_punct!($($tt)+))
            }

            fn display() -> &'static $crate::export::str {
                custom_punctuation_concat!("`", stringify_punct!($($tt)+), "`")
            }
        }

        impl $crate::parse::Parse for $ident {
            fn parse(input: $crate::parse::ParseStream) -> $crate::parse::Result<$ident> {
                let spans: custom_punctuation_repr!($($tt)+) =
                    $crate::token::parsing::punct(input, stringify_punct!($($tt)+))?;
                Ok($ident(spans))
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "parsing"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_parse_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "printing")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_to_tokens_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::ToTokens for $ident {
            fn to_tokens(&self, tokens: &mut $crate::export::TokenStream2) {
                $crate::token::printing::punct(stringify_punct!($($tt)+), &self.spans, tokens)
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "printing"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_to_tokens_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "clone-impls")]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_clone_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::Copy for $ident {}

        impl $crate::export::Clone for $ident {
            fn clone(&self) -> Self {
                *self
            }
        }
    };
}

// Not public API.
#[cfg(not(feature = "clone-impls"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_clone_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[cfg(feature = "extra-traits")]
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! impl_extra_traits_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {
        impl $crate::export::Debug for $ident {
            fn fmt(&self, f: &mut $crate::export::Formatter) -> $crate::export::fmt::Result {
                $crate::export::Formatter::write_str(f, stringify_punct!($($tt)+))
            }
        }

        impl $crate::export::Eq for $ident {}

        impl $crate::export::PartialEq for $ident {
            fn eq(&self, _other: &Self) -> $crate::export::bool {
                true
            }
        }

        impl $crate::export::Hash for $ident {
            fn hash<__H: $crate::export::Hasher>(&self, _state: &mut __H) {}
        }
    };
}

// Not public API.
#[cfg(not(feature = "extra-traits"))]
#[doc(hidden)]
#[macro_export]
macro_rules! impl_extra_traits_for_custom_punctuation {
    ($ident:ident, $($tt:tt)+) => {};
}

// Not public API.
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation_repr {
    ($($tt:tt)+) => {
        [$crate::export::Span; 0 $(+ custom_punctuation_len!(lenient, $tt))+]
    };
}

// Not public API.
#[doc(hidden)]
#[macro_export(local_inner_macros)]
macro_rules! custom_punctuation_len {
    ($mode:ident, +)     => { 1 };
    ($mode:ident, +=)    => { 2 };
    ($mode:ident, &)     => { 1 };
    ($mode:ident, &&)    => { 2 };
    ($mode:ident, &=)    => { 2 };
    ($mode:ident, @)     => { 1 };
    ($mode:ident, !)     => { 1 };
    ($mode:ident, ^)     => { 1 };
    ($mode:ident, ^=)    => { 2 };
    ($mode:ident, :)     => { 1 };
    ($mode:ident, ::)    => { 2 };
    ($mode:ident, ,)     => { 1 };
    ($mode:ident, /)     => { 1 };
    ($mode:ident, /=)    => { 2 };
    ($mode:ident, .)     => { 1 };
    ($mode:ident, ..)    => { 2 };
    ($mode:ident, ...)   => { 3 };
    ($mode:ident, ..=)   => { 3 };
    ($mode:ident, =)     => { 1 };
    ($mode:ident, ==)    => { 2 };
    ($mode:ident, >=)    => { 2 };
    ($mode:ident, >)     => { 1 };
    ($mode:ident, <=)    => { 2 };
    ($mode:ident, <)     => { 1 };
    ($mode:ident, *=)    => { 2 };
    ($mode:ident, !=)    => { 2 };
    ($mode:ident, |)     => { 1 };
    ($mode:ident, |=)    => { 2 };
    ($mode:ident, ||)    => { 2 };
    ($mode:ident, #)     => { 1 };
    ($mode:ident, ?)     => { 1 };
    ($mode:ident, ->)    => { 2 };
    ($mode:ident, <-)    => { 2 };
    ($mode:ident, %)     => { 1 };
    ($mode:ident, %=)    => { 2 };
    ($mode:ident, =>)    => { 2 };
    ($mode:ident, ;)     => { 1 };
    ($mode:ident, <<)    => { 2 };
    ($mode:ident, <<=)   => { 3 };
    ($mode:ident, >>)    => { 2 };
    ($mode:ident, >>=)   => { 3 };
    ($mode:ident, *)     => { 1 };
    ($mode:ident, -)     => { 1 };
    ($mode:ident, -=)    => { 2 };
    ($mode:ident, ~)     => { 1 };
    (lenient, $tt:tt)    => { 0 };
    (strict, $tt:tt)     => {{ custom_punctuation_unexpected!($tt); 0 }};
}

// Not public API.
#[doc(hidden)]
#[macro_export]
macro_rules! custom_punctuation_unexpected {
    () => {};
}

// Not public API.
#[doc(hidden)]
#[macro_export]
macro_rules! stringify_punct {
    ($($tt:tt)+) => {
        concat!($(stringify!($tt)),+)
    };
}

// Not public API.
// Without this, local_inner_macros breaks when looking for concat!
#[doc(hidden)]
#[macro_export]
macro_rules! custom_punctuation_concat {
    ($($tt:tt)*) => {
        concat!($($tt)*)
    };
}
<|endoftext|>