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

use arg_enum_proc_macro::ArgEnum;
use num_derive::FromPrimitive;

/// Sample position for subsampled chroma
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSamplePosition {
  /// The source video transfer function must be signaled
  /// outside the AV1 bitstream.
  Unknown,
  /// Horizontally co-located with (0, 0) luma sample, vertically positioned
  /// in the middle between two luma samples.
  Vertical,
  /// Co-located with (0, 0) luma sample.
  Colocated,
}

impl Default for ChromaSamplePosition {
  fn default() -> Self {
    ChromaSamplePosition::Unknown
  }
}

/// Chroma subsampling format
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSampling {
  /// Both vertically and horizontally subsampled.
  Cs420,
  /// Horizontally subsampled.
  Cs422,
  /// Not subsampled.
  Cs444,
  /// Monochrome.
  Cs400,
}

impl Default for ChromaSampling {
  fn default() -> Self {
    ChromaSampling::Cs420
  }
}

impl ChromaSampling {
  /// Provides the amount to right shift the luma plane dimensions to get the
  ///  chroma plane dimensions.
  /// Only values 0 or 1 are ever returned.
  /// The plane dimensions must also be rounded up to accommodate odd luma plane
  ///  sizes.
  /// Cs400 returns None, as there are no chroma planes.
  pub fn get_decimation(self) -> Option<(usize, usize)> {
    use self::ChromaSampling::*;
    match self {
      Cs420 => Some((1, 1)),
      Cs422 => Some((1, 0)),
      Cs444 => Some((0, 0)),
      Cs400 => None,
    }
  }

  /// Calculates the size of a chroma plane for this sampling type, given the luma plane dimensions.
  pub fn get_chroma_dimensions(
    self, luma_width: usize, luma_height: usize,
  ) -> (usize, usize) {
    if let Some((ss_x, ss_y)) = self.get_decimation() {
      ((luma_width + ss_x) >> ss_x, (luma_height + ss_y) >> ss_y)
    } else {
      (0, 0)
    }
  }
}

/// Supported Color Primaries
///
/// As defined by “Color primaries” section of ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ColorPrimaries {
  /// BT.709
  BT709 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240M (historical)
  ST240M,
  /// Generic film
  Film,
  /// BT.2020, BT.2100
  BT2020,
  /// SMPTE 248 (CIE 1921 XYZ)
  ST428,
  /// SMPTE RP 431-2
  P3DCI,
  /// SMPTE EG 432-1
  P3Display,
  /// EBU Tech. 3213-E
  Tech3213 = 22,
}

impl Default for ColorPrimaries {
  fn default() -> Self {
    ColorPrimaries::Unspecified
  }
}

/// Supported Transfer Characteristics
///
/// As defined by “Transfer characteristics” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum TransferCharacteristics {
  /// BT.709
  BT1886 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240 M
  ST240M,
  /// Linear
  Linear,
  /// Logarithmic (100:1 range)
  Logarithmic100,
  /// Logarithmic ((100 * √10):1 range)
  Logarithmic316,
  /// IEC 61966-2-4
  XVYCC,
  /// BT.1361 extended color gamut system (historical)
  BT1361E,
  /// sRGB or sYCC
  SRGB,
  /// BT.2020 10-bit systems
  BT2020Ten,
  /// BT.2020 12-bit systems
  BT2020Twelve,
  /// SMPTE ST 2084, ITU BT.2100 PQ
  PerceptualQuantizer,
  /// SMPTE ST 428
  ST428,
  /// BT.2100 HLG, ARIB STD-B67
  HybridLogGamma,
}

impl Default for TransferCharacteristics {
  fn default() -> Self {
    TransferCharacteristics::Unspecified
  }
}

/// Matrix coefficients
///
/// As defined by the “Matrix coefficients” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum MatrixCoefficients {
  /// Identity matrix
  Identity = 0,
  /// BT.709
  BT709,
  /// Unspecified, must be signaled or inferred outside of the bitstream.
  Unspecified,
  /// US FCC 73.628
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  ST170M,
  /// SMPTE 240 M
  ST240M,
  /// YCgCo
  YCgCo,
  /// BT.2020 non-constant luminance, BT.2100 YCbCr
  BT2020NonConstantLuminance,
  /// BT.2020 constant luminance
  BT2020ConstantLuminance,
  /// SMPTE ST 2085 YDzDx
  ST2085,
  /// Chromaticity-derived non-constant luminance
  ChromaticityDerivedNonConstantLuminance,
  /// Chromaticity-derived constant luminance
  ChromaticityDerivedConstantLuminance,
  /// BT.2020 ICtCp
  ICtCp,
}

impl Default for MatrixCoefficients {
  fn default() -> Self {
    MatrixCoefficients::Unspecified
  }
}

/// Signal the content color description
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
  pub color_primaries: ColorPrimaries,
  pub transfer_characteristics: TransferCharacteristics,
  pub matrix_coefficients: MatrixCoefficients,
}

/// Allowed pixel value range
///
/// C.f. VideoFullRangeFlag variable specified in ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum PixelRange {
  /// Studio swing representation
  Limited,
  /// Full swing representation
  Full,
}

impl Default for PixelRange {
  fn default() -> Self {
    PixelRange::Limited
  }
}

/// High dynamic range content light level
///
/// As defined by CEA-861.3, Appendix A.
#[derive(Copy, Clone, Debug)]
pub struct ContentLight {
  /// Maximum content light level
  pub max_content_light_level: u16,
  /// Maximum frame-average light level
  pub max_frame_average_light_level: u16,
}

/// Chromaticity coordinates expressed as 0.16 fixed-point values
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16,
}

/// High dynamic range mastering display color volume
///
/// As defined by CIE 1931
#[derive(Copy, Clone, Debug)]
pub struct MasteringDisplay {
  /// Chromaticity coordinates in Red, Green, Blue order
  /// expressed as 0.16 fixed-point
  pub primaries: [Point; 3],
  /// Chromaticity coordinates expressed as 0.16 fixed-point
  pub white_point: Point,
  /// 24.8 fixed-point maximum luminance in candelas per square meter
  pub max_luminance: u32,
  /// 18.14 fixed-point minimum luminance in candelas per square meter
  pub min_luminance: u32,
}
<commit_msg>Use AV1 spec names for color config enum variants (#1544)<commit_after>// Copyright (c) 2019, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use arg_enum_proc_macro::ArgEnum;
use num_derive::FromPrimitive;

/// Sample position for subsampled chroma
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSamplePosition {
  /// The source video transfer function must be signaled
  /// outside the AV1 bitstream.
  Unknown,
  /// Horizontally co-located with (0, 0) luma sample, vertically positioned
  /// in the middle between two luma samples.
  Vertical,
  /// Co-located with (0, 0) luma sample.
  Colocated,
}

impl Default for ChromaSamplePosition {
  fn default() -> Self {
    ChromaSamplePosition::Unknown
  }
}

/// Chroma subsampling format
#[derive(Copy, Clone, Debug, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ChromaSampling {
  /// Both vertically and horizontally subsampled.
  Cs420,
  /// Horizontally subsampled.
  Cs422,
  /// Not subsampled.
  Cs444,
  /// Monochrome.
  Cs400,
}

impl Default for ChromaSampling {
  fn default() -> Self {
    ChromaSampling::Cs420
  }
}

impl ChromaSampling {
  /// Provides the amount to right shift the luma plane dimensions to get the
  ///  chroma plane dimensions.
  /// Only values 0 or 1 are ever returned.
  /// The plane dimensions must also be rounded up to accommodate odd luma plane
  ///  sizes.
  /// Cs400 returns None, as there are no chroma planes.
  pub fn get_decimation(self) -> Option<(usize, usize)> {
    use self::ChromaSampling::*;
    match self {
      Cs420 => Some((1, 1)),
      Cs422 => Some((1, 0)),
      Cs444 => Some((0, 0)),
      Cs400 => None,
    }
  }

  /// Calculates the size of a chroma plane for this sampling type, given the luma plane dimensions.
  pub fn get_chroma_dimensions(
    self, luma_width: usize, luma_height: usize,
  ) -> (usize, usize) {
    if let Some((ss_x, ss_y)) = self.get_decimation() {
      ((luma_width + ss_x) >> ss_x, (luma_height + ss_y) >> ss_y)
    } else {
      (0, 0)
    }
  }
}

/// Supported Color Primaries
///
/// As defined by “Color primaries” section of ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum ColorPrimaries {
  /// BT.709
  BT709 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  BT601,
  /// SMPTE 240M (historical)
  SMPTE240,
  /// Generic film
  GenericFilm,
  /// BT.2020, BT.2100
  BT2020,
  /// SMPTE 248 (CIE 1921 XYZ)
  XYZ,
  /// SMPTE RP 431-2
  SMPTE431,
  /// SMPTE EG 432-1
  SMPTE432,
  /// EBU Tech. 3213-E
  EBU3213 = 22,
}

impl Default for ColorPrimaries {
  fn default() -> Self {
    ColorPrimaries::Unspecified
  }
}

/// Supported Transfer Characteristics
///
/// As defined by “Transfer characteristics” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum TransferCharacteristics {
  /// BT.709
  BT709 = 1,
  /// Unspecified, must be signaled or inferred outside of the bitstream
  Unspecified,
  /// BT.470 System M (historical)
  BT470M = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  BT601,
  /// SMPTE 240 M
  SMPTE240,
  /// Linear
  Linear,
  /// Logarithmic (100:1 range)
  Log100,
  /// Logarithmic ((100 * √10):1 range)
  Log100Sqrt10,
  /// IEC 61966-2-4
  IEC61966,
  /// BT.1361 extended color gamut system (historical)
  BT1361,
  /// sRGB or sYCC
  SRGB,
  /// BT.2020 10-bit systems
  BT2020_10Bit,
  /// BT.2020 12-bit systems
  BT2020_12Bit,
  /// SMPTE ST 2084, ITU BT.2100 PQ
  SMPTE2084,
  /// SMPTE ST 428
  SMPTE428,
  /// BT.2100 HLG (Hybrid Log Gamma), ARIB STD-B67
  HLG,
}

impl Default for TransferCharacteristics {
  fn default() -> Self {
    TransferCharacteristics::Unspecified
  }
}

/// Matrix coefficients
///
/// As defined by the “Matrix coefficients” section of ISO/IEC 23091-4/ITU-TH.273.
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum MatrixCoefficients {
  /// Identity matrix
  Identity = 0,
  /// BT.709
  BT709,
  /// Unspecified, must be signaled or inferred outside of the bitstream.
  Unspecified,
  /// US FCC 73.628
  FCC = 4,
  /// BT.470 System B, G (historical)
  BT470BG,
  /// BT.601-7 525 (SMPTE 170 M)
  BT601,
  /// SMPTE 240 M
  SMPTE240,
  /// YCgCo
  YCgCo,
  /// BT.2020 non-constant luminance, BT.2100 YCbCr
  BT2020NCL,
  /// BT.2020 constant luminance
  BT2020CL,
  /// SMPTE ST 2085 YDzDx
  SMPTE2085,
  /// Chromaticity-derived non-constant luminance
  ChromatNCL,
  /// Chromaticity-derived constant luminance
  ChromatCL,
  /// BT.2020 ICtCp
  ICtCp,
}

impl Default for MatrixCoefficients {
  fn default() -> Self {
    MatrixCoefficients::Unspecified
  }
}

/// Signal the content color description
#[derive(Copy, Clone, Debug)]
pub struct ColorDescription {
  pub color_primaries: ColorPrimaries,
  pub transfer_characteristics: TransferCharacteristics,
  pub matrix_coefficients: MatrixCoefficients,
}

/// Allowed pixel value range
///
/// C.f. VideoFullRangeFlag variable specified in ISO/IEC 23091-4/ITU-T H.273
#[derive(ArgEnum, Debug, Clone, Copy, PartialEq, FromPrimitive)]
#[repr(C)]
pub enum PixelRange {
  /// Studio swing representation
  Limited,
  /// Full swing representation
  Full,
}

impl Default for PixelRange {
  fn default() -> Self {
    PixelRange::Limited
  }
}

/// High dynamic range content light level
///
/// As defined by CEA-861.3, Appendix A.
#[derive(Copy, Clone, Debug)]
pub struct ContentLight {
  /// Maximum content light level
  pub max_content_light_level: u16,
  /// Maximum frame-average light level
  pub max_frame_average_light_level: u16,
}

/// Chromaticity coordinates expressed as 0.16 fixed-point values
#[derive(Clone, Copy, Debug)]
#[repr(C)]
pub struct Point {
  pub x: u16,
  pub y: u16,
}

/// High dynamic range mastering display color volume
///
/// As defined by CIE 1931
#[derive(Copy, Clone, Debug)]
pub struct MasteringDisplay {
  /// Chromaticity coordinates in Red, Green, Blue order
  /// expressed as 0.16 fixed-point
  pub primaries: [Point; 3],
  /// Chromaticity coordinates expressed as 0.16 fixed-point
  pub white_point: Point,
  /// 24.8 fixed-point maximum luminance in candelas per square meter
  pub max_luminance: u32,
  /// 18.14 fixed-point minimum luminance in candelas per square meter
  pub min_luminance: u32,
}
<|endoftext|>


<commit_before>use self::StartsWithStateInternal::*;

/// Automaton describes types that behave as a finite automaton.
///
/// All implementators of this trait are represented by *byte based* automata.
/// Stated differently, all transitions in the automata correspond to a single
/// byte in the input.
///
/// This implementation choice is important for a couple reasons:
///
/// 1. The set of possible transitions in each node is small, which may make
///    efficient memory usage easier.
/// 2. The finite state transducers in this crate are all byte based, so any
///    automata used on them must also be byte based.
///
/// In practice, this does present somewhat of a problem, for example, if
/// you're storing UTF-8 encoded strings in a finite state transducer. Consider
/// using a `Levenshtein` automaton, which accepts a query string and an edit
/// distance. The edit distance should apply to some notion of *character*,
/// which could be represented by at least 1-4 bytes in a UTF-8 encoding (for
/// some definition of "character"). Therefore, the automaton must have UTF-8
/// decoding built into it. This can be tricky to implement, so you may find
/// the [`utf8-ranges`](https://crates.io/crates/utf8-ranges) crate useful.
pub trait Automaton {
    /// The type of the state used in the automaton.
    type State;

    /// Returns a single start state for this automaton.
    ///
    /// This method should always return the same value for each
    /// implementation.
    fn start(&self) -> Self::State;

    /// Returns true if and only if `state` is a match state.
    fn is_match(&self, state: &Self::State) -> bool;

    /// Returns true if and only if `state` can lead to a match in zero or more
    /// steps.
    ///
    /// If this returns `false`, then no sequence of inputs from this state
    /// should ever produce a match. If this does not follow, then those match
    /// states may never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `true` even when no match is possible, then behavior
    /// will be correct, but callers may be forced to do additional work.
    fn can_match(&self, _state: &Self::State) -> bool {
        true
    }

    /// Returns true if and only if `state` matches and must match no matter what
    /// steps are taken.
    ///
    /// If this returns `true`, then every sequence of inputs from this state
    /// produces a match. If this does not follow, then those match states may
    /// never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `false` even when every sequence of inputs will lead to
    /// a match, then behavior will be correct, but callers may be forced to do
    /// additional work.
    fn will_always_match(&self, _state: &Self::State) -> bool {
        false
    }

    /// Return the next state given `state` and an input.
    fn accept(&self, state: &Self::State, byte: u8) -> Self::State;

    /// Returns an automaton that matches the strings that start with something
    /// this automaton matches.
    fn starts_with(self) -> StartsWith<Self> where Self: Sized {
        StartsWith(self)
    }

    /// Returns an automaton that matches the strings matched by either this or
    /// the other automaton.
    fn union<Rhs: Automaton>(self, rhs: Rhs) -> Union<Self, Rhs> where Self: Sized {
        Union(self, rhs)
    }

    /// Returns an automaton that matches the strings matched by both this and
    /// the other automaton.
    fn intersection<Rhs: Automaton>(self, rhs: Rhs) -> Intersection<Self, Rhs> where Self: Sized {
        Intersection(self, rhs)
    }

    /// Returns an automaton that matches the strings not matched by this automaton.
    fn complement(self) -> Complement<Self> where Self: Sized {
        Complement(self)
    }
}

impl<'a, T: Automaton> Automaton for &'a T {
    type State = T::State;

    fn start(&self) -> Self::State {
        (*self).start()
    }

    fn is_match(&self, state: &Self::State) -> bool {
        (*self).is_match(state)
    }

    fn can_match(&self, state: &Self::State) -> bool {
        (*self).can_match(state)
    }

    fn will_always_match(&self, state: &Self::State) -> bool {
        (*self).will_always_match(state)
    }

    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
        (*self).accept(state, byte)
    }
}

/// An automaton that always matches.
///
/// This is useful in a generic context as a way to express that no automaton
/// should be used.
pub struct AlwaysMatch;

impl Automaton for AlwaysMatch {
    type State = ();

    fn start(&self) -> () { () }
    fn is_match(&self, _: &()) -> bool { true }
    fn can_match(&self, _: &()) -> bool { true }
    fn will_always_match(&self, _: &()) -> bool { true }
    fn accept(&self, _: &(), _: u8) -> () { () }
}

/// An automaton that matches and string that begins with something that the
/// wrapped automaton matches.
pub struct StartsWith<A>(A);

/// The `Automaton` state for `StartsWith<A>`.
pub struct StartsWithState<A: Automaton>(StartsWithStateInternal<A>);

enum StartsWithStateInternal<A: Automaton> {
    Done,
    Running(A::State)
}

impl<A: Automaton> Automaton for StartsWith<A> {
    type State = StartsWithState<A>;
    fn start(&self) -> StartsWithState<A> {
        StartsWithState({
            let inner = self.0.start();
            if self.0.is_match(&inner) {
                Done
            } else {
                Running(inner)
            }
        })
    }

    fn is_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn can_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(ref inner) => self.0.can_match(inner)
        }
    }

    fn will_always_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn accept(&self, state: &StartsWithState<A>, byte: u8) -> StartsWithState<A> {
        StartsWithState(
            match state.0 {
                Done => Done,
                Running(ref inner) => {
                    let next_inner = self.0.accept(inner, byte);
                    if self.0.is_match(&next_inner) {
                        Done
                    } else {
                        Running(next_inner)
                    }
                }
            }
        )
    }
}

/// An automaton that matches when one of its component automata match.
pub struct Union<A, B>(A, B);

/// The `Automaton` state for `Union<A, B>`.
pub struct UnionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Union<A, B> {
    type State = UnionState<A, B>;

    fn start(&self) -> UnionState<A, B> {
        UnionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.is_match(&state.0) || self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.can_match(&state.0) || self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) || self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &UnionState<A, B>, byte: u8) -> UnionState<A, B> {
        UnionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches when both of its component automata match.
pub struct Intersection<A, B>(A, B);

/// The `Automaton` state for `Intersection<A, B>`.
pub struct IntersectionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Intersection<A, B> {
    type State = IntersectionState<A, B>;

    fn start(&self) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.is_match(&state.0) && self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.can_match(&state.0) && self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) && self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &IntersectionState<A, B>, byte: u8) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches exactly when the automaton it wraps does not.
pub struct Complement<A>(A);

/// The `Automaton` state for `Complement<A>`.
pub struct ComplementState<A: Automaton>(A::State);

impl<A: Automaton> Automaton for Complement<A> {
    type State = ComplementState<A>;

    fn start(&self) -> ComplementState<A> {
        ComplementState(self.0.start())
    }

    fn is_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.is_match(&state.0)
    }

    fn can_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.will_always_match(&state.0)
    }

    fn will_always_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.can_match(&state.0)
    }

    fn accept(&self, state: &ComplementState<A>, byte: u8) -> ComplementState<A> {
        ComplementState(self.0.accept(&state.0, byte))
    }
}
<commit_msg>Add tests for the automaton adapters<commit_after>use self::StartsWithStateInternal::*;

/// Automaton describes types that behave as a finite automaton.
///
/// All implementators of this trait are represented by *byte based* automata.
/// Stated differently, all transitions in the automata correspond to a single
/// byte in the input.
///
/// This implementation choice is important for a couple reasons:
///
/// 1. The set of possible transitions in each node is small, which may make
///    efficient memory usage easier.
/// 2. The finite state transducers in this crate are all byte based, so any
///    automata used on them must also be byte based.
///
/// In practice, this does present somewhat of a problem, for example, if
/// you're storing UTF-8 encoded strings in a finite state transducer. Consider
/// using a `Levenshtein` automaton, which accepts a query string and an edit
/// distance. The edit distance should apply to some notion of *character*,
/// which could be represented by at least 1-4 bytes in a UTF-8 encoding (for
/// some definition of "character"). Therefore, the automaton must have UTF-8
/// decoding built into it. This can be tricky to implement, so you may find
/// the [`utf8-ranges`](https://crates.io/crates/utf8-ranges) crate useful.
pub trait Automaton {
    /// The type of the state used in the automaton.
    type State;

    /// Returns a single start state for this automaton.
    ///
    /// This method should always return the same value for each
    /// implementation.
    fn start(&self) -> Self::State;

    /// Returns true if and only if `state` is a match state.
    fn is_match(&self, state: &Self::State) -> bool;

    /// Returns true if and only if `state` can lead to a match in zero or more
    /// steps.
    ///
    /// If this returns `false`, then no sequence of inputs from this state
    /// should ever produce a match. If this does not follow, then those match
    /// states may never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `true` even when no match is possible, then behavior
    /// will be correct, but callers may be forced to do additional work.
    fn can_match(&self, _state: &Self::State) -> bool {
        true
    }

    /// Returns true if and only if `state` matches and must match no matter what
    /// steps are taken.
    ///
    /// If this returns `true`, then every sequence of inputs from this state
    /// produces a match. If this does not follow, then those match states may
    /// never be reached. In other words, behavior may be incorrect.
    ///
    /// If this returns `false` even when every sequence of inputs will lead to
    /// a match, then behavior will be correct, but callers may be forced to do
    /// additional work.
    fn will_always_match(&self, _state: &Self::State) -> bool {
        false
    }

    /// Return the next state given `state` and an input.
    fn accept(&self, state: &Self::State, byte: u8) -> Self::State;

    /// Returns an automaton that matches the strings that start with something
    /// this automaton matches.
    fn starts_with(self) -> StartsWith<Self> where Self: Sized {
        StartsWith(self)
    }

    /// Returns an automaton that matches the strings matched by either this or
    /// the other automaton.
    fn union<Rhs: Automaton>(self, rhs: Rhs) -> Union<Self, Rhs> where Self: Sized {
        Union(self, rhs)
    }

    /// Returns an automaton that matches the strings matched by both this and
    /// the other automaton.
    fn intersection<Rhs: Automaton>(self, rhs: Rhs) -> Intersection<Self, Rhs> where Self: Sized {
        Intersection(self, rhs)
    }

    /// Returns an automaton that matches the strings not matched by this automaton.
    fn complement(self) -> Complement<Self> where Self: Sized {
        Complement(self)
    }
}

impl<'a, T: Automaton> Automaton for &'a T {
    type State = T::State;

    fn start(&self) -> Self::State {
        (*self).start()
    }

    fn is_match(&self, state: &Self::State) -> bool {
        (*self).is_match(state)
    }

    fn can_match(&self, state: &Self::State) -> bool {
        (*self).can_match(state)
    }

    fn will_always_match(&self, state: &Self::State) -> bool {
        (*self).will_always_match(state)
    }

    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
        (*self).accept(state, byte)
    }
}

/// An automaton that always matches.
///
/// This is useful in a generic context as a way to express that no automaton
/// should be used.
pub struct AlwaysMatch;

impl Automaton for AlwaysMatch {
    type State = ();

    fn start(&self) -> () { () }
    fn is_match(&self, _: &()) -> bool { true }
    fn can_match(&self, _: &()) -> bool { true }
    fn will_always_match(&self, _: &()) -> bool { true }
    fn accept(&self, _: &(), _: u8) -> () { () }
}

/// An automaton that matches and string that begins with something that the
/// wrapped automaton matches.
pub struct StartsWith<A>(A);

/// The `Automaton` state for `StartsWith<A>`.
pub struct StartsWithState<A: Automaton>(StartsWithStateInternal<A>);

enum StartsWithStateInternal<A: Automaton> {
    Done,
    Running(A::State)
}

impl<A: Automaton> Automaton for StartsWith<A> {
    type State = StartsWithState<A>;
    fn start(&self) -> StartsWithState<A> {
        StartsWithState({
            let inner = self.0.start();
            if self.0.is_match(&inner) {
                Done
            } else {
                Running(inner)
            }
        })
    }

    fn is_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn can_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(ref inner) => self.0.can_match(inner)
        }
    }

    fn will_always_match(&self, state: &StartsWithState<A>) -> bool {
        match state.0 {
            Done => true,
            Running(_) => false
        }
    }

    fn accept(&self, state: &StartsWithState<A>, byte: u8) -> StartsWithState<A> {
        StartsWithState(
            match state.0 {
                Done => Done,
                Running(ref inner) => {
                    let next_inner = self.0.accept(inner, byte);
                    if self.0.is_match(&next_inner) {
                        Done
                    } else {
                        Running(next_inner)
                    }
                }
            }
        )
    }
}

/// An automaton that matches when one of its component automata match.
pub struct Union<A, B>(A, B);

/// The `Automaton` state for `Union<A, B>`.
pub struct UnionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Union<A, B> {
    type State = UnionState<A, B>;

    fn start(&self) -> UnionState<A, B> {
        UnionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.is_match(&state.0) || self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.can_match(&state.0) || self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &UnionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) || self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &UnionState<A, B>, byte: u8) -> UnionState<A, B> {
        UnionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches when both of its component automata match.
pub struct Intersection<A, B>(A, B);

/// The `Automaton` state for `Intersection<A, B>`.
pub struct IntersectionState<A: Automaton, B: Automaton>(A::State, B::State);

impl<A: Automaton, B: Automaton> Automaton for Intersection<A, B> {
    type State = IntersectionState<A, B>;

    fn start(&self) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.start(),
            self.1.start()
        )
    }

    fn is_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.is_match(&state.0) && self.1.is_match(&state.1)
    }

    fn can_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.can_match(&state.0) && self.1.can_match(&state.1)
    }

    fn will_always_match(&self, state: &IntersectionState<A, B>) -> bool {
        self.0.will_always_match(&state.0) && self.1.will_always_match(&state.1)
    }

    fn accept(&self, state: &IntersectionState<A, B>, byte: u8) -> IntersectionState<A, B> {
        IntersectionState(
            self.0.accept(&state.0, byte),
            self.1.accept(&state.1, byte)
        )
    }
}

/// An automaton that matches exactly when the automaton it wraps does not.
pub struct Complement<A>(A);

/// The `Automaton` state for `Complement<A>`.
pub struct ComplementState<A: Automaton>(A::State);

impl<A: Automaton> Automaton for Complement<A> {
    type State = ComplementState<A>;

    fn start(&self) -> ComplementState<A> {
        ComplementState(self.0.start())
    }

    fn is_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.is_match(&state.0)
    }

    fn can_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.will_always_match(&state.0)
    }

    fn will_always_match(&self, state: &ComplementState<A>) -> bool {
        !self.0.can_match(&state.0)
    }

    fn accept(&self, state: &ComplementState<A>, byte: u8) -> ComplementState<A> {
        ComplementState(self.0.accept(&state.0, byte))
    }
}

#[cfg(test)]
mod test {
    use ::{IntoStreamer, Streamer, Levenshtein, Regex, Set, Automaton};
    use set::OpBuilder;

    static WORDS: &'static str = include_str!("../data/words-10000");

    fn get_set() -> Set {
        Set::from_iter(WORDS.lines()).unwrap()
    }

    #[test]
    fn complement_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.complement()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "focus", "foul"]);
    }

    #[test]
    fn startswith_small() {
        let keys = vec!["", "cooing", "fa", "fo", "fob", "focus", "foo", "food", "foul", "fritter", "frothing"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let stream = set.search(lev.starts_with()).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["cooing", "fo", "fob", "focus", "foo", "food", "foul", "frothing"]);
    }

    #[test]
    fn intersection_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.intersection(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fo", "food"]);
    }

    #[test]
    fn union_small() {
        let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
        let set = Set::from_iter(keys).unwrap();
        let lev = Levenshtein::new("foo", 1).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let stream = set.search(lev.union(reg)).into_stream();

        let keys = stream.into_strs().unwrap();
        assert_eq!(keys, vec!["fa", "fo", "fob", "foo", "food", "foul"]);
    }

    #[test]
    fn intersection_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).intersection(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).intersection();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }

    #[test]
    fn union_large() {
        let set = get_set();
        let lev = Levenshtein::new("foo", 3).unwrap();
        let reg = Regex::new("(..)*").unwrap();
        let mut stream1 = set.search((&lev).union(&reg)).into_stream();
        let mut stream2 = OpBuilder::new().add(set.search(&lev)).add(set.search(&reg)).union();
        while let Some(key1) = stream1.next() {
            assert_eq!(stream2.next(), Some(key1));
        }
        assert_eq!(stream2.next(), None);
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Write};
use std::iter::repeat;
use std::time::Duration;
use time::{self, SteadyTime};
use std::io::Stdout;
use tty::{Width, terminal_size};

macro_rules! printfl {
   ($w:expr, $($tt:tt)*) => {{
        $w.write(&format!($($tt)*).as_bytes()).ok().expect("write() fail");
        $w.flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

const FORMAT: &'static str = "[=>-]";
const TICK_FORMAT: &'static str = "\\|/-";
const NANOS_PER_SEC: u32 = 1_000_000_000;

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

pub struct ProgressBar<T: Write> {
    start_time: SteadyTime,
    units: Units,
    pub total: u64,
    current: u64,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
    tick: Vec<String>,
    tick_state: usize,
    tick_len: usize,
    message: String,
    pub is_finish: bool,
    pub show_bar: bool,
    pub show_speed: bool,
    pub show_percent: bool,
    pub show_counter: bool,
    pub show_time_left: bool,
    pub show_tick: bool,
    pub show_message: bool,
    handle: T,
}

impl ProgressBar<Stdout> {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: u64) -> ProgressBar<Stdout> {
        let handle = ::std::io::stdout();
        ProgressBar::on(handle, total)
    }
}

impl<T: Write> ProgressBar<T> {
    /// Create a new ProgressBar with default configuration but
    /// pass an arbitrary writer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use std::io::stderr;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::on(stderr(), count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn on(handle: T, total: u64) -> ProgressBar<T> {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: SteadyTime::now(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: true,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            show_tick: false,
            show_message: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
            tick: Vec::new(),
            tick_state: 0,
            tick_len: 4,
            message: String::new(),
            handle: handle,
        };
        pb.format(FORMAT);
        pb.tick_format(TICK_FORMAT);
        pb
    }

    /// Set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let n_bytes = 100;
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_owned();
            self.bar_current = v[2].to_owned();
            self.bar_current_n = v[3].to_owned();
            self.bar_remain = v[4].to_owned();
            self.bar_end = v[5].to_owned();
        }
    }

    /// Set message to display in the prefix, call with "" to stop printing a message.
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(20);
    ///
    /// for x in 0..20 {
    ///    match x {
    ///       0 => pb.message("Doing 1st Quarter"),
    ///       5 => pb.message("Doing 2nd Quarter"),
    ///       10 => pb.message("Doing 3rd Quarter"),
    ///       15 => pb.message("Doing 4th Quarter"),
    ///    }
    ///    pb.inc().
    /// }
    ///
    /// ```
    pub fn message(&mut self, message: &str) {
        self.message = message.to_owned()
    }

    /// Set tick format for the progressBar, default is \\|/-
    ///
    /// Format is not limited to 4 characters, any string can
    /// be used as a tick format (the tick will successively
    /// take the value of each char but won't loop backwards).
    ///
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.tick_format("▀▐▄▌")
    /// ```
    pub fn tick_format(&mut self, tick_fmt: &str) {
        if tick_fmt != TICK_FORMAT {
            self.show_tick = true;
        };
        self.tick = tick_fmt.split("").map(|x| x.to_owned()).filter(|x| x != "").collect();
        self.tick_len = self.tick.len();
    }

    /// Update progress bar even though no progress are made
    /// Useful to see if a program is bricked or just
    /// not doing any progress.
    ///
    /// tick is not needed with add or inc
    /// as performed operation take place
    /// in draw function.
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.inc();
    /// for _ in ... {
    ///    ...do something
    ///    pb.tick();
    /// }
    /// pb.finish();
    /// ```
    pub fn tick(&mut self) {
        self.tick_state = (self.tick_state + 1) % self.tick_len;
        if self.current <= self.total {
            self.draw()
        };
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::ProgressBar;
    ///
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: u64) -> u64 {
        self.current += i;
        self.tick();
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> u64 {
        self.add(1)
    }

    fn draw(&mut self) {
        let time_elapsed = time_to_std(SteadyTime::now() - self.start_time);
        let speed = self.current as f64 / fract_dur(time_elapsed);

        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;

        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if self.show_speed {
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left && self.current > 0 {
            if self.total > self.current {
                let left = 1. / speed * (self.total - self.current) as f64;
                if left < 60. {
                    suffix = suffix + &format!("{:.0}s", left);
                } else {
                    suffix = suffix + &format!("{:.0}m", left / 60.);
                }
            }
        }
        // message box
        if self.show_message {
            prefix = prefix + &format!("{}", self.message)
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = prefix +
                     &match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // tick box
        if self.show_tick {
            prefix = prefix + &format!("{} ", self.tick[self.tick_state]);
        }
        // bar box
        if self.show_bar {
            let p = prefix.len() + suffix.len() + 3;
            if p <= width {
                let size = width - p;
                if size > 0 {
                    let curr_count = ((self.current as f64 / self.total as f64) * size as f64)
                        .ceil() as usize;
                    if size > curr_count {
                        let rema_count = size - curr_count;
                        base = self.bar_start.clone();
                        if rema_count > 0 && curr_count > 0 {
                            base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                                   &self.bar_current_n;
                        } else {
                            base = base + repeat!(self.bar_current.as_ref(), curr_count);
                        }
                        base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
                    }
                }
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!(self.handle, "\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        printfl!(self.handle, "");
        self.is_finish = true;
    }

    /// Call finish and write string 's'
    pub fn finish_print(&mut self, s: &str) {
        self.finish();
        printfl!(self.handle, "{}", s)
    }
}

// Implement io::Writer
impl<T: Write> Write for ProgressBar<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n as u64);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

fn time_to_std(d: time::Duration) -> Duration {
    assert!(d > time::Duration::zero());

    let secs = d.num_seconds();
    let nsecs = (d - time::Duration::seconds(secs)).num_nanoseconds().unwrap();
    Duration::new(secs as u64, nsecs as u32)
}

fn fract_dur(d: Duration) -> f64 {
    d.as_secs() as f64 + d.subsec_nanos() as f64 / NANOS_PER_SEC as f64
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2, "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total, "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<commit_msg>fix(pb): draw the bar box in the last iteration<commit_after>use std::io::{self, Write};
use std::iter::repeat;
use std::time::Duration;
use time::{self, SteadyTime};
use std::io::Stdout;
use tty::{Width, terminal_size};

macro_rules! printfl {
   ($w:expr, $($tt:tt)*) => {{
        $w.write(&format!($($tt)*).as_bytes()).ok().expect("write() fail");
        $w.flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

const FORMAT: &'static str = "[=>-]";
const TICK_FORMAT: &'static str = "\\|/-";
const NANOS_PER_SEC: u32 = 1_000_000_000;

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

pub struct ProgressBar<T: Write> {
    start_time: SteadyTime,
    units: Units,
    pub total: u64,
    current: u64,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
    tick: Vec<String>,
    tick_state: usize,
    tick_len: usize,
    message: String,
    pub is_finish: bool,
    pub show_bar: bool,
    pub show_speed: bool,
    pub show_percent: bool,
    pub show_counter: bool,
    pub show_time_left: bool,
    pub show_tick: bool,
    pub show_message: bool,
    handle: T,
}

impl ProgressBar<Stdout> {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: u64) -> ProgressBar<Stdout> {
        let handle = ::std::io::stdout();
        ProgressBar::on(handle, total)
    }
}

impl<T: Write> ProgressBar<T> {
    /// Create a new ProgressBar with default configuration but
    /// pass an arbitrary writer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use std::io::stderr;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::on(stderr(), count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn on(handle: T, total: u64) -> ProgressBar<T> {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: SteadyTime::now(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: true,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            show_tick: false,
            show_message: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
            tick: Vec::new(),
            tick_state: 0,
            tick_len: 4,
            message: String::new(),
            handle: handle,
        };
        pb.format(FORMAT);
        pb.tick_format(TICK_FORMAT);
        pb
    }

    /// Set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let n_bytes = 100;
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_owned();
            self.bar_current = v[2].to_owned();
            self.bar_current_n = v[3].to_owned();
            self.bar_remain = v[4].to_owned();
            self.bar_end = v[5].to_owned();
        }
    }

    /// Set message to display in the prefix, call with "" to stop printing a message.
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(20);
    ///
    /// for x in 0..20 {
    ///    match x {
    ///       0 => pb.message("Doing 1st Quarter"),
    ///       5 => pb.message("Doing 2nd Quarter"),
    ///       10 => pb.message("Doing 3rd Quarter"),
    ///       15 => pb.message("Doing 4th Quarter"),
    ///    }
    ///    pb.inc().
    /// }
    ///
    /// ```
    pub fn message(&mut self, message: &str) {
        self.message = message.to_owned()
    }

    /// Set tick format for the progressBar, default is \\|/-
    ///
    /// Format is not limited to 4 characters, any string can
    /// be used as a tick format (the tick will successively
    /// take the value of each char but won't loop backwards).
    ///
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.tick_format("▀▐▄▌")
    /// ```
    pub fn tick_format(&mut self, tick_fmt: &str) {
        if tick_fmt != TICK_FORMAT {
            self.show_tick = true;
        };
        self.tick = tick_fmt.split("").map(|x| x.to_owned()).filter(|x| x != "").collect();
        self.tick_len = self.tick.len();
    }

    /// Update progress bar even though no progress are made
    /// Useful to see if a program is bricked or just
    /// not doing any progress.
    ///
    /// tick is not needed with add or inc
    /// as performed operation take place
    /// in draw function.
    ///
    /// # Examples
    /// ```ignore
    /// let mut pb = ProgressBar::new(...);
    /// pb.inc();
    /// for _ in ... {
    ///    ...do something
    ///    pb.tick();
    /// }
    /// pb.finish();
    /// ```
    pub fn tick(&mut self) {
        self.tick_state = (self.tick_state + 1) % self.tick_len;
        if self.current <= self.total {
            self.draw()
        };
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::ProgressBar;
    ///
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: u64) -> u64 {
        self.current += i;
        self.tick();
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> u64 {
        self.add(1)
    }

    fn draw(&mut self) {
        let time_elapsed = time_to_std(SteadyTime::now() - self.start_time);
        let speed = self.current as f64 / fract_dur(time_elapsed);

        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;

        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if self.show_speed {
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left && self.current > 0 {
            if self.total > self.current {
                let left = 1. / speed * (self.total - self.current) as f64;
                if left < 60. {
                    suffix = suffix + &format!("{:.0}s", left);
                } else {
                    suffix = suffix + &format!("{:.0}m", left / 60.);
                }
            }
        }
        // message box
        if self.show_message {
            prefix = prefix + &format!("{}", self.message)
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = prefix +
                     &match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // tick box
        if self.show_tick {
            prefix = prefix + &format!("{} ", self.tick[self.tick_state]);
        }
        // bar box
        if self.show_bar {
            let p = prefix.len() + suffix.len() + 3;
            if p < width {
                let size = width - p;
                let curr_count = ((self.current as f64 / self.total as f64) * size as f64)
                    .ceil() as usize;
                if size >= curr_count {
                    let rema_count = size - curr_count;
                    base = self.bar_start.clone();
                    if rema_count > 0 && curr_count > 0 {
                        base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                                &self.bar_current_n;
                    } else {
                        base = base + repeat!(self.bar_current.as_ref(), curr_count);
                    }
                    base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
                }
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!(self.handle, "\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        printfl!(self.handle, "");
        self.is_finish = true;
    }

    /// Call finish and write string 's'
    pub fn finish_print(&mut self, s: &str) {
        self.finish();
        printfl!(self.handle, "{}", s)
    }
}

// Implement io::Writer
impl<T: Write> Write for ProgressBar<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n as u64);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

fn time_to_std(d: time::Duration) -> Duration {
    assert!(d > time::Duration::zero());

    let secs = d.num_seconds();
    let nsecs = (d - time::Duration::seconds(secs)).num_nanoseconds().unwrap();
    Duration::new(secs as u64, nsecs as u32)
}

fn fract_dur(d: Duration) -> f64 {
    d.as_secs() as f64 + d.subsec_nanos() as f64 / NANOS_PER_SEC as f64
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2, "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total, "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<|endoftext|>


<commit_before>extern crate cargo;
extern crate env_logger;
extern crate git2_curl;
extern crate rustc_serialize;
extern crate toml;
#[macro_use] extern crate log;

use std::collections::BTreeSet;
use std::env;
use std::fs;
use std::path::PathBuf;

use cargo::execute_main_without_stdin;
use cargo::util::{self, CliResult, lev_distance, Config, human, CargoResult, ChainError};

#[derive(RustcDecodable)]
pub struct Flags {
    flag_list: bool,
    flag_version: bool,
    flag_verbose: Option<bool>,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    arg_command: String,
    arg_args: Vec<String>,
}

const USAGE: &'static str = "
Rust's package manager

Usage:
    cargo <command> [<args>...]
    cargo [options]

Options:
    -h, --help          Display this message
    -V, --version       Print version info and exit
    --list              List installed commands
    -v, --verbose       Use verbose output
    -q, --quiet         No output printed to stdout
    --color WHEN        Coloring: auto, always, never

Some common cargo commands are:
    build       Compile the current project
    clean       Remove the target directory
    doc         Build this project's and its dependencies' documentation
    new         Create a new cargo project
    init        Create a new cargo project in an existing directory
    run         Build and execute src/main.rs
    test        Run the tests
    bench       Run the benchmarks
    update      Update dependencies listed in Cargo.lock
    search      Search registry for crates
    publish     Package and upload this project to the registry
    install     Install a Rust binary

See 'cargo help <command>' for more information on a specific command.
";

fn main() {
    env_logger::init().unwrap();
    execute_main_without_stdin(execute, true, USAGE)
}

macro_rules! each_subcommand{
    ($mac:ident) => {
        $mac!(bench);
        $mac!(build);
        $mac!(clean);
        $mac!(doc);
        $mac!(fetch);
        $mac!(generate_lockfile);
        $mac!(git_checkout);
        $mac!(help);
        $mac!(init);
        $mac!(install);
        $mac!(locate_project);
        $mac!(login);
        $mac!(metadata);
        $mac!(new);
        $mac!(owner);
        $mac!(package);
        $mac!(pkgid);
        $mac!(publish);
        $mac!(read_manifest);
        $mac!(run);
        $mac!(rustc);
        $mac!(rustdoc);
        $mac!(search);
        $mac!(test);
        $mac!(uninstall);
        $mac!(update);
        $mac!(verify_project);
        $mac!(version);
        $mac!(yank);
    }
}

mod subcommands {
    macro_rules! declare_mod {
        ($name:ident) => ( pub mod $name; )
    }
    each_subcommand!(declare_mod);
}

/**
  The top-level `cargo` command handles configuration and project location
  because they are fundamental (and intertwined). Other commands can rely
  on this top-level information.
*/
fn execute(flags: Flags, config: &Config) -> CliResult<Option<()>> {
    try!(config.configure_shell(flags.flag_verbose,
                                flags.flag_quiet,
                                &flags.flag_color));

    init_git_transports(config);
    cargo::util::job::setup();

    if flags.flag_version {
        println!("{}", cargo::version());
        return Ok(None)
    }

    if flags.flag_list {
        println!("Installed Commands:");
        for command in list_commands(config) {
            println!("    {}", command);
        };
        return Ok(None)
    }

    let args = match &flags.arg_command[..] {
        // For the commands `cargo` and `cargo help`, re-execute ourselves as
        // `cargo -h` so we can go through the normal process of printing the
        // help message.
        "" | "help" if flags.arg_args.is_empty() => {
            config.shell().set_verbose(true);
            let args = &["cargo".to_string(), "-h".to_string()];
            let r = cargo::call_main_without_stdin(execute, config, USAGE, args,
                                                   false);
            cargo::process_executed(r, &mut config.shell());
            return Ok(None)
        }

        // For `cargo help -h` and `cargo help --help`, print out the help
        // message for `cargo help`
        "help" if flags.arg_args[0] == "-h" ||
                  flags.arg_args[0] == "--help" => {
            vec!["cargo".to_string(), "help".to_string(), "-h".to_string()]
        }

        // For `cargo help foo`, print out the usage message for the specified
        // subcommand by executing the command with the `-h` flag.
        "help" => vec!["cargo".to_string(), flags.arg_args[0].clone(),
                       "-h".to_string()],

        // For all other invocations, we're of the form `cargo foo args...`. We
        // use the exact environment arguments to preserve tokens like `--` for
        // example.
        _ => env::args().collect(),
    };

    macro_rules! cmd{
        ($name:ident) => (if args[1] == stringify!($name).replace("_", "-") {
            config.shell().set_verbose(true);
            let r = cargo::call_main_without_stdin(subcommands::$name::execute, config,
                                                   subcommands::$name::USAGE,
                                                   &args,
                                                   false);
            cargo::process_executed(r, &mut config.shell());
            return Ok(None)
        })
    }
    each_subcommand!(cmd);

    try!(execute_subcommand(config, &args[1], &args));
    Ok(None)
}

fn find_closest(config: &Config, cmd: &str) -> Option<String> {
    let cmds = list_commands(config);
    // Only consider candidates with a lev_distance of 3 or less so we don't
    // suggest out-of-the-blue options.
    let mut filtered = cmds.iter().map(|c| (lev_distance(&c, cmd), c))
                                  .filter(|&(d, _)| d < 4)
                                  .collect::<Vec<_>>();
    filtered.sort_by(|a, b| a.0.cmp(&b.0));
    filtered.get(0).map(|slot| slot.1.clone())
}

fn execute_subcommand(config: &Config,
                      cmd: &str,
                      args: &[String]) -> CargoResult<()> {
    let command_exe = format!("cargo-{}{}", cmd, env::consts::EXE_SUFFIX);
    let path = search_directories(config)
                    .iter()
                    .map(|dir| dir.join(&command_exe))
                    .filter_map(|dir| fs::metadata(&dir).ok().map(|m| (dir, m)))
                    .find(|&(_, ref meta)| is_executable(meta));
    let command = match path {
        Some((command, _)) => command,
        None => {
            return Err(human(match find_closest(config, cmd) {
                Some(closest) => format!("no such subcommand\n\n\t\
                                          Did you mean `{}`?\n", closest),
                None => "no such subcommand".to_string()
            }))
        }
    };
    try!(util::process(&command).args(&args[1..]).exec().chain_error(|| {
        human(format!("third party subcommand `{}` exited unsuccessfully", command_exe))
    }));
    Ok(())
}

/// List all runnable commands. find_command should always succeed
/// if given one of returned command.
fn list_commands(config: &Config) -> BTreeSet<String> {
    let prefix = "cargo-";
    let suffix = env::consts::EXE_SUFFIX;
    let mut commands = BTreeSet::new();
    for dir in search_directories(config) {
        let entries = match fs::read_dir(dir) {
            Ok(entries) => entries,
            _ => continue
        };
        for entry in entries.filter_map(|e| e.ok()) {
            let path = entry.path();
            let filename = match path.file_name().and_then(|s| s.to_str()) {
                Some(filename) => filename,
                _ => continue
            };
            if !filename.starts_with(prefix) || !filename.ends_with(suffix) {
                continue
            }
            if let Ok(meta) = entry.metadata() {
                if is_executable(&meta) {
                    let end = filename.len() - suffix.len();
                    commands.insert(filename[prefix.len()..end].to_string());
                }
            }
        }
    }

    macro_rules! add_cmd {
        ($cmd:ident) => ({ commands.insert(stringify!($cmd).replace("_", "-")); })
    }
    each_subcommand!(add_cmd);
    commands
}

#[cfg(unix)]
fn is_executable(metadata: &fs::Metadata) -> bool {
    use std::os::unix::prelude::*;
    metadata.is_file() && metadata.permissions().mode() & 0o111 != 0
}
#[cfg(windows)]
fn is_executable(metadata: &fs::Metadata) -> bool {
    metadata.is_file()
}

fn search_directories(config: &Config) -> Vec<PathBuf> {
    let mut dirs = vec![config.home().join("bin")];
    if let Some(val) = env::var_os("PATH") {
        dirs.extend(env::split_paths(&val));
    }
    dirs
}

fn init_git_transports(config: &Config) {
    // Only use a custom transport if a proxy is configured, right now libgit2
    // doesn't support proxies and we have to use a custom transport in this
    // case. The custom transport, however, is not as well battle-tested.
    match cargo::ops::http_proxy_exists(config) {
        Ok(true) => {}
        _ => return
    }

    let handle = match cargo::ops::http_handle(config) {
        Ok(handle) => handle,
        Err(..) => return,
    };

    // The unsafety of the registration function derives from two aspects:
    //
    // 1. This call must be synchronized with all other registration calls as
    //    well as construction of new transports.
    // 2. The argument is leaked.
    //
    // We're clear on point (1) because this is only called at the start of this
    // binary (we know what the state of the world looks like) and we're mostly
    // clear on point (2) because we'd only free it after everything is done
    // anyway
    unsafe {
        git2_curl::register(handle);
    }
}
<commit_msg>Auto merge of #2438 - jseyfried:subcommands, r=alexcrichton<commit_after>extern crate cargo;
extern crate env_logger;
extern crate git2_curl;
extern crate rustc_serialize;
extern crate toml;
#[macro_use] extern crate log;

use std::collections::BTreeSet;
use std::env;
use std::fs;
use std::path::PathBuf;

use cargo::execute_main_without_stdin;
use cargo::util::{self, CliResult, lev_distance, Config, human, CargoResult, ChainError};

#[derive(RustcDecodable)]
pub struct Flags {
    flag_list: bool,
    flag_version: bool,
    flag_verbose: Option<bool>,
    flag_quiet: Option<bool>,
    flag_color: Option<String>,
    arg_command: String,
    arg_args: Vec<String>,
}

const USAGE: &'static str = "
Rust's package manager

Usage:
    cargo <command> [<args>...]
    cargo [options]

Options:
    -h, --help          Display this message
    -V, --version       Print version info and exit
    --list              List installed commands
    -v, --verbose       Use verbose output
    -q, --quiet         No output printed to stdout
    --color WHEN        Coloring: auto, always, never

Some common cargo commands are:
    build       Compile the current project
    clean       Remove the target directory
    doc         Build this project's and its dependencies' documentation
    new         Create a new cargo project
    init        Create a new cargo project in an existing directory
    run         Build and execute src/main.rs
    test        Run the tests
    bench       Run the benchmarks
    update      Update dependencies listed in Cargo.lock
    search      Search registry for crates
    publish     Package and upload this project to the registry
    install     Install a Rust binary

See 'cargo help <command>' for more information on a specific command.
";

fn main() {
    env_logger::init().unwrap();
    execute_main_without_stdin(execute, true, USAGE)
}

macro_rules! each_subcommand{
    ($mac:ident) => {
        $mac!(bench);
        $mac!(build);
        $mac!(clean);
        $mac!(doc);
        $mac!(fetch);
        $mac!(generate_lockfile);
        $mac!(git_checkout);
        $mac!(help);
        $mac!(init);
        $mac!(install);
        $mac!(locate_project);
        $mac!(login);
        $mac!(metadata);
        $mac!(new);
        $mac!(owner);
        $mac!(package);
        $mac!(pkgid);
        $mac!(publish);
        $mac!(read_manifest);
        $mac!(run);
        $mac!(rustc);
        $mac!(rustdoc);
        $mac!(search);
        $mac!(test);
        $mac!(uninstall);
        $mac!(update);
        $mac!(verify_project);
        $mac!(version);
        $mac!(yank);
    }
}

macro_rules! declare_mod {
    ($name:ident) => ( pub mod $name; )
}
each_subcommand!(declare_mod);

/**
  The top-level `cargo` command handles configuration and project location
  because they are fundamental (and intertwined). Other commands can rely
  on this top-level information.
*/
fn execute(flags: Flags, config: &Config) -> CliResult<Option<()>> {
    try!(config.configure_shell(flags.flag_verbose,
                                flags.flag_quiet,
                                &flags.flag_color));

    init_git_transports(config);
    cargo::util::job::setup();

    if flags.flag_version {
        println!("{}", cargo::version());
        return Ok(None)
    }

    if flags.flag_list {
        println!("Installed Commands:");
        for command in list_commands(config) {
            println!("    {}", command);
        };
        return Ok(None)
    }

    let args = match &flags.arg_command[..] {
        // For the commands `cargo` and `cargo help`, re-execute ourselves as
        // `cargo -h` so we can go through the normal process of printing the
        // help message.
        "" | "help" if flags.arg_args.is_empty() => {
            config.shell().set_verbose(true);
            let args = &["cargo".to_string(), "-h".to_string()];
            let r = cargo::call_main_without_stdin(execute, config, USAGE, args,
                                                   false);
            cargo::process_executed(r, &mut config.shell());
            return Ok(None)
        }

        // For `cargo help -h` and `cargo help --help`, print out the help
        // message for `cargo help`
        "help" if flags.arg_args[0] == "-h" ||
                  flags.arg_args[0] == "--help" => {
            vec!["cargo".to_string(), "help".to_string(), "-h".to_string()]
        }

        // For `cargo help foo`, print out the usage message for the specified
        // subcommand by executing the command with the `-h` flag.
        "help" => vec!["cargo".to_string(), flags.arg_args[0].clone(),
                       "-h".to_string()],

        // For all other invocations, we're of the form `cargo foo args...`. We
        // use the exact environment arguments to preserve tokens like `--` for
        // example.
        _ => env::args().collect(),
    };

    macro_rules! cmd{
        ($name:ident) => (if args[1] == stringify!($name).replace("_", "-") {
            config.shell().set_verbose(true);
            let r = cargo::call_main_without_stdin($name::execute, config,
                                                   $name::USAGE,
                                                   &args,
                                                   false);
            cargo::process_executed(r, &mut config.shell());
            return Ok(None)
        })
    }
    each_subcommand!(cmd);

    try!(execute_subcommand(config, &args[1], &args));
    Ok(None)
}

fn find_closest(config: &Config, cmd: &str) -> Option<String> {
    let cmds = list_commands(config);
    // Only consider candidates with a lev_distance of 3 or less so we don't
    // suggest out-of-the-blue options.
    let mut filtered = cmds.iter().map(|c| (lev_distance(&c, cmd), c))
                                  .filter(|&(d, _)| d < 4)
                                  .collect::<Vec<_>>();
    filtered.sort_by(|a, b| a.0.cmp(&b.0));
    filtered.get(0).map(|slot| slot.1.clone())
}

fn execute_subcommand(config: &Config,
                      cmd: &str,
                      args: &[String]) -> CargoResult<()> {
    let command_exe = format!("cargo-{}{}", cmd, env::consts::EXE_SUFFIX);
    let path = search_directories(config)
                    .iter()
                    .map(|dir| dir.join(&command_exe))
                    .filter_map(|dir| fs::metadata(&dir).ok().map(|m| (dir, m)))
                    .find(|&(_, ref meta)| is_executable(meta));
    let command = match path {
        Some((command, _)) => command,
        None => {
            return Err(human(match find_closest(config, cmd) {
                Some(closest) => format!("no such subcommand\n\n\t\
                                          Did you mean `{}`?\n", closest),
                None => "no such subcommand".to_string()
            }))
        }
    };
    try!(util::process(&command).args(&args[1..]).exec().chain_error(|| {
        human(format!("third party subcommand `{}` exited unsuccessfully", command_exe))
    }));
    Ok(())
}

/// List all runnable commands. find_command should always succeed
/// if given one of returned command.
fn list_commands(config: &Config) -> BTreeSet<String> {
    let prefix = "cargo-";
    let suffix = env::consts::EXE_SUFFIX;
    let mut commands = BTreeSet::new();
    for dir in search_directories(config) {
        let entries = match fs::read_dir(dir) {
            Ok(entries) => entries,
            _ => continue
        };
        for entry in entries.filter_map(|e| e.ok()) {
            let path = entry.path();
            let filename = match path.file_name().and_then(|s| s.to_str()) {
                Some(filename) => filename,
                _ => continue
            };
            if !filename.starts_with(prefix) || !filename.ends_with(suffix) {
                continue
            }
            if let Ok(meta) = entry.metadata() {
                if is_executable(&meta) {
                    let end = filename.len() - suffix.len();
                    commands.insert(filename[prefix.len()..end].to_string());
                }
            }
        }
    }

    macro_rules! add_cmd {
        ($cmd:ident) => ({ commands.insert(stringify!($cmd).replace("_", "-")); })
    }
    each_subcommand!(add_cmd);
    commands
}

#[cfg(unix)]
fn is_executable(metadata: &fs::Metadata) -> bool {
    use std::os::unix::prelude::*;
    metadata.is_file() && metadata.permissions().mode() & 0o111 != 0
}
#[cfg(windows)]
fn is_executable(metadata: &fs::Metadata) -> bool {
    metadata.is_file()
}

fn search_directories(config: &Config) -> Vec<PathBuf> {
    let mut dirs = vec![config.home().join("bin")];
    if let Some(val) = env::var_os("PATH") {
        dirs.extend(env::split_paths(&val));
    }
    dirs
}

fn init_git_transports(config: &Config) {
    // Only use a custom transport if a proxy is configured, right now libgit2
    // doesn't support proxies and we have to use a custom transport in this
    // case. The custom transport, however, is not as well battle-tested.
    match cargo::ops::http_proxy_exists(config) {
        Ok(true) => {}
        _ => return
    }

    let handle = match cargo::ops::http_handle(config) {
        Ok(handle) => handle,
        Err(..) => return,
    };

    // The unsafety of the registration function derives from two aspects:
    //
    // 1. This call must be synchronized with all other registration calls as
    //    well as construction of new transports.
    // 2. The argument is leaked.
    //
    // We're clear on point (1) because this is only called at the start of this
    // binary (we know what the state of the world looks like) and we're mostly
    // clear on point (2) because we'd only free it after everything is done
    // anyway
    unsafe {
        git2_curl::register(handle);
    }
}
<|endoftext|>