<commit_before>extern crate sdl2;
use self::sdl2::audio::{AudioQueue};

mod pulse_channel;
mod triangle_channel;
mod noise_channel;
mod dmc_channel;
mod envelope;
mod length_counter;
mod timer;

use memory::Memory;

use self::pulse_channel::{PulseChannel, Complement};
use self::triangle_channel::TriangleChannel;
use self::noise_channel::NoiseChannel;
use self::dmc_channel::DmcChannel;
use self::envelope::Envelope;

use std::collections::VecDeque;
use std::cell::RefCell;
use std::rc::Rc;

const APU_STATUS_REGISTER : u16 = 0x4015;
const FRAME_COUNTER_REGISTER : u16 = 0x4017;
/*

TODO:
    * Implement triangle wave, noise, dmc channels
    * Implement SDL backend
    * Verify that register writes into pulse channels do the right things
        * Write into 0x4003/0x4007 should reset phase
*/

enum FrameMode {
    Mode0, // 4 step mode
    Mode1 // 5 step mode
}

enum CycleState {
    QuarterFrameCycle,
    HalfFrameCycle,
    NormalCycle,
}

struct FrameCounter {
    mode: FrameMode,
    cycle: u16,
    interrupt_disabled: bool,
    interrupt_flag: bool,
    even_cycle: bool
}

impl FrameCounter {
    fn new() -> FrameCounter {
         FrameCounter {
            mode: FrameMode::Mode0,
            cycle: 0,
            interrupt_disabled: true,
            interrupt_flag: false,
            even_cycle: false
        }
    }

    fn cycle(&mut self) -> CycleState {
        self.cycle += 1;
        // apu cycles at uneven cycles (like at cycle 3728.5)
        // the mod ensures that the apu remains in sync
        let cycle_mod = if self.even_cycle {
            1
        } else {
            0
        };

        match self.mode {
            FrameMode::Mode0 => {
                if self.cycle == 0 {
                    self.interrupt();
                } else if (self.cycle == 3728 + cycle_mod) ||
                    (self.cycle == 11185 + cycle_mod) {
                    return CycleState::QuarterFrameCycle;
                } else if self.cycle == 7456 + cycle_mod {
                    return CycleState::HalfFrameCycle;
                } else if self.cycle == 14914 {
                    self.interrupt();
                }

                if self.cycle == 14914 + cycle_mod {
                    self.cycle = 0;
                    self.even_cycle = !self.even_cycle;
                    self.interrupt();
                    return CycleState::HalfFrameCycle;
                }

            },
            FrameMode::Mode1 => {
                if (self.cycle == 3728 + cycle_mod) ||
                    (self.cycle == 11185 + cycle_mod) {
                    return CycleState::QuarterFrameCycle;
                } else if self.cycle == 7456 + cycle_mod {
                    return CycleState::HalfFrameCycle;
                } else if self.cycle == 18640 + cycle_mod {
                    self.cycle = 0;
                    self.even_cycle = !self.even_cycle;
                    return CycleState::HalfFrameCycle;
                }
            }
        }

        return CycleState::NormalCycle
    }

    fn interrupt(&mut self) {
        if !self.interrupt_disabled {
            self.interrupt_flag = true;
        } else {

        }
    }

    fn clear_interrupt(&mut self) {
        self.interrupt_flag = false;
    }
}

pub struct Apu<'a> {
    pulse_channel_1: PulseChannel,
    pulse_channel_2: PulseChannel,
    triangle_channel: TriangleChannel,
    noise_channel: NoiseChannel,
    dmc_channel: DmcChannel<'a>,
    frame_counter: FrameCounter,
    buffer: Vec<f32>,
    sample_cycle: f64,
    cycles_per_sample: f64,
    max_samples_before_clearing_buffer: usize,
    audio_queue: AudioQueue<f32>
}

impl<'a> Memory for Apu<'a> {
    fn read(&mut self,  address: u16) -> u8 {
        if address == APU_STATUS_REGISTER {

            let mut ret = 0;

            if self.dmc_channel.pending_interrupt() {
                ret |= 0b1000_0000;
            }

            if self.frame_counter.interrupt_flag {
                ret |= 0b0100_0000;
            }

            if self.dmc_channel.active() {
                ret |= 0b0001_0000;
            }

            if self.noise_channel.length_counter_nonzero() {
                ret |= 0b0000_1000;
            }

            if self.triangle_channel.length_counter_nonzero() {
                ret |= 0b0000_0100;
            }

            if self.pulse_channel_2.length_counter_nonzero() {
                ret |= 0b0000_0010;
            }

            if self.pulse_channel_1.length_counter_nonzero() {
                ret |= 0b0000_0001;
            }

            self.frame_counter.clear_interrupt();
            return ret;

        } else {
            panic!("Invalid APU register read for {:0x}", address);
        }
    }

    fn write(&mut self, address: u16, value: u8) {

        if address >= 0x4000 && address <= 0x4003 {
            self.pulse_channel_1.write(address, value);
        } else if address >= 0x4004 && address <= 0x4007 {
            self.pulse_channel_2.write(address, value);
        } else if address >= 0x4008 && address <= 0x400B {
            self.triangle_channel.write(address, value);
        } else if address >= 0x400C && address <= 0x400F {
            self.noise_channel.write(address, value);
        } else if address >= 0x4010 && address <= 0x4013 {
            self.dmc_channel.write(address, value);
        } else if address == APU_STATUS_REGISTER {
            let enable_dmc = (0b0001_0000 & value) != 0;
            let enable_noise = (0b0000_1000 & value) != 0;
            let enable_triangle = (0b0000_0100 & value) != 0;
            let enable_pulse_2 = (0b0000_0010 & value) != 0;
            let enable_pulse_1 = (0b0000_0001 & value) != 0;

            self.pulse_channel_1.enable_channel(enable_pulse_1);
            self.pulse_channel_2.enable_channel(enable_pulse_2);
            self.triangle_channel.enable_channel(enable_triangle);
            self.noise_channel.enable_channel(enable_noise);
            self.dmc_channel.enable_channel(enable_dmc);

            self.dmc_channel.clear_interrupt();
        } else if address == FRAME_COUNTER_REGISTER {
            let mode = (0b1000_0000 & value) != 0;
            let interrupt_disable_flag = (0b0100_0000 & value) != 0;

            self.frame_counter.cycle = 0;
            if mode {
                self.frame_counter.mode = FrameMode::Mode1;
            } else {
                self.frame_counter.mode = FrameMode::Mode0;
            }

            self.frame_counter.interrupt_disabled = interrupt_disable_flag;

            if interrupt_disable_flag {
                self.frame_counter.interrupt_flag = false;
            }

        } else {
            panic!("Invalid write address {:0x} for APU: ", address);
        }
    }
}

impl<'a> Apu<'a> {
    pub fn new(audio_queue: AudioQueue<f32>) -> Apu<'a> {
        Apu {
            pulse_channel_1: PulseChannel::new(Complement::One),
            pulse_channel_2: PulseChannel::new(Complement::Two),
            triangle_channel: TriangleChannel::new(),
            noise_channel: NoiseChannel::new(),
            dmc_channel: DmcChannel::new(),
            frame_counter: FrameCounter::new(),
            buffer: vec![],
            sample_cycle: 0.0,
            cycles_per_sample: 0.0,
            max_samples_before_clearing_buffer: 0,
            audio_queue: audio_queue,
        }
    }

    pub fn samples(&mut self, samples: u16) {
        self.buffer.resize(samples as usize, 0.0);
        self.max_samples_before_clearing_buffer = samples as usize;
    }

    pub fn set_sampling_rate(&mut self, cpu_frequency: f64, sample_rate: i32) {
        let cpu_cycles_per_apu_cycle = 2.0;
        self.cycles_per_sample =
            ((cpu_frequency*1000_000.0 / cpu_cycles_per_apu_cycle)
            / sample_rate as f64);
    }

    pub fn execute_cycle(&mut self) {
        self.emulate_channels();
        self.gather_sample();
    }

    fn emulate_channels(&mut self) {
        match self.frame_counter.cycle() {
            CycleState::QuarterFrameCycle => {
                self.pulse_channel_1.cycle_envelope();
                self.pulse_channel_2.cycle_envelope();

                self.triangle_channel.cycle_linear_counter();

            },
            CycleState::HalfFrameCycle => {
                self.pulse_channel_1.cycle_envelope();
                self.pulse_channel_2.cycle_envelope();
                self.noise_channel.cycle_envelope();

                self.pulse_channel_1.cycle_length_counter();
                self.pulse_channel_2.cycle_length_counter();
                self.triangle_channel.cycle_length_counter();
                self.noise_channel.cycle_length_counter();

                self.pulse_channel_1.cycle_sweep_unit();
                self.pulse_channel_2.cycle_sweep_unit();
            },
            CycleState::NormalCycle => {}
        }

        self.pulse_channel_1.cycle_timer();
        self.pulse_channel_2.cycle_timer();

        // triangle channel timer cycles twice for each apu cycle
        self.triangle_channel.cycle_timer();
        self.triangle_channel.cycle_timer();

        self.noise_channel.cycle_timer();

        self.dmc_channel.cycle_timer();
    }

    fn gather_sample(&mut self) {
        self.sample_cycle += 1.0;
        // get samples every ~ (apu cycle) / (sample rate) / 2
        // (apu cycle -> 2 cpu cycles)
        if self.sample_cycle >= self.cycles_per_sample {
            let output = self.output() as f32;
            self.buffer.push(output);
            self.sample_cycle -= self.cycles_per_sample;

            if self.buffer.len() >= self.max_samples_before_clearing_buffer {
                self.audio_queue.queue(self.buffer.as_slice());
                self.buffer.clear();
            }
        }
    }

    pub fn pending_interrupt(&self) -> bool {
       /*self.frame_counter.interrupt_flag || */self.dmc_channel.pending_interrupt()
    }

    pub fn set_memory(&mut self, mem: Rc<RefCell<Box<Memory + 'a>>>) {
        self.dmc_channel.set_memory(mem);
    }

    fn output(&self) -> f64 {
        let pulse_output =
            0.00752*(
                self.pulse_channel_1.output() +
                self.pulse_channel_2.output());

        let tnd_output =
            0.00851*self.triangle_channel.output()
            + 0.00494*self.noise_channel.output()
            + 0.00335*self.dmc_channel.output();
        pulse_output + tnd_output
    }
}<commit_msg>APU: Remove comment<commit_after>extern crate sdl2;
use self::sdl2::audio::{AudioQueue};

mod pulse_channel;
mod triangle_channel;
mod noise_channel;
mod dmc_channel;
mod envelope;
mod length_counter;
mod timer;

use memory::Memory;

use self::pulse_channel::{PulseChannel, Complement};
use self::triangle_channel::TriangleChannel;
use self::noise_channel::NoiseChannel;
use self::dmc_channel::DmcChannel;
use self::envelope::Envelope;

use std::collections::VecDeque;
use std::cell::RefCell;
use std::rc::Rc;

const APU_STATUS_REGISTER : u16 = 0x4015;
const FRAME_COUNTER_REGISTER : u16 = 0x4017;

enum FrameMode {
    Mode0, // 4 step mode
    Mode1 // 5 step mode
}

enum CycleState {
    QuarterFrameCycle,
    HalfFrameCycle,
    NormalCycle,
}

struct FrameCounter {
    mode: FrameMode,
    cycle: u16,
    interrupt_disabled: bool,
    interrupt_flag: bool,
    even_cycle: bool
}

impl FrameCounter {
    fn new() -> FrameCounter {
         FrameCounter {
            mode: FrameMode::Mode0,
            cycle: 0,
            interrupt_disabled: true,
            interrupt_flag: false,
            even_cycle: false
        }
    }

    fn cycle(&mut self) -> CycleState {
        self.cycle += 1;
        // apu cycles at uneven cycles (like at cycle 3728.5)
        // the mod ensures that the apu remains in sync
        let cycle_mod = if self.even_cycle {
            1
        } else {
            0
        };

        match self.mode {
            FrameMode::Mode0 => {
                if self.cycle == 0 {
                    self.interrupt();
                } else if (self.cycle == 3728 + cycle_mod) ||
                    (self.cycle == 11185 + cycle_mod) {
                    return CycleState::QuarterFrameCycle;
                } else if self.cycle == 7456 + cycle_mod {
                    return CycleState::HalfFrameCycle;
                } else if self.cycle == 14914 {
                    self.interrupt();
                }

                if self.cycle == 14914 + cycle_mod {
                    self.cycle = 0;
                    self.even_cycle = !self.even_cycle;
                    self.interrupt();
                    return CycleState::HalfFrameCycle;
                }

            },
            FrameMode::Mode1 => {
                if (self.cycle == 3728 + cycle_mod) ||
                    (self.cycle == 11185 + cycle_mod) {
                    return CycleState::QuarterFrameCycle;
                } else if self.cycle == 7456 + cycle_mod {
                    return CycleState::HalfFrameCycle;
                } else if self.cycle == 18640 + cycle_mod {
                    self.cycle = 0;
                    self.even_cycle = !self.even_cycle;
                    return CycleState::HalfFrameCycle;
                }
            }
        }

        return CycleState::NormalCycle
    }

    fn interrupt(&mut self) {
        if !self.interrupt_disabled {
            self.interrupt_flag = true;
        } else {

        }
    }

    fn clear_interrupt(&mut self) {
        self.interrupt_flag = false;
    }
}

pub struct Apu<'a> {
    pulse_channel_1: PulseChannel,
    pulse_channel_2: PulseChannel,
    triangle_channel: TriangleChannel,
    noise_channel: NoiseChannel,
    dmc_channel: DmcChannel<'a>,
    frame_counter: FrameCounter,
    buffer: Vec<f32>,
    sample_cycle: f64,
    cycles_per_sample: f64,
    max_samples_before_clearing_buffer: usize,
    audio_queue: AudioQueue<f32>
}

impl<'a> Memory for Apu<'a> {
    fn read(&mut self,  address: u16) -> u8 {
        if address == APU_STATUS_REGISTER {

            let mut ret = 0;

            if self.dmc_channel.pending_interrupt() {
                ret |= 0b1000_0000;
            }

            if self.frame_counter.interrupt_flag {
                ret |= 0b0100_0000;
            }

            if self.dmc_channel.active() {
                ret |= 0b0001_0000;
            }

            if self.noise_channel.length_counter_nonzero() {
                ret |= 0b0000_1000;
            }

            if self.triangle_channel.length_counter_nonzero() {
                ret |= 0b0000_0100;
            }

            if self.pulse_channel_2.length_counter_nonzero() {
                ret |= 0b0000_0010;
            }

            if self.pulse_channel_1.length_counter_nonzero() {
                ret |= 0b0000_0001;
            }

            self.frame_counter.clear_interrupt();
            return ret;

        } else {
            panic!("Invalid APU register read for {:0x}", address);
        }
    }

    fn write(&mut self, address: u16, value: u8) {

        if address >= 0x4000 && address <= 0x4003 {
            self.pulse_channel_1.write(address, value);
        } else if address >= 0x4004 && address <= 0x4007 {
            self.pulse_channel_2.write(address, value);
        } else if address >= 0x4008 && address <= 0x400B {
            self.triangle_channel.write(address, value);
        } else if address >= 0x400C && address <= 0x400F {
            self.noise_channel.write(address, value);
        } else if address >= 0x4010 && address <= 0x4013 {
            self.dmc_channel.write(address, value);
        } else if address == APU_STATUS_REGISTER {
            let enable_dmc = (0b0001_0000 & value) != 0;
            let enable_noise = (0b0000_1000 & value) != 0;
            let enable_triangle = (0b0000_0100 & value) != 0;
            let enable_pulse_2 = (0b0000_0010 & value) != 0;
            let enable_pulse_1 = (0b0000_0001 & value) != 0;

            self.pulse_channel_1.enable_channel(enable_pulse_1);
            self.pulse_channel_2.enable_channel(enable_pulse_2);
            self.triangle_channel.enable_channel(enable_triangle);
            self.noise_channel.enable_channel(enable_noise);
            self.dmc_channel.enable_channel(enable_dmc);

            self.dmc_channel.clear_interrupt();
        } else if address == FRAME_COUNTER_REGISTER {
            let mode = (0b1000_0000 & value) != 0;
            let interrupt_disable_flag = (0b0100_0000 & value) != 0;

            self.frame_counter.cycle = 0;
            if mode {
                self.frame_counter.mode = FrameMode::Mode1;
            } else {
                self.frame_counter.mode = FrameMode::Mode0;
            }

            self.frame_counter.interrupt_disabled = interrupt_disable_flag;

            if interrupt_disable_flag {
                self.frame_counter.interrupt_flag = false;
            }

        } else {
            panic!("Invalid write address {:0x} for APU: ", address);
        }
    }
}

impl<'a> Apu<'a> {
    pub fn new(audio_queue: AudioQueue<f32>) -> Apu<'a> {
        Apu {
            pulse_channel_1: PulseChannel::new(Complement::One),
            pulse_channel_2: PulseChannel::new(Complement::Two),
            triangle_channel: TriangleChannel::new(),
            noise_channel: NoiseChannel::new(),
            dmc_channel: DmcChannel::new(),
            frame_counter: FrameCounter::new(),
            buffer: vec![],
            sample_cycle: 0.0,
            cycles_per_sample: 0.0,
            max_samples_before_clearing_buffer: 0,
            audio_queue: audio_queue,
        }
    }

    pub fn samples(&mut self, samples: u16) {
        self.buffer.resize(samples as usize, 0.0);
        self.max_samples_before_clearing_buffer = samples as usize;
    }

    pub fn set_sampling_rate(&mut self, cpu_frequency: f64, sample_rate: i32) {
        let cpu_cycles_per_apu_cycle = 2.0;
        self.cycles_per_sample =
            ((cpu_frequency*1000_000.0 / cpu_cycles_per_apu_cycle)
            / sample_rate as f64);
    }

    pub fn execute_cycle(&mut self) {
        self.emulate_channels();
        self.gather_sample();
    }

    fn emulate_channels(&mut self) {
        match self.frame_counter.cycle() {
            CycleState::QuarterFrameCycle => {
                self.pulse_channel_1.cycle_envelope();
                self.pulse_channel_2.cycle_envelope();

                self.triangle_channel.cycle_linear_counter();

            },
            CycleState::HalfFrameCycle => {
                self.pulse_channel_1.cycle_envelope();
                self.pulse_channel_2.cycle_envelope();
                self.noise_channel.cycle_envelope();

                self.pulse_channel_1.cycle_length_counter();
                self.pulse_channel_2.cycle_length_counter();
                self.triangle_channel.cycle_length_counter();
                self.noise_channel.cycle_length_counter();

                self.pulse_channel_1.cycle_sweep_unit();
                self.pulse_channel_2.cycle_sweep_unit();
            },
            CycleState::NormalCycle => {}
        }

        self.pulse_channel_1.cycle_timer();
        self.pulse_channel_2.cycle_timer();

        // triangle channel timer cycles twice for each apu cycle
        self.triangle_channel.cycle_timer();
        self.triangle_channel.cycle_timer();

        self.noise_channel.cycle_timer();

        self.dmc_channel.cycle_timer();
    }

    fn gather_sample(&mut self) {
        self.sample_cycle += 1.0;
        // get samples every ~ (apu cycle) / (sample rate) / 2
        // (apu cycle -> 2 cpu cycles)
        if self.sample_cycle >= self.cycles_per_sample {
            let output = self.output() as f32;
            self.buffer.push(output);
            self.sample_cycle -= self.cycles_per_sample;

            if self.buffer.len() >= self.max_samples_before_clearing_buffer {
                self.audio_queue.queue(self.buffer.as_slice());
                self.buffer.clear();
            }
        }
    }

    pub fn pending_interrupt(&self) -> bool {
       /*self.frame_counter.interrupt_flag || */self.dmc_channel.pending_interrupt()
    }

    pub fn set_memory(&mut self, mem: Rc<RefCell<Box<Memory + 'a>>>) {
        self.dmc_channel.set_memory(mem);
    }

    fn output(&self) -> f64 {
        let pulse_output =
            0.00752*(
                self.pulse_channel_1.output() +
                self.pulse_channel_2.output());

        let tnd_output =
            0.00851*self.triangle_channel.output()
            + 0.00494*self.noise_channel.output()
            + 0.00335*self.dmc_channel.output();
        pulse_output + tnd_output
    }
}<|endoftext|>


<commit_before>use crate::{generate::*, inner::*, norm::Norm, types::*};
use ndarray::*;

/// Iterative orthogonalizer using modified Gram-Schmit procedure
#[derive(Debug, Clone)]
pub struct MGS<A> {
    /// Dimension of base space
    dimension: usize,
    /// Basis of spanned space
    q: Vec<Array1<A>>,
}

pub type Residual<S> = ArrayBase<S, Ix1>;
pub type Coefficient<A> = Array1<A>;

impl<A: Scalar + Lapack> MGS<A> {
    pub fn new(dimension: usize) -> Self {
        Self {
            dimension,
            q: Vec::new(),
        }
    }

    pub fn dim(&self) -> usize {
        self.dimension
    }

    pub fn len(&self) -> usize {
        self.q.len()
    }

    /// Orthogonalize given vector using current basis
    ///
    /// Panic
    /// -------
    /// - if the size of the input array mismaches to the dimension
    pub fn orthogonalize<S>(&self, mut a: ArrayBase<S, Ix1>) -> (Residual<S>, Coefficient<A>)
    where
        S: DataMut<Elem = A>,
    {
        assert_eq!(a.len(), self.dim());
        let mut coef = Array1::zeros(self.len() + 1);
        for i in 0..self.len() {
            let q = &self.q[i];
            let c = q.inner(&a);
            azip!(mut a, q (q) in { *a = *a - c * q } );
            coef[i] = c;
        }
        let nrm = a.norm_l2();
        coef[self.len()] = A::from_real(nrm);
        (a, coef)
    }

    /// Add new vector if the residual is larger than relative tolerance
    ///
    /// Panic
    /// -------
    /// - if the size of the input array mismaches to the dimension
    pub fn append_if_independent<S>(&mut self, a: ArrayBase<S, Ix1>, rtol: A::Real) -> Option<Coefficient<A>>
    where
        S: Data<Elem = A>,
    {
        let a = a.into_owned();
        let (mut a, coef) = self.orthogonalize(a);
        let nrm = coef[coef.len()].re();
        if nrm < rtol {
            // Linearly dependent
            return None;
        }
        azip!(mut a in { *a = *a / A::from_real(nrm) });
        self.q.push(a);
        Some(coef)
    }

    /// Get orthogonal basis as Q matrix
    pub fn get_q(&self) -> Array2<A> {
        hstack(&self.q).unwrap()
    }
}

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

    const N: usize = 5;

    #[test]
    fn new() {
        let mgs: MGS<f32> = MGS::new(N);
        assert_eq!(mgs.dim(), N);
        assert_eq!(mgs.len(), 0);
    }

    fn test_append<A: Scalar + Lapack>(rtol: A::Real) {
        let mut mgs: MGS<A> = MGS::new(N);
        let a: Array2<A> = crate::generate::random((N, 3));
        dbg!(&a);
        for col in a.axis_iter(Axis(1)) {
            let res = mgs.append(col);
            dbg!(res);
        }
        let q = mgs.get_q();
        dbg!(&q);
        let r = mgs.get_r();
        dbg!(&r);

        dbg!(q.dot(&r));
        close_l2(&q.dot(&r), &a, rtol).unwrap();

        let qt: Array2<_> = conjugate(&q);
        dbg!(qt.dot(&q));
        close_l2(&qt.dot(&q), &Array2::eye(3), rtol).unwrap();
    }

    #[test]
    fn append() {
        test_append::<f32>(1e-5);
        test_append::<c32>(1e-5);
        test_append::<f64>(1e-9);
        test_append::<c64>(1e-9);
    }

    fn test_append_if<A: Scalar + Lapack>(rtol: A::Real) {
        let mut mgs: MGS<A> = MGS::new(N);
        let a: Array2<A> = crate::generate::random((N, 8));
        dbg!(&a);
        for col in a.axis_iter(Axis(1)) {
            match mgs.append_if(col, rtol) {
                Some(res) => {
                    dbg!(res);
                }
                None => break,
            }
        }
        let q = mgs.get_q();
        dbg!(&q);
        let r = mgs.get_r();
        dbg!(&r);

        dbg!(q.dot(&r));
        close_l2(&q.dot(&r), &a.slice(s![.., 0..N]), rtol).unwrap();

        let qt: Array2<_> = conjugate(&q);
        dbg!(qt.dot(&q));
        close_l2(&qt.dot(&q), &Array2::eye(N), rtol).unwrap();
    }

    #[test]
    fn append_if() {
        test_append_if::<f32>(1e-5);
        test_append_if::<c32>(1e-5);
        test_append_if::<f64>(1e-9);
        test_append_if::<c64>(1e-9);
    }
}
<commit_msg>Rewrite tests for MGS::append<commit_after>use crate::{generate::*, inner::*, norm::Norm, types::*};
use ndarray::*;

/// Iterative orthogonalizer using modified Gram-Schmit procedure
#[derive(Debug, Clone)]
pub struct MGS<A> {
    /// Dimension of base space
    dimension: usize,
    /// Basis of spanned space
    q: Vec<Array1<A>>,
}

/// Residual vector of orthogonalization
pub type Residual<S> = ArrayBase<S, Ix1>;
/// Residual vector of orthogonalization
pub type Coefficient<A> = Array1<A>;
/// Q-matrix (unitary matrix)
pub type Q<A> = Array2<A>;
/// R-matrix (upper triangle)
pub type R<A> = Array2<A>;

impl<A: Scalar> MGS<A> {
    /// Create empty linear space
    ///
    /// ```rust
    /// # use ndarray_linalg::*;
    /// const N: usize = 5;
    /// let mgs = arnoldi::MGS::<f32>::new(N);
    /// assert_eq!(mgs.dim(), N);
    /// assert_eq!(mgs.len(), 0);
    /// ```
    pub fn new(dimension: usize) -> Self {
        Self {
            dimension,
            q: Vec::new(),
        }
    }

    pub fn dim(&self) -> usize {
        self.dimension
    }

    pub fn len(&self) -> usize {
        self.q.len()
    }

    /// Orthogonalize given vector using current basis
    ///
    /// Panic
    /// -------
    /// - if the size of the input array mismaches to the dimension
    pub fn orthogonalize<S>(&self, mut a: ArrayBase<S, Ix1>) -> (Residual<S>, Coefficient<A>)
    where
        A: Lapack,
        S: DataMut<Elem = A>,
    {
        assert_eq!(a.len(), self.dim());
        let mut coef = Array1::zeros(self.len() + 1);
        for i in 0..self.len() {
            let q = &self.q[i];
            let c = q.inner(&a);
            azip!(mut a, q (q) in { *a = *a - c * q } );
            coef[i] = c;
        }
        let nrm = a.norm_l2();
        coef[self.len()] = A::from_real(nrm);
        (a, coef)
    }

    /// Add new vector if the residual is larger than relative tolerance
    ///
    /// Panic
    /// -------
    /// - if the size of the input array mismaches to the dimension
    ///
    /// ```rust
    /// # use ndarray::*;
    /// # use ndarray_linalg::*;
    /// let mut mgs = arnoldi::MGS::new(3);
    /// let coef = mgs.append(array![1.0, 0.0, 0.0], 1e-9).unwrap();
    /// close_l2(&coef, &array![1.0], 1e-9).unwrap();
    ///
    /// let coef = mgs.append(array![1.0, 1.0, 0.0], 1e-9).unwrap();
    /// close_l2(&coef, &array![1.0, 1.0], 1e-9).unwrap();
    ///
    /// assert!(mgs.append(array![1.0, 1.0, 0.0], 1e-9).is_none());  // Cannot append dependent vector
    /// ```
    pub fn append<S>(&mut self, a: ArrayBase<S, Ix1>, rtol: A::Real) -> Option<Coefficient<A>>
    where
        A: Lapack,
        S: Data<Elem = A>,
    {
        let a = a.into_owned();
        let (mut a, coef) = self.orthogonalize(a);
        let nrm = coef[coef.len() - 1].re();
        if nrm < rtol {
            // Linearly dependent
            return None;
        }
        azip!(mut a in { *a = *a / A::from_real(nrm) });
        self.q.push(a);
        Some(coef)
    }

    /// Get orthogonal basis as Q matrix
    pub fn get_q(&self) -> Q<A> {
        hstack(&self.q).unwrap()
    }
}

pub fn mgs<A, S>(iter: impl Iterator<Item = ArrayBase<S, Ix1>>, dim: usize, rtol: A::Real) -> (Q<A>, R<A>)
where
    A: Scalar + Lapack,
    S: Data<Elem = A>,
{
    let mut ortho = MGS::new(dim);
    let mut coefs = Vec::new();
    for a in iter {
        match ortho.append(a, rtol) {
            Some(coef) => coefs.push(coef),
            None => break,
        }
    }
    let m = coefs.len();
    let mut r = Array2::zeros((m, m));
    for i in 0..m {
        for j in 0..=i {
            r[(j, i)] = coefs[i][j];
        }
    }
    (ortho.get_q(), r)
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    pub ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    pub context: Context,
}

pub struct Frame {
    pub ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    pub ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<commit_msg>Allocate frame data in Frame::new()<commit_after>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ptr;

use errors::*;

#[derive(Debug, Clone, Copy)]
pub struct Rational {
    pub numerator: i32,
    pub denominator: i32,
}

#[derive(Clone, Copy)]
pub struct Codec {
    ptr: *mut ffmpeg_sys::AVCodec,
}

unsafe impl Send for Codec {}
unsafe impl Sync for Codec {}

pub struct PixelFormats {
    codec: Codec,
    index: isize,
}

pub struct Context {
    // The Codec which was used to initialize this Context.
    codec: Codec,
    pub ptr: *mut ffmpeg_sys::AVCodecContext,
}

pub struct OpenContext {
    pub context: Context,
}

pub struct Frame {
    pub ptr: *mut ffmpeg_sys::AVFrame,
}

pub struct OutputContext {
    pub ptr: *mut ffmpeg_sys::AVFormatContext,
}

impl From<(i32, i32)> for Rational {
    fn from((numerator, denominator): (i32, i32)) -> Self {
        Self {
            numerator,
            denominator,
        }
    }
}

impl From<ffmpeg_sys::AVRational> for Rational {
    fn from(rational: ffmpeg_sys::AVRational) -> Self {
        Self {
            numerator: rational.num,
            denominator: rational.den,
        }
    }
}

impl From<Rational> for ffmpeg_sys::AVRational {
    fn from(rational: Rational) -> Self {
        Self {
            num: rational.numerator,
            den: rational.denominator,
        }
    }
}

impl Codec {
    pub fn description(self) -> String {
        unsafe {
            CStr::from_ptr((*self.ptr).long_name).to_string_lossy().into_owned()
        }
    }

    #[inline]
    pub fn kind(self) -> ffmpeg_sys::AVMediaType {
        unsafe {
            (*self.ptr).kind
        }
    }

    #[inline]
    pub fn is_video(self) -> bool {
        self.kind() == ffmpeg_sys::AVMediaType::AVMEDIA_TYPE_VIDEO
    }

    pub fn pixel_formats(self) -> Option<PixelFormats> {
        unsafe {
            if (*self.ptr).pix_fmts.is_null() {
                None
            } else {
                Some(PixelFormats::new(self))
            }
        }
    }

    pub fn context(self) -> Result<Context> {
        let ptr = unsafe {
            ffmpeg_sys::avcodec_alloc_context3(self.ptr)
        };

        ensure!(!ptr.is_null(), "unable to allocate the codec context");

        Ok(Context {
            codec: self,
            ptr,
        })
    }
}

impl PixelFormats {
    fn new(codec: Codec) -> Self {
        Self {
            codec,
            index: 0,
        }
    }
}

impl Iterator for PixelFormats {
    type Item = ffmpeg_sys::AVPixelFormat;

    fn next(&mut self) -> Option<Self::Item> {
        let format = unsafe {
            *(*self.codec.ptr).pix_fmts.offset(self.index)
        };

        if format == ffmpeg_sys::AVPixelFormat::AV_PIX_FMT_NONE {
            None
        } else {
            self.index += 1;
            Some(format)
        }
    }
}

impl Context {
    #[inline]
    pub fn width(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).width) as u32
        }
    }

    #[inline]
    pub fn set_width(&mut self, width: u32) {
        unsafe {
            (*self.ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn height(&self) -> u32 {
        unsafe {
            cmp::max(0, (*self.ptr).height) as u32
        }
    }

    #[inline]
    pub fn set_height(&mut self, height: u32) {
        unsafe {
            (*self.ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }
    }

    #[inline]
    pub fn time_base(&self) -> Rational {
        unsafe {
            (*self.ptr).time_base.into()
        }
    }

    #[inline]
    pub fn set_time_base(&mut self, time_base: &Rational) {
        unsafe {
            (*self.ptr).time_base = (*time_base).into();
        }
    }

    #[inline]
    pub fn pixel_format(&self) -> ffmpeg_sys::AVPixelFormat {
        unsafe {
            (*self.ptr).pix_fmt
        }
    }

    #[inline]
    pub fn set_pixel_format(&mut self, pixel_format: ffmpeg_sys::AVPixelFormat) {
        unsafe {
            (*self.ptr).pix_fmt = pixel_format;
        }
    }

    pub fn open(self) -> Result<OpenContext> {
        let rv = unsafe {
            ffmpeg_sys::avcodec_open2(self.ptr, self.codec.ptr, ptr::null_mut())
        };

        ensure!(rv == 0, "error opening context");

        Ok(OpenContext { context: self })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avcodec_free_context(&mut self.ptr);
        }
    }
}

impl Frame {
    pub fn new(pixel_format: ffmpeg_sys::AVPixelFormat,
               width: u32,
               height: u32) -> Result<Self> {
        let ptr = unsafe {
            ffmpeg_sys::av_frame_alloc()
        };

        ensure!(!ptr.is_null(), "unable to allocate a frame");

        unsafe {
            (*ptr).format = pixel_format as c_int;
            (*ptr).width = cmp::min(width, c_int::max_value() as u32) as c_int;
            (*ptr).height = cmp::min(height, c_int::max_value() as u32) as c_int;
        }

        let rv = unsafe {
            ffmpeg_sys::av_frame_get_buffer(ptr, 32)
        };

        ensure!(rv == 0, "unable to allocate frame data buffer");

        Ok(Self { ptr })
    }
}

impl Drop for Frame {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::av_frame_free(&mut self.ptr);
        }
    }
}

impl OutputContext {
    pub fn new(filename: &str) -> Result<Self> {
        let filename = CString::new(filename)
            .chain_err(|| "unable to convert filename to a CString")?;

        let mut ptr = ptr::null_mut();

        let rv = unsafe {
            ffmpeg_sys::avformat_alloc_output_context2(&mut ptr,
                                                       ptr::null_mut(),
                                                       ptr::null_mut(),
                                                       filename.as_ptr())
        };

        // TODO: check and report the error code.
        ensure!(rv >= 0, "unable to allocate the output context");

        Ok(Self { ptr })
    }
}

impl Drop for OutputContext {
    fn drop(&mut self) {
        unsafe {
            ffmpeg_sys::avformat_free_context(self.ptr);
        }
    }
}

pub fn initialize() {
    unsafe {
        ffmpeg_sys::av_register_all();
        ffmpeg_sys::avcodec_register_all();
    }
}

pub fn find_encoder_by_name(name: &str) -> Result<Option<Codec>> {
    let name = CString::new(name).chain_err(|| "could not convert name to CString")?;

    let codec = unsafe {
        ffmpeg_sys::avcodec_find_encoder_by_name(name.as_ptr())
    };

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<|endoftext|>


<commit_before>// TODO: cache values: Create new struct and use it within Backend

use std::collections::HashMap;

use conf::{RemoteBranch};
use constants::{CHANGED_KEY,NEW_KEY,STAGED_KEY,CONFLICTS_KEY};

use git2::*;

pub struct Backend {
    pub repo: Repository,
    pub debug: bool,
    pub current_branch_name: Option<String>  // cache
}

pub struct BranchAheadBehind {
    pub local_branch_name: Option<String>,
    pub remote_branch_name: Option<String>,
    pub remote_name: Option<String>,
    pub ahead: usize,
    pub behind: usize
}

impl BranchAheadBehind {
    fn new(l: Option<String>) -> BranchAheadBehind {
        BranchAheadBehind{ local_branch_name: l, remote_branch_name: None, remote_name: None,
                           ahead: 0, behind: 0 }
    }
}

#[derive(Clone)]
struct RefPair {
    remote_name: String,
    branch_name: String,
    oid: Oid,
}

impl Backend {
    pub fn new(repo: Repository, debug: bool) -> Backend {
        Backend{ repo: repo, debug: debug, current_branch_name: None }
    }

    fn get_head(&self) -> Option<Reference> {
        match self.repo.head() {
            Ok(head) => Some(head),
            Err(e2) => {
                log!(self, "Can't get HEAD: {}", e2);
                match self.repo.find_reference("HEAD") {
                    Ok(x) => {
                        log!(self, "Found HEAD directly: {:?}", x.name());
                        Some(x)
                    },
                    Err(e) => {
                        log!(self, "reference HEAD: {}", e);
                        return None
                    },
                }
            },
        }
    }

    fn get_current_branch_oid(&self) -> Option<Oid> {
        let head = self.get_head();
        match head {
            Some(v) => v.target(),
            None => {
                log!(self, "Failed to find Oid for HEAD.");
                return None
            },
        }
    }

    fn resolve_symbolic_reference<'a>(&self, reference: Option<Reference<'a>>) -> Option<Reference<'a>> {
        return reference;
        let s_reference = match reference {
            Some(v) => {
                match v.resolve() {
                    Ok(y) => {
                        log!(self, "Branch name is {:?}", y.name());
                        y
                    },
                    Err(e) => {
                        log!(self, "Branch name is {:?}, error: {:?}", v.name(), e);
                        v
                    }
                }
            }
            None => {
                log!(self, "No branch name found");
                return None;
            }
        };
        Some(s_reference)
    }

    fn get_branch_name_for_reference(&self, r: Reference) -> Option<String> {
        if let Some(ref_name_string) = r.shorthand() {
            if ref_name_string != "HEAD" {
                let s = ref_name_string.to_string();
                log!(self, "Shorthand for reference is: {}", s);
                return Some(s);
            } else {
                if let Some(ref_name) = r.symbolic_target() {
                    let ref_name_string = ref_name.to_string();
                    log!(self, "shorthand = HEAD, links to: {}", ref_name_string);
                    let mut path: Vec<&str> = ref_name_string.split('/').collect();
                    if let Some(branch_short) = path.pop() {
                        let s = branch_short.to_string();
                        log!(self, "Last part of full name is: {}", s);
                        return Some(s);
                    }
                }
            }
        }

        match r.target() {
            Some(v) => {
                let hash_str = v.to_string();
                if hash_str.len() >= 8 {
                    let (s, _) = hash_str.split_at(7);
                    Some(s.to_string())
                } else {
                    Some(hash_str)
                }
            },
            None => None,
        }
    }

    fn set_current_branch_name(&mut self, n: Option<String>) {
        self.current_branch_name = n;
    }

    pub fn get_current_branch_name(&mut self) -> Option<String> {
        if self.current_branch_name.is_some() {
            return self.current_branch_name.clone();
        }
        let mut head: Option<Reference> = None;
        if head.is_none() {
            head = self.get_head();
        }
        let h = match self.resolve_symbolic_reference(head) {
            Some(v) => v,
            None => return None,
        };
        let current_branch_name = self.get_branch_name_for_reference(h);
        current_branch_name
    }

    fn get_branch_remote(&self, reference: Reference) -> Option<RefPair> {
        let b = Branch::wrap(reference);
        let upstream = match b.upstream() {
            Ok(u) => u,
            Err(e) => {
                log!(self, "Can't get upstream branch for {:?}", b.name());
                return None;
            }
        };
        let branch_name = match upstream.name() {
            Ok(o_n) => match o_n {
                Some(n) => n,
                None => {
                    log!(self, "Invalid branch name");
                    return None;
                },
            },
            Err(e) => {
                log!(self, "Error while getting name for upstream branch");
                return None;
            }
        };
        let upstream_reference = upstream.get();
        let oid = match upstream_reference.target() {
            Some(o) => o,
            None => {
                log!(self, "Can't get oid of upstream branch");
                return None;
            }
        };
        let remote_name = match upstream_reference.name() {
            Some(n) => {
                let v: Vec<&str> = n.splitn(3, "/").collect();
                if v.len() >= 3 {
                    v[1]
                } else {
                    log!(self, "Can't figure out remote name: {:?}", v);
                    return None;
                }
            },
            None => {
                log!(self, "Can't get full name of upstream branch.");
                return None;
            }
        };
        Some(RefPair{ remote_name: remote_name.to_string(),
                      branch_name: branch_name.to_string(), oid: oid })
    }

    fn get_current_branch_remote_oid(&self) -> Option<RefPair> {
        match self.get_head() {
            Some(r) => self.get_branch_remote(r),
            None => None
        }
    }

    pub fn get_branch_ahead_behind(&mut self, remote_branch: Option<RemoteBranch>) -> Option<BranchAheadBehind> {
        let current_branch_name = self.get_current_branch_name();
        log!(self, "Current branch name = {:?}", current_branch_name);
        let mut ab = BranchAheadBehind::new(current_branch_name);
        let ref_pair_option = self.get_remote_branch(remote_branch);
        let ref_pair = match ref_pair_option {
            Some(u) => {
                u.clone()
            },
            None => return Some(ab),
        };
        ab.remote_branch_name = Some(ref_pair.branch_name.clone());
        ab.remote_name = Some(ref_pair.remote_name);

        let oid = match self.get_current_branch_oid() {
            Some(r) => r,
            None => return None
        };
        let res = self.repo.graph_ahead_behind(oid, ref_pair.oid);
        match res {
            Ok((a, b)) => {
                ab.ahead = a;
                ab.behind = b;
            },
            Err(e) => {
                log!(self, "Can't get ahead & behind stats for branch {}", ref_pair.branch_name);
            }
        };
        Some(ab)
    }

    // find remote branch if branch_name is specified
    // if not, get remote tracking branch for current branch
    fn get_remote_branch(&self, remote_branch: Option<RemoteBranch>) -> Option<RefPair> {
        match remote_branch {
            Some(b) => {
                match self.repo.find_branch(&b.remote_branch, BranchType::Remote) {
                    Ok(o) => Some(RefPair{ branch_name: b.remote_branch_name,
                                           remote_name: b.remote_name,
                                           oid: o.into_reference().target().unwrap() }),
                    Err(e) => {
                        // don't panic here - it doesn't exist, we don't care
                        log!(self, "No remote branch found for {}", b.remote_branch_name);
                        None
                    }
                }
            },
            None => {
                self.get_current_branch_remote_oid()
            }
        }
    }

    pub fn get_status(&self) -> Option<Statuses> {
        let mut so = StatusOptions::new();
        let mut opts = so.show(StatusShow::IndexAndWorkdir);
        opts.include_untracked(true);
        match self.repo.statuses(Some(&mut opts)) {
            Ok(s) => Some(s),
            Err(e) => {
                log!(self, "Unable to get status of repository: {:?}", e);
                None
            }
        }
    }

    pub fn get_repository_state(&self) -> String {
        let state = self.repo.state();
        match state {
            RepositoryState::Clean => String::from(""),  // XXX: this seems to be hardcoded
            RepositoryState::Merge => String::from("merge"),
            RepositoryState::Revert | RepositoryState::RevertSequence => String::from("revert"),
            RepositoryState::CherryPick | RepositoryState::CherryPickSequence => String::from("cherry-pick"),
            RepositoryState::Bisect => String::from("bisect"),
            RepositoryState::Rebase |
                RepositoryState::RebaseInteractive |
                RepositoryState::RebaseMerge => String::from("rebase"),
            RepositoryState::ApplyMailbox | RepositoryState::ApplyMailboxOrRebase => String::from("apply"),
        }
    }

    pub fn get_file_status(&self) -> Option<HashMap<&str, u32>> {
        let mut d = HashMap::new();

        let changed = STATUS_WT_MODIFIED | STATUS_WT_DELETED | STATUS_WT_TYPECHANGE | STATUS_WT_RENAMED;
        let staged = STATUS_INDEX_MODIFIED | STATUS_INDEX_DELETED | STATUS_INDEX_TYPECHANGE | STATUS_INDEX_RENAMED | STATUS_INDEX_NEW;

        let statuses = match self.get_status() {
            Some(x) => x,
            None => return None,
        };

        for s in statuses.iter() {
            let file_status = s.status();
            log!(self, "{}", s.path().unwrap());

            if file_status.intersects(changed) {
                let counter = d.entry(CHANGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.contains(STATUS_WT_NEW) {
                let counter = d.entry(NEW_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(staged) {
                let counter = d.entry(STAGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(STATUS_CONFLICTED) {
                let counter = d.entry(CONFLICTS_KEY).or_insert(0);
                *counter += 1;
            };
        }
        Some(d)
    }
}
<commit_msg>cache current_branch_name<commit_after>use std::collections::HashMap;
use std::cell::RefCell;

use conf::{RemoteBranch};
use constants::{CHANGED_KEY,NEW_KEY,STAGED_KEY,CONFLICTS_KEY};

use git2::*;

struct Cache {
    current_branch_name: RefCell<Option<String>>,
    // TODO: Reference can't be cached (can't be cloned)
    //       implement via OID
    // head: RefCell<Option<Oid>>,
}

pub struct Backend {
    cache: Cache,
    pub repo: Repository,
    pub debug: bool,
}

pub struct BranchAheadBehind {
    pub local_branch_name: Option<String>,
    pub remote_branch_name: Option<String>,
    pub remote_name: Option<String>,
    pub ahead: usize,
    pub behind: usize
}

impl BranchAheadBehind {
    fn new(l: Option<String>) -> BranchAheadBehind {
        BranchAheadBehind{ local_branch_name: l, remote_branch_name: None, remote_name: None,
                           ahead: 0, behind: 0 }
    }
}

#[derive(Clone)]
struct RefPair {
    remote_name: String,
    branch_name: String,
    oid: Oid,
}

impl Cache {
    pub fn new() -> Cache {
        Cache{ current_branch_name: RefCell::new(None) }
    }

    fn set_current_branch_name(&self, n: Option<String>) {
        let mut c = self.current_branch_name.borrow_mut();
        *c = n;
    }

    fn is_current_branch_set(&self) -> bool {
        self.current_branch_name.borrow().is_some()
    }

    fn get_current_branch(&self) -> Option<String> {
        self.current_branch_name.borrow().clone()
    }
}

impl Backend {
    pub fn new(repo: Repository, debug: bool) -> Backend {
        Backend{ repo: repo, debug: debug, cache: Cache::new() }
    }

    fn get_head(&self) -> Option<Reference> {
        match self.repo.head() {
            Ok(head) => Some(head),
            Err(e2) => {
                log!(self, "Can't get HEAD: {}", e2);
                match self.repo.find_reference("HEAD") {
                    Ok(x) => {
                        log!(self, "Found HEAD directly: {:?}", x.name());
                        Some(x)
                    },
                    Err(e) => {
                        log!(self, "reference HEAD: {}", e);
                        return None
                    },
                }
            },
        }
    }

    fn get_current_branch_oid(&self) -> Option<Oid> {
        let head = self.get_head();
        match head {
            Some(v) => v.target(),
            None => {
                log!(self, "Failed to find Oid for HEAD.");
                return None
            },
        }
    }

    fn resolve_symbolic_reference<'a>(&self, reference: Option<Reference<'a>>) -> Option<Reference<'a>> {
        return reference;
        let s_reference = match reference {
            Some(v) => {
                match v.resolve() {
                    Ok(y) => {
                        log!(self, "Branch name is {:?}", y.name());
                        y
                    },
                    Err(e) => {
                        log!(self, "Branch name is {:?}, error: {:?}", v.name(), e);
                        v
                    }
                }
            }
            None => {
                log!(self, "No branch name found");
                return None;
            }
        };
        Some(s_reference)
    }

    fn get_branch_name_for_reference(&self, r: Reference) -> Option<String> {
        if let Some(ref_name_string) = r.shorthand() {
            if ref_name_string != "HEAD" {
                let s = ref_name_string.to_string();
                log!(self, "Shorthand for reference is: {}", s);
                return Some(s);
            } else {
                if let Some(ref_name) = r.symbolic_target() {
                    let ref_name_string = ref_name.to_string();
                    log!(self, "shorthand = HEAD, links to: {}", ref_name_string);
                    let mut path: Vec<&str> = ref_name_string.split('/').collect();
                    if let Some(branch_short) = path.pop() {
                        let s = branch_short.to_string();
                        log!(self, "Last part of full name is: {}", s);
                        return Some(s);
                    }
                }
            }
        }

        match r.target() {
            Some(v) => {
                let hash_str = v.to_string();
                if hash_str.len() >= 8 {
                    let (s, _) = hash_str.split_at(7);
                    Some(s.to_string())
                } else {
                    Some(hash_str)
                }
            },
            None => None,
        }
    }

    pub fn get_current_branch_name(&mut self) -> Option<String> {
        if self.cache.is_current_branch_set() {
            return self.cache.get_current_branch();
        }
        let mut head: Option<Reference> = None;
        if head.is_none() {
            head = self.get_head();
        }
        let h = match self.resolve_symbolic_reference(head) {
            Some(v) => v,
            None => return None,
        };
        let current_branch_name = self.get_branch_name_for_reference(h);
        self.cache.set_current_branch_name(current_branch_name.clone());
        current_branch_name
    }

    fn get_branch_remote(&self, reference: Reference) -> Option<RefPair> {
        let b = Branch::wrap(reference);
        let upstream = match b.upstream() {
            Ok(u) => u,
            Err(e) => {
                log!(self, "Can't get upstream branch for {:?}", b.name());
                return None;
            }
        };
        let branch_name = match upstream.name() {
            Ok(o_n) => match o_n {
                Some(n) => n,
                None => {
                    log!(self, "Invalid branch name");
                    return None;
                },
            },
            Err(e) => {
                log!(self, "Error while getting name for upstream branch");
                return None;
            }
        };
        let upstream_reference = upstream.get();
        let oid = match upstream_reference.target() {
            Some(o) => o,
            None => {
                log!(self, "Can't get oid of upstream branch");
                return None;
            }
        };
        let remote_name = match upstream_reference.name() {
            Some(n) => {
                let v: Vec<&str> = n.splitn(3, "/").collect();
                if v.len() >= 3 {
                    v[1]
                } else {
                    log!(self, "Can't figure out remote name: {:?}", v);
                    return None;
                }
            },
            None => {
                log!(self, "Can't get full name of upstream branch.");
                return None;
            }
        };
        Some(RefPair{ remote_name: remote_name.to_string(),
                      branch_name: branch_name.to_string(), oid: oid })
    }

    fn get_current_branch_remote_oid(&self) -> Option<RefPair> {
        match self.get_head() {
            Some(r) => self.get_branch_remote(r),
            None => None
        }
    }

    pub fn get_branch_ahead_behind(&mut self, remote_branch: Option<RemoteBranch>) -> Option<BranchAheadBehind> {
        let current_branch_name = self.get_current_branch_name();
        log!(self, "Current branch name = {:?}", current_branch_name);
        let mut ab = BranchAheadBehind::new(current_branch_name);
        let ref_pair_option = self.get_remote_branch(remote_branch);
        let ref_pair = match ref_pair_option {
            Some(u) => {
                u.clone()
            },
            None => return Some(ab),
        };
        ab.remote_branch_name = Some(ref_pair.branch_name.clone());
        ab.remote_name = Some(ref_pair.remote_name);

        let oid = match self.get_current_branch_oid() {
            Some(r) => r,
            None => return None
        };
        let res = self.repo.graph_ahead_behind(oid, ref_pair.oid);
        match res {
            Ok((a, b)) => {
                ab.ahead = a;
                ab.behind = b;
            },
            Err(e) => {
                log!(self, "Can't get ahead & behind stats for branch {}", ref_pair.branch_name);
            }
        };
        Some(ab)
    }

    // find remote branch if branch_name is specified
    // if not, get remote tracking branch for current branch
    fn get_remote_branch(&self, remote_branch: Option<RemoteBranch>) -> Option<RefPair> {
        match remote_branch {
            Some(b) => {
                match self.repo.find_branch(&b.remote_branch, BranchType::Remote) {
                    Ok(o) => Some(RefPair{ branch_name: b.remote_branch_name,
                                           remote_name: b.remote_name,
                                           oid: o.into_reference().target().unwrap() }),
                    Err(e) => {
                        // don't panic here - it doesn't exist, we don't care
                        log!(self, "No remote branch found for {}", b.remote_branch_name);
                        None
                    }
                }
            },
            None => {
                self.get_current_branch_remote_oid()
            }
        }
    }

    pub fn get_status(&self) -> Option<Statuses> {
        let mut so = StatusOptions::new();
        let mut opts = so.show(StatusShow::IndexAndWorkdir);
        opts.include_untracked(true);
        match self.repo.statuses(Some(&mut opts)) {
            Ok(s) => Some(s),
            Err(e) => {
                log!(self, "Unable to get status of repository: {:?}", e);
                None
            }
        }
    }

    pub fn get_repository_state(&self) -> String {
        let state = self.repo.state();
        match state {
            RepositoryState::Clean => String::from(""),  // XXX: this seems to be hardcoded
            RepositoryState::Merge => String::from("merge"),
            RepositoryState::Revert | RepositoryState::RevertSequence => String::from("revert"),
            RepositoryState::CherryPick | RepositoryState::CherryPickSequence => String::from("cherry-pick"),
            RepositoryState::Bisect => String::from("bisect"),
            RepositoryState::Rebase |
                RepositoryState::RebaseInteractive |
                RepositoryState::RebaseMerge => String::from("rebase"),
            RepositoryState::ApplyMailbox | RepositoryState::ApplyMailboxOrRebase => String::from("apply"),
        }
    }

    pub fn get_file_status(&self) -> Option<HashMap<&str, u32>> {
        let mut d = HashMap::new();

        let changed = STATUS_WT_MODIFIED | STATUS_WT_DELETED | STATUS_WT_TYPECHANGE | STATUS_WT_RENAMED;
        let staged = STATUS_INDEX_MODIFIED | STATUS_INDEX_DELETED | STATUS_INDEX_TYPECHANGE | STATUS_INDEX_RENAMED | STATUS_INDEX_NEW;

        let statuses = match self.get_status() {
            Some(x) => x,
            None => return None,
        };

        for s in statuses.iter() {
            let file_status = s.status();
            log!(self, "{}", s.path().unwrap());

            if file_status.intersects(changed) {
                let counter = d.entry(CHANGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.contains(STATUS_WT_NEW) {
                let counter = d.entry(NEW_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(staged) {
                let counter = d.entry(STAGED_KEY).or_insert(0);
                *counter += 1;
            };
            if file_status.intersects(STATUS_CONFLICTED) {
                let counter = d.entry(CONFLICTS_KEY).or_insert(0);
                *counter += 1;
            };
        }
        Some(d)
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate extra;

use std::env;
use std::error::Error;
use std::fs;
use std::io::{self, Read, Stderr, StdoutLock, Write};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = r#"NAME
    cat - concatenate files and print on the standard output

SYNOPSIS
    cat [-h | --help] [-A | --show-all] [-b | --number-nonblank] [-e] [-E | --show-ends] [-n | --number] [-s | --squeeze-blank] [-t] [-T] FILES...

DESCRIPTION
    Concatenates all files to the standard output.

    If no file is given, or if FILE is '-', read from standard input.

OPTIONS
    -A
    --show-all
        equivalent to -v -E -T

    -b
    --number-nonblank
        number nonempty output lines, overriding -n

    -e
        equivalent to -v -E

    -E
    --show-ends
        display $ at the end of each line

    -n
    --number
        number all output lines

    -s
    --squeeze-blank
        supress repeated empty output lines

    -t
        equivalent to -v -T

    -T
    --show_tabs
        display TAB characters as ^I

    -v
    --show-nonprinting
        use ^ and M- notation,except for LFD and TAB.

    -h
    --help
        display this help and exit

AUTHOR
    Written by Michael Murphy.
"#;

struct Program {
    number:          bool,
    number_nonblank: bool,
    show_ends:       bool,
    show_tabs:       bool,
    squeeze_blank:   bool,
    paths:           Vec<String>,
}

impl Program {
    /// Initialize the program's arguments and flags.
    fn initialize(stdout: &mut StdoutLock, stderr: &mut Stderr) -> Program {
        let mut cat = Program {
            number:          false,
            number_nonblank: false,
            show_ends:       false,
            show_tabs:       false,
            squeeze_blank:   false,
            paths:           Vec::new()
        };
        for arg in env::args().skip(1) {
            if arg.starts_with('-') {
                match arg.as_str() {
                    "-h" | "--help" => {
                        stdout.write(MAN_PAGE.as_bytes()).try(stderr);
                        stdout.flush().try(stderr);
                        std::process::exit(0);
                    }
                    "-A" | "--show-all" => {
                        // TODO: implement -v
                        cat.show_ends = true;
                        cat.show_tabs = true;
                    },
                    "-b" | "--number-nonblank" => {
                        cat.number_nonblank = true;
                        cat.number = false;
                    },
                    "-e" => {
                        // TODO: implement -v
                        cat.show_ends = true;
                    },
                    "-E" | "--show-ends" => cat.show_ends = true,
                    "-n" | "--number" => {
                        cat.number = true;
                        cat.number_nonblank = false;
                    },
                    "-s" | "--squeeze-blank" => cat.squeeze_blank = true,
                    "-t" => {
                        // TODO: implement -v
                        cat.show_tabs = true;
                    },
                    "-T" => cat.show_tabs = true,
                    "-v" | "--show-nonprinting" => {
                        // TODO: implement -v, --show-nonprinting
                        continue
                    },
                    _ => {
                        stderr.write(b"invalid option -- '").try(stderr);
                        stderr.write(arg.as_bytes()).try(stderr);
                        stderr.write(b"'\nTry 'cat --help' for more information.\n").try(stderr);
                        std::process::exit(1);
                    }
                }
            } else {
                cat.paths.push(arg);
            }
        }
        cat
    }

    /// Execute the parameters given to the program.
    fn execute(&self, stdout: &mut StdoutLock, stderr: &mut Stderr) {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut exit_status = 0i32;

        if self.paths.len() == 0 {
            io::copy(&mut stdin, stdout).try(stderr);
        } else {
            let mut line_count = 1;
            for path in &self.paths {
                if self.number || self.number_nonblank || self.show_ends || self.show_tabs || self.squeeze_blank {
                    let file = match fs::File::open(&path) {
                        Ok(file) => file,
                        Err(message) => {
                            stderr.write(&path.as_bytes()).try(stderr);
                            stderr.write(b": ").try(stderr);
                            stderr.write(message.description().as_bytes()).try(stderr);
                            stderr.write(b"\n").try(stderr);
                            stderr.flush().try(stderr);
                            exit_status = 1;
                            continue
                        }
                    };
                    let mut current_line: Vec<u8> = Vec::new();
                    let mut last_line_was_blank = false;
                    for byte in file.bytes().map(|x| x.unwrap_or(b' ')) {
                        if byte == b'\n' {
                            if current_line.is_empty() {
                                if last_line_was_blank && self.squeeze_blank {
                                    continue
                                } else if !last_line_was_blank {
                                    last_line_was_blank = true;
                                }
                            }
                            if self.number || (self.number_nonblank && !current_line.is_empty()) {
                                stdout.write(format!("     {}  ", line_count).as_bytes()).try(stderr);
                                line_count += 1;
                            }
                            stdout.write(current_line.as_slice()).try(stderr);
                            if self.show_ends { stdout.write(b"$\n").try(stderr); } else { stdout.write(b"\n").try(stderr); }
                            stdout.flush().try(stderr);
                            current_line.clear();
                        } else {
                            if self.show_tabs && byte == b'\t' {
                                current_line.push(b'^');
                                current_line.push(b'I');
                            } else {
                                current_line.push(byte);
                            }
                        }
                    }
                } else {
                    if path == "-" {
                        io::copy(&mut stdin, stdout).try(stderr);
                    } else {
                        let file = fs::File::open(&path);
                        let mut file = file.try(stderr);

                        io::copy(&mut file, stdout).try(stderr);
                    }
                }
            }
        }
        std::process::exit(exit_status);
    }
}

fn main() {
    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut stderr = io::stderr();
    Program::initialize(&mut stdout, &mut stderr).execute(&mut stdout, &mut stderr);
}
<commit_msg>forgot to flush<commit_after>#![deny(warnings)]

extern crate extra;

use std::env;
use std::error::Error;
use std::fs;
use std::io::{self, Read, Stderr, StdoutLock, Write};
use extra::option::OptionalExt;

const MAN_PAGE: &'static str = r#"NAME
    cat - concatenate files and print on the standard output

SYNOPSIS
    cat [-h | --help] [-A | --show-all] [-b | --number-nonblank] [-e] [-E | --show-ends] [-n | --number] [-s | --squeeze-blank] [-t] [-T] FILES...

DESCRIPTION
    Concatenates all files to the standard output.

    If no file is given, or if FILE is '-', read from standard input.

OPTIONS
    -A
    --show-all
        equivalent to -v -E -T

    -b
    --number-nonblank
        number nonempty output lines, overriding -n

    -e
        equivalent to -v -E

    -E
    --show-ends
        display $ at the end of each line

    -n
    --number
        number all output lines

    -s
    --squeeze-blank
        supress repeated empty output lines

    -t
        equivalent to -v -T

    -T
    --show_tabs
        display TAB characters as ^I

    -v
    --show-nonprinting
        use ^ and M- notation,except for LFD and TAB.

    -h
    --help
        display this help and exit

AUTHOR
    Written by Michael Murphy.
"#;

struct Program {
    number:          bool,
    number_nonblank: bool,
    show_ends:       bool,
    show_tabs:       bool,
    squeeze_blank:   bool,
    paths:           Vec<String>,
}

impl Program {
    /// Initialize the program's arguments and flags.
    fn initialize(stdout: &mut StdoutLock, stderr: &mut Stderr) -> Program {
        let mut cat = Program {
            number:          false,
            number_nonblank: false,
            show_ends:       false,
            show_tabs:       false,
            squeeze_blank:   false,
            paths:           Vec::new()
        };
        for arg in env::args().skip(1) {
            if arg.starts_with('-') {
                match arg.as_str() {
                    "-h" | "--help" => {
                        stdout.write(MAN_PAGE.as_bytes()).try(stderr);
                        stdout.flush().try(stderr);
                        std::process::exit(0);
                    }
                    "-A" | "--show-all" => {
                        // TODO: implement -v
                        cat.show_ends = true;
                        cat.show_tabs = true;
                    },
                    "-b" | "--number-nonblank" => {
                        cat.number_nonblank = true;
                        cat.number = false;
                    },
                    "-e" => {
                        // TODO: implement -v
                        cat.show_ends = true;
                    },
                    "-E" | "--show-ends" => cat.show_ends = true,
                    "-n" | "--number" => {
                        cat.number = true;
                        cat.number_nonblank = false;
                    },
                    "-s" | "--squeeze-blank" => cat.squeeze_blank = true,
                    "-t" => {
                        // TODO: implement -v
                        cat.show_tabs = true;
                    },
                    "-T" => cat.show_tabs = true,
                    "-v" | "--show-nonprinting" => {
                        // TODO: implement -v, --show-nonprinting
                        continue
                    },
                    _ => {
                        stderr.write(b"invalid option -- '").try(stderr);
                        stderr.write(arg.as_bytes()).try(stderr);
                        stderr.write(b"'\nTry 'cat --help' for more information.\n").try(stderr);
                        stderr.flush().try(stderr);
                        std::process::exit(1);
                    }
                }
            } else {
                cat.paths.push(arg);
            }
        }
        cat
    }

    /// Execute the parameters given to the program.
    fn execute(&self, stdout: &mut StdoutLock, stderr: &mut Stderr) {
        let stdin = io::stdin();
        let mut stdin = stdin.lock();
        let mut exit_status = 0i32;

        if self.paths.len() == 0 {
            io::copy(&mut stdin, stdout).try(stderr);
        } else {
            let mut line_count = 1;
            for path in &self.paths {
                if self.number || self.number_nonblank || self.show_ends || self.show_tabs || self.squeeze_blank {
                    let file = match fs::File::open(&path) {
                        Ok(file) => file,
                        Err(message) => {
                            stderr.write(&path.as_bytes()).try(stderr);
                            stderr.write(b": ").try(stderr);
                            stderr.write(message.description().as_bytes()).try(stderr);
                            stderr.write(b"\n").try(stderr);
                            stderr.flush().try(stderr);
                            exit_status = 1;
                            continue
                        }
                    };
                    let mut current_line: Vec<u8> = Vec::new();
                    let mut last_line_was_blank = false;
                    for byte in file.bytes().map(|x| x.unwrap_or(b' ')) {
                        if byte == b'\n' {
                            if current_line.is_empty() {
                                if last_line_was_blank && self.squeeze_blank {
                                    continue
                                } else if !last_line_was_blank {
                                    last_line_was_blank = true;
                                }
                            }
                            if self.number || (self.number_nonblank && !current_line.is_empty()) {
                                stdout.write(format!("     {}  ", line_count).as_bytes()).try(stderr);
                                line_count += 1;
                            }
                            stdout.write(current_line.as_slice()).try(stderr);
                            if self.show_ends { stdout.write(b"$\n").try(stderr); } else { stdout.write(b"\n").try(stderr); }
                            stdout.flush().try(stderr);
                            current_line.clear();
                        } else {
                            if self.show_tabs && byte == b'\t' {
                                current_line.push(b'^');
                                current_line.push(b'I');
                            } else {
                                current_line.push(byte);
                            }
                        }
                    }
                } else {
                    if path == "-" {
                        io::copy(&mut stdin, stdout).try(stderr);
                    } else {
                        let file = fs::File::open(&path);
                        let mut file = file.try(stderr);

                        io::copy(&mut file, stdout).try(stderr);
                    }
                }
            }
        }
        std::process::exit(exit_status);
    }
}

fn main() {
    let stdout = io::stdout();
    let mut stdout = stdout.lock();
    let mut stderr = io::stderr();
    Program::initialize(&mut stdout, &mut stderr).execute(&mut stdout, &mut stderr);
}
<|endoftext|>


<commit_before>extern crate clap;

use clap::{AppSettings, Arg, ArgMatches};

use cargo::{self, CliResult, Config};

use super::list_commands;
use super::commands;
use command_prelude::*;

pub fn main(config: &mut Config) -> CliResult {
    let args = cli().get_matches_safe()?;

    if args.value_of("unstable-features") == Some("help") {
        println!(
            "
Available unstable (nightly-only) flags:

    -Z avoid-dev-deps   -- Avoid installing dev-dependencies if possible
    -Z minimal-versions -- Install minimal dependency versions instead of maximum
    -Z no-index-update  -- Do not update the registry, avoids a network request for benchmarking
    -Z offline          -- Offline mode that does not perform network requests
    -Z unstable-options -- Allow the usage of unstable options such as --registry

Run with 'cargo -Z [FLAG] [SUBCOMMAND]'"
        );
        return Ok(());
    }

    let is_verbose = args.occurrences_of("verbose") > 0;
    if args.is_present("version") {
        let version = cargo::version();
        println!("{}", version);
        if is_verbose {
            println!(
                "release: {}.{}.{}",
                version.major, version.minor, version.patch
            );
            if let Some(ref cfg) = version.cfg_info {
                if let Some(ref ci) = cfg.commit_info {
                    println!("commit-hash: {}", ci.commit_hash);
                    println!("commit-date: {}", ci.commit_date);
                }
            }
        }
        return Ok(());
    }

    if let Some(ref code) = args.value_of("explain") {
        let mut procss = config.rustc()?.process();
        procss.arg("--explain").arg(code).exec()?;
        return Ok(());
    }

    if args.is_present("list") {
        println!("Installed Commands:");
        for command in list_commands(config) {
            let (command, path) = command;
            if is_verbose {
                match path {
                    Some(p) => println!("    {:<20} {}", command, p),
                    None => println!("    {:<20}", command),
                }
            } else {
                println!("    {}", command);
            }
        }
        return Ok(());
    }

    if args.subcommand_name().is_none() {}

    execute_subcommand(config, args)
}

fn execute_subcommand(config: &mut Config, args: ArgMatches) -> CliResult {
    config.configure(
        args.occurrences_of("verbose") as u32,
        if args.is_present("quiet") {
            Some(true)
        } else {
            None
        },
        &args.value_of("color").map(|s| s.to_string()),
        args.is_present("frozen"),
        args.is_present("locked"),
        &args.values_of_lossy("unstable-features")
            .unwrap_or_default(),
    )?;

    let (cmd, args) = match args.subcommand() {
        (cmd, Some(args)) => (cmd, args),
        _ => {
            cli().print_help()?;
            return Ok(());
        }
    };

    if let Some(exec) = commands::builtin_exec(cmd) {
        return exec(config, args);
    }

    if let Some(mut alias) = super::aliased_command(config, cmd)? {
        alias.extend(
            args.values_of("")
                .unwrap_or_default()
                .map(|s| s.to_string()),
        );
        let args = cli()
            .setting(AppSettings::NoBinaryName)
            .get_matches_from_safe(alias)?;
        return execute_subcommand(config, args);
    }
    let mut ext_args: Vec<&str> = vec![cmd];
    ext_args.extend(args.values_of("").unwrap_or_default());
    super::execute_external_subcommand(config, cmd, &ext_args)
}

fn cli() -> App {
    let app = App::new("cargo")
        .settings(&[
            AppSettings::UnifiedHelpMessage,
            AppSettings::DeriveDisplayOrder,
            AppSettings::VersionlessSubcommands,
            AppSettings::AllowExternalSubcommands,
        ])
        .about("")
        .template(
            "\
Rust's package manager

USAGE:
    {usage}

OPTIONS:
{unified}

Some common cargo commands are (see all commands with --list):
    build       Compile the current project
    check       Analyze the current project and report errors, but don't build object files
    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
    uninstall   Uninstall a Rust binary

See 'cargo help <command>' for more information on a specific command.\n",
        )
        .arg(opt("version", "Print version info and exit").short("V"))
        .arg(opt("list", "List installed commands"))
        .arg(opt("explain", "Run `rustc --explain CODE`").value_name("CODE"))
        .arg(
            opt(
                "verbose",
                "Use verbose output (-vv very verbose/build.rs output)",
            ).short("v")
                .multiple(true)
                .global(true),
        )
        .arg(
            opt("quiet", "No output printed to stdout")
                .short("q")
                .global(true),
        )
        .arg(
            opt("color", "Coloring: auto, always, never")
                .value_name("WHEN")
                .global(true),
        )
        .arg(opt("frozen", "Require Cargo.lock and cache are up to date").global(true))
        .arg(opt("locked", "Require Cargo.lock is up to date").global(true))
        .arg(
            Arg::with_name("unstable-features")
                .help("Unstable (nightly-only) flags to Cargo, see 'cargo -Z help' for details")
                .short("Z")
                .value_name("FLAG")
                .multiple(true)
                .number_of_values(1)
                .global(true),
        )
        .subcommands(commands::builtin());
    app
}
<commit_msg>Auto merge of #5276 - klausi:main-dead-code, r=matklad<commit_after>extern crate clap;

use clap::{AppSettings, Arg, ArgMatches};

use cargo::{self, CliResult, Config};

use super::list_commands;
use super::commands;
use command_prelude::*;

pub fn main(config: &mut Config) -> CliResult {
    let args = cli().get_matches_safe()?;

    if args.value_of("unstable-features") == Some("help") {
        println!(
            "
Available unstable (nightly-only) flags:

    -Z avoid-dev-deps   -- Avoid installing dev-dependencies if possible
    -Z minimal-versions -- Install minimal dependency versions instead of maximum
    -Z no-index-update  -- Do not update the registry, avoids a network request for benchmarking
    -Z offline          -- Offline mode that does not perform network requests
    -Z unstable-options -- Allow the usage of unstable options such as --registry

Run with 'cargo -Z [FLAG] [SUBCOMMAND]'"
        );
        return Ok(());
    }

    let is_verbose = args.occurrences_of("verbose") > 0;
    if args.is_present("version") {
        let version = cargo::version();
        println!("{}", version);
        if is_verbose {
            println!(
                "release: {}.{}.{}",
                version.major, version.minor, version.patch
            );
            if let Some(ref cfg) = version.cfg_info {
                if let Some(ref ci) = cfg.commit_info {
                    println!("commit-hash: {}", ci.commit_hash);
                    println!("commit-date: {}", ci.commit_date);
                }
            }
        }
        return Ok(());
    }

    if let Some(ref code) = args.value_of("explain") {
        let mut procss = config.rustc()?.process();
        procss.arg("--explain").arg(code).exec()?;
        return Ok(());
    }

    if args.is_present("list") {
        println!("Installed Commands:");
        for command in list_commands(config) {
            let (command, path) = command;
            if is_verbose {
                match path {
                    Some(p) => println!("    {:<20} {}", command, p),
                    None => println!("    {:<20}", command),
                }
            } else {
                println!("    {}", command);
            }
        }
        return Ok(());
    }

    execute_subcommand(config, args)
}

fn execute_subcommand(config: &mut Config, args: ArgMatches) -> CliResult {
    config.configure(
        args.occurrences_of("verbose") as u32,
        if args.is_present("quiet") {
            Some(true)
        } else {
            None
        },
        &args.value_of("color").map(|s| s.to_string()),
        args.is_present("frozen"),
        args.is_present("locked"),
        &args.values_of_lossy("unstable-features")
            .unwrap_or_default(),
    )?;

    let (cmd, args) = match args.subcommand() {
        (cmd, Some(args)) => (cmd, args),
        _ => {
            cli().print_help()?;
            return Ok(());
        }
    };

    if let Some(exec) = commands::builtin_exec(cmd) {
        return exec(config, args);
    }

    if let Some(mut alias) = super::aliased_command(config, cmd)? {
        alias.extend(
            args.values_of("")
                .unwrap_or_default()
                .map(|s| s.to_string()),
        );
        let args = cli()
            .setting(AppSettings::NoBinaryName)
            .get_matches_from_safe(alias)?;
        return execute_subcommand(config, args);
    }
    let mut ext_args: Vec<&str> = vec![cmd];
    ext_args.extend(args.values_of("").unwrap_or_default());
    super::execute_external_subcommand(config, cmd, &ext_args)
}

fn cli() -> App {
    let app = App::new("cargo")
        .settings(&[
            AppSettings::UnifiedHelpMessage,
            AppSettings::DeriveDisplayOrder,
            AppSettings::VersionlessSubcommands,
            AppSettings::AllowExternalSubcommands,
        ])
        .about("")
        .template(
            "\
Rust's package manager

USAGE:
    {usage}

OPTIONS:
{unified}

Some common cargo commands are (see all commands with --list):
    build       Compile the current project
    check       Analyze the current project and report errors, but don't build object files
    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
    uninstall   Uninstall a Rust binary

See 'cargo help <command>' for more information on a specific command.\n",
        )
        .arg(opt("version", "Print version info and exit").short("V"))
        .arg(opt("list", "List installed commands"))
        .arg(opt("explain", "Run `rustc --explain CODE`").value_name("CODE"))
        .arg(
            opt(
                "verbose",
                "Use verbose output (-vv very verbose/build.rs output)",
            ).short("v")
                .multiple(true)
                .global(true),
        )
        .arg(
            opt("quiet", "No output printed to stdout")
                .short("q")
                .global(true),
        )
        .arg(
            opt("color", "Coloring: auto, always, never")
                .value_name("WHEN")
                .global(true),
        )
        .arg(opt("frozen", "Require Cargo.lock and cache are up to date").global(true))
        .arg(opt("locked", "Require Cargo.lock is up to date").global(true))
        .arg(
            Arg::with_name("unstable-features")
                .help("Unstable (nightly-only) flags to Cargo, see 'cargo -Z help' for details")
                .short("Z")
                .value_name("FLAG")
                .multiple(true)
                .number_of_values(1)
                .global(true),
        )
        .subcommands(commands::builtin());
    app
}
<|endoftext|>


<commit_before>#![cfg(feature = "tools")]

extern crate octavo;
extern crate clap;
extern crate rustc_serialize;

use std::io::{self, Read};
use std::fs::File;

use clap::{App, Arg};
use rustc_serialize::hex::ToHex;

use octavo::digest::*;

const ENGINES: &'static [&'static str] = &["MD4",
                                           "MD5",
                                           "RIPEMD-160",
                                           "SHA1",
                                           "SHA-224",
                                           "SHA-256",
                                           "SHA-384",
                                           "SHA-512",
                                           "SHA3-224",
                                           "SHA3-256",
                                           "SHA3-384",
                                           "SHA3-512",
                                           "Tiger"];

fn checksum<Eng: Digest + Default>(input: &mut Read, buffer_size: usize) -> String {
    let mut eng = Eng::default();
    let mut buf = vec![0; buffer_size];

    while let Ok(len) = input.read(&mut buf[..]) {
        if len == 0 {
            break;
        }

        eng.update(&buf[..len]);
    }

    let mut result = vec![0; Eng::output_bytes()];
    eng.result(&mut result[..]);

    result[..].to_hex()
}

macro_rules! digest_for {
    (match $name:ident with $size:expr, $input:ident { $($val:pat => $eng:ty,)+ }) => {
        match $name {
            $($val => checksum::<$eng>(&mut $input, $size),)*
                _ => unreachable!()
        }
    }
}

fn input(name: &str) -> io::Result<Box<Read>> {
    if name == "-" {
        return Ok(Box::new(io::stdin()));
    }

    Ok(Box::new(try!(File::open(name))))
}

fn main() {
    let matches = App::new("sum")
                      .arg(Arg::with_name("digest")
                               .help("Select checksum engine. Defaults to SHA1.")
                               .possible_values(ENGINES)
                               .short("D")
                               .long("digest")
                               .global(true)
                               .empty_values(false)
                               .takes_value(true))
                      .arg(Arg::with_name("FILE")
                               .help("Files to compute checksums. `-` means STDIN. Defaults to \
                                      STDIN.")
                               .multiple(true))
                      .arg(Arg::with_name("buffer size")
                               .help("Size of read buffer in bytes. Defaults to 1MiB.")
                               .takes_value(true)
                               .long("buffer-size"))
                      .get_matches();

    let engine_name = matches.value_of("digest").unwrap_or("SHA1");
    let files = matches.values_of("FILE").unwrap_or(vec!["-"]);
    let buffer_size = matches.value_of("buffer size")
                             .and_then(|val| val.parse().ok())
                             .unwrap_or(1024 * 1024);

    for file in files {
        let mut input: Box<Read> = input(file).unwrap();

        let result = digest_for! {
            match engine_name with buffer_size, input {
                "MD4" => md4::Md4,
                "MD5" => md5::Md5,
                "RIPEMD-160" => ripemd::Ripemd160,
                "SHA1" => sha1::Sha1,
                "SHA-224" => sha2::Sha224,
                "SHA-256" => sha2::Sha256,
                "SHA-384" => sha2::Sha384,
                "SHA-512" => sha2::Sha512,
                "SHA3-224" => sha3::Sha3224,
                "SHA3-256" => sha3::Sha3256,
                "SHA3-384" => sha3::Sha3384,
                "SHA3-512" => sha3::Sha3512,
                "Tiger" => tiger::Tiger,
            }
        };

        println!("{} {}", result, file);
    }
}
<commit_msg>perf(bin/sum): Speed up bin/sum a little<commit_after>#![cfg(feature = "tools")]

extern crate octavo;
extern crate clap;
extern crate rustc_serialize;

use std::io::{self, Read, BufReader};
use std::fs::File;

use clap::{App, Arg};
use rustc_serialize::hex::ToHex;

use octavo::digest::*;

const ENGINES: &'static [&'static str] = &["MD4",
                                           "MD5",
                                           "RIPEMD-160",
                                           "SHA1",
                                           "SHA-224",
                                           "SHA-256",
                                           "SHA-384",
                                           "SHA-512",
                                           "SHA3-224",
                                           "SHA3-256",
                                           "SHA3-384",
                                           "SHA3-512",
                                           "Tiger"];

fn checksum<Eng: Digest + Default>(input: &mut Read) -> String {
    let mut eng = Eng::default();
    let mut buf = vec![0; Eng::block_size()];

    while let Ok(len) = input.read(&mut buf) {
        if len == 0 {
            break;
        }

        eng.update(&buf[..len]);
    }

    let mut result = vec![0; Eng::output_bytes()];
    eng.result(&mut result[..]);

    result.to_hex()
}

macro_rules! digest_for {
    (match $name:ident with $input:ident { $($val:pat => $eng:ty,)+ }) => {
        match $name {
            $($val => checksum::<$eng>(&mut $input),)*
                _ => unreachable!()
        }
    }
}

fn input(name: &str) -> io::Result<Box<Read>> {
    if name == "-" {
        return Ok(Box::new(io::stdin()));
    }

    Ok(Box::new(try!(File::open(name))))
}

fn main() {
    let matches = App::new("sum")
                      .arg(Arg::with_name("digest")
                               .help("Select checksum engine. Defaults to SHA1.")
                               .possible_values(ENGINES)
                               .short("D")
                               .long("digest")
                               .global(true)
                               .empty_values(false)
                               .takes_value(true))
                      .arg(Arg::with_name("FILE")
                               .help("Files to compute checksums. `-` means STDIN. Defaults to \
                                      STDIN.")
                               .multiple(true))
                      .get_matches();

    let engine_name = matches.value_of("digest").unwrap_or("SHA1");
    let files = matches.values_of("FILE").unwrap_or(vec!["-"]);

    for file in files {
        let mut input = BufReader::new(input(file).unwrap());

        let result = digest_for! {
            match engine_name with input {
                "MD4" => md4::Md4,
                "MD5" => md5::Md5,
                "RIPEMD-160" => ripemd::Ripemd160,
                "SHA1" => sha1::Sha1,
                "SHA-224" => sha2::Sha224,
                "SHA-256" => sha2::Sha256,
                "SHA-384" => sha2::Sha384,
                "SHA-512" => sha2::Sha512,
                "SHA3-224" => sha3::Sha3224,
                "SHA3-256" => sha3::Sha3256,
                "SHA3-384" => sha3::Sha3384,
                "SHA3-512" => sha3::Sha3512,
                "Tiger" => tiger::Tiger,
            }
        };

        println!("{} {}", result, file);
    }
}
<|endoftext|>