<commit_before>// Internal
use crate::builder::StyledStr;
use crate::builder::{ArgPredicate, Command};
use crate::output::TAB;
use crate::parser::ArgMatcher;
use crate::util::ChildGraph;
use crate::util::FlatSet;
use crate::util::Id;

static DEFAULT_SUB_VALUE_NAME: &str = "COMMAND";

pub(crate) struct Usage<'cmd> {
    cmd: &'cmd Command,
    required: Option<&'cmd ChildGraph<Id>>,
}

impl<'cmd> Usage<'cmd> {
    pub(crate) fn new(cmd: &'cmd Command) -> Self {
        Usage {
            cmd,
            required: None,
        }
    }

    pub(crate) fn required(mut self, required: &'cmd ChildGraph<Id>) -> Self {
        self.required = Some(required);
        self
    }

    // Creates a usage string for display. This happens just after all arguments were parsed, but before
    // any subcommands have been parsed (so as to give subcommands their own usage recursively)
    pub(crate) fn create_usage_with_title(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_usage_with_title");
        let mut styled = StyledStr::new();
        styled.header("Usage:");
        styled.none("\n");
        styled.none(TAB);
        styled.extend(self.create_usage_no_title(used).into_iter());
        styled
    }

    // Creates a usage string (*without title*) if one was not provided by the user manually.
    pub(crate) fn create_usage_no_title(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_usage_no_title");
        if let Some(u) = self.cmd.get_override_usage() {
            u.clone()
        } else if used.is_empty() {
            self.create_help_usage(true)
        } else {
            self.create_smart_usage(used)
        }
    }

    // Creates a usage string for display in help messages (i.e. not for errors)
    fn create_help_usage(&self, incl_reqs: bool) -> StyledStr {
        debug!("Usage::create_help_usage; incl_reqs={:?}", incl_reqs);
        let mut styled = StyledStr::new();
        let name = self
            .cmd
            .get_usage_name()
            .or_else(|| self.cmd.get_bin_name())
            .unwrap_or_else(|| self.cmd.get_name());
        styled.literal(name);

        if self.needs_options_tag() {
            styled.placeholder(" [OPTIONS]");
        }

        self.write_args(&[], !incl_reqs, &mut styled);

        // incl_reqs is only false when this function is called recursively
        if self.cmd.has_visible_subcommands() && incl_reqs
            || self.cmd.is_allow_external_subcommands_set()
        {
            let placeholder = self
                .cmd
                .get_subcommand_value_name()
                .unwrap_or(DEFAULT_SUB_VALUE_NAME);
            if self.cmd.is_subcommand_negates_reqs_set()
                || self.cmd.is_args_conflicts_with_subcommands_set()
            {
                styled.none("\n    ");
                if self.cmd.is_args_conflicts_with_subcommands_set() {
                    // Short-circuit full usage creation since no args will be relevant
                    styled.literal(name);
                } else {
                    styled.extend(self.create_help_usage(false).into_iter());
                }
                styled.placeholder(" <");
                styled.placeholder(placeholder);
                styled.placeholder(">");
            } else if self.cmd.is_subcommand_required_set() {
                styled.placeholder(" <");
                styled.placeholder(placeholder);
                styled.placeholder(">");
            } else {
                styled.placeholder(" [");
                styled.placeholder(placeholder);
                styled.placeholder("]");
            }
        }
        styled.trim();
        debug!("Usage::create_help_usage: usage={}", styled);
        styled
    }

    // Creates a context aware usage string, or "smart usage" from currently used
    // args, and requirements
    fn create_smart_usage(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_smart_usage");
        let mut styled = StyledStr::new();

        styled.literal(
            self.cmd
                .get_usage_name()
                .or_else(|| self.cmd.get_bin_name())
                .unwrap_or_else(|| self.cmd.get_name()),
        );

        self.write_args(used, false, &mut styled);

        if self.cmd.is_subcommand_required_set() {
            styled.placeholder(" <");
            styled.placeholder(
                self.cmd
                    .get_subcommand_value_name()
                    .unwrap_or(DEFAULT_SUB_VALUE_NAME),
            );
            styled.placeholder(">");
        }
        styled
    }

    // Determines if we need the `[OPTIONS]` tag in the usage string
    fn needs_options_tag(&self) -> bool {
        debug!("Usage::needs_options_tag");
        'outer: for f in self.cmd.get_non_positionals() {
            debug!("Usage::needs_options_tag:iter: f={}", f.get_id());

            // Don't print `[OPTIONS]` just for help or version
            if f.get_long() == Some("help") || f.get_long() == Some("version") {
                debug!("Usage::needs_options_tag:iter Option is built-in");
                continue;
            }

            if f.is_hide_set() {
                debug!("Usage::needs_options_tag:iter Option is hidden");
                continue;
            }
            if f.is_required_set() {
                debug!("Usage::needs_options_tag:iter Option is required");
                continue;
            }
            for grp_s in self.cmd.groups_for_arg(&f.id) {
                debug!("Usage::needs_options_tag:iter:iter: grp_s={:?}", grp_s);
                if self.cmd.get_groups().any(|g| g.id == grp_s && g.required) {
                    debug!("Usage::needs_options_tag:iter:iter: Group is required");
                    continue 'outer;
                }
            }

            debug!("Usage::needs_options_tag:iter: [OPTIONS] required");
            return true;
        }

        debug!("Usage::needs_options_tag: [OPTIONS] not required");
        false
    }

    // Returns the required args in usage string form by fully unrolling all groups
    pub(crate) fn write_args(&self, incls: &[Id], force_optional: bool, styled: &mut StyledStr) {
        for required in self.get_args(incls, force_optional) {
            styled.none(" ");
            styled.extend(required.into_iter());
        }
    }

    pub(crate) fn get_args(&self, incls: &[Id], force_optional: bool) -> Vec<StyledStr> {
        debug!("Usage::get_required_usage_from: incls={:?}", incls,);

        let required_owned;
        let required = if let Some(required) = self.required {
            required
        } else {
            required_owned = self.cmd.required_graph();
            &required_owned
        };

        let mut unrolled_reqs = Vec::new();
        for a in required.iter() {
            let is_relevant = |(val, req_arg): &(ArgPredicate, Id)| -> Option<Id> {
                let required = match val {
                    ArgPredicate::Equals(_) => false,
                    ArgPredicate::IsPresent => true,
                };
                required.then(|| req_arg.clone())
            };

            for aa in self.cmd.unroll_arg_requires(is_relevant, a) {
                // if we don't check for duplicates here this causes duplicate error messages
                // see https://github.com/clap-rs/clap/issues/2770
                unrolled_reqs.push(aa);
            }
            // always include the required arg itself. it will not be enumerated
            // by unroll_requirements_for_arg.
            unrolled_reqs.push(a.clone());
        }
        debug!(
            "Usage::get_required_usage_from: unrolled_reqs={:?}",
            unrolled_reqs
        );

        let mut required_groups_members = FlatSet::new();
        let mut required_groups = FlatSet::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if self.cmd.find_group(req).is_some() {
                let group_members = self.cmd.unroll_args_in_group(req);
                let elem = self.cmd.format_group(req);
                required_groups.insert(elem);
                required_groups_members.extend(group_members);
            } else {
                debug_assert!(self.cmd.find(req).is_some());
            }
        }

        let mut required_opts = FlatSet::new();
        let mut required_positionals = Vec::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if let Some(arg) = self.cmd.find(req) {
                if required_groups_members.contains(arg.get_id()) {
                    continue;
                }

                let stylized = arg.stylized(Some(!force_optional));
                if let Some(index) = arg.get_index() {
                    let new_len = index + 1;
                    if required_positionals.len() < new_len {
                        required_positionals.resize(new_len, None);
                    }
                    required_positionals[index] = Some(stylized);
                } else {
                    required_opts.insert(stylized);
                }
            } else {
                debug_assert!(self.cmd.find_group(req).is_some());
            }
        }

        for pos in self.cmd.get_positionals() {
            if pos.is_hide_set() {
                continue;
            }
            if required_groups_members.contains(pos.get_id()) {
                continue;
            }

            let index = pos.get_index().unwrap();
            let new_len = index + 1;
            if required_positionals.len() < new_len {
                required_positionals.resize(new_len, None);
            }
            if required_positionals[index].is_some() {
                if pos.is_last_set() {
                    let styled = required_positionals[index].take().unwrap();
                    let mut new = StyledStr::new();
                    new.literal("-- ");
                    new.extend(styled.into_iter());
                    required_positionals[index] = Some(new);
                }
            } else {
                let mut styled;
                if pos.is_last_set() {
                    styled = StyledStr::new();
                    styled.literal("[-- ");
                    styled.extend(pos.stylized(Some(true)).into_iter());
                    styled.literal("]");
                } else {
                    styled = pos.stylized(Some(false));
                }
                required_positionals[index] = Some(styled);
            }
            if pos.is_last_set() && force_optional {
                required_positionals[index] = None;
            }
        }

        let mut ret_val = Vec::new();
        if !force_optional {
            ret_val.extend(required_opts);
            ret_val.extend(required_groups);
        }
        for pos in required_positionals.into_iter().flatten() {
            ret_val.push(pos);
        }

        debug!("Usage::get_required_usage_from: ret_val={:?}", ret_val);
        ret_val
    }

    pub(crate) fn get_required_usage_from(
        &self,
        incls: &[Id],
        matcher: Option<&ArgMatcher>,
        incl_last: bool,
    ) -> Vec<StyledStr> {
        debug!(
            "Usage::get_required_usage_from: incls={:?}, matcher={:?}, incl_last={:?}",
            incls,
            matcher.is_some(),
            incl_last
        );

        let required_owned;
        let required = if let Some(required) = self.required {
            required
        } else {
            required_owned = self.cmd.required_graph();
            &required_owned
        };

        let mut unrolled_reqs = Vec::new();
        for a in required.iter() {
            let is_relevant = |(val, req_arg): &(ArgPredicate, Id)| -> Option<Id> {
                let required = match val {
                    ArgPredicate::Equals(_) => {
                        if let Some(matcher) = matcher {
                            matcher.check_explicit(a, val)
                        } else {
                            false
                        }
                    }
                    ArgPredicate::IsPresent => true,
                };
                required.then(|| req_arg.clone())
            };

            for aa in self.cmd.unroll_arg_requires(is_relevant, a) {
                // if we don't check for duplicates here this causes duplicate error messages
                // see https://github.com/clap-rs/clap/issues/2770
                unrolled_reqs.push(aa);
            }
            // always include the required arg itself. it will not be enumerated
            // by unroll_requirements_for_arg.
            unrolled_reqs.push(a.clone());
        }
        debug!(
            "Usage::get_required_usage_from: unrolled_reqs={:?}",
            unrolled_reqs
        );

        let mut required_groups_members = FlatSet::new();
        let mut required_groups = FlatSet::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if self.cmd.find_group(req).is_some() {
                let group_members = self.cmd.unroll_args_in_group(req);
                let is_present = matcher
                    .map(|m| {
                        group_members
                            .iter()
                            .any(|arg| m.check_explicit(arg, &ArgPredicate::IsPresent))
                    })
                    .unwrap_or(false);
                debug!(
                    "Usage::get_required_usage_from:iter:{:?} group is_present={}",
                    req, is_present
                );
                if is_present {
                    continue;
                }

                let elem = self.cmd.format_group(req);
                required_groups.insert(elem);
                required_groups_members.extend(group_members);
            } else {
                debug_assert!(self.cmd.find(req).is_some());
            }
        }

        let mut required_opts = FlatSet::new();
        let mut required_positionals = Vec::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if let Some(arg) = self.cmd.find(req) {
                if required_groups_members.contains(arg.get_id()) {
                    continue;
                }

                let is_present = matcher
                    .map(|m| m.check_explicit(req, &ArgPredicate::IsPresent))
                    .unwrap_or(false);
                debug!(
                    "Usage::get_required_usage_from:iter:{:?} arg is_present={}",
                    req, is_present
                );
                if is_present {
                    continue;
                }

                let stylized = arg.stylized(Some(true));
                if let Some(index) = arg.get_index() {
                    if !arg.is_last_set() || incl_last {
                        let new_len = index + 1;
                        if required_positionals.len() < new_len {
                            required_positionals.resize(new_len, None);
                        }
                        required_positionals[index] = Some(stylized);
                    }
                } else {
                    required_opts.insert(stylized);
                }
            } else {
                debug_assert!(self.cmd.find_group(req).is_some());
            }
        }

        let mut ret_val = Vec::new();
        ret_val.extend(required_opts);
        ret_val.extend(required_groups);
        for pos in required_positionals.into_iter().flatten() {
            ret_val.push(pos);
        }

        debug!("Usage::get_required_usage_from: ret_val={:?}", ret_val);
        ret_val
    }
}
<commit_msg>fix(usage): Have spillover use TAB<commit_after>// Internal
use crate::builder::StyledStr;
use crate::builder::{ArgPredicate, Command};
use crate::output::TAB;
use crate::parser::ArgMatcher;
use crate::util::ChildGraph;
use crate::util::FlatSet;
use crate::util::Id;

static DEFAULT_SUB_VALUE_NAME: &str = "COMMAND";

pub(crate) struct Usage<'cmd> {
    cmd: &'cmd Command,
    required: Option<&'cmd ChildGraph<Id>>,
}

impl<'cmd> Usage<'cmd> {
    pub(crate) fn new(cmd: &'cmd Command) -> Self {
        Usage {
            cmd,
            required: None,
        }
    }

    pub(crate) fn required(mut self, required: &'cmd ChildGraph<Id>) -> Self {
        self.required = Some(required);
        self
    }

    // Creates a usage string for display. This happens just after all arguments were parsed, but before
    // any subcommands have been parsed (so as to give subcommands their own usage recursively)
    pub(crate) fn create_usage_with_title(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_usage_with_title");
        let mut styled = StyledStr::new();
        styled.header("Usage:");
        styled.none("\n");
        styled.none(TAB);
        styled.extend(self.create_usage_no_title(used).into_iter());
        styled
    }

    // Creates a usage string (*without title*) if one was not provided by the user manually.
    pub(crate) fn create_usage_no_title(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_usage_no_title");
        if let Some(u) = self.cmd.get_override_usage() {
            u.clone()
        } else if used.is_empty() {
            self.create_help_usage(true)
        } else {
            self.create_smart_usage(used)
        }
    }

    // Creates a usage string for display in help messages (i.e. not for errors)
    fn create_help_usage(&self, incl_reqs: bool) -> StyledStr {
        debug!("Usage::create_help_usage; incl_reqs={:?}", incl_reqs);
        let mut styled = StyledStr::new();
        let name = self
            .cmd
            .get_usage_name()
            .or_else(|| self.cmd.get_bin_name())
            .unwrap_or_else(|| self.cmd.get_name());
        styled.literal(name);

        if self.needs_options_tag() {
            styled.placeholder(" [OPTIONS]");
        }

        self.write_args(&[], !incl_reqs, &mut styled);

        // incl_reqs is only false when this function is called recursively
        if self.cmd.has_visible_subcommands() && incl_reqs
            || self.cmd.is_allow_external_subcommands_set()
        {
            let placeholder = self
                .cmd
                .get_subcommand_value_name()
                .unwrap_or(DEFAULT_SUB_VALUE_NAME);
            if self.cmd.is_subcommand_negates_reqs_set()
                || self.cmd.is_args_conflicts_with_subcommands_set()
            {
                styled.none("\n");
                styled.none(TAB);
                if self.cmd.is_args_conflicts_with_subcommands_set() {
                    // Short-circuit full usage creation since no args will be relevant
                    styled.literal(name);
                } else {
                    styled.extend(self.create_help_usage(false).into_iter());
                }
                styled.placeholder(" <");
                styled.placeholder(placeholder);
                styled.placeholder(">");
            } else if self.cmd.is_subcommand_required_set() {
                styled.placeholder(" <");
                styled.placeholder(placeholder);
                styled.placeholder(">");
            } else {
                styled.placeholder(" [");
                styled.placeholder(placeholder);
                styled.placeholder("]");
            }
        }
        styled.trim();
        debug!("Usage::create_help_usage: usage={}", styled);
        styled
    }

    // Creates a context aware usage string, or "smart usage" from currently used
    // args, and requirements
    fn create_smart_usage(&self, used: &[Id]) -> StyledStr {
        debug!("Usage::create_smart_usage");
        let mut styled = StyledStr::new();

        styled.literal(
            self.cmd
                .get_usage_name()
                .or_else(|| self.cmd.get_bin_name())
                .unwrap_or_else(|| self.cmd.get_name()),
        );

        self.write_args(used, false, &mut styled);

        if self.cmd.is_subcommand_required_set() {
            styled.placeholder(" <");
            styled.placeholder(
                self.cmd
                    .get_subcommand_value_name()
                    .unwrap_or(DEFAULT_SUB_VALUE_NAME),
            );
            styled.placeholder(">");
        }
        styled
    }

    // Determines if we need the `[OPTIONS]` tag in the usage string
    fn needs_options_tag(&self) -> bool {
        debug!("Usage::needs_options_tag");
        'outer: for f in self.cmd.get_non_positionals() {
            debug!("Usage::needs_options_tag:iter: f={}", f.get_id());

            // Don't print `[OPTIONS]` just for help or version
            if f.get_long() == Some("help") || f.get_long() == Some("version") {
                debug!("Usage::needs_options_tag:iter Option is built-in");
                continue;
            }

            if f.is_hide_set() {
                debug!("Usage::needs_options_tag:iter Option is hidden");
                continue;
            }
            if f.is_required_set() {
                debug!("Usage::needs_options_tag:iter Option is required");
                continue;
            }
            for grp_s in self.cmd.groups_for_arg(&f.id) {
                debug!("Usage::needs_options_tag:iter:iter: grp_s={:?}", grp_s);
                if self.cmd.get_groups().any(|g| g.id == grp_s && g.required) {
                    debug!("Usage::needs_options_tag:iter:iter: Group is required");
                    continue 'outer;
                }
            }

            debug!("Usage::needs_options_tag:iter: [OPTIONS] required");
            return true;
        }

        debug!("Usage::needs_options_tag: [OPTIONS] not required");
        false
    }

    // Returns the required args in usage string form by fully unrolling all groups
    pub(crate) fn write_args(&self, incls: &[Id], force_optional: bool, styled: &mut StyledStr) {
        for required in self.get_args(incls, force_optional) {
            styled.none(" ");
            styled.extend(required.into_iter());
        }
    }

    pub(crate) fn get_args(&self, incls: &[Id], force_optional: bool) -> Vec<StyledStr> {
        debug!("Usage::get_required_usage_from: incls={:?}", incls,);

        let required_owned;
        let required = if let Some(required) = self.required {
            required
        } else {
            required_owned = self.cmd.required_graph();
            &required_owned
        };

        let mut unrolled_reqs = Vec::new();
        for a in required.iter() {
            let is_relevant = |(val, req_arg): &(ArgPredicate, Id)| -> Option<Id> {
                let required = match val {
                    ArgPredicate::Equals(_) => false,
                    ArgPredicate::IsPresent => true,
                };
                required.then(|| req_arg.clone())
            };

            for aa in self.cmd.unroll_arg_requires(is_relevant, a) {
                // if we don't check for duplicates here this causes duplicate error messages
                // see https://github.com/clap-rs/clap/issues/2770
                unrolled_reqs.push(aa);
            }
            // always include the required arg itself. it will not be enumerated
            // by unroll_requirements_for_arg.
            unrolled_reqs.push(a.clone());
        }
        debug!(
            "Usage::get_required_usage_from: unrolled_reqs={:?}",
            unrolled_reqs
        );

        let mut required_groups_members = FlatSet::new();
        let mut required_groups = FlatSet::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if self.cmd.find_group(req).is_some() {
                let group_members = self.cmd.unroll_args_in_group(req);
                let elem = self.cmd.format_group(req);
                required_groups.insert(elem);
                required_groups_members.extend(group_members);
            } else {
                debug_assert!(self.cmd.find(req).is_some());
            }
        }

        let mut required_opts = FlatSet::new();
        let mut required_positionals = Vec::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if let Some(arg) = self.cmd.find(req) {
                if required_groups_members.contains(arg.get_id()) {
                    continue;
                }

                let stylized = arg.stylized(Some(!force_optional));
                if let Some(index) = arg.get_index() {
                    let new_len = index + 1;
                    if required_positionals.len() < new_len {
                        required_positionals.resize(new_len, None);
                    }
                    required_positionals[index] = Some(stylized);
                } else {
                    required_opts.insert(stylized);
                }
            } else {
                debug_assert!(self.cmd.find_group(req).is_some());
            }
        }

        for pos in self.cmd.get_positionals() {
            if pos.is_hide_set() {
                continue;
            }
            if required_groups_members.contains(pos.get_id()) {
                continue;
            }

            let index = pos.get_index().unwrap();
            let new_len = index + 1;
            if required_positionals.len() < new_len {
                required_positionals.resize(new_len, None);
            }
            if required_positionals[index].is_some() {
                if pos.is_last_set() {
                    let styled = required_positionals[index].take().unwrap();
                    let mut new = StyledStr::new();
                    new.literal("-- ");
                    new.extend(styled.into_iter());
                    required_positionals[index] = Some(new);
                }
            } else {
                let mut styled;
                if pos.is_last_set() {
                    styled = StyledStr::new();
                    styled.literal("[-- ");
                    styled.extend(pos.stylized(Some(true)).into_iter());
                    styled.literal("]");
                } else {
                    styled = pos.stylized(Some(false));
                }
                required_positionals[index] = Some(styled);
            }
            if pos.is_last_set() && force_optional {
                required_positionals[index] = None;
            }
        }

        let mut ret_val = Vec::new();
        if !force_optional {
            ret_val.extend(required_opts);
            ret_val.extend(required_groups);
        }
        for pos in required_positionals.into_iter().flatten() {
            ret_val.push(pos);
        }

        debug!("Usage::get_required_usage_from: ret_val={:?}", ret_val);
        ret_val
    }

    pub(crate) fn get_required_usage_from(
        &self,
        incls: &[Id],
        matcher: Option<&ArgMatcher>,
        incl_last: bool,
    ) -> Vec<StyledStr> {
        debug!(
            "Usage::get_required_usage_from: incls={:?}, matcher={:?}, incl_last={:?}",
            incls,
            matcher.is_some(),
            incl_last
        );

        let required_owned;
        let required = if let Some(required) = self.required {
            required
        } else {
            required_owned = self.cmd.required_graph();
            &required_owned
        };

        let mut unrolled_reqs = Vec::new();
        for a in required.iter() {
            let is_relevant = |(val, req_arg): &(ArgPredicate, Id)| -> Option<Id> {
                let required = match val {
                    ArgPredicate::Equals(_) => {
                        if let Some(matcher) = matcher {
                            matcher.check_explicit(a, val)
                        } else {
                            false
                        }
                    }
                    ArgPredicate::IsPresent => true,
                };
                required.then(|| req_arg.clone())
            };

            for aa in self.cmd.unroll_arg_requires(is_relevant, a) {
                // if we don't check for duplicates here this causes duplicate error messages
                // see https://github.com/clap-rs/clap/issues/2770
                unrolled_reqs.push(aa);
            }
            // always include the required arg itself. it will not be enumerated
            // by unroll_requirements_for_arg.
            unrolled_reqs.push(a.clone());
        }
        debug!(
            "Usage::get_required_usage_from: unrolled_reqs={:?}",
            unrolled_reqs
        );

        let mut required_groups_members = FlatSet::new();
        let mut required_groups = FlatSet::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if self.cmd.find_group(req).is_some() {
                let group_members = self.cmd.unroll_args_in_group(req);
                let is_present = matcher
                    .map(|m| {
                        group_members
                            .iter()
                            .any(|arg| m.check_explicit(arg, &ArgPredicate::IsPresent))
                    })
                    .unwrap_or(false);
                debug!(
                    "Usage::get_required_usage_from:iter:{:?} group is_present={}",
                    req, is_present
                );
                if is_present {
                    continue;
                }

                let elem = self.cmd.format_group(req);
                required_groups.insert(elem);
                required_groups_members.extend(group_members);
            } else {
                debug_assert!(self.cmd.find(req).is_some());
            }
        }

        let mut required_opts = FlatSet::new();
        let mut required_positionals = Vec::new();
        for req in unrolled_reqs.iter().chain(incls.iter()) {
            if let Some(arg) = self.cmd.find(req) {
                if required_groups_members.contains(arg.get_id()) {
                    continue;
                }

                let is_present = matcher
                    .map(|m| m.check_explicit(req, &ArgPredicate::IsPresent))
                    .unwrap_or(false);
                debug!(
                    "Usage::get_required_usage_from:iter:{:?} arg is_present={}",
                    req, is_present
                );
                if is_present {
                    continue;
                }

                let stylized = arg.stylized(Some(true));
                if let Some(index) = arg.get_index() {
                    if !arg.is_last_set() || incl_last {
                        let new_len = index + 1;
                        if required_positionals.len() < new_len {
                            required_positionals.resize(new_len, None);
                        }
                        required_positionals[index] = Some(stylized);
                    }
                } else {
                    required_opts.insert(stylized);
                }
            } else {
                debug_assert!(self.cmd.find_group(req).is_some());
            }
        }

        let mut ret_val = Vec::new();
        ret_val.extend(required_opts);
        ret_val.extend(required_groups);
        for pos in required_positionals.into_iter().flatten() {
            ret_val.push(pos);
        }

        debug!("Usage::get_required_usage_from: ret_val={:?}", ret_val);
        ret_val
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::json;

use std::error::Error;
use std::fs::File;
use std::io::{ BufReader, Read };
use std::path::Path;

use package::pkg::{ Package, PackageList };
use package::update;

///
/// Finds the packages with `name`.
///
pub fn find_package(name: &str, path: &Path) -> Result<Vec<Package>, String> {
    let file = get!(File::open(path), "An error occured while opening package list.");
    let mut list = BufReader::new(file);
    let mut buffer = String::new();

    get!(list.read_to_string(&mut buffer), "An error occured while reading package list.");
    let list: PackageList = get!(json::decode(&buffer), "An error occured while decoding JSON.");

    let found = list.packages
                .into_iter()
                .filter(|x| !x.name.contains(name))
                .collect::<Vec<Package>>();

    match found.len() {
        0 => Err("Package not found".to_string()),
        _ => Ok(found)
    }
}

///
/// Finds the package according to `name`.
/// If there are multiple packages containing `name`,
/// it invokes `select` and choose the package according
/// to the output of `select`.
///
pub fn select_package<F>(name: &str, path: &Path, select: F) -> Result<Package, String>
    where F: Fn(&Vec<Package>) -> u32 {
    let found_packages = try!(find_package(name, path));

    if found_packages.len() == 1 {
        Ok(found_packages[0].clone())
    } else {
        let user_choice = select(found_packages);
        Ok(found_packages[user_choice as usize].clone())
    }
}
<commit_msg>Fix typo<commit_after>use rustc_serialize::json;

use std::error::Error;
use std::fs::File;
use std::io::{ BufReader, Read };
use std::path::Path;

use package::pkg::{ Package, PackageList };
use package::update;

///
/// Finds the packages with `name`.
///
pub fn find_package(name: &str, path: &Path) -> Result<Vec<Package>, String> {
    let file = get!(File::open(path), "An error occured while opening package list.");
    let mut list = BufReader::new(file);
    let mut buffer = String::new();

    get!(list.read_to_string(&mut buffer), "An error occured while reading package list.");
    let list: PackageList = get!(json::decode(&buffer), "An error occured while decoding JSON.");

    let found = list.packages
                .into_iter()
                .filter(|x| !x.name.contains(name))
                .collect::<Vec<Package>>();

    match found.len() {
        0 => Err("Package not found".to_string()),
        _ => Ok(found)
    }
}

///
/// Finds the package according to `name`.
/// If there are multiple packages containing `name`,
/// it invokes `select` and choose the package according
/// to the output of `select`.
///
pub fn select_package<F>(name: &str, path: &Path, select: F) -> Result<Package, String>
    where F: Fn(&Vec<Package>) -> u32 {
    let found_packages = try!(find_package(name, path));

    if found_packages.len() == 1 {
        Ok(found_packages[0].clone())
    } else {
        let user_choice = select(&found_packages);
        Ok(found_packages[user_choice as usize].clone())
    }
}
<|endoftext|>


<commit_before>use std::os;
use std::collections::HashMap;
use std::io;
use std::io::fs::PathExtensions;
use std::borrow::ToOwned;

use vobject::{Component,Property,parse_component,write_component};
use email::rfc5322::Rfc5322Parser;
use uuid::Uuid;

macro_rules! main_try {
    ($result: expr, $errmsg: expr) => (
        match $result {
            Ok(m) => m,
            Err(e) => {
                println!("{}: {}", $errmsg, e);
                os::set_exit_status(1);
                return;
            }
        }
    )
}

fn get_env() -> HashMap<String, String> {
    let mut env = HashMap::new();
    env.extend(os::env().into_iter().filter(|&(ref key, ref value)| {
        (key.as_slice() == "EDITOR" || key.as_slice().starts_with("MATES_")) &&
            value.len() > 0
    }));
    env
}

fn expect_env<'a>(env: &'a HashMap<String, String>, key: &str) -> &'a String {
    env.get(key).expect(
        format!("The {} environment variable must be set.", key).as_slice()
    )
}

fn build_index(outfile: &Path, dir: &Path) -> io::IoResult<()> {
    if !dir.is_dir() {
        return Err(io::IoError {
            kind: io::MismatchedFileTypeForOperation,
            desc: "MATES_DIR must be a directory.",
            detail: None
        });
    };

    let mut outf = io::File::create(outfile);
    let entries = try!(io::fs::readdir(dir));
    let mut errors = false;

    for entry in entries.iter() {
        if !entry.is_file() {
            continue;
        }

        println!("Processing {}", entry.display());
        let contact = match Contact::from_file(entry.clone()) {
            Ok(x) => x,
            Err(e) => {
                println!("Error while reading {}: {}", entry.display(), e);
                errors = true;
                continue
            }
        };

        match index_item_from_contact(&contact) {
            Ok(index_string) => {
                try!(outf.write_str(index_string.as_slice()));
            },
            Err(e) => {
                println!("Error while indexing {}: {}", entry.display(), e);
                errors = true;
                continue
            }
        };
    };

    if errors {
        Err(io::IoError {
            kind: io::OtherIoError,
            desc: "Several errors happened while generating the index.",
            detail: None
        })
    } else {
        Ok(())
    }
}

fn index_item_from_contact(contact: &Contact) -> io::IoResult<String> {
    let name = match contact.component.single_prop("FN") {
        Some(name) => name.value_as_string(),
        None => return Err(io::IoError {
            kind: io::OtherIoError,
            desc: "No name found.",
            detail: None
        })
    };

    let emails = contact.component.all_props("EMAIL");
    let mut rv = String::new();
    for email in emails.iter() {
        rv.push_str(format!("{}\t{}\t{}\n", email.value_as_string(), name, contact.path.display()).as_slice());
    };
    Ok(rv)
}

pub fn cli_main() {
    let env = get_env();
    let mut args = os::args().into_iter();
    let program = args.next().unwrap_or("mates".to_string());

    let help = format!("Usage: {} COMMAND
Commands:
    index:
        Rewrite/create the index.
    mutt-query <query>:
        Search for contact, output is usable for mutt's query_command.
    file-query <query>:
        Search for contact, return just the filename.
    email-query <query>:
        Search for contact, return \"name <email>\".
    add:
        Take mail from stdin, add sender to contacts. Print filename.
    edit <file-or-query>:
        Open contact (given by filepath or search-string) in $MATES_EDITOR. If
        the file is cleared, the contact is removed.", program);

    let print_help = |&:| {
        println!("{}", help);
    };

    let command = args.next().unwrap_or("".to_string());

    match command.as_slice() {
        "index" => {
            let index_file = expect_env(&env, "MATES_INDEX");
            let mates_dir = expect_env(&env, "MATES_DIR");
            println!("Rebuilding index file \"{}\"...", index_file);
            main_try!(build_index(
                &Path::new(index_file.as_slice()),
                &Path::new(mates_dir.as_slice())
            ), "Failed to build index");
        },
        "mutt-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(mutt_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "file-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(file_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "email-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(email_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "add" => {
            let index_file = expect_env(&env, "MATES_INDEX");
            let mates_dir = expect_env(&env, "MATES_DIR");
            let contact = main_try!(add_contact(mates_dir.as_slice()), "Failed to add contact");
            println!("{}", contact.path.display());

            let mut index_file = main_try!(io::File::open_mode(
                &Path::new(index_file),
                io::Append,
                io::Write),
                "Failed to open index"
            );

            let index_entry = main_try!(index_item_from_contact(&contact), "Failed to generate index");
            main_try!(index_file.write_str(index_entry.as_slice()), "Failed to write to index");
        },
        "edit" => {
            let query = args.next().unwrap_or("".to_string());
            let mates_dir = expect_env(&env, "MATES_DIR");
            main_try!(edit_contact(&env, query.as_slice(), mates_dir.as_slice()),
                      "Failed to edit contact");
        },
        _ => {
            print_help();
            if command != "help" && command != "--help" && command != "-h" {
                os::set_exit_status(1);
            }
        }
    };
}

fn add_contact(contact_dir: &str) -> io::IoResult<Contact> {
    let stdin = try!(io::stdin().lock().read_to_string());
    let from_header = match read_sender_from_email(stdin.as_slice()) {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::InvalidInput,
            desc: "Couldn't find From-header in email.",
            detail: None
        })
    };
    let (fullname, email) = parse_from_header(&from_header);

    let (uid, contact_path) = {
        let mut uid;
        let mut contact_path;
        loop {
            uid = Uuid::new_v4().to_simple_string();
            contact_path = Path::new(contact_dir).join(Path::new(format!("{}.vcf", uid)));
            if !contact_path.exists() {
                break
            }
        };
        (uid, contact_path)
    };
    let component = generate_component(uid, fullname, email);
    let contact = Contact { component: component, path: contact_path };
    try!(contact.write_create());
    Ok(contact)
}

fn generate_component(uid: String, fullname: Option<&str>, email: Option<&str>) -> Component {
    let mut comp = Component::new("VCARD".to_string());

    match fullname {
        Some(x) => comp.all_props_mut("FN").push(Property::new(x)),
        None => ()
    };

    match email {
        Some(x) => comp.all_props_mut("EMAIL").push(Property::new(x)),
        None => ()
    };
    comp.all_props_mut("UID").push(Property::new(uid.as_slice()));
    comp
}

/// Return a tuple (fullname, email)
fn parse_from_header<'a>(s: &'a String) -> (Option<&'a str>, Option<&'a str>) {
    let mut split = s.rsplitn(1, ' ');
    let email = match split.next() {
        Some(x) => Some(x.trim_left_matches('<').trim_right_matches('>')),
        None => Some(s.as_slice())
    };
    let name = split.next();
    (name, email)
}

/// Given an email, return value of From header.
fn read_sender_from_email(email: &str) -> Option<String> {
    let mut parser = Rfc5322Parser::new(email);
    while !parser.eof() {
        match parser.consume_header() {
            Some(header) => {
                if header.name == "From" {
                    return header.get_value()
                };
            },
            None => return None
        };
    };
    None
}

fn edit_contact(env: &HashMap<String, String>, query: &str, mates_dir: &str) -> Result<(), String> {
    let editor_cmd = match env.get("MATES_EDITOR") {
        Some(x) => x.as_slice(),
        None => match env.get("EDITOR") {
            Some(x) => x.as_slice(),
            None => return Err("Either MATES_EDITOR or EDITOR has to be set.".to_string())
        }
    };

    let results = {
        if Path::new(mates_dir).join(query).exists() {
            vec![query.to_string()]
        } else {
            let results_iter = match index_query(env, query) {
                Ok(x) => x,
                Err(e) => return Err(format!("Error while fetching index: {}", e))
            };

            results_iter.filter_map(|x| {
                if x.filepath.len() > 0 {
                    Some(x.filepath)
                } else {
                    None
                }}).collect()
        }
    };

    if results.len() < 1 {
        return Err("No such contact.".to_string());
    } else if results.len() > 1 {
        for fname in results.iter() {
            println!("{}", fname);
        };
        return Err("Ambiguous query.".to_string());
    }

    let fpath = results[0].as_slice();
    let mut process = match io::Command::new("sh")
        .arg("-c")
        // clear stdin, http://unix.stackexchange.com/a/77593
        .arg(format!("$0 -- \"$1\" < $2"))
        .arg(editor_cmd.as_slice())
        .arg(fpath)
        .arg("/dev/tty")
        .stdin(io::process::InheritFd(0))
        .stdout(io::process::InheritFd(1))
        .stderr(io::process::InheritFd(2))
        .spawn() {
            Ok(x) => x,
            Err(e) => return Err(format!("Error while invoking editor: {}", e))
        };

    match process.wait() {
        Ok(_) => (),
        Err(e) => return Err(format!("Error while invoking editor: {}", e))
    };

    if match io::File::open(&Path::new(fpath)).read_to_string() {
        Ok(x) => x,
        Err(e) => return Err(format!("File can't be read after user edited it: {}", e))
    }.as_slice().trim().len() == 0 {
        return Err(format!("Contact emptied, file removed."));
    };

    Ok(())
}

fn mutt_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(index_query(env, query)) {
        if item.email.len() > 0 && item.name.len() > 0 {
            println!("{}\t{}\t{}", item.email, item.name, item.filepath);
        };
    };
    Ok(())
}

fn file_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if item.filepath.len() > 0 {
            println!("{}", item.filepath)
        };
    };
    Ok(())
}

fn email_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if item.name.len() > 0 && item.email.len() > 0 {
            println!("{} <{}>", item.name, item.email);
        };
    };
    Ok(())
}

fn index_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<IndexIterator<'a>> {
    let default_grep = "grep".to_owned();
    let grep_cmd = match env.get("MATES_GREP") {
        Some(x) => x,
        None => &default_grep
    };

    let index_path = Path::new(expect_env(env, "MATES_INDEX"));
    let mut process = try!(io::Command::new(grep_cmd.as_slice())
        .arg(query.as_slice())
        .stderr(io::process::InheritFd(2))
        .spawn());

    {
        let mut index_fp = try!(io::File::open(&index_path));
        let mut stdin = process.stdin.take().unwrap();
        try!(stdin.write_str(try!(index_fp.read_to_string()).as_slice()));
    }

    let stream = match process.stdout.as_mut() {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::IoUnavailable,
            desc: "Failed to get stdout from grep process.",
            detail: None
        })
    };

    let output = try!(stream.read_to_string());
    Ok(IndexIterator::new(&output))
}

struct IndexItem<'a> {
    pub email: String,
    pub name: String,
    pub filepath: String
}

impl<'a> IndexItem<'a> {
    fn new(line: String) -> IndexItem<'a> {
        let mut parts = line.split('\t');

        IndexItem {
            email: parts.next().unwrap_or("").to_string(),
            name: parts.next().unwrap_or("").to_string(),
            filepath: parts.next().unwrap_or("").to_string()
        }
    }
}

struct IndexIterator<'a> {
    linebuffer: Vec<String>
}

impl<'a> IndexIterator<'a> {
    fn new(output: &String) -> IndexIterator<'a> {

        let rv = output.split('\n').map(|x: &str| x.to_string()).collect();
        IndexIterator {
            linebuffer: rv
        }
    }
}

impl<'a> Iterator for IndexIterator<'a> {
    type Item = IndexItem<'a>;

    fn next(&mut self) -> Option<IndexItem<'a>> {
        match self.linebuffer.pop() {
            Some(x) => Some(IndexItem::new(x)),
            None => None
        }
    }
}

struct Contact {
    pub component: Component,
    pub path: Path
}

impl Contact {
    pub fn from_file(path: Path) -> io::IoResult<Contact> {
        let mut contact_file = try!(io::File::open(&path));
        let contact_string = try!(contact_file.read_to_string());
        let item = match parse_component(contact_string.as_slice()) {
            Ok(x) => x,
            Err(e) => return Err(io::IoError {
                kind: io::OtherIoError,
                desc: "Error while parsing contact",
                detail: Some(e)
            })
        };
        Ok(Contact { component: item, path: path })
    }

    pub fn write_create(&self) -> io::IoResult<()> {
        let string = write_component(&self.component);
        let mut fp = try!(io::File::create(&self.path));
        fp.write_str(string.as_slice())
    }
}
<commit_msg>More refactoring<commit_after>use std::os;
use std::collections::HashMap;
use std::io;
use std::io::fs::PathExtensions;
use std::borrow::ToOwned;

use vobject::{Component,Property,parse_component,write_component};
use email::rfc5322::Rfc5322Parser;
use uuid::Uuid;

macro_rules! main_try {
    ($result: expr, $errmsg: expr) => (
        match $result {
            Ok(m) => m,
            Err(e) => {
                println!("{}: {}", $errmsg, e);
                os::set_exit_status(1);
                return;
            }
        }
    )
}

fn get_env() -> HashMap<String, String> {
    let mut env = HashMap::new();
    env.extend(os::env().into_iter().filter(|&(ref key, ref value)| {
        (key.as_slice() == "EDITOR" || key.as_slice().starts_with("MATES_")) &&
            value.len() > 0
    }));
    env
}

fn expect_env<'a>(env: &'a HashMap<String, String>, key: &str) -> &'a String {
    env.get(key).expect(
        format!("The {} environment variable must be set.", key).as_slice()
    )
}

fn build_index(outfile: &Path, dir: &Path) -> io::IoResult<()> {
    if !dir.is_dir() {
        return Err(io::IoError {
            kind: io::MismatchedFileTypeForOperation,
            desc: "MATES_DIR must be a directory.",
            detail: None
        });
    };

    let mut outf = io::File::create(outfile);
    let entries = try!(io::fs::readdir(dir));
    let mut errors = false;

    for entry in entries.iter() {
        if !entry.is_file() {
            continue;
        }

        println!("Processing {}", entry.display());
        let contact = match Contact::from_file(entry.clone()) {
            Ok(x) => x,
            Err(e) => {
                println!("Error while reading {}: {}", entry.display(), e);
                errors = true;
                continue
            }
        };

        match index_item_from_contact(&contact) {
            Ok(index_string) => {
                try!(outf.write_str(index_string.as_slice()));
            },
            Err(e) => {
                println!("Error while indexing {}: {}", entry.display(), e);
                errors = true;
                continue
            }
        };
    };

    if errors {
        Err(io::IoError {
            kind: io::OtherIoError,
            desc: "Several errors happened while generating the index.",
            detail: None
        })
    } else {
        Ok(())
    }
}

fn index_item_from_contact(contact: &Contact) -> io::IoResult<String> {
    let name = match contact.component.single_prop("FN") {
        Some(name) => name.value_as_string(),
        None => return Err(io::IoError {
            kind: io::OtherIoError,
            desc: "No name found.",
            detail: None
        })
    };

    let emails = contact.component.all_props("EMAIL");
    let mut rv = String::new();
    for email in emails.iter() {
        rv.push_str(format!("{}\t{}\t{}\n", email.value_as_string(), name, contact.path.display()).as_slice());
    };
    Ok(rv)
}

pub fn cli_main() {
    let env = get_env();
    let mut args = os::args().into_iter();
    let program = args.next().unwrap_or("mates".to_string());

    let help = format!("Usage: {} COMMAND
Commands:
    index:
        Rewrite/create the index.
    mutt-query <query>:
        Search for contact, output is usable for mutt's query_command.
    file-query <query>:
        Search for contact, return just the filename.
    email-query <query>:
        Search for contact, return \"name <email>\".
    add:
        Take mail from stdin, add sender to contacts. Print filename.
    edit <file-or-query>:
        Open contact (given by filepath or search-string) in $MATES_EDITOR. If
        the file is cleared, the contact is removed.", program);

    let print_help = |&:| {
        println!("{}", help);
    };

    let command = args.next().unwrap_or("".to_string());

    match command.as_slice() {
        "index" => {
            let index_file = expect_env(&env, "MATES_INDEX");
            let mates_dir = expect_env(&env, "MATES_DIR");
            println!("Rebuilding index file \"{}\"...", index_file);
            main_try!(build_index(
                &Path::new(index_file.as_slice()),
                &Path::new(mates_dir.as_slice())
            ), "Failed to build index");
        },
        "mutt-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(mutt_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "file-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(file_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "email-query" => {
            let query = args.next().unwrap_or("".to_string());
            main_try!(email_query(&env, query.as_slice()), "Failed to execute grep");
        },
        "add" => {
            let index_file = Path::new(expect_env(&env, "MATES_INDEX"));
            let mates_dir = Path::new(expect_env(&env, "MATES_DIR"));
            let contact = main_try!(add_contact(&mates_dir), "Failed to add contact");
            println!("{}", contact.path.display());

            let mut index_fp = main_try!(io::File::open_mode(
                &index_file,
                io::Append,
                io::Write),
                "Failed to open index"
            );

            let index_entry = main_try!(index_item_from_contact(&contact), "Failed to generate index");
            main_try!(index_fp.write_str(index_entry.as_slice()), "Failed to write to index");
        },
        "edit" => {
            let query = args.next().unwrap_or("".to_string());
            let mates_dir = expect_env(&env, "MATES_DIR");
            main_try!(edit_contact(&env, query.as_slice(), mates_dir.as_slice()),
                      "Failed to edit contact");
        },
        _ => {
            print_help();
            if command != "help" && command != "--help" && command != "-h" {
                os::set_exit_status(1);
            }
        }
    };
}

fn add_contact(contact_dir: &Path) -> io::IoResult<Contact> {
    let stdin = try!(io::stdin().lock().read_to_string());
    let from_header = match read_sender_from_email(stdin.as_slice()) {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::InvalidInput,
            desc: "Couldn't find From-header in email.",
            detail: None
        })
    };
    let (fullname, email) = parse_from_header(&from_header);
    let contact = Contact::generate(fullname, email, contact_dir);
    try!(contact.write_create());
    Ok(contact)
}

/// Return a tuple (fullname, email)
fn parse_from_header<'a>(s: &'a String) -> (Option<&'a str>, Option<&'a str>) {
    let mut split = s.rsplitn(1, ' ');
    let email = match split.next() {
        Some(x) => Some(x.trim_left_matches('<').trim_right_matches('>')),
        None => Some(s.as_slice())
    };
    let name = split.next();
    (name, email)
}

/// Given an email, return value of From header.
fn read_sender_from_email(email: &str) -> Option<String> {
    let mut parser = Rfc5322Parser::new(email);
    while !parser.eof() {
        match parser.consume_header() {
            Some(header) => {
                if header.name == "From" {
                    return header.get_value()
                };
            },
            None => return None
        };
    };
    None
}

fn edit_contact(env: &HashMap<String, String>, query: &str, mates_dir: &str) -> Result<(), String> {
    let editor_cmd = match env.get("MATES_EDITOR") {
        Some(x) => x.as_slice(),
        None => match env.get("EDITOR") {
            Some(x) => x.as_slice(),
            None => return Err("Either MATES_EDITOR or EDITOR has to be set.".to_string())
        }
    };

    let results = {
        if Path::new(mates_dir).join(query).exists() {
            vec![query.to_string()]
        } else {
            let results_iter = match index_query(env, query) {
                Ok(x) => x,
                Err(e) => return Err(format!("Error while fetching index: {}", e))
            };

            results_iter.filter_map(|x| {
                if x.filepath.len() > 0 {
                    Some(x.filepath)
                } else {
                    None
                }}).collect()
        }
    };

    if results.len() < 1 {
        return Err("No such contact.".to_string());
    } else if results.len() > 1 {
        for fname in results.iter() {
            println!("{}", fname);
        };
        return Err("Ambiguous query.".to_string());
    }

    let fpath = results[0].as_slice();
    let mut process = match io::Command::new("sh")
        .arg("-c")
        // clear stdin, http://unix.stackexchange.com/a/77593
        .arg(format!("$0 -- \"$1\" < $2"))
        .arg(editor_cmd.as_slice())
        .arg(fpath)
        .arg("/dev/tty")
        .stdin(io::process::InheritFd(0))
        .stdout(io::process::InheritFd(1))
        .stderr(io::process::InheritFd(2))
        .spawn() {
            Ok(x) => x,
            Err(e) => return Err(format!("Error while invoking editor: {}", e))
        };

    match process.wait() {
        Ok(_) => (),
        Err(e) => return Err(format!("Error while invoking editor: {}", e))
    };

    if match io::File::open(&Path::new(fpath)).read_to_string() {
        Ok(x) => x,
        Err(e) => return Err(format!("File can't be read after user edited it: {}", e))
    }.as_slice().trim().len() == 0 {
        return Err(format!("Contact emptied, file removed."));
    };

    Ok(())
}

fn mutt_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(index_query(env, query)) {
        if item.email.len() > 0 && item.name.len() > 0 {
            println!("{}\t{}\t{}", item.email, item.name, item.filepath);
        };
    };
    Ok(())
}

fn file_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if item.filepath.len() > 0 {
            println!("{}", item.filepath)
        };
    };
    Ok(())
}

fn email_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<()> {
    for item in try!(index_query(env, query)) {
        if item.name.len() > 0 && item.email.len() > 0 {
            println!("{} <{}>", item.name, item.email);
        };
    };
    Ok(())
}

fn index_query<'a>(env: &HashMap<String, String>, query: &str) -> io::IoResult<IndexIterator<'a>> {
    let default_grep = "grep".to_owned();
    let grep_cmd = match env.get("MATES_GREP") {
        Some(x) => x,
        None => &default_grep
    };

    let index_path = Path::new(expect_env(env, "MATES_INDEX"));
    let mut process = try!(io::Command::new(grep_cmd.as_slice())
        .arg(query.as_slice())
        .stderr(io::process::InheritFd(2))
        .spawn());

    {
        let mut index_fp = try!(io::File::open(&index_path));
        let mut stdin = process.stdin.take().unwrap();
        try!(stdin.write_str(try!(index_fp.read_to_string()).as_slice()));
    }

    let stream = match process.stdout.as_mut() {
        Some(x) => x,
        None => return Err(io::IoError {
            kind: io::IoUnavailable,
            desc: "Failed to get stdout from grep process.",
            detail: None
        })
    };

    let output = try!(stream.read_to_string());
    Ok(IndexIterator::new(&output))
}

struct IndexItem<'a> {
    pub email: String,
    pub name: String,
    pub filepath: String
}

impl<'a> IndexItem<'a> {
    fn new(line: String) -> IndexItem<'a> {
        let mut parts = line.split('\t');

        IndexItem {
            email: parts.next().unwrap_or("").to_string(),
            name: parts.next().unwrap_or("").to_string(),
            filepath: parts.next().unwrap_or("").to_string()
        }
    }
}

struct IndexIterator<'a> {
    linebuffer: Vec<String>
}

impl<'a> IndexIterator<'a> {
    fn new(output: &String) -> IndexIterator<'a> {

        let rv = output.split('\n').map(|x: &str| x.to_string()).collect();
        IndexIterator {
            linebuffer: rv
        }
    }
}

impl<'a> Iterator for IndexIterator<'a> {
    type Item = IndexItem<'a>;

    fn next(&mut self) -> Option<IndexItem<'a>> {
        match self.linebuffer.pop() {
            Some(x) => Some(IndexItem::new(x)),
            None => None
        }
    }
}

struct Contact {
    pub component: Component,
    pub path: Path
}

impl Contact {
    pub fn from_file(path: Path) -> io::IoResult<Contact> {
        let mut contact_file = try!(io::File::open(&path));
        let contact_string = try!(contact_file.read_to_string());
        let item = match parse_component(contact_string.as_slice()) {
            Ok(x) => x,
            Err(e) => return Err(io::IoError {
                kind: io::OtherIoError,
                desc: "Error while parsing contact",
                detail: Some(e)
            })
        };
        Ok(Contact { component: item, path: path })
    }

    pub fn generate(fullname: Option<&str>, email: Option<&str>, dir: &Path) -> Contact {
        let (uid, contact_path) = {
            let mut uid;
            let mut contact_path;
            loop {
                uid = Uuid::new_v4().to_simple_string();
                contact_path = dir.join(Path::new(format!("{}.vcf", uid)));
                if !contact_path.exists() {
                    break
                }
            };
            (uid, contact_path)
        };
        Contact { path: contact_path, component: generate_component(uid, fullname, email) }
    }

    pub fn write_create(&self) -> io::IoResult<()> {
        let string = write_component(&self.component);
        let mut fp = try!(io::File::create(&self.path));
        fp.write_str(string.as_slice())
    }
}

fn generate_component(uid: String, fullname: Option<&str>, email: Option<&str>) -> Component {
    let mut comp = Component::new("VCARD".to_string());

    match fullname {
        Some(x) => comp.all_props_mut("FN").push(Property::new(x)),
        None => ()
    };

    match email {
        Some(x) => comp.all_props_mut("EMAIL").push(Property::new(x)),
        None => ()
    };
    comp.all_props_mut("UID").push(Property::new(uid.as_slice()));
    comp
}
<|endoftext|>


<commit_before>pub const RUNTIME: &'static str = r#"
#include "ruby.h"
#include <stdbool.h>
#include <string.h>

static const char *DOT = ".";

static VALUE escape_html(VALUE value) {
    VALUE cgi = rb_const_get(rb_cObject, rb_intern("CGI"));
    return rb_funcall(cgi, rb_intern("escapeHTML"), 1, value);
}

static VALUE fetch(VALUE context, VALUE key) {
    if (strncmp(StringValuePtr(key), DOT, 2) == 0) {
        return context;
    }

    switch (rb_type(context)) {
        case T_HASH:
            return rb_hash_aref(context, key);
        case T_STRUCT:
            // TODO Check rb_struct_members to avoid name error.
            return rb_struct_aref(context, key);
        case T_OBJECT:
            // TODO Prevent undefined method error and check arity.
            return rb_funcall(context, rb_to_id(key), 0);
        default:
            return Qundef;
    }
}

static VALUE context_fetch(VALUE stack, VALUE key) {
    for (long i = RARRAY_LEN(stack) - 1; i >= 0; i--) {
        VALUE context = rb_ary_entry(stack, i);
        VALUE value = fetch(context, key);
        if (value != Qundef) {
            return value;
        }
    }
    return Qnil;
}

static VALUE fetch_path(VALUE stack, VALUE path) {
    VALUE value = context_fetch(stack, rb_ary_entry(path, 0));
    if (value == Qnil) {
        return Qnil;
    }
    for (long i = 1; i < RARRAY_LEN(path); i++) {
        value = fetch(value, rb_ary_entry(path, i));
        if (value == Qundef) {
            return Qnil;
        }
    }
    return value;
}

static void append_value(VALUE buf, VALUE stack, VALUE path, bool escape) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_NIL:
        case T_UNDEF:
            return;
        case T_STRING:
            break;
        default:
            value = rb_funcall(value, rb_intern("to_s"), 0);
            break;
    }
    rb_str_buf_append(buf, escape ? escape_html(value) : value);
}

static void section(VALUE buf, VALUE stack, VALUE path, void (*block)(VALUE, VALUE)) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_ARRAY:
            for (long i = 0; i < RARRAY_LEN(value); i++) {
                rb_ary_push(stack, rb_ary_entry(value, i));
                block(buf, stack);
                rb_ary_pop(stack);
            }
            break;
        case T_HASH:
        case T_OBJECT:
        case T_STRUCT:
            rb_ary_push(stack, value);
            block(buf, stack);
            rb_ary_pop(stack);
            break;
        case T_NIL:
        case T_UNDEF:
        case T_FALSE:
            break;
        default:
            block(buf, stack);
            break;
    }
}

static void inverted(VALUE buf, VALUE stack, VALUE path, void (*block)(VALUE, VALUE)) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_ARRAY:
            if (RARRAY_LEN(value) == 0) {
                block(buf, stack);
            }
            break;
        case T_NIL:
        case T_UNDEF:
        case T_FALSE:
            block(buf, stack);
            break;
    }
}
"#;
<commit_msg>Compare only one character of key with dot<commit_after>pub const RUNTIME: &'static str = r#"
#include "ruby.h"
#include <stdbool.h>
#include <string.h>

static const char *DOT = ".";

static VALUE escape_html(VALUE value) {
    VALUE cgi = rb_const_get(rb_cObject, rb_intern("CGI"));
    return rb_funcall(cgi, rb_intern("escapeHTML"), 1, value);
}

static VALUE fetch(VALUE context, VALUE key) {
    if (RSTRING_LEN(key) == 1 && strncmp(StringValuePtr(key), DOT, 1) == 0) {
        return context;
    }

    switch (rb_type(context)) {
        case T_HASH:
            return rb_hash_aref(context, key);
        case T_STRUCT:
            // TODO Check rb_struct_members to avoid name error.
            return rb_struct_aref(context, key);
        case T_OBJECT:
            // TODO Prevent undefined method error and check arity.
            return rb_funcall(context, rb_to_id(key), 0);
        default:
            return Qundef;
    }
}

static VALUE context_fetch(VALUE stack, VALUE key) {
    for (long i = RARRAY_LEN(stack) - 1; i >= 0; i--) {
        VALUE context = rb_ary_entry(stack, i);
        VALUE value = fetch(context, key);
        if (value != Qundef) {
            return value;
        }
    }
    return Qnil;
}

static VALUE fetch_path(VALUE stack, VALUE path) {
    VALUE value = context_fetch(stack, rb_ary_entry(path, 0));
    if (value == Qnil) {
        return Qnil;
    }
    for (long i = 1; i < RARRAY_LEN(path); i++) {
        value = fetch(value, rb_ary_entry(path, i));
        if (value == Qundef) {
            return Qnil;
        }
    }
    return value;
}

static void append_value(VALUE buf, VALUE stack, VALUE path, bool escape) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_NIL:
        case T_UNDEF:
            return;
        case T_STRING:
            break;
        default:
            value = rb_funcall(value, rb_intern("to_s"), 0);
            break;
    }
    rb_str_buf_append(buf, escape ? escape_html(value) : value);
}

static void section(VALUE buf, VALUE stack, VALUE path, void (*block)(VALUE, VALUE)) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_ARRAY:
            for (long i = 0; i < RARRAY_LEN(value); i++) {
                rb_ary_push(stack, rb_ary_entry(value, i));
                block(buf, stack);
                rb_ary_pop(stack);
            }
            break;
        case T_HASH:
        case T_OBJECT:
        case T_STRUCT:
            rb_ary_push(stack, value);
            block(buf, stack);
            rb_ary_pop(stack);
            break;
        case T_NIL:
        case T_UNDEF:
        case T_FALSE:
            break;
        default:
            block(buf, stack);
            break;
    }
}

static void inverted(VALUE buf, VALUE stack, VALUE path, void (*block)(VALUE, VALUE)) {
    VALUE value = fetch_path(stack, path);
    switch (rb_type(value)) {
        case T_ARRAY:
            if (RARRAY_LEN(value) == 0) {
                block(buf, stack);
            }
            break;
        case T_NIL:
        case T_UNDEF:
        case T_FALSE:
            block(buf, stack);
            break;
    }
}
"#;
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::convert::TryFrom;
use std::ffi::{OsStr, OsString};
use std::iter;
#[cfg(unix)]
use std::os::unix::ffi::{OsStrExt, OsStringExt};
#[cfg(windows)]
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::path::{Path, PathBuf};

use rustc::ty::layout::LayoutOf;

use crate::*;

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Helper function to read an OsString from a null-terminated sequence of bytes, which is what
    /// the Unix APIs usually handle.
    fn read_os_str_from_c_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, &'a OsStr>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        #[cfg(unix)]
        fn bytes_to_os_str<'tcx, 'a>(bytes: &'a [u8]) -> InterpResult<'tcx, &'a OsStr> {
            Ok(OsStr::from_bytes(bytes))
        }
        #[cfg(not(unix))]
        fn bytes_to_os_str<'tcx, 'a>(bytes: &'a [u8]) -> InterpResult<'tcx, &'a OsStr> {
            let s = std::str::from_utf8(bytes)
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", bytes))?;
            Ok(OsStr::new(s))
        }

        let this = self.eval_context_ref();
        let bytes = this.memory.read_c_str(scalar)?;
        bytes_to_os_str(bytes)
    }

    /// Helper function to read an OsString from a 0x0000-terminated sequence of u16,
    /// which is what the Windows APIs usually handle.
    fn read_os_str_from_wide_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, OsString>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        #[cfg(windows)]
        pub fn u16vec_to_osstring<'tcx, 'a>(u16_vec: Vec<u16>) -> InterpResult<'tcx, OsString> {
            Ok(OsString::from_wide(&u16_vec[..]))
        }
        #[cfg(not(windows))]
        pub fn u16vec_to_osstring<'tcx, 'a>(u16_vec: Vec<u16>) -> InterpResult<'tcx, OsString> {
            let s = String::from_utf16(&u16_vec[..])
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-16 string", u16_vec))?;
            Ok(s.into())
        }

        let u16_vec = self.eval_context_ref().memory.read_wide_str(scalar)?;
        u16vec_to_osstring(u16_vec)
    }

    /// Helper function to write an OsStr as a null-terminated sequence of bytes, which is what
    /// the Unix APIs usually handle. This function returns `Ok((false, length))` without trying
    /// to write if `size` is not large enough to fit the contents of `os_string` plus a null
    /// terminator. It returns `Ok((true, length))` if the writing process was successful. The
    /// string length returned does not include the null terminator.
    fn write_os_str_to_c_str(
        &mut self,
        os_str: &OsStr,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        #[cfg(unix)]
        fn os_str_to_bytes<'tcx, 'a>(os_str: &'a OsStr) -> InterpResult<'tcx, &'a [u8]> {
            Ok(os_str.as_bytes())
        }
        #[cfg(not(unix))]
        fn os_str_to_bytes<'tcx, 'a>(os_str: &'a OsStr) -> InterpResult<'tcx, &'a [u8]> {
            // On non-unix platforms the best we can do to transform bytes from/to OS strings is to do the
            // intermediate transformation into strings. Which invalidates non-utf8 paths that are actually
            // valid.
            os_str
                .to_str()
                .map(|s| s.as_bytes())
                .ok_or_else(|| err_unsup_format!("{:?} is not a valid utf-8 string", os_str).into())
        }

        let bytes = os_str_to_bytes(os_str)?;
        // If `size` is smaller or equal than `bytes.len()`, writing `bytes` plus the required null
        // terminator to memory using the `ptr` pointer would cause an out-of-bounds access.
        let string_length = u64::try_from(bytes.len()).unwrap();
        if size <= string_length {
            return Ok((false, string_length));
        }
        self.eval_context_mut()
            .memory
            .write_bytes(scalar, bytes.iter().copied().chain(iter::once(0u8)))?;
        Ok((true, string_length))
    }

    /// Helper function to write an OsStr as a 0x0000-terminated u16-sequence, which is what
    /// the Windows APIs usually handle. This function returns `Ok((false, length))` without trying
    /// to write if `size` is not large enough to fit the contents of `os_string` plus a null
    /// terminator. It returns `Ok((true, length))` if the writing process was successful. The
    /// string length returned does not include the null terminator.
    fn write_os_str_to_wide_str(
        &mut self,
        os_str: &OsStr,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        #[cfg(windows)]
        fn os_str_to_u16vec<'tcx>(os_str: &OsStr) -> InterpResult<'tcx, Vec<u16>> {
            Ok(os_str.encode_wide().collect())
        }
        #[cfg(not(windows))]
        fn os_str_to_u16vec<'tcx>(os_str: &OsStr) -> InterpResult<'tcx, Vec<u16>> {
            // On non-Windows platforms the best we can do to transform Vec<u16> from/to OS strings is to do the
            // intermediate transformation into strings. Which invalidates non-utf8 paths that are actually
            // valid.
            os_str
                .to_str()
                .map(|s| s.encode_utf16().collect())
                .ok_or_else(|| err_unsup_format!("{:?} is not a valid utf-8 string", os_str).into())
        }

        let u16_vec = os_str_to_u16vec(os_str)?;
        // If `size` is smaller or equal than `bytes.len()`, writing `bytes` plus the required
        // 0x0000 terminator to memory would cause an out-of-bounds access.
        let string_length = u64::try_from(u16_vec.len()).unwrap();
        if size <= string_length {
            return Ok((false, string_length));
        }

        // Store the UTF-16 string.
        self.eval_context_mut()
            .memory
            .write_u16s(scalar, u16_vec.into_iter().chain(iter::once(0x0000)))?;
        Ok((true, string_length))
    }

    /// Allocate enough memory to store the given `OsStr` as a null-terminated sequence of bytes.
    fn alloc_os_str_as_c_str(
        &mut self,
        os_str: &OsStr,
        memkind: MemoryKind<MiriMemoryKind>,
    ) -> Pointer<Tag> {
        let size = u64::try_from(os_str.len()).unwrap().checked_add(1).unwrap(); // Make space for `0` terminator.
        let this = self.eval_context_mut();

        let arg_type = this.tcx.mk_array(this.tcx.types.u8, size);
        let arg_place = this.allocate(this.layout_of(arg_type).unwrap(), memkind);
        assert!(self.write_os_str_to_c_str(os_str, arg_place.ptr, size).unwrap().0);
        arg_place.ptr.assert_ptr()
    }

    /// Allocate enough memory to store the given `OsStr` as a null-terminated sequence of `u16`.
    fn alloc_os_str_as_wide_str(
        &mut self,
        os_str: &OsStr,
        memkind: MemoryKind<MiriMemoryKind>,
    ) -> Pointer<Tag> {
        let size = u64::try_from(os_str.len()).unwrap().checked_add(1).unwrap(); // Make space for `0x0000` terminator.
        let this = self.eval_context_mut();

        let arg_type = this.tcx.mk_array(this.tcx.types.u16, size);
        let arg_place = this.allocate(this.layout_of(arg_type).unwrap(), memkind);
        assert!(self.write_os_str_to_wide_str(os_str, arg_place.ptr, size).unwrap().0);
        arg_place.ptr.assert_ptr()
    }

    /// Read a null-terminated sequence of bytes, and perform path separator conversion if needed.
    fn read_path_from_c_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, Cow<'a, Path>>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        let this = self.eval_context_ref();
        let os_str = this.read_os_str_from_c_str(scalar)?;

        #[cfg(windows)]
        return Ok(if this.tcx.sess.target.target.target_os == "windows" {
            // Windows-on-Windows, all fine.
            Cow::Borrowed(Path::new(os_str))
        } else {
            // Unix target, Windows host. Need to convert target '/' to host '\'.
            let converted = os_str
                .encode_wide()
                .map(|wchar| if wchar == '/' as u16 { '\\' as u16 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(PathBuf::from(OsString::from_wide(&converted)))
        });
        #[cfg(unix)]
        return Ok(if this.tcx.sess.target.target.target_os == "windows" {
            // Windows target, Unix host. Need to convert target '\' to host '/'.
            let converted = os_str
                .as_bytes()
                .iter()
                .map(|&wchar| if wchar == '/' as u8 { '\\' as u8 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(PathBuf::from(OsString::from_vec(converted)))
        } else {
            // Unix-on-Unix, all is fine.
            Cow::Borrowed(Path::new(os_str))
        });
    }

    /// Write a Path to the machine memory, adjusting path separators if needed.
    fn write_path_to_c_str(
        &mut self,
        path: &Path,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        let this = self.eval_context_mut();

        #[cfg(windows)]
        let os_str = if this.tcx.sess.target.target.target_os == "windows" {
            // Windows-on-Windows, all fine.
            Cow::Borrowed(path.as_os_str())
        } else {
            // Unix target, Windows host. Need to convert host '\\' to target '/'.
            let converted = path
                .as_os_str()
                .encode_wide()
                .map(|wchar| if wchar == '\\' as u16 { '/' as u16 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(OsString::from_wide(&converted))
        };
        #[cfg(unix)]
        let os_str = if this.tcx.sess.target.target.target_os == "windows" {
            // Windows target, Unix host. Need to convert host '/' to target '\'.
            let converted = path
                .as_os_str()
                .as_bytes()
                .iter()
                .map(|&wchar| if wchar == '/' as u8 { '\\' as u8 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(OsString::from_vec(converted))
        } else {
            // Unix-on-Unix, all is fine.
            Cow::Borrowed(path.as_os_str())
        };

        this.write_os_str_to_c_str(&os_str, scalar, size)
    }
}
<commit_msg>visually separate conditional imports in 'os_str.rs'<commit_after>use std::borrow::Cow;
use std::convert::TryFrom;
use std::ffi::{OsStr, OsString};
use std::iter;
use std::path::{Path, PathBuf};

#[cfg(unix)]
use std::os::unix::ffi::{OsStrExt, OsStringExt};
#[cfg(windows)]
use std::os::windows::ffi::{OsStrExt, OsStringExt};

use rustc::ty::layout::LayoutOf;

use crate::*;

impl<'mir, 'tcx> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx> {
    /// Helper function to read an OsString from a null-terminated sequence of bytes, which is what
    /// the Unix APIs usually handle.
    fn read_os_str_from_c_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, &'a OsStr>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        #[cfg(unix)]
        fn bytes_to_os_str<'tcx, 'a>(bytes: &'a [u8]) -> InterpResult<'tcx, &'a OsStr> {
            Ok(OsStr::from_bytes(bytes))
        }
        #[cfg(not(unix))]
        fn bytes_to_os_str<'tcx, 'a>(bytes: &'a [u8]) -> InterpResult<'tcx, &'a OsStr> {
            let s = std::str::from_utf8(bytes)
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-8 string", bytes))?;
            Ok(OsStr::new(s))
        }

        let this = self.eval_context_ref();
        let bytes = this.memory.read_c_str(scalar)?;
        bytes_to_os_str(bytes)
    }

    /// Helper function to read an OsString from a 0x0000-terminated sequence of u16,
    /// which is what the Windows APIs usually handle.
    fn read_os_str_from_wide_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, OsString>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        #[cfg(windows)]
        pub fn u16vec_to_osstring<'tcx, 'a>(u16_vec: Vec<u16>) -> InterpResult<'tcx, OsString> {
            Ok(OsString::from_wide(&u16_vec[..]))
        }
        #[cfg(not(windows))]
        pub fn u16vec_to_osstring<'tcx, 'a>(u16_vec: Vec<u16>) -> InterpResult<'tcx, OsString> {
            let s = String::from_utf16(&u16_vec[..])
                .map_err(|_| err_unsup_format!("{:?} is not a valid utf-16 string", u16_vec))?;
            Ok(s.into())
        }

        let u16_vec = self.eval_context_ref().memory.read_wide_str(scalar)?;
        u16vec_to_osstring(u16_vec)
    }

    /// Helper function to write an OsStr as a null-terminated sequence of bytes, which is what
    /// the Unix APIs usually handle. This function returns `Ok((false, length))` without trying
    /// to write if `size` is not large enough to fit the contents of `os_string` plus a null
    /// terminator. It returns `Ok((true, length))` if the writing process was successful. The
    /// string length returned does not include the null terminator.
    fn write_os_str_to_c_str(
        &mut self,
        os_str: &OsStr,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        #[cfg(unix)]
        fn os_str_to_bytes<'tcx, 'a>(os_str: &'a OsStr) -> InterpResult<'tcx, &'a [u8]> {
            Ok(os_str.as_bytes())
        }
        #[cfg(not(unix))]
        fn os_str_to_bytes<'tcx, 'a>(os_str: &'a OsStr) -> InterpResult<'tcx, &'a [u8]> {
            // On non-unix platforms the best we can do to transform bytes from/to OS strings is to do the
            // intermediate transformation into strings. Which invalidates non-utf8 paths that are actually
            // valid.
            os_str
                .to_str()
                .map(|s| s.as_bytes())
                .ok_or_else(|| err_unsup_format!("{:?} is not a valid utf-8 string", os_str).into())
        }

        let bytes = os_str_to_bytes(os_str)?;
        // If `size` is smaller or equal than `bytes.len()`, writing `bytes` plus the required null
        // terminator to memory using the `ptr` pointer would cause an out-of-bounds access.
        let string_length = u64::try_from(bytes.len()).unwrap();
        if size <= string_length {
            return Ok((false, string_length));
        }
        self.eval_context_mut()
            .memory
            .write_bytes(scalar, bytes.iter().copied().chain(iter::once(0u8)))?;
        Ok((true, string_length))
    }

    /// Helper function to write an OsStr as a 0x0000-terminated u16-sequence, which is what
    /// the Windows APIs usually handle. This function returns `Ok((false, length))` without trying
    /// to write if `size` is not large enough to fit the contents of `os_string` plus a null
    /// terminator. It returns `Ok((true, length))` if the writing process was successful. The
    /// string length returned does not include the null terminator.
    fn write_os_str_to_wide_str(
        &mut self,
        os_str: &OsStr,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        #[cfg(windows)]
        fn os_str_to_u16vec<'tcx>(os_str: &OsStr) -> InterpResult<'tcx, Vec<u16>> {
            Ok(os_str.encode_wide().collect())
        }
        #[cfg(not(windows))]
        fn os_str_to_u16vec<'tcx>(os_str: &OsStr) -> InterpResult<'tcx, Vec<u16>> {
            // On non-Windows platforms the best we can do to transform Vec<u16> from/to OS strings is to do the
            // intermediate transformation into strings. Which invalidates non-utf8 paths that are actually
            // valid.
            os_str
                .to_str()
                .map(|s| s.encode_utf16().collect())
                .ok_or_else(|| err_unsup_format!("{:?} is not a valid utf-8 string", os_str).into())
        }

        let u16_vec = os_str_to_u16vec(os_str)?;
        // If `size` is smaller or equal than `bytes.len()`, writing `bytes` plus the required
        // 0x0000 terminator to memory would cause an out-of-bounds access.
        let string_length = u64::try_from(u16_vec.len()).unwrap();
        if size <= string_length {
            return Ok((false, string_length));
        }

        // Store the UTF-16 string.
        self.eval_context_mut()
            .memory
            .write_u16s(scalar, u16_vec.into_iter().chain(iter::once(0x0000)))?;
        Ok((true, string_length))
    }

    /// Allocate enough memory to store the given `OsStr` as a null-terminated sequence of bytes.
    fn alloc_os_str_as_c_str(
        &mut self,
        os_str: &OsStr,
        memkind: MemoryKind<MiriMemoryKind>,
    ) -> Pointer<Tag> {
        let size = u64::try_from(os_str.len()).unwrap().checked_add(1).unwrap(); // Make space for `0` terminator.
        let this = self.eval_context_mut();

        let arg_type = this.tcx.mk_array(this.tcx.types.u8, size);
        let arg_place = this.allocate(this.layout_of(arg_type).unwrap(), memkind);
        assert!(self.write_os_str_to_c_str(os_str, arg_place.ptr, size).unwrap().0);
        arg_place.ptr.assert_ptr()
    }

    /// Allocate enough memory to store the given `OsStr` as a null-terminated sequence of `u16`.
    fn alloc_os_str_as_wide_str(
        &mut self,
        os_str: &OsStr,
        memkind: MemoryKind<MiriMemoryKind>,
    ) -> Pointer<Tag> {
        let size = u64::try_from(os_str.len()).unwrap().checked_add(1).unwrap(); // Make space for `0x0000` terminator.
        let this = self.eval_context_mut();

        let arg_type = this.tcx.mk_array(this.tcx.types.u16, size);
        let arg_place = this.allocate(this.layout_of(arg_type).unwrap(), memkind);
        assert!(self.write_os_str_to_wide_str(os_str, arg_place.ptr, size).unwrap().0);
        arg_place.ptr.assert_ptr()
    }

    /// Read a null-terminated sequence of bytes, and perform path separator conversion if needed.
    fn read_path_from_c_str<'a>(&'a self, scalar: Scalar<Tag>) -> InterpResult<'tcx, Cow<'a, Path>>
    where
        'tcx: 'a,
        'mir: 'a,
    {
        let this = self.eval_context_ref();
        let os_str = this.read_os_str_from_c_str(scalar)?;

        #[cfg(windows)]
        return Ok(if this.tcx.sess.target.target.target_os == "windows" {
            // Windows-on-Windows, all fine.
            Cow::Borrowed(Path::new(os_str))
        } else {
            // Unix target, Windows host. Need to convert target '/' to host '\'.
            let converted = os_str
                .encode_wide()
                .map(|wchar| if wchar == '/' as u16 { '\\' as u16 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(PathBuf::from(OsString::from_wide(&converted)))
        });
        #[cfg(unix)]
        return Ok(if this.tcx.sess.target.target.target_os == "windows" {
            // Windows target, Unix host. Need to convert target '\' to host '/'.
            let converted = os_str
                .as_bytes()
                .iter()
                .map(|&wchar| if wchar == '/' as u8 { '\\' as u8 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(PathBuf::from(OsString::from_vec(converted)))
        } else {
            // Unix-on-Unix, all is fine.
            Cow::Borrowed(Path::new(os_str))
        });
    }

    /// Write a Path to the machine memory, adjusting path separators if needed.
    fn write_path_to_c_str(
        &mut self,
        path: &Path,
        scalar: Scalar<Tag>,
        size: u64,
    ) -> InterpResult<'tcx, (bool, u64)> {
        let this = self.eval_context_mut();

        #[cfg(windows)]
        let os_str = if this.tcx.sess.target.target.target_os == "windows" {
            // Windows-on-Windows, all fine.
            Cow::Borrowed(path.as_os_str())
        } else {
            // Unix target, Windows host. Need to convert host '\\' to target '/'.
            let converted = path
                .as_os_str()
                .encode_wide()
                .map(|wchar| if wchar == '\\' as u16 { '/' as u16 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(OsString::from_wide(&converted))
        };
        #[cfg(unix)]
        let os_str = if this.tcx.sess.target.target.target_os == "windows" {
            // Windows target, Unix host. Need to convert host '/' to target '\'.
            let converted = path
                .as_os_str()
                .as_bytes()
                .iter()
                .map(|&wchar| if wchar == '/' as u8 { '\\' as u8 } else { wchar })
                .collect::<Vec<_>>();
            Cow::Owned(OsString::from_vec(converted))
        } else {
            // Unix-on-Unix, all is fine.
            Cow::Borrowed(path.as_os_str())
        };

        this.write_os_str_to_c_str(&os_str, scalar, size)
    }
}
<|endoftext|>


<commit_before>//! Treiber stacks.

use std::sync::atomic;
use std::{mem, ptr};
use {Atomic, Guard};

/// A Treiber stack.
pub struct Treiber<T> {
    head: Atomic<Node<T>>,
}

struct Node<T> {
    data: T,
    next: *const Node<T>,
}

impl<T> Treiber<T> {
    /// Create a new, empty Treiber stack.
    pub fn new() -> Treiber<T> {
        Treiber {
            head: Atomic::default(),
        }
    }

    /// Pop an item from the stack.
    pub fn pop(&self) -> Option<Guard<T>> {
        // TODO: Use `catch {}` here when it lands.
        // Read the head snapshot.
        let mut snapshot = self.head.load(atomic::Ordering::Acquire);

        // Unless the head snapshot is `None`, try to replace it with the tail.
        while let Some(node) = snapshot {
            // Attempt to replace the head with the tail of the head.
            match unsafe {
                self.head.compare_and_swap_raw(
                    &*node,
                    node.next as *mut Node<T>,
                    atomic::Ordering::Release,
                )
            } {
                // It succeeded; return the item.
                Ok(_) => return Some(node.map(|x| &x.data)),
                // It failed, update the head snapshot and continue.
                Err(new) => snapshot = new,
            }
        }

        // As the head was empty, there is nothing to pop.
        None
    }

    /// Push an item to the stack.
    pub fn push(&self, item: T)
    where T: 'static {
        // Load the head snapshot.
        let mut snapshot = self.head.load(atomic::Ordering::Acquire);
        let mut snapshot_ptr: Option<*const Node<T>>;

        // TODO: Use `catch {}` here when it lands.
        // Construct a node, which will be the new head.
        let mut node = Box::new(Node {
            data: item,
            // Placeholder; we will replace it with an actual value in the loop.
            next: ptr::null(),
        });

        loop {
            // Derive the nullable snapshot pointer from the head snapshot.
            snapshot_ptr = snapshot.as_ref().map(Guard::as_raw);
            // Construct the next-pointer of the new node from the head snapshot.
            node.next = snapshot_ptr.unwrap_or(ptr::null());

            // TODO: This should be something that ignores the guard creation when the CAS
            //       succeeds, because it's expensive to do and not used anyway. It should be easy
            //       enough to implement, but I am struggling to come up with a good name for the
            //       method.

            // CAS from the read pointer to the new head.
            match self.head.compare_and_swap(snapshot_ptr, Some(node), atomic::Ordering::Release) {
                // If it succeeds, the item has been pushed.
                Ok(_) => break,
                // If it fails, we will retry the CAS with updated values.
                Err((new_head, Some(node2))) => {
                    // Update the head snapshot.
                    snapshot = new_head;
                    // Put the box we gave back to the variable where it belongs.
                    node = node2;
                },
                // This should never be reached as we gave an argument which was unconditionally
                // `Some`.
                _ => unreachable!(),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::sync::Arc;

    #[test]
    fn single_thread() {
        let stack = Treiber::new();

        for i in 0..10000 {
            stack.push(i);
        }

        for i in (0..10000).rev() {
            assert_eq!(*stack.pop().unwrap(), i);
        }

        for i in 0..10000 {
            stack.push(i);
        }

        for i in (0..10000).rev() {
            assert_eq!(*stack.pop().unwrap(), i);
        }

        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
    }

    #[test]
    fn push_pop() {
        let stack = Arc::new(Treiber::new());
        let mut j = Vec::new();
        for _ in 0..16 {
            let s = stack.clone();
            j.push(thread::spawn(move || {
                for _ in 0..10_000_000 {
                    s.push(23);
                    assert_eq!(*s.pop().unwrap(), 23);
                }
            }));
        }

        for i in j {
            i.join().unwrap();
        }
    }
}
<commit_msg>Extend documentation for `conc::sync::Treiber`.<commit_after>//! Treiber stacks.

use std::sync::atomic;
use std::{mem, ptr};
use {Atomic, Guard};

/// A Treiber stack.
///
/// Treiber stacks are one way to implement concurrent LIFO stack.
///
/// Treiber stacks builds on linked lists. They are lock-free and non-blocking. It can be compared
/// to transactional memory in that it repeats operations, if another thread changes it while.
///
/// The ABA problem is of course addressed through the API of this crate.
pub struct Treiber<T> {
    /// The head node.
    head: Atomic<Node<T>>,
}

/// A node in the stack.
struct Node<T> {
    /// The data this node holds.
    data: T,
    /// The next node.
    next: *const Node<T>,
}

impl<T> Treiber<T> {
    /// Create a new, empty Treiber stack.
    pub fn new() -> Treiber<T> {
        Treiber {
            head: Atomic::default(),
        }
    }

    /// Pop an item from the stack.
    pub fn pop(&self) -> Option<Guard<T>> {
        // TODO: Use `catch {}` here when it lands.
        // Read the head snapshot.
        let mut snapshot = self.head.load(atomic::Ordering::Acquire);

        // Unless the head snapshot is `None`, try to replace it with the tail.
        while let Some(node) = snapshot {
            // Attempt to replace the head with the tail of the head.
            match unsafe {
                self.head.compare_and_swap_raw(
                    &*node,
                    node.next as *mut Node<T>,
                    atomic::Ordering::Release,
                )
            } {
                // It succeeded; return the item.
                Ok(_) => return Some(node.map(|x| &x.data)),
                // It failed, update the head snapshot and continue.
                Err(new) => snapshot = new,
            }
        }

        // As the head was empty, there is nothing to pop.
        None
    }

    /// Push an item to the stack.
    pub fn push(&self, item: T)
    where T: 'static {
        // Load the head snapshot.
        let mut snapshot = self.head.load(atomic::Ordering::Acquire);
        let mut snapshot_ptr: Option<*const Node<T>>;

        // TODO: Use `catch {}` here when it lands.
        // Construct a node, which will be the new head.
        let mut node = Box::new(Node {
            data: item,
            // Placeholder; we will replace it with an actual value in the loop.
            next: ptr::null(),
        });

        loop {
            // Derive the nullable snapshot pointer from the head snapshot.
            snapshot_ptr = snapshot.as_ref().map(Guard::as_raw);
            // Construct the next-pointer of the new node from the head snapshot.
            node.next = snapshot_ptr.unwrap_or(ptr::null());

            // TODO: This should be something that ignores the guard creation when the CAS
            //       succeeds, because it's expensive to do and not used anyway. It should be easy
            //       enough to implement, but I am struggling to come up with a good name for the
            //       method.

            // CAS from the read pointer to the new head.
            match self.head.compare_and_swap(snapshot_ptr, Some(node), atomic::Ordering::Release) {
                // If it succeeds, the item has been pushed.
                Ok(_) => break,
                // If it fails, we will retry the CAS with updated values.
                Err((new_head, Some(node2))) => {
                    // Update the head snapshot.
                    snapshot = new_head;
                    // Put the box we gave back to the variable where it belongs.
                    node = node2;
                },
                // This should never be reached as we gave an argument which was unconditionally
                // `Some`.
                _ => unreachable!(),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::sync::Arc;

    #[test]
    fn single_thread() {
        let stack = Treiber::new();

        for i in 0..10000 {
            stack.push(i);
        }

        for i in (0..10000).rev() {
            assert_eq!(*stack.pop().unwrap(), i);
        }

        for i in 0..10000 {
            stack.push(i);
        }

        for i in (0..10000).rev() {
            assert_eq!(*stack.pop().unwrap(), i);
        }

        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
        assert!(stack.pop().is_none());
    }

    #[test]
    fn push_pop() {
        let stack = Arc::new(Treiber::new());
        let mut j = Vec::new();
        for _ in 0..16 {
            let s = stack.clone();
            j.push(thread::spawn(move || {
                for _ in 0..10_000_000 {
                    s.push(23);
                    assert_eq!(*s.pop().unwrap(), 23);
                }
            }));
        }

        for i in j {
            i.join().unwrap();
        }
    }
}
<|endoftext|>