<commit_before>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        let mut pos_index = 0;
        let mut is_escaped = false;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir, pos_index, is_escaped);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                    pos_index = 0;
                    continue;
                }
            }

            if arg.is_escape() {
                is_escaped = true;
            } else {
                pos_index += 1;
            }
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
        pos_index: usize,
        is_escaped: bool,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        if !is_escaped {
            completions.extend(
                crate::generator::utils::longs_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("--{}", f).into()),
            );
            completions.extend(
                crate::generator::utils::shorts_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("-{}", f).into()),
            );
        }

        let mut positionals = Vec::new();
        if let Some(positional) = cmd
            .get_positionals()
            .find(|p| p.get_index() == Some(pos_index))
        {
            if let Some(possible_values) = positional.get_possible_values() {
                positionals.extend(possible_values.into_iter().map(|p| p.get_name().into()));
            } else {
                match positional.get_value_hint() {
                    clap::ValueHint::Other => {
                        // Should not complete
                    }
                    clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                    clap::ValueHint::FilePath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_file()));
                    }
                    clap::ValueHint::DirPath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_dir()));
                    }
                    clap::ValueHint::ExecutablePath => {
                        use is_executable::IsExecutable;
                        positionals.extend(complete_path(current_dir, |p| p.is_executable()));
                    }
                    clap::ValueHint::CommandName
                    | clap::ValueHint::CommandString
                    | clap::ValueHint::CommandWithArguments
                    | clap::ValueHint::Username
                    | clap::ValueHint::Hostname
                    | clap::ValueHint::Url
                    | clap::ValueHint::EmailAddress => {
                        // No completion implementation
                    }
                    _ => {
                        // Safe-ish fallback
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                }
                positionals.sort();
            }
            completions.extend(positionals);
        }

        if !is_escaped {
            completions.extend(all_subcommands(cmd).into_iter().map(|s| s));
        }

        Ok(completions)
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<commit_msg>refactor(complete): Extract arg value completion<commit_after>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        let mut pos_index = 0;
        let mut is_escaped = false;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir, pos_index, is_escaped);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                    pos_index = 0;
                    continue;
                }
            }

            if arg.is_escape() {
                is_escaped = true;
            } else {
                pos_index += 1;
            }
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
        pos_index: usize,
        is_escaped: bool,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        if !is_escaped {
            completions.extend(
                crate::generator::utils::longs_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("--{}", f).into()),
            );
            completions.extend(
                crate::generator::utils::shorts_and_visible_aliases(cmd)
                    .into_iter()
                    .map(|f| format!("-{}", f).into()),
            );
        }

        if let Some(positional) = cmd
            .get_positionals()
            .find(|p| p.get_index() == Some(pos_index))
        {
            completions.extend(complete_arg_value(positional, current_dir));
        }

        if !is_escaped {
            completions.extend(all_subcommands(cmd).into_iter().map(|s| s));
        }

        Ok(completions)
    }

    fn complete_arg_value(
        arg: &clap::Arg<'_>,
        current_dir: Option<&std::path::Path>,
    ) -> Vec<OsString> {
        let mut values = Vec::new();

        if let Some(possible_values) = arg.get_possible_values() {
            values.extend(possible_values.into_iter().map(|p| p.get_name().into()));
        } else {
            match arg.get_value_hint() {
                clap::ValueHint::Other => {
                    // Should not complete
                }
                clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
                    values.extend(complete_path(current_dir, |_| true));
                }
                clap::ValueHint::FilePath => {
                    values.extend(complete_path(current_dir, |p| p.is_file()));
                }
                clap::ValueHint::DirPath => {
                    values.extend(complete_path(current_dir, |p| p.is_dir()));
                }
                clap::ValueHint::ExecutablePath => {
                    use is_executable::IsExecutable;
                    values.extend(complete_path(current_dir, |p| p.is_executable()));
                }
                clap::ValueHint::CommandName
                | clap::ValueHint::CommandString
                | clap::ValueHint::CommandWithArguments
                | clap::ValueHint::Username
                | clap::ValueHint::Hostname
                | clap::ValueHint::Url
                | clap::ValueHint::EmailAddress => {
                    // No completion implementation
                }
                _ => {
                    // Safe-ish fallback
                    values.extend(complete_path(current_dir, |_| true));
                }
            }
            values.sort();
        }

        values
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::*;
use crate::mc::MotionVector;
use crate::partition::*;
use crate::predict::PredictionMode;
use crate::transform::*;

use std::cmp;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
use std::slice;

/// Tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocks<'a> {
  data: *const Block,
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  frame_cols: usize,
  frame_rows: usize,
  phantom: PhantomData<&'a Block>,
}

/// Mutable tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocksMut<'a> {
  data: *mut Block,
  // private to guarantee borrowing rules
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  frame_cols: usize,
  frame_rows: usize,
  phantom: PhantomData<&'a mut Block>,
}

// common impl for TileBlocks and TileBlocksMut
macro_rules! tile_blocks_common {
  // $name: TileBlocks or TileBlocksMut
  // $opt_mut: nothing or mut
  ($name:ident $(,$opt_mut:tt)?) => {
    impl<'a> $name<'a> {

      #[inline(always)]
      pub fn new(
        frame_blocks: &'a $($opt_mut)? FrameBlocks,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> Self {
        Self {
          data: & $($opt_mut)? frame_blocks[y][x],
          x,
          y,
          cols,
          rows,
          frame_cols: frame_blocks.cols,
          frame_rows: frame_blocks.rows,
          phantom: PhantomData,
        }
      }

      pub fn subregion(
        &mut self,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> TileBlocks<'_> {
        TileBlocks {
          data: &self[y][x],
          x: self.x+x,
          y: self.y+y,
          cols: cmp::min(cols, self.cols - x),
          rows: cmp::min(rows, self.rows - y),
          frame_cols: self.frame_cols,
          frame_rows: self.frame_rows,
          phantom: PhantomData,
        }
      }

      #[inline(always)]
      pub const fn x(&self) -> usize {
        self.x
      }

      #[inline(always)]
      pub const fn y(&self) -> usize {
        self.y
      }

      #[inline(always)]
      pub const fn cols(&self) -> usize {
        self.cols
      }

      #[inline(always)]
      pub const fn rows(&self) -> usize {
        self.rows
      }

      #[inline(always)]
      pub const fn frame_cols(&self) -> usize {
        self.frame_cols
      }

      #[inline(always)]
      pub const fn frame_rows(&self) -> usize {
        self.frame_rows
      }

      #[inline(always)]
      pub fn above_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y - 1][bo.0.x]
      }

      #[inline(always)]
      pub fn left_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y][bo.0.x - 1]
      }

      #[inline(always)]
      pub fn above_left_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y - 1][bo.0.x - 1]
      }

      pub fn get_cdef(&self, sbo: TileSuperBlockOffset) -> u8 {
        let bo = sbo.block_offset(0, 0).0;
        self[bo.y][bo.x].cdef_index
      }
    }

    unsafe impl Send for $name<'_> {}
    unsafe impl Sync for $name<'_> {}

    impl Index<usize> for $name<'_> {
      type Output = [Block];
      #[inline(always)]
      fn index(&self, index: usize) -> &Self::Output {
        assert!(index < self.rows);
        unsafe {
          let ptr = self.data.add(index * self.frame_cols);
          slice::from_raw_parts(ptr, self.cols)
        }
      }
    }

    // for convenience, also index by TileBlockOffset
    impl Index<TileBlockOffset> for $name<'_> {
      type Output = Block;
      #[inline(always)]
      fn index(&self, bo: TileBlockOffset) -> &Self::Output {
        &self[bo.0.y][bo.0.x]
      }
    }
  }
}

tile_blocks_common!(TileBlocks);
tile_blocks_common!(TileBlocksMut, mut);

impl TileBlocksMut<'_> {
  #[inline(always)]
  pub const fn as_const(&self) -> TileBlocks<'_> {
    TileBlocks {
      data: self.data,
      x: self.x,
      y: self.y,
      cols: self.cols,
      rows: self.rows,
      frame_cols: self.frame_cols,
      frame_rows: self.frame_rows,
      phantom: PhantomData,
    }
  }

  pub fn subregion_mut(
    &mut self, x: usize, y: usize, cols: usize, rows: usize,
  ) -> TileBlocksMut<'_> {
    TileBlocksMut {
      data: &mut self[y][x],
      x: self.x + x,
      y: self.y + y,
      cols: cmp::min(cols, self.cols - x),
      rows: cmp::min(rows, self.rows - y),
      frame_cols: self.frame_cols,
      frame_rows: self.frame_rows,
      phantom: PhantomData,
    }
  }

  #[inline(always)]
  pub fn for_each<F>(&mut self, bo: TileBlockOffset, bsize: BlockSize, f: F)
  where
    F: Fn(&mut Block) -> (),
  {
    let bw = bsize.width_mi();
    let bh = bsize.height_mi();
    for y in 0..bh {
      for x in 0..bw {
        f(&mut self[bo.0.y + y as usize][bo.0.x + x as usize]);
      }
    }
  }

  #[inline(always)]
  pub fn set_mode(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, mode: PredictionMode,
  ) {
    self.for_each(bo, bsize, |block| block.mode = mode);
  }

  #[inline(always)]
  pub fn set_block_size(&mut self, bo: TileBlockOffset, bsize: BlockSize) {
    let n4_w = bsize.width_mi();
    let n4_h = bsize.height_mi();
    self.for_each(bo, bsize, |block| {
      block.bsize = bsize;
      block.n4_w = n4_w;
      block.n4_h = n4_h
    });
  }

  #[inline(always)]
  pub fn set_tx_size(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize,
  ) {
    self.for_each(bo, bsize, |block| block.txsize = tx_size);
  }

  #[inline(always)]
  pub fn set_skip(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, skip: bool,
  ) {
    self.for_each(bo, bsize, |block| block.skip = skip);
  }

  #[inline(always)]
  pub fn set_segmentation_idx(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, idx: u8,
  ) {
    self.for_each(bo, bsize, |block| block.segmentation_idx = idx);
  }

  #[inline(always)]
  pub fn set_ref_frames(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, r: [RefType; 2],
  ) {
    self.for_each(bo, bsize, |block| block.ref_frames = r);
  }

  #[inline(always)]
  pub fn set_motion_vectors(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, mvs: [MotionVector; 2],
  ) {
    self.for_each(bo, bsize, |block| block.mv = mvs);
  }

  #[inline(always)]
  pub fn set_cdef(&mut self, sbo: TileSuperBlockOffset, cdef_index: u8) {
    let bo = sbo.block_offset(0, 0).0;
    // Checkme: Is 16 still the right block unit for 128x128 superblocks?
    let bw = cmp::min(bo.x + MIB_SIZE, self.cols);
    let bh = cmp::min(bo.y + MIB_SIZE, self.rows);
    for y in bo.y..bh {
      for x in bo.x..bw {
        self[y as usize][x as usize].cdef_index = cdef_index;
      }
    }
  }
}

impl IndexMut<usize> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
    assert!(index < self.rows);
    unsafe {
      let ptr = self.data.add(index * self.frame_cols);
      slice::from_raw_parts_mut(ptr, self.cols)
    }
  }
}

impl IndexMut<TileBlockOffset> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, bo: TileBlockOffset) -> &mut Self::Output {
    &mut self[bo.0.y][bo.0.x]
  }
}
<commit_msg>Reduce manual indexing in for_each in TileBlocks<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::context::*;
use crate::mc::MotionVector;
use crate::partition::*;
use crate::predict::PredictionMode;
use crate::transform::*;

use std::cmp;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
use std::slice;

/// Tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocks<'a> {
  data: *const Block,
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  frame_cols: usize,
  frame_rows: usize,
  phantom: PhantomData<&'a Block>,
}

/// Mutable tiled view of FrameBlocks
#[derive(Debug)]
pub struct TileBlocksMut<'a> {
  data: *mut Block,
  // private to guarantee borrowing rules
  x: usize,
  y: usize,
  cols: usize,
  rows: usize,
  frame_cols: usize,
  frame_rows: usize,
  phantom: PhantomData<&'a mut Block>,
}

// common impl for TileBlocks and TileBlocksMut
macro_rules! tile_blocks_common {
  // $name: TileBlocks or TileBlocksMut
  // $opt_mut: nothing or mut
  ($name:ident $(,$opt_mut:tt)?) => {
    impl<'a> $name<'a> {

      #[inline(always)]
      pub fn new(
        frame_blocks: &'a $($opt_mut)? FrameBlocks,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> Self {
        Self {
          data: & $($opt_mut)? frame_blocks[y][x],
          x,
          y,
          cols,
          rows,
          frame_cols: frame_blocks.cols,
          frame_rows: frame_blocks.rows,
          phantom: PhantomData,
        }
      }

      pub fn subregion(
        &mut self,
        x: usize,
        y: usize,
        cols: usize,
        rows: usize,
      ) -> TileBlocks<'_> {
        TileBlocks {
          data: &self[y][x],
          x: self.x+x,
          y: self.y+y,
          cols: cmp::min(cols, self.cols - x),
          rows: cmp::min(rows, self.rows - y),
          frame_cols: self.frame_cols,
          frame_rows: self.frame_rows,
          phantom: PhantomData,
        }
      }

      #[inline(always)]
      pub const fn x(&self) -> usize {
        self.x
      }

      #[inline(always)]
      pub const fn y(&self) -> usize {
        self.y
      }

      #[inline(always)]
      pub const fn cols(&self) -> usize {
        self.cols
      }

      #[inline(always)]
      pub const fn rows(&self) -> usize {
        self.rows
      }

      #[inline(always)]
      pub const fn frame_cols(&self) -> usize {
        self.frame_cols
      }

      #[inline(always)]
      pub const fn frame_rows(&self) -> usize {
        self.frame_rows
      }

      #[inline(always)]
      pub fn above_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y - 1][bo.0.x]
      }

      #[inline(always)]
      pub fn left_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y][bo.0.x - 1]
      }

      #[inline(always)]
      pub fn above_left_of(&self, bo: TileBlockOffset) -> &Block {
        &self[bo.0.y - 1][bo.0.x - 1]
      }

      pub fn get_cdef(&self, sbo: TileSuperBlockOffset) -> u8 {
        let bo = sbo.block_offset(0, 0).0;
        self[bo.y][bo.x].cdef_index
      }
    }

    unsafe impl Send for $name<'_> {}
    unsafe impl Sync for $name<'_> {}

    impl Index<usize> for $name<'_> {
      type Output = [Block];
      #[inline(always)]
      fn index(&self, index: usize) -> &Self::Output {
        assert!(index < self.rows);
        unsafe {
          let ptr = self.data.add(index * self.frame_cols);
          slice::from_raw_parts(ptr, self.cols)
        }
      }
    }

    // for convenience, also index by TileBlockOffset
    impl Index<TileBlockOffset> for $name<'_> {
      type Output = Block;
      #[inline(always)]
      fn index(&self, bo: TileBlockOffset) -> &Self::Output {
        &self[bo.0.y][bo.0.x]
      }
    }
  }
}

tile_blocks_common!(TileBlocks);
tile_blocks_common!(TileBlocksMut, mut);

impl TileBlocksMut<'_> {
  #[inline(always)]
  pub const fn as_const(&self) -> TileBlocks<'_> {
    TileBlocks {
      data: self.data,
      x: self.x,
      y: self.y,
      cols: self.cols,
      rows: self.rows,
      frame_cols: self.frame_cols,
      frame_rows: self.frame_rows,
      phantom: PhantomData,
    }
  }

  pub fn subregion_mut(
    &mut self, x: usize, y: usize, cols: usize, rows: usize,
  ) -> TileBlocksMut<'_> {
    TileBlocksMut {
      data: &mut self[y][x],
      x: self.x + x,
      y: self.y + y,
      cols: cmp::min(cols, self.cols - x),
      rows: cmp::min(rows, self.rows - y),
      frame_cols: self.frame_cols,
      frame_rows: self.frame_rows,
      phantom: PhantomData,
    }
  }

  #[inline(always)]
  pub fn for_each<F>(&mut self, bo: TileBlockOffset, bsize: BlockSize, f: F)
  where
    F: Fn(&mut Block) -> (),
  {
    let bw = bsize.width_mi();
    let bh = bsize.height_mi();
    for y in 0..bh {
      for block in self[bo.0.y + y][bo.0.x..bo.0.x + bw].iter_mut() {
        f(block);
      }
    }
  }

  #[inline(always)]
  pub fn set_mode(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, mode: PredictionMode,
  ) {
    self.for_each(bo, bsize, |block| block.mode = mode);
  }

  #[inline(always)]
  pub fn set_block_size(&mut self, bo: TileBlockOffset, bsize: BlockSize) {
    let n4_w = bsize.width_mi();
    let n4_h = bsize.height_mi();
    self.for_each(bo, bsize, |block| {
      block.bsize = bsize;
      block.n4_w = n4_w;
      block.n4_h = n4_h
    });
  }

  #[inline(always)]
  pub fn set_tx_size(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, tx_size: TxSize,
  ) {
    self.for_each(bo, bsize, |block| block.txsize = tx_size);
  }

  #[inline(always)]
  pub fn set_skip(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, skip: bool,
  ) {
    self.for_each(bo, bsize, |block| block.skip = skip);
  }

  #[inline(always)]
  pub fn set_segmentation_idx(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, idx: u8,
  ) {
    self.for_each(bo, bsize, |block| block.segmentation_idx = idx);
  }

  #[inline(always)]
  pub fn set_ref_frames(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, r: [RefType; 2],
  ) {
    self.for_each(bo, bsize, |block| block.ref_frames = r);
  }

  #[inline(always)]
  pub fn set_motion_vectors(
    &mut self, bo: TileBlockOffset, bsize: BlockSize, mvs: [MotionVector; 2],
  ) {
    self.for_each(bo, bsize, |block| block.mv = mvs);
  }

  #[inline(always)]
  pub fn set_cdef(&mut self, sbo: TileSuperBlockOffset, cdef_index: u8) {
    let bo = sbo.block_offset(0, 0).0;
    // Checkme: Is 16 still the right block unit for 128x128 superblocks?
    let bw = cmp::min(bo.x + MIB_SIZE, self.cols);
    let bh = cmp::min(bo.y + MIB_SIZE, self.rows);
    for y in bo.y..bh {
      for x in bo.x..bw {
        self[y as usize][x as usize].cdef_index = cdef_index;
      }
    }
  }
}

impl IndexMut<usize> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, index: usize) -> &mut Self::Output {
    assert!(index < self.rows);
    unsafe {
      let ptr = self.data.add(index * self.frame_cols);
      slice::from_raw_parts_mut(ptr, self.cols)
    }
  }
}

impl IndexMut<TileBlockOffset> for TileBlocksMut<'_> {
  #[inline(always)]
  fn index_mut(&mut self, bo: TileBlockOffset) -> &mut Self::Output {
    &mut self[bo.0.y][bo.0.x]
  }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use std::cmp;
use torrent::{Info, Peer, Bitfield};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    info: Arc<Info>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 5;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: Arc<Info>, pieces: &Bitfield) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = sequential::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        Picker {
            picker: PickerKind::Sequential(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading: HashMap::new(),
            info,
        }
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        piece.and_then(|p| self.pick_piece(p, peer.id())).or_else(
            || {
                self.pick_downloading(peer)
            },
        )
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize ||
            (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick an already requested block
    fn pick_downloading<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in &mut self.downloading {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS &&
                            r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut().find(|r| r.offset == b.offset).map(
                    |r| r.complete(),
                )
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps)) &&
                    r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter().find(|r| r.offset == b.offset).map(
                    |r| r.completed,
                )
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        }
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.add_peer(peer);
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.remove_peer(peer);
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool) {
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        for i in self.downloading.keys() {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.dec_avail(*i);
            }
        }
    }

    pub fn refresh_picker(&mut self, pieces: &Bitfield, pri: &[u8]) {
        // Map piece -> priority
        let mut piece_map = HashMap::new();
        // If a piece is completely in a file, just assign that pri.
        // Otherwise mark it as the higher pri piece
        for p in 0..self.info.pieces() {
            let locs = self.info.piece_disk_locs(p);
            let mp = locs.iter().fold(0, |mp, loc| {
                cmp::max(mp, pri[self.info.file_idx[&loc.file]] as usize)
            });
            piece_map.insert(p, mp);
        }

        self.unpicked = pieces.clone();
        self.picker = if self.is_sequential() {
            let mut picker = rarest::Picker::new(&self.unpicked);
            for (piece, pri) in piece_map {
                for _ in 0..pri {
                    picker.piece_available(piece);
                }
            }
            PickerKind::Rarest(picker)
        } else {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in piece_map {
                pieces[pri].push(piece);
            }
            PickerKind::Sequential(sequential::Picker::new_pri(pieces))
        };
    }
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(Arc::new(info.clone()), pieces)
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(Arc::new(info.clone()), pieces);
        p.change_picker(true);
        p
    }
}
impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<commit_msg>Use rarest picker by default<commit_after>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use std::cmp;
use torrent::{Info, Peer, Bitfield};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    info: Arc<Info>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 5;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: Arc<Info>, pieces: &Bitfield) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = rarest::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        Picker {
            picker: PickerKind::Rarest(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading: HashMap::new(),
            info,
        }
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        piece.and_then(|p| self.pick_piece(p, peer.id())).or_else(
            || {
                self.pick_downloading(peer)
            },
        )
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize ||
            (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick an already requested block
    fn pick_downloading<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in &mut self.downloading {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS &&
                            r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut().find(|r| r.offset == b.offset).map(
                    |r| r.complete(),
                )
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps)) &&
                    r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter().find(|r| r.offset == b.offset).map(
                    |r| r.completed,
                )
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        }
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.add_peer(peer);
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.remove_peer(peer);
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool) {
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        for i in self.downloading.keys() {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.dec_avail(*i);
            }
        }
    }

    pub fn refresh_picker(&mut self, pieces: &Bitfield, pri: &[u8]) {
        // Map piece -> priority
        let mut piece_map = HashMap::new();
        // If a piece is completely in a file, just assign that pri.
        // Otherwise mark it as the higher pri piece
        for p in 0..self.info.pieces() {
            let locs = self.info.piece_disk_locs(p);
            let mp = locs.iter().fold(0, |mp, loc| {
                cmp::max(mp, pri[self.info.file_idx[&loc.file]] as usize)
            });
            piece_map.insert(p, mp);
        }

        self.unpicked = pieces.clone();
        self.picker = if self.is_sequential() {
            let mut picker = rarest::Picker::new(&self.unpicked);
            for (piece, pri) in piece_map {
                for _ in 0..pri {
                    picker.piece_available(piece);
                }
            }
            PickerKind::Rarest(picker)
        } else {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in piece_map {
                pieces[pri].push(piece);
            }
            PickerKind::Sequential(sequential::Picker::new_pri(pieces))
        };
    }
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(Arc::new(info.clone()), pieces)
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(Arc::new(info.clone()), pieces);
        p.change_picker(true);
        p
    }
}
impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<|endoftext|>


<commit_before>pub trait ToHomogeneous<U>
{
  fn to_homogeneous(&self) -> U;
}

pub trait FromHomogeneous<U>
{
  fn from_homogeneous(&U) -> Self;
}
<commit_msg>from_homogeneous -> from<commit_after>pub trait ToHomogeneous<U>
{
  fn to_homogeneous(&self) -> U;
}

pub trait FromHomogeneous<U>
{
  fn from(&U) -> Self;
}
<|endoftext|>


<commit_before>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

use ::types::Vector3D;

/// The Particle type hold basic data about a particle in the system. It is self
/// contained, so that it will be easy to send data between parrallels
/// processes.
#[derive(Clone, Debug)]
pub struct Particle {
    /// Particle name
    name: String,
    /// Particle mass
    mass: f64,
    /// Particle positions
    position: Vector3D,
    /// Particle velocity, if needed
    velocity: Vector3D,
}

impl Particle {
    pub fn new<S: Into<String>>(name: S) -> Particle {
        // TODO: get the mass here
        Particle{name: name.into(),
                 mass: 1.0,
                 position: Vector3D::new(0.0, 0.0, 0.0),
                 velocity: Vector3D::new(0.0, 0.0, 0.0)}
    }

    pub fn name<'a>(&'a self) -> &'a str {
        &self.name
    }

    pub fn set_name<S>(&mut self, name: S) where S: Into<String> {
        self.name = name.into();
    }

    pub fn mass(&self) -> f64 {
        self.mass
    }
    pub fn set_mass(&mut self, mass: f64) {
        self.mass = mass;
    }

    pub fn position<'a>(&'a self) -> &'a Vector3D {
        &self.position
    }
    pub fn set_position(&mut self, pos: Vector3D) {
        self.position = pos;
    }
    pub fn add_position(&mut self, new_pos: Vector3D) {
        self.position = self.position + new_pos;
    }

    pub fn velocity<'a>(&'a self) -> &'a Vector3D {
        &self.velocity
    }
    pub fn set_velocity(&mut self, vel: Vector3D) {
        self.velocity = vel;
    }
    pub fn add_velocity(&mut self, new_vel: Vector3D) {
        self.velocity = self.velocity + new_vel;
    }
}

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

    #[test]
    fn set_mass() {
        let mut part = Particle::new("O");
        assert_eq!(part.name, "O");

        part.set_mass(16.0);
        assert_eq!(part.mass(), 16.0);
    }

    #[test]
    fn set_coords() {
        let mut part = Particle::new("O");
        assert_eq!(part.position(), &Vector3D::new(0.0, 0.0, 0.0));

        part.set_position(Vector3D::new(1.0, 2.0, 3.0));
        assert_eq!(part.position(), &Vector3D::new(1.0, 2.0, 3.0));

        part.set_velocity(Vector3D::new(1.0, 2.0, 3.0));
        assert_eq!(part.velocity(), &Vector3D::new(1.0, 2.0, 3.0));
    }
}
<commit_msg>Refactor tests for Particle<commit_after>/*
 * Cymbalum, Molecular Simulation in Rust
 * Copyright (C) 2015 Guillaume Fraux
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/
*/

use ::types::Vector3D;

/// The Particle type hold basic data about a particle in the system. It is self
/// contained, so that it will be easy to send data between parrallels
/// processes.
#[derive(Clone, Debug)]
pub struct Particle {
    /// Particle name
    name: String,
    /// Particle mass
    mass: f64,
    /// Particle positions
    position: Vector3D,
    /// Particle velocity, if needed
    velocity: Vector3D,
}

impl Particle {
    pub fn new<S: Into<String>>(name: S) -> Particle {
        // TODO: get the mass here
        Particle{name: name.into(),
                 mass: 1.0,
                 position: Vector3D::new(0.0, 0.0, 0.0),
                 velocity: Vector3D::new(0.0, 0.0, 0.0)}
    }

    pub fn name<'a>(&'a self) -> &'a str {
        &self.name
    }

    pub fn set_name<S>(&mut self, name: S) where S: Into<String> {
        self.name = name.into();
    }

    pub fn mass(&self) -> f64 {
        self.mass
    }
    pub fn set_mass(&mut self, mass: f64) {
        self.mass = mass;
    }

    pub fn position<'a>(&'a self) -> &'a Vector3D {
        &self.position
    }
    pub fn set_position(&mut self, pos: Vector3D) {
        self.position = pos;
    }
    pub fn add_position(&mut self, new_pos: Vector3D) {
        self.position = self.position + new_pos;
    }

    pub fn velocity<'a>(&'a self) -> &'a Vector3D {
        &self.velocity
    }
    pub fn set_velocity(&mut self, vel: Vector3D) {
        self.velocity = vel;
    }
    pub fn add_velocity(&mut self, new_vel: Vector3D) {
        self.velocity = self.velocity + new_vel;
    }
}

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

    #[test]
    fn name() {
        let mut part = Particle::new("O");
        assert_eq!(part.name(), "O");

        part.set_name("H");
        assert_eq!(part.name(), "H");

        let name = "Zn";
        part.set_name(name);
        assert_eq!(part.name(), "Zn");
    }

    #[test]
    fn mass() {
        let mut part = Particle::new("O");
        part.set_mass(16.0);
        assert_eq!(part.mass(), 16.0);
    }

    #[test]
    fn coordinates() {
        let mut part = Particle::new("O");
        assert_eq!(part.position(), &Vector3D::new(0.0, 0.0, 0.0));

        part.set_position(Vector3D::new(1.0, 2.0, 3.0));
        assert_eq!(part.position(), &Vector3D::new(1.0, 2.0, 3.0));

        part.set_velocity(Vector3D::new(1.0, 2.0, 3.0));
        assert_eq!(part.velocity(), &Vector3D::new(1.0, 2.0, 3.0));
    }
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Nick Platt <platt.nicholas@gmail.com>
// (c) Jian Zeng <anonymousknight96 AT gmail.com>
//
// For the full copyright and license information, please view the LICENSE file
// that was distributed with this source code.

// spell-checker:ignore (ToDO) filetime strptime utcoff strs datetime MMDDhhmm clapv

pub extern crate filetime;

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, AppSettings, Arg, ArgGroup};
use filetime::*;
use std::fs::{self, File};
use std::path::Path;
use uucore::display::Quotable;
use uucore::error::{FromIo, UError, UResult, USimpleError};

static ABOUT: &str = "Update the access and modification times of each FILE to the current time.";
pub mod options {
    // Both SOURCES and sources are needed as we need to be able to refer to the ArgGroup.
    pub static SOURCES: &str = "sources";
    pub mod sources {
        pub static DATE: &str = "date";
        pub static REFERENCE: &str = "reference";
        pub static CURRENT: &str = "current";
    }
    pub static ACCESS: &str = "access";
    pub static MODIFICATION: &str = "modification";
    pub static NO_CREATE: &str = "no-create";
    pub static NO_DEREF: &str = "no-dereference";
    pub static TIME: &str = "time";
}

static ARG_FILES: &str = "files";

fn to_local(mut tm: time::Tm) -> time::Tm {
    tm.tm_utcoff = time::now().tm_utcoff;
    tm
}

fn local_tm_to_filetime(tm: time::Tm) -> FileTime {
    let ts = tm.to_timespec();
    FileTime::from_unix_time(ts.sec as i64, ts.nsec as u32)
}

fn usage() -> String {
    format!("{0} [OPTION]... [USER]", uucore::execution_phrase())
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let usage = usage();

    let matches = uu_app().override_usage(&usage[..]).get_matches_from(args);

    let files = matches.values_of_os(ARG_FILES).unwrap();

    let (mut atime, mut mtime) =
        if let Some(reference) = matches.value_of_os(options::sources::REFERENCE) {
            stat(Path::new(reference), !matches.is_present(options::NO_DEREF))?
        } else {
            let timestamp = if let Some(date) = matches.value_of(options::sources::DATE) {
                parse_date(date)?
            } else if let Some(current) = matches.value_of(options::sources::CURRENT) {
                parse_timestamp(current)?
            } else {
                local_tm_to_filetime(time::now())
            };
            (timestamp, timestamp)
        };

    for filename in files {
        let path = Path::new(filename);
        if !path.exists() {
            // no-dereference included here for compatibility
            if matches.is_present(options::NO_CREATE) || matches.is_present(options::NO_DEREF) {
                continue;
            }

            if let Err(e) = File::create(path) {
                show!(e.map_err_context(|| format!("cannot touch {}", path.quote())));
                continue;
            };

            // Minor optimization: if no reference time was specified, we're done.
            if !matches.is_present(options::SOURCES) {
                continue;
            }
        }

        // If changing "only" atime or mtime, grab the existing value of the other.
        // Note that "-a" and "-m" may be passed together; this is not an xor.
        if matches.is_present(options::ACCESS)
            || matches.is_present(options::MODIFICATION)
            || matches.is_present(options::TIME)
        {
            let st = stat(path, !matches.is_present(options::NO_DEREF))?;
            let time = matches.value_of(options::TIME).unwrap_or("");

            if !(matches.is_present(options::ACCESS)
                || time.contains(&"access".to_owned())
                || time.contains(&"atime".to_owned())
                || time.contains(&"use".to_owned()))
            {
                atime = st.0;
            }

            if !(matches.is_present(options::MODIFICATION)
                || time.contains(&"modify".to_owned())
                || time.contains(&"mtime".to_owned()))
            {
                mtime = st.1;
            }
        }

        if matches.is_present(options::NO_DEREF) {
            set_symlink_file_times(path, atime, mtime)
        } else {
            filetime::set_file_times(path, atime, mtime)
        }
        .map_err_context(|| format!("setting times of {}", path.quote()))?;
    }

    Ok(())
}

pub fn uu_app<'a>() -> App<'a> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .setting(AppSettings::InferLongArgs)
        .arg(
            Arg::new(options::ACCESS)
                .short('a')
                .help("change only the access time"),
        )
        .arg(
            Arg::new(options::sources::CURRENT)
                .short('t')
                .help("use [[CC]YY]MMDDhhmm[.ss] instead of the current time")
                .value_name("STAMP")
                .takes_value(true),
        )
        .arg(
            Arg::new(options::sources::DATE)
                .short('d')
                .long(options::sources::DATE)
                .help("parse argument and use it instead of current time")
                .value_name("STRING"),
        )
        .arg(
            Arg::new(options::MODIFICATION)
                .short('m')
                .help("change only the modification time"),
        )
        .arg(
            Arg::new(options::NO_CREATE)
                .short('c')
                .long(options::NO_CREATE)
                .help("do not create any files"),
        )
        .arg(
            Arg::new(options::NO_DEREF)
                .short('h')
                .long(options::NO_DEREF)
                .help(
                    "affect each symbolic link instead of any referenced file \
                     (only for systems that can change the timestamps of a symlink)",
                ),
        )
        .arg(
            Arg::new(options::sources::REFERENCE)
                .short('r')
                .long(options::sources::REFERENCE)
                .help("use this file's times instead of the current time")
                .value_name("FILE")
                .allow_invalid_utf8(true),
        )
        .arg(
            Arg::new(options::TIME)
                .long(options::TIME)
                .help(
                    "change only the specified time: \"access\", \"atime\", or \
                     \"use\" are equivalent to -a; \"modify\" or \"mtime\" are \
                     equivalent to -m",
                )
                .value_name("WORD")
                .possible_values(&["access", "atime", "use"])
                .takes_value(true),
        )
        .arg(
            Arg::new(ARG_FILES)
                .multiple_occurrences(true)
                .takes_value(true)
                .min_values(1)
                .allow_invalid_utf8(true),
        )
        .group(ArgGroup::new(options::SOURCES).args(&[
            options::sources::CURRENT,
            options::sources::DATE,
            options::sources::REFERENCE,
        ]))
}

fn stat(path: &Path, follow: bool) -> UResult<(FileTime, FileTime)> {
    let metadata = if follow {
        fs::symlink_metadata(path)
    } else {
        fs::metadata(path)
    }
    .map_err_context(|| format!("failed to get attributes of {}", path.quote()))?;

    Ok((
        FileTime::from_last_access_time(&metadata),
        FileTime::from_last_modification_time(&metadata),
    ))
}

fn parse_date(str: &str) -> UResult<FileTime> {
    // This isn't actually compatible with GNU touch, but there doesn't seem to
    // be any simple specification for what format this parameter allows and I'm
    // not about to implement GNU parse_datetime.
    // http://git.savannah.gnu.org/gitweb/?p=gnulib.git;a=blob_plain;f=lib/parse-datetime.y
    let formats = vec!["%c", "%F"];
    for f in formats {
        if let Ok(tm) = time::strptime(str, f) {
            return Ok(local_tm_to_filetime(to_local(tm)));
        }
    }

    if let Ok(tm) = time::strptime(str, "@%s") {
        // Don't convert to local time in this case - seconds since epoch are not time-zone dependent
        return Ok(local_tm_to_filetime(tm));
    }

    Err(USimpleError::new(
        1,
        format!("Unable to parse date: {}", str),
    ))
}

fn parse_timestamp(s: &str) -> UResult<FileTime> {
    let now = time::now();
    let (format, ts) = match s.chars().count() {
        15 => ("%Y%m%d%H%M.%S", s.to_owned()),
        12 => ("%Y%m%d%H%M", s.to_owned()),
        13 => ("%y%m%d%H%M.%S", s.to_owned()),
        10 => ("%y%m%d%H%M", s.to_owned()),
        11 => ("%Y%m%d%H%M.%S", format!("{}{}", now.tm_year + 1900, s)),
        8 => ("%Y%m%d%H%M", format!("{}{}", now.tm_year + 1900, s)),
        _ => {
            return Err(USimpleError::new(
                1,
                format!("invalid date format {}", s.quote()),
            ))
        }
    };

    let tm = time::strptime(&ts, format)
        .map_err(|_| USimpleError::new(1, format!("invalid date format {}", s.quote())))?;

    let mut local = to_local(tm);
    local.tm_isdst = -1;
    let ft = local_tm_to_filetime(local);

    // We have to check that ft is valid time. Due to daylight saving
    // time switch, local time can jump from 1:59 AM to 3:00 AM,
    // in which case any time between 2:00 AM and 2:59 AM is not valid.
    // Convert back to local time and see if we got the same value back.
    let ts = time::Timespec {
        sec: ft.unix_seconds(),
        nsec: 0,
    };
    let tm2 = time::at(ts);
    if tm.tm_hour != tm2.tm_hour {
        return Err(USimpleError::new(
            1,
            format!("invalid date format {}", s.quote()),
        ));
    }

    Ok(ft)
}
<commit_msg>touch: Better error message when no args is given<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Nick Platt <platt.nicholas@gmail.com>
// (c) Jian Zeng <anonymousknight96 AT gmail.com>
//
// For the full copyright and license information, please view the LICENSE file
// that was distributed with this source code.

// spell-checker:ignore (ToDO) filetime strptime utcoff strs datetime MMDDhhmm clapv

pub extern crate filetime;

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, AppSettings, Arg, ArgGroup};
use filetime::*;
use std::fs::{self, File};
use std::path::Path;
use uucore::display::Quotable;
use uucore::error::{FromIo, UError, UResult, USimpleError};

static ABOUT: &str = "Update the access and modification times of each FILE to the current time.";
pub mod options {
    // Both SOURCES and sources are needed as we need to be able to refer to the ArgGroup.
    pub static SOURCES: &str = "sources";
    pub mod sources {
        pub static DATE: &str = "date";
        pub static REFERENCE: &str = "reference";
        pub static CURRENT: &str = "current";
    }
    pub static ACCESS: &str = "access";
    pub static MODIFICATION: &str = "modification";
    pub static NO_CREATE: &str = "no-create";
    pub static NO_DEREF: &str = "no-dereference";
    pub static TIME: &str = "time";
}

static ARG_FILES: &str = "files";

fn to_local(mut tm: time::Tm) -> time::Tm {
    tm.tm_utcoff = time::now().tm_utcoff;
    tm
}

fn local_tm_to_filetime(tm: time::Tm) -> FileTime {
    let ts = tm.to_timespec();
    FileTime::from_unix_time(ts.sec as i64, ts.nsec as u32)
}

fn usage() -> String {
    format!("{0} [OPTION]... [USER]", uucore::execution_phrase())
}

#[uucore::main]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let usage = usage();

    let matches = uu_app().override_usage(&usage[..]).get_matches_from(args);

    let files = matches.values_of_os(ARG_FILES).ok_or(USimpleError::new(
        1,
        r##"missing file operand
Try 'touch --help' for more information."##,
    ))?;

    let (mut atime, mut mtime) =
        if let Some(reference) = matches.value_of_os(options::sources::REFERENCE) {
            stat(Path::new(reference), !matches.is_present(options::NO_DEREF))?
        } else {
            let timestamp = if let Some(date) = matches.value_of(options::sources::DATE) {
                parse_date(date)?
            } else if let Some(current) = matches.value_of(options::sources::CURRENT) {
                parse_timestamp(current)?
            } else {
                local_tm_to_filetime(time::now())
            };
            (timestamp, timestamp)
        };

    for filename in files {
        let path = Path::new(filename);
        if !path.exists() {
            // no-dereference included here for compatibility
            if matches.is_present(options::NO_CREATE) || matches.is_present(options::NO_DEREF) {
                continue;
            }

            if let Err(e) = File::create(path) {
                show!(e.map_err_context(|| format!("cannot touch {}", path.quote())));
                continue;
            };

            // Minor optimization: if no reference time was specified, we're done.
            if !matches.is_present(options::SOURCES) {
                continue;
            }
        }

        // If changing "only" atime or mtime, grab the existing value of the other.
        // Note that "-a" and "-m" may be passed together; this is not an xor.
        if matches.is_present(options::ACCESS)
            || matches.is_present(options::MODIFICATION)
            || matches.is_present(options::TIME)
        {
            let st = stat(path, !matches.is_present(options::NO_DEREF))?;
            let time = matches.value_of(options::TIME).unwrap_or("");

            if !(matches.is_present(options::ACCESS)
                || time.contains(&"access".to_owned())
                || time.contains(&"atime".to_owned())
                || time.contains(&"use".to_owned()))
            {
                atime = st.0;
            }

            if !(matches.is_present(options::MODIFICATION)
                || time.contains(&"modify".to_owned())
                || time.contains(&"mtime".to_owned()))
            {
                mtime = st.1;
            }
        }

        if matches.is_present(options::NO_DEREF) {
            set_symlink_file_times(path, atime, mtime)
        } else {
            filetime::set_file_times(path, atime, mtime)
        }
        .map_err_context(|| format!("setting times of {}", path.quote()))?;
    }

    Ok(())
}

pub fn uu_app<'a>() -> App<'a> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .setting(AppSettings::InferLongArgs)
        .arg(
            Arg::new(options::ACCESS)
                .short('a')
                .help("change only the access time"),
        )
        .arg(
            Arg::new(options::sources::CURRENT)
                .short('t')
                .help("use [[CC]YY]MMDDhhmm[.ss] instead of the current time")
                .value_name("STAMP")
                .takes_value(true),
        )
        .arg(
            Arg::new(options::sources::DATE)
                .short('d')
                .long(options::sources::DATE)
                .help("parse argument and use it instead of current time")
                .value_name("STRING"),
        )
        .arg(
            Arg::new(options::MODIFICATION)
                .short('m')
                .help("change only the modification time"),
        )
        .arg(
            Arg::new(options::NO_CREATE)
                .short('c')
                .long(options::NO_CREATE)
                .help("do not create any files"),
        )
        .arg(
            Arg::new(options::NO_DEREF)
                .short('h')
                .long(options::NO_DEREF)
                .help(
                    "affect each symbolic link instead of any referenced file \
                     (only for systems that can change the timestamps of a symlink)",
                ),
        )
        .arg(
            Arg::new(options::sources::REFERENCE)
                .short('r')
                .long(options::sources::REFERENCE)
                .help("use this file's times instead of the current time")
                .value_name("FILE")
                .allow_invalid_utf8(true),
        )
        .arg(
            Arg::new(options::TIME)
                .long(options::TIME)
                .help(
                    "change only the specified time: \"access\", \"atime\", or \
                     \"use\" are equivalent to -a; \"modify\" or \"mtime\" are \
                     equivalent to -m",
                )
                .value_name("WORD")
                .possible_values(&["access", "atime", "use"])
                .takes_value(true),
        )
        .arg(
            Arg::new(ARG_FILES)
                .multiple_occurrences(true)
                .takes_value(true)
                .min_values(1)
                .allow_invalid_utf8(true),
        )
        .group(ArgGroup::new(options::SOURCES).args(&[
            options::sources::CURRENT,
            options::sources::DATE,
            options::sources::REFERENCE,
        ]))
}

fn stat(path: &Path, follow: bool) -> UResult<(FileTime, FileTime)> {
    let metadata = if follow {
        fs::symlink_metadata(path)
    } else {
        fs::metadata(path)
    }
    .map_err_context(|| format!("failed to get attributes of {}", path.quote()))?;

    Ok((
        FileTime::from_last_access_time(&metadata),
        FileTime::from_last_modification_time(&metadata),
    ))
}

fn parse_date(str: &str) -> UResult<FileTime> {
    // This isn't actually compatible with GNU touch, but there doesn't seem to
    // be any simple specification for what format this parameter allows and I'm
    // not about to implement GNU parse_datetime.
    // http://git.savannah.gnu.org/gitweb/?p=gnulib.git;a=blob_plain;f=lib/parse-datetime.y
    let formats = vec!["%c", "%F"];
    for f in formats {
        if let Ok(tm) = time::strptime(str, f) {
            return Ok(local_tm_to_filetime(to_local(tm)));
        }
    }

    if let Ok(tm) = time::strptime(str, "@%s") {
        // Don't convert to local time in this case - seconds since epoch are not time-zone dependent
        return Ok(local_tm_to_filetime(tm));
    }

    Err(USimpleError::new(
        1,
        format!("Unable to parse date: {}", str),
    ))
}

fn parse_timestamp(s: &str) -> UResult<FileTime> {
    let now = time::now();
    let (format, ts) = match s.chars().count() {
        15 => ("%Y%m%d%H%M.%S", s.to_owned()),
        12 => ("%Y%m%d%H%M", s.to_owned()),
        13 => ("%y%m%d%H%M.%S", s.to_owned()),
        10 => ("%y%m%d%H%M", s.to_owned()),
        11 => ("%Y%m%d%H%M.%S", format!("{}{}", now.tm_year + 1900, s)),
        8 => ("%Y%m%d%H%M", format!("{}{}", now.tm_year + 1900, s)),
        _ => {
            return Err(USimpleError::new(
                1,
                format!("invalid date format {}", s.quote()),
            ))
        }
    };

    let tm = time::strptime(&ts, format)
        .map_err(|_| USimpleError::new(1, format!("invalid date format {}", s.quote())))?;

    let mut local = to_local(tm);
    local.tm_isdst = -1;
    let ft = local_tm_to_filetime(local);

    // We have to check that ft is valid time. Due to daylight saving
    // time switch, local time can jump from 1:59 AM to 3:00 AM,
    // in which case any time between 2:00 AM and 2:59 AM is not valid.
    // Convert back to local time and see if we got the same value back.
    let ts = time::Timespec {
        sec: ft.unix_seconds(),
        nsec: 0,
    };
    let tm2 = time::at(ts);
    if tm.tm_hour != tm2.tm_hour {
        return Err(USimpleError::new(
            1,
            format!("invalid date format {}", s.quote()),
        ));
    }

    Ok(ft)
}
<|endoftext|>