<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use smallvec::SmallVec;
use std::cell::RefCell;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

type Exponent = u8;

#[derive(Clone, Debug)]
struct Decomposition(SmallVec<[(u64, Exponent); NUM_FACTORS_INLINE]>);

// The number of factors to inline directly into a `Decomposition` object.
// As a consequence of the Erdős–Kac theorem, the average number of prime factors
// of integers < 10²⁵ ≃ 2⁸³ is 4, so we can use a slightly higher value.
const NUM_FACTORS_INLINE: usize = 5;

impl Decomposition {
    fn one() -> Decomposition {
        Decomposition(SmallVec::new())
    }

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);

        if let Some((_, e)) = self.0.iter_mut().find(|(f, _)| *f == factor) {
            *e += exp;
        } else {
            self.0.push((factor, exp))
        }
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }

    fn get(&self, p: u64) -> Option<&(u64, u8)> {
        self.0.iter().find(|(q, _)| *q == p)
    }
}

impl PartialEq for Decomposition {
    fn eq(&self, other: &Decomposition) -> bool {
        for p in &self.0 {
            if other.get(p.0) != Some(p) {
                return false;
            }
        }

        for p in &other.0 {
            if self.get(p.0) != Some(p) {
                return false;
            }
        }

        true
    }
}
impl Eq for Decomposition {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors(RefCell<Decomposition>);

impl Factors {
    pub fn one() -> Factors {
        Factors(RefCell::new(Decomposition::one()))
    }

    pub fn add(&mut self, prime: u64, exp: Exponent) {
        debug_assert!(miller_rabin::is_prime(prime));
        self.0.borrow_mut().add(prime, exp)
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let v = &mut (self.0).borrow_mut().0;
        v.sort_unstable();

        for (p, exp) in v.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            #[cfg(feature = "coz")]
            coz::progress!("factor found");
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    #[cfg(feature = "coz")]
    coz::begin!("factorization");
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let n_zeros = n.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }

    if n == 1 {
        #[cfg(feature = "coz")]
        coz::end!("factorization");
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    let r = if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    };

    #[cfg(feature = "coz")]
    coz::end!("factorization");

    return r;
}

#[cfg(test)]
mod tests {
    use super::{factor, Decomposition, Exponent, Factors};
    use quickcheck::quickcheck;
    use smallvec::smallvec;
    use std::cell::RefCell;

    #[test]
    fn factor_2044854919485649() {
        let f = Factors(RefCell::new(Decomposition(smallvec![
            (503, 1),
            (2423, 1),
            (40961, 2)
        ])));
        assert_eq!(factor(f.product()), f);
    }

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's code path handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> () {
            assert_eq!(factor(f.product()), f);
        }

        fn exponentiate_factors(f: Factors, e: Exponent) -> () {
            if e == 0 { return; }
            if let Some(fe) = f.product().checked_pow(e.into()) {
                assert_eq!(factor(fe), f ^ e);
            }
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
        use rand::Rng;
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorization.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = gen.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}

#[cfg(test)]
impl std::ops::BitXor<Exponent> for Factors {
    type Output = Self;

    fn bitxor(self, rhs: Exponent) -> Factors {
        debug_assert_ne!(rhs, 0);
        let mut r = Factors::one();
        for (p, e) in self.0.borrow().0.iter() {
            r.add(*p, rhs * e);
        }

        debug_assert_eq!(r.product(), self.product().pow(rhs.into()));
        return r;
    }
}
<commit_msg>factor: unneeded  statement<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use smallvec::SmallVec;
use std::cell::RefCell;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

type Exponent = u8;

#[derive(Clone, Debug)]
struct Decomposition(SmallVec<[(u64, Exponent); NUM_FACTORS_INLINE]>);

// The number of factors to inline directly into a `Decomposition` object.
// As a consequence of the Erdős–Kac theorem, the average number of prime factors
// of integers < 10²⁵ ≃ 2⁸³ is 4, so we can use a slightly higher value.
const NUM_FACTORS_INLINE: usize = 5;

impl Decomposition {
    fn one() -> Decomposition {
        Decomposition(SmallVec::new())
    }

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);

        if let Some((_, e)) = self.0.iter_mut().find(|(f, _)| *f == factor) {
            *e += exp;
        } else {
            self.0.push((factor, exp))
        }
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }

    fn get(&self, p: u64) -> Option<&(u64, u8)> {
        self.0.iter().find(|(q, _)| *q == p)
    }
}

impl PartialEq for Decomposition {
    fn eq(&self, other: &Decomposition) -> bool {
        for p in &self.0 {
            if other.get(p.0) != Some(p) {
                return false;
            }
        }

        for p in &other.0 {
            if self.get(p.0) != Some(p) {
                return false;
            }
        }

        true
    }
}
impl Eq for Decomposition {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors(RefCell<Decomposition>);

impl Factors {
    pub fn one() -> Factors {
        Factors(RefCell::new(Decomposition::one()))
    }

    pub fn add(&mut self, prime: u64, exp: Exponent) {
        debug_assert!(miller_rabin::is_prime(prime));
        self.0.borrow_mut().add(prime, exp)
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let v = &mut (self.0).borrow_mut().0;
        v.sort_unstable();

        for (p, exp) in v.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            #[cfg(feature = "coz")]
            coz::progress!("factor found");
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    #[cfg(feature = "coz")]
    coz::begin!("factorization");
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let n_zeros = n.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }

    if n == 1 {
        #[cfg(feature = "coz")]
        coz::end!("factorization");
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    let r = if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    };

    #[cfg(feature = "coz")]
    coz::end!("factorization");

    r
}

#[cfg(test)]
mod tests {
    use super::{factor, Decomposition, Exponent, Factors};
    use quickcheck::quickcheck;
    use smallvec::smallvec;
    use std::cell::RefCell;

    #[test]
    fn factor_2044854919485649() {
        let f = Factors(RefCell::new(Decomposition(smallvec![
            (503, 1),
            (2423, 1),
            (40961, 2)
        ])));
        assert_eq!(factor(f.product()), f);
    }

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's code path handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> () {
            assert_eq!(factor(f.product()), f);
        }

        fn exponentiate_factors(f: Factors, e: Exponent) -> () {
            if e == 0 { return; }
            if let Some(fe) = f.product().checked_pow(e.into()) {
                assert_eq!(factor(fe), f ^ e);
            }
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
        use rand::Rng;
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorization.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = gen.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}

#[cfg(test)]
impl std::ops::BitXor<Exponent> for Factors {
    type Output = Self;

    fn bitxor(self, rhs: Exponent) -> Factors {
        debug_assert_ne!(rhs, 0);
        let mut r = Factors::one();
        for (p, e) in self.0.borrow().0.iter() {
            r.add(*p, rhs * e);
        }

        debug_assert_eq!(r.product(), self.product().pow(rhs.into()));
        return r;
    }
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Sunrin SHIMURA
// Collaborator: Jian Zeng
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (paths) GPGHome

use clap::{crate_version, Arg, Command};
use uucore::display::{println_verbatim, Quotable};
use uucore::error::{FromIo, UError, UResult};
use uucore::format_usage;

use std::env;
use std::error::Error;
use std::fmt::Display;
use std::iter;
use std::path::{is_separator, Path, PathBuf};

use rand::Rng;
use tempfile::Builder;

static ABOUT: &str = "create a temporary file or directory.";
const USAGE: &str = "{} [OPTION]... [TEMPLATE]";

static DEFAULT_TEMPLATE: &str = "tmp.XXXXXXXXXX";

static OPT_DIRECTORY: &str = "directory";
static OPT_DRY_RUN: &str = "dry-run";
static OPT_QUIET: &str = "quiet";
static OPT_SUFFIX: &str = "suffix";
static OPT_TMPDIR: &str = "tmpdir";
static OPT_T: &str = "t";

static ARG_TEMPLATE: &str = "template";

#[derive(Debug)]
enum MkTempError {
    PersistError(PathBuf),
    MustEndInX(String),
    TooFewXs(String),
    ContainsDirSeparator(String),
    InvalidTemplate(String),
}

impl UError for MkTempError {}

impl Error for MkTempError {}

impl Display for MkTempError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use MkTempError::*;
        match self {
            PersistError(p) => write!(f, "could not persist file {}", p.quote()),
            MustEndInX(s) => write!(f, "with --suffix, template {} must end in X", s.quote()),
            TooFewXs(s) => write!(f, "too few X's in template {}", s.quote()),
            ContainsDirSeparator(s) => {
                write!(
                    f,
                    "invalid suffix {}, contains directory separator",
                    s.quote()
                )
            }
            InvalidTemplate(s) => write!(
                f,
                "invalid template, {}; with --tmpdir, it may not be absolute",
                s.quote()
            ),
        }
    }
}

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

    let template = matches.value_of(ARG_TEMPLATE).unwrap();
    let tmpdir = matches.value_of(OPT_TMPDIR).unwrap_or_default();

    // Treat the template string as a path to get the directory
    // containing the last component.
    let path = PathBuf::from(template);

    let (template, tmpdir) = if matches.is_present(OPT_TMPDIR)
        && !PathBuf::from(tmpdir).is_dir() // if a temp dir is provided, it must be an actual path
        && tmpdir.contains("XXX")
    // If this is a template, it has to contain at least 3 X
        && template == DEFAULT_TEMPLATE
    // That means that clap does not think we provided a template
    {
        // Special case to workaround a limitation of clap when doing
        // mktemp --tmpdir apt-key-gpghome.XXX
        // The behavior should be
        // mktemp --tmpdir $TMPDIR apt-key-gpghome.XX
        // As --tmpdir is empty
        //
        // Fixed in clap 3
        // See https://github.com/clap-rs/clap/pull/1587
        let tmp = env::temp_dir();
        (tmpdir, tmp)
    } else if !matches.is_present(OPT_TMPDIR) {
        // In this case, the command line was `mktemp -t XXX`, so we
        // treat the argument `XXX` as though it were a filename
        // regardless of whether it has path separators in it.
        if matches.is_present(OPT_T) {
            let tmp = env::temp_dir();
            (template, tmp)
        // In this case, the command line was `mktemp XXX`, so we need
        // to parse out the parent directory and the filename from the
        // argument `XXX`, since it may be include path separators.
        } else {
            let tmp = match path.parent() {
                None => PathBuf::from("."),
                Some(d) => PathBuf::from(d),
            };
            let filename = path.file_name();
            let template = filename.unwrap().to_str().unwrap();
            (template, tmp)
        }
    } else {
        (template, PathBuf::from(tmpdir))
    };

    let make_dir = matches.is_present(OPT_DIRECTORY);
    let dry_run = matches.is_present(OPT_DRY_RUN);
    let suppress_file_err = matches.is_present(OPT_QUIET);

    let (prefix, rand, suffix) = parse_template(template, matches.value_of(OPT_SUFFIX))?;

    if matches.is_present(OPT_TMPDIR) && PathBuf::from(prefix).is_absolute() {
        return Err(MkTempError::InvalidTemplate(template.into()).into());
    }

    let res = if dry_run {
        dry_exec(tmpdir, prefix, rand, suffix)
    } else {
        exec(&tmpdir, prefix, rand, suffix, make_dir)
    };

    if suppress_file_err {
        // Mapping all UErrors to ExitCodes prevents the errors from being printed
        res.map_err(|e| e.code().into())
    } else {
        res
    }
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(OPT_DIRECTORY)
                .short('d')
                .long(OPT_DIRECTORY)
                .help("Make a directory instead of a file"),
        )
        .arg(
            Arg::new(OPT_DRY_RUN)
                .short('u')
                .long(OPT_DRY_RUN)
                .help("do not create anything; merely print a name (unsafe)"),
        )
        .arg(
            Arg::new(OPT_QUIET)
                .short('q')
                .long("quiet")
                .help("Fail silently if an error occurs."),
        )
        .arg(
            Arg::new(OPT_SUFFIX)
                .long(OPT_SUFFIX)
                .help(
                    "append SUFFIX to TEMPLATE; SUFFIX must not contain a path separator. \
                     This option is implied if TEMPLATE does not end with X.",
                )
                .value_name("SUFFIX"),
        )
        .arg(
            Arg::new(OPT_TMPDIR)
                .short('p')
                .long(OPT_TMPDIR)
                .help(
                    "interpret TEMPLATE relative to DIR; if DIR is not specified, use \
                     $TMPDIR ($TMP on windows) if set, else /tmp. With this option, TEMPLATE must not \
                     be an absolute name; unlike with -t, TEMPLATE may contain \
                     slashes, but mktemp creates only the final component",
                )
                .value_name("DIR"),
        )
        .arg(Arg::new(OPT_T).short('t').help(
            "Generate a template (using the supplied prefix and TMPDIR (TMP on windows) if set) \
             to create a filename template [deprecated]",
        ))
        .arg(
            Arg::new(ARG_TEMPLATE)
                .multiple_occurrences(false)
                .takes_value(true)
                .max_values(1)
                .default_value(DEFAULT_TEMPLATE),
        )
}

fn parse_template<'a>(
    temp: &'a str,
    suffix: Option<&'a str>,
) -> UResult<(&'a str, usize, &'a str)> {
    let right = match temp.rfind('X') {
        Some(r) => r + 1,
        None => return Err(MkTempError::TooFewXs(temp.into()).into()),
    };
    let left = temp[..right].rfind(|c| c != 'X').map_or(0, |i| i + 1);
    let prefix = &temp[..left];
    let rand = right - left;

    if rand < 3 {
        return Err(MkTempError::TooFewXs(temp.into()).into());
    }

    let mut suf = &temp[right..];

    if let Some(s) = suffix {
        if suf.is_empty() {
            suf = s;
        } else {
            return Err(MkTempError::MustEndInX(temp.into()).into());
        }
    };

    if suf.chars().any(is_separator) {
        return Err(MkTempError::ContainsDirSeparator(suf.into()).into());
    }

    Ok((prefix, rand, suf))
}

pub fn dry_exec(mut tmpdir: PathBuf, prefix: &str, rand: usize, suffix: &str) -> UResult<()> {
    let len = prefix.len() + suffix.len() + rand;
    let mut buf = Vec::with_capacity(len);
    buf.extend(prefix.as_bytes());
    buf.extend(iter::repeat(b'X').take(rand));
    buf.extend(suffix.as_bytes());

    // Randomize.
    let bytes = &mut buf[prefix.len()..prefix.len() + rand];
    rand::thread_rng().fill(bytes);
    for byte in bytes.iter_mut() {
        *byte = match *byte % 62 {
            v @ 0..=9 => (v + b'0'),
            v @ 10..=35 => (v - 10 + b'a'),
            v @ 36..=61 => (v - 36 + b'A'),
            _ => unreachable!(),
        }
    }
    // We guarantee utf8.
    let buf = String::from_utf8(buf).unwrap();
    tmpdir.push(buf);
    println_verbatim(tmpdir).map_err_context(|| "failed to print directory name".to_owned())
}

fn exec(dir: &Path, prefix: &str, rand: usize, suffix: &str, make_dir: bool) -> UResult<()> {
    let context = || {
        format!(
            "failed to create file via template '{}{}{}'",
            prefix,
            "X".repeat(rand),
            suffix
        )
    };

    let mut builder = Builder::new();
    builder.prefix(prefix).rand_bytes(rand).suffix(suffix);

    let path = if make_dir {
        builder
            .tempdir_in(&dir)
            .map_err_context(context)?
            .into_path() // `into_path` consumes the TempDir without removing it
    } else {
        builder
            .tempfile_in(&dir)
            .map_err_context(context)?
            .keep() // `keep` ensures that the file is not deleted
            .map_err(|e| MkTempError::PersistError(e.file.path().to_path_buf()))?
            .1
    };

    // Get just the last component of the path to the created
    // temporary file or directory.
    let filename = path.file_name();
    let filename = filename.unwrap().to_str().unwrap();

    // Join the directory to the path to get the path to print. We
    // cannot use the path returned by the `Builder` because it gives
    // the absolute path and we need to return a filename that matches
    // the template given on the command-line which might be a
    // relative path.
    let mut path = dir.to_path_buf();
    path.push(filename);

    println_verbatim(path).map_err_context(|| "failed to print directory name".to_owned())
}
<commit_msg>mktemp: return MkTempError from parse_template()<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Sunrin SHIMURA
// Collaborator: Jian Zeng
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

// spell-checker:ignore (paths) GPGHome

use clap::{crate_version, Arg, Command};
use uucore::display::{println_verbatim, Quotable};
use uucore::error::{FromIo, UError, UResult};
use uucore::format_usage;

use std::env;
use std::error::Error;
use std::fmt::Display;
use std::iter;
use std::path::{is_separator, Path, PathBuf};

use rand::Rng;
use tempfile::Builder;

static ABOUT: &str = "create a temporary file or directory.";
const USAGE: &str = "{} [OPTION]... [TEMPLATE]";

static DEFAULT_TEMPLATE: &str = "tmp.XXXXXXXXXX";

static OPT_DIRECTORY: &str = "directory";
static OPT_DRY_RUN: &str = "dry-run";
static OPT_QUIET: &str = "quiet";
static OPT_SUFFIX: &str = "suffix";
static OPT_TMPDIR: &str = "tmpdir";
static OPT_T: &str = "t";

static ARG_TEMPLATE: &str = "template";

#[derive(Debug)]
enum MkTempError {
    PersistError(PathBuf),
    MustEndInX(String),
    TooFewXs(String),
    ContainsDirSeparator(String),
    InvalidTemplate(String),
}

impl UError for MkTempError {}

impl Error for MkTempError {}

impl Display for MkTempError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use MkTempError::*;
        match self {
            PersistError(p) => write!(f, "could not persist file {}", p.quote()),
            MustEndInX(s) => write!(f, "with --suffix, template {} must end in X", s.quote()),
            TooFewXs(s) => write!(f, "too few X's in template {}", s.quote()),
            ContainsDirSeparator(s) => {
                write!(
                    f,
                    "invalid suffix {}, contains directory separator",
                    s.quote()
                )
            }
            InvalidTemplate(s) => write!(
                f,
                "invalid template, {}; with --tmpdir, it may not be absolute",
                s.quote()
            ),
        }
    }
}

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

    let template = matches.value_of(ARG_TEMPLATE).unwrap();
    let tmpdir = matches.value_of(OPT_TMPDIR).unwrap_or_default();

    // Treat the template string as a path to get the directory
    // containing the last component.
    let path = PathBuf::from(template);

    let (template, tmpdir) = if matches.is_present(OPT_TMPDIR)
        && !PathBuf::from(tmpdir).is_dir() // if a temp dir is provided, it must be an actual path
        && tmpdir.contains("XXX")
    // If this is a template, it has to contain at least 3 X
        && template == DEFAULT_TEMPLATE
    // That means that clap does not think we provided a template
    {
        // Special case to workaround a limitation of clap when doing
        // mktemp --tmpdir apt-key-gpghome.XXX
        // The behavior should be
        // mktemp --tmpdir $TMPDIR apt-key-gpghome.XX
        // As --tmpdir is empty
        //
        // Fixed in clap 3
        // See https://github.com/clap-rs/clap/pull/1587
        let tmp = env::temp_dir();
        (tmpdir, tmp)
    } else if !matches.is_present(OPT_TMPDIR) {
        // In this case, the command line was `mktemp -t XXX`, so we
        // treat the argument `XXX` as though it were a filename
        // regardless of whether it has path separators in it.
        if matches.is_present(OPT_T) {
            let tmp = env::temp_dir();
            (template, tmp)
        // In this case, the command line was `mktemp XXX`, so we need
        // to parse out the parent directory and the filename from the
        // argument `XXX`, since it may be include path separators.
        } else {
            let tmp = match path.parent() {
                None => PathBuf::from("."),
                Some(d) => PathBuf::from(d),
            };
            let filename = path.file_name();
            let template = filename.unwrap().to_str().unwrap();
            (template, tmp)
        }
    } else {
        (template, PathBuf::from(tmpdir))
    };

    let make_dir = matches.is_present(OPT_DIRECTORY);
    let dry_run = matches.is_present(OPT_DRY_RUN);
    let suppress_file_err = matches.is_present(OPT_QUIET);

    let (prefix, rand, suffix) = parse_template(template, matches.value_of(OPT_SUFFIX))?;

    if matches.is_present(OPT_TMPDIR) && PathBuf::from(prefix).is_absolute() {
        return Err(MkTempError::InvalidTemplate(template.into()).into());
    }

    let res = if dry_run {
        dry_exec(tmpdir, prefix, rand, suffix)
    } else {
        exec(&tmpdir, prefix, rand, suffix, make_dir)
    };

    if suppress_file_err {
        // Mapping all UErrors to ExitCodes prevents the errors from being printed
        res.map_err(|e| e.code().into())
    } else {
        res
    }
}

pub fn uu_app<'a>() -> Command<'a> {
    Command::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .override_usage(format_usage(USAGE))
        .infer_long_args(true)
        .arg(
            Arg::new(OPT_DIRECTORY)
                .short('d')
                .long(OPT_DIRECTORY)
                .help("Make a directory instead of a file"),
        )
        .arg(
            Arg::new(OPT_DRY_RUN)
                .short('u')
                .long(OPT_DRY_RUN)
                .help("do not create anything; merely print a name (unsafe)"),
        )
        .arg(
            Arg::new(OPT_QUIET)
                .short('q')
                .long("quiet")
                .help("Fail silently if an error occurs."),
        )
        .arg(
            Arg::new(OPT_SUFFIX)
                .long(OPT_SUFFIX)
                .help(
                    "append SUFFIX to TEMPLATE; SUFFIX must not contain a path separator. \
                     This option is implied if TEMPLATE does not end with X.",
                )
                .value_name("SUFFIX"),
        )
        .arg(
            Arg::new(OPT_TMPDIR)
                .short('p')
                .long(OPT_TMPDIR)
                .help(
                    "interpret TEMPLATE relative to DIR; if DIR is not specified, use \
                     $TMPDIR ($TMP on windows) if set, else /tmp. With this option, TEMPLATE must not \
                     be an absolute name; unlike with -t, TEMPLATE may contain \
                     slashes, but mktemp creates only the final component",
                )
                .value_name("DIR"),
        )
        .arg(Arg::new(OPT_T).short('t').help(
            "Generate a template (using the supplied prefix and TMPDIR (TMP on windows) if set) \
             to create a filename template [deprecated]",
        ))
        .arg(
            Arg::new(ARG_TEMPLATE)
                .multiple_occurrences(false)
                .takes_value(true)
                .max_values(1)
                .default_value(DEFAULT_TEMPLATE),
        )
}

/// Parse a template string into prefix, suffix, and random components.
///
/// `temp` is the template string, with three or more consecutive `X`s
/// representing a placeholder for randomly generated characters (for
/// example, `"abc_XXX.txt"`). If `temp` ends in an `X`, then a suffix
/// can be specified by `suffix` instead.
///
/// # Errors
///
/// * If there are fewer than three consecutive `X`s in `temp`.
/// * If `suffix` is a [`Some`] object but `temp` does not end in `X`.
/// * If the suffix (specified either way) contains a path separator.
///
/// # Examples
///
/// ```rust,ignore
/// assert_eq!(parse_template("XXX", None).unwrap(), ("", 3, ""));
/// assert_eq!(parse_template("abcXXX", None).unwrap(), ("abc", 3, ""));
/// assert_eq!(parse_template("XXXdef", None).unwrap(), ("", 3, "def"));
/// assert_eq!(parse_template("abcXXXdef", None).unwrap(), ("abc", 3, "def"));
/// ```
fn parse_template<'a>(
    temp: &'a str,
    suffix: Option<&'a str>,
) -> Result<(&'a str, usize, &'a str), MkTempError> {
    let right = match temp.rfind('X') {
        Some(r) => r + 1,
        None => return Err(MkTempError::TooFewXs(temp.into())),
    };
    let left = temp[..right].rfind(|c| c != 'X').map_or(0, |i| i + 1);
    let prefix = &temp[..left];
    let rand = right - left;

    if rand < 3 {
        return Err(MkTempError::TooFewXs(temp.into()));
    }

    let mut suf = &temp[right..];

    if let Some(s) = suffix {
        if suf.is_empty() {
            suf = s;
        } else {
            return Err(MkTempError::MustEndInX(temp.into()));
        }
    };

    if suf.chars().any(is_separator) {
        return Err(MkTempError::ContainsDirSeparator(suf.into()));
    }

    Ok((prefix, rand, suf))
}

pub fn dry_exec(mut tmpdir: PathBuf, prefix: &str, rand: usize, suffix: &str) -> UResult<()> {
    let len = prefix.len() + suffix.len() + rand;
    let mut buf = Vec::with_capacity(len);
    buf.extend(prefix.as_bytes());
    buf.extend(iter::repeat(b'X').take(rand));
    buf.extend(suffix.as_bytes());

    // Randomize.
    let bytes = &mut buf[prefix.len()..prefix.len() + rand];
    rand::thread_rng().fill(bytes);
    for byte in bytes.iter_mut() {
        *byte = match *byte % 62 {
            v @ 0..=9 => (v + b'0'),
            v @ 10..=35 => (v - 10 + b'a'),
            v @ 36..=61 => (v - 36 + b'A'),
            _ => unreachable!(),
        }
    }
    // We guarantee utf8.
    let buf = String::from_utf8(buf).unwrap();
    tmpdir.push(buf);
    println_verbatim(tmpdir).map_err_context(|| "failed to print directory name".to_owned())
}

fn exec(dir: &Path, prefix: &str, rand: usize, suffix: &str, make_dir: bool) -> UResult<()> {
    let context = || {
        format!(
            "failed to create file via template '{}{}{}'",
            prefix,
            "X".repeat(rand),
            suffix
        )
    };

    let mut builder = Builder::new();
    builder.prefix(prefix).rand_bytes(rand).suffix(suffix);

    let path = if make_dir {
        builder
            .tempdir_in(&dir)
            .map_err_context(context)?
            .into_path() // `into_path` consumes the TempDir without removing it
    } else {
        builder
            .tempfile_in(&dir)
            .map_err_context(context)?
            .keep() // `keep` ensures that the file is not deleted
            .map_err(|e| MkTempError::PersistError(e.file.path().to_path_buf()))?
            .1
    };

    // Get just the last component of the path to the created
    // temporary file or directory.
    let filename = path.file_name();
    let filename = filename.unwrap().to_str().unwrap();

    // Join the directory to the path to get the path to print. We
    // cannot use the path returned by the `Builder` because it gives
    // the absolute path and we need to return a filename that matches
    // the template given on the command-line which might be a
    // relative path.
    let mut path = dir.to_path_buf();
    path.push(filename);

    println_verbatim(path).map_err_context(|| "failed to print directory name".to_owned())
}

#[cfg(test)]
mod tests {
    use crate::parse_template;

    #[test]
    fn test_parse_template_no_suffix() {
        assert_eq!(parse_template("XXX", None).unwrap(), ("", 3, ""));
        assert_eq!(parse_template("abcXXX", None).unwrap(), ("abc", 3, ""));
        assert_eq!(parse_template("XXXdef", None).unwrap(), ("", 3, "def"));
        assert_eq!(
            parse_template("abcXXXdef", None).unwrap(),
            ("abc", 3, "def")
        );
    }

    #[test]
    fn test_parse_template_suffix() {
        assert_eq!(parse_template("XXX", Some("def")).unwrap(), ("", 3, "def"));
        assert_eq!(
            parse_template("abcXXX", Some("def")).unwrap(),
            ("abc", 3, "def")
        );
    }

    #[test]
    fn test_parse_template_errors() {
        // TODO This should be an error as well, but we are not
        // catching it just yet. A future commit will correct this.
        //
        //     assert!(parse_template("a/bXXX", None).is_err());
        //
        assert!(parse_template("XXXa/b", None).is_err());
        assert!(parse_template("XX", None).is_err());
        assert!(parse_template("XXXabc", Some("def")).is_err());
        assert!(parse_template("XXX", Some("a/b")).is_err());
    }
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Yury Krivopalov <ykrivopalov@yandex.ru>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore N'th M'th

use crate::errors::*;
use crate::format::format_and_print;
use crate::options::*;
use crate::units::{Result, Unit};
use clap::{crate_version, App, AppSettings, Arg, ArgMatches};
use std::io::{BufRead, Write};
use uucore::display::Quotable;
use uucore::error::UResult;
use uucore::ranges::Range;

pub mod errors;
pub mod format;
pub mod options;
mod units;

static ABOUT: &str = "Convert numbers from/to human-readable strings";
static LONG_HELP: &str = "UNIT options:
   none   no auto-scaling is done; suffixes will trigger an error

   auto   accept optional single/two letter suffix:

          1K = 1000, 1Ki = 1024, 1M = 1000000, 1Mi = 1048576,

   si     accept optional single letter suffix:

          1K = 1000, 1M = 1000000, ...

   iec    accept optional single letter suffix:

          1K = 1024, 1M = 1048576, ...

   iec-i  accept optional two-letter suffix:

          1Ki = 1024, 1Mi = 1048576, ...

FIELDS supports cut(1) style field ranges:
  N    N'th field, counted from 1
  N-   from N'th field, to end of line
  N-M  from N'th to M'th field (inclusive)
  -M   from first to M'th field (inclusive)
  -    all fields
Multiple fields/ranges can be separated with commas
";

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

fn handle_args<'a>(args: impl Iterator<Item = &'a str>, options: NumfmtOptions) -> UResult<()> {
    for l in args {
        match format_and_print(l, &options) {
            Ok(_) => Ok(()),
            Err(e) => Err(NumfmtError::FormattingError(e.to_string())),
        }?;
    }

    Ok(())
}

fn handle_buffer<R>(input: R, options: NumfmtOptions) -> UResult<()>
where
    R: BufRead,
{
    let mut lines = input.lines();
    for (idx, line) in lines.by_ref().enumerate() {
        match line {
            Ok(l) if idx < options.header => {
                println!("{}", l);
                Ok(())
            }
            Ok(l) => match format_and_print(&l, &options) {
                Ok(_) => Ok(()),
                Err(e) => Err(NumfmtError::FormattingError(e.to_string())),
            },
            Err(e) => Err(NumfmtError::IoError(e.to_string())),
        }?;
    }
    Ok(())
}

fn parse_unit(s: &str) -> Result<Unit> {
    match s {
        "auto" => Ok(Unit::Auto),
        "si" => Ok(Unit::Si),
        "iec" => Ok(Unit::Iec(false)),
        "iec-i" => Ok(Unit::Iec(true)),
        "none" => Ok(Unit::None),
        _ => Err("Unsupported unit is specified".to_owned()),
    }
}

fn parse_options(args: &ArgMatches) -> Result<NumfmtOptions> {
    let from = parse_unit(args.value_of(options::FROM).unwrap())?;
    let to = parse_unit(args.value_of(options::TO).unwrap())?;

    let transform = TransformOptions { from, to };

    let padding = match args.value_of(options::PADDING) {
        Some(s) => s.parse::<isize>().map_err(|err| err.to_string()),
        None => Ok(0),
    }?;

    let header = match args.occurrences_of(options::HEADER) {
        0 => Ok(0),
        _ => {
            let value = args.value_of(options::HEADER).unwrap();

            value
                .parse::<usize>()
                .map_err(|_| value)
                .and_then(|n| match n {
                    0 => Err(value),
                    _ => Ok(n),
                })
                .map_err(|value| format!("invalid header value {}", value.quote()))
        }
    }?;

    let fields = match args.value_of(options::FIELD).unwrap() {
        "-" => vec![Range {
            low: 1,
            high: std::usize::MAX,
        }],
        v => Range::from_list(v)?,
    };

    let delimiter = args.value_of(options::DELIMITER).map_or(Ok(None), |arg| {
        if arg.len() == 1 {
            Ok(Some(arg.to_string()))
        } else {
            Err("the delimiter must be a single character".to_string())
        }
    })?;

    // unwrap is fine because the argument has a default value
    let round = match args.value_of(options::ROUND).unwrap() {
        "up" => RoundMethod::Up,
        "down" => RoundMethod::Down,
        "from-zero" => RoundMethod::FromZero,
        "towards-zero" => RoundMethod::TowardsZero,
        "nearest" => RoundMethod::Nearest,
        _ => unreachable!("Should be restricted by clap"),
    };

    let suffix = args.value_of(options::SUFFIX).map(|s| s.to_owned());

    Ok(NumfmtOptions {
        transform,
        padding,
        header,
        fields,
        delimiter,
        round,
        suffix,
    })
}

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

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

    let options = parse_options(&matches).map_err(NumfmtError::IllegalArgument)?;

    let result = match matches.values_of(options::NUMBER) {
        Some(values) => handle_args(values, options),
        None => {
            let stdin = std::io::stdin();
            let mut locked_stdin = stdin.lock();
            handle_buffer(&mut locked_stdin, options)
        }
    };

    match result {
        Err(e) => {
            std::io::stdout().flush().expect("error flushing stdout");
            Err(e)
        }
        _ => Ok(()),
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .after_help(LONG_HELP)
        .setting(AppSettings::AllowNegativeNumbers)
        .arg(
            Arg::with_name(options::DELIMITER)
                .short("d")
                .long(options::DELIMITER)
                .value_name("X")
                .help("use X instead of whitespace for field delimiter"),
        )
        .arg(
            Arg::with_name(options::FIELD)
                .long(options::FIELD)
                .help("replace the numbers in these input fields (default=1) see FIELDS below")
                .value_name("FIELDS")
                .default_value(options::FIELD_DEFAULT),
        )
        .arg(
            Arg::with_name(options::FROM)
                .long(options::FROM)
                .help("auto-scale input numbers to UNITs; see UNIT below")
                .value_name("UNIT")
                .default_value(options::FROM_DEFAULT),
        )
        .arg(
            Arg::with_name(options::TO)
                .long(options::TO)
                .help("auto-scale output numbers to UNITs; see UNIT below")
                .value_name("UNIT")
                .default_value(options::TO_DEFAULT),
        )
        .arg(
            Arg::with_name(options::PADDING)
                .long(options::PADDING)
                .help(
                    "pad the output to N characters; positive N will \
                     right-align; negative N will left-align; padding is \
                     ignored if the output is wider than N; the default is \
                     to automatically pad if a whitespace is found",
                )
                .value_name("N"),
        )
        .arg(
            Arg::with_name(options::HEADER)
                .long(options::HEADER)
                .help(
                    "print (without converting) the first N header lines; \
                     N defaults to 1 if not specified",
                )
                .value_name("N")
                .default_value(options::HEADER_DEFAULT)
                .hide_default_value(true),
        )
        .arg(
            Arg::with_name(options::ROUND)
                .long(options::ROUND)
                .help(
                    "use METHOD for rounding when scaling; METHOD can be: up,\
                    down, from-zero (default), towards-zero, nearest",
                )
                .value_name("METHOD")
                .default_value("from-zero")
                .possible_values(&["up", "down", "from-zero", "towards-zero", "nearest"]),
        )
        .arg(
            Arg::with_name(options::SUFFIX)
                .long(options::SUFFIX)
                .help(
                    "print SUFFIX after each formatted number, and accept \
                    inputs optionally ending with SUFFIX",
                )
                .value_name("SUFFIX"),
        )
        .arg(Arg::with_name(options::NUMBER).hidden(true).multiple(true))
}

#[cfg(test)]
mod tests {
    use super::{handle_buffer, NumfmtOptions, Range, RoundMethod, TransformOptions, Unit};
    use std::io::{BufReader, Error, ErrorKind, Read};
    struct MockBuffer {}

    impl Read for MockBuffer {
        fn read(&mut self, _: &mut [u8]) -> Result<usize, Error> {
            Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
        }
    }

    fn get_valid_options() -> NumfmtOptions {
        NumfmtOptions {
            transform: TransformOptions {
                from: Unit::None,
                to: Unit::None,
            },
            padding: 10,
            header: 0,
            fields: vec![Range { low: 0, high: 1 }],
            delimiter: None,
            round: RoundMethod::Nearest,
            suffix: None,
        }
    }

    fn get_invalid_options() -> NumfmtOptions {
        NumfmtOptions {
            transform: TransformOptions {
                from: Unit::Auto,
                to: Unit::Auto,
            },
            padding: 10,
            header: 0,
            fields: vec![Range { low: 0, high: 1 }],
            delimiter: None,
            round: RoundMethod::Nearest,
            suffix: None,
        }
    }

    #[test]
    fn broken_buffer_returns_io_error() {
        let mock_buffer = MockBuffer {};
        let result = handle_buffer(BufReader::new(mock_buffer), get_valid_options())
            .expect_err("returned Ok after receiving IO error");
        let result_debug = format!("{:?}", result);
        let result_display = format!("{}", result);
        assert_eq!(result_debug, "IoError(\"broken pipe\")");
        assert_eq!(result_display, "broken pipe");
        assert_eq!(result.code(), 1);
    }

    #[test]
    fn non_numeric_returns_formatting_error() {
        let input_value = b"hello";
        let result = handle_buffer(BufReader::new(&input_value[..]), get_valid_options())
            .expect_err("returned Ok after receiving improperly formatted input");
        let result_debug = format!("{:?}", result);
        let result_display = format!("{}", result);
        assert_eq!(
            result_debug,
            "FormattingError(\"invalid suffix in input: 'hello'\")"
        );
        assert_eq!(result_display, "invalid suffix in input: 'hello'");
        assert_eq!(result.code(), 2);
    }

    #[test]
    fn valid_input_returns_ok() {
        let input_value = b"165";
        let result = handle_buffer(BufReader::new(&input_value[..]), get_valid_options());
        println!("{:?}", result);
        assert!(result.is_ok(), "did not return Ok for valid input");
    }
}
<commit_msg>numfmt: remove unused function<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Yury Krivopalov <ykrivopalov@yandex.ru>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore N'th M'th

use crate::errors::*;
use crate::format::format_and_print;
use crate::options::*;
use crate::units::{Result, Unit};
use clap::{crate_version, App, AppSettings, Arg, ArgMatches};
use std::io::{BufRead, Write};
use uucore::display::Quotable;
use uucore::error::UResult;
use uucore::ranges::Range;

pub mod errors;
pub mod format;
pub mod options;
mod units;

static ABOUT: &str = "Convert numbers from/to human-readable strings";
static LONG_HELP: &str = "UNIT options:
   none   no auto-scaling is done; suffixes will trigger an error

   auto   accept optional single/two letter suffix:

          1K = 1000, 1Ki = 1024, 1M = 1000000, 1Mi = 1048576,

   si     accept optional single letter suffix:

          1K = 1000, 1M = 1000000, ...

   iec    accept optional single letter suffix:

          1K = 1024, 1M = 1048576, ...

   iec-i  accept optional two-letter suffix:

          1Ki = 1024, 1Mi = 1048576, ...

FIELDS supports cut(1) style field ranges:
  N    N'th field, counted from 1
  N-   from N'th field, to end of line
  N-M  from N'th to M'th field (inclusive)
  -M   from first to M'th field (inclusive)
  -    all fields
Multiple fields/ranges can be separated with commas
";

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

fn handle_args<'a>(args: impl Iterator<Item = &'a str>, options: NumfmtOptions) -> UResult<()> {
    for l in args {
        match format_and_print(l, &options) {
            Ok(_) => Ok(()),
            Err(e) => Err(NumfmtError::FormattingError(e.to_string())),
        }?;
    }

    Ok(())
}

fn handle_buffer<R>(input: R, options: NumfmtOptions) -> UResult<()>
where
    R: BufRead,
{
    let mut lines = input.lines();
    for (idx, line) in lines.by_ref().enumerate() {
        match line {
            Ok(l) if idx < options.header => {
                println!("{}", l);
                Ok(())
            }
            Ok(l) => match format_and_print(&l, &options) {
                Ok(_) => Ok(()),
                Err(e) => Err(NumfmtError::FormattingError(e.to_string())),
            },
            Err(e) => Err(NumfmtError::IoError(e.to_string())),
        }?;
    }
    Ok(())
}

fn parse_unit(s: &str) -> Result<Unit> {
    match s {
        "auto" => Ok(Unit::Auto),
        "si" => Ok(Unit::Si),
        "iec" => Ok(Unit::Iec(false)),
        "iec-i" => Ok(Unit::Iec(true)),
        "none" => Ok(Unit::None),
        _ => Err("Unsupported unit is specified".to_owned()),
    }
}

fn parse_options(args: &ArgMatches) -> Result<NumfmtOptions> {
    let from = parse_unit(args.value_of(options::FROM).unwrap())?;
    let to = parse_unit(args.value_of(options::TO).unwrap())?;

    let transform = TransformOptions { from, to };

    let padding = match args.value_of(options::PADDING) {
        Some(s) => s.parse::<isize>().map_err(|err| err.to_string()),
        None => Ok(0),
    }?;

    let header = match args.occurrences_of(options::HEADER) {
        0 => Ok(0),
        _ => {
            let value = args.value_of(options::HEADER).unwrap();

            value
                .parse::<usize>()
                .map_err(|_| value)
                .and_then(|n| match n {
                    0 => Err(value),
                    _ => Ok(n),
                })
                .map_err(|value| format!("invalid header value {}", value.quote()))
        }
    }?;

    let fields = match args.value_of(options::FIELD).unwrap() {
        "-" => vec![Range {
            low: 1,
            high: std::usize::MAX,
        }],
        v => Range::from_list(v)?,
    };

    let delimiter = args.value_of(options::DELIMITER).map_or(Ok(None), |arg| {
        if arg.len() == 1 {
            Ok(Some(arg.to_string()))
        } else {
            Err("the delimiter must be a single character".to_string())
        }
    })?;

    // unwrap is fine because the argument has a default value
    let round = match args.value_of(options::ROUND).unwrap() {
        "up" => RoundMethod::Up,
        "down" => RoundMethod::Down,
        "from-zero" => RoundMethod::FromZero,
        "towards-zero" => RoundMethod::TowardsZero,
        "nearest" => RoundMethod::Nearest,
        _ => unreachable!("Should be restricted by clap"),
    };

    let suffix = args.value_of(options::SUFFIX).map(|s| s.to_owned());

    Ok(NumfmtOptions {
        transform,
        padding,
        header,
        fields,
        delimiter,
        round,
        suffix,
    })
}

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

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

    let options = parse_options(&matches).map_err(NumfmtError::IllegalArgument)?;

    let result = match matches.values_of(options::NUMBER) {
        Some(values) => handle_args(values, options),
        None => {
            let stdin = std::io::stdin();
            let mut locked_stdin = stdin.lock();
            handle_buffer(&mut locked_stdin, options)
        }
    };

    match result {
        Err(e) => {
            std::io::stdout().flush().expect("error flushing stdout");
            Err(e)
        }
        _ => Ok(()),
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .after_help(LONG_HELP)
        .setting(AppSettings::AllowNegativeNumbers)
        .arg(
            Arg::with_name(options::DELIMITER)
                .short("d")
                .long(options::DELIMITER)
                .value_name("X")
                .help("use X instead of whitespace for field delimiter"),
        )
        .arg(
            Arg::with_name(options::FIELD)
                .long(options::FIELD)
                .help("replace the numbers in these input fields (default=1) see FIELDS below")
                .value_name("FIELDS")
                .default_value(options::FIELD_DEFAULT),
        )
        .arg(
            Arg::with_name(options::FROM)
                .long(options::FROM)
                .help("auto-scale input numbers to UNITs; see UNIT below")
                .value_name("UNIT")
                .default_value(options::FROM_DEFAULT),
        )
        .arg(
            Arg::with_name(options::TO)
                .long(options::TO)
                .help("auto-scale output numbers to UNITs; see UNIT below")
                .value_name("UNIT")
                .default_value(options::TO_DEFAULT),
        )
        .arg(
            Arg::with_name(options::PADDING)
                .long(options::PADDING)
                .help(
                    "pad the output to N characters; positive N will \
                     right-align; negative N will left-align; padding is \
                     ignored if the output is wider than N; the default is \
                     to automatically pad if a whitespace is found",
                )
                .value_name("N"),
        )
        .arg(
            Arg::with_name(options::HEADER)
                .long(options::HEADER)
                .help(
                    "print (without converting) the first N header lines; \
                     N defaults to 1 if not specified",
                )
                .value_name("N")
                .default_value(options::HEADER_DEFAULT)
                .hide_default_value(true),
        )
        .arg(
            Arg::with_name(options::ROUND)
                .long(options::ROUND)
                .help(
                    "use METHOD for rounding when scaling; METHOD can be: up,\
                    down, from-zero (default), towards-zero, nearest",
                )
                .value_name("METHOD")
                .default_value("from-zero")
                .possible_values(&["up", "down", "from-zero", "towards-zero", "nearest"]),
        )
        .arg(
            Arg::with_name(options::SUFFIX)
                .long(options::SUFFIX)
                .help(
                    "print SUFFIX after each formatted number, and accept \
                    inputs optionally ending with SUFFIX",
                )
                .value_name("SUFFIX"),
        )
        .arg(Arg::with_name(options::NUMBER).hidden(true).multiple(true))
}

#[cfg(test)]
mod tests {
    use super::{handle_buffer, NumfmtOptions, Range, RoundMethod, TransformOptions, Unit};
    use std::io::{BufReader, Error, ErrorKind, Read};
    struct MockBuffer {}

    impl Read for MockBuffer {
        fn read(&mut self, _: &mut [u8]) -> Result<usize, Error> {
            Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
        }
    }

    fn get_valid_options() -> NumfmtOptions {
        NumfmtOptions {
            transform: TransformOptions {
                from: Unit::None,
                to: Unit::None,
            },
            padding: 10,
            header: 0,
            fields: vec![Range { low: 0, high: 1 }],
            delimiter: None,
            round: RoundMethod::Nearest,
            suffix: None,
        }
    }

    #[test]
    fn broken_buffer_returns_io_error() {
        let mock_buffer = MockBuffer {};
        let result = handle_buffer(BufReader::new(mock_buffer), get_valid_options())
            .expect_err("returned Ok after receiving IO error");
        let result_debug = format!("{:?}", result);
        let result_display = format!("{}", result);
        assert_eq!(result_debug, "IoError(\"broken pipe\")");
        assert_eq!(result_display, "broken pipe");
        assert_eq!(result.code(), 1);
    }

    #[test]
    fn non_numeric_returns_formatting_error() {
        let input_value = b"hello";
        let result = handle_buffer(BufReader::new(&input_value[..]), get_valid_options())
            .expect_err("returned Ok after receiving improperly formatted input");
        let result_debug = format!("{:?}", result);
        let result_display = format!("{}", result);
        assert_eq!(
            result_debug,
            "FormattingError(\"invalid suffix in input: 'hello'\")"
        );
        assert_eq!(result_display, "invalid suffix in input: 'hello'");
        assert_eq!(result.code(), 2);
    }

    #[test]
    fn valid_input_returns_ok() {
        let input_value = b"165";
        let result = handle_buffer(BufReader::new(&input_value[..]), get_valid_options());
        println!("{:?}", result);
        assert!(result.is_ok(), "did not return Ok for valid input");
    }
}
<|endoftext|>


<commit_before>//! Contains implementation of ASCII FBX emitter.
extern crate rustc_serialize;

use std::io::Write;
use writer::error::{Result, Error};
use common::Property;

fn indent<W: Write>(sink: &mut W, depth: usize) -> Result<()> {
    for _ in 0..depth {
        try!(sink.write(b"\t"));
    }
    Ok(())
}

fn print_property<W: Write>(sink: &mut W, property: &Property, prop_depth: usize) -> Result<()> {
    assert!(prop_depth > 0);

    // TODO: I've never seen vector of booleans (in binary or ascii FBX)... How should it be?
    // TODO: How will it be when other properties follows a property of array value?
    // TODO: Implement folding of large array.
    match *property {
        Property::Bool(false) => {
            try!(sink.write(b"T"));
        },
        Property::Bool(true) => {
            try!(sink.write(b"Y"));
        },
        Property::I16(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::I32(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::I64(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::F32(v) => {
            // NOTE: Is outputted data accurate enough?
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::F64(v) => {
            // NOTE: Is outputted data accurate enough?
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::VecBool(vec) => {
            warn!("ASCII representation of vector of boolean values may be wrong.");
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write(if v { b"Y" } else { b"T" }));
            }
            for &v in iter {
                try!(sink.write(if v { b",Y" } else { b",T" }));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::VecI32(vec) => {
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write_fmt(format_args!("{}", v)));
            }
            for &v in iter {
                try!(sink.write_fmt(format_args!(",{}", v)));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::VecI64(vec) => {
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write_fmt(format_args!("{}", v)));
            }
            for &v in iter {
                try!(sink.write_fmt(format_args!(",{}", v)));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::VecF32(vec) => {
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write_fmt(format_args!("{}", v)));
            }
            for &v in iter {
                try!(sink.write_fmt(format_args!(",{}", v)));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::VecF64(vec) => {
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write_fmt(format_args!("{}", v)));
            }
            for &v in iter {
                try!(sink.write_fmt(format_args!(",{}", v)));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::String(v) => {
            try!(sink.write(b"\""));
            for c in v.chars() {
                match c {
                    '"' => {
                        try!(sink.write(b"""));
                    },
                    '\n' => {
                        try!(sink.write(b"&lf;"));
                    },
                    '\r' => {
                        try!(sink.write(b"&cr;"));
                    },
                    _ => {
                        try!(sink.write_fmt(format_args!("{}", c)));
                    }
                }
            }
            try!(sink.write(b"\""));
        },
        Property::Binary(v) => {
            // TODO: Implement folding of long line.
            use self::rustc_serialize::base64::{ToBase64, STANDARD};
            // base64 conversion.
            try!(sink.write_fmt(format_args!("\"{}\"", v.to_base64(STANDARD))));
        },
    }
    Ok(())
}

/// A writer for ASCII FBX.
#[derive(Debug, Clone)]
pub struct AsciiEmitter {
    prop_child_existence: Vec<(bool, bool)>,
}

impl AsciiEmitter {
    /// Constructs ASCII FBX writer.
    pub fn new() -> Self {
        AsciiEmitter {
            prop_child_existence: vec![],
        }
    }

    pub fn emit_start_fbx<W: Write>(&mut self, sink: &mut W, ver: u32) -> Result<()> {
        if (ver < 7000) || (ver >= 8000) {
            error!("Unsupported version: {}", ver);
            return Err(Error::UnsupportedFbxVersion(ver));
        }
        {
            let (major, minor) = (ver / 1000, ver % 1000);
            let (minor, revision) = (minor / 100, minor % 100);
            // Write magic for ASCII FBX.
            try!(sink.write_fmt(format_args!("; FBX {}.{}.{} project file\n", major, minor, revision)));
        }

        Ok(())
    }

    pub fn emit_end_fbx<W: Write>(&mut self, _sink: &mut W) -> Result<()> {
        Ok(())
    }

    pub fn emit_start_node<W: Write>(&mut self, sink: &mut W, name: &str, properties: &[Property]) -> Result<()> {
        if let Some((prop_exist, child_exist)) = self.prop_child_existence.pop() {
            // Print brace for *parent node*, if the current node is the first child.
            // (i.e. `child_exist` of parent is `false`.)
            if !child_exist {
                try!(sink.write(b" {\n"));
            }
            self.prop_child_existence.push((prop_exist, true));
        }
        try!(indent(sink, self.prop_child_existence.len()));
        self.prop_child_existence.push((!properties.is_empty(), false));
        try!(sink.write_fmt(format_args!("{}: ", name)));

        let prop_depth = self.prop_child_existence.len();
        let mut prop_iter = properties.iter();
        if let Some(prop) = prop_iter.next() {
            try!(print_property(sink, prop, prop_depth));
        }
        for prop in prop_iter {
            try!(sink.write(b", "));
            try!(print_property(sink, prop, prop_depth));
        }

        Ok(())
    }

    pub fn emit_end_node<W: Write>(&mut self, sink: &mut W) -> Result<()> {
        let (prop_exist, child_exist) = self.prop_child_existence.pop().unwrap();
        if !prop_exist || child_exist {
            if !prop_exist && !child_exist {
                try!(sink.write(b" {\n"));
            }
            try!(indent(sink, self.prop_child_existence.len()));
            try!(sink.write(b"}\n"));
        } else {
            try!(sink.write(b"\n"));
        }

        Ok(())
    }

    pub fn emit_comment<W: Write>(&mut self, sink: &mut W, comment: &str) -> Result<()> {
        for line in comment.lines() {
            try!(indent(sink, self.prop_child_existence.len()));
            try!(sink.write(line.as_bytes()));
            try!(sink.write(b"\n"));
        }

        Ok(())
    }
}
<commit_msg>Simplify code<commit_after>//! Contains implementation of ASCII FBX emitter.
extern crate rustc_serialize;

use std::io::Write;
use writer::error::{Result, Error};
use common::Property;

fn indent<W: Write>(sink: &mut W, depth: usize) -> Result<()> {
    for _ in 0..depth {
        try!(sink.write(b"\t"));
    }
    Ok(())
}

fn print_property<W: Write>(sink: &mut W, property: &Property, prop_depth: usize) -> Result<()> {
    assert!(prop_depth > 0);

    // TODO: I've never seen vector of booleans (in binary or ascii FBX)... How should it be?
    // TODO: How will it be when other properties follows a property of array value?
    // TODO: Implement folding of large array.
    macro_rules! generic_vec_print {
        ($vec:ident) => ({
            try!(sink.write_fmt(format_args!("*{} {{\n", $vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = $vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write_fmt(format_args!("{}", v)));
            }
            for &v in iter {
                try!(sink.write_fmt(format_args!(",{}", v)));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        })
    }
    match *property {
        Property::Bool(false) => {
            try!(sink.write(b"T"));
        },
        Property::Bool(true) => {
            try!(sink.write(b"Y"));
        },
        Property::I16(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::I32(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::I64(v) => {
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::F32(v) => {
            // NOTE: Is outputted data accurate enough?
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::F64(v) => {
            // NOTE: Is outputted data accurate enough?
            try!(sink.write_fmt(format_args!("{}", v)));
        },
        Property::VecBool(vec) => {
            warn!("ASCII representation of vector of boolean values may be wrong.");
            try!(sink.write_fmt(format_args!("*{} {{\n", vec.len())));
            try!(indent(sink, prop_depth));
            try!(sink.write(b"a: "));
            let mut iter = vec.iter();
            if let Some(&v) = iter.next() {
                try!(sink.write(if v { b"Y" } else { b"T" }));
            }
            for &v in iter {
                try!(sink.write(if v { b",Y" } else { b",T" }));
            }
            try!(sink.write(b"\n"));
            try!(indent(sink, prop_depth-1));
            try!(sink.write(b"}"));
        },
        Property::VecI32(vec) => {
            generic_vec_print!(vec);
        },
        Property::VecI64(vec) => {
            generic_vec_print!(vec);
        },
        Property::VecF32(vec) => {
            generic_vec_print!(vec);
        },
        Property::VecF64(vec) => {
            generic_vec_print!(vec);
        },
        Property::String(v) => {
            try!(sink.write(b"\""));
            for c in v.chars() {
                match c {
                    '"' => {
                        try!(sink.write(b"""));
                    },
                    '\n' => {
                        try!(sink.write(b"&lf;"));
                    },
                    '\r' => {
                        try!(sink.write(b"&cr;"));
                    },
                    _ => {
                        try!(sink.write_fmt(format_args!("{}", c)));
                    }
                }
            }
            try!(sink.write(b"\""));
        },
        Property::Binary(v) => {
            // TODO: Implement folding of long line.
            use self::rustc_serialize::base64::{ToBase64, STANDARD};
            // base64 conversion.
            try!(sink.write_fmt(format_args!("\"{}\"", v.to_base64(STANDARD))));
        },
    }
    Ok(())
}

/// A writer for ASCII FBX.
#[derive(Debug, Clone)]
pub struct AsciiEmitter {
    prop_child_existence: Vec<(bool, bool)>,
}

impl AsciiEmitter {
    /// Constructs ASCII FBX writer.
    pub fn new() -> Self {
        AsciiEmitter {
            prop_child_existence: vec![],
        }
    }

    pub fn emit_start_fbx<W: Write>(&mut self, sink: &mut W, ver: u32) -> Result<()> {
        if (ver < 7000) || (ver >= 8000) {
            error!("Unsupported version: {}", ver);
            return Err(Error::UnsupportedFbxVersion(ver));
        }
        {
            let (major, minor) = (ver / 1000, ver % 1000);
            let (minor, revision) = (minor / 100, minor % 100);
            // Write magic for ASCII FBX.
            try!(sink.write_fmt(format_args!("; FBX {}.{}.{} project file\n", major, minor, revision)));
        }

        Ok(())
    }

    pub fn emit_end_fbx<W: Write>(&mut self, _sink: &mut W) -> Result<()> {
        Ok(())
    }

    pub fn emit_start_node<W: Write>(&mut self, sink: &mut W, name: &str, properties: &[Property]) -> Result<()> {
        if let Some((prop_exist, child_exist)) = self.prop_child_existence.pop() {
            // Print brace for *parent node*, if the current node is the first child.
            // (i.e. `child_exist` of parent is `false`.)
            if !child_exist {
                try!(sink.write(b" {\n"));
            }
            self.prop_child_existence.push((prop_exist, true));
        }
        try!(indent(sink, self.prop_child_existence.len()));
        self.prop_child_existence.push((!properties.is_empty(), false));
        try!(sink.write_fmt(format_args!("{}: ", name)));

        let prop_depth = self.prop_child_existence.len();
        let mut prop_iter = properties.iter();
        if let Some(prop) = prop_iter.next() {
            try!(print_property(sink, prop, prop_depth));
        }
        for prop in prop_iter {
            try!(sink.write(b", "));
            try!(print_property(sink, prop, prop_depth));
        }

        Ok(())
    }

    pub fn emit_end_node<W: Write>(&mut self, sink: &mut W) -> Result<()> {
        let (prop_exist, child_exist) = self.prop_child_existence.pop().unwrap();
        if !prop_exist || child_exist {
            if !prop_exist && !child_exist {
                try!(sink.write(b" {\n"));
            }
            try!(indent(sink, self.prop_child_existence.len()));
            try!(sink.write(b"}\n"));
        } else {
            try!(sink.write(b"\n"));
        }

        Ok(())
    }

    pub fn emit_comment<W: Write>(&mut self, sink: &mut W, comment: &str) -> Result<()> {
        for line in comment.lines() {
            try!(indent(sink, self.prop_child_existence.len()));
            try!(sink.write(line.as_bytes()));
            try!(sink.write(b"\n"));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>fn main() {
    let mut p = ~[1, 2, 3];
    incrementMut(p);
    for &x in q.iter() {
        print!("{:d} ", x);
    }
    print("\n");
}

fn incrementMut(vec: &[int]) -> &[int] {
    for &x in vec.iter() {
        x = x + 1;
    }
}
<commit_msg>working ex2.2<commit_after>fn main() {
    let mut p = ~[1, 2, 3];
    incrementMut(p);
    for &x in p.iter() {
        print!("{:d} ", x);
    }
    print("\n");
}

// Increment each value in `vec` by one.
fn incrementMut(vec: &mut [int]) {
    for i in range(0, vec.len()) {
        vec[i] += 1;
    }
}
<|endoftext|>


<commit_before>extern crate kafka;
extern crate getopts;
extern crate env_logger;

use std::{env, fmt, process};
use std::io::{self, Write};
use std::ascii::AsciiExt;

use kafka::consumer::{Consumer, GroupOffsetStorage};

/// This is a very simple command line application reading from a
/// specific kafka topic and dumping the messages to standard output.
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    };
    if let Err(e) = process(cfg) {
        println!("{}", e);
        process::exit(1);
    }
}

fn process(cfg: Config) -> Result<(), Error> {
    let mut c = {
        let mut cb = Consumer::from_hosts(cfg.brokers)
            .with_group(cfg.group)
            .with_fetch_max_wait_time(1000)
            .with_fetch_min_bytes(1_000)
            .with_fetch_max_bytes_per_partition(100_000)
            .with_retry_max_bytes_limit(1_000_000)
            .with_offset_storage(cfg.offset_storage);
        for topic in cfg.topics {
            cb = cb.with_topic(topic);
        }
        try!(cb.create())
    };

    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut buf = Vec::with_capacity(1024);

    let do_commit = !cfg.no_commit;
    loop {
        for ms in try!(c.poll()).iter() {
            for m in ms.messages() {
                // ~ clear the output buffer
                unsafe { buf.set_len(0) };
                // ~ format the message for output
                let _ = write!(buf, "{}:{}@{}:\n", ms.topic(), ms.partition(), m.offset);
                buf.extend_from_slice(m.value);
                buf.push(b'\n');
                // ~ write to output channel
                try!(stdout.write_all(&buf));
            }
            let _ = c.consume_messageset(ms);
        }
        if do_commit {
            try!(c.commit_consumed());
        }
    }
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    group: String,
    topics: Vec<String>,
    no_commit: bool,
    offset_storage: GroupOffsetStorage,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<_> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify topics (comma separated)", "NAMES");
        opts.optopt("", "group", "Specify the consumer group", "NAME");
        opts.optflag("", "no-commit", "Do not commit consumed messages");
        opts.optopt("", "offsets", "Specify the offset store [zookeeper, kafka]", "STORE");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }

        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("offsets") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            group: m.opt_str("group")
                .unwrap_or_else(|| String::new()),
            topics: m.opt_str("topic")
                .unwrap_or_else(|| String::new())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            no_commit: m.opt_present("no-commit"),
            offset_storage: offset_storage,
        })
   }
}
<commit_msg>console-consumer: make --brokers and --topics options required<commit_after>extern crate kafka;
extern crate getopts;
extern crate env_logger;

use std::{env, fmt, process};
use std::io::{self, Write};
use std::ascii::AsciiExt;

use kafka::consumer::{Consumer, GroupOffsetStorage};

/// This is a very simple command line application reading from a
/// specific kafka topic and dumping the messages to standard output.
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    };
    if let Err(e) = process(cfg) {
        println!("{}", e);
        process::exit(1);
    }
}

fn process(cfg: Config) -> Result<(), Error> {
    let mut c = {
        let mut cb = Consumer::from_hosts(cfg.brokers)
            .with_group(cfg.group)
            .with_fetch_max_wait_time(1000)
            .with_fetch_min_bytes(1_000)
            .with_fetch_max_bytes_per_partition(100_000)
            .with_retry_max_bytes_limit(1_000_000)
            .with_offset_storage(cfg.offset_storage);
        for topic in cfg.topics {
            cb = cb.with_topic(topic);
        }
        try!(cb.create())
    };

    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut buf = Vec::with_capacity(1024);

    let do_commit = !cfg.no_commit;
    loop {
        for ms in try!(c.poll()).iter() {
            for m in ms.messages() {
                // ~ clear the output buffer
                unsafe { buf.set_len(0) };
                // ~ format the message for output
                let _ = write!(buf, "{}:{}@{}:\n", ms.topic(), ms.partition(), m.offset);
                buf.extend_from_slice(m.value);
                buf.push(b'\n');
                // ~ write to output channel
                try!(stdout.write_all(&buf));
            }
            let _ = c.consume_messageset(ms);
        }
        if do_commit {
            try!(c.commit_consumed());
        }
    }
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    group: String,
    topics: Vec<String>,
    no_commit: bool,
    offset_storage: GroupOffsetStorage,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<_> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka brokers (comma separated)", "HOSTS");
        opts.optopt("", "topics", "Specify topics (comma separated)", "NAMES");
        opts.optopt("", "group", "Specify the consumer group", "NAME");
        opts.optflag("", "no-commit", "Do not commit consumed messages");
        opts.optopt("", "offsets", "Specify the offset store [zookeeper, kafka]", "STORE");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }

        macro_rules! required_list {
            ($name:expr) => {{
                let opt = $name;
                let xs: Vec<_> = match m.opt_str(opt) {
                    None => return Err(Error::Literal(format!("Required option --{} missing", opt))),
                    Some(s) => s.split(',').map(|s| s.trim().to_owned()).filter(|s| !s.is_empty()).collect(),
                };
                if xs.is_empty() {
                    return Err(Error::Literal(format!("Invalid --{} specified!", opt)));
                }
                xs
            }}
        };

        let brokers = required_list!("brokers");
        let topics = required_list!("topics");

        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("offsets") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        Ok(Config {
            brokers: brokers,
            group: m.opt_str("group").unwrap_or_else(|| String::new()),
            topics: topics,
            no_commit: m.opt_present("no-commit"),
            offset_storage: offset_storage,
        })
   }
}
<|endoftext|>


<commit_before>extern crate hyper;

use std::sync::{Arc, Mutex};

pub mod client;
mod contentlength;
pub mod download;
pub mod http_version;
pub mod write;

/// Represents a number of bytes, in `u64`.
type Byte = u64;
/// Represents a 'chunk', which is just a piece of bytes.
type Chunk = Vec<u8>;
/// Represents a list of chunks
type Chunks = Vec<Chunk>;
/// Represents a shared mutable reference of chunks
type SChunks = Arc<Mutex<Chunks>>;
<commit_msg>SChunks is now public<commit_after>extern crate hyper;

use std::sync::{Arc, Mutex};

pub mod client;
mod contentlength;
pub mod download;
pub mod http_version;
pub mod write;

/// Represents a number of bytes, in `u64`.
type Byte = u64;
/// Represents a 'chunk', which is just a piece of bytes.
type Chunk = Vec<u8>;
/// Represents a list of chunks
type Chunks = Vec<Chunk>;
/// Represents a shared mutable reference of chunks
pub type SChunks = Arc<Mutex<Chunks>>;
<|endoftext|>


<commit_before>extern crate libc;

use std::io::{Error, Result};

fn get_limit() -> Result<libc::rlimit> {
    unsafe {
        let mut rlim = libc::rlimit{rlim_cur: 0, rlim_max: 0};
        if libc::getrlimit(libc::RLIMIT_NOFILE, &mut rlim) != 0 {
            return Err(Error::last_os_error());
        }
        return Ok(rlim)
    }
}

pub fn get() -> Result<usize> {
    match get_limit() {
        Ok(rlim) => Ok(rlim.rlim_cur as usize),
        Err(err) => Err(err),
    }
}

pub enum MaxLimit {
    Val(usize),
    NoLimit,
}

pub fn max() -> Result<MaxLimit> {
    let rlim = match get_limit() {
        Ok(rlim) => rlim,
        Err(err) => return Err(err),
    };
    if rlim.rlim_max == libc::RLIM_INFINITY {
        Ok( MaxLimit::NoLimit )
    } else {
        Ok( MaxLimit::Val(rlim.rlim_max as usize))
    }
}

pub fn set_to_max() -> Result<usize> {
    let mut r = match get_limit() {
        Ok(l) => l,
        Err(err) => return Err(err),
    };
    if r.rlim_max != r.rlim_cur {
        if r.rlim_max == libc::RLIM_INFINITY { // no limit restriction
            r.rlim_cur *= 8;
        } else if r.rlim_cur < r.rlim_max{
            r.rlim_cur = r.rlim_max;
        }
        unsafe {
            if libc::setrlimit(libc::RLIMIT_NOFILE, &r) < 0 {
                let err = Error::last_os_error();
                return Err(err);
            }
        }
    }

    Ok(r.rlim_cur as usize)
}

#[test]
fn get_file_limit() {
    let lim = get().unwrap();
    assert!(lim > 0);
}

#[test]
fn max_limit() {
    match max() {
        Ok(MaxLimit::Val(v)) => {
            assert!(v > 0);
        },
        Ok(MaxLimit::NoLimit) => assert!(true),
        Err(err) => assert!(false),
    }
}

#[test]
fn set_max() {
    let cur_limit: usize = get().unwrap();
    let max_limit = set_to_max().unwrap();
    let new_limit: usize = get().unwrap();
    assert_eq!(max_limit, new_limit);
    assert!(new_limit >= cur_limit);
}
<commit_msg>add annotations for cfg and inline<commit_after>#![crate_name = "file_limit"]
#![cfg(linux)]

extern crate libc;

use std::io::{Error, Result};

// private fn to get the kernal struct for file limits
#[inline]
fn get_limit() -> Result<libc::rlimit> {
    unsafe {
        let mut rlim = libc::rlimit{rlim_cur: 0, rlim_max: 0};
        if libc::getrlimit(libc::RLIMIT_NOFILE, &mut rlim) != 0 {
            return Err(Error::last_os_error());
        }
        return Ok(rlim)
    }
}

#[inline]
pub fn get() -> Result<usize> {
    match get_limit() {
        Ok(rlim) => Ok(rlim.rlim_cur as usize),
        Err(err) => Err(err),
    }
}

pub enum MaxLimit {
    Val(usize),
    NoLimit,
}

#[inline]
pub fn max() -> Result<MaxLimit> {
    let rlim = match get_limit() {
        Ok(rlim) => rlim,
        Err(err) => return Err(err),
    };
    if rlim.rlim_max == libc::RLIM_INFINITY {
        Ok( MaxLimit::NoLimit )
    } else {
        Ok( MaxLimit::Val(rlim.rlim_max as usize))
    }
}

#[inline]
pub fn set_to_max() -> Result<usize> {
    let mut r = match get_limit() {
        Ok(l) => l,
        Err(err) => return Err(err),
    };
    if r.rlim_max != r.rlim_cur {
        if r.rlim_max == libc::RLIM_INFINITY { // no limit restriction
            r.rlim_cur *= 8;
        } else if r.rlim_cur < r.rlim_max{
            r.rlim_cur = r.rlim_max;
        }
        unsafe {
            if libc::setrlimit(libc::RLIMIT_NOFILE, &r) < 0 {
                let err = Error::last_os_error();
                return Err(err);
            }
        }
    }

    Ok(r.rlim_cur as usize)
}

#[test]
fn get_file_limit() {
    let lim = get().unwrap();
    assert!(lim > 0);
}

#[test]
fn max_limit() {
    match max() {
        Ok(MaxLimit::Val(v)) => {
            assert!(v > 0);
        },
        Ok(MaxLimit::NoLimit) => assert!(true),
        Err(err) => assert!(false),
    }
}

#[test]
fn set_max() {
    let cur_limit: usize = get().unwrap();
    let max_limit = set_to_max().unwrap();
    let new_limit: usize = get().unwrap();
    assert_eq!(max_limit, new_limit);
    assert!(new_limit >= cur_limit);
}
<|endoftext|>


<commit_before>//! A crate that provides support for half-precision 16-bit floating point types.
//!
//! This crate provides the [`f16`] type, which is an implementation of the IEEE 754-2008 standard
//! [`binary16`] a.k.a `half` floating point type. This 16-bit floating point type is intended for
//! efficient storage where the full range and precision of a larger floating point value is not
//! required. This is especially useful for image storage formats.
//!
//! This crate also provides a [`bf16`] type in the [`bfloat`] module, an alternative 16-bit floating
//! point format. The [`bfloat16`] format is a truncated IEEE 754 standard `binary32` float that
//! preserves the exponent to allow the same range as `f32` but with only 8 bits of precision
//! (instead of 11 bits for [`f16`]). See the [`bfloat`] module for details.
//!
//! Because [`f16`] and [`bf16`] are primarily for efficient storage, floating point operations such as
//! addition, multiplication, etc. are not implemented. Operations should be performed with `f32`
//! or higher-precision types and converted to/from [`f16`] or [`bf16`] as necessary.
//!
//! This crate also provides a [`slice`] module for zero-copy in-place conversions of `u16` slices
//! to both [`f16`] and [`bf16`], as well as efficient vectorized conversions of larger buffers of
//! floating point values to and from these half formats.
//!
//! A [`prelude`] module is provided for easy importing of available utility traits.
//!
//! Some hardware architectures provide support for 16-bit floating point conversions. Enable the
//! `use-intrinsics` feature to use LLVM intrinsics for hardware conversions. This crate does no
//! checks on whether the hardware supports the feature. This feature currently only works on
//! nightly Rust due to a compiler feature gate. When this feature is enabled and the hardware
//! supports it, the [`slice`] trait conversions will use vectorized SIMD intructions for
//! increased efficiency.
//!
//! Support for `serde` crate `Serialize` and `Deserialize` traits is provided when the `serde`
//! feature is enabled. This adds a dependency on `serde` crate so is an optional cargo feature.
//!
//! The crate uses `#[no_std]` by default, so can be used in embedded environments without using the
//! Rust `std` library. A `std` feature is available, which enables additional utilities using the
//! `std` library, such as the `vec` module that provides zero-copy `Vec` conversions.
//!
//! [`f16`]: struct.f16.html
//! [`binary16`]: https://en.wikipedia.org/wiki/Half-precision_floating-point_format
//! [`bf16`]: bfloat/struct.bf16.html
//! [`bfloat`]: bfloat/index.html
//! [`bfloat16`]: https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
//! [`slice`]: slice/index.html
//! [`prelude`]: prelude/index.html

#![warn(
    missing_docs,
    missing_copy_implementations,
    missing_debug_implementations,
    trivial_numeric_casts,
    unused_extern_crates,
    unused_import_braces,
    future_incompatible,
    rust_2018_compatibility,
    rust_2018_idioms,
    clippy::all
)]
#![allow(clippy::verbose_bit_mask, clippy::cast_lossless)]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(
    all(
        feature = "use-intrinsics",
        any(target_arch = "x86", target_arch = "x86_64")
    ),
    feature(stdsimd, f16c_target_feature)
)]
#![doc(html_root_url = "https://docs.rs/half/1.4.0")]

pub mod bfloat;
mod binary16;
pub mod slice;
#[cfg(feature = "std")]
pub mod vec;

pub use binary16::{consts, f16};

/// A collection of the most used items and traits in this crate for easy importing.
///
/// # Examples
///
/// ```rust
/// use half::prelude::*;
/// ```
pub mod prelude {
    #[doc(no_inline)]
    pub use crate::{
        bfloat::bf16,
        f16,
        slice::{HalfBitsSliceExt, HalfFloatSliceExt},
    };

    #[cfg(feature = "std")]
    pub use crate::vec::{HalfBitsVecExt, HalfFloatVecExt};
}

// Keep this module private to crate
pub(crate) mod private {
    use crate::{bfloat::bf16, f16};

    pub trait SealedHalf {}

    impl SealedHalf for f16 {}
    impl SealedHalf for bf16 {}
}
<commit_msg>add link so serde in crate doc<commit_after>//! A crate that provides support for half-precision 16-bit floating point types.
//!
//! This crate provides the [`f16`] type, which is an implementation of the IEEE 754-2008 standard
//! [`binary16`] a.k.a `half` floating point type. This 16-bit floating point type is intended for
//! efficient storage where the full range and precision of a larger floating point value is not
//! required. This is especially useful for image storage formats.
//!
//! This crate also provides a [`bf16`] type in the [`bfloat`] module, an alternative 16-bit floating
//! point format. The [`bfloat16`] format is a truncated IEEE 754 standard `binary32` float that
//! preserves the exponent to allow the same range as `f32` but with only 8 bits of precision
//! (instead of 11 bits for [`f16`]). See the [`bfloat`] module for details.
//!
//! Because [`f16`] and [`bf16`] are primarily for efficient storage, floating point operations such as
//! addition, multiplication, etc. are not implemented. Operations should be performed with `f32`
//! or higher-precision types and converted to/from [`f16`] or [`bf16`] as necessary.
//!
//! This crate also provides a [`slice`] module for zero-copy in-place conversions of `u16` slices
//! to both [`f16`] and [`bf16`], as well as efficient vectorized conversions of larger buffers of
//! floating point values to and from these half formats.
//!
//! A [`prelude`] module is provided for easy importing of available utility traits.
//!
//! Some hardware architectures provide support for 16-bit floating point conversions. Enable the
//! `use-intrinsics` feature to use LLVM intrinsics for hardware conversions. This crate does no
//! checks on whether the hardware supports the feature. This feature currently only works on
//! nightly Rust due to a compiler feature gate. When this feature is enabled and the hardware
//! supports it, the [`slice`] trait conversions will use vectorized SIMD intructions for
//! increased efficiency.
//!
//! Support for [`serde`] crate `Serialize` and `Deserialize` traits is provided when the `serde`
//! feature is enabled. This adds a dependency on [`serde`] crate so is an optional cargo feature.
//!
//! The crate uses `#[no_std]` by default, so can be used in embedded environments without using the
//! Rust `std` library. A `std` feature is available, which enables additional utilities using the
//! `std` library, such as the `vec` module that provides zero-copy `Vec` conversions.
//!
//! [`f16`]: struct.f16.html
//! [`binary16`]: https://en.wikipedia.org/wiki/Half-precision_floating-point_format
//! [`bf16`]: bfloat/struct.bf16.html
//! [`bfloat`]: bfloat/index.html
//! [`bfloat16`]: https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
//! [`slice`]: slice/index.html
//! [`prelude`]: prelude/index.html
//! [`serde`]: https://crates.io/crates/serde

#![warn(
    missing_docs,
    missing_copy_implementations,
    missing_debug_implementations,
    trivial_numeric_casts,
    unused_extern_crates,
    unused_import_braces,
    future_incompatible,
    rust_2018_compatibility,
    rust_2018_idioms,
    clippy::all
)]
#![allow(clippy::verbose_bit_mask, clippy::cast_lossless)]
#![cfg_attr(not(feature = "std"), no_std)]
#![cfg_attr(
    all(
        feature = "use-intrinsics",
        any(target_arch = "x86", target_arch = "x86_64")
    ),
    feature(stdsimd, f16c_target_feature)
)]
#![doc(html_root_url = "https://docs.rs/half/1.4.0")]

pub mod bfloat;
mod binary16;
pub mod slice;
#[cfg(feature = "std")]
pub mod vec;

pub use binary16::{consts, f16};

/// A collection of the most used items and traits in this crate for easy importing.
///
/// # Examples
///
/// ```rust
/// use half::prelude::*;
/// ```
pub mod prelude {
    #[doc(no_inline)]
    pub use crate::{
        bfloat::bf16,
        f16,
        slice::{HalfBitsSliceExt, HalfFloatSliceExt},
    };

    #[cfg(feature = "std")]
    pub use crate::vec::{HalfBitsVecExt, HalfFloatVecExt};
}

// Keep this module private to crate
pub(crate) mod private {
    use crate::{bfloat::bf16, f16};

    pub trait SealedHalf {}

    impl SealedHalf for f16 {}
    impl SealedHalf for bf16 {}
}
<|endoftext|>