<commit_before>use std::collections::HashMap;

pub struct Translator
{
	rules: HashMap<(char, char), char>,
}

impl Translator
{
	pub fn new() -> Translator
	{
		Translator
		{
			rules: get_default_rules(),
		}
	}

	// Returns translated characted and boolean value indicating that previous character should be erased.
	fn translate_char(&self, current : char, previous : char) -> (char, bool)
	{
		match self.rules.get(&(current, previous))
		{
			Some(result) => (*result, true),
			None => match self.rules.get(&(current, ' '))
			{
				Some(result) => (*result, false),
				None => (current, false),
			}
		}
	}

	pub fn translate(&self, text : &str) -> String
	{
		let mut result = String::with_capacity(text.len());

		let mut previous = ' ';
		for c in text.chars()
		{
			let (translated, erase_previous) = self.translate_char(c, previous);
			if erase_previous
			{
				result.pop();
			}

			result.push(translated);
			previous = translated;
		}

		result
	}
}

fn get_default_rules() -> HashMap<(char, char), char>
{
	let mut result = HashMap::with_capacity(50);

	result.insert(('a', ' '), 'а');
	result.insert(('b', ' '), 'б');
	result.insert(('v', ' '), 'в');
	result.insert(('g', ' '), 'г');
	result.insert(('d', ' '), 'д');
	result.insert(('e', ' '), 'е');
	result.insert(('o', 'j'), 'ё');
	result.insert(('o', 'й'), 'ё');
	result.insert(('h', 'z'), 'ж');
	result.insert(('h', 'з'), 'ж');
	result.insert(('z', ' '), 'з');
	result.insert(('i', ' '), 'и');
	result.insert(('j', ' '), 'й');
	result.insert(('k', ' '), 'к');
	result.insert(('l', ' '), 'л');
	result.insert(('m', ' '), 'м');
	result.insert(('n', ' '), 'н');
	result.insert(('o', ' '), 'о');
	result.insert(('p', ' '), 'п');
	result.insert(('r', ' '), 'р');
	result.insert(('s', ' '), 'с');
	result.insert(('t', ' '), 'т');
	result.insert(('u', ' '), 'у');
	result.insert(('f', ' '), 'ф');
	result.insert(('h', ' '), 'х');
	result.insert(('x', ' '), 'х');
	result.insert(('c', ' '), 'ц');
	result.insert(('h', 'c'), 'ч');
	result.insert(('h', 'ц'), 'ч');
	result.insert(('h', 'h'), 'ш');
	result.insert(('h', 'с'), 'ш');
	result.insert(('h', 'ш'), 'щ');
	result.insert(('w', ' '), 'щ');
	result.insert(('#', ' '), 'ъ');
	result.insert(('y', ' '), 'ы');
	result.insert(('\'', ' '), 'ь');
	result.insert(('e', 'j'), 'э');
	result.insert(('e', 'й'), 'э');
	result.insert(('u', 'j'), 'ю');
	result.insert(('u', 'й'), 'ю');
	result.insert(('a', 'j'), 'я');
	result.insert(('a', 'й'), 'я');

	result.insert(('A', ' '), 'А');
	result.insert(('B', ' '), 'Б');
	result.insert(('V', ' '), 'В');
	result.insert(('G', ' '), 'Г');
	result.insert(('D', ' '), 'Д');
	result.insert(('E', ' '), 'Е');
	result.insert(('O', 'J'), 'Ё');
	result.insert(('O', 'Й'), 'Ё');
	result.insert(('H', 'Z'), 'Ж');
	result.insert(('H', 'З'), 'Ж');
	result.insert(('Z', ' '), 'З');
	result.insert(('I', ' '), 'И');
	result.insert(('J', ' '), 'Й');
	result.insert(('K', ' '), 'К');
	result.insert(('L', ' '), 'Л');
	result.insert(('M', ' '), 'М');
	result.insert(('N', ' '), 'Н');
	result.insert(('O', ' '), 'О');
	result.insert(('P', ' '), 'П');
	result.insert(('R', ' '), 'Р');
	result.insert(('S', ' '), 'С');
	result.insert(('T', ' '), 'Т');
	result.insert(('U', ' '), 'У');
	result.insert(('F', ' '), 'Ф');
	result.insert(('H', ' '), 'Х');
	result.insert(('X', ' '), 'Х');
	result.insert(('C', ' '), 'Ц');
	result.insert(('H', 'C'), 'Ч');
	result.insert(('H', 'Ц'), 'Ч');
	result.insert(('H', 'H'), 'Ш');
	result.insert(('H', 'С'), 'Ш');
	result.insert(('H', 'Ш'), 'Щ');
	result.insert(('W', ' '), 'Щ');
	result.insert(('#', 'ъ'), 'Ъ');
	result.insert(('Y', ' '), 'Ы');
	result.insert(('\'', 'ь'), 'Ь');
	result.insert(('E', 'J'), 'Э');
	result.insert(('E', 'Й'), 'Э');
	result.insert(('U', 'J'), 'Ю');
	result.insert(('U', 'Й'), 'Ю');
	result.insert(('A', 'J'), 'Я');
	result.insert(('A', 'Й'), 'Я');

	result
}
<commit_msg>Spaces processing has been fixed<commit_after>use std::collections::HashMap;

pub struct Translator
{
	rules: HashMap<(char, char), char>,
}

impl Translator
{
	pub fn new() -> Translator
	{
		Translator
		{
			rules: get_default_rules(),
		}
	}

	// Returns translated characted and boolean value indicating that previous character should be erased.
	fn translate_char(&self, current : char, previous : char) -> (char, bool)
	{
		match self.rules.get(&(current, previous))
		{
			Some(result) => (*result, previous != ' '),
			None => match self.rules.get(&(current, ' '))
			{
				Some(result) => (*result, false),
				None => (current, false),
			}
		}
	}

	pub fn translate(&self, text : &str) -> String
	{
		let mut result = String::with_capacity(text.len());

		let mut previous = ' ';
		for c in text.chars()
		{
			let (translated, erase_previous) = self.translate_char(c, previous);
			if erase_previous
			{
				result.pop();
			}

			result.push(translated);
			previous = translated;
		}

		result
	}
}

fn get_default_rules() -> HashMap<(char, char), char>
{
	let mut result = HashMap::with_capacity(50);

	result.insert(('a', ' '), 'а');
	result.insert(('b', ' '), 'б');
	result.insert(('v', ' '), 'в');
	result.insert(('g', ' '), 'г');
	result.insert(('d', ' '), 'д');
	result.insert(('e', ' '), 'е');
	result.insert(('o', 'j'), 'ё');
	result.insert(('o', 'й'), 'ё');
	result.insert(('h', 'z'), 'ж');
	result.insert(('h', 'з'), 'ж');
	result.insert(('z', ' '), 'з');
	result.insert(('i', ' '), 'и');
	result.insert(('j', ' '), 'й');
	result.insert(('k', ' '), 'к');
	result.insert(('l', ' '), 'л');
	result.insert(('m', ' '), 'м');
	result.insert(('n', ' '), 'н');
	result.insert(('o', ' '), 'о');
	result.insert(('p', ' '), 'п');
	result.insert(('r', ' '), 'р');
	result.insert(('s', ' '), 'с');
	result.insert(('t', ' '), 'т');
	result.insert(('u', ' '), 'у');
	result.insert(('f', ' '), 'ф');
	result.insert(('h', ' '), 'х');
	result.insert(('x', ' '), 'х');
	result.insert(('c', ' '), 'ц');
	result.insert(('h', 'c'), 'ч');
	result.insert(('h', 'ц'), 'ч');
	result.insert(('h', 'h'), 'ш');
	result.insert(('h', 'с'), 'ш');
	result.insert(('h', 'ш'), 'щ');
	result.insert(('w', ' '), 'щ');
	result.insert(('#', ' '), 'ъ');
	result.insert(('y', ' '), 'ы');
	result.insert(('\'', ' '), 'ь');
	result.insert(('e', 'j'), 'э');
	result.insert(('e', 'й'), 'э');
	result.insert(('u', 'j'), 'ю');
	result.insert(('u', 'й'), 'ю');
	result.insert(('a', 'j'), 'я');
	result.insert(('a', 'й'), 'я');

	result.insert(('A', ' '), 'А');
	result.insert(('B', ' '), 'Б');
	result.insert(('V', ' '), 'В');
	result.insert(('G', ' '), 'Г');
	result.insert(('D', ' '), 'Д');
	result.insert(('E', ' '), 'Е');
	result.insert(('O', 'J'), 'Ё');
	result.insert(('O', 'Й'), 'Ё');
	result.insert(('H', 'Z'), 'Ж');
	result.insert(('H', 'З'), 'Ж');
	result.insert(('Z', ' '), 'З');
	result.insert(('I', ' '), 'И');
	result.insert(('J', ' '), 'Й');
	result.insert(('K', ' '), 'К');
	result.insert(('L', ' '), 'Л');
	result.insert(('M', ' '), 'М');
	result.insert(('N', ' '), 'Н');
	result.insert(('O', ' '), 'О');
	result.insert(('P', ' '), 'П');
	result.insert(('R', ' '), 'Р');
	result.insert(('S', ' '), 'С');
	result.insert(('T', ' '), 'Т');
	result.insert(('U', ' '), 'У');
	result.insert(('F', ' '), 'Ф');
	result.insert(('H', ' '), 'Х');
	result.insert(('X', ' '), 'Х');
	result.insert(('C', ' '), 'Ц');
	result.insert(('H', 'C'), 'Ч');
	result.insert(('H', 'Ц'), 'Ч');
	result.insert(('H', 'H'), 'Ш');
	result.insert(('H', 'С'), 'Ш');
	result.insert(('H', 'Ш'), 'Щ');
	result.insert(('W', ' '), 'Щ');
	result.insert(('#', 'ъ'), 'Ъ');
	result.insert(('Y', ' '), 'Ы');
	result.insert(('\'', 'ь'), 'Ь');
	result.insert(('E', 'J'), 'Э');
	result.insert(('E', 'Й'), 'Э');
	result.insert(('U', 'J'), 'Ю');
	result.insert(('U', 'Й'), 'Ю');
	result.insert(('A', 'J'), 'Я');
	result.insert(('A', 'Й'), 'Я');

	result
}
<|endoftext|>


<commit_before>// Copyright 2017 Alexander Reece
//
// 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::hash::{Hash, Hasher};

use phf::{ordered_map, OrderedMap, PhfHash};
use {Spec, Class, Constant, Version};

pub struct Classes(ordered_map::Values<'static, &'static str, Class>);
pub struct Constants(ordered_map::Values<'static, &'static str, Constant>);
pub struct Domains(ordered_map::Entries<'static, &'static str, &'static str>);

impl Spec {
    pub fn name(&self) -> &'static str {
        self.name
    }
    pub fn classes(&self) -> Classes {
        Classes(self.classes.values())
    }
    pub fn constants(&self) -> Constants {
        Constants(self.constants.values())
    }
    pub fn domains(&self) -> Domains {
        Domains(self.domains.entries())
    }
    pub fn frame_types(&self) -> Constants {
        Constants(self.frame_types.values())
    }
    pub fn response_codes(&self) -> Constants {
        Constants(self.response_codes.values())
    }
    pub fn version(&self) -> &Version {
        &self.version
    }
}

impl PartialEq for Spec {
    fn eq(&self, other: &Spec) -> bool {
        fn is_map_eq<K, V>(map1: &OrderedMap<K, V>, map2: &OrderedMap<K, V>) -> bool
            where K: Eq + PhfHash,
                  V: Eq
        {
            map1.entries().all(|(k, v1)| {
                match map2.get(k) {
                    Some(v2) if v1 == v2 => true,
                    _ => false,
                }
            })
        }

        self.name == other.name
            && self.version == other.version
            && is_map_eq(self.classes, other.classes)
            && is_map_eq(self.constants, other.constants)
            && is_map_eq(self.domains, other.domains)
            && is_map_eq(self.frame_types, other.frame_types)
            && is_map_eq(self.response_codes, other.response_codes)
    }
}

impl Eq for Spec {}

impl Hash for Spec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for (name, class) in self.classes {
            name.hash(state);
            class.hash(state);
        }
        for (name, value) in self.constants {
            name.hash(state);
            value.hash(state);
        }
        for (name, mapping) in self.domains {
            name.hash(state);
            mapping.hash(state);
        }
        for (name, value) in self.frame_types {
            name.hash(state);
            value.hash(state);
        }
        for (name, response_code) in self.response_codes {
            name.hash(state);
            response_code.hash(state);
        }
        self.version.hash(state)
    }
}

impl Iterator for Classes {
    type Item = &'static Class;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl Iterator for Constants {
    type Item = &'static Constant;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl Iterator for Domains {
    type Item = (&'static str, &'static str);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(name, mapping)| (*name, *mapping))
    }
}
<commit_msg>[specs] Add `Spec::domain()`<commit_after>// Copyright 2017 Alexander Reece
//
// 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::hash::{Hash, Hasher};

use phf::{ordered_map, OrderedMap, PhfHash};
use {Spec, Class, Constant, Version};

pub struct Classes(ordered_map::Values<'static, &'static str, Class>);
pub struct Constants(ordered_map::Values<'static, &'static str, Constant>);
pub struct Domains(ordered_map::Entries<'static, &'static str, &'static str>);

impl Spec {
    pub fn name(&self) -> &'static str {
        self.name
    }
    pub fn classes(&self) -> Classes {
        Classes(self.classes.values())
    }
    pub fn constants(&self) -> Constants {
        Constants(self.constants.values())
    }
    pub fn domain(&self, domain: &str) -> Option<&str> {
        self.domains.get(domain).map(|mapping| *mapping)
    }
    pub fn domains(&self) -> Domains {
        Domains(self.domains.entries())
    }
    pub fn frame_types(&self) -> Constants {
        Constants(self.frame_types.values())
    }
    pub fn response_codes(&self) -> Constants {
        Constants(self.response_codes.values())
    }
    pub fn version(&self) -> &Version {
        &self.version
    }
}

impl PartialEq for Spec {
    fn eq(&self, other: &Spec) -> bool {
        fn is_map_eq<K, V>(map1: &OrderedMap<K, V>, map2: &OrderedMap<K, V>) -> bool
            where K: Eq + PhfHash,
                  V: Eq
        {
            map1.entries().all(|(k, v1)| {
                match map2.get(k) {
                    Some(v2) if v1 == v2 => true,
                    _ => false,
                }
            })
        }

        self.name == other.name
            && self.version == other.version
            && is_map_eq(self.classes, other.classes)
            && is_map_eq(self.constants, other.constants)
            && is_map_eq(self.domains, other.domains)
            && is_map_eq(self.frame_types, other.frame_types)
            && is_map_eq(self.response_codes, other.response_codes)
    }
}

impl Eq for Spec {}

impl Hash for Spec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        for (name, class) in self.classes {
            name.hash(state);
            class.hash(state);
        }
        for (name, value) in self.constants {
            name.hash(state);
            value.hash(state);
        }
        for (name, mapping) in self.domains {
            name.hash(state);
            mapping.hash(state);
        }
        for (name, value) in self.frame_types {
            name.hash(state);
            value.hash(state);
        }
        for (name, response_code) in self.response_codes {
            name.hash(state);
            response_code.hash(state);
        }
        self.version.hash(state)
    }
}

impl Iterator for Classes {
    type Item = &'static Class;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl Iterator for Constants {
    type Item = &'static Constant;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl Iterator for Domains {
    type Item = (&'static str, &'static str);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(name, mapping)| (*name, *mapping))
    }
}
<|endoftext|>


<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 crate::command::studio::enter::ARTIFACT_PATH_ENVVAR;
use crate::common::ui::UI;
use crate::hcore::crypto::default_cache_key_path;
use crate::hcore::env as henv;
use crate::hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use crate::hcore::os::process;
use crate::hcore::package::target;

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

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

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

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    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 = if is_serving_windows_containers(&docker_cmd) {
        "c:"
    } else {
        ""
    };
    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("DO_CHECK"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_NOCOLORING"),
        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_else(|_| 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);
    let image = image_identifier_for_active_target(&docker_cmd);

    cmd.arg("images").arg(&image).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_for_active_target(&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.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.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()];
    let image = image_identifier_for_active_target(&docker_cmd);

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.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 using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") | Some("run") => {}
        _ => {
            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.into());
    cmd_args.extend_from_slice(args.as_slice());
    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    process::become_command(docker_cmd, cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(docker_cmd: &Path) -> String {
    image_identifier(
        is_serving_windows_containers(docker_cmd),
        target::PackageTarget::active_target(),
    )
}

/// 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(using_windows_containers: bool, target: &target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS.as_ref())
    } else {
        let t = match *target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2.as_ref(),
            _ => target::X86_64_LINUX.as_ref(),
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR)
        .unwrap_or_else(|_| format!("{}-{}:{}", img, studio_target, version[0]))
}

#[cfg(test)]
mod tests {
    use super::{image_identifier, DOCKER_IMAGE, DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        let windows_container = true;
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_DARWIN),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_LINUX),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_LINUX_KERNEL2),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_WINDOWS),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(windows_container, &target::X86_64_WINDOWS),
            format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION)
        );
        assert_eq!(
            image_identifier(windows_container, &target::X86_64_LINUX),
            format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION)
        );
    }
}
<commit_msg>include HAB_FEAT_INSTALL_HOOK in docker studio env_vars<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 crate::command::studio::enter::ARTIFACT_PATH_ENVVAR;
use crate::common::ui::UI;
use crate::hcore::crypto::default_cache_key_path;
use crate::hcore::env as henv;
use crate::hcore::fs::{find_command, CACHE_ARTIFACT_PATH, CACHE_KEY_PATH};
use crate::hcore::os::process;
use crate::hcore::package::target;

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

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

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

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    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 = if is_serving_windows_containers(&docker_cmd) {
        "c:"
    } else {
        ""
    };
    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("DO_CHECK"),
        String::from("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        String::from("HAB_FEAT_INSTALL_HOOK"),
        String::from("HAB_NOCOLORING"),
        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_else(|_| 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);
    let image = image_identifier_for_active_target(&docker_cmd);

    cmd.arg("images").arg(&image).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_for_active_target(&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.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.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()];
    let image = image_identifier_for_active_target(&docker_cmd);

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.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 using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") | Some("run") => {}
        _ => {
            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.into());
    cmd_args.extend_from_slice(args.as_slice());
    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    process::become_command(docker_cmd, cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(docker_cmd: &Path) -> String {
    image_identifier(
        is_serving_windows_containers(docker_cmd),
        target::PackageTarget::active_target(),
    )
}

/// 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(using_windows_containers: bool, target: &target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS.as_ref())
    } else {
        let t = match *target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2.as_ref(),
            _ => target::X86_64_LINUX.as_ref(),
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR)
        .unwrap_or_else(|_| format!("{}-{}:{}", img, studio_target, version[0]))
}

#[cfg(test)]
mod tests {
    use super::{image_identifier, DOCKER_IMAGE, DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        let windows_container = true;
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_DARWIN),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_LINUX),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_LINUX_KERNEL2),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION)
        );
        assert_eq!(
            image_identifier(!windows_container, &target::X86_64_WINDOWS),
            format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION)
        );
        assert_eq!(
            image_identifier(windows_container, &target::X86_64_WINDOWS),
            format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION)
        );
        assert_eq!(
            image_identifier(windows_container, &target::X86_64_LINUX),
            format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION)
        );
    }
}
<|endoftext|>


<commit_before>use crate::reflect::MessageDescriptor;
use crate::reflect::ReflectEqMode;
use crate::stream::WithCodedOutputStream;
use crate::CodedInputStream;
use crate::CodedOutputStream;
use crate::Message;
use crate::ProtobufError;
use crate::ProtobufResult;
use crate::UnknownFields;
use std::any::Any;
use std::any::TypeId;
use std::fmt;
use std::io::Write;

/// Dynamic-dispatch version of [`Message`].
pub trait MessageDyn: fmt::Debug + Send + Sync + 'static {
    /// Message descriptor for this message, used for reflection.
    fn descriptor_dyn(&self) -> MessageDescriptor;

    /// Update this message fields with contents of given stream.
    fn merge_from_dyn(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write the message.
    fn write_to_with_cached_sizes_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute (and cache) the message size.
    fn compute_size_dyn(&self) -> u32;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized_dyn(&self) -> bool;

    /// Get a reference to unknown fields.
    fn get_unknown_fields_dyn(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields_dyn(&mut self) -> &mut UnknownFields;
}

impl<M: Message> MessageDyn for M {
    fn descriptor_dyn(&self) -> MessageDescriptor {
        self.descriptor_by_instance()
    }

    fn merge_from_dyn(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()> {
        self.merge_from(is)
    }

    fn write_to_with_cached_sizes_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.write_to_with_cached_sizes(os)
    }

    fn compute_size_dyn(&self) -> u32 {
        self.compute_size()
    }

    fn is_initialized_dyn(&self) -> bool {
        self.is_initialized()
    }

    fn get_unknown_fields_dyn(&self) -> &UnknownFields {
        self.get_unknown_fields()
    }

    fn mut_unknown_fields_dyn(&mut self) -> &mut UnknownFields {
        self.mut_unknown_fields()
    }
}

impl dyn MessageDyn {
    /// Check if all required fields of this object are initialized.
    pub fn check_initialized_dyn(&self) -> ProtobufResult<()> {
        if !self.is_initialized_dyn() {
            Err(ProtobufError::MessageNotInitialized(
                self.descriptor_dyn().name().to_owned(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    pub fn write_to_writer_dyn(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to_dyn(os))
    }

    /// Write the message to bytes vec.
    pub fn write_to_vec_dyn(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to_dyn(os))
    }

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    pub fn write_to_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized_dyn()?;

        // cache sizes
        self.compute_size_dyn();
        // TODO: reserve additional
        self.write_to_with_cached_sizes_dyn(os)?;

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_vec_dyn(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to_dyn(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    pub fn merge_from_bytes_dyn(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from_dyn(&mut is)
    }

    /// Write the message to bytes vec.
    ///
    /// > **Note**: You can use [`Message::parse_from_bytes`]
    /// to do the reverse.
    pub fn write_to_bytes_dyn(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized_dyn()?;

        let size = self.compute_size_dyn() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes_dyn(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size_dyn();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes_dyn(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_writer_dyn(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to_dyn(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_bytes_dyn(&self) -> ProtobufResult<Vec<u8>> {
        let mut v = Vec::new();
        v.with_coded_output_stream(|os| self.write_length_delimited_to_dyn(os))?;
        Ok(v)
    }

    /// Downcast `Box<dyn Message>` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: Box<dyn Message>) {
    /// let m: Box<dyn Message> = message;
    /// let m: Box<MyMessage> = Message::downcast_box(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_box<T: Any>(self: Box<dyn MessageDyn>) -> Result<Box<T>, Box<dyn MessageDyn>> {
        if Any::type_id(&*self) == TypeId::of::<T>() {
            unsafe {
                let raw: *mut dyn MessageDyn = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }

    /// Downcast `&dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &dyn Message) {
    /// let m: &dyn Message = message;
    /// let m: &MyMessage = Message::downcast_ref(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_ref<'a, M: Message + 'a>(&'a self) -> Option<&'a M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&*(self as *const dyn MessageDyn as *const M)) }
        } else {
            None
        }
    }

    /// Downcast `&mut dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &mut dyn Message) {
    /// let m: &mut dyn Message = message;
    /// let m: &mut MyMessage = Message::downcast_mut(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_mut<'a, M: Message + 'a>(&'a mut self) -> Option<&'a mut M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&mut *(self as *mut dyn MessageDyn as *mut M)) }
        } else {
            None
        }
    }

    /// Clone from a `dyn Message` reference.
    pub fn clone_box(&self) -> Box<dyn MessageDyn> {
        self.descriptor_dyn().clone_message(self)
    }

    /// Reflectively compare the messages.
    pub fn reflect_eq_dyn(&self, other: &dyn MessageDyn, mode: &ReflectEqMode) -> bool {
        MessageDescriptor::reflect_eq_maybe_unrelated(self, other, mode)
    }
}

impl Clone for Box<dyn MessageDyn> {
    fn clone(&self) -> Self {
        (*self).clone_box()
    }
}

impl PartialEq for Box<dyn MessageDyn> {
    fn eq(&self, other: &Box<dyn MessageDyn>) -> bool {
        MessageDescriptor::reflect_eq_maybe_unrelated(&**self, &**other, &ReflectEqMode::default())
    }
}
<commit_msg>Fix MessageDyn::downcast*<commit_after>use crate::reflect::MessageDescriptor;
use crate::reflect::ReflectEqMode;
use crate::stream::WithCodedOutputStream;
use crate::CodedInputStream;
use crate::CodedOutputStream;
use crate::Message;
use crate::ProtobufError;
use crate::ProtobufResult;
use crate::UnknownFields;
use std::any::Any;
use std::any::TypeId;
use std::fmt;
use std::io::Write;

/// Dynamic-dispatch version of [`Message`].
pub trait MessageDyn: Any + fmt::Debug + Send + Sync + 'static {
    /// Message descriptor for this message, used for reflection.
    fn descriptor_dyn(&self) -> MessageDescriptor;

    /// Update this message fields with contents of given stream.
    fn merge_from_dyn(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()>;

    /// Write the message.
    fn write_to_with_cached_sizes_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()>;

    /// Compute (and cache) the message size.
    fn compute_size_dyn(&self) -> u32;

    /// True iff all required fields are initialized.
    /// Always returns `true` for protobuf 3.
    fn is_initialized_dyn(&self) -> bool;

    /// Get a reference to unknown fields.
    fn get_unknown_fields_dyn(&self) -> &UnknownFields;
    /// Get a mutable reference to unknown fields.
    fn mut_unknown_fields_dyn(&mut self) -> &mut UnknownFields;
}

impl<M: Message> MessageDyn for M {
    fn descriptor_dyn(&self) -> MessageDescriptor {
        self.descriptor_by_instance()
    }

    fn merge_from_dyn(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()> {
        self.merge_from(is)
    }

    fn write_to_with_cached_sizes_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.write_to_with_cached_sizes(os)
    }

    fn compute_size_dyn(&self) -> u32 {
        self.compute_size()
    }

    fn is_initialized_dyn(&self) -> bool {
        self.is_initialized()
    }

    fn get_unknown_fields_dyn(&self) -> &UnknownFields {
        self.get_unknown_fields()
    }

    fn mut_unknown_fields_dyn(&mut self) -> &mut UnknownFields {
        self.mut_unknown_fields()
    }
}

impl dyn MessageDyn {
    /// Check if all required fields of this object are initialized.
    pub fn check_initialized_dyn(&self) -> ProtobufResult<()> {
        if !self.is_initialized_dyn() {
            Err(ProtobufError::MessageNotInitialized(
                self.descriptor_dyn().name().to_owned(),
            ))
        } else {
            Ok(())
        }
    }

    /// Write the message to the writer.
    pub fn write_to_writer_dyn(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_to_dyn(os))
    }

    /// Write the message to bytes vec.
    pub fn write_to_vec_dyn(&self, v: &mut Vec<u8>) -> ProtobufResult<()> {
        v.with_coded_output_stream(|os| self.write_to_dyn(os))
    }

    /// Write the message to the stream.
    ///
    /// Results in error if message is not fully initialized.
    pub fn write_to_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        self.check_initialized_dyn()?;

        // cache sizes
        self.compute_size_dyn();
        // TODO: reserve additional
        self.write_to_with_cached_sizes_dyn(os)?;

        Ok(())
    }

    /// Write the message to the vec, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_vec_dyn(&self, vec: &mut Vec<u8>) -> ProtobufResult<()> {
        let mut os = CodedOutputStream::vec(vec);
        self.write_length_delimited_to_dyn(&mut os)?;
        os.flush()?;
        Ok(())
    }

    /// Update this message object with fields read from given stream.
    pub fn merge_from_bytes_dyn(&mut self, bytes: &[u8]) -> ProtobufResult<()> {
        let mut is = CodedInputStream::from_bytes(bytes);
        self.merge_from_dyn(&mut is)
    }

    /// Write the message to bytes vec.
    ///
    /// > **Note**: You can use [`Message::parse_from_bytes`]
    /// to do the reverse.
    pub fn write_to_bytes_dyn(&self) -> ProtobufResult<Vec<u8>> {
        self.check_initialized_dyn()?;

        let size = self.compute_size_dyn() as usize;
        let mut v = Vec::with_capacity(size);
        // skip zerofill
        unsafe {
            v.set_len(size);
        }
        {
            let mut os = CodedOutputStream::bytes(&mut v);
            self.write_to_with_cached_sizes_dyn(&mut os)?;
            os.check_eof();
        }
        Ok(v)
    }

    /// Write the message to the stream prepending the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_dyn(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
        let size = self.compute_size_dyn();
        os.write_raw_varint32(size)?;
        self.write_to_with_cached_sizes_dyn(os)?;

        // TODO: assert we've written same number of bytes as computed

        Ok(())
    }

    /// Write the message to the writer, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_writer_dyn(&self, w: &mut dyn Write) -> ProtobufResult<()> {
        w.with_coded_output_stream(|os| self.write_length_delimited_to_dyn(os))
    }

    /// Write the message to the bytes vec, prepend the message with message length
    /// encoded as varint.
    pub fn write_length_delimited_to_bytes_dyn(&self) -> ProtobufResult<Vec<u8>> {
        let mut v = Vec::new();
        v.with_coded_output_stream(|os| self.write_length_delimited_to_dyn(os))?;
        Ok(v)
    }

    /// Downcast `Box<dyn Message>` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: Box<dyn Message>) {
    /// let m: Box<dyn Message> = message;
    /// let m: Box<MyMessage> = Message::downcast_box(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_box<T: Any>(self: Box<dyn MessageDyn>) -> Result<Box<T>, Box<dyn MessageDyn>> {
        if Any::type_id(&*self) == TypeId::of::<T>() {
            unsafe {
                let raw: *mut dyn MessageDyn = Box::into_raw(self);
                Ok(Box::from_raw(raw as *mut T))
            }
        } else {
            Err(self)
        }
    }

    /// Downcast `&dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &dyn Message) {
    /// let m: &dyn Message = message;
    /// let m: &MyMessage = Message::downcast_ref(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_ref<'a, M: Message + 'a>(&'a self) -> Option<&'a M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&*(self as *const dyn MessageDyn as *const M)) }
        } else {
            None
        }
    }

    /// Downcast `&mut dyn Message` to specific message type.
    ///
    /// ```
    /// # use protobuf::Message;
    /// # fn foo<MyMessage: Message>(message: &mut dyn Message) {
    /// let m: &mut dyn Message = message;
    /// let m: &mut MyMessage = Message::downcast_mut(m).unwrap();
    /// # }
    /// ```
    pub fn downcast_mut<'a, M: Message + 'a>(&'a mut self) -> Option<&'a mut M> {
        if Any::type_id(&*self) == TypeId::of::<M>() {
            unsafe { Some(&mut *(self as *mut dyn MessageDyn as *mut M)) }
        } else {
            None
        }
    }

    /// Clone from a `dyn Message` reference.
    pub fn clone_box(&self) -> Box<dyn MessageDyn> {
        self.descriptor_dyn().clone_message(self)
    }

    /// Reflectively compare the messages.
    pub fn reflect_eq_dyn(&self, other: &dyn MessageDyn, mode: &ReflectEqMode) -> bool {
        MessageDescriptor::reflect_eq_maybe_unrelated(self, other, mode)
    }
}

impl Clone for Box<dyn MessageDyn> {
    fn clone(&self) -> Self {
        (*self).clone_box()
    }
}

impl PartialEq for Box<dyn MessageDyn> {
    fn eq(&self, other: &Box<dyn MessageDyn>) -> bool {
        MessageDescriptor::reflect_eq_maybe_unrelated(&**self, &**other, &ReflectEqMode::default())
    }
}

#[cfg(test)]
mod test {
    use crate::descriptor::FileDescriptorProto;
    use crate::message::Message;
    use crate::MessageDyn;

    #[test]
    fn downcast_ref() {
        let mut m = FileDescriptorProto::new();
        let mut d = &m as &dyn MessageDyn;
        let mut c: &FileDescriptorProto = d.downcast_ref().unwrap();
        assert_eq!(
            c as *const FileDescriptorProto,
            &m as *const FileDescriptorProto
        );
    }

    #[test]
    fn downcast_mut() {
        let mut m = FileDescriptorProto::new();
        let mut d = &mut m as &mut dyn MessageDyn;
        let mut c: &mut FileDescriptorProto = d.downcast_mut().unwrap();
        assert_eq!(
            c as *const FileDescriptorProto,
            &m as *const FileDescriptorProto
        );
    }

    #[test]
    fn downcast_box() {
        let mut m = FileDescriptorProto::new();
        let mut d: Box<dyn MessageDyn> = Box::new(m);
        let mut _c: Box<FileDescriptorProto> = d.downcast_box().unwrap();
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The vulkano developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::sync::Arc;
use std::ops::Range;
use std::ptr;
use smallvec::SmallVec;

use command_buffer::RawCommandBufferPrototype;
use command_buffer::CommandsList;
use command_buffer::CommandsListSink;
use device::Device;
use format::ClearValue;
use framebuffer::traits::TrackedFramebuffer;
use framebuffer::RenderPass;
use framebuffer::RenderPassClearValues;
use VulkanObject;
use VulkanPointers;
use vk;

/// Wraps around a commands list and adds to the end of it a command that enters a render pass.
pub struct CmdBeginRenderPass<L, Rp, F> where L: CommandsList {
    // Parent commands list.
    previous: L,
    // True if only secondary command buffers can be added.
    secondary: bool,
    // The draw area.
    rect: [Range<u32>; 2],
    // The clear values for the clear attachments.
    clear_values: SmallVec<[vk::ClearValue; 6]>,
    // The raw render pass handle to bind.
    raw_render_pass: vk::RenderPass,
    // The raw framebuffer handle to bind.
    raw_framebuffer: vk::Framebuffer,
    // The device.
    device: Arc<Device>,
    // The render pass. Can be `None` if same as framebuffer.
    render_pass: Option<Rp>,
    // The framebuffer.
    framebuffer: F,
}

impl<L, F> CmdBeginRenderPass<L, F::RenderPass, F>
    where L: CommandsList, F: TrackedFramebuffer
{
    /// See the documentation of the `begin_render_pass` method.
    // TODO: allow setting more parameters
    pub fn new<C>(previous: L, framebuffer: F, secondary: bool, clear_values: C)
                  -> CmdBeginRenderPass<L, F::RenderPass, F>
        where F::RenderPass: RenderPassClearValues<C>
    {
        let raw_render_pass = framebuffer.render_pass().inner().internal_object();
        let device = framebuffer.render_pass().inner().device().clone();
        let raw_framebuffer = framebuffer.internal_object();

        let clear_values = framebuffer.render_pass().convert_clear_values(clear_values)
                                      .map(|clear_value|
        {
            match clear_value {
                ClearValue::None => {
                    vk::ClearValue::color(vk::ClearColorValue::float32([0.0; 4]))
                },
                ClearValue::Float(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::float32(val))
                },
                ClearValue::Int(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::int32(val))
                },
                ClearValue::Uint(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::uint32(val))
                },
                ClearValue::Depth(val) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: val, stencil: 0
                    })
                },
                ClearValue::Stencil(val) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: 0.0, stencil: val
                    })
                },
                ClearValue::DepthStencil((depth, stencil)) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: depth, stencil: stencil,
                    })
                },
            }
        }).collect();

        let rect = [0 .. framebuffer.dimensions()[0], 0 .. framebuffer.dimensions()[1]];

        CmdBeginRenderPass {
            previous: previous,
            secondary: secondary,
            rect: rect,
            clear_values: clear_values,
            raw_render_pass: raw_render_pass,
            raw_framebuffer: raw_framebuffer,
            device: device,
            render_pass: None,
            framebuffer: framebuffer,
        }
    }
}

unsafe impl<L, Rp, F> CommandsList for CmdBeginRenderPass<L, Rp, F> where L: CommandsList {
    #[inline]
    fn append<'a>(&'a self, builder: &mut CommandsListSink<'a>) {
        self.previous.append(builder);

        assert_eq!(self.device.internal_object(), builder.device().internal_object());

        debug_assert!(self.rect[0].start <= self.rect[0].end);
        debug_assert!(self.rect[1].start <= self.rect[1].end);

        builder.add_command(Box::new(move |raw: &mut RawCommandBufferPrototype| {
            unsafe {
                let vk = raw.device.pointers();
                let cmd = raw.command_buffer.clone().take().unwrap();

                let begin = vk::RenderPassBeginInfo {
                    sType: vk::STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
                    pNext: ptr::null(),
                    renderPass: self.raw_render_pass,
                    framebuffer: self.raw_framebuffer,
                    renderArea: vk::Rect2D {
                        offset: vk::Offset2D {
                            x: self.rect[0].start as i32,
                            y: self.rect[1].start as i32,
                        },
                        extent: vk::Extent2D {
                            width: self.rect[0].end - self.rect[0].start,
                            height: self.rect[1].end - self.rect[1].start,
                        },
                    },
                    clearValueCount: self.clear_values.len() as u32,
                    pClearValues: self.clear_values.as_ptr(),
                };

                let contents = if self.secondary { vk::SUBPASS_CONTENTS_SECONDARY_COMMAND_BUFFERS }
                               else { vk::SUBPASS_CONTENTS_INLINE };

                vk.CmdBeginRenderPass(cmd, &begin, contents);
            }
        }));
    }
}
<commit_msg>Perform transitions in begin render pass<commit_after>// Copyright (c) 2016 The vulkano developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::sync::Arc;
use std::ops::Range;
use std::ptr;
use smallvec::SmallVec;

use command_buffer::RawCommandBufferPrototype;
use command_buffer::CommandsList;
use command_buffer::CommandsListSink;
use device::Device;
use format::ClearValue;
use framebuffer::traits::TrackedFramebuffer;
use framebuffer::RenderPass;
use framebuffer::RenderPassClearValues;
use VulkanObject;
use VulkanPointers;
use vk;

/// Wraps around a commands list and adds to the end of it a command that enters a render pass.
pub struct CmdBeginRenderPass<L, Rp, F> where L: CommandsList {
    // Parent commands list.
    previous: L,
    // True if only secondary command buffers can be added.
    secondary: bool,
    // The draw area.
    rect: [Range<u32>; 2],
    // The clear values for the clear attachments.
    clear_values: SmallVec<[vk::ClearValue; 6]>,
    // The raw render pass handle to bind.
    raw_render_pass: vk::RenderPass,
    // The raw framebuffer handle to bind.
    raw_framebuffer: vk::Framebuffer,
    // The device.
    device: Arc<Device>,
    // The render pass. Can be `None` if same as framebuffer.
    render_pass: Option<Rp>,
    // The framebuffer.
    framebuffer: F,
}

impl<L, F> CmdBeginRenderPass<L, F::RenderPass, F>
    where L: CommandsList, F: TrackedFramebuffer
{
    /// See the documentation of the `begin_render_pass` method.
    // TODO: allow setting more parameters
    pub fn new<C>(previous: L, framebuffer: F, secondary: bool, clear_values: C)
                  -> CmdBeginRenderPass<L, F::RenderPass, F>
        where F::RenderPass: RenderPassClearValues<C>
    {
        let raw_render_pass = framebuffer.render_pass().inner().internal_object();
        let device = framebuffer.render_pass().inner().device().clone();
        let raw_framebuffer = framebuffer.internal_object();

        let clear_values = framebuffer.render_pass().convert_clear_values(clear_values)
                                      .map(|clear_value|
        {
            match clear_value {
                ClearValue::None => {
                    vk::ClearValue::color(vk::ClearColorValue::float32([0.0; 4]))
                },
                ClearValue::Float(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::float32(val))
                },
                ClearValue::Int(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::int32(val))
                },
                ClearValue::Uint(val) => {
                    vk::ClearValue::color(vk::ClearColorValue::uint32(val))
                },
                ClearValue::Depth(val) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: val, stencil: 0
                    })
                },
                ClearValue::Stencil(val) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: 0.0, stencil: val
                    })
                },
                ClearValue::DepthStencil((depth, stencil)) => {
                    vk::ClearValue::depth_stencil(vk::ClearDepthStencilValue {
                        depth: depth, stencil: stencil,
                    })
                },
            }
        }).collect();

        let rect = [0 .. framebuffer.dimensions()[0], 0 .. framebuffer.dimensions()[1]];

        CmdBeginRenderPass {
            previous: previous,
            secondary: secondary,
            rect: rect,
            clear_values: clear_values,
            raw_render_pass: raw_render_pass,
            raw_framebuffer: raw_framebuffer,
            device: device,
            render_pass: None,
            framebuffer: framebuffer,
        }
    }
}

unsafe impl<L, Rp, F> CommandsList for CmdBeginRenderPass<L, Rp, F>
    where L: CommandsList, F: TrackedFramebuffer
{
    #[inline]
    fn append<'a>(&'a self, builder: &mut CommandsListSink<'a>) {
        self.previous.append(builder);

        assert_eq!(self.device.internal_object(), builder.device().internal_object());

        debug_assert!(self.rect[0].start <= self.rect[0].end);
        debug_assert!(self.rect[1].start <= self.rect[1].end);

        self.framebuffer.add_transition(builder);

        builder.add_command(Box::new(move |raw: &mut RawCommandBufferPrototype| {
            unsafe {
                let vk = raw.device.pointers();
                let cmd = raw.command_buffer.clone().take().unwrap();

                let begin = vk::RenderPassBeginInfo {
                    sType: vk::STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
                    pNext: ptr::null(),
                    renderPass: self.raw_render_pass,
                    framebuffer: self.raw_framebuffer,
                    renderArea: vk::Rect2D {
                        offset: vk::Offset2D {
                            x: self.rect[0].start as i32,
                            y: self.rect[1].start as i32,
                        },
                        extent: vk::Extent2D {
                            width: self.rect[0].end - self.rect[0].start,
                            height: self.rect[1].end - self.rect[1].start,
                        },
                    },
                    clearValueCount: self.clear_values.len() as u32,
                    pClearValues: self.clear_values.as_ptr(),
                };

                let contents = if self.secondary { vk::SUBPASS_CONTENTS_SECONDARY_COMMAND_BUFFERS }
                               else { vk::SUBPASS_CONTENTS_INLINE };

                vk.CmdBeginRenderPass(cmd, &begin, contents);
            }
        }));
    }
}
<|endoftext|>


<commit_before>mod inputs;

use super::PartialSeries;
use crate::err::{Error, Result};
use crate::series::info::{InfoSelector, SeriesInfo};
use crate::series::{SeriesParams, SeriesPath};
use crate::tui::component::{Component, Draw};
use crate::tui::{UIBackend, UIState};
use crate::{try_opt_r, try_opt_ret};
use anime::local::Episodes;
use inputs::{InputSet, ParsedValue, ValidatedInput};
use std::borrow::Cow;
use std::mem;
use std::time::Instant;
use termion::event::Key;
use tui::backend::Backend;
use tui::layout::{Alignment, Constraint, Direction, Layout, Rect};
use tui::style::{Color, Modifier, Style};
use tui::terminal::Frame;
use tui::widgets::{Block, Borders, Paragraph, Text};

const SECS_BETWEEN_SERIES_UPDATES: f32 = 0.25;

pub struct AddSeriesPanel {
    inputs: InputSet,
    selected_input: usize,
    error: Option<Cow<'static, str>>,
    last_update: Option<Instant>,
    series_builder: SeriesBuilder,
}

impl AddSeriesPanel {
    pub fn new() -> Self {
        Self {
            inputs: InputSet::new(),
            selected_input: 0,
            error: None,
            last_update: None,
            series_builder: SeriesBuilder::new(),
        }
    }

    #[inline(always)]
    fn current_input(&mut self) -> &mut dyn ValidatedInput {
        self.inputs.index_mut(self.selected_input)
    }

    fn validate_selected(&mut self) {
        self.current_input().validate();

        for input in &self.inputs.all_mut() {
            if let value @ Some(_) = input.error() {
                self.error = value;
                return;
            }
        }

        self.error = None;
    }

    fn draw_add_series_panel<B>(&mut self, rect: Rect, frame: &mut Frame<B>)
    where
        B: Backend,
    {
        let vert_fields = Layout::default()
            .direction(Direction::Vertical)
            .constraints(
                [
                    // Field
                    Constraint::Length(4),
                    // Spacer
                    Constraint::Length(1),
                    // Field
                    Constraint::Length(4),
                    // Remaining space
                    Constraint::Percentage(100),
                ]
                .as_ref(),
            )
            .vertical_margin(2)
            .split(rect);

        let horiz_fields = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref());

        let horiz_fields_top = horiz_fields.clone().split(vert_fields[0]);
        let horiz_fields_bottom = horiz_fields.split(vert_fields[2]);

        let field_positions = horiz_fields_top
            .into_iter()
            .chain(horiz_fields_bottom.into_iter());

        for (input, pos) in self.inputs.all_mut().iter_mut().zip(field_positions) {
            let layout = Layout::default()
                .direction(Direction::Vertical)
                .constraints([Constraint::Length(1), Constraint::Length(3)].as_ref())
                .horizontal_margin(3)
                .split(pos);

            let text = [bold_header(input.label())];
            let widget = Paragraph::new(text.iter())
                .wrap(false)
                .alignment(Alignment::Center);
            frame.render_widget(widget, layout[0]);

            input.input_mut().draw(&(), layout[1], frame);
        }
    }

    fn draw_detected_panel<B>(&mut self, rect: Rect, frame: &mut Frame<B>)
    where
        B: Backend,
    {
        macro_rules! info_label {
            ($label:expr, $value:expr, $rect:expr) => {{
                let label = concat!($label, "\n");
                let text = [
                    bold_header(label),
                    Text::styled($value, Style::default().modifier(Modifier::ITALIC)),
                ];

                let widget = Paragraph::new(text.iter())
                    .wrap(false)
                    .alignment(Alignment::Center);
                frame.render_widget(widget, $rect);
            }};
        }

        let (header_text, has_error) = match (&self.error, &self.series_builder.params) {
            (Some(err), Some(_)) | (Some(err), None) => {
                ([bold_header_with(err.as_ref(), |s| s.fg(Color::Red))], true)
            }
            (None, Some(_)) => ([bold_header("Detected")], false),
            (None, None) => (
                [bold_header_with("Nothing Detected", |s| s.fg(Color::Red))],
                false,
            ),
        };

        let header = Paragraph::new(header_text.iter())
            .wrap(true)
            .alignment(Alignment::Center);
        frame.render_widget(header, rect);

        let fields = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())
            .vertical_margin(2)
            .split(rect);

        if has_error {
            return;
        }

        let (params, episodes) = try_opt_ret!(&self.series_builder.params);

        info_label!(
            "Relative Path",
            format!("{}", params.path.display()),
            fields[0]
        );
        info_label!("Episodes", episodes.len().to_string(), fields[1]);
    }
}

impl Component for AddSeriesPanel {
    type State = UIState;
    type KeyResult = AddSeriesResult;

    fn tick(&mut self, state: &mut UIState) -> Result<()> {
        let last_update = try_opt_r!(self.last_update);

        if last_update.elapsed().as_secs_f32() < SECS_BETWEEN_SERIES_UPDATES {
            return Ok(());
        }

        self.series_builder.update(&self.inputs, state).ok();
        self.last_update = None;

        Ok(())
    }

    fn process_key(&mut self, key: Key, state: &mut Self::State) -> Self::KeyResult {
        match key {
            Key::Esc => AddSeriesResult::Reset,
            Key::Char('\n') => {
                self.validate_selected();

                if self.error.is_some() {
                    return AddSeriesResult::Ok;
                }

                match self.series_builder.build(&self.inputs, state) {
                    Ok(partial) => AddSeriesResult::AddSeries(Box::new(partial)),
                    Err(err) => AddSeriesResult::Error(err),
                }
            }
            Key::Char('\t') => {
                self.validate_selected();

                self.current_input().input_mut().selected = false;
                self.selected_input = (self.selected_input + 1) % self.inputs.len();
                self.current_input().input_mut().selected = true;

                AddSeriesResult::Ok
            }
            key => {
                self.current_input().input_mut().process_key(key);
                self.validate_selected();

                self.last_update = Some(Instant::now());
                AddSeriesResult::Ok
            }
        }
    }
}

impl<B> Draw<B> for AddSeriesPanel
where
    B: Backend,
{
    type State = ();

    fn draw(&mut self, _: &Self::State, rect: Rect, frame: &mut Frame<B>) {
        let split = Layout::default()
            .direction(Direction::Vertical)
            .constraints([Constraint::Min(7), Constraint::Length(6)].as_ref())
            .horizontal_margin(2)
            .split(rect);

        let outline = Block::default().title("Add Series").borders(Borders::ALL);
        frame.render_widget(outline, rect);

        self.draw_add_series_panel(split[0], frame);
        self.draw_detected_panel(split[1], frame);
    }

    fn after_draw(&mut self, backend: &mut UIBackend<B>, _: &Self::State) {
        for input in &mut self.inputs.all_mut() {
            input.input_mut().after_draw(backend, &());
        }
    }
}

pub enum AddSeriesResult {
    Ok,
    Reset,
    AddSeries(Box<PartialSeries>),
    Error(Error),
}

struct SeriesBuilder {
    params: Option<(SeriesParams, Episodes)>,
}

impl SeriesBuilder {
    fn new() -> Self {
        Self { params: None }
    }

    fn path(&self, inputs: &InputSet, state: &UIState) -> Result<SeriesPath> {
        match &inputs.path.parsed_value() {
            Some(path) => Ok(SeriesPath::new(path, &state.config)),
            None => SeriesPath::closest_matching(inputs.name.parsed_value(), &state.config),
        }
    }

    fn update(&mut self, inputs: &InputSet, state: &UIState) -> Result<()> {
        match self.update_internal(inputs, state) {
            ok @ Ok(_) => ok,
            err @ Err(_) => {
                self.params = None;
                err
            }
        }
    }

    fn update_internal(&mut self, inputs: &InputSet, state: &UIState) -> Result<()> {
        let path = self.path(inputs, state)?;

        let parser = inputs.parser.parsed_value();
        let episodes = Episodes::parse(path.absolute(&state.config), parser)?;
        let name = inputs.name.parsed_value();

        match &mut self.params {
            Some((params, cur_episodes)) => {
                params.update(name, path, parser);
                *cur_episodes = episodes;
                Ok(())
            }
            None => {
                let params = SeriesParams::new(name, path, parser.clone());
                self.params = Some((params, episodes));
                Ok(())
            }
        }
    }

    fn build(&mut self, inputs: &InputSet, state: &UIState) -> Result<PartialSeries> {
        let (params, episodes) = match self.update(inputs, state) {
            Ok(_) => mem::take(&mut self.params).unwrap(),
            Err(err) => return Err(err),
        };

        let info = {
            let id = inputs.id.parsed_value();
            let sel = id.map_or_else(
                || InfoSelector::from_path_or_name(&params.path, &params.name),
                InfoSelector::ID,
            );

            SeriesInfo::from_remote(sel, state.remote.as_ref())?
        };

        let params = SeriesParams::new(params.name, params.path, params.parser);

        Ok(PartialSeries::new(info, params, episodes))
    }
}

#[inline(always)]
fn bold_header_with<F>(header: &str, extra_style: F) -> Text
where
    F: FnOnce(Style) -> Style,
{
    Text::styled(
        header,
        extra_style(Style::default().modifier(Modifier::BOLD)),
    )
}

#[inline(always)]
fn bold_header(header: &str) -> Text {
    bold_header_with(header, |s| s)
}
<commit_msg>anup: Display detected episode ranges in add series panel<commit_after>mod inputs;

use super::PartialSeries;
use crate::err::{Error, Result};
use crate::series::info::{InfoSelector, SeriesInfo};
use crate::series::{SeriesParams, SeriesPath};
use crate::tui::component::{Component, Draw};
use crate::tui::{UIBackend, UIState};
use crate::{try_opt_r, try_opt_ret};
use anime::local::Episodes;
use inputs::{InputSet, ParsedValue, ValidatedInput};
use std::borrow::Cow;
use std::mem;
use std::time::Instant;
use termion::event::Key;
use tui::backend::Backend;
use tui::layout::{Alignment, Constraint, Direction, Layout, Rect};
use tui::style::{Color, Modifier, Style};
use tui::terminal::Frame;
use tui::widgets::{Block, Borders, Paragraph, Text};

const SECS_BETWEEN_SERIES_UPDATES: f32 = 0.25;

pub struct AddSeriesPanel {
    inputs: InputSet,
    selected_input: usize,
    error: Option<Cow<'static, str>>,
    last_update: Option<Instant>,
    series_builder: SeriesBuilder,
}

impl AddSeriesPanel {
    pub fn new() -> Self {
        Self {
            inputs: InputSet::new(),
            selected_input: 0,
            error: None,
            last_update: None,
            series_builder: SeriesBuilder::new(),
        }
    }

    #[inline(always)]
    fn current_input(&mut self) -> &mut dyn ValidatedInput {
        self.inputs.index_mut(self.selected_input)
    }

    fn validate_selected(&mut self) {
        self.current_input().validate();

        for input in &self.inputs.all_mut() {
            if let value @ Some(_) = input.error() {
                self.error = value;
                return;
            }
        }

        self.error = None;
    }

    fn draw_add_series_panel<B>(&mut self, rect: Rect, frame: &mut Frame<B>)
    where
        B: Backend,
    {
        let vert_fields = Layout::default()
            .direction(Direction::Vertical)
            .constraints(
                [
                    // Field
                    Constraint::Length(4),
                    // Spacer
                    Constraint::Length(1),
                    // Field
                    Constraint::Length(4),
                    // Remaining space
                    Constraint::Percentage(100),
                ]
                .as_ref(),
            )
            .vertical_margin(2)
            .split(rect);

        let horiz_fields = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref());

        let horiz_fields_top = horiz_fields.clone().split(vert_fields[0]);
        let horiz_fields_bottom = horiz_fields.split(vert_fields[2]);

        let field_positions = horiz_fields_top
            .into_iter()
            .chain(horiz_fields_bottom.into_iter());

        for (input, pos) in self.inputs.all_mut().iter_mut().zip(field_positions) {
            let layout = Layout::default()
                .direction(Direction::Vertical)
                .constraints([Constraint::Length(1), Constraint::Length(3)].as_ref())
                .horizontal_margin(3)
                .split(pos);

            let text = [bold_header(input.label())];
            let widget = Paragraph::new(text.iter())
                .wrap(false)
                .alignment(Alignment::Center);
            frame.render_widget(widget, layout[0]);

            input.input_mut().draw(&(), layout[1], frame);
        }
    }

    fn draw_detected_panel<B>(&mut self, rect: Rect, frame: &mut Frame<B>)
    where
        B: Backend,
    {
        macro_rules! info_label {
            ($label:expr, $value:expr, $rect:expr) => {{
                let label = concat!($label, "\n");
                let text = [bold_header(label), $value];

                let widget = Paragraph::new(text.iter())
                    .wrap(false)
                    .alignment(Alignment::Center);
                frame.render_widget(widget, $rect);
            }};
        }

        let (header_text, has_error) = match (&self.error, &self.series_builder.params) {
            (Some(err), Some(_)) | (Some(err), None) => {
                ([bold_header_with(err.as_ref(), |s| s.fg(Color::Red))], true)
            }
            (None, Some(_)) => ([bold_header("Detected")], false),
            (None, None) => (
                [bold_header_with("Nothing Detected", |s| s.fg(Color::Red))],
                false,
            ),
        };

        let header = Paragraph::new(header_text.iter())
            .wrap(true)
            .alignment(Alignment::Center);
        frame.render_widget(header, rect);

        let fields = Layout::default()
            .direction(Direction::Horizontal)
            .constraints([Constraint::Percentage(50), Constraint::Percentage(50)].as_ref())
            .vertical_margin(2)
            .split(rect);

        if has_error {
            return;
        }

        let built = try_opt_ret!(&self.series_builder.params);

        info_label!(
            "Relative Path",
            italic(format!("{}", built.params.path.display())),
            fields[0]
        );

        let episodes_text = match &built.found_episodes {
            Some(text) => italic(text),
            None => italic_with("none", |s| s.fg(Color::Yellow)),
        };

        info_label!("Found Episodes", episodes_text, fields[1]);
    }
}

impl Component for AddSeriesPanel {
    type State = UIState;
    type KeyResult = AddSeriesResult;

    fn tick(&mut self, state: &mut UIState) -> Result<()> {
        let last_update = try_opt_r!(self.last_update);

        if last_update.elapsed().as_secs_f32() < SECS_BETWEEN_SERIES_UPDATES {
            return Ok(());
        }

        self.series_builder.update(&self.inputs, state).ok();
        self.last_update = None;

        Ok(())
    }

    fn process_key(&mut self, key: Key, state: &mut Self::State) -> Self::KeyResult {
        match key {
            Key::Esc => AddSeriesResult::Reset,
            Key::Char('\n') => {
                self.validate_selected();

                if self.error.is_some() {
                    return AddSeriesResult::Ok;
                }

                match self.series_builder.build(&self.inputs, state) {
                    Ok(partial) => AddSeriesResult::AddSeries(Box::new(partial)),
                    Err(err) => AddSeriesResult::Error(err),
                }
            }
            Key::Char('\t') => {
                self.validate_selected();

                self.current_input().input_mut().selected = false;
                self.selected_input = (self.selected_input + 1) % self.inputs.len();
                self.current_input().input_mut().selected = true;

                AddSeriesResult::Ok
            }
            key => {
                self.current_input().input_mut().process_key(key);
                self.validate_selected();

                self.last_update = Some(Instant::now());
                AddSeriesResult::Ok
            }
        }
    }
}

impl<B> Draw<B> for AddSeriesPanel
where
    B: Backend,
{
    type State = ();

    fn draw(&mut self, _: &Self::State, rect: Rect, frame: &mut Frame<B>) {
        let split = Layout::default()
            .direction(Direction::Vertical)
            .constraints([Constraint::Min(7), Constraint::Length(6)].as_ref())
            .horizontal_margin(2)
            .split(rect);

        let outline = Block::default().title("Add Series").borders(Borders::ALL);
        frame.render_widget(outline, rect);

        self.draw_add_series_panel(split[0], frame);
        self.draw_detected_panel(split[1], frame);
    }

    fn after_draw(&mut self, backend: &mut UIBackend<B>, _: &Self::State) {
        for input in &mut self.inputs.all_mut() {
            input.input_mut().after_draw(backend, &());
        }
    }
}

pub enum AddSeriesResult {
    Ok,
    Reset,
    AddSeries(Box<PartialSeries>),
    Error(Error),
}

struct SeriesBuilder {
    params: Option<BuiltSeriesParams>,
}

impl SeriesBuilder {
    fn new() -> Self {
        Self { params: None }
    }

    fn path(&self, inputs: &InputSet, state: &UIState) -> Result<SeriesPath> {
        match &inputs.path.parsed_value() {
            Some(path) => Ok(SeriesPath::new(path, &state.config)),
            None => SeriesPath::closest_matching(inputs.name.parsed_value(), &state.config),
        }
    }

    fn update(&mut self, inputs: &InputSet, state: &UIState) -> Result<()> {
        match self.update_internal(inputs, state) {
            ok @ Ok(_) => ok,
            err @ Err(_) => {
                self.params = None;
                err
            }
        }
    }

    fn update_internal(&mut self, inputs: &InputSet, state: &UIState) -> Result<()> {
        let path = self.path(inputs, state)?;

        let parser = inputs.parser.parsed_value();
        let episodes = Episodes::parse(path.absolute(&state.config), parser)?;
        let name = inputs.name.parsed_value();

        match &mut self.params {
            Some(built) => {
                built.params.update(name, path, parser);
                built.set_episodes(episodes);
                Ok(())
            }
            None => {
                let params = SeriesParams::new(name, path, parser.clone());
                let built = BuiltSeriesParams::new(params, episodes);
                self.params = Some(built);
                Ok(())
            }
        }
    }

    fn build(&mut self, inputs: &InputSet, state: &UIState) -> Result<PartialSeries> {
        let built = match self.update(inputs, state) {
            Ok(_) => mem::take(&mut self.params).unwrap(),
            Err(err) => return Err(err),
        };

        let params = built.params;

        let info = {
            let id = inputs.id.parsed_value();
            let sel = id.map_or_else(
                || InfoSelector::from_path_or_name(&params.path, &params.name),
                InfoSelector::ID,
            );

            SeriesInfo::from_remote(sel, state.remote.as_ref())?
        };

        let params = SeriesParams::new(params.name, params.path, params.parser);

        Ok(PartialSeries::new(info, params, built.episodes))
    }
}

struct BuiltSeriesParams {
    params: SeriesParams,
    episodes: Episodes,
    found_episodes: Option<String>,
}

impl BuiltSeriesParams {
    fn new(params: SeriesParams, episodes: Episodes) -> Self {
        let found_episodes = Self::episode_range_str(&episodes);

        Self {
            params,
            episodes,
            found_episodes,
        }
    }

    fn set_episodes(&mut self, episodes: Episodes) {
        self.episodes = episodes;
        self.found_episodes = Self::episode_range_str(&self.episodes);
    }

    /// Build a string that displays ranges and holes within a set of episodes.
    /// A hole is considered to be an episode that is not sequential.
    fn episode_range_str(episodes: &Episodes) -> Option<String> {
        use std::ops::Range;

        const RANGE_SEPARATOR: char = '-';
        const HOLE_SEPARATOR: char = '|';

        fn push_range(result: &mut String, range: Range<u32>) {
            result.push_str(&range.start.to_string());

            if range.end != range.start {
                result.push(RANGE_SEPARATOR);
                result.push_str(&range.end.to_string());
            }
        }

        if episodes.is_empty() {
            return None;
        }

        // Having sorted episode numbers allows us to very easily find the next hole
        let episode_nums = {
            let mut nums = episodes.keys().cloned().collect::<Vec<_>>();
            nums.sort_unstable();
            nums
        };

        if episode_nums.len() < 2 {
            return Some(episode_nums[0].to_string());
        }

        let mut result = String::new();
        let mut range = episode_nums[0]..episode_nums[0];

        for &num in &episode_nums[1..] {
            if num - range.end > 1 {
                push_range(&mut result, range);
                result.push(HOLE_SEPARATOR);
                range = num..num;

                continue;
            }

            range.end = num;
        }

        push_range(&mut result, range);
        Some(result)
    }
}

#[inline(always)]
fn bold_header_with<F>(header: &str, extra_style: F) -> Text
where
    F: FnOnce(Style) -> Style,
{
    Text::styled(
        header,
        extra_style(Style::default().modifier(Modifier::BOLD)),
    )
}

#[inline(always)]
fn bold_header(header: &str) -> Text {
    bold_header_with(header, |s| s)
}

#[inline(always)]
fn italic_with<'a, S, F>(text: S, extra_style: F) -> Text<'a>
where
    S: Into<Cow<'a, str>>,
    F: FnOnce(Style) -> Style,
{
    Text::styled(
        text,
        extra_style(Style::default().modifier(Modifier::ITALIC)),
    )
}

#[inline(always)]
fn italic<'a, S>(text: S) -> Text<'a>
where
    S: Into<Cow<'a, str>>,
{
    italic_with(text, |s| s)
}

#[cfg(test)]
mod tests {
    use super::*;
    use anime::local::EpisodeMap;
    use std::ops::RangeInclusive;

    fn insert_range(map: &mut EpisodeMap, range: RangeInclusive<u32>) {
        for i in range {
            map.insert(i, String::new());
        }
    }

    macro_rules! episodes {
        () => {{
            use std::collections::HashMap;
            Episodes::new(HashMap::new())
        }};

        ($($range:expr),+) => {{
            use std::collections::HashMap;

            let mut episodes = HashMap::new();
            $(
            insert_range(&mut episodes, $range);
            )+
            Episodes::new(episodes)
        }};
    }

    #[test]
    fn found_episodes_str_detection() {
        let test_sets = vec![
            (episodes!(), None),
            (episodes!(6..=6), Some("6")),
            (episodes!(1..=6, 12..=16), Some("1-6|12-16")),
            (episodes!(1..=1), Some("1")),
            (episodes!(1..=2), Some("1-2")),
            (episodes!(1..=6, 8..=9), Some("1-6|8-9")),
            (episodes!(2..=2, 6..=12), Some("2|6-12")),
            (episodes!(1..=12, 14..=14), Some("1-12|14")),
            (
                episodes!(1..=12, 16..=24, 32..=48),
                Some("1-12|16-24|32-48"),
            ),
            (episodes!(1..=12, 16..=16, 24..=32), Some("1-12|16|24-32")),
            (episodes!(2..=2, 6..=6, 12..=12), Some("2|6|12")),
        ];

        for (episodes, expected) in test_sets {
            match BuiltSeriesParams::episode_range_str(&episodes) {
                Some(result) => assert_eq!(Some(result.as_str()), expected),
                None => assert_eq!(None, expected),
            }
        }
    }
}
<|endoftext|>