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

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

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

type Exponent = u8;

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

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

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

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);
        // Assert the factor doesn't already exist in the Decomposition object
        debug_assert_eq!(self.0.iter_mut().find(|(f, _)| *f == factor), None);

        self.0.push((factor, exp))
    }

    fn is_one(&self) -> bool {
        self.0.is_empty()
    }

    fn pop(&mut self) -> Option<(u64, Exponent)> {
        self.0.pop()
    }

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

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

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

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

        true
    }
}
impl Eq for Decomposition {}

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

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

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

    #[cfg(test)]
    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

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

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

        Ok(())
    }
}

fn _find_factor<A: Arithmetic + miller_rabin::Basis>(num: u64) -> Option<u64> {
    use miller_rabin::Result::*;

    let n = A::new(num);
    match miller_rabin::test::<A>(n) {
        Prime => None,
        Composite(d) => Some(d),
        Pseudoprime => Some(rho::find_divisor::<A>(n)),
    }
}

fn find_factor(num: u64) -> Option<u64> {
    if num < (1 << 32) {
        _find_factor::<Montgomery<u32>>(num)
    } else {
        _find_factor::<Montgomery<u64>>(num)
    }
}

pub fn factor(num: u64) -> Factors {
    let mut factors = Factors::one();

    if num < 2 {
        return factors;
    }

    let mut n = num;
    let n_zeros = num.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }
    debug_assert_eq!(num, n * factors.product());

    if n == 1 {
        return factors;
    }

    table::factor(&mut n, &mut factors);
    debug_assert_eq!(num, n * factors.product());

    if n == 1 {
        return factors;
    }

    let mut dec = Decomposition::one();
    dec.add(n, 1);

    while !dec.is_one() {
        // Check correctness invariant
        debug_assert_eq!(num, factors.product() * dec.product());

        let (factor, exp) = dec.pop().unwrap();

        if let Some(divisor) = find_factor(factor) {
            let mut gcd_queue = Decomposition::one();

            let quotient = factor / divisor;
            let mut trivial_gcd = quotient == divisor;
            if trivial_gcd {
                gcd_queue.add(divisor, exp + 1);
            } else {
                gcd_queue.add(divisor, exp);
                gcd_queue.add(quotient, exp);
            }

            while !trivial_gcd {
                debug_assert_eq!(factor, gcd_queue.product());

                let mut tmp = Decomposition::one();
                trivial_gcd = true;
                for i in 0..gcd_queue.0.len() - 1 {
                    let (mut a, exp_a) = gcd_queue.0[i];
                    let (mut b, exp_b) = gcd_queue.0[i + 1];

                    if a == 1 {
                        continue;
                    }

                    let g = gcd(a, b);
                    if g != 1 {
                        trivial_gcd = false;
                        a /= g;
                        b /= g;
                    }
                    if a != 1 {
                        tmp.add(a, exp_a);
                    }
                    if g != 1 {
                        tmp.add(g, exp_a + exp_b);
                    }

                    if i + 1 != gcd_queue.0.len() - 1 {
                        gcd_queue.0[i + 1].0 = b;
                    } else if b != 1 {
                        tmp.add(b, exp_b);
                    }
                }
                gcd_queue = tmp;
            }

            debug_assert_eq!(factor, gcd_queue.product());
            dec.0.extend(gcd_queue.0);
        } else {
            // factor is prime
            factors.add(factor, exp);
        }
    }

    factors
}

#[cfg(test)]
mod tests {
    use super::{factor, Factors};
    use quickcheck::quickcheck;

    #[test]
    fn factor_correctly_recombines_prior_test_failures() {
        let prior_failures = [
            // * integers with duplicate factors (ie, N.pow(M))
            4566769_u64, // == 2137.pow(2)
            2044854919485649_u64,
            18446739546814299361_u64,
            18446738440860217487_u64,
            18446736729316206481_u64,
        ];
        assert!(prior_failures.iter().all(|i| factor(*i).product() == *i));
    }

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

    #[test]
    fn factor_recombines_small_squares() {
        // factor(18446736729316206481) == 4294966441 ** 2 ; causes debug_assert fault for repeated decomposition factor in add()
        // ToDO: explain/combine with factor_18446736729316206481 and factor_18446739546814299361 tests
        assert!((1..10_000)
            .map(|i| (2 * i + 1) * (2 * i + 1))
            .all(|i| factor(i).product() == i));
    }

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

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

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

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

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

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

            return f;
        }
    }
}
<commit_msg>factor::tests: Check that powers of known-factorization numbers are factored correctly (#1831)<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

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

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

type Exponent = u8;

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

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

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

    fn add(&mut self, factor: u64, exp: Exponent) {
        debug_assert!(exp > 0);
        // Assert the factor doesn't already exist in the Decomposition object
        debug_assert_eq!(self.0.iter_mut().find(|(f, _)| *f == factor), None);

        self.0.push((factor, exp))
    }

    fn is_one(&self) -> bool {
        self.0.is_empty()
    }

    fn pop(&mut self) -> Option<(u64, Exponent)> {
        self.0.pop()
    }

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

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

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

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

        true
    }
}
impl Eq for Decomposition {}

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

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

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

    #[cfg(test)]
    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    fn product(&self) -> u64 {
        self.0.borrow().product()
    }
}

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

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

        Ok(())
    }
}

fn _find_factor<A: Arithmetic + miller_rabin::Basis>(num: u64) -> Option<u64> {
    use miller_rabin::Result::*;

    let n = A::new(num);
    match miller_rabin::test::<A>(n) {
        Prime => None,
        Composite(d) => Some(d),
        Pseudoprime => Some(rho::find_divisor::<A>(n)),
    }
}

fn find_factor(num: u64) -> Option<u64> {
    if num < (1 << 32) {
        _find_factor::<Montgomery<u32>>(num)
    } else {
        _find_factor::<Montgomery<u64>>(num)
    }
}

pub fn factor(num: u64) -> Factors {
    let mut factors = Factors::one();

    if num < 2 {
        return factors;
    }

    let mut n = num;
    let n_zeros = num.trailing_zeros();
    if n_zeros > 0 {
        factors.add(2, n_zeros as Exponent);
        n >>= n_zeros;
    }
    debug_assert_eq!(num, n * factors.product());

    if n == 1 {
        return factors;
    }

    table::factor(&mut n, &mut factors);
    debug_assert_eq!(num, n * factors.product());

    if n == 1 {
        return factors;
    }

    let mut dec = Decomposition::one();
    dec.add(n, 1);

    while !dec.is_one() {
        // Check correctness invariant
        debug_assert_eq!(num, factors.product() * dec.product());

        let (factor, exp) = dec.pop().unwrap();

        if let Some(divisor) = find_factor(factor) {
            let mut gcd_queue = Decomposition::one();

            let quotient = factor / divisor;
            let mut trivial_gcd = quotient == divisor;
            if trivial_gcd {
                gcd_queue.add(divisor, exp + 1);
            } else {
                gcd_queue.add(divisor, exp);
                gcd_queue.add(quotient, exp);
            }

            while !trivial_gcd {
                debug_assert_eq!(factor, gcd_queue.product());

                let mut tmp = Decomposition::one();
                trivial_gcd = true;
                for i in 0..gcd_queue.0.len() - 1 {
                    let (mut a, exp_a) = gcd_queue.0[i];
                    let (mut b, exp_b) = gcd_queue.0[i + 1];

                    if a == 1 {
                        continue;
                    }

                    let g = gcd(a, b);
                    if g != 1 {
                        trivial_gcd = false;
                        a /= g;
                        b /= g;
                    }
                    if a != 1 {
                        tmp.add(a, exp_a);
                    }
                    if g != 1 {
                        tmp.add(g, exp_a + exp_b);
                    }

                    if i + 1 != gcd_queue.0.len() - 1 {
                        gcd_queue.0[i + 1].0 = b;
                    } else if b != 1 {
                        tmp.add(b, exp_b);
                    }
                }
                gcd_queue = tmp;
            }

            debug_assert_eq!(factor, gcd_queue.product());
            dec.0.extend(gcd_queue.0);
        } else {
            // factor is prime
            factors.add(factor, exp);
        }
    }

    factors
}

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

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

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

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

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

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

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

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

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

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

            return f;
        }
    }
}

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

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

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


<commit_before>use byteorder::{BigEndian, WriteBytesExt};
use generic_array::{ArrayLength, GenericArray};
use std::io::{Result, Error, ErrorKind, Read, Write};
use std::mem::size_of;

pub trait Message {
    fn encode_with_hash(&self) -> Result<Vec<u8>> {
        let hash = self.hash();
        let size = hash.size() + self.size();
        let mut buffer = Vec::with_capacity(size);
        hash.encode(&mut buffer)?;
        self.encode(&mut buffer)?;
        Ok(buffer)
    }
    fn hash(&self) -> i64 { 0 }
    fn encode(&self, buffer: &mut Write) -> Result<()>;
    fn decode(&mut self, buffer: &mut Read) -> Result<()> { Err(Error::new(ErrorKind::Other, "Unimplemented")) }
    fn size(&self) -> usize;
}

impl Message for bool {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        let value : i8 = match *self {
            true => 1,
            false => 0
        };
        value.encode(buffer)
    }

    fn size(&self) -> usize {
        size_of::<i8>()
    }
}

impl Message for u8 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_u8(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i8 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i8(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i16 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i16::<BigEndian>(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i32 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i32::<BigEndian>(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i64 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i64::<BigEndian>(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for f32 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_f32::<BigEndian>(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for f64 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_f64::<BigEndian>(*self)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for String {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        let len : i32 = self.len() as i32 + 1;
        len.encode(buffer)?;
        for &b in self.as_bytes() {
            b.encode(buffer)?;
        }
        (0 as u8).encode(buffer)?;
        Ok(())
    }

    fn size(&self) -> usize {
        size_of::<i32>() + self.len() + 1
    }
}

impl<T> Message for Vec<T> where
    T : Message {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        for val in self.iter() {
            val.encode(buffer)?;
        }
        Ok(())
    }

    fn size(&self) -> usize {
        self.iter().map(Message::size).sum()
    }
}

impl<T,N> Message for GenericArray<T, N> where
    T : Message, N : ArrayLength<T> {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        for val in self.iter() {
            val.encode(buffer)?;
        }
        Ok(())
    }

    fn size(&self) -> usize {
        self.iter().map(Message::size).sum()
    }
}
<commit_msg>Add decode functions to Message trait.<commit_after>use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use generic_array::{ArrayLength, GenericArray};
use std::io::{Result, Error, ErrorKind, Read, Write};
use std::mem::size_of;
use std::default::Default;

/// TODO: Document
pub trait Message {
    /// Encodes a message into a buffer, with the message hash at the beginning.
    fn encode_with_hash(&self) -> Result<Vec<u8>> {
        let hash = self.hash();
        let size = hash.size() + self.size();
        let mut buffer = Vec::with_capacity(size);
        hash.encode(&mut buffer)?;
        self.encode(&mut buffer)?;
        Ok(buffer)
    }

    /// Decodes a message from a buffer,
    /// and also checks that the hash at the beginning is correct.
    fn decode_with_hash(&mut self, mut buffer: &mut Read) -> Result<()>
        where Self: Sized + Default {

        let mut hash : i64 = 0;
        hash.decode(&mut buffer)?;
        if hash != self.hash() {
            return Err(Error::new(ErrorKind::Other, "Invalid hash"));
        }
        self.decode(buffer)
    }

    /// Returns the message hash for this type.
    /// Returns `0` for all primitive types.
    /// Generated `Lcm` types should implement this function.
    fn hash(&self) -> i64 { 0 }

    /// Encodes a message into a buffer.
    /// `Lcm` uses a `Vec<u8>` with its capacity set to the value returned by [`size()`].
    fn encode(&self, buffer: &mut Write) -> Result<()>;

    /// Decodes a message from a buffer.
    fn decode(&mut self, buffer: &mut Read) -> Result<()>;

    /// Returns the number of bytes this message is expected to take when encoded.
    fn size(&self) -> usize;
}

impl Message for bool {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        let value : i8 = match *self {
            true => 1,
            false => 0
        };
        value.encode(buffer)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        let value = buffer.read_i8()?;
        *self = match value {
            0 => false,
            1 => true,
            _ => return Err(Error::new(ErrorKind::InvalidData, "Booleans should be encoded as 0 or 1"))
        };
        Ok(())
    }

    fn size(&self) -> usize {
        size_of::<i8>()
    }
}

impl Message for u8 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_u8(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_u8()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i8 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i8(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_i8()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i16 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i16::<BigEndian>(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_i16::<BigEndian>()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i32 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i32::<BigEndian>(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_i32::<BigEndian>()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for i64 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_i64::<BigEndian>(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_i64::<BigEndian>()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for f32 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_f32::<BigEndian>(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_f32::<BigEndian>()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for f64 {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        buffer.write_f64::<BigEndian>(*self)
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        Ok(*self = buffer.read_f64::<BigEndian>()?)
    }

    fn size(&self) -> usize {
        size_of::<Self>()
    }
}

impl Message for String {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        let len : i32 = self.len() as i32 + 1;
        len.encode(buffer)?;
        for &b in self.as_bytes() {
            b.encode(buffer)?;
        }
        (0 as u8).encode(buffer)?;
        Ok(())
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        let len = buffer.read_i32::<BigEndian>()? - 1;
        let mut buf = Vec::with_capacity(len as usize);
        buf.decode(buffer)?;
        *self = String::from_utf8(buf).map_err(|e| Error::new(ErrorKind::Other, e))?;
        match buffer.read_u8() {
            Ok(0) => Ok(()),
            Ok(_) => Err(Error::new(ErrorKind::InvalidData, "Expected null terminator")),
            Err(e) => Err(e)
        }
    }

    fn size(&self) -> usize {
        size_of::<i32>() + self.len() + 1
    }
}

impl<T> Message for Vec<T> where
    T : Message + Default {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        for val in self.iter() {
            val.encode(buffer)?;
        }
        Ok(())
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        while self.len() != self.capacity() {
            let mut element = T::default();
            element.decode(buffer)?;
            self.push(element);
        }
        Ok(())
    }

    fn size(&self) -> usize {
        self.iter().map(Message::size).sum()
    }
}

impl<T,N> Message for GenericArray<T, N> where
    T : Message + Default, N : ArrayLength<T> {
    fn encode(&self, buffer: &mut Write) -> Result<()> {
        for val in self.iter() {
            val.encode(buffer)?;
        }
        Ok(())
    }

    fn decode(&mut self, buffer: &mut Read) -> Result<()> {
        for i in 0..self.len() {
            let mut element = T::default();
            element.decode(buffer)?;
            self[i] = element;
        }
        Ok(())
    }

    fn size(&self) -> usize {
        self.iter().map(Message::size).sum()
    }
}
<|endoftext|>


<commit_before>use std::io::{Result, Write};
use {Class, Function, Type, TypeKind};
use std::collections::BTreeMap;

fn get_hl_type(ty: &Type) -> String {
    if ty.is_custom {
        match ty.kind {
            TypeKind::Ref => format!("{}Ref", ty.name),
            TypeKind::RefMut => format!("{}RefMut", ty.name),
            TypeKind::Value => ty.name.clone(),
        }
    } else {
        get_ll_type(ty).to_string()
    }
}

fn get_ll_type(ty: &Type) -> &str {
    match (ty.kind, ty.name.as_str()) {
        (TypeKind::Ref, "c_char") => "c_char_p",
        (TypeKind::Ref, _) | (TypeKind::RefMut, _) => "c_void_p",
        (_, t) if !ty.is_custom => match t {
            "i8" => "c_int8",
            "i16" => "c_int16",
            "i32" => "c_int32",
            "i64" => "c_int64",
            "u8" => "c_uint8",
            "u16" => "c_uint16",
            "u32" => "c_uint32",
            "u64" => "c_uint64",
            "usize" => "c_size_t",
            "f32" => "c_float",
            "f64" => "c_double",
            "bool" => "c_bool",
            "()" => "None",
            "c_char" => "c_char",
            "Json" => "c_char_p",
            x => x,
        },
        _ => "c_void_p",
    }
}

fn map_var(var: &str) -> &str {
    if var == "this" {
        "self"
    } else {
        var
    }
}

fn write_class_comments<W: Write>(mut writer: W, comments: &[String]) -> Result<()> {
    write!(
        writer,
        r#"
    """"#
    )?;

    for comment in comments {
        write!(
            writer,
            r#"{}
    "#,
            comment
                .replace("<NULL>", "None")
                .replace("<TRUE>", "True")
                .replace("<FALSE>", "False")
        )?;
    }

    write!(
        writer,
        r#""""
"#
    )
}

fn write_fn<W: Write>(mut writer: W, function: &Function) -> Result<()> {
    let is_static = function.is_static();
    let has_return_type = function.has_return_type();
    let return_type = get_hl_type(&function.output);

    if is_static {
        write!(
            writer,
            r#"
    @staticmethod
    def {}("#,
            function.method
        )?;
    } else {
        write!(
            writer,
            r#"
    def {}("#,
            function.method
        )?;
    }

    for (i, &(ref name, _)) in function.inputs.iter().enumerate() {
        if i != 0 {
            write!(writer, ", ")?;
        }
        write!(writer, "{}", map_var(name))?;
    }

    write!(
        writer,
        r#"):
        "#
    )?;

    write!(writer, r#"""""#)?;

    for comment in &function.comments {
        write!(
            writer,
            r#"{}
        "#,
            comment
                .replace("<NULL>", "None")
                .replace("<TRUE>", "True")
                .replace("<FALSE>", "False")
        )?;
    }

    write!(
        writer,
        r#""""
        "#
    )?;

    for &(ref name, ref typ) in function.inputs.iter() {
        if typ.is_custom {
            write!(
                writer,
                r#"if {name}.ptr == None:
            raise Exception("{name} is disposed")
        "#,
                name = map_var(name)
            )?;
        }
    }

    if has_return_type {
        if function.output.is_custom {
            write!(writer, r#"result = {}("#, return_type)?;
        } else {
            write!(writer, "result = ")?;
        }
    }

    write!(writer, r#"livesplit_core_native.{}("#, &function.name)?;

    for (i, &(ref name, ref typ)) in function.inputs.iter().enumerate() {
        if i != 0 {
            write!(writer, ", ")?;
        }
        write!(
            writer,
            "{}",
            if name == "this" {
                "self.ptr".to_string()
            } else if typ.is_custom {
                format!("{}.ptr", name)
            } else {
                name.to_string()
            }
        )?;
    }

    write!(writer, ")")?;

    if has_return_type && function.output.is_custom {
        write!(writer, r#")"#)?;
    }

    for &(ref name, ref typ) in function.inputs.iter() {
        if typ.is_custom && typ.kind == TypeKind::Value {
            write!(
                writer,
                r#"
        {}.ptr = None"#,
                map_var(name)
            )?;
        }
    }

    if has_return_type {
        if function.output.is_nullable && function.output.is_custom {
            write!(
                writer,
                r#"
        if result.ptr == None:
            return None"#
            )?;
        }
        write!(
            writer,
            r#"
        return result"#
        )?;
    }

    write!(
        writer,
        r#"
"#
    )?;

    Ok(())
}

pub fn write<W: Write>(mut writer: W, classes: &BTreeMap<String, Class>) -> Result<()> {
    write!(writer,
           "{}",
           r#"#!/usr/bin/env python3
# coding: utf-8

import sys, ctypes
from ctypes import c_char_p, c_void_p, c_int8, c_int16, c_int32, c_int64, c_uint8, c_uint16, c_uint32, c_uint64, c_size_t, c_float, c_double, c_bool, c_char, c_byte

prefix = {'win32': ''}.get(sys.platform, './lib')
extension = {'darwin': '.dylib', 'win32': '.dll'}.get(sys.platform, '.so')
livesplit_core_native = ctypes.cdll.LoadLibrary(prefix + "livesplit_core" + extension)
"#)?;

    for class in classes.values() {
        for function in class
            .static_fns
            .iter()
            .chain(class.own_fns.iter())
            .chain(class.shared_fns.iter())
            .chain(class.mut_fns.iter())
        {
            write!(
                writer,
                r#"
livesplit_core_native.{}.argtypes = ("#,
                function.name
            )?;

            for &(_, ref typ) in &function.inputs {
                write!(writer, "{}, ", get_ll_type(typ))?;
            }

            write!(
                writer,
                r#")
livesplit_core_native.{}.restype = {}"#,
                function.name,
                get_ll_type(&function.output)
            )?;
        }
    }

    writeln!(writer)?;

    for (class_name, class) in classes {
        let class_name_ref = format!("{}Ref", class_name);
        let class_name_ref_mut = format!("{}RefMut", class_name);

        write!(
            writer,
            r#"
class {class}:"#,
            class = class_name_ref
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        for function in &class.shared_fns {
            write_fn(&mut writer, function)?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr

class {class}({base_class}):"#,
            class = class_name_ref_mut,
            base_class = class_name_ref
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        for function in &class.mut_fns {
            write_fn(&mut writer, function)?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr

class {class}({base_class}):"#,
            class = class_name,
            base_class = class_name_ref_mut
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        write!(
            writer,
            r#"
    def drop(self):
        if self.ptr != None:"#,
        )?;

        if let Some(function) = class.own_fns.iter().find(|f| f.method == "drop") {
            write!(
                writer,
                r#"
            livesplit_core_native.{}(self.ptr)"#,
                function.name
            )?;
        }

        write!(
            writer,
            r#"
            self.ptr = None

    def __del__(self):
        self.drop()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.drop()
"#
        )?;

        for function in class.static_fns.iter().chain(class.own_fns.iter()) {
            if function.method != "drop" {
                write_fn(&mut writer, function)?;
            }
        }

        if class_name == "Run" {
            writeln!(
                writer,
                r#"
    @staticmethod
    def parse_file(file, path, load_files):
        data = file.read()
        if sys.version_info[0] > 2:
            if isinstance(data, str):
                raise TypeError("File must be opened in binary mode!")
        bytes = bytearray(data)
        bufferType = c_byte * len(bytes)
        buffer = bufferType(*bytes)
        return Run.parse(buffer, len(bytes), path, load_files)"#
            )?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr
"#
        )?;
    }

    Ok(())
}
<commit_msg>Automatically convert C strings to and from str type<commit_after>use std::io::{Result, Write};
use {Class, Function, Type, TypeKind};
use std::collections::BTreeMap;

fn get_hl_type(ty: &Type) -> String {
    if ty.is_custom {
        match ty.kind {
            TypeKind::Ref => format!("{}Ref", ty.name),
            TypeKind::RefMut => format!("{}RefMut", ty.name),
            TypeKind::Value => ty.name.clone(),
        }
    } else {
        get_ll_type(ty).to_string()
    }
}

fn get_ll_type(ty: &Type) -> &str {
    match (ty.kind, ty.name.as_str()) {
        (TypeKind::Ref, "c_char") => "c_char_p",
        (TypeKind::Ref, _) | (TypeKind::RefMut, _) => "c_void_p",
        (_, t) if !ty.is_custom => match t {
            "i8" => "c_int8",
            "i16" => "c_int16",
            "i32" => "c_int32",
            "i64" => "c_int64",
            "u8" => "c_uint8",
            "u16" => "c_uint16",
            "u32" => "c_uint32",
            "u64" => "c_uint64",
            "usize" => "c_size_t",
            "f32" => "c_float",
            "f64" => "c_double",
            "bool" => "c_bool",
            "()" => "None",
            "c_char" => "c_char",
            "Json" => "c_char_p",
            x => x,
        },
        _ => "c_void_p",
    }
}

fn map_var(var: &str) -> &str {
    if var == "this" {
        "self"
    } else {
        var
    }
}

fn write_class_comments<W: Write>(mut writer: W, comments: &[String]) -> Result<()> {
    write!(
        writer,
        r#"
    """"#
    )?;

    for comment in comments {
        write!(
            writer,
            r#"{}
    "#,
            comment
                .replace("<NULL>", "None")
                .replace("<TRUE>", "True")
                .replace("<FALSE>", "False")
        )?;
    }

    write!(
        writer,
        r#""""
"#
    )
}

fn write_fn<W: Write>(mut writer: W, function: &Function) -> Result<()> {
    let is_static = function.is_static();
    let has_return_type = function.has_return_type();
    let return_type = get_hl_type(&function.output);

    if is_static {
        write!(
            writer,
            r#"
    @staticmethod
    def {}("#,
            function.method
        )?;
    } else {
        write!(
            writer,
            r#"
    def {}("#,
            function.method
        )?;
    }

    for (i, &(ref name, _)) in function.inputs.iter().enumerate() {
        if i != 0 {
            write!(writer, ", ")?;
        }
        write!(writer, "{}", map_var(name))?;
    }

    write!(
        writer,
        r#"):
        "#
    )?;

    write!(writer, r#"""""#)?;

    for comment in &function.comments {
        write!(
            writer,
            r#"{}
        "#,
            comment
                .replace("<NULL>", "None")
                .replace("<TRUE>", "True")
                .replace("<FALSE>", "False")
        )?;
    }

    write!(
        writer,
        r#""""
        "#
    )?;

    for &(ref name, ref typ) in function.inputs.iter() {
        if typ.is_custom {
            write!(
                writer,
                r#"if {name}.ptr == None:
            raise Exception("{name} is disposed")
        "#,
                name = map_var(name)
            )?;
        }
    }

    if has_return_type {
        if function.output.is_custom {
            write!(writer, r#"result = {}("#, return_type)?;
        } else {
            write!(writer, "result = ")?;
        }
    }

    write!(writer, r#"livesplit_core_native.{}("#, &function.name)?;

    for (i, &(ref name, ref typ)) in function.inputs.iter().enumerate() {
        if i != 0 {
            write!(writer, ", ")?;
        }
        write!(
            writer,
            "{}",
            if name == "this" {
                "self.ptr".to_string()
            } else if typ.is_custom {
                format!("{}.ptr", name)
            } else if typ.name == "c_char" {
                format!("{}.encode()", name)
            } else {
                name.to_string()
            }
        )?;
    }

    write!(writer, ")")?;

    if has_return_type {
        if function.output.is_custom {
            write!(writer, r#")"#)?;
        } else if function.output.name == "c_char" {
            write!(writer, r#".decode()"#)?;
        }
    }

    for &(ref name, ref typ) in function.inputs.iter() {
        if typ.is_custom && typ.kind == TypeKind::Value {
            write!(
                writer,
                r#"
        {}.ptr = None"#,
                map_var(name)
            )?;
        }
    }

    if has_return_type {
        if function.output.is_nullable && function.output.is_custom {
            write!(
                writer,
                r#"
        if result.ptr == None:
            return None"#
            )?;
        }
        write!(
            writer,
            r#"
        return result"#
        )?;
    }

    write!(
        writer,
        r#"
"#
    )?;

    Ok(())
}

pub fn write<W: Write>(mut writer: W, classes: &BTreeMap<String, Class>) -> Result<()> {
    write!(writer,
           "{}",
           r#"#!/usr/bin/env python3
# coding: utf-8

import sys, ctypes
from ctypes import c_char_p, c_void_p, c_int8, c_int16, c_int32, c_int64, c_uint8, c_uint16, c_uint32, c_uint64, c_size_t, c_float, c_double, c_bool, c_char, c_byte

prefix = {'win32': ''}.get(sys.platform, './lib')
extension = {'darwin': '.dylib', 'win32': '.dll'}.get(sys.platform, '.so')
livesplit_core_native = ctypes.cdll.LoadLibrary(prefix + "livesplit_core" + extension)
"#)?;

    for class in classes.values() {
        for function in class
            .static_fns
            .iter()
            .chain(class.own_fns.iter())
            .chain(class.shared_fns.iter())
            .chain(class.mut_fns.iter())
        {
            write!(
                writer,
                r#"
livesplit_core_native.{}.argtypes = ("#,
                function.name
            )?;

            for &(_, ref typ) in &function.inputs {
                write!(writer, "{}, ", get_ll_type(typ))?;
            }

            write!(
                writer,
                r#")
livesplit_core_native.{}.restype = {}"#,
                function.name,
                get_ll_type(&function.output)
            )?;
        }
    }

    writeln!(writer)?;

    for (class_name, class) in classes {
        let class_name_ref = format!("{}Ref", class_name);
        let class_name_ref_mut = format!("{}RefMut", class_name);

        write!(
            writer,
            r#"
class {class}:"#,
            class = class_name_ref
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        for function in &class.shared_fns {
            write_fn(&mut writer, function)?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr

class {class}({base_class}):"#,
            class = class_name_ref_mut,
            base_class = class_name_ref
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        for function in &class.mut_fns {
            write_fn(&mut writer, function)?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr

class {class}({base_class}):"#,
            class = class_name,
            base_class = class_name_ref_mut
        )?;

        write_class_comments(&mut writer, &class.comments)?;

        write!(
            writer,
            r#"
    def drop(self):
        if self.ptr != None:"#,
        )?;

        if let Some(function) = class.own_fns.iter().find(|f| f.method == "drop") {
            write!(
                writer,
                r#"
            livesplit_core_native.{}(self.ptr)"#,
                function.name
            )?;
        }

        write!(
            writer,
            r#"
            self.ptr = None

    def __del__(self):
        self.drop()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.drop()
"#
        )?;

        for function in class.static_fns.iter().chain(class.own_fns.iter()) {
            if function.method != "drop" {
                write_fn(&mut writer, function)?;
            }
        }

        if class_name == "Run" {
            writeln!(
                writer,
                r#"
    @staticmethod
    def parse_file(file, path, load_files):
        data = file.read()
        if sys.version_info[0] > 2:
            if isinstance(data, str):
                raise TypeError("File must be opened in binary mode!")
        bytes = bytearray(data)
        bufferType = c_byte * len(bytes)
        buffer = bufferType(*bytes)
        return Run.parse(buffer, len(bytes), path, load_files)"#
            )?;
        }

        write!(
            writer,
            r#"
    def __init__(self, ptr):
        self.ptr = ptr
"#
        )?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::fmt;
use {encode_config_slice, Config};
use std::io::{Result, Write};

/// A `Write` proxy that base64-encodes written data and hands the result off to another writer.
pub struct Base64Encoder<'a> {
    config: Config,
    w: &'a mut Write,
    extra: [u8; 3],
    extra_len: usize,
    output: [u8; 1024],
}

impl<'a> fmt::Debug for Base64Encoder<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "extra:{:?} extra_len:{:?} output[..5]: {:?}",
            self.extra,
            self.extra_len,
            &self.output[0..5]
        )
    }
}

impl<'a> Base64Encoder<'a> {
    /// Create a new encoder around an existing writer.
    pub fn new(w: &'a mut Write, config: Config) -> Base64Encoder<'a> {
        Base64Encoder {
            config,
            w,                   // writer to write encoded data to
            extra: [0u8; 3],     // extra data left over from previous write
            extra_len: 0,        // how much extra data
            output: [0u8; 1024], // output buffer
        }
    }
}

impl<'a> Write for Base64Encoder<'a> {
    fn write(&mut self, input: &[u8]) -> Result<usize> {
        let mut p = input;
        let mut input_read_cnt = 0;

        // process leftover stuff from last write
        if self.extra_len > 0 {
            let mut i = 0;
            while i < p.len() && self.extra_len < 3 {
                self.extra[self.extra_len] = p[i];
                self.extra_len += 1;
                i += 1;
            }
            input_read_cnt += i;
            p = &p[i..];
            if self.extra_len < 3 {
                // not enough to actually encode, yet.
                return Ok(input_read_cnt);
            }
            let sz = encode_config_slice(&self.extra[..3], self.config, &mut self.output[..]);
            self.extra_len = 0;
            let _ = self.w.write(&self.output[..sz])?;
        }

        // process chunks, 768 bytes -> 1024 (encoded) bytes, at a time
        let mut output_cnt = 0;
        while p.len() >= 3 {
            let mut nn = (self.output.len() / 4) * 3;
            if nn > p.len() {
                // force nn to be a multiple of three o we can encode cleanly
                nn = p.len();
                nn -= nn % 3;
            }
            let sz = encode_config_slice(&p[..nn], self.config, &mut self.output[..]);
            input_read_cnt += nn;
            output_cnt += sz;
            p = &p[nn..];
        }

        // stash leftover bytes
        let mut i = 0;
        while i < p.len() {
            self.extra[i] = p[i];
            i += 1;
        }
        input_read_cnt += p.len();
        self.extra_len = p.len();

        let _ = self.w.write(&self.output[..output_cnt])?;

        return Ok(input_read_cnt);
    }

    fn flush(&mut self) -> Result<()> {
        if self.extra_len > 0 {
            let sz = encode_config_slice(
                &self.extra[..self.extra_len],
                self.config,
                &mut self.output[..],
            );
            let _ = self.w.write(&self.output[..sz])?;
        }

        self.w.flush()
    }
}

#[cfg(test)]
mod tests {
    use encoder::Base64Encoder;
    use {encode_config, URL_SAFE};
    use std::io::{Cursor, Write};

    #[test]
    fn encode_three_bytes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abc").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_nine_bytes_two_writes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abcdef").unwrap();
            assert_eq!(sz, 6);
            let sz = enc.write(b"ghi").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdefghi", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_one_two() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bc").unwrap();
            assert_eq!(sz, 2);
        }
        assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_one_five() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bcdef").unwrap();
            assert_eq!(sz, 5);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdef", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_1_2_3() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bc").unwrap();
            assert_eq!(sz, 2);
            let sz = enc.write(b"def").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdef", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_with_padding() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abcd").unwrap();
            assert_eq!(sz, 4);

            enc.flush().unwrap();
        }
        assert_eq!(&c.get_ref()[..], encode_config("abcd", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_with_padding_multiple_writes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let _ = enc.write(b"a").unwrap();
            let _ = enc.write(b"bcd").unwrap();
            let _ = enc.write(b"ef").unwrap();
            let _ = enc.write(b"g").unwrap();

            enc.flush().unwrap();
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdefg", URL_SAFE).as_bytes()
        );
    }
}
<commit_msg>Add a randomized test. Fix bug that choked on large input slices. Reformat and use longer var names.<commit_after>use std::fmt;
use encode::encode_to_slice;
use {encode_config_slice, Config};
use std::io::{Result, Write};

/// A `Write` proxy that base64-encodes written data and hands the result off to another writer.
pub struct Base64Encoder<'a> {
    config: Config,
    w: &'a mut Write,
    extra: [u8; 3],
    extra_len: usize,
    output: [u8; 1024],
}

impl<'a> fmt::Debug for Base64Encoder<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "extra:{:?} extra_len:{:?} output[..5]: {:?}",
            self.extra,
            self.extra_len,
            &self.output[0..5]
        )
    }
}

impl<'a> Base64Encoder<'a> {
    /// Create a new encoder around an existing writer.
    pub fn new(w: &'a mut Write, config: Config) -> Base64Encoder<'a> {
        Base64Encoder {
            config,
            w,                   // writer to write encoded data to
            extra: [0u8; 3],     // extra data left over from previous write
            extra_len: 0,        // how much extra data
            output: [0u8; 1024], // output buffer
        }
    }
}

impl<'a> Write for Base64Encoder<'a> {
    fn write(&mut self, input: &[u8]) -> Result<usize> {
        // TODO handle line breaks
        let mut input = input;
        let mut input_read_cnt = 0;

        // process leftover stuff from last write
        if self.extra_len > 0 {
            let mut i = 0;
            while i < input.len() && self.extra_len < 3 {
                self.extra[self.extra_len] = input[i];
                self.extra_len += 1;
                i += 1;
            }
            input_read_cnt += i;
            input = &input[i..];

            if self.extra_len < 3 {
                // not enough to actually encode, yet.
                return Ok(input_read_cnt);
            }

            let encoded_size = encode_to_slice(
                &self.extra[..3],
                &mut self.output[..],
                self.config.char_set.encode_table(),
            );
            self.extra_len = 0;
            let _ = self.w.write(&self.output[..encoded_size])?;
        }

        // process chunks, 768 bytes -> 1024 (encoded) bytes, at a time
        // TODO try using `chunks()`
        while input.len() >= 3 {
            let mut bytes_to_encode = (self.output.len() / 4) * 3;
            if bytes_to_encode > input.len() {
                // force nn to be a multiple of three o we can encode cleanly
                bytes_to_encode = input.len();
                bytes_to_encode -= bytes_to_encode % 3;
            }
            let encoded_size = encode_to_slice(
                &input[..bytes_to_encode],
                &mut self.output[..],
                self.config.char_set.encode_table(),
            );
            input_read_cnt += bytes_to_encode;
            input = &input[bytes_to_encode..];
            let _ = self.w.write(&self.output[..encoded_size])?;
        }

        // stash leftover bytes
        let mut i = 0;
        while i < input.len() {
            self.extra[i] = input[i];
            i += 1;
        }
        input_read_cnt += input.len();
        self.extra_len = input.len();

        return Ok(input_read_cnt);
    }

    fn flush(&mut self) -> Result<()> {
        if self.extra_len > 0 {
            // TODO should not write padding on each flush() -- should be able to flush() > 1x
            let sz = encode_config_slice(
                &self.extra[..self.extra_len],
                self.config,
                &mut self.output[..],
            );
            let _ = self.w.write(&self.output[..sz])?;
        }

        self.w.flush()
    }
}

#[cfg(test)]
mod tests {
    extern crate rand;

    use super::Base64Encoder;
    use {encode_config, encode_config_buf, URL_SAFE};
    use tests::random_config;

    use std::io::{Cursor, Write};
    use std::str;

    use self::rand::Rng;
    use self::rand::distributions::range;

    #[test]
    fn encode_three_bytes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abc").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_nine_bytes_two_writes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abcdef").unwrap();
            assert_eq!(sz, 6);
            let sz = enc.write(b"ghi").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdefghi", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_one_two() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bc").unwrap();
            assert_eq!(sz, 2);
        }
        assert_eq!(&c.get_ref()[..], encode_config("abc", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_one_five() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bcdef").unwrap();
            assert_eq!(sz, 5);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdef", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_1_2_3() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"a").unwrap();
            assert_eq!(sz, 1);
            let sz = enc.write(b"bc").unwrap();
            assert_eq!(sz, 2);
            let sz = enc.write(b"def").unwrap();
            assert_eq!(sz, 3);
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdef", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_with_padding() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let sz = enc.write(b"abcd").unwrap();
            assert_eq!(sz, 4);

            enc.flush().unwrap();
        }
        assert_eq!(&c.get_ref()[..], encode_config("abcd", URL_SAFE).as_bytes());
    }

    #[test]
    fn encode_with_padding_multiple_writes() {
        let mut c = Cursor::new(Vec::new());
        {
            let mut enc = Base64Encoder::new(&mut c, URL_SAFE);

            let _ = enc.write(b"a").unwrap();
            let _ = enc.write(b"bcd").unwrap();
            let _ = enc.write(b"ef").unwrap();
            let _ = enc.write(b"g").unwrap();

            enc.flush().unwrap();
        }
        assert_eq!(
            &c.get_ref()[..],
            encode_config("abcdefg", URL_SAFE).as_bytes()
        );
    }

    #[test]
    fn encode_random_config_matches_normal_encode() {
        let mut rng = rand::thread_rng();
        let mut orig_data = Vec::<u8>::new();
        let mut stream_encoded = Vec::<u8>::new();
        let mut normal_encoded = String::new();
        let line_len_range = range::Range::new(1, 2000);

        for _ in 0..1_000 {
            orig_data.clear();
            stream_encoded.clear();
            normal_encoded.clear();

            // TODO for now, ignore configs with line wraps
            let mut config = random_config(&mut rng, &line_len_range);
            while let ::LineWrap::Wrap(_, _) = config.line_wrap {
                config = random_config(&mut rng, &line_len_range)
            }

            let orig_len: usize = rng.gen_range(100, 10_000);
            for _ in 0..orig_len {
                orig_data.push(rng.gen());
            }

            // encode the normal way
            encode_config_buf(&orig_data, config, &mut normal_encoded);

            // encode via the stream encoder
            {
                let mut stream_encoder = Base64Encoder::new(&mut stream_encoded, config);
                let mut bytes_consumed = 0;
                while bytes_consumed < orig_len {
                    let input_len: usize = rng.gen_range(0, orig_len - bytes_consumed + 1);

                    // write a little bit of the data
                    bytes_consumed += stream_encoder
                        .write(&orig_data[bytes_consumed..bytes_consumed + input_len])
                        .unwrap();
                }

                // TODO final write should be done by drop()
                stream_encoder.flush().unwrap();

                assert_eq!(orig_len, bytes_consumed);
            }

            assert_eq!(normal_encoded, str::from_utf8(&stream_encoded).unwrap());
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::hash_map::Values;
use quadtree::{Span, Dir};

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Position {
    x: i32,
    y: i32
}

impl Position {
    pub fn new(x: i32, y: i32) -> Position {
        return Position { x: x, y: y };
    }
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct PosSpan {
    nw: Position,
    width: i32,
    height: i32
}

impl PosSpan {

    /// Creates a new span starting at `(x, y)` inclusive
    /// ending at `(x + w, y + h)` exclusive.
    pub fn new(x: i32, y: i32, w: i32, h: i32) -> PosSpan {
        return PosSpan {
            nw: Position { x: x, y: y },
            width: w,
            height: h
        };
    }
}

impl Span<PosSpan, Position> for PosSpan {
    fn dir_of(&self, t: &Position) -> Option<Dir> {
        if t.x < self.nw.x {
            if t.y < self.nw.y {
                return Some(Dir::NW);
            } else if t.y < self.nw.y + self.height {
                return Some(Dir::W);
            } else {
                return Some(Dir::SW);
            }
        } else if t.x < self.nw.x + self.width {
            if t.y < self.nw.y {
                return Some(Dir::N);
            } else if t.y < self.nw.y + self.height {
                return None;
            } else {
                return Some(Dir::S);
            }
        } else {
            if t.y < self.nw.y {
                return Some(Dir::NE);
            } else if t.y < self.nw.y + self.height {
                return Some(Dir::E);
            } else {
                return Some(Dir::SE);
            }
        }
    }

    fn north_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x, self.nw.y - self.height, self.width, self.height);
    }

    fn south_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x, self.nw.y + self.height, self.width, self.height);
    }

    fn west_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x - self.width, self.nw.y, self.width, self.height);
    }

    fn east_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x + self.width, self.nw.y, self.width, self.height);
    }

    fn can_split(&self) -> bool {
        return self.width > 1 && self.height > 1;
    }

    fn split(&self) -> HashMap<Dir, PosSpan> {
        let left_x = self.nw.x;
        let left_y = self.nw.y;
        let width_mid = self.width / 2;
        let height_mid = self.height / 2;
        let mut result = HashMap::new();
        result.insert(Dir::NW, PosSpan::new(left_x, left_y, width_mid, height_mid));
        result.insert(Dir::NE, PosSpan::new(left_x + width_mid, left_y, self.width - width_mid, height_mid));
        result.insert(Dir::SW, PosSpan::new(left_x, left_y + height_mid, width_mid, self.height - height_mid));
        result.insert(Dir::SE, PosSpan::new(left_x + width_mid, left_y + height_mid, self.width - width_mid, self.height - height_mid));
        return result;
    }

    fn expand(&self, dir: &Dir) -> PosSpan {
        return match dir {
            &Dir::N => PosSpan::new(self.nw.x, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::S => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::E => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::W => PosSpan::new(self.nw.x - self.width, self.nw.y, self.width * 2, self.height * 2),
            &Dir::NE => PosSpan::new(self.nw.x, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::NW => PosSpan::new(self.nw.x - self.width, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::SE => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::SW => PosSpan::new(self.nw.x - self.width, self.nw.y, self.width * 2, self.height * 2)
        };
    }

    fn merge(spans: Values<Dir, PosSpan>) -> PosSpan {
        unreachable!("Doesn't work");

        let mut min_x = None;
        let mut min_y = None;
        let mut max_x = None;
        let mut max_y = None;

        for span in spans {
            if min_x == None || min_x.unwrap() > span.nw.x {
                min_x = Some(span.nw.x);
            }
            if min_y == None || min_y.unwrap() > span.nw.y {
                min_y = Some(span.nw.y);
            }
            if max_x == None || max_x.unwrap() < span.nw.x + span.width {
                max_x = Some(span.nw.x + span.width);
            }
            if max_y == None || max_y.unwrap() < span.nw.y + span.height {
                max_y = Some(span.nw.y + span.height);
            }
        }

        return PosSpan::new(
                    min_x.unwrap(),
                    min_y.unwrap(),
                    max_x.unwrap() - min_x.unwrap(),
                    max_y.unwrap() - min_y.unwrap());
    }

    fn overlaps(&self, other: &PosSpan) -> bool {
        if self.nw.x + self.width <= other.nw.x {
            return false;
        }
        if other.nw.x + other.width <= self.nw.x {
            return false;
        }
        if self.nw.y + self.height <= other.nw.y {
            return false;
        }
        if other.nw.y + other.height <= self.nw.y {
            return false;
        }
        return true;
    }
}
<commit_msg>bettero debug for position<commit_after>use std::collections::HashMap;
use std::collections::hash_map::Values;
use quadtree::{Span, Dir};
use std::fmt;

#[derive(PartialEq, Eq, Clone)]
pub struct Position {
    x: i32,
    y: i32
}

impl Position {
    pub fn new(x: i32, y: i32) -> Position {
        return Position { x: x, y: y };
    }
}

impl fmt::Debug for Position {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({}, {})", self.x, self.y)
    }
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct PosSpan {
    nw: Position,
    width: i32,
    height: i32
}

impl PosSpan {

    /// Creates a new span starting at `(x, y)` inclusive
    /// ending at `(x + w, y + h)` exclusive.
    pub fn new(x: i32, y: i32, w: i32, h: i32) -> PosSpan {
        return PosSpan {
            nw: Position { x: x, y: y },
            width: w,
            height: h
        };
    }
}

impl Span<PosSpan, Position> for PosSpan {
    fn dir_of(&self, t: &Position) -> Option<Dir> {
        if t.x < self.nw.x {
            if t.y < self.nw.y {
                return Some(Dir::NW);
            } else if t.y < self.nw.y + self.height {
                return Some(Dir::W);
            } else {
                return Some(Dir::SW);
            }
        } else if t.x < self.nw.x + self.width {
            if t.y < self.nw.y {
                return Some(Dir::N);
            } else if t.y < self.nw.y + self.height {
                return None;
            } else {
                return Some(Dir::S);
            }
        } else {
            if t.y < self.nw.y {
                return Some(Dir::NE);
            } else if t.y < self.nw.y + self.height {
                return Some(Dir::E);
            } else {
                return Some(Dir::SE);
            }
        }
    }

    fn north_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x, self.nw.y - self.height, self.width, self.height);
    }

    fn south_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x, self.nw.y + self.height, self.width, self.height);
    }

    fn west_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x - self.width, self.nw.y, self.width, self.height);
    }

    fn east_span(&self) -> PosSpan {
        return PosSpan::new(self.nw.x + self.width, self.nw.y, self.width, self.height);
    }

    fn can_split(&self) -> bool {
        return self.width > 1 && self.height > 1;
    }

    fn split(&self) -> HashMap<Dir, PosSpan> {
        let left_x = self.nw.x;
        let left_y = self.nw.y;
        let width_mid = self.width / 2;
        let height_mid = self.height / 2;
        let mut result = HashMap::new();
        result.insert(Dir::NW, PosSpan::new(left_x, left_y, width_mid, height_mid));
        result.insert(Dir::NE, PosSpan::new(left_x + width_mid, left_y, self.width - width_mid, height_mid));
        result.insert(Dir::SW, PosSpan::new(left_x, left_y + height_mid, width_mid, self.height - height_mid));
        result.insert(Dir::SE, PosSpan::new(left_x + width_mid, left_y + height_mid, self.width - width_mid, self.height - height_mid));
        return result;
    }

    fn expand(&self, dir: &Dir) -> PosSpan {
        return match dir {
            &Dir::N => PosSpan::new(self.nw.x, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::S => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::E => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::W => PosSpan::new(self.nw.x - self.width, self.nw.y, self.width * 2, self.height * 2),
            &Dir::NE => PosSpan::new(self.nw.x, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::NW => PosSpan::new(self.nw.x - self.width, self.nw.y - self.height, self.width * 2, self.height * 2),
            &Dir::SE => PosSpan::new(self.nw.x, self.nw.y, self.width * 2, self.height * 2),
            &Dir::SW => PosSpan::new(self.nw.x - self.width, self.nw.y, self.width * 2, self.height * 2)
        };
    }

    fn merge(spans: Values<Dir, PosSpan>) -> PosSpan {
        unreachable!("Doesn't work");

        let mut min_x = None;
        let mut min_y = None;
        let mut max_x = None;
        let mut max_y = None;

        for span in spans {
            if min_x == None || min_x.unwrap() > span.nw.x {
                min_x = Some(span.nw.x);
            }
            if min_y == None || min_y.unwrap() > span.nw.y {
                min_y = Some(span.nw.y);
            }
            if max_x == None || max_x.unwrap() < span.nw.x + span.width {
                max_x = Some(span.nw.x + span.width);
            }
            if max_y == None || max_y.unwrap() < span.nw.y + span.height {
                max_y = Some(span.nw.y + span.height);
            }
        }

        return PosSpan::new(
                    min_x.unwrap(),
                    min_y.unwrap(),
                    max_x.unwrap() - min_x.unwrap(),
                    max_y.unwrap() - min_y.unwrap());
    }

    fn overlaps(&self, other: &PosSpan) -> bool {
        if self.nw.x + self.width <= other.nw.x {
            return false;
        }
        if other.nw.x + other.width <= self.nw.x {
            return false;
        }
        if self.nw.y + self.height <= other.nw.y {
            return false;
        }
        if other.nw.y + other.height <= self.nw.y {
            return false;
        }
        return true;
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::sync::Arc;

use regex::Regex;

use config::Config;
use errors::*;
use git::Git;
use git_clone_manager::GitCloneManager;
use github;
use github::api::{GithubSessionFactory, Session};
use messenger::{self, Messenger};
use slack::{SlackAttachmentBuilder, SlackRequest};
use worker;

fn clone_and_merge_pull_request(
    github_app: &GithubSessionFactory,
    clone_mgr: &GitCloneManager,
    owner: &str,
    repo: &str,
    pull_request: &github::PullRequest,
    target_branch: &str,
) -> Result<github::PullRequest> {

    let session = github_app.new_session(owner, repo)?;
    let held_clone_dir = clone_mgr.clone(owner, repo)?;
    let clone_dir = held_clone_dir.dir();
    let git = Git::new(session.github_host(), session.github_token(), clone_dir);

    merge_pull_request(&git, &session, owner, repo, pull_request, target_branch)
}

pub fn merge_pull_request(
    git: &Git,
    session: &Session,
    owner: &str,
    repo: &str,
    pull_request: &github::PullRequest,
    target_branch: &str,
) -> Result<github::PullRequest> {
    if !pull_request.is_merged() {
        return Err(format_err!("Pull Request #{} is not yet merged.", pull_request.number));
    }

    let merge_commit_sha;
    if let Some(ref sha) = pull_request.merge_commit_sha {
        merge_commit_sha = sha;
    } else {
        return Err(format_err!("Pull Request #{} has no merge commit.", pull_request.number));
    }

    // strip everything before last slash
    let regex = Regex::new(r".*/").unwrap();
    let pr_branch_name =
        format!("{}-{}", regex.replace(&pull_request.head.ref_name, ""), regex.replace(&target_branch, ""));

    // make sure there isn't already such a branch
    let current_remotes = git.run(&["ls-remote", "--heads"])?;
    if current_remotes.contains(&format!("refs/heads/{}", pr_branch_name)) {
        return Err(format_err!("PR branch already exists on origin: '{}'", pr_branch_name));
    }

    let (title, body, whitespace_mode) = cherry_pick(
        &git,
        &merge_commit_sha,
        &pr_branch_name,
        pull_request.number,
        &target_branch,
        &pull_request.base.ref_name,
    )?;

    git.run(&["push", "origin", &format!("HEAD:{}", pr_branch_name)])?;

    let new_pr = session.create_pull_request(owner, repo, &title, &body, &pr_branch_name, &target_branch)?;

    let mut assignees: Vec<String> = pull_request.assignees.iter().map(|a| a.login().to_string()).collect();
    assignees.retain(|r| r != pull_request.user.login());
    if !assignees.is_empty() {
        session.assign_pull_request(owner, repo, new_pr.number, assignees)?;
    }

    let mut reviewers: Vec<String> = pull_request.all_reviewers().into_iter().map(|a| a.login().to_string()).collect();
    reviewers.retain(|r| r != pull_request.user.login());
    if !reviewers.is_empty() {
        session.request_review(owner, repo, new_pr.number, reviewers)?;
    }

    if whitespace_mode.len() > 0 {
        let msg = format!("Cherry-pick required option `{}`. Please verify correctness.", whitespace_mode);
        if let Err(e) = session.comment_pull_request(owner, repo, new_pr.number, &msg) {
            error!("Error making whitespace comment on pull request: {}", e);
        }
    }

    Ok(new_pr)
}

pub fn cherry_pick(
    git: &Git,
    commit_hash: &str,
    pr_branch_name: &str,
    pr_number: u32,
    target_branch: &str,
    orig_base_branch: &str,
) -> Result<(String, String, String)> {
    git.checkout_branch(pr_branch_name, &format!("origin/{}", target_branch))?;

    let (user, email) = git.get_commit_author(commit_hash)?;
    let email = format!("user.email={}", email);
    let user = format!("user.name={}", user);
    let user_opts = ["-c", &email, "-c", &user];

    // cherry-pick!

    let mut whitespace_mode = "";
    if let Err(e) = do_cherry_pick(git, &commit_hash, &[], &user_opts) {
        info!("Could not cherry-pick normally. Ignoring changed whitespace. {}", e);

        whitespace_mode = "ignore-space-change";
        if let Err(e) = do_cherry_pick(git, &commit_hash, &["-X", whitespace_mode], &user_opts) {
            info!("Could not cherry-pick with `-X {}`. Ignoring all whitespace. {}", whitespace_mode, e);

            whitespace_mode = "ignore-all-space";
            if let Err(e) = do_cherry_pick(git, &commit_hash, &["-X", whitespace_mode], &user_opts) {
                info!("Could not cherry-pick with `-X {}`: {}", whitespace_mode, e);
                return Err(e);
            }
        }
    }

    let desc = git.get_commit_desc(commit_hash)?;
    let (title, body) = make_merge_desc(desc, commit_hash, pr_number, target_branch, orig_base_branch);

    // change commit message
    let mut amend_args = vec![];
    amend_args.extend(user_opts.iter());
    amend_args.extend(["commit", "--amend", "-F", "-"].iter());
    git.run_with_stdin(&amend_args, &format!("{}\n\n{}", &title, &body))?;

    Ok((title, body, whitespace_mode.into()))
}

fn do_cherry_pick(git: &Git, commit_hash: &str, opts: &[&str], user_opts: &[&str]) -> Result<String> {
    git.run(&vec!["reset", "--hard"])?;

    let mut args = vec!["-c", "merge.renameLimit=999999"];
    args.extend(user_opts.iter());
    args.extend(vec!["cherry-pick", "--allow-empty"].iter());

    for i in 0..opts.len() {
        args.push(opts[i]);
    }
    args.push(commit_hash);

    git.run(&args)
}

fn make_merge_desc(
    orig_desc: (String, String),
    commit_hash: &str,
    pr_number: u32,
    target_branch: &str,
    orig_base_branch: &str,
) -> (String, String) {
    // grab original title and strip out the PR number at the end
    let pr_regex = Regex::new(r"(\s*\(#\d+\))+$").unwrap();
    let orig_title = pr_regex.replace(&orig_desc.0, "");
    // strip out 'release' from the prefix to keep titles shorter
    let release_branch_regex = Regex::new(r"^release/").unwrap();
    let title = format!("{}->{}: {}", orig_base_branch, release_branch_regex.replace(target_branch, ""), orig_title);
    let mut body = orig_desc.1;

    if body.len() != 0 {
        body += "\n\n";
    }
    body += format!("(cherry-picked from {}, PR #{})", commit_hash, pr_number).as_str();

    (title, body)
}

#[derive(Debug, PartialEq)]
pub struct PRMergeRequest {
    pub repo: github::Repo,
    pub pull_request: github::PullRequest,
    pub target_branch: String,
}

struct Runner {
    config: Arc<Config>,
    github_app: Arc<GithubSessionFactory>,
    clone_mgr: Arc<GitCloneManager>,
    slack: Arc<worker::Worker<SlackRequest>>,
}

pub fn req(repo: &github::Repo, pull_request: &github::PullRequest, target_branch: &str) -> PRMergeRequest {
    PRMergeRequest {
        repo: repo.clone(),
        pull_request: pull_request.clone(),
        target_branch: target_branch.to_string(),
    }
}

pub fn new_runner(
    config: Arc<Config>,
    github_app: Arc<GithubSessionFactory>,
    clone_mgr: Arc<GitCloneManager>,
    slack: Arc<worker::Worker<SlackRequest>>,
) -> Arc<worker::Runner<PRMergeRequest>> {
    Arc::new(Runner {
        config: config,
        github_app: github_app,
        clone_mgr: clone_mgr.clone(),
        slack: slack,
    })
}

impl worker::Runner<PRMergeRequest> for Runner {
    fn handle(&self, req: PRMergeRequest) {
        let merge_result = clone_and_merge_pull_request(
            self.github_app.borrow(),
            self.clone_mgr.borrow(),
            &req.repo.owner.login(),
            &req.repo.name,
            &req.pull_request,
            &req.target_branch,
        );

        if let Err(e) = merge_result {
            let attach = SlackAttachmentBuilder::new(&format!("{}", e))
                .title(
                    format!("Source PR: #{}: \"{}\"", req.pull_request.number, req.pull_request.title)
                        .as_str(),
                )
                .title_link(req.pull_request.html_url.clone())
                .color("danger")
                .build();

            let messenger = messenger::new(self.config.clone(), self.slack.clone());
            messenger.send_to_owner(
                &format!(
                    "Error creating merge PR from {} to {}",
                    req.pull_request.head.ref_name,
                    req.target_branch
                ),
                &vec![attach],
                &req.pull_request.user,
                &req.repo,
            );
        }
    }
}

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

    #[test]
    fn test_make_merge_desc() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("here is more data about it")),
            "abcdef",
            99,
            "release/target_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->target_branch: Yay, I made a change");
        assert_eq!(desc.1, "here is more data about it\n\n(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_no_body() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "release/target_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->target_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_no_release_branch() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "other_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->other_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }
}
<commit_msg>Strip out previous merge title prefixes (#230)<commit_after>use std::borrow::Borrow;
use std::sync::Arc;

use regex::Regex;

use config::Config;
use errors::*;
use git::Git;
use git_clone_manager::GitCloneManager;
use github;
use github::api::{GithubSessionFactory, Session};
use messenger::{self, Messenger};
use slack::{SlackAttachmentBuilder, SlackRequest};
use worker;

fn clone_and_merge_pull_request(
    github_app: &GithubSessionFactory,
    clone_mgr: &GitCloneManager,
    owner: &str,
    repo: &str,
    pull_request: &github::PullRequest,
    target_branch: &str,
) -> Result<github::PullRequest> {

    let session = github_app.new_session(owner, repo)?;
    let held_clone_dir = clone_mgr.clone(owner, repo)?;
    let clone_dir = held_clone_dir.dir();
    let git = Git::new(session.github_host(), session.github_token(), clone_dir);

    merge_pull_request(&git, &session, owner, repo, pull_request, target_branch)
}

pub fn merge_pull_request(
    git: &Git,
    session: &Session,
    owner: &str,
    repo: &str,
    pull_request: &github::PullRequest,
    target_branch: &str,
) -> Result<github::PullRequest> {
    if !pull_request.is_merged() {
        return Err(format_err!("Pull Request #{} is not yet merged.", pull_request.number));
    }

    let merge_commit_sha;
    if let Some(ref sha) = pull_request.merge_commit_sha {
        merge_commit_sha = sha;
    } else {
        return Err(format_err!("Pull Request #{} has no merge commit.", pull_request.number));
    }

    // strip everything before last slash
    let regex = Regex::new(r".*/").unwrap();
    let pr_branch_name =
        format!("{}-{}", regex.replace(&pull_request.head.ref_name, ""), regex.replace(&target_branch, ""));

    // make sure there isn't already such a branch
    let current_remotes = git.run(&["ls-remote", "--heads"])?;
    if current_remotes.contains(&format!("refs/heads/{}", pr_branch_name)) {
        return Err(format_err!("PR branch already exists on origin: '{}'", pr_branch_name));
    }

    let (title, body, whitespace_mode) = cherry_pick(
        &git,
        &merge_commit_sha,
        &pr_branch_name,
        pull_request.number,
        &target_branch,
        &pull_request.base.ref_name,
    )?;

    git.run(&["push", "origin", &format!("HEAD:{}", pr_branch_name)])?;

    let new_pr = session.create_pull_request(owner, repo, &title, &body, &pr_branch_name, &target_branch)?;

    let mut assignees: Vec<String> = pull_request.assignees.iter().map(|a| a.login().to_string()).collect();
    assignees.retain(|r| r != pull_request.user.login());
    if !assignees.is_empty() {
        session.assign_pull_request(owner, repo, new_pr.number, assignees)?;
    }

    let mut reviewers: Vec<String> = pull_request.all_reviewers().into_iter().map(|a| a.login().to_string()).collect();
    reviewers.retain(|r| r != pull_request.user.login());
    if !reviewers.is_empty() {
        session.request_review(owner, repo, new_pr.number, reviewers)?;
    }

    if whitespace_mode.len() > 0 {
        let msg = format!("Cherry-pick required option `{}`. Please verify correctness.", whitespace_mode);
        if let Err(e) = session.comment_pull_request(owner, repo, new_pr.number, &msg) {
            error!("Error making whitespace comment on pull request: {}", e);
        }
    }

    Ok(new_pr)
}

pub fn cherry_pick(
    git: &Git,
    commit_hash: &str,
    pr_branch_name: &str,
    pr_number: u32,
    target_branch: &str,
    orig_base_branch: &str,
) -> Result<(String, String, String)> {
    git.checkout_branch(pr_branch_name, &format!("origin/{}", target_branch))?;

    let (user, email) = git.get_commit_author(commit_hash)?;
    let email = format!("user.email={}", email);
    let user = format!("user.name={}", user);
    let user_opts = ["-c", &email, "-c", &user];

    // cherry-pick!

    let mut whitespace_mode = "";
    if let Err(e) = do_cherry_pick(git, &commit_hash, &[], &user_opts) {
        info!("Could not cherry-pick normally. Ignoring changed whitespace. {}", e);

        whitespace_mode = "ignore-space-change";
        if let Err(e) = do_cherry_pick(git, &commit_hash, &["-X", whitespace_mode], &user_opts) {
            info!("Could not cherry-pick with `-X {}`. Ignoring all whitespace. {}", whitespace_mode, e);

            whitespace_mode = "ignore-all-space";
            if let Err(e) = do_cherry_pick(git, &commit_hash, &["-X", whitespace_mode], &user_opts) {
                info!("Could not cherry-pick with `-X {}`: {}", whitespace_mode, e);
                return Err(e);
            }
        }
    }

    let desc = git.get_commit_desc(commit_hash)?;
    let (title, body) = make_merge_desc(desc, commit_hash, pr_number, target_branch, orig_base_branch);

    // change commit message
    let mut amend_args = vec![];
    amend_args.extend(user_opts.iter());
    amend_args.extend(["commit", "--amend", "-F", "-"].iter());
    git.run_with_stdin(&amend_args, &format!("{}\n\n{}", &title, &body))?;

    Ok((title, body, whitespace_mode.into()))
}

fn do_cherry_pick(git: &Git, commit_hash: &str, opts: &[&str], user_opts: &[&str]) -> Result<String> {
    git.run(&vec!["reset", "--hard"])?;

    let mut args = vec!["-c", "merge.renameLimit=999999"];
    args.extend(user_opts.iter());
    args.extend(vec!["cherry-pick", "--allow-empty"].iter());

    for i in 0..opts.len() {
        args.push(opts[i]);
    }
    args.push(commit_hash);

    git.run(&args)
}

fn make_merge_desc(
    orig_desc: (String, String),
    commit_hash: &str,
    pr_number: u32,
    target_branch: &str,
    orig_base_branch: &str,
) -> (String, String) {
    // grab original title and strip out the PR number at the end
    let pr_regex = Regex::new(r"(\s*\(#\d+\))+$").unwrap();
    let prev_merge_regex = Regex::new(r"^([^:]+->[^:]+: )+").unwrap();

    // strip out PR from title
    let orig_title = pr_regex.replace(&orig_desc.0, "");
    // strip out previous merge title prefixes
    let orig_title = prev_merge_regex.replace(&orig_title, "");

    // strip out 'release' from the prefix to keep titles shorter
    let release_branch_regex = Regex::new(r"^release/").unwrap();
    let title = format!("{}->{}: {}", orig_base_branch, release_branch_regex.replace(target_branch, ""), orig_title);
    let mut body = orig_desc.1;

    if body.len() != 0 {
        body += "\n\n";
    }
    body += format!("(cherry-picked from {}, PR #{})", commit_hash, pr_number).as_str();

    (title, body)
}

#[derive(Debug, PartialEq)]
pub struct PRMergeRequest {
    pub repo: github::Repo,
    pub pull_request: github::PullRequest,
    pub target_branch: String,
}

struct Runner {
    config: Arc<Config>,
    github_app: Arc<GithubSessionFactory>,
    clone_mgr: Arc<GitCloneManager>,
    slack: Arc<worker::Worker<SlackRequest>>,
}

pub fn req(repo: &github::Repo, pull_request: &github::PullRequest, target_branch: &str) -> PRMergeRequest {
    PRMergeRequest {
        repo: repo.clone(),
        pull_request: pull_request.clone(),
        target_branch: target_branch.to_string(),
    }
}

pub fn new_runner(
    config: Arc<Config>,
    github_app: Arc<GithubSessionFactory>,
    clone_mgr: Arc<GitCloneManager>,
    slack: Arc<worker::Worker<SlackRequest>>,
) -> Arc<worker::Runner<PRMergeRequest>> {
    Arc::new(Runner {
        config: config,
        github_app: github_app,
        clone_mgr: clone_mgr.clone(),
        slack: slack,
    })
}

impl worker::Runner<PRMergeRequest> for Runner {
    fn handle(&self, req: PRMergeRequest) {
        let merge_result = clone_and_merge_pull_request(
            self.github_app.borrow(),
            self.clone_mgr.borrow(),
            &req.repo.owner.login(),
            &req.repo.name,
            &req.pull_request,
            &req.target_branch,
        );

        if let Err(e) = merge_result {
            let attach = SlackAttachmentBuilder::new(&format!("{}", e))
                .title(
                    format!("Source PR: #{}: \"{}\"", req.pull_request.number, req.pull_request.title)
                        .as_str(),
                )
                .title_link(req.pull_request.html_url.clone())
                .color("danger")
                .build();

            let messenger = messenger::new(self.config.clone(), self.slack.clone());
            messenger.send_to_owner(
                &format!(
                    "Error creating merge PR from {} to {}",
                    req.pull_request.head.ref_name,
                    req.target_branch
                ),
                &vec![attach],
                &req.pull_request.user,
                &req.repo,
            );
        }
    }
}

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

    #[test]
    fn test_make_merge_desc() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("here is more data about it")),
            "abcdef",
            99,
            "release/target_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->target_branch: Yay, I made a change");
        assert_eq!(desc.1, "here is more data about it\n\n(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_no_body() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "release/target_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->target_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_no_release_branch() {
        let desc = make_merge_desc(
            (String::from("Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "other_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->other_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_multi1() {
        let desc = make_merge_desc(
            (String::from("prev_branch->source_branch: Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "other_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->other_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }

    #[test]
    fn test_make_merge_desc_multi2() {
        let desc = make_merge_desc(
            (String::from("prev_branch->source_branch: more_branches->prev_branch: Yay, I made a change (#99)"), String::from("")),
            "abcdef",
            99,
            "other_branch",
            "source_branch",
        );

        assert_eq!(desc.0, "source_branch->other_branch: Yay, I made a change");
        assert_eq!(desc.1, "(cherry-picked from abcdef, PR #99)");
    }
}
<|endoftext|>


<commit_before>use std::fmt;

#[derive(PartialEq, Debug, Clone)]
/// A Term can represent a Terminal or Nonterminal node
pub enum Term {
    Terminal(String),
    Nonterminal(String),
}

impl fmt::Display for Term {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Term::Terminal(ref s) => write!(f, "\"{}\"", s),
            Term::Nonterminal(ref s) => write!(f, "<{}>", s),
        }
    }
}

#[derive(PartialEq, Debug, Clone)]
/// An Expression is comprised of any number of Terms
pub struct Expression {
    pub terms: Vec<Term>,
}

impl Expression {
    #[allow(dead_code)]
    pub fn new() -> Expression {
        Expression { terms: vec![] }
    }

    #[allow(dead_code)]
    pub fn from_parts(v: Vec<Term>) -> Expression {
        Expression { terms: v }
    }
}

impl fmt::Display for Expression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let display = self.terms
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<_>>()
            .join(" ");

        write!(f, "{}", display)
    }
}

#[derive(PartialEq, Debug, Clone)]
/// A Production is comprised of any number of Expressions
pub struct Production {
    pub lhs: Term,
    pub rhs: Vec<Expression>,
}

impl Production {
    #[allow(dead_code)]
    pub fn new() -> Production {
        Production {
            lhs: Term::Nonterminal(String::new()),
            rhs: vec![],
        }
    }

    #[allow(dead_code)]
    pub fn from_parts(t: Term, e: Vec<Expression>) -> Production {
        Production { lhs: t, rhs: e }
    }
}

impl fmt::Display for Production {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} ::= {};",
            self.lhs.to_string(),
            self.rhs
                .iter()
                .map(|s| s.to_string())
                .collect::<Vec<_>>()
                .join(" | ")
        )
    }
}

#[derive(PartialEq, Debug, Clone)]
/// A Grammar is comprised of any number of Productions
pub struct Grammar {
    pub productions: Vec<Production>,
}

impl Grammar {
    pub fn new() -> Grammar {
        Grammar {
            productions: vec![],
        }
    }

    #[allow(dead_code)]
    pub fn from_parts(v: Vec<Production>) -> Grammar {
        Grammar { productions: v }
    }
}

impl fmt::Display for Grammar {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}\n",
            self.productions
                .iter()
                .map(|s| s.to_string())
                .collect::<Vec<_>>()
                .join("\n")
        )
    }
}
<commit_msg>Remove dead code allowances<commit_after>use std::fmt;

#[derive(PartialEq, Debug, Clone)]
/// A Term can represent a Terminal or Nonterminal node
pub enum Term {
    Terminal(String),
    Nonterminal(String),
}

impl fmt::Display for Term {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Term::Terminal(ref s) => write!(f, "\"{}\"", s),
            Term::Nonterminal(ref s) => write!(f, "<{}>", s),
        }
    }
}

#[derive(PartialEq, Debug, Clone)]
/// An Expression is comprised of any number of Terms
pub struct Expression {
    pub terms: Vec<Term>,
}

impl Expression {
    pub fn new() -> Expression {
        Expression { terms: vec![] }
    }

    pub fn from_parts(v: Vec<Term>) -> Expression {
        Expression { terms: v }
    }
}

impl fmt::Display for Expression {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let display = self.terms
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<_>>()
            .join(" ");

        write!(f, "{}", display)
    }
}

#[derive(PartialEq, Debug, Clone)]
/// A Production is comprised of any number of Expressions
pub struct Production {
    pub lhs: Term,
    pub rhs: Vec<Expression>,
}

impl Production {
    pub fn new() -> Production {
        Production {
            lhs: Term::Nonterminal(String::new()),
            rhs: vec![],
        }
    }

    pub fn from_parts(t: Term, e: Vec<Expression>) -> Production {
        Production { lhs: t, rhs: e }
    }
}

impl fmt::Display for Production {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{} ::= {};",
            self.lhs.to_string(),
            self.rhs
                .iter()
                .map(|s| s.to_string())
                .collect::<Vec<_>>()
                .join(" | ")
        )
    }
}

#[derive(PartialEq, Debug, Clone)]
/// A Grammar is comprised of any number of Productions
pub struct Grammar {
    pub productions: Vec<Production>,
}

impl Grammar {
    pub fn new() -> Grammar {
        Grammar {
            productions: vec![],
        }
    }

    pub fn from_parts(v: Vec<Production>) -> Grammar {
        Grammar { productions: v }
    }
}

impl fmt::Display for Grammar {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}\n",
            self.productions
                .iter()
                .map(|s| s.to_string())
                .collect::<Vec<_>>()
                .join("\n")
        )
    }
}
<|endoftext|>