<commit_before>//! Provides an efficient Fast Fourier transform.
//!
//! A `Transformer` must first be created that stores precomputed information to
//! speed up the transform. This `Transformer` can then be used only to perform
//! FFTs of the chosen size.

use std::f32::consts::PI;
use std::vec::Vec;

use utils::complex::Complex;

/// A container for precomputed values to perform FFTs of a fixed size.
pub struct Transformer {
    size: usize,
    bit_reverses: Vec<usize>,
    twiddles: Vec<Complex>,
}

impl Transformer {
    /// Returns a set precomputed information used to perform FFTs of the
    /// provided size.
    pub fn new(size: usize) -> Transformer {
        // Only operate in powers of two
        let bufsize = size.next_power_of_two();

        // Populate the bit reverses
        // We only use the lower log2(size) bits to express the index
        let mut bit_reverses = Vec::with_capacity(bufsize);
        for i in (0 .. size) {
            let br = bit_reverse(i as u32, int_log(bufsize as u32));
            bit_reverses.push(br as usize);
        }

        // Populate the twiddle factors w_n^i
        // w_n = exp(-j*2*pi*n/N)
        let mut twiddles = Vec::with_capacity(bufsize);
        let exponent = Complex::new(0.0, -2.0*PI/(bufsize as f32));
        for i in (0 .. bufsize) {
            twiddles.push(exponent.scale(i as f32).exp());
        }

        Transformer { size: bufsize, bit_reverses: bit_reverses,
            twiddles: twiddles }
    }

    /// Returns the size FFTs this Transformer performs
    pub fn get_size(&self) -> usize {
        self.size
    }

    /// Performs an FFT on `input`, and places the result in `output`.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    ///
    /// Returns `Error` if `output` is too small to hold the result.
    pub fn fft(&self, input: &Vec<Complex>, output: &mut Vec<Complex>) ->
        Result<(),()> {
        self.transform(input, output, false)
    }

    /// Performs an inverse FFT on `input`, and places the result in `output`.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    ///
    /// Returns `Error` if `output` is too small to hold the result.
    pub fn ifft(&self, input: &Vec<Complex>, output: &mut Vec<Complex>) ->
        Result<(),()> {
        self.transform(input, output, true)
    }

    /// Performs the actual transform on `input`, placing the result in
    /// `output`.
    ///
    /// This function performs both forward and backwards transforms, since
    /// there are only minor algorithmic differences in the beginning and end
    /// of transformation.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    ///
    /// Returns `Error` if `output` is too small to hold the result.
    fn transform(&self, input: &Vec<Complex>, output: &mut Vec<Complex>,
                     inverse: bool) -> Result<(),()> {
        // Verify the provided vector is big enough for the result
        if output.len() < self.size {
            return Err(())
        }

        // Copy the input into bit reverse order, zero padding if necessary,
        // conjugating if we are inverse transforming
        for i in (0 .. input.len()) {
            output[self.bit_reverses[i]] =
                if inverse {
                    input[i].conj()
                } else {
                    input[i]
                }
        }
        for i in self.bit_reverses[input.len() .. self.size].iter() {
            output[*i] = Complex::from_real(0.0);
        }

        // Iteratively perform FFT, starting at 2 points
        let mut n = 2;
        while n <= self.size {
            // For each of the small FFTs
            for set in (0 .. self.size/n) {
                // For each pair of n
                for i in (0 .. n/2) {
                    let ilo = n*set + i;
                    let ihi = ilo + n/2;

                    // Grab out the lower and upper n
                    let lower = output[ilo];
                    let upper = output[ihi] * self.twiddles[self.size/n * i];

                    // Assign them back using a butterfly
                    output[ilo] = lower + upper;
                    output[ihi] = lower - upper;
                }
            }

            // Double the number of points per FFT
            n *= 2;
        }

        // If we are inverse transforming, conjugate and normalize the output
        if inverse {
            for i in (0 .. self.size) {
                output[i] = output[i].conj().scale(1.0/(self.size as f32));
            }
        }

        Ok(())
    }
}

/// Returns the bit reverse of `n`, for the lower `bits` bits.
///
/// For small examples, the bit reverse of 0b00011010 is 0b01011000, but the bit
/// reverse of just the lower 5 bits is 0b00001011.
fn bit_reverse(n: u32, bits: u32) -> u32 {
    let mut i = n;
    i = (i >> 16) | (i << 16);
    i = ((i & 0xFF00FF00) >> 8) | ((i & 0x00FF00FF) << 8);
    i = ((i & 0xF0F0F0F0) >> 4) | ((i & 0x0F0F0F0F) << 4);
    i = ((i & 0xCCCCCCCC) >> 2) | ((i & 0x33333333) << 2);
    i = ((i & 0xAAAAAAAA) >> 1) | ((i & 0x55555555) << 1);
    i >> ((32 - bits) as u32)
}

/// Returns the log base 2 of n, rounded up.
fn int_log(n: u32) -> u32 {
    let mut i = n-1; // correct for exact powers of 2
    let mut res = 0;
    while i > 0 {
        i = i >> 1;
        res += 1;
    }
    res
}

// Unit tests...
#[cfg(test)]
mod test {
    use std::vec::Vec;

    use utils::complex::Complex;
    use super::{Transformer, int_log, bit_reverse};

    /// Tests int_log with many values
    #[test]
    fn test_int_log() {
        assert_eq!(int_log(1), 0);
        assert_eq!(int_log(2), 1);
        assert_eq!(int_log(3), 2);
        assert_eq!(int_log(4), 2);
        assert_eq!(int_log(7), 3);
        assert_eq!(int_log(8), 3);
        assert_eq!(int_log(31), 5);
        assert_eq!(int_log(32), 5);
    }

    /// Tests bit_reverse.
    #[test]
    fn test_bit_reverse() {
        assert_eq!(bit_reverse(0x00000000, 32), 0x00000000);
        assert_eq!(bit_reverse(0xFFFFFFFF, 32), 0xFFFFFFFF);
        assert_eq!(bit_reverse(0x00000001, 32), 0x80000000);
        assert_eq!(bit_reverse(0x11111111, 32), 0x88888888);
        assert_eq!(bit_reverse(0x234f9e01, 32), 0x8079f2c4); //random
        assert_eq!(bit_reverse(0x00000001, 4), 0x00000008);
        assert_eq!(bit_reverse(0x0000000F, 4), 0x0000000F);
    }

    /// Tests the FFT of an impulse function.
    ///
    /// Analytically, an impulse function has a constant fourier transform.
    #[test]
    fn test_fft_impulse() {
        let zero = Complex::zero();
        let one = Complex::from_real(1.0);

        let mut impulse = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            impulse.push(if i == 0 { one } else { zero });
            out.push(Complex::zero());
        }

        let t = Transformer::new(8);
        t.fft(&impulse, &mut out).unwrap();

        for c in out.iter() {
            assert!(c.eq(one))
        }
    }

    /// Tests the IFFT of an impulse function.
    ///
    /// Analytically, a constant frequency domain results in an impulse
    /// function.
    #[test]
    fn test_ifft_impulse() {
        let zero = Complex::zero();
        let one = Complex::from_real(1.0);

        let mut impulse = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for _i in (0 .. 8) {
            impulse.push(one);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.ifft(&impulse, &mut out).unwrap();

        assert!(out[0].eq(one));
        for c in out[1 .. 8].iter() {
            assert!(c.eq(zero));
        }
    }

    /// Tests that the identify property, i.e. IFFT(FTT(f)) == f
    #[test]
    fn test_fft_identity() {
        let zero = Complex::zero();
        let epsilon = 1e-6;

        let mut input = Vec::with_capacity(8);
        let mut fft = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            input.push(Complex::from_real((i+1) as f32));
            fft.push(zero);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.fft(&input, &mut fft).unwrap();
        t.ifft(&fft, &mut out).unwrap();

        for i in (0 .. 7) {
            println!("{}",out[i].real() - ((i+1) as f32));
            assert!(out[i].real() - ((i+1) as f32) < epsilon);
            assert!(out[i].imag() < epsilon);
        }
    }

    /// Tests that the Transformer properly handles input buffers that are too
    /// short by zero padding them.
    #[test]
    fn test_fft_zero_pad() {
        let zero = Complex::zero();
        let epsilon = 1e-6;

        let mut input = Vec::with_capacity(7);
        let mut fft = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            if i < 7 {
                input.push(Complex::from_real((i+1) as f32));
            }
            fft.push(zero);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.fft(&input, &mut fft).unwrap();
        t.ifft(&fft, &mut out).unwrap();

        for i in (0 .. 7) {
            println!("{}",out[i].real() - ((i+1) as f32));
            assert!(out[i].real() - ((i+1) as f32) < epsilon);
            assert!(out[i].imag() < epsilon);
        }
    }

    /// Tests that the transformer fails when the output buffer is too short.
    #[test]
    fn test_fft_output_buffer_too_small() {
        let zero = Complex::zero();

        let mut input = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(7);
        for i in (0 .. 8) {
            input.push(zero);
            if i < 7 {
                out.push(zero);
            }
        }

        let t = Transformer::new(8);
        assert!(t.fft(&input, &mut out).is_err());
    }
}
<commit_msg>Strip Result from FFT.<commit_after>//! Provides an efficient Fast Fourier transform.
//!
//! A `Transformer` must first be created that stores precomputed information to
//! speed up the transform. This `Transformer` can then be used only to perform
//! FFTs of the chosen size.

use std::f32::consts::PI;
use std::vec::Vec;

use utils::complex::Complex;

/// A container for precomputed values to perform FFTs of a fixed size.
pub struct Transformer {
    size: usize,
    bit_reverses: Vec<usize>,
    twiddles: Vec<Complex>,
}

impl Transformer {
    /// Returns a set precomputed information used to perform FFTs of the
    /// provided size. The size is rounded up to the nearest power of two.
    pub fn new(size: usize) -> Transformer {
        // Only operate in powers of two
        let bufsize = size.next_power_of_two();

        // Populate the bit reverses
        // We only use the lower log2(size) bits to express the index
        let mut bit_reverses = Vec::with_capacity(bufsize);
        for i in (0 .. size) {
            let br = bit_reverse(i as u32, int_log(bufsize as u32));
            bit_reverses.push(br as usize);
        }

        // Populate the twiddle factors w_n^i
        // w_n = exp(-j*2*pi*n/N)
        let mut twiddles = Vec::with_capacity(bufsize);
        let exponent = Complex::new(0.0, -2.0*PI/(bufsize as f32));
        for i in (0 .. bufsize) {
            twiddles.push(exponent.scale(i as f32).exp());
        }

        Transformer { size: bufsize, bit_reverses: bit_reverses,
            twiddles: twiddles }
    }

    /// Returns the size FFTs this Transformer performs
    pub fn get_size(&self) -> usize {
        self.size
    }

    /// Performs an FFT on `input`, and places the result in `output`.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    pub fn fft(&self, input: &Vec<Complex>, output: &mut Vec<Complex>) {
        self.transform(input, output, false);
    }

    /// Performs an inverse FFT on `input`, and places the result in `output`.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    pub fn ifft(&self, input: &Vec<Complex>, output: &mut Vec<Complex>) {
        self.transform(input, output, true);
    }

    /// Performs the actual transform on `input`, placing the result in
    /// `output`.
    ///
    /// This function performs both forward and backwards transforms, since
    /// there are only minor algorithmic differences in the beginning and end
    /// of transformation.
    ///
    /// The input is zero padded if less than `size` samples are provided, and
    /// truncated if more than `size` samples are provided.
    fn transform(&self, input: &Vec<Complex>, output: &mut Vec<Complex>,
                 inverse: bool) {
        // Copy the input into bit reverse order, zero padding if necessary,
        // conjugating if we are inverse transforming
        for i in (0 .. input.len()) {
            output[self.bit_reverses[i]] =
                if inverse {
                    input[i].conj()
                } else {
                    input[i]
                }
        }
        for i in self.bit_reverses[input.len() .. self.size].iter() {
            output[*i] = Complex::from_real(0.0);
        }

        // Iteratively perform FFT, starting at 2 points
        let mut n = 2;
        while n <= self.size {
            // For each of the small FFTs
            for set in (0 .. self.size/n) {
                // For each pair of n
                for i in (0 .. n/2) {
                    let ilo = n*set + i;
                    let ihi = ilo + n/2;

                    // Grab out the lower and upper n
                    let lower = output[ilo];
                    let upper = output[ihi] * self.twiddles[self.size/n * i];

                    // Assign them back using a butterfly
                    output[ilo] = lower + upper;
                    output[ihi] = lower - upper;
                }
            }

            // Double the number of points per FFT
            n *= 2;
        }

        // If we are inverse transforming, conjugate and normalize the output
        if inverse {
            for i in (0 .. self.size) {
                output[i] = output[i].conj().scale(1.0/(self.size as f32));
            }
        }
    }
}

/// Returns the bit reverse of `n`, for the lower `bits` bits.
///
/// For small examples, the bit reverse of 0b00011010 is 0b01011000, but the bit
/// reverse of just the lower 5 bits is 0b00001011.
fn bit_reverse(n: u32, bits: u32) -> u32 {
    let mut i = n;
    i = (i >> 16) | (i << 16);
    i = ((i & 0xFF00FF00) >> 8) | ((i & 0x00FF00FF) << 8);
    i = ((i & 0xF0F0F0F0) >> 4) | ((i & 0x0F0F0F0F) << 4);
    i = ((i & 0xCCCCCCCC) >> 2) | ((i & 0x33333333) << 2);
    i = ((i & 0xAAAAAAAA) >> 1) | ((i & 0x55555555) << 1);
    i >> ((32 - bits) as u32)
}

/// Returns the log base 2 of n, rounded up.
fn int_log(n: u32) -> u32 {
    let mut i = n-1; // correct for exact powers of 2
    let mut res = 0;
    while i > 0 {
        i = i >> 1;
        res += 1;
    }
    res
}

// Unit tests...
#[cfg(test)]
mod test {
    use std::vec::Vec;

    use utils::complex::Complex;
    use super::{Transformer, int_log, bit_reverse};

    /// Tests int_log with many values
    #[test]
    fn test_int_log() {
        assert_eq!(int_log(1), 0);
        assert_eq!(int_log(2), 1);
        assert_eq!(int_log(3), 2);
        assert_eq!(int_log(4), 2);
        assert_eq!(int_log(7), 3);
        assert_eq!(int_log(8), 3);
        assert_eq!(int_log(31), 5);
        assert_eq!(int_log(32), 5);
    }

    /// Tests bit_reverse.
    #[test]
    fn test_bit_reverse() {
        assert_eq!(bit_reverse(0x00000000, 32), 0x00000000);
        assert_eq!(bit_reverse(0xFFFFFFFF, 32), 0xFFFFFFFF);
        assert_eq!(bit_reverse(0x00000001, 32), 0x80000000);
        assert_eq!(bit_reverse(0x11111111, 32), 0x88888888);
        assert_eq!(bit_reverse(0x234f9e01, 32), 0x8079f2c4); //random
        assert_eq!(bit_reverse(0x00000001, 4), 0x00000008);
        assert_eq!(bit_reverse(0x0000000F, 4), 0x0000000F);
    }

    /// Tests the FFT of an impulse function.
    ///
    /// Analytically, an impulse function has a constant fourier transform.
    #[test]
    fn test_fft_impulse() {
        let zero = Complex::zero();
        let one = Complex::from_real(1.0);

        let mut impulse = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            impulse.push(if i == 0 { one } else { zero });
            out.push(Complex::zero());
        }

        let t = Transformer::new(8);
        t.fft(&impulse, &mut out);

        for c in out.iter() {
            assert!(c.eq(one))
        }
    }

    /// Tests the IFFT of an impulse function.
    ///
    /// Analytically, a constant frequency domain results in an impulse
    /// function.
    #[test]
    fn test_ifft_impulse() {
        let zero = Complex::zero();
        let one = Complex::from_real(1.0);

        let mut impulse = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for _i in (0 .. 8) {
            impulse.push(one);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.ifft(&impulse, &mut out);

        assert!(out[0].eq(one));
        for c in out[1 .. 8].iter() {
            assert!(c.eq(zero));
        }
    }

    /// Tests that the identify property, i.e. IFFT(FTT(f)) == f
    #[test]
    fn test_fft_identity() {
        let zero = Complex::zero();
        let epsilon = 1e-6;

        let mut input = Vec::with_capacity(8);
        let mut fft = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            input.push(Complex::from_real((i+1) as f32));
            fft.push(zero);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.fft(&input, &mut fft);
        t.ifft(&fft, &mut out);

        for i in (0 .. 7) {
            println!("{}",out[i].real() - ((i+1) as f32));
            assert!(out[i].real() - ((i+1) as f32) < epsilon);
            assert!(out[i].imag() < epsilon);
        }
    }

    /// Tests that the Transformer properly handles input buffers that are too
    /// short by zero padding them.
    #[test]
    fn test_fft_zero_pad() {
        let zero = Complex::zero();
        let epsilon = 1e-6;

        let mut input = Vec::with_capacity(7);
        let mut fft = Vec::with_capacity(8);
        let mut out = Vec::with_capacity(8);
        for i in (0 .. 8) {
            if i < 7 {
                input.push(Complex::from_real((i+1) as f32));
            }
            fft.push(zero);
            out.push(zero);
        }

        let t = Transformer::new(8);
        t.fft(&input, &mut fft);
        t.ifft(&fft, &mut out);

        for i in (0 .. 7) {
            println!("{}",out[i].real() - ((i+1) as f32));
            assert!(out[i].real() - ((i+1) as f32) < epsilon);
            assert!(out[i].imag() < epsilon);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{account_config::CORE_CODE_ADDRESS, event::EventKey};
use anyhow::Result;
use move_core_types::move_resource::MoveResource;
use serde::{Deserialize, Serialize};

/// Struct that represents a UpgradeEvent.
#[derive(Debug, Serialize, Deserialize)]
pub struct UpgradeEvent {
    write_set: Vec<u8>,
}

impl UpgradeEvent {
    /// Get the applied writeset.
    pub fn write_set(&self) -> &[u8] {
        &self.write_set
    }

    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self> {
        lcs::from_bytes(bytes).map_err(Into::into)
    }

    pub fn event_key() -> EventKey {
        EventKey::new_from_address(&CORE_CODE_ADDRESS, 15)
    }
}

impl MoveResource for UpgradeEvent {
    const MODULE_NAME: &'static str = "LibraWriteSetManager";
    const STRUCT_NAME: &'static str = "UpgradeEvent";
}
<commit_msg>Removing event_key from types/src/account_config/events/upgrade.rs<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use move_core_types::move_resource::MoveResource;
use serde::{Deserialize, Serialize};

/// Struct that represents a UpgradeEvent.
#[derive(Debug, Serialize, Deserialize)]
pub struct UpgradeEvent {
    write_set: Vec<u8>,
}

impl UpgradeEvent {
    /// Get the applied writeset.
    pub fn write_set(&self) -> &[u8] {
        &self.write_set
    }

    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self> {
        lcs::from_bytes(bytes).map_err(Into::into)
    }
}

impl MoveResource for UpgradeEvent {
    const MODULE_NAME: &'static str = "LibraWriteSetManager";
    const STRUCT_NAME: &'static str = "UpgradeEvent";
}
<|endoftext|>


<commit_before>#[derive(PartialEq)]
enum OneOrHalf {
    O,
    H,
}

fn sum_matches(ns: &[u32], oh: &OneOrHalf) -> u32 {
    let len: usize = ns.len();
    let step: usize = if oh == &OneOrHalf::O { 1 } else { len / 2 };
    (0..len)
        .filter(|&i| ns[i] == ns[(i + step) % len])
        .map(|i| ns[i])
        .sum()
}

fn parse_input(s: &str) -> Vec<u32> {
    s.chars().filter_map(|c| c.to_digit(10)).collect()
}

fn part1(input: &str) -> u32 {
    sum_matches(&parse_input(input), &OneOrHalf::O)
}

fn part2(input: &str) -> u32 {
    sum_matches(&parse_input(input), &OneOrHalf::H)
}

fn main() {
    let input: &str =
        "68376334795224855827459835293967497295464175589881588256882344699473595413912688278647235862566123233983921662578792917453912795352746426512649965615919588512125567186837411371179875287621488759761429629174886972298349197722423458299323141529413191327622485249495864168181327197661454464926326248274999448373741839963155646828842752761293142356422964355349521987483211496361289666375779728345952231649453711684539164893151811849653331845998998597991146881361717234517911759893792348815818755262456378627116779495435596139617246571678531183335956244163871445674244765586446362529159854137535962117184875192273872222899887357292312978286182636232921252574738118347521187637829623831872437381979223955675634257889137823684924127338433248519515211796732599314921611399736571277222546332369461136277417419794865524123989722492356536832313937597437717873787593849468836733642529378547151146397532997237439387663769334722979172954835154486382983716698212694357398153392926255272961384626131829678171219569288685597141132355322788254163923888378155573948753185423158997877718687642446457446643422536541238979761725496426292359382168535641216124211741896552562128941824172241913873537828976172738276983915232241451589421911121567228899853934667954786256223614621554618294467191255153395256524786159758429643756586457639177183891162214163549688595416893383194995824534247841414247526268212761954913719452114876764745799982792594753759626334319631191917894368116738893548797661111899664138398354818931135486984944719992393148681724116616741428937687985152658296679845474766477741553632712968679175356452987459761126437216758171182395219393289199148996813762849991484678429793578629331215796996751484375784895561682156658579887518746862371751372692472765217374791324656745291574784495299477362964676351148183676897122366838656342745944945275263617729359831466565694983217252594237828187612857523344265418227883219383138893873384775659548637662867572687198263688597865118173921615178165442133987362382721444844952715592955744739873677838847693982379696776";
    println!("{}", part1(input));
    println!("{}", part2(input));
}

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

    #[test]
    fn test_sum_matches_one() {
        assert_eq!(sum_matches(&[1, 1, 2, 2], &OneOrHalf::O), 3);
        assert_eq!(sum_matches(&[1, 1, 1, 1], &OneOrHalf::O), 4);
        assert_eq!(sum_matches(&[1, 2, 3, 4], &OneOrHalf::O), 0);
        assert_eq!(sum_matches(&[9, 1, 2, 1, 2, 1, 2, 9], &OneOrHalf::O), 9);
    }

    #[test]
    fn test_parse_input() {
        assert_eq!(parse_input("1122"), &[1, 1, 2, 2]);
        assert_eq!(parse_input("1111"), &[1, 1, 1, 1]);
        assert_eq!(parse_input("1234"), &[1, 2, 3, 4]);
        assert_eq!(parse_input("91212129"), &[9, 1, 2, 1, 2, 1, 2, 9]);
    }

    #[test]
    fn test_part1() {
        assert_eq!(part1("1122"), 3);
        assert_eq!(part1("1111"), 4);
        assert_eq!(part1("1234"), 0);
        assert_eq!(part1("91212129"), 9);
    }

    #[test]
    fn test_part2() {
        assert_eq!(part2("1212"), 6);
        assert_eq!(part2("1221"), 0);
        assert_eq!(part2("123425"), 4);
        assert_eq!(part2("123123"), 12);
        assert_eq!(part2("12131415"), 4);
    }

}
<commit_msg>cleaner day1<commit_after>#[derive(PartialEq)]
enum Step {
    One,
    HalfOfLen,
}

fn sum_matches(ns: &[u32], step: &Step) -> u32 {
    let len: usize = ns.len();
    let j: usize = match *step {
        Step::One => 1,
        Step::HalfOfLen => len / 2,
    };
    (0..len)
        .filter(|&i| ns[i] == ns[(i + j) % len])
        .map(|i| ns[i])
        .sum()
}

fn parse_input(s: &str) -> Vec<u32> {
    s.chars().filter_map(|c| c.to_digit(10)).collect()
}

fn part1(input: &str) -> u32 {
    sum_matches(&parse_input(input), &Step::One)
}

fn part2(input: &str) -> u32 {
    sum_matches(&parse_input(input), &Step::HalfOfLen)
}

fn main() {
    let input: &str =
        "68376334795224855827459835293967497295464175589881588256882344699473595413912688278647235862566123233983921662578792917453912795352746426512649965615919588512125567186837411371179875287621488759761429629174886972298349197722423458299323141529413191327622485249495864168181327197661454464926326248274999448373741839963155646828842752761293142356422964355349521987483211496361289666375779728345952231649453711684539164893151811849653331845998998597991146881361717234517911759893792348815818755262456378627116779495435596139617246571678531183335956244163871445674244765586446362529159854137535962117184875192273872222899887357292312978286182636232921252574738118347521187637829623831872437381979223955675634257889137823684924127338433248519515211796732599314921611399736571277222546332369461136277417419794865524123989722492356536832313937597437717873787593849468836733642529378547151146397532997237439387663769334722979172954835154486382983716698212694357398153392926255272961384626131829678171219569288685597141132355322788254163923888378155573948753185423158997877718687642446457446643422536541238979761725496426292359382168535641216124211741896552562128941824172241913873537828976172738276983915232241451589421911121567228899853934667954786256223614621554618294467191255153395256524786159758429643756586457639177183891162214163549688595416893383194995824534247841414247526268212761954913719452114876764745799982792594753759626334319631191917894368116738893548797661111899664138398354818931135486984944719992393148681724116616741428937687985152658296679845474766477741553632712968679175356452987459761126437216758171182395219393289199148996813762849991484678429793578629331215796996751484375784895561682156658579887518746862371751372692472765217374791324656745291574784495299477362964676351148183676897122366838656342745944945275263617729359831466565694983217252594237828187612857523344265418227883219383138893873384775659548637662867572687198263688597865118173921615178165442133987362382721444844952715592955744739873677838847693982379696776";
    println!("{}", part1(input));
    println!("{}", part2(input));
}

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

    #[test]
    fn test_sum_matches_one() {
        assert_eq!(sum_matches(&[1, 1, 2, 2], &Step::One), 3);
        assert_eq!(sum_matches(&[1, 1, 1, 1], &Step::One), 4);
        assert_eq!(sum_matches(&[1, 2, 3, 4], &Step::One), 0);
        assert_eq!(sum_matches(&[9, 1, 2, 1, 2, 1, 2, 9], &Step::One), 9);
    }

    #[test]
    fn test_parse_input() {
        assert_eq!(parse_input("1122"), &[1, 1, 2, 2]);
        assert_eq!(parse_input("1111"), &[1, 1, 1, 1]);
        assert_eq!(parse_input("1234"), &[1, 2, 3, 4]);
        assert_eq!(parse_input("91212129"), &[9, 1, 2, 1, 2, 1, 2, 9]);
    }

    #[test]
    fn test_part1() {
        assert_eq!(part1("1122"), 3);
        assert_eq!(part1("1111"), 4);
        assert_eq!(part1("1234"), 0);
        assert_eq!(part1("91212129"), 9);
    }

    #[test]
    fn test_part2() {
        assert_eq!(part2("1212"), 6);
        assert_eq!(part2("1221"), 0);
        assert_eq!(part2("123425"), 4);
        assert_eq!(part2("123123"), 12);
        assert_eq!(part2("12131415"), 4);
    }

}
<|endoftext|>


<commit_before>macro_rules! vk_functions {
    ($($raw_name: expr, $name: ident ($($param_name: ident: $param: ty),*,) -> $ret: ty;)+) => {
        $(
            pub unsafe extern "system" fn $name (
                $($param_name: $param),*
                ) -> $ret {
                use std::mem;

                mem::transmute::<_, $name::FnType>($name::fn_ptr)($($param_name),*)
            }

            pub mod $name {

                use super::super::*;
                #[allow(unused_imports)]
                use super::super::libc_reexports::*;
                pub const RAW_NAME: &'static str = $raw_name;
                pub static mut fn_ptr: *const () = unloaded_function_panic as *const ();
                #[doc(hidden)]
                pub type FnType = unsafe extern "system" fn($($param),*) -> $ret;

                pub fn is_loaded() -> bool {
                    unsafe{ fn_ptr == unloaded_function_panic as *const () }
                }
            }
        )+

        pub fn load_with<F: FnMut(&str) -> *const ()>(mut load_fn: F) -> ::std::result::Result<(), Vec<&'static str>> {unsafe{
            use std::ptr;
            let mut fn_buf: *const ();
            let mut unloaded_fns = Vec::new();

            $(
                fn_buf = load_fn($raw_name);
                if ptr::null() != fn_buf {
                    $name::fn_ptr = fn_buf;
                } else if $name::fn_ptr != unloaded_function_panic as *const () {
                    unloaded_fns.push($raw_name)
                }
            )+

            if 0 == unloaded_fns.len() {
                Ok(())
            } else {
                Err(unloaded_fns)
            }
        }}
    }
}
<commit_msg>Fixed lowercase global in global generator<commit_after>macro_rules! vk_functions {
    ($($raw_name: expr, $name: ident ($($param_name: ident: $param: ty),*,) -> $ret: ty;)+) => {
        $(
            pub unsafe extern "system" fn $name (
                $($param_name: $param),*
                ) -> $ret {
                use std::mem;

                mem::transmute::<_, $name::FnType>($name::FN_PTR)($($param_name),*)
            }

            pub mod $name {

                use super::super::*;
                #[allow(unused_imports)]
                use super::super::libc_reexports::*;
                pub const RAW_NAME: &'static str = $raw_name;
                pub static mut FN_PTR: *const () = unloaded_function_panic as *const ();
                #[doc(hidden)]
                pub type FnType = unsafe extern "system" fn($($param),*) -> $ret;

                pub fn is_loaded() -> bool {
                    unsafe{ FN_PTR == unloaded_function_panic as *const () }
                }
            }
        )+

        pub fn load_with<F: FnMut(&str) -> *const ()>(mut load_fn: F) -> ::std::result::Result<(), Vec<&'static str>> {unsafe{
            use std::ptr;
            let mut fn_buf: *const ();
            let mut unloaded_fns = Vec::new();

            $(
                fn_buf = load_fn($raw_name);
                if ptr::null() != fn_buf {
                    $name::FN_PTR = fn_buf;
                } else if $name::FN_PTR != unloaded_function_panic as *const () {
                    unloaded_fns.push($raw_name)
                }
            )+

            if 0 == unloaded_fns.len() {
                Ok(())
            } else {
                Err(unloaded_fns)
            }
        }}
    }
}
<|endoftext|>


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

use std::ascii::AsciiExt;
use std::env;
use std::fmt;
use std::io::{self, stdout, Write, BufWriter};
use std::process;
use std::thread;
use std::time as stdtime;

use kafka::client::{KafkaClient, FetchOffset, GroupOffsetStorage};

/// A very simple offset monitor for a particular topic able to show
/// the lag for a particular consumer group. Dumps the offset/lag of
/// the monitored topic/group to stdout every few seconds.
fn main() {
    env_logger::init().unwrap();

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

    if let Err(e) = run(cfg) {
        println!("error: {}", e);
        process::exit(1);
    }
}

#[derive(Copy, Clone)]
struct Offsets {
    latest: i64,
    group: i64,
}

impl Default for Offsets {
    fn default() -> Self {
        Offsets { latest: -1, group: -1 }
    }
}

fn run(cfg: Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers);
    client.set_group_offset_storage(cfg.offset_storage);
    try!(client.load_metadata_all());

    let mut stdout = stdout();

    if cfg.topic.is_empty() {
        let ts = client.topics();
        let num_topics = ts.len();
        if num_topics == 0 {
            return Err(Error::Literal("no topics available".to_owned()));
        }
        let mut names: Vec<&str> = Vec::with_capacity(ts.len());
        names.extend(ts.names());
        names.sort();

        let mut buf = BufWriter::with_capacity(4096, stdout);
        for name in names {
            try!(write!(buf, "topic: {}\n", name));
        }
        return Err(Error::Literal("choose a topic".to_owned()));
    }

    let num_partitions = match client.topics().partitions(&cfg.topic) {
        None => return Err(Error::Literal(format!("no such topic: {}", &cfg.topic))),
        Some(partitions) => partitions.len(),
    };
    let mut offs = vec![Offsets::default(); num_partitions];
    let mut offs = offs.as_mut_slice();

    let mut fmt_buf = String::with_capacity(16);
    let mut out_buf = String::with_capacity(128);

    // ~ shall we produce the lag of the specified group?
    let show_lag = !cfg.group.is_empty();

    // ~ print a header
    {
        use std::fmt::Write;

        let _ = write!(out_buf, "{:>8} ", "time");
        for i in 0..num_partitions {
            fmt_buf.clear();
            let _ = write!(fmt_buf, "p-{}", i);
            let _ = write!(out_buf, " {:>10}", fmt_buf);
            if show_lag {
                let _ = write!(out_buf, " {:6}", "(lag)");
            }
        }
        if cfg.summary {
            let _ = write!(out_buf, " {:>12}", "summary");
            if show_lag {
                let _ = write!(out_buf, " {:8}", "(lag)");
            }
        }
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());
    }
    // ~ now periodically fetch the offsets and dump them to stdout
    loop {
        let t = time::now();
        // ~ latest topic offsets
        let latests = try!(client.fetch_topic_offsets(&cfg.topic, FetchOffset::Latest));
        for l in latests {
            let v = l.offset.unwrap_or(-1);
            offs.get_mut(l.partition as usize).expect("available partition").latest = v;
        }
        if show_lag {
            // ~ group offsets
            let groups = try!(client.fetch_group_topic_offsets(&cfg.group, &cfg.topic));
            for g in groups {
                let v = g.offset.unwrap_or(-1);
                offs.get_mut(g.partition as usize).expect("available partition").group = v;
            }
        }
        // ~ format the output
        {
            use std::fmt::Write;

            let mut total_latest = 0;
            let mut total_lag = 0;

            out_buf.clear();
            let _ = write!(out_buf, "{} ", t.strftime("%H:%M:%S").unwrap());
            for o in &*offs {
                total_latest += o.latest;
                let _ = write!(out_buf, " {:>10}", o.latest);
                if show_lag {
                    fmt_buf.clear();
                    if o.group < 0 {
                        let _ = write!(fmt_buf, "()");
                        total_lag = -1;
                    } else {
                        let mut lag = if cfg.commited_not_consumed {
                            o.latest - o.group
                        } else {
                            o.latest - o.group - 1
                        };
                        if lag < 0 {
                            // ~ it's quite likely that we fetched group offsets
                            // which are a bit ahead of the topic's latest offset
                            // since we issued the fetch-latest-offset request
                            // earlier than the request for the group-offsets
                            lag = 0;
                        }
                        if total_lag >= 0 {
                            total_lag += lag;
                        }
                        let _ = write!(fmt_buf, "({:<})", lag);
                    }
                    let _ = write!(out_buf, " {:6}", fmt_buf);
                }
            }
            if cfg.summary {
                let _ = write!(out_buf, " {:12}", total_latest);
                if show_lag {
                    fmt_buf.clear();
                    if total_lag >= 0 {
                        let _ = write!(fmt_buf, "({:<})", total_lag);
                    } else {
                        let _ = write!(fmt_buf, "()");
                    }
                    let _ = write!(out_buf, " {:8}", fmt_buf);
                }
            }
        }
        // ~ write the output
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());

        thread::sleep(cfg.period);
    }
}

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

struct Config {
    brokers: Vec<String>,
    topic: String,
    group: String,
    offset_storage: GroupOffsetStorage,
    period: stdtime::Duration,
    commited_not_consumed: bool,
    summary: bool,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka bootstrap brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify the topic to monitor", "TOPIC");
        opts.optopt("", "group", "Specify the group to monitor", "GROUP");
        opts.optopt("", "offsets", "Specify offset store [zookeeper, kafka]", "STORE");
        opts.optopt("", "sleep", "Specify the sleep time", "SECS");
        opts.optflag("", "summary", "Print a summary for the topic");
        opts.optflag("", "committed-not-yet-consumed",
                     "Assume commited group offsets specify \
                      messages the group will start consuming \
                      (including those at these offsets)");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("offsets") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        let mut period = stdtime::Duration::from_secs(5);
        if let Some(s) = m.opt_str("sleep") {
            match s.parse::<u64>() {
                Ok(n) if n != 0 => period = stdtime::Duration::from_secs(n),
                _ => return Err(Error::Literal(format!("not a number greater than zero: {}", s))),
            }
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| String::new()),
            group: m.opt_str("group")
                .unwrap_or_else(|| String::new()),
            offset_storage: offset_storage,
            period: period,
            commited_not_consumed: m.opt_present("committed-not-yet-consumed"),
            summary: m.opt_present("summary"),
        })
    }
}

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

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

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

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

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}
<commit_msg>Rework offset-monitor to be more useful<commit_after>extern crate kafka;
extern crate getopts;
extern crate env_logger;
extern crate time;

use std::ascii::AsciiExt;
use std::cmp;
use std::env;
use std::fmt;
use std::io::{self, stdout, stderr, BufWriter, Write};
use std::process;
use std::thread;
use std::time as stdtime;

use kafka::client::{KafkaClient, FetchOffset, GroupOffsetStorage};

/// A very simple offset monitor for a particular topic able to show
/// the lag for a particular consumer group. Dumps the offset/lag of
/// the monitored topic/group to stdout every few seconds.
fn main() {
    env_logger::init().unwrap();

    macro_rules! abort {
        ($e:expr) => {{
            let mut out = stderr();
            let _ = write!(out, "error: {}\n", $e);
            let _ = out.flush();
            process::exit(1);
        }}
    };

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => abort!(e),
    };

    if let Err(e) = run(cfg) {
        abort!(e);
    }
}

fn run(cfg: Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers.clone());
    client.set_group_offset_storage(cfg.offset_storage);
    try!(client.load_metadata_all());

    // ~ if no topic specified, print all available and be done.
    if cfg.topic.is_empty() {
        let ts = client.topics();
        let num_topics = ts.len();
        if num_topics == 0 {
            return Err(Error::Literal("no topics available".to_owned()));
        }
        let mut names: Vec<&str> = Vec::with_capacity(ts.len());
        names.extend(ts.names());
        names.sort();

        let mut buf = BufWriter::with_capacity(1024, stdout());
        for name in names {
            let _ = write!(buf, "topic: {}\n", name);
        }
        return Err(Error::Literal("choose a topic".to_owned()));
    }

    // ~ otherwise let's loop over the topic partition offsets
    let num_partitions = match client.topics().partitions(&cfg.topic) {
        None => return Err(Error::Literal(format!("no such topic: {}", &cfg.topic))),
        Some(partitions) => partitions.len(),
    };
    let mut state = State::new(num_partitions, cfg.commited_not_consumed);
    let mut printer = Printer::new(stdout(), &cfg);
    try!(printer.print_head(num_partitions));

    // ~ initialize the state
    let mut first_time = true;
    loop {
        let t = time::now();
        try!(state.update_partitions(&mut client, &cfg.topic, &cfg.group));
        if first_time {
            state.curr_to_prev();
            first_time = false;
        }
        try!(printer.print_offsets(&t, &state.offsets));
        thread::sleep(cfg.period);
    }
}

#[derive(Copy, Clone)]
struct Partition {
    prev_latest: i64,
    curr_latest: i64,
    curr_lag: i64,
}

impl Default for Partition {
    fn default() -> Self {
        Partition { prev_latest: -1, curr_latest: -1, curr_lag: -1 }
    }
}

struct State {
    offsets: Vec<Partition>,
    lag_decr: i64,
}

impl State {
    fn new(num_partitions: usize, commited_not_consumed: bool) -> State {
        State {
            offsets: vec![Default::default(); num_partitions],
            lag_decr: if commited_not_consumed { 0 } else { 1 },
        }
    }

    fn update_partitions(&mut self, client: &mut KafkaClient, topic: &str, group: &str)
                         -> Result<(), Error>
    {
        // ~ get the latest topic offsets
        let latests = try!(client.fetch_topic_offsets(topic, FetchOffset::Latest));
        for l in latests {
            let off = self.offsets.get_mut(l.partition as usize)
                .expect("[topic offset] non-existent partition");
            off.prev_latest = off.curr_latest;
            off.curr_latest = l.offset.unwrap_or(-1);
        }
        if !group.is_empty() {
            // ~ get the current group offsets
            let groups = try!(client.fetch_group_topic_offsets(group, topic));
            for g in groups {
                let off = self.offsets.get_mut(g.partition as usize)
                    .expect("[group offset] non-existent partition");
                match g.offset {
                    Ok(v) => {
                        // ~ it's quite likely that we fetched group offsets
                        // which are a bit ahead of the topic's latest offset
                        // since we issued the fetch-latest-offset request
                        // earlier than the request for the group offsets
                        off.curr_lag = cmp::max(0, off.curr_latest - v - self.lag_decr);
                    }
                    Err(_) => {
                        off.curr_lag = -1;
                    }
                }
            }
        }
        Ok(())
    }

    fn curr_to_prev(&mut self) {
        for o in &mut self.offsets {
            o.prev_latest = o.curr_latest;
        }
    }
}

struct Printer<W> {
    out: W,

    timefmt: String,
    fmt_buf: String,
    out_buf: String,

    time_width: usize,
    offset_width: usize,
    diff_width: usize,
    lag_width: usize,

    print_diff: bool,
    print_lag: bool,
    print_summary: bool,
}

impl<W: Write> Printer<W> {
    fn new(out: W, cfg: &Config) -> Printer<W> {
        Printer {
            out: out,
            timefmt: "%H:%M:%S".into(),
            fmt_buf: String::with_capacity(30),
            out_buf: String::with_capacity(160),
            time_width: 10,
            offset_width: 11,
            diff_width: 8,
            lag_width: 6,
            print_diff: cfg.diff,
            print_lag: !cfg.group.is_empty(),
            print_summary: cfg.summary,
        }
    }

    fn print_head(&mut self, num_partitions: usize) -> Result<(), io::Error> {
        self.out_buf.clear();
        {   // ~ format
            use std::fmt::Write;
            let _ = write!(self.out_buf, "{1:<0$}", self.time_width, "time");
            if self.print_summary {
                let _ = write!(self.out_buf, " {1:>0$}", self.offset_width, "topic");
                if self.print_diff {
                    let _ = write!(self.out_buf, " [{1:>0$}]", self.diff_width, "growth");
                }
                if self.print_lag {
                    let _ = write!(self.out_buf, " {1:0$}", self.lag_width, "(lag)");
                }
            } else {
                for i in 0..num_partitions {
                    self.fmt_buf.clear();
                    let _ = write!(self.fmt_buf, "p-{}", i);
                    let _ = write!(self.out_buf, " {1:>0$}", self.offset_width, self.fmt_buf);
                    if self.print_diff {
                        let _ = write!(self.out_buf, " [{1:>0$}]", self.diff_width, "growth");
                    }
                    if self.print_lag {
                        let _ = write!(self.out_buf, " {1:0$}", self.lag_width, "(lag)");
                    }
                }
            }
            self.out_buf.push('\n');
        }
        {   // ~ print
            self.out.write_all(self.out_buf.as_bytes())
        }
    }

    fn print_offsets(&mut self, time: &time::Tm, partitions: &[Partition]) -> Result<(), io::Error> {
        self.out_buf.clear();
        {   // ~ format
            use std::fmt::Write;

            self.fmt_buf.clear();
            let _ = write!(self.fmt_buf, "{}", time.strftime(&self.timefmt).expect("invalid timefmt"));
            let _ = write!(self.out_buf, "{1:<0$}", self.time_width, self.fmt_buf);
            if self.print_summary {
                let mut prev_latest = 0;
                let mut curr_latest = 0;
                let mut curr_lag = 0;
                for p in partitions {
                    macro_rules! cond_add {
                        ($v:ident) => {
                            if $v != -1 {
                                if p.$v < 0 { $v = -1; } else { $v += p.$v; }
                            }
                        }
                    };
                    cond_add!(prev_latest);
                    cond_add!(curr_latest);
                    cond_add!(curr_lag);
                }
                let _ = write!(self.out_buf, " {1:>0$}", self.offset_width, curr_latest);
                if self.print_diff {
                    self.fmt_buf.clear();
                    let _ = write!(self.fmt_buf, "{:+}", curr_latest - prev_latest);
                    let _ = write!(self.out_buf, " [{1:>0$}]", self.diff_width, self.fmt_buf);
                }
                if self.print_lag {
                    self.fmt_buf.clear();
                    let _ = write!(self.fmt_buf, "({})", curr_lag);
                    let _ = write!(self.out_buf, " {1:<0$}", self.lag_width, self.fmt_buf);
                }
            } else {
                for p in partitions {
                    let _ = write!(self.out_buf, " {1:>0$}", self.offset_width, p.curr_latest);
                    if self.print_diff {
                        self.fmt_buf.clear();
                        let _ = write!(self.fmt_buf, "{:+}", p.curr_latest - p.prev_latest);
                        let _ = write!(self.out_buf, " [{1:>0$}]", self.diff_width, self.fmt_buf);
                    }
                    if self.print_lag {
                        self.fmt_buf.clear();
                        let _ = write!(self.fmt_buf, "({})", p.curr_lag);
                        let _ = write!(self.out_buf, " {1:<0$}", self.lag_width, self.fmt_buf);
                    }
                }
            }
        }
        self.out_buf.push('\n');
        self.out.write_all(self.out_buf.as_bytes())
    }
}

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

struct Config {
    brokers: Vec<String>,
    topic: String,
    group: String,
    offset_storage: GroupOffsetStorage,
    period: stdtime::Duration,
    commited_not_consumed: bool,
    summary: bool,
    diff: bool,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka bootstrap brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify the topic to monitor", "TOPIC");
        opts.optopt("", "group", "Specify the group to monitor", "GROUP");
        opts.optopt("", "storage", "Specify offset store [zookeeper, kafka]", "STORE");
        opts.optopt("", "sleep", "Specify the sleep time", "SECS");
        opts.optflag("", "partitions", "Print each partition instead of the summary");
        opts.optflag("", "no-growth", "Don't print offset growth");
        opts.optflag("", "committed-not-yet-consumed",
                     "Assume commited group offsets specify \
                      messages the group will start consuming \
                      (including those at these offsets)");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("storage") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        let mut period = stdtime::Duration::from_secs(5);
        if let Some(s) = m.opt_str("sleep") {
            match s.parse::<u64>() {
                Ok(n) if n != 0 => period = stdtime::Duration::from_secs(n),
                _ => return Err(Error::Literal(format!("not a number greater than zero: {}", s))),
            }
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| String::new()),
            group: m.opt_str("group")
                .unwrap_or_else(|| String::new()),
            offset_storage: offset_storage,
            period: period,
            commited_not_consumed: m.opt_present("committed-not-yet-consumed"),
            summary: !m.opt_present("partitions"),
            diff: !m.opt_present("no-growth"),
        })
    }
}

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

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

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

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

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


<commit_before>pub mod doc_id_set;
pub mod boolean_retrieval;
pub mod scoring;
pub mod plan;

use std::fmt;
use std::io::{Cursor, Read, Write};
use std::collections::HashMap;
use std::rc::Rc;

use kite::schema::{FieldRef, SchemaRead};
use kite::query::Query;
use kite::query::term_scorer::TermScorer;
use kite::collectors::{Collector, DocumentMatch};
use rocksdb::DBVector;
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use term_dictionary::TermRef;
use super::{RocksDBIndexReader, DocRef};
use search::doc_id_set::DocIdSet;
use search::boolean_retrieval::BooleanQueryOp;
use search::scoring::{CombinatorScorer, ScoreFunctionOp};
use search::plan::SearchPlan;
use search::plan::boolean_query_builder::BooleanQueryBuilder;

impl<'a> RocksDBIndexReader<'a> {
    fn plan_query_combinator(&self, mut plan: &mut SearchPlan, queries: &Vec<Query>, join_op: BooleanQueryOp, score: bool, scorer: CombinatorScorer) {
        match queries.len() {
            0 => plan.boolean_query.push(BooleanQueryOp::PushEmpty),
            1 =>  self.plan_query(&mut plan, &queries[0], score),
            _ => {
                let mut query_iter = queries.iter();
                self.plan_query(&mut plan, query_iter.next().unwrap(), score);

                for query in query_iter {
                    self.plan_query(&mut plan, query, score);
                    plan.boolean_query.push(join_op.clone());
                }
            }
        }

        plan.score_function.push(ScoreFunctionOp::CombinatorScorer(queries.len() as u32, scorer));
    }

    fn plan_query(&self, mut plan: &mut SearchPlan, query: &Query, score: bool) {
        match *query {
            Query::MatchAll{ref score} => {
                plan.boolean_query.push(BooleanQueryOp::PushFull);
                plan.score_function.push(ScoreFunctionOp::Literal(*score));
            }
            Query::MatchNone => {
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::MatchTerm{ref field, ref term, ref scorer} => {
                // Get term
                let term_bytes = term.to_bytes();
                let term_ref = match self.store.term_dictionary.get(&term_bytes) {
                    Some(term_ref) => term_ref,
                    None => {
                        // Term doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                // Get field
                let field_ref = match self.schema().get_field_by_name(field) {
                    Some(field_ref) => field_ref,
                    None => {
                        // Field doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        return
                    }
                };

                let tag = plan.allocate_tag().unwrap_or(0);
                plan.boolean_query.push(BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag));
                plan.score_function.push(ScoreFunctionOp::TermScorer(field_ref, term_ref, scorer.clone(), tag));
            }
            Query::MatchMultiTerm{ref field, ref term_selector, ref scorer} => {
                // TODO
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::Conjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::And, score, CombinatorScorer::Avg);
            }
            Query::Disjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);
            }
            Query::NDisjunction{ref queries, minimum_should_match} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);  // FIXME
            }
            Query::DisjunctionMax{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Max);
            }
            Query::Filter{ref query, ref filter} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, filter, false);
                plan.boolean_query.push(BooleanQueryOp::And);
            }
            Query::Exclude{ref query, ref exclude} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, exclude, false);
                plan.boolean_query.push(BooleanQueryOp::AndNot);
            }
        }
    }

    fn search_chunk_boolean_phase(&self, plan: &SearchPlan, chunk: u32) -> (DocIdSet, HashMap<u8, DocIdSet>) {
        let mut tagged_docid_sets = HashMap::new();

        // Execute boolean query
        let mut stack = Vec::new();
        for op in plan.boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            tagged_docid_sets.insert(tag, data.clone());
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if plan.boolean_query_is_negated {
            let kb = KeyBuilder::chunk_stat(chunk, b"total_docs");
            let total_docs = match self.snapshot.get(&kb.key()) {
                Ok(Some(total_docs)) => BigEndian::read_i64(&total_docs) as u16,
                Ok(None) => 0,
                Err(e) => 0,  // FIXME
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        (matches, tagged_docid_sets)
    }

    fn score_doc(&self, doc_id: u16, tagged_docid_sets: &HashMap<u8, DocIdSet>, plan: &SearchPlan) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in plan.score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer, tag) => {
                    match tagged_docid_sets.get(&tag) {
                        Some(docid_set) => {
                            if docid_set.contains_doc(doc_id) {
                                stack.push(1.0f64);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64)
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for i in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for i in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32) {
        let (matches, tagged_docid_sets) = self.search_chunk_boolean_phase(plan, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, &tagged_docid_sets, plan);

            let doc_ref = DocRef(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        let mut plan = SearchPlan::new();
        self.plan_query(&mut plan, query, true);

        // Add operations to exclude deleted documents to boolean query
        plan.boolean_query.push(BooleanQueryOp::PushDeletionList);
        plan.boolean_query.push(BooleanQueryOp::AndNot);

        // Optimise boolean query
        let mut optimiser = BooleanQueryBuilder::new();
        for op in plan.boolean_query.iter() {
            optimiser.push_op(op);
        }
        let (boolean_query, boolean_query_is_negated) = optimiser.build();
        plan.boolean_query = boolean_query;
        plan.boolean_query_is_negated = boolean_query_is_negated;

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk);
        }
    }
}
<commit_msg>Always push a score op<commit_after>pub mod doc_id_set;
pub mod boolean_retrieval;
pub mod scoring;
pub mod plan;

use std::fmt;
use std::io::{Cursor, Read, Write};
use std::collections::HashMap;
use std::rc::Rc;

use kite::schema::{FieldRef, SchemaRead};
use kite::query::Query;
use kite::query::term_scorer::TermScorer;
use kite::collectors::{Collector, DocumentMatch};
use rocksdb::DBVector;
use byteorder::{ByteOrder, BigEndian};

use key_builder::KeyBuilder;
use term_dictionary::TermRef;
use super::{RocksDBIndexReader, DocRef};
use search::doc_id_set::DocIdSet;
use search::boolean_retrieval::BooleanQueryOp;
use search::scoring::{CombinatorScorer, ScoreFunctionOp};
use search::plan::SearchPlan;
use search::plan::boolean_query_builder::BooleanQueryBuilder;

impl<'a> RocksDBIndexReader<'a> {
    fn plan_query_combinator(&self, mut plan: &mut SearchPlan, queries: &Vec<Query>, join_op: BooleanQueryOp, score: bool, scorer: CombinatorScorer) {
        match queries.len() {
            0 => {
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            1 =>  self.plan_query(&mut plan, &queries[0], score),
            _ => {
                let mut query_iter = queries.iter();
                self.plan_query(&mut plan, query_iter.next().unwrap(), score);

                for query in query_iter {
                    self.plan_query(&mut plan, query, score);
                    plan.boolean_query.push(join_op.clone());
                }
            }
        }

        plan.score_function.push(ScoreFunctionOp::CombinatorScorer(queries.len() as u32, scorer));
    }

    fn plan_query(&self, mut plan: &mut SearchPlan, query: &Query, score: bool) {
        match *query {
            Query::MatchAll{ref score} => {
                plan.boolean_query.push(BooleanQueryOp::PushFull);
                plan.score_function.push(ScoreFunctionOp::Literal(*score));
            }
            Query::MatchNone => {
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::MatchTerm{ref field, ref term, ref scorer} => {
                // Get term
                let term_bytes = term.to_bytes();
                let term_ref = match self.store.term_dictionary.get(&term_bytes) {
                    Some(term_ref) => term_ref,
                    None => {
                        // Term doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
                        return
                    }
                };

                // Get field
                let field_ref = match self.schema().get_field_by_name(field) {
                    Some(field_ref) => field_ref,
                    None => {
                        // Field doesn't exist, so will never match
                        plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                        plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
                        return
                    }
                };

                let tag = plan.allocate_tag().unwrap_or(0);
                plan.boolean_query.push(BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag));
                plan.score_function.push(ScoreFunctionOp::TermScorer(field_ref, term_ref, scorer.clone(), tag));
            }
            Query::MatchMultiTerm{ref field, ref term_selector, ref scorer} => {
                // TODO
                plan.boolean_query.push(BooleanQueryOp::PushEmpty);
                plan.score_function.push(ScoreFunctionOp::Literal(0.0f64));
            }
            Query::Conjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::And, score, CombinatorScorer::Avg);
            }
            Query::Disjunction{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);
            }
            Query::NDisjunction{ref queries, minimum_should_match} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Avg);  // FIXME
            }
            Query::DisjunctionMax{ref queries} => {
                self.plan_query_combinator(&mut plan, queries, BooleanQueryOp::Or, score, CombinatorScorer::Max);
            }
            Query::Filter{ref query, ref filter} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, filter, false);
                plan.boolean_query.push(BooleanQueryOp::And);
            }
            Query::Exclude{ref query, ref exclude} => {
                self.plan_query(&mut plan, query, score);
                self.plan_query(&mut plan, exclude, false);
                plan.boolean_query.push(BooleanQueryOp::AndNot);
            }
        }
    }

    fn search_chunk_boolean_phase(&self, plan: &SearchPlan, chunk: u32) -> (DocIdSet, HashMap<u8, DocIdSet>) {
        let mut tagged_docid_sets = HashMap::new();

        // Execute boolean query
        let mut stack = Vec::new();
        for op in plan.boolean_query.iter() {
            match *op {
                BooleanQueryOp::PushEmpty => {
                    stack.push(DocIdSet::new_filled(0));
                }
                BooleanQueryOp::PushFull => {
                    stack.push(DocIdSet::new_filled(65536));
                }
                BooleanQueryOp::PushTermDirectory(field_ref, term_ref, tag) => {
                    let kb = KeyBuilder::chunk_dir_list(chunk, field_ref.ord(), term_ref.ord());
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(docid_set)) => {
                            let data = DocIdSet::FromRDB(docid_set);
                            tagged_docid_sets.insert(tag, data.clone());
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::PushDeletionList => {
                    let kb = KeyBuilder::chunk_del_list(chunk);
                    match self.snapshot.get(&kb.key()) {
                        Ok(Some(deletion_list)) => {
                            let data = DocIdSet::FromRDB(deletion_list);
                            stack.push(data);
                        }
                        Ok(None) => stack.push(DocIdSet::new_filled(0)),
                        Err(e) => {},  // FIXME
                    }
                }
                BooleanQueryOp::And => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.intersection(&b));
                }
                BooleanQueryOp::Or => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.union(&b));
                }
                BooleanQueryOp::AndNot => {
                    let b = stack.pop().expect("stack underflow");
                    let a = stack.pop().expect("stack underflow");
                    stack.push(a.exclusion(&b));
                }
            }
        }

        if !stack.len() == 1 {
            // TODO: Error
        }
        let mut matches = stack.pop().unwrap();

        // Invert the list if the query is negated
        if plan.boolean_query_is_negated {
            let kb = KeyBuilder::chunk_stat(chunk, b"total_docs");
            let total_docs = match self.snapshot.get(&kb.key()) {
                Ok(Some(total_docs)) => BigEndian::read_i64(&total_docs) as u16,
                Ok(None) => 0,
                Err(e) => 0,  // FIXME
            };

            let all_docs = DocIdSet::new_filled(total_docs);
            matches = all_docs.exclusion(&matches);
        }

        (matches, tagged_docid_sets)
    }

    fn score_doc(&self, doc_id: u16, tagged_docid_sets: &HashMap<u8, DocIdSet>, plan: &SearchPlan) -> f64 {
        // Execute score function
        let mut stack = Vec::new();
        for op in plan.score_function.iter() {
            match *op {
                ScoreFunctionOp::Literal(val) => stack.push(val),
                ScoreFunctionOp::TermScorer(field_ref, term_ref, ref scorer, tag) => {
                    match tagged_docid_sets.get(&tag) {
                        Some(docid_set) => {
                            if docid_set.contains_doc(doc_id) {
                                stack.push(1.0f64);
                            } else {
                                stack.push(0.0f64);
                            }
                        }
                        None => stack.push(0.0f64)
                    }
                }
                ScoreFunctionOp::CombinatorScorer(num_vals, ref scorer) => {
                    let score = match *scorer {
                        CombinatorScorer::Avg => {
                            let mut total_score = 0.0f64;

                            for i in 0..num_vals {
                                total_score += stack.pop().expect("stack underflow");
                            }

                            total_score / num_vals as f64
                        }
                        CombinatorScorer::Max => {
                            let mut max_score = 0.0f64;

                            for i in 0..num_vals {
                                let score = stack.pop().expect("stack underflow");
                                if score > max_score {
                                    max_score = score
                                }
                            }

                            max_score
                        }
                    };

                    stack.push(score);
                }
            }
        }

        stack.pop().expect("stack underflow")
    }

    fn search_chunk<C: Collector>(&self, collector: &mut C, plan: &SearchPlan, chunk: u32) {
        let (matches, tagged_docid_sets) = self.search_chunk_boolean_phase(plan, chunk);

        // Score documents and pass to collector
        for doc in matches.iter() {
            let score = self.score_doc(doc, &tagged_docid_sets, plan);

            let doc_ref = DocRef(chunk, doc);
            let doc_match = DocumentMatch::new_scored(doc_ref.as_u64(), score);
            collector.collect(doc_match);
        }
    }

    pub fn search<C: Collector>(&self, collector: &mut C, query: &Query) {
        let mut plan = SearchPlan::new();
        self.plan_query(&mut plan, query, true);

        // Add operations to exclude deleted documents to boolean query
        plan.boolean_query.push(BooleanQueryOp::PushDeletionList);
        plan.boolean_query.push(BooleanQueryOp::AndNot);

        // Optimise boolean query
        let mut optimiser = BooleanQueryBuilder::new();
        for op in plan.boolean_query.iter() {
            optimiser.push_op(op);
        }
        let (boolean_query, boolean_query_is_negated) = optimiser.build();
        plan.boolean_query = boolean_query;
        plan.boolean_query_is_negated = boolean_query_is_negated;

        // Run query on each chunk
        for chunk in self.store.chunks.iter_active(&self.snapshot) {
            self.search_chunk(collector, &plan, chunk);
        }
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fmt;
use std::fs::{self, PathExt};
use std::io::{self, Read};
use std::path::{Path, PathBuf};
use std::process;
use std::str::FromStr;
use std::sync::atomic;

use csv;

use Csv;

static XSV_INTEGRATION_TEST_DIR: &'static str = "xit";

static NEXT_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;

pub struct Workdir {
    root: PathBuf,
    dir: PathBuf,
    flexible: bool,
}

impl Workdir {
    pub fn new(name: &str) -> Workdir {
        let id = NEXT_ID.fetch_add(1, atomic::Ordering::SeqCst);
        let root = env::current_exe().unwrap()
                       .parent()
                       .expect("executable's directory")
                       .to_path_buf();
        let dir = root.join(XSV_INTEGRATION_TEST_DIR)
                      .join(name)
                      .join(&format!("test-{}", id));

        // I don't get why this is necessary, but Travis seems to need it?
        // let md = fs::metadata(&dir);
        // if fs::metadata(&dir).map(|md| md.is_dir()).unwrap_or(false) {
            // if let Err(err) = fs::remove_dir_all(&dir) {
                // panic!("Could not remove directory '{:?}': {}", dir, err);
            // }
        // }
        // if fs::metadata(&dir).map(|md| md.is_file()).unwrap_or(false) {
            // if let Err(err) = fs::remove_file(&dir) {
                // panic!("Could not remove file '{:?}': {}", dir, err);
            // }
        // }
        if let Err(err) = create_dir_all(&dir) {
            panic!("Could not create '{:?}': {}", dir, err);
        }
        Workdir { root: root, dir: dir, flexible: false }
    }

    pub fn flexible(mut self, yes: bool) -> Workdir {
        self.flexible = yes;
        self
    }

    pub fn create<T: Csv>(&self, name: &str, rows: T) {
        let mut wtr = match csv::Writer::from_file(&self.path(name)) {
            Ok(wtr) => wtr.flexible(self.flexible),
            Err(err) => panic!("Could not open '{:?}': {}",
                                self.path(name), err),
        };
        for row in rows.to_vecs().into_iter() {
            wtr.write(row.iter()).unwrap();
        }
        wtr.flush().unwrap();
    }

    pub fn create_indexed<T: Csv>(&self, name: &str, rows: T) {
        self.create(name, rows);

        let mut cmd = self.command("index");
        cmd.arg(name);
        self.run(&mut cmd);
    }

    pub fn read_stdout<T: Csv>(&self, cmd: &mut process::Command) -> T {
        let mut rdr = csv::Reader::from_string(self.stdout::<String>(cmd))
                                  .has_headers(false);
        Csv::from_vecs(rdr.records().collect::<Result<_, _>>().unwrap())
    }

    pub fn command(&self, sub_command: &str) -> process::Command {
        let mut cmd = process::Command::new(&self.xsv_bin());
        cmd.current_dir(&self.dir).arg(sub_command);
        cmd
    }

    pub fn output(&self, cmd: &mut process::Command) -> process::Output {
        debug!("[{}]: {:?}", self.dir.display(), cmd);
        let o = cmd.output().unwrap();
        if !o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command failed but expected success!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr))
        }
        o
    }

    pub fn run(&self, cmd: &mut process::Command) {
        self.output(cmd);
    }

    pub fn stdout<T: FromStr>(&self, cmd: &mut process::Command) -> T {
        let o = self.output(cmd);
        let stdout = String::from_utf8_lossy(&o.stdout);
        stdout.trim().parse().ok().expect(
            &format!("Could not convert from string: '{}'", stdout))
    }

    pub fn assert_err(&self, cmd: &mut process::Command) {
        let o = cmd.output().unwrap();
        if o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command succeeded but expected failure!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr));
        }
    }

    pub fn from_str<T: FromStr>(&self, name: &Path) -> T {
        let mut o = String::new();
        fs::File::open(name).unwrap().read_to_string(&mut o).unwrap();
        o.parse().ok().expect("fromstr")
    }

    pub fn path(&self, name: &str) -> PathBuf {
        self.dir.join(name)
    }

    pub fn xsv_bin(&self) -> PathBuf {
        self.root.join("xsv")
    }
}

impl fmt::Debug for Workdir {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "path={}", self.dir.display())
    }
}

fn create_dir_all<P: AsRef<Path>>(p: P) -> io::Result<()> {
    let mut last_err = None;
    for _ in 0..10 {
        if let Err(err) = fs::create_dir_all(&p) {
            last_err = Some(err);
            ::std::thread::sleep_ms(500);
        } else {
            return Ok(())
        }
    }
    Err(last_err.unwrap())
}
<commit_msg>Yay. Finally fixed Travis with a hack.<commit_after>use std::env;
use std::fmt;
use std::fs;
use std::io::{self, Read};
use std::path::{Path, PathBuf};
use std::process;
use std::str::FromStr;
use std::sync::atomic;

use csv;

use Csv;

static XSV_INTEGRATION_TEST_DIR: &'static str = "xit";

static NEXT_ID: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;

pub struct Workdir {
    root: PathBuf,
    dir: PathBuf,
    flexible: bool,
}

impl Workdir {
    pub fn new(name: &str) -> Workdir {
        let id = NEXT_ID.fetch_add(1, atomic::Ordering::SeqCst);
        let root = env::current_exe().unwrap()
                       .parent()
                       .expect("executable's directory")
                       .to_path_buf();
        let dir = root.join(XSV_INTEGRATION_TEST_DIR)
                      .join(name)
                      .join(&format!("test-{}", id));
        if let Err(err) = create_dir_all(&dir) {
            panic!("Could not create '{:?}': {}", dir, err);
        }
        Workdir { root: root, dir: dir, flexible: false }
    }

    pub fn flexible(mut self, yes: bool) -> Workdir {
        self.flexible = yes;
        self
    }

    pub fn create<T: Csv>(&self, name: &str, rows: T) {
        let mut wtr = csv::Writer::from_file(&self.path(name)).unwrap();
        for row in rows.to_vecs().into_iter() {
            wtr.write(row.iter()).unwrap();
        }
        wtr.flush().unwrap();
    }

    pub fn create_indexed<T: Csv>(&self, name: &str, rows: T) {
        self.create(name, rows);

        let mut cmd = self.command("index");
        cmd.arg(name);
        self.run(&mut cmd);
    }

    pub fn read_stdout<T: Csv>(&self, cmd: &mut process::Command) -> T {
        let mut rdr = csv::Reader::from_string(self.stdout::<String>(cmd))
                                  .has_headers(false);
        Csv::from_vecs(rdr.records().collect::<Result<_, _>>().unwrap())
    }

    pub fn command(&self, sub_command: &str) -> process::Command {
        let mut cmd = process::Command::new(&self.xsv_bin());
        cmd.current_dir(&self.dir).arg(sub_command);
        cmd
    }

    pub fn output(&self, cmd: &mut process::Command) -> process::Output {
        debug!("[{}]: {:?}", self.dir.display(), cmd);
        let o = cmd.output().unwrap();
        if !o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command failed but expected success!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr))
        }
        o
    }

    pub fn run(&self, cmd: &mut process::Command) {
        self.output(cmd);
    }

    pub fn stdout<T: FromStr>(&self, cmd: &mut process::Command) -> T {
        let o = self.output(cmd);
        let stdout = String::from_utf8_lossy(&o.stdout);
        stdout.trim().parse().ok().expect(
            &format!("Could not convert from string: '{}'", stdout))
    }

    pub fn assert_err(&self, cmd: &mut process::Command) {
        let o = cmd.output().unwrap();
        if o.status.success() {
            panic!("\n\n===== {:?} =====\n\
                    command succeeded but expected failure!\
                    \n\ncwd: {}\
                    \n\nstatus: {}\
                    \n\nstdout: {}\n\nstderr: {}\
                    \n\n=====\n",
                   cmd, self.dir.display(), o.status,
                   String::from_utf8_lossy(&o.stdout),
                   String::from_utf8_lossy(&o.stderr));
        }
    }

    pub fn from_str<T: FromStr>(&self, name: &Path) -> T {
        let mut o = String::new();
        fs::File::open(name).unwrap().read_to_string(&mut o).unwrap();
        o.parse().ok().expect("fromstr")
    }

    pub fn path(&self, name: &str) -> PathBuf {
        self.dir.join(name)
    }

    pub fn xsv_bin(&self) -> PathBuf {
        self.root.join("xsv")
    }
}

impl fmt::Debug for Workdir {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "path={}", self.dir.display())
    }
}

fn create_dir_all<P: AsRef<Path>>(p: P) -> io::Result<()> {
    let mut last_err = None;
    for _ in 0..10 {
        if let Err(err) = fs::create_dir_all(&p) {
            last_err = Some(err);
            ::std::thread::sleep_ms(500);
        } else {
            return Ok(())
        }
    }
    Err(last_err.unwrap())
}
<|endoftext|>


<commit_before>use rustc::hir::def_id::LOCAL_CRATE;
use rustc::mir::mono::MonoItem;
use rustc::session::config::OptLevel;
use rustc::ty::{self, TyCtxt, Instance};
use rustc::ty::subst::InternalSubsts;
use rustc::ty::print::obsolete::DefPathBasedNames;
use syntax::attr::InlineAttr;
use std::fmt;
use rustc::mir::mono::Linkage;
use syntax_pos::symbol::InternedString;
use syntax::source_map::Span;

/// Describes how a monomorphization will be instantiated in object files.
#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
pub enum InstantiationMode {
    /// There will be exactly one instance of the given MonoItem. It will have
    /// external linkage so that it can be linked to from other codegen units.
    GloballyShared {
        /// In some compilation scenarios we may decide to take functions that
        /// are typically `LocalCopy` and instead move them to `GloballyShared`
        /// to avoid codegenning them a bunch of times. In this situation,
        /// however, our local copy may conflict with other crates also
        /// inlining the same function.
        ///
        /// This flag indicates that this situation is occurring, and informs
        /// symbol name calculation that some extra mangling is needed to
        /// avoid conflicts. Note that this may eventually go away entirely if
        /// ThinLTO enables us to *always* have a globally shared instance of a
        /// function within one crate's compilation.
        may_conflict: bool,
    },

    /// Each codegen unit containing a reference to the given MonoItem will
    /// have its own private copy of the function (with internal linkage).
    LocalCopy,
}

pub trait MonoItemExt<'tcx>: fmt::Debug {
    fn as_mono_item(&self) -> &MonoItem<'tcx>;

    fn is_generic_fn(&self) -> bool {
        match *self.as_mono_item() {
            MonoItem::Fn(ref instance) => {
                instance.substs.non_erasable_generics().next().is_some()
            }
            MonoItem::Static(..) |
            MonoItem::GlobalAsm(..) => false,
        }
    }

    fn symbol_name(&self, tcx: TyCtxt<'tcx>) -> ty::SymbolName {
        match *self.as_mono_item() {
            MonoItem::Fn(instance) => tcx.symbol_name(instance),
            MonoItem::Static(def_id) => {
                tcx.symbol_name(Instance::mono(tcx, def_id))
            }
            MonoItem::GlobalAsm(hir_id) => {
                let def_id = tcx.hir().local_def_id(hir_id);
                ty::SymbolName {
                    name: InternedString::intern(&format!("global_asm_{:?}", def_id))
                }
            }
        }
    }
    fn instantiation_mode(&self, tcx: TyCtxt<'tcx>) -> InstantiationMode {
        let inline_in_all_cgus =
            tcx.sess.opts.debugging_opts.inline_in_all_cgus.unwrap_or_else(|| {
                tcx.sess.opts.optimize != OptLevel::No
            }) && !tcx.sess.opts.cg.link_dead_code;

        match *self.as_mono_item() {
            MonoItem::Fn(ref instance) => {
                let entry_def_id = tcx.entry_fn(LOCAL_CRATE).map(|(id, _)| id);
                // If this function isn't inlined or otherwise has explicit
                // linkage, then we'll be creating a globally shared version.
                if self.explicit_linkage(tcx).is_some() ||
                    !instance.def.requires_local(tcx) ||
                    Some(instance.def_id()) == entry_def_id
                {
                    return InstantiationMode::GloballyShared  { may_conflict: false }
                }

                // At this point we don't have explicit linkage and we're an
                // inlined function. If we're inlining into all CGUs then we'll
                // be creating a local copy per CGU
                if inline_in_all_cgus {
                    return InstantiationMode::LocalCopy
                }

                // Finally, if this is `#[inline(always)]` we're sure to respect
                // that with an inline copy per CGU, but otherwise we'll be
                // creating one copy of this `#[inline]` function which may
                // conflict with upstream crates as it could be an exported
                // symbol.
                match tcx.codegen_fn_attrs(instance.def_id()).inline {
                    InlineAttr::Always => InstantiationMode::LocalCopy,
                    _ => {
                        InstantiationMode::GloballyShared  { may_conflict: true }
                    }
                }
            }
            MonoItem::Static(..) |
            MonoItem::GlobalAsm(..) => {
                InstantiationMode::GloballyShared { may_conflict: false }
            }
        }
    }

    fn explicit_linkage(&self, tcx: TyCtxt<'tcx>) -> Option<Linkage> {
        let def_id = match *self.as_mono_item() {
            MonoItem::Fn(ref instance) => instance.def_id(),
            MonoItem::Static(def_id) => def_id,
            MonoItem::GlobalAsm(..) => return None,
        };

        let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
        codegen_fn_attrs.linkage
    }

    /// Returns `true` if this instance is instantiable - whether it has no unsatisfied
    /// predicates.
    ///
    /// In order to codegen an item, all of its predicates must hold, because
    /// otherwise the item does not make sense. Type-checking ensures that
    /// the predicates of every item that is *used by* a valid item *do*
    /// hold, so we can rely on that.
    ///
    /// However, we codegen collector roots (reachable items) and functions
    /// in vtables when they are seen, even if they are not used, and so they
    /// might not be instantiable. For example, a programmer can define this
    /// public function:
    ///
    ///     pub fn foo<'a>(s: &'a mut ()) where &'a mut (): Clone {
    ///         <&mut () as Clone>::clone(&s);
    ///     }
    ///
    /// That function can't be codegened, because the method `<&mut () as Clone>::clone`
    /// does not exist. Luckily for us, that function can't ever be used,
    /// because that would require for `&'a mut (): Clone` to hold, so we
    /// can just not emit any code, or even a linker reference for it.
    ///
    /// Similarly, if a vtable method has such a signature, and therefore can't
    /// be used, we can just not emit it and have a placeholder (a null pointer,
    /// which will never be accessed) in its place.
    fn is_instantiable(&self, tcx: TyCtxt<'tcx>) -> bool {
        debug!("is_instantiable({:?})", self);
        let (def_id, substs) = match *self.as_mono_item() {
            MonoItem::Fn(ref instance) => (instance.def_id(), instance.substs),
            MonoItem::Static(def_id) => (def_id, InternalSubsts::empty()),
            // global asm never has predicates
            MonoItem::GlobalAsm(..) => return true
        };

        tcx.substitute_normalize_and_test_predicates((def_id, &substs))
    }

    fn to_string(&self, tcx: TyCtxt<'tcx>, debug: bool) -> String {
        return match *self.as_mono_item() {
            MonoItem::Fn(instance) => {
                to_string_internal(tcx, "fn ", instance, debug)
            },
            MonoItem::Static(def_id) => {
                let instance = Instance::new(def_id, tcx.intern_substs(&[]));
                to_string_internal(tcx, "static ", instance, debug)
            },
            MonoItem::GlobalAsm(..) => {
                "global_asm".to_string()
            }
        };

        fn to_string_internal<'a, 'tcx>(
            tcx: TyCtxt<'tcx>,
            prefix: &str,
            instance: Instance<'tcx>,
            debug: bool,
        ) -> String {
            let mut result = String::with_capacity(32);
            result.push_str(prefix);
            let printer = DefPathBasedNames::new(tcx, false, false);
            printer.push_instance_as_string(instance, &mut result, debug);
            result
        }
    }

    fn local_span(&self, tcx: TyCtxt<'tcx>) -> Option<Span> {
        match *self.as_mono_item() {
            MonoItem::Fn(Instance { def, .. }) => {
                tcx.hir().as_local_hir_id(def.def_id())
            }
            MonoItem::Static(def_id) => {
                tcx.hir().as_local_hir_id(def_id)
            }
            MonoItem::GlobalAsm(hir_id) => {
                Some(hir_id)
            }
        }.map(|hir_id| tcx.hir().span(hir_id))
    }
}

impl MonoItemExt<'tcx> for MonoItem<'tcx> {
    fn as_mono_item(&self) -> &MonoItem<'tcx> {
        self
    }
}
<commit_msg>Rollup merge of #65158 - ishitatsuyuki:remove-dead, r=Mark-Simulacrum<commit_after><|endoftext|>