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

use num_complex::Complex;
use common::{FFTnum, verify_length, verify_length_divisible};

use math_utils;
use array_utils;

use ::{Length, IsInverse, FFT};
use algorithm::butterflies::FFTButterfly;

/// Implementation of the [Good-Thomas Algorithm (AKA Prime Factor Algorithm)](https://en.wikipedia.org/wiki/Prime-factor_FFT_algorithm)
///
/// This algorithm factors a size n FFT into n1 * n2, where GCD(n1, n2) == 1
///
/// Conceptually, this algorithm is very similar to the Mixed-Radix FFT, except because GCD(n1, n2) == 1 we can do some
/// number theory trickery to reduce the number of floating-point multiplications and additions. Additionally, It can
/// be faster than Mixed-Radix at sizes below 10,000 or so.
///
/// ~~~
/// // Computes a forward FFT of size 1200, using the Good-Thomas Algorithm
/// use rustfft::algorithm::GoodThomasAlgorithm;
/// use rustfft::{FFT, FFTplanner};
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 1200];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 1200];
///
/// // we need to find an n1 and n2 such that n1 * n2 == 1200 and GCD(n1, n2) == 1
/// // n1 = 48 and n2 = 25 satisfies this
/// let mut planner = FFTplanner::new(false);
/// let inner_fft_n1 = planner.plan_fft(48);
/// let inner_fft_n2 = planner.plan_fft(25);
///
/// // the good-thomas FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 1200
/// let fft = GoodThomasAlgorithm::new(inner_fft_n1, inner_fft_n2);
/// fft.process(&mut input, &mut output);
/// ~~~
pub struct GoodThomasAlgorithm<T> {
    width: usize,
    width_size_fft: Arc<FFT<T>>,

    height: usize,
    height_size_fft: Arc<FFT<T>>,

    input_output_map: Box<[usize]>,

    inverse: bool,
}

impl<T: FFTnum> GoodThomasAlgorithm<T> {
    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
    ///
    /// GCD(width_fft.len(), height_fft.len()) must be equal to 1
    pub fn new(width_fft: Arc<FFT<T>>, height_fft: Arc<FFT<T>>) -> Self {
        assert_eq!(
            width_fft.is_inverse(), height_fft.is_inverse(),
            "width_fft and height_fft must both be inverse, or neither. got width inverse={}, height inverse={}",
            width_fft.is_inverse(), height_fft.is_inverse());

        let width = width_fft.len();
        let height = height_fft.len();
        let len = width * height;

        // compute the nultiplicative inverse of width mod height and vice versa
        let (gcd, mut width_inverse, mut height_inverse) =
            math_utils::extended_euclidean_algorithm(width as i64, height as i64);
        assert!(gcd == 1,
                "Invalid input width and height to Good-Thomas Algorithm: ({},{}): Inputs must be coprime",
                width,
                height);

        // width_inverse or height_inverse might be negative, make it positive
        if width_inverse < 0 {
            width_inverse += height as i64;
        }
        if height_inverse < 0 {
            height_inverse += width as i64;
        }

        // NOTE: we are precomputing the input and output reordering indexes, because benchmarking shows that it's 10-20% faster
        // If we wanted to optimize for memory use or setup time instead of multiple-FFT speed, we could compute these on the fly in the perform_fft() method
        let input_iter = (0..len)
                .map(|i| (i % width, i / width))
                .map(|(x, y)| (x * height + y * width) % len);
        let output_iter = (0..len)
                .map(|i| (i % height, i / height))
                .map(|(y, x)| (x * height * height_inverse as usize + y * width * width_inverse as usize) % len);

        let input_output_map: Vec<usize> = input_iter.chain(output_iter).collect();

        GoodThomasAlgorithm {
            inverse: width_fft.is_inverse(),

            width: width,
            width_size_fft: width_fft,

            height: height,
            height_size_fft: height_fft,

            input_output_map: input_output_map.into_boxed_slice(),
        }
    }

    fn perform_fft(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        let (input_map, output_map) = self.input_output_map.split_at(self.len());

        // copy the input into the output buffer
        for (output_element, &input_index) in output.iter_mut().zip(input_map.iter()) {
            *output_element = input[input_index];
        }

        // run FFTs of size `width`
        self.width_size_fft.process_multi(output, input);

        // transpose
        array_utils::transpose(self.width, self.height, input, output);

        // run FFTs of size 'height'
        self.height_size_fft.process_multi(output, input);

        // copy to the output, using our output redordeing mapping
        for (input_element, &output_index) in input.iter().zip(output_map.iter()) {
            output[output_index] = *input_element;
        }
    }
}

impl<T: FFTnum> FFT<T> for GoodThomasAlgorithm<T> {
    fn process(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length(input, output, self.len());

        self.perform_fft(input, output);
    }
    fn process_multi(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length_divisible(input, output, self.len());

        for (in_chunk, out_chunk) in input.chunks_mut(self.len()).zip(output.chunks_mut(self.len())) {
            self.perform_fft(in_chunk, out_chunk);
        }
    }
}
impl<T> Length for GoodThomasAlgorithm<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.width * self.height
    }
}
impl<T> IsInverse for GoodThomasAlgorithm<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inverse
    }
}

/// Implementation of the Good-Thomas Algorithm, specialized for the case where both inner FFTs are butterflies
///
/// This algorithm factors a size n FFT into n1 * n2, where GCD(n1, n2) == 1
///
/// Conceptually, this algorithm is very similar to the Mixed-Radix FFT, except because GCD(n1, n2) == 1 we can do some
/// number theory trickery to reduce the number of floating-point multiplications and additions. It typically performs
/// better than Mixed-Radix Double Butterfly Algorithm, especially at small sizes.
///
/// ~~~
/// // Computes a forward FFT of size 56, using the Good-Thoma Butterfly Algorithm
/// use std::sync::Arc;
/// use rustfft::algorithm::GoodThomasAlgorithmDoubleButterfly;
/// use rustfft::algorithm::butterflies::{Butterfly7, Butterfly8};
/// use rustfft::FFT;
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 56];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 56];
///
/// // we need to find an n1 and n2 such that n1 * n2 == 56 and GCD(n1, n2) == 1
/// // n1 = 7 and n2 = 8 satisfies this
/// let inner_fft_n1 = Arc::new(Butterfly7::new(false));
/// let inner_fft_n2 = Arc::new(Butterfly8::new(false));
///
/// // the good-thomas FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 56
/// let fft = GoodThomasAlgorithmDoubleButterfly::new(inner_fft_n1, inner_fft_n2);
/// fft.process(&mut input, &mut output);
/// ~~~
pub struct GoodThomasAlgorithmDoubleButterfly<T> {
    width: usize,
    width_size_fft: Arc<FFTButterfly<T>>,

    height: usize,
    height_size_fft: Arc<FFTButterfly<T>>,

    input_output_map: Box<[usize]>,

    inverse: bool,
}

impl<T: FFTnum> GoodThomasAlgorithmDoubleButterfly<T> {
    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
    ///
    /// GCD(n1.len(), n2.len()) must be equal to 1
    pub fn new(width_fft: Arc<FFTButterfly<T>>, height_fft: Arc<FFTButterfly<T>>) -> Self {
        assert_eq!(
            width_fft.is_inverse(), height_fft.is_inverse(),
            "n1_fft and height_fft must both be inverse, or neither. got width inverse={}, height inverse={}",
            width_fft.is_inverse(), height_fft.is_inverse());

        let width = width_fft.len();
        let height = height_fft.len();
        let len = width * height;

        // compute the nultiplicative inverse of n1 mod height and vice versa
        let (gcd, mut width_inverse, mut height_inverse) =
            math_utils::extended_euclidean_algorithm(width as i64, height as i64);
        assert!(gcd == 1,
                "Invalid input n1 and height to Good-Thomas Algorithm: ({},{}): Inputs must be coprime",
                width,
                height);

        // width_inverse or height_inverse might be negative, make it positive
        if width_inverse < 0 {
            width_inverse += height as i64;
        }
        if height_inverse < 0 {
            height_inverse += width as i64;
        }

        // NOTE: we are precomputing the input and output reordering indexes, because benchmarking shows that it's 10-20% faster
        // If we wanted to optimize for memory use or setup time instead of multiple-FFT speed, we could compute these on the fly in the perform_fft() method
        let input_iter = (0..len)
                .map(|i| (i % width, i / width))
                .map(|(x, y)| (x * height + y * width) % len);
        let output_iter = (0..len)
                .map(|i| (i % height, i / height))
                .map(|(y, x)| (x * height * height_inverse as usize + y * width * width_inverse as usize) % len);

        let input_output_map: Vec<usize> = input_iter.chain(output_iter).collect();

        GoodThomasAlgorithmDoubleButterfly {
            inverse: width_fft.is_inverse(),

            width: width,
            width_size_fft: width_fft,

            height: height,
            height_size_fft: height_fft,

            input_output_map: input_output_map.into_boxed_slice(),
        }
    }

    unsafe fn perform_fft(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {

        let (input_map, output_map) = self.input_output_map.split_at(self.len());

        // copy the input using our reordering mapping
        for (output_element, &input_index) in output.iter_mut().zip(input_map.iter()) {
            *output_element = input[input_index];
        }

        // run FFTs of size `width`
        self.width_size_fft.process_multi_inplace(output);

        // transpose
        array_utils::transpose_small(self.width, self.height, output, input);

        // run FFTs of size 'height'
        self.height_size_fft.process_multi_inplace(input);

        // copy to the output, using our output redordeing mapping
        for (input_element, &output_index) in input.iter().zip(output_map.iter()) {
            output[output_index] = *input_element;
        }
    }
}

impl<T: FFTnum> FFT<T> for GoodThomasAlgorithmDoubleButterfly<T> {
    fn process(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length(input, output, self.len());

        unsafe { self.perform_fft(input, output) };
    }
    fn process_multi(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length_divisible(input, output, self.len());

        for (in_chunk, out_chunk) in input.chunks_mut(self.len()).zip(output.chunks_mut(self.len())) {
             unsafe { self.perform_fft(in_chunk, out_chunk) };
        }
    }
}
impl<T> Length for GoodThomasAlgorithmDoubleButterfly<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.width * self.height
    }
}
impl<T> IsInverse for GoodThomasAlgorithmDoubleButterfly<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inverse
    }
}

#[cfg(test)]
mod unit_tests {
    use super::*;
    use std::sync::Arc;
    use test_utils::{check_fft_algorithm, make_butterfly};
    use algorithm::DFT;
    use num_integer::gcd;

    #[test]
    fn test_good_thomas() {
        for width in 2..20 {
            for height in 2..20 {
                if gcd(width, height) == 1 {
                    test_good_thomas_with_lengths(width, height, false);
                    test_good_thomas_with_lengths(width, height, true);
                }
            }
        }

        //verify that it works correctly when width and/or height are 1
        test_good_thomas_with_lengths(1, 10, false);
        test_good_thomas_with_lengths(10, 1, false);
        test_good_thomas_with_lengths(1, 1, false);
    }

    #[test]
    fn test_good_thomas_double_butterfly() {
        let butterfly_sizes = [2,3,4,5,6,7,8,16];
        for width in &butterfly_sizes {
            for height in &butterfly_sizes {
                if gcd(*width, *height) == 1 {
                    test_good_thomas_butterfly_with_lengths(*width, *height, false);
                    test_good_thomas_butterfly_with_lengths(*width, *height, true);
                }
            }
        }
    }

    fn test_good_thomas_with_lengths(width: usize, height: usize, inverse: bool) {
        let width_fft = Arc::new(DFT::new(width, inverse)) as Arc<FFT<f32>>;
        let height_fft = Arc::new(DFT::new(height, inverse)) as Arc<FFT<f32>>;

        let fft = GoodThomasAlgorithm::new(width_fft, height_fft);

        check_fft_algorithm(&fft, width * height, inverse);
    }

    fn test_good_thomas_butterfly_with_lengths(width: usize, height: usize, inverse: bool) {
        let width_fft = make_butterfly(width, inverse);
        let height_fft = make_butterfly(height, inverse);

        let fft = GoodThomasAlgorithmDoubleButterfly::new(width_fft, height_fft);

        check_fft_algorithm(&fft, width * height, inverse);
    }
}
<commit_msg>Simplified and sped up good-thomas test<commit_after>use std::sync::Arc;

use num_complex::Complex;
use common::{FFTnum, verify_length, verify_length_divisible};

use math_utils;
use array_utils;

use ::{Length, IsInverse, FFT};
use algorithm::butterflies::FFTButterfly;

/// Implementation of the [Good-Thomas Algorithm (AKA Prime Factor Algorithm)](https://en.wikipedia.org/wiki/Prime-factor_FFT_algorithm)
///
/// This algorithm factors a size n FFT into n1 * n2, where GCD(n1, n2) == 1
///
/// Conceptually, this algorithm is very similar to the Mixed-Radix FFT, except because GCD(n1, n2) == 1 we can do some
/// number theory trickery to reduce the number of floating-point multiplications and additions. Additionally, It can
/// be faster than Mixed-Radix at sizes below 10,000 or so.
///
/// ~~~
/// // Computes a forward FFT of size 1200, using the Good-Thomas Algorithm
/// use rustfft::algorithm::GoodThomasAlgorithm;
/// use rustfft::{FFT, FFTplanner};
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 1200];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 1200];
///
/// // we need to find an n1 and n2 such that n1 * n2 == 1200 and GCD(n1, n2) == 1
/// // n1 = 48 and n2 = 25 satisfies this
/// let mut planner = FFTplanner::new(false);
/// let inner_fft_n1 = planner.plan_fft(48);
/// let inner_fft_n2 = planner.plan_fft(25);
///
/// // the good-thomas FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 1200
/// let fft = GoodThomasAlgorithm::new(inner_fft_n1, inner_fft_n2);
/// fft.process(&mut input, &mut output);
/// ~~~
pub struct GoodThomasAlgorithm<T> {
    width: usize,
    width_size_fft: Arc<FFT<T>>,

    height: usize,
    height_size_fft: Arc<FFT<T>>,

    input_output_map: Box<[usize]>,

    inverse: bool,
}

impl<T: FFTnum> GoodThomasAlgorithm<T> {
    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
    ///
    /// GCD(width_fft.len(), height_fft.len()) must be equal to 1
    pub fn new(width_fft: Arc<FFT<T>>, height_fft: Arc<FFT<T>>) -> Self {
        assert_eq!(
            width_fft.is_inverse(), height_fft.is_inverse(),
            "width_fft and height_fft must both be inverse, or neither. got width inverse={}, height inverse={}",
            width_fft.is_inverse(), height_fft.is_inverse());

        let width = width_fft.len();
        let height = height_fft.len();
        let len = width * height;

        // compute the nultiplicative inverse of width mod height and vice versa
        let (gcd, mut width_inverse, mut height_inverse) =
            math_utils::extended_euclidean_algorithm(width as i64, height as i64);
        assert!(gcd == 1,
                "Invalid input width and height to Good-Thomas Algorithm: ({},{}): Inputs must be coprime",
                width,
                height);

        // width_inverse or height_inverse might be negative, make it positive
        if width_inverse < 0 {
            width_inverse += height as i64;
        }
        if height_inverse < 0 {
            height_inverse += width as i64;
        }

        // NOTE: we are precomputing the input and output reordering indexes, because benchmarking shows that it's 10-20% faster
        // If we wanted to optimize for memory use or setup time instead of multiple-FFT speed, we could compute these on the fly in the perform_fft() method
        let input_iter = (0..len)
                .map(|i| (i % width, i / width))
                .map(|(x, y)| (x * height + y * width) % len);
        let output_iter = (0..len)
                .map(|i| (i % height, i / height))
                .map(|(y, x)| (x * height * height_inverse as usize + y * width * width_inverse as usize) % len);

        let input_output_map: Vec<usize> = input_iter.chain(output_iter).collect();

        GoodThomasAlgorithm {
            inverse: width_fft.is_inverse(),

            width: width,
            width_size_fft: width_fft,

            height: height,
            height_size_fft: height_fft,

            input_output_map: input_output_map.into_boxed_slice(),
        }
    }

    fn perform_fft(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        let (input_map, output_map) = self.input_output_map.split_at(self.len());

        // copy the input into the output buffer
        for (output_element, &input_index) in output.iter_mut().zip(input_map.iter()) {
            *output_element = input[input_index];
        }

        // run FFTs of size `width`
        self.width_size_fft.process_multi(output, input);

        // transpose
        array_utils::transpose(self.width, self.height, input, output);

        // run FFTs of size 'height'
        self.height_size_fft.process_multi(output, input);

        // copy to the output, using our output redordeing mapping
        for (input_element, &output_index) in input.iter().zip(output_map.iter()) {
            output[output_index] = *input_element;
        }
    }
}

impl<T: FFTnum> FFT<T> for GoodThomasAlgorithm<T> {
    fn process(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length(input, output, self.len());

        self.perform_fft(input, output);
    }
    fn process_multi(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length_divisible(input, output, self.len());

        for (in_chunk, out_chunk) in input.chunks_mut(self.len()).zip(output.chunks_mut(self.len())) {
            self.perform_fft(in_chunk, out_chunk);
        }
    }
}
impl<T> Length for GoodThomasAlgorithm<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.width * self.height
    }
}
impl<T> IsInverse for GoodThomasAlgorithm<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inverse
    }
}

/// Implementation of the Good-Thomas Algorithm, specialized for the case where both inner FFTs are butterflies
///
/// This algorithm factors a size n FFT into n1 * n2, where GCD(n1, n2) == 1
///
/// Conceptually, this algorithm is very similar to the Mixed-Radix FFT, except because GCD(n1, n2) == 1 we can do some
/// number theory trickery to reduce the number of floating-point multiplications and additions. It typically performs
/// better than Mixed-Radix Double Butterfly Algorithm, especially at small sizes.
///
/// ~~~
/// // Computes a forward FFT of size 56, using the Good-Thoma Butterfly Algorithm
/// use std::sync::Arc;
/// use rustfft::algorithm::GoodThomasAlgorithmDoubleButterfly;
/// use rustfft::algorithm::butterflies::{Butterfly7, Butterfly8};
/// use rustfft::FFT;
/// use rustfft::num_complex::Complex;
/// use rustfft::num_traits::Zero;
///
/// let mut input:  Vec<Complex<f32>> = vec![Zero::zero(); 56];
/// let mut output: Vec<Complex<f32>> = vec![Zero::zero(); 56];
///
/// // we need to find an n1 and n2 such that n1 * n2 == 56 and GCD(n1, n2) == 1
/// // n1 = 7 and n2 = 8 satisfies this
/// let inner_fft_n1 = Arc::new(Butterfly7::new(false));
/// let inner_fft_n2 = Arc::new(Butterfly8::new(false));
///
/// // the good-thomas FFT length will be inner_fft_n1.len() * inner_fft_n2.len() = 56
/// let fft = GoodThomasAlgorithmDoubleButterfly::new(inner_fft_n1, inner_fft_n2);
/// fft.process(&mut input, &mut output);
/// ~~~
pub struct GoodThomasAlgorithmDoubleButterfly<T> {
    width: usize,
    width_size_fft: Arc<FFTButterfly<T>>,

    height: usize,
    height_size_fft: Arc<FFTButterfly<T>>,

    input_output_map: Box<[usize]>,

    inverse: bool,
}

impl<T: FFTnum> GoodThomasAlgorithmDoubleButterfly<T> {
    /// Creates a FFT instance which will process inputs/outputs of size `width_fft.len() * height_fft.len()`
    ///
    /// GCD(n1.len(), n2.len()) must be equal to 1
    pub fn new(width_fft: Arc<FFTButterfly<T>>, height_fft: Arc<FFTButterfly<T>>) -> Self {
        assert_eq!(
            width_fft.is_inverse(), height_fft.is_inverse(),
            "n1_fft and height_fft must both be inverse, or neither. got width inverse={}, height inverse={}",
            width_fft.is_inverse(), height_fft.is_inverse());

        let width = width_fft.len();
        let height = height_fft.len();
        let len = width * height;

        // compute the nultiplicative inverse of n1 mod height and vice versa
        let (gcd, mut width_inverse, mut height_inverse) =
            math_utils::extended_euclidean_algorithm(width as i64, height as i64);
        assert!(gcd == 1,
                "Invalid input n1 and height to Good-Thomas Algorithm: ({},{}): Inputs must be coprime",
                width,
                height);

        // width_inverse or height_inverse might be negative, make it positive
        if width_inverse < 0 {
            width_inverse += height as i64;
        }
        if height_inverse < 0 {
            height_inverse += width as i64;
        }

        // NOTE: we are precomputing the input and output reordering indexes, because benchmarking shows that it's 10-20% faster
        // If we wanted to optimize for memory use or setup time instead of multiple-FFT speed, we could compute these on the fly in the perform_fft() method
        let input_iter = (0..len)
                .map(|i| (i % width, i / width))
                .map(|(x, y)| (x * height + y * width) % len);
        let output_iter = (0..len)
                .map(|i| (i % height, i / height))
                .map(|(y, x)| (x * height * height_inverse as usize + y * width * width_inverse as usize) % len);

        let input_output_map: Vec<usize> = input_iter.chain(output_iter).collect();

        GoodThomasAlgorithmDoubleButterfly {
            inverse: width_fft.is_inverse(),

            width: width,
            width_size_fft: width_fft,

            height: height,
            height_size_fft: height_fft,

            input_output_map: input_output_map.into_boxed_slice(),
        }
    }

    unsafe fn perform_fft(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {

        let (input_map, output_map) = self.input_output_map.split_at(self.len());

        // copy the input using our reordering mapping
        for (output_element, &input_index) in output.iter_mut().zip(input_map.iter()) {
            *output_element = input[input_index];
        }

        // run FFTs of size `width`
        self.width_size_fft.process_multi_inplace(output);

        // transpose
        array_utils::transpose_small(self.width, self.height, output, input);

        // run FFTs of size 'height'
        self.height_size_fft.process_multi_inplace(input);

        // copy to the output, using our output redordeing mapping
        for (input_element, &output_index) in input.iter().zip(output_map.iter()) {
            output[output_index] = *input_element;
        }
    }
}

impl<T: FFTnum> FFT<T> for GoodThomasAlgorithmDoubleButterfly<T> {
    fn process(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length(input, output, self.len());

        unsafe { self.perform_fft(input, output) };
    }
    fn process_multi(&self, input: &mut [Complex<T>], output: &mut [Complex<T>]) {
        verify_length_divisible(input, output, self.len());

        for (in_chunk, out_chunk) in input.chunks_mut(self.len()).zip(output.chunks_mut(self.len())) {
             unsafe { self.perform_fft(in_chunk, out_chunk) };
        }
    }
}
impl<T> Length for GoodThomasAlgorithmDoubleButterfly<T> {
    #[inline(always)]
    fn len(&self) -> usize {
        self.width * self.height
    }
}
impl<T> IsInverse for GoodThomasAlgorithmDoubleButterfly<T> {
    #[inline(always)]
    fn is_inverse(&self) -> bool {
        self.inverse
    }
}

#[cfg(test)]
mod unit_tests {
    use super::*;
    use std::sync::Arc;
    use test_utils::{check_fft_algorithm, make_butterfly};
    use algorithm::DFT;
    use num_integer::gcd;

    #[test]
    fn test_good_thomas() {
        for width in 1..12 {
            for height in 1..12 {
                if gcd(width, height) == 1 {
                    test_good_thomas_with_lengths(width, height, false);
                    test_good_thomas_with_lengths(width, height, true);
                }
            }
        }
    }

    #[test]
    fn test_good_thomas_double_butterfly() {
        let butterfly_sizes = [2,3,4,5,6,7,8,16];
        for width in &butterfly_sizes {
            for height in &butterfly_sizes {
                if gcd(*width, *height) == 1 {
                    test_good_thomas_butterfly_with_lengths(*width, *height, false);
                    test_good_thomas_butterfly_with_lengths(*width, *height, true);
                }
            }
        }
    }

    fn test_good_thomas_with_lengths(width: usize, height: usize, inverse: bool) {
        let width_fft = Arc::new(DFT::new(width, inverse)) as Arc<FFT<f32>>;
        let height_fft = Arc::new(DFT::new(height, inverse)) as Arc<FFT<f32>>;

        let fft = GoodThomasAlgorithm::new(width_fft, height_fft);

        check_fft_algorithm(&fft, width * height, inverse);
    }

    fn test_good_thomas_butterfly_with_lengths(width: usize, height: usize, inverse: bool) {
        let width_fft = make_butterfly(width, inverse);
        let height_fft = make_butterfly(height, inverse);

        let fft = GoodThomasAlgorithmDoubleButterfly::new(width_fft, height_fft);

        check_fft_algorithm(&fft, width * height, inverse);
    }
}
<|endoftext|>


<commit_before>use crate::check::{FnCtxt, Needs, PlaceOp};
use rustc_hir as hir;
use rustc_middle::ty::adjustment::{Adjust, Adjustment, OverloadedDeref, PointerCast};
use rustc_middle::ty::adjustment::{AllowTwoPhase, AutoBorrow, AutoBorrowMutability};
use rustc_middle::ty::{self, Ty};

impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
    /// Convert auto-derefs, indices, etc of an expression from `Deref` and `Index`
    /// into `DerefMut` and `IndexMut` respectively.
    ///
    /// This is a second pass of typechecking derefs/indices. We need this we do not
    /// always know whether a place needs to be mutable or not in the first pass.
    /// This happens whether there is an implicit mutable reborrow, e.g. when the type
    /// is used as the receiver of a method call.
    pub fn convert_place_derefs_to_mutable(&self, expr: &hir::Expr<'_>) {
        // Gather up expressions we want to munge.
        let mut exprs = vec![expr];

        loop {
            match exprs.last().unwrap().kind {
                hir::ExprKind::Field(ref expr, _)
                | hir::ExprKind::Index(ref expr, _)
                | hir::ExprKind::Unary(hir::UnOp::UnDeref, ref expr) => exprs.push(&expr),
                _ => break,
            }
        }

        debug!("convert_place_derefs_to_mutable: exprs={:?}", exprs);

        // Fix up autoderefs and derefs.
        for (i, &expr) in exprs.iter().rev().enumerate() {
            debug!("convert_place_derefs_to_mutable: i={} expr={:?}", i, expr);

            // Fix up the autoderefs. Autorefs can only occur immediately preceding
            // overloaded place ops, and will be fixed by them in order to get
            // the correct region.
            let mut source = self.node_ty(expr.hir_id);
            // Do not mutate adjustments in place, but rather take them,
            // and replace them after mutating them, to avoid having the
            // tables borrowed during (`deref_mut`) method resolution.
            let previous_adjustments =
                self.tables.borrow_mut().adjustments_mut().remove(expr.hir_id);
            if let Some(mut adjustments) = previous_adjustments {
                let needs = Needs::MutPlace;
                for adjustment in &mut adjustments {
                    if let Adjust::Deref(Some(ref mut deref)) = adjustment.kind {
                        if let Some(ok) = self.try_overloaded_deref(expr.span, source, needs) {
                            let method = self.register_infer_ok_obligations(ok);
                            if let ty::Ref(region, _, mutbl) = method.sig.output().kind {
                                *deref = OverloadedDeref { region, mutbl };
                            }
                        }
                    }
                    source = adjustment.target;
                }
                self.tables.borrow_mut().adjustments_mut().insert(expr.hir_id, adjustments);
            }

            match expr.kind {
                hir::ExprKind::Index(ref base_expr, ref index_expr) => {
                    // We need to get the final type in case dereferences were needed for the trait
                    // to apply (#72002).
                    let index_expr_ty = self.tables.borrow().expr_ty_adjusted(index_expr);
                    self.convert_place_op_to_mutable(
                        PlaceOp::Index,
                        expr,
                        base_expr,
                        &[index_expr_ty],
                    );
                }
                hir::ExprKind::Unary(hir::UnOp::UnDeref, ref base_expr) => {
                    self.convert_place_op_to_mutable(PlaceOp::Deref, expr, base_expr, &[]);
                }
                _ => {}
            }
        }
    }

    fn convert_place_op_to_mutable(
        &self,
        op: PlaceOp,
        expr: &hir::Expr<'_>,
        base_expr: &hir::Expr<'_>,
        arg_tys: &[Ty<'tcx>],
    ) {
        debug!("convert_place_op_to_mutable({:?}, {:?}, {:?}, {:?})", op, expr, base_expr, arg_tys);
        if !self.tables.borrow().is_method_call(expr) {
            debug!("convert_place_op_to_mutable - builtin, nothing to do");
            return;
        }

        let base_ty = self
            .tables
            .borrow()
            .expr_adjustments(base_expr)
            .last()
            .map_or_else(|| self.node_ty(expr.hir_id), |adj| adj.target);
        let base_ty = self.resolve_vars_if_possible(&base_ty);

        // Need to deref because overloaded place ops take self by-reference.
        let base_ty =
            base_ty.builtin_deref(false).expect("place op takes something that is not a ref").ty;

        let method = self.try_overloaded_place_op(expr.span, base_ty, arg_tys, Needs::MutPlace, op);
        let method = match method {
            Some(ok) => self.register_infer_ok_obligations(ok),
            None => return self.tcx.sess.delay_span_bug(expr.span, "re-trying op failed"),
        };
        debug!("convert_place_op_to_mutable: method={:?}", method);
        self.write_method_call(expr.hir_id, method);

        let (region, mutbl) = if let ty::Ref(r, _, mutbl) = method.sig.inputs()[0].kind {
            (r, mutbl)
        } else {
            span_bug!(expr.span, "input to place op is not a ref?");
        };

        // Convert the autoref in the base expr to mutable with the correct
        // region and mutability.
        let base_expr_ty = self.node_ty(base_expr.hir_id);
        if let Some(adjustments) =
            self.tables.borrow_mut().adjustments_mut().get_mut(base_expr.hir_id)
        {
            let mut source = base_expr_ty;
            for adjustment in &mut adjustments[..] {
                if let Adjust::Borrow(AutoBorrow::Ref(..)) = adjustment.kind {
                    debug!("convert_place_op_to_mutable: converting autoref {:?}", adjustment);
                    let mutbl = match mutbl {
                        hir::Mutability::Not => AutoBorrowMutability::Not,
                        hir::Mutability::Mut => AutoBorrowMutability::Mut {
                            // For initial two-phase borrow
                            // deployment, conservatively omit
                            // overloaded operators.
                            allow_two_phase_borrow: AllowTwoPhase::No,
                        },
                    };
                    adjustment.kind = Adjust::Borrow(AutoBorrow::Ref(region, mutbl));
                    adjustment.target =
                        self.tcx.mk_ref(region, ty::TypeAndMut { ty: source, mutbl: mutbl.into() });
                }
                source = adjustment.target;
            }

            // If we have an autoref followed by unsizing at the end, fix the unsize target.

            if let [.., Adjustment { kind: Adjust::Borrow(AutoBorrow::Ref(..)), .. }, Adjustment { kind: Adjust::Pointer(PointerCast::Unsize), ref mut target }] =
                adjustments[..]
            {
                *target = method.sig.inputs()[0];
            }
        }
    }
}
<commit_msg>Use expr_ty_adjusted in convert_place_op_to_mutable<commit_after>use crate::check::{FnCtxt, Needs, PlaceOp};
use rustc_hir as hir;
use rustc_middle::ty::adjustment::{Adjust, Adjustment, OverloadedDeref, PointerCast};
use rustc_middle::ty::adjustment::{AllowTwoPhase, AutoBorrow, AutoBorrowMutability};
use rustc_middle::ty::{self, Ty};

impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
    /// Convert auto-derefs, indices, etc of an expression from `Deref` and `Index`
    /// into `DerefMut` and `IndexMut` respectively.
    ///
    /// This is a second pass of typechecking derefs/indices. We need this we do not
    /// always know whether a place needs to be mutable or not in the first pass.
    /// This happens whether there is an implicit mutable reborrow, e.g. when the type
    /// is used as the receiver of a method call.
    pub fn convert_place_derefs_to_mutable(&self, expr: &hir::Expr<'_>) {
        // Gather up expressions we want to munge.
        let mut exprs = vec![expr];

        loop {
            match exprs.last().unwrap().kind {
                hir::ExprKind::Field(ref expr, _)
                | hir::ExprKind::Index(ref expr, _)
                | hir::ExprKind::Unary(hir::UnOp::UnDeref, ref expr) => exprs.push(&expr),
                _ => break,
            }
        }

        debug!("convert_place_derefs_to_mutable: exprs={:?}", exprs);

        // Fix up autoderefs and derefs.
        for (i, &expr) in exprs.iter().rev().enumerate() {
            debug!("convert_place_derefs_to_mutable: i={} expr={:?}", i, expr);

            // Fix up the autoderefs. Autorefs can only occur immediately preceding
            // overloaded place ops, and will be fixed by them in order to get
            // the correct region.
            let mut source = self.node_ty(expr.hir_id);
            // Do not mutate adjustments in place, but rather take them,
            // and replace them after mutating them, to avoid having the
            // tables borrowed during (`deref_mut`) method resolution.
            let previous_adjustments =
                self.tables.borrow_mut().adjustments_mut().remove(expr.hir_id);
            if let Some(mut adjustments) = previous_adjustments {
                let needs = Needs::MutPlace;
                for adjustment in &mut adjustments {
                    if let Adjust::Deref(Some(ref mut deref)) = adjustment.kind {
                        if let Some(ok) = self.try_overloaded_deref(expr.span, source, needs) {
                            let method = self.register_infer_ok_obligations(ok);
                            if let ty::Ref(region, _, mutbl) = method.sig.output().kind {
                                *deref = OverloadedDeref { region, mutbl };
                            }
                        }
                    }
                    source = adjustment.target;
                }
                self.tables.borrow_mut().adjustments_mut().insert(expr.hir_id, adjustments);
            }

            match expr.kind {
                hir::ExprKind::Index(ref base_expr, ref index_expr) => {
                    // We need to get the final type in case dereferences were needed for the trait
                    // to apply (#72002).
                    let index_expr_ty = self.tables.borrow().expr_ty_adjusted(index_expr);
                    self.convert_place_op_to_mutable(
                        PlaceOp::Index,
                        expr,
                        base_expr,
                        &[index_expr_ty],
                    );
                }
                hir::ExprKind::Unary(hir::UnOp::UnDeref, ref base_expr) => {
                    self.convert_place_op_to_mutable(PlaceOp::Deref, expr, base_expr, &[]);
                }
                _ => {}
            }
        }
    }

    fn convert_place_op_to_mutable(
        &self,
        op: PlaceOp,
        expr: &hir::Expr<'_>,
        base_expr: &hir::Expr<'_>,
        arg_tys: &[Ty<'tcx>],
    ) {
        debug!("convert_place_op_to_mutable({:?}, {:?}, {:?}, {:?})", op, expr, base_expr, arg_tys);
        if !self.tables.borrow().is_method_call(expr) {
            debug!("convert_place_op_to_mutable - builtin, nothing to do");
            return;
        }

        // Need to deref because overloaded place ops take self by-reference.
        let base_ty = self
            .tables
            .borrow()
            .expr_ty_adjusted(base_expr)
            .builtin_deref(false)
            .expect("place op takes something that is not a ref")
            .ty;

        let method = self.try_overloaded_place_op(expr.span, base_ty, arg_tys, Needs::MutPlace, op);
        let method = match method {
            Some(ok) => self.register_infer_ok_obligations(ok),
            None => return self.tcx.sess.delay_span_bug(expr.span, "re-trying op failed"),
        };
        debug!("convert_place_op_to_mutable: method={:?}", method);
        self.write_method_call(expr.hir_id, method);

        let (region, mutbl) = if let ty::Ref(r, _, mutbl) = method.sig.inputs()[0].kind {
            (r, mutbl)
        } else {
            span_bug!(expr.span, "input to place op is not a ref?");
        };

        // Convert the autoref in the base expr to mutable with the correct
        // region and mutability.
        let base_expr_ty = self.node_ty(base_expr.hir_id);
        if let Some(adjustments) =
            self.tables.borrow_mut().adjustments_mut().get_mut(base_expr.hir_id)
        {
            let mut source = base_expr_ty;
            for adjustment in &mut adjustments[..] {
                if let Adjust::Borrow(AutoBorrow::Ref(..)) = adjustment.kind {
                    debug!("convert_place_op_to_mutable: converting autoref {:?}", adjustment);
                    let mutbl = match mutbl {
                        hir::Mutability::Not => AutoBorrowMutability::Not,
                        hir::Mutability::Mut => AutoBorrowMutability::Mut {
                            // For initial two-phase borrow
                            // deployment, conservatively omit
                            // overloaded operators.
                            allow_two_phase_borrow: AllowTwoPhase::No,
                        },
                    };
                    adjustment.kind = Adjust::Borrow(AutoBorrow::Ref(region, mutbl));
                    adjustment.target =
                        self.tcx.mk_ref(region, ty::TypeAndMut { ty: source, mutbl: mutbl.into() });
                }
                source = adjustment.target;
            }

            // If we have an autoref followed by unsizing at the end, fix the unsize target.

            if let [.., Adjustment { kind: Adjust::Borrow(AutoBorrow::Ref(..)), .. }, Adjustment { kind: Adjust::Pointer(PointerCast::Unsize), ref mut target }] =
                adjustments[..]
            {
                *target = method.sig.inputs()[0];
            }
        }
    }
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/earliest-possible-day-of-full-bloom/
pub fn earliest_full_bloom(plant_time: Vec<i32>, grow_time: Vec<i32>) -> i32 {
    todo!()
}

fn main() {
    println!("{}", earliest_full_bloom(vec![1, 4, 3], vec![2, 3, 1])); // 9
    println!(
        "{}",
        earliest_full_bloom(vec![1, 2, 3, 2], vec![2, 1, 2, 1])
    ); // 9
    println!("{}", earliest_full_bloom(vec![1], vec![1])); // 2
}
<commit_msg>Update problem 2136<commit_after>// https://leetcode.com/problems/earliest-possible-day-of-full-bloom/
pub fn earliest_full_bloom(plant_time: Vec<i32>, grow_time: Vec<i32>) -> i32 {
    todo!()
}

fn main() {
    println!("{}", earliest_full_bloom(vec![1, 4, 3], vec![2, 3, 1])); // 9
    println!(
        "{}",
        earliest_full_bloom(vec![1, 2, 3, 2], vec![2, 1, 2, 1])
    ); // 9
    println!("{}", earliest_full_bloom(vec![1], vec![1])); // 2
    println!(
        "{}",
        earliest_full_bloom(vec![3, 1, 5, 3], vec![2, 3, 1, 4])
    ); // 13
}
<|endoftext|>


<commit_before>use std::fs::{create_dir_all, File, remove_file};
use std::io::Write;

use clap::ArgMatches;
use rusoto::ChainProvider;

use aws::credentials_provider;
use encryption::Encryptor;
use error::KawsResult;
use process::execute_child_process;

pub struct Cluster<'a> {
    aws_account_id: Option<&'a str>,
    aws_credentials_provider: ChainProvider,
    ca_cert_path: String,
    ca_key_path: String,
    coreos_ami: Option<&'a str>,
    domain: Option<&'a str>,
    encrypted_ca_key_path: String,
    encrypted_master_key_path: String,
    encrypted_node_key_path: String,
    iam_users: Option<Vec<&'a str>>,
    instance_size: Option<&'a str>,
    kms_master_key_id: Option<&'a str>,
    kms_policy_path: String,
    kubernetes_version: Option<&'a str>,
    master_cert_path: String,
    master_csr_path: String,
    master_key_path: String,
    masters_max_size: Option<&'a str>,
    masters_min_size: Option<&'a str>,
    name: &'a str,
    node_cert_path: String,
    node_csr_path: String,
    node_key_path: String,
    nodes_max_size: Option<&'a str>,
    nodes_min_size: Option<&'a str>,
    openssl_config_path: String,
    region: &'a str,
    ssh_key: Option<&'a str>,
    tfvars_path: String,
    zone_id: Option<&'a str>,
}

impl<'a> Cluster<'a> {
    pub fn new(matches: &'a ArgMatches) -> Self {
        let name =  matches.value_of("cluster").expect("clap should have required cluster");

        Cluster {
            aws_account_id: matches.value_of("aws-account-id"),
            aws_credentials_provider: credentials_provider(
                matches.value_of("aws-credentials-path"),
                matches.value_of("aws-credentials-profile"),
            ),
            ca_cert_path: format!("clusters/{}/ca.pem", name),
            ca_key_path: format!("clusters/{}/ca-key.pem", name),
            coreos_ami: matches.value_of("ami"),
            domain: matches.value_of("domain"),
            encrypted_ca_key_path: format!("clusters/{}/ca-key-encrypted.base64", name),
            encrypted_master_key_path: format!("clusters/{}/master-key-encrypted.base64", name),
            encrypted_node_key_path: format!("clusters/{}/node-key-encrypted.base64", name),
            name: name,
            iam_users: matches.values_of("iam-users").map(|values| values.collect()),
            instance_size: matches.value_of("size"),
            kms_master_key_id: matches.value_of("kms-key"),
            kms_policy_path: format!("clusters/{}/kms-policy.json", name),
            kubernetes_version: matches.value_of("k8s-version"),
            master_cert_path: format!("clusters/{}/master.pem", name),
            master_csr_path: format!("clusters/{}/master.csr", name),
            master_key_path: format!("clusters/{}/master-key.pem", name),
            masters_max_size: matches.value_of("masters-max-size"),
            masters_min_size: matches.value_of("masters-min-size"),
            node_cert_path: format!("clusters/{}/node.pem", name),
            node_csr_path: format!("clusters/{}/node.csr", name),
            node_key_path: format!("clusters/{}/node-key.pem", name),
            nodes_max_size: matches.value_of("nodes-max-size"),
            nodes_min_size: matches.value_of("nodes-min-size"),
            openssl_config_path: format!("clusters/{}/openssl.cnf", name),
            region: matches.value_of("region").expect("clap should have required region"),
            ssh_key: matches.value_of("ssh-key"),
            tfvars_path: format!("clusters/{}/terraform.tfvars", name),
            zone_id: matches.value_of("zone-id"),
        }
    }

    pub fn init(&mut self) -> KawsResult {
        try!(self.create_directories());
        try!(self.create_tfvars());
        try!(self.create_openssl_config());
        try!(self.create_kms_policy());
        try!(self.create_pki_stubs());

        Ok(Some(format!(
            "Cluster \"{}\" initialized!",
            self.name,
        )))
    }

    pub fn generate_pki(&mut self) -> KawsResult {
        try!(self.create_ca());
        try!(self.create_master_credentials());
        try!(self.create_node_credentials());
        try!(self.encrypt_secrets(
            self.kms_master_key_id.expect("clap should have required kms-key"),
        ));

        Ok(None)
    }

    fn create_directories(&self) -> KawsResult {
        log_wrap!("Creating directories for the new cluster", {
            try!(create_dir_all(format!("clusters/{}", self.name)));
        });

        Ok(None)
    }

    fn create_kms_policy(&self) -> KawsResult {
        log_wrap!("Creating KMS policy file", {
            let mut file = try!(File::create(&self.kms_policy_path));

            let aws_account_id = self.aws_account_id.expect(
                "AWS account ID should have been required by clap"
            );

            let iam_user_arns = self.iam_users.as_ref().expect(
                "IAM users should have been required by clap"
            ).iter().map(|iam_user| {
                format!("\"arn:aws:iam::{}:user/{}\"", aws_account_id, iam_user)
            }).collect::<Vec<String>>().join(",");

            try!(write!(
                file,
                r#"{{
  "Version": "2012-10-17",
  "Statement": [
    {{
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:root"
        ]
      }},
      "Action": "kms:*",
      "Resource": "*"
    }},
    {{
      "Sid": "Allow access for Key Administrators",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Create*",
        "kms:Describe*",
        "kms:Enable*",
        "kms:List*",
        "kms:Put*",
        "kms:Update*",
        "kms:Revoke*",
        "kms:Disable*",
        "kms:Get*",
        "kms:Delete*",
        "kms:ScheduleKeyDeletion",
        "kms:CancelKeyDeletion"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow use of the key",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws_k8s_master_{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws_k8s_node_{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Encrypt",
        "kms:Decrypt",
        "kms:ReEncrypt*",
        "kms:GenerateDataKey*",
        "kms:DescribeKey"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow attachment of persistent resources",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws_k8s_master_{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws_k8s_node_{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:CreateGrant",
        "kms:ListGrants",
        "kms:RevokeGrant"
      ],
      "Resource": "*",
      "Condition": {{
        "Bool": {{
          "kms:GrantIsForAWSResource": true
        }}
      }}
    }}
  ]
}}
"#,
                aws_account_id = aws_account_id,
                cluster = self.name,
                iam_user_arns = iam_user_arns,
            ));
        });

        Ok(None)
    }

    fn create_master_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes master private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.master_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.master_key_path,
                "-out",
                &self.master_csr_path,
                "-subj",
                &format!("/CN=kaws-{}-master", self.name),
                "-config",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.master_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.master_cert_path,
                "-days",
                "365",
                "-extensions",
                "v3_req",
                "-extfile",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Removing Kubernetes master certificate signing request", {
            try!(remove_file(&self.master_csr_path));
        });

        Ok(None)
    }

    fn create_node_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes node private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.node_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.node_key_path,
                "-out",
                &self.node_csr_path,
                "-subj",
                &format!("/CN=kaws-{}-node", self.name),
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.node_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.node_cert_path,
                "-days",
                "365",
            ]));
        });

        log_wrap!("Removing Kubernetes node certificate signing request", {
            try!(remove_file(&self.node_csr_path));
        });

        Ok(None)
    }

    fn create_openssl_config(&self) -> KawsResult {
        log_wrap!("Creating OpenSSL config file", {
            let mut file = try!(File::create(&self.openssl_config_path));

            try!(write!(
                file,
                "\
[req]
req_extensions = v3_req
distinguished_name = req_distinguished_name
[req_distinguished_name]
[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment
subjectAltName = @alt_names
[alt_names]
DNS.1 = kubernetes
DNS.2 = kubernetes.default
DNS.3 = kubernetes.{}
IP.1 = 10.3.0.1
",
                self.domain.expect("domain should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn create_ca(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes certificate authority private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.ca_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes certificate authority certificate", {
            try!(execute_child_process("openssl", &[
                "req",
                "-x509",
                "-new",
                "-nodes",
                "-key",
                &self.ca_key_path,
                "-days",
                "10000",
                "-out",
                &self.ca_cert_path,
                "-subj",
                &format!("/CN=kaws-{}-ca", self.name),
            ]));
        });

        Ok(None)
    }

    fn create_pki_stubs(&self) -> KawsResult {
        let paths = [
            &self.ca_cert_path,
            &self.master_cert_path,
            &self.encrypted_master_key_path,
            &self.node_cert_path,
            &self.encrypted_node_key_path,
        ];

        for path in paths.iter() {
            try!(File::create(path));
        }

        Ok(None)
    }

    fn create_tfvars(&self) -> KawsResult {
        log_wrap!("Creating tfvars file", {
            let mut file = try!(File::create(&self.tfvars_path));

            try!(write!(
                file,
                "\
cluster = \"{}\"
coreos_ami = \"{}\"
domain = \"{}\"
etcd_01_initial_cluster_state = \"new\"
etcd_02_initial_cluster_state = \"new\"
etcd_03_initial_cluster_state = \"new\"
instance_size = \"{}\"
masters_max_size = \"{}\"
masters_min_size = \"{}\"
nodes_max_size = \"{}\"
nodes_min_size = \"{}\"
region = \"{}\"
ssh_key = \"{}\"
version = \"{}\"
zone_id = \"{}\"
",
                self.name,
                self.coreos_ami.expect("AMI should have been required by clap"),
                self.domain.expect("domain should have been required by clap"),
                self.instance_size.expect("instance size should have been required by clap"),
                self.masters_max_size.expect("masters max size should have been required by clap"),
                self.masters_min_size.expect("masters min size should have been required by clap"),
                self.nodes_max_size.expect("nodes max size should have been required by clap"),
                self.nodes_min_size.expect("nodes min size should have been required by clap"),
                self.region,
                self.ssh_key.expect("ssh key should have been required by clap"),
                self.kubernetes_version.expect("k8s version should have been required by clap"),
                self.zone_id.expect("zone ID should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn encrypt_secrets<'b>(&self, kms_key_id: &'b str) -> KawsResult {
        let mut encryptor = Encryptor::new(
            self.aws_credentials_provider.clone(),
            try!(self.region.parse()),
            kms_key_id,
        );

        log_wrap!("Encrypting Kubernetes certificate authority private key", {
            try!(encryptor.encrypt_file(&self.ca_key_path, &self.encrypted_ca_key_path));
        });

        log_wrap!("Encrypting Kubernetes master private key", {
            try!(encryptor.encrypt_file(&self.master_key_path, &self.encrypted_master_key_path));
        });

        log_wrap!("Encrypting Kubernetes node private key", {
            try!(encryptor.encrypt_file(&self.node_key_path, &self.encrypted_node_key_path));
        });

        Ok(None)
    }
}
<commit_msg>Use new role names in kms-policy.json.<commit_after>use std::fs::{create_dir_all, File, remove_file};
use std::io::Write;

use clap::ArgMatches;
use rusoto::ChainProvider;

use aws::credentials_provider;
use encryption::Encryptor;
use error::KawsResult;
use process::execute_child_process;

pub struct Cluster<'a> {
    aws_account_id: Option<&'a str>,
    aws_credentials_provider: ChainProvider,
    ca_cert_path: String,
    ca_key_path: String,
    coreos_ami: Option<&'a str>,
    domain: Option<&'a str>,
    encrypted_ca_key_path: String,
    encrypted_master_key_path: String,
    encrypted_node_key_path: String,
    iam_users: Option<Vec<&'a str>>,
    instance_size: Option<&'a str>,
    kms_master_key_id: Option<&'a str>,
    kms_policy_path: String,
    kubernetes_version: Option<&'a str>,
    master_cert_path: String,
    master_csr_path: String,
    master_key_path: String,
    masters_max_size: Option<&'a str>,
    masters_min_size: Option<&'a str>,
    name: &'a str,
    node_cert_path: String,
    node_csr_path: String,
    node_key_path: String,
    nodes_max_size: Option<&'a str>,
    nodes_min_size: Option<&'a str>,
    openssl_config_path: String,
    region: &'a str,
    ssh_key: Option<&'a str>,
    tfvars_path: String,
    zone_id: Option<&'a str>,
}

impl<'a> Cluster<'a> {
    pub fn new(matches: &'a ArgMatches) -> Self {
        let name =  matches.value_of("cluster").expect("clap should have required cluster");

        Cluster {
            aws_account_id: matches.value_of("aws-account-id"),
            aws_credentials_provider: credentials_provider(
                matches.value_of("aws-credentials-path"),
                matches.value_of("aws-credentials-profile"),
            ),
            ca_cert_path: format!("clusters/{}/ca.pem", name),
            ca_key_path: format!("clusters/{}/ca-key.pem", name),
            coreos_ami: matches.value_of("ami"),
            domain: matches.value_of("domain"),
            encrypted_ca_key_path: format!("clusters/{}/ca-key-encrypted.base64", name),
            encrypted_master_key_path: format!("clusters/{}/master-key-encrypted.base64", name),
            encrypted_node_key_path: format!("clusters/{}/node-key-encrypted.base64", name),
            name: name,
            iam_users: matches.values_of("iam-users").map(|values| values.collect()),
            instance_size: matches.value_of("size"),
            kms_master_key_id: matches.value_of("kms-key"),
            kms_policy_path: format!("clusters/{}/kms-policy.json", name),
            kubernetes_version: matches.value_of("k8s-version"),
            master_cert_path: format!("clusters/{}/master.pem", name),
            master_csr_path: format!("clusters/{}/master.csr", name),
            master_key_path: format!("clusters/{}/master-key.pem", name),
            masters_max_size: matches.value_of("masters-max-size"),
            masters_min_size: matches.value_of("masters-min-size"),
            node_cert_path: format!("clusters/{}/node.pem", name),
            node_csr_path: format!("clusters/{}/node.csr", name),
            node_key_path: format!("clusters/{}/node-key.pem", name),
            nodes_max_size: matches.value_of("nodes-max-size"),
            nodes_min_size: matches.value_of("nodes-min-size"),
            openssl_config_path: format!("clusters/{}/openssl.cnf", name),
            region: matches.value_of("region").expect("clap should have required region"),
            ssh_key: matches.value_of("ssh-key"),
            tfvars_path: format!("clusters/{}/terraform.tfvars", name),
            zone_id: matches.value_of("zone-id"),
        }
    }

    pub fn init(&mut self) -> KawsResult {
        try!(self.create_directories());
        try!(self.create_tfvars());
        try!(self.create_openssl_config());
        try!(self.create_kms_policy());
        try!(self.create_pki_stubs());

        Ok(Some(format!(
            "Cluster \"{}\" initialized!",
            self.name,
        )))
    }

    pub fn generate_pki(&mut self) -> KawsResult {
        try!(self.create_ca());
        try!(self.create_master_credentials());
        try!(self.create_node_credentials());
        try!(self.encrypt_secrets(
            self.kms_master_key_id.expect("clap should have required kms-key"),
        ));

        Ok(None)
    }

    fn create_directories(&self) -> KawsResult {
        log_wrap!("Creating directories for the new cluster", {
            try!(create_dir_all(format!("clusters/{}", self.name)));
        });

        Ok(None)
    }

    fn create_kms_policy(&self) -> KawsResult {
        log_wrap!("Creating KMS policy file", {
            let mut file = try!(File::create(&self.kms_policy_path));

            let aws_account_id = self.aws_account_id.expect(
                "AWS account ID should have been required by clap"
            );

            let iam_user_arns = self.iam_users.as_ref().expect(
                "IAM users should have been required by clap"
            ).iter().map(|iam_user| {
                format!("\"arn:aws:iam::{}:user/{}\"", aws_account_id, iam_user)
            }).collect::<Vec<String>>().join(",");

            try!(write!(
                file,
                r#"{{
  "Version": "2012-10-17",
  "Statement": [
    {{
      "Sid": "Enable IAM User Permissions",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:root"
        ]
      }},
      "Action": "kms:*",
      "Resource": "*"
    }},
    {{
      "Sid": "Allow access for Key Administrators",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Create*",
        "kms:Describe*",
        "kms:Enable*",
        "kms:List*",
        "kms:Put*",
        "kms:Update*",
        "kms:Revoke*",
        "kms:Disable*",
        "kms:Get*",
        "kms:Delete*",
        "kms:ScheduleKeyDeletion",
        "kms:CancelKeyDeletion"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow use of the key",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:Encrypt",
        "kms:Decrypt",
        "kms:ReEncrypt*",
        "kms:GenerateDataKey*",
        "kms:DescribeKey"
      ],
      "Resource": "*"
    }},
    {{
      "Sid": "Allow attachment of persistent resources",
      "Effect": "Allow",
      "Principal": {{
        "AWS": [
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-master-{cluster}",
          "arn:aws:iam::{aws_account_id}:role/kaws-k8s-node-{cluster}",
          {iam_user_arns}
        ]
      }},
      "Action": [
        "kms:CreateGrant",
        "kms:ListGrants",
        "kms:RevokeGrant"
      ],
      "Resource": "*",
      "Condition": {{
        "Bool": {{
          "kms:GrantIsForAWSResource": true
        }}
      }}
    }}
  ]
}}
"#,
                aws_account_id = aws_account_id,
                cluster = self.name,
                iam_user_arns = iam_user_arns,
            ));
        });

        Ok(None)
    }

    fn create_master_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes master private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.master_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.master_key_path,
                "-out",
                &self.master_csr_path,
                "-subj",
                &format!("/CN=kaws-master-{}", self.name),
                "-config",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Creating Kubernetes master certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.master_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.master_cert_path,
                "-days",
                "365",
                "-extensions",
                "v3_req",
                "-extfile",
                &self.openssl_config_path,
            ]));
        });

        log_wrap!("Removing Kubernetes master certificate signing request", {
            try!(remove_file(&self.master_csr_path));
        });

        Ok(None)
    }

    fn create_node_credentials(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes node private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.node_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate signing request", {
            try!(execute_child_process("openssl", &[
                "req",
                "-new",
                "-key",
                &self.node_key_path,
                "-out",
                &self.node_csr_path,
                "-subj",
                &format!("/CN=kaws-node-{}", self.name),
            ]));
        });

        log_wrap!("Creating Kubernetes node certificate", {
            try!(execute_child_process("openssl", &[
                "x509",
                "-req",
                "-in",
                &self.node_csr_path,
                "-CA",
                &self.ca_cert_path,
                "-CAkey",
                &self.ca_key_path,
                "-CAcreateserial",
                "-out",
                &self.node_cert_path,
                "-days",
                "365",
            ]));
        });

        log_wrap!("Removing Kubernetes node certificate signing request", {
            try!(remove_file(&self.node_csr_path));
        });

        Ok(None)
    }

    fn create_openssl_config(&self) -> KawsResult {
        log_wrap!("Creating OpenSSL config file", {
            let mut file = try!(File::create(&self.openssl_config_path));

            try!(write!(
                file,
                "\
[req]
req_extensions = v3_req
distinguished_name = req_distinguished_name
[req_distinguished_name]
[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment
subjectAltName = @alt_names
[alt_names]
DNS.1 = kubernetes
DNS.2 = kubernetes.default
DNS.3 = kubernetes.{}
IP.1 = 10.3.0.1
",
                self.domain.expect("domain should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn create_ca(&self) -> KawsResult {
        log_wrap!("Creating Kubernetes certificate authority private key", {
            try!(execute_child_process("openssl", &[
                "genrsa",
                "-out",
                &self.ca_key_path,
                "2048",
            ]));
        });

        log_wrap!("Creating Kubernetes certificate authority certificate", {
            try!(execute_child_process("openssl", &[
                "req",
                "-x509",
                "-new",
                "-nodes",
                "-key",
                &self.ca_key_path,
                "-days",
                "10000",
                "-out",
                &self.ca_cert_path,
                "-subj",
                &format!("/CN=kaws-ca-{}", self.name),
            ]));
        });

        Ok(None)
    }

    fn create_pki_stubs(&self) -> KawsResult {
        let paths = [
            &self.ca_cert_path,
            &self.master_cert_path,
            &self.encrypted_master_key_path,
            &self.node_cert_path,
            &self.encrypted_node_key_path,
        ];

        for path in paths.iter() {
            try!(File::create(path));
        }

        Ok(None)
    }

    fn create_tfvars(&self) -> KawsResult {
        log_wrap!("Creating tfvars file", {
            let mut file = try!(File::create(&self.tfvars_path));

            try!(write!(
                file,
                "\
cluster = \"{}\"
coreos_ami = \"{}\"
domain = \"{}\"
etcd_01_initial_cluster_state = \"new\"
etcd_02_initial_cluster_state = \"new\"
etcd_03_initial_cluster_state = \"new\"
instance_size = \"{}\"
masters_max_size = \"{}\"
masters_min_size = \"{}\"
nodes_max_size = \"{}\"
nodes_min_size = \"{}\"
region = \"{}\"
ssh_key = \"{}\"
version = \"{}\"
zone_id = \"{}\"
",
                self.name,
                self.coreos_ami.expect("AMI should have been required by clap"),
                self.domain.expect("domain should have been required by clap"),
                self.instance_size.expect("instance size should have been required by clap"),
                self.masters_max_size.expect("masters max size should have been required by clap"),
                self.masters_min_size.expect("masters min size should have been required by clap"),
                self.nodes_max_size.expect("nodes max size should have been required by clap"),
                self.nodes_min_size.expect("nodes min size should have been required by clap"),
                self.region,
                self.ssh_key.expect("ssh key should have been required by clap"),
                self.kubernetes_version.expect("k8s version should have been required by clap"),
                self.zone_id.expect("zone ID should have been required by clap"),
            ));
        });

        Ok(None)
    }

    fn encrypt_secrets<'b>(&self, kms_key_id: &'b str) -> KawsResult {
        let mut encryptor = Encryptor::new(
            self.aws_credentials_provider.clone(),
            try!(self.region.parse()),
            kms_key_id,
        );

        log_wrap!("Encrypting Kubernetes certificate authority private key", {
            try!(encryptor.encrypt_file(&self.ca_key_path, &self.encrypted_ca_key_path));
        });

        log_wrap!("Encrypting Kubernetes master private key", {
            try!(encryptor.encrypt_file(&self.master_key_path, &self.encrypted_master_key_path));
        });

        log_wrap!("Encrypting Kubernetes node private key", {
            try!(encryptor.encrypt_file(&self.node_key_path, &self.encrypted_node_key_path));
        });

        Ok(None)
    }
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;

pub struct CommandInfo<T: Copy> {
  key: char,
  desc: String,
  cmd: T,
  hidden: bool,
}

impl<T: Copy> CommandInfo<T> {
  pub fn new<S: AsRef<str>>(key: char, desc: S, cmd: T) -> Self {
    Self {
      key,
      // TODO: We might be making unnecessary copies of strings here.
      desc: String::from(desc.as_ref()),
      cmd,
      hidden: false
    }
  }

  pub fn hidden(mut self) -> Self {
    self.hidden = true;
    self
  }
}

pub trait CommandProcessor<T: Copy> {
  fn from_char(c: char) -> Option<T> {
    for info in Self::get_command_info().iter() {
      if c == info.key { return Some(info.cmd); }
    }
    None
  }

  fn get_command_info() -> Vec<CommandInfo<T>>;

  fn show_help() {
    for info in Self::get_command_info().iter().filter(|i| !i.hidden) {
      println!("  {} - {}", info.key, info.desc);
    }
  }

  fn get_from_input(input: String) -> Option<T> {
    match input.chars().next() {
      Some(k) => {
        let k = k.to_ascii_lowercase();
        if k == 'h' || k == '?' {
          println!("Here's what I understand right now:\n");
          Self::show_help();
          println!("");
          return None;
        } else if let Some(cmd) = Self::from_char(k) {
          return Some(cmd);
        }
      },
      None => {}
    }
    println!("I have no idea what you're talking about.");
    return None;
  }
}
<commit_msg>Add a bit of docs.<commit_after>use std::ascii::AsciiExt;

pub struct CommandInfo<T: Copy> {
  key: char,
  desc: String,
  cmd: T,
  hidden: bool,
}

impl<T: Copy> CommandInfo<T> {
  pub fn new<S: AsRef<str>>(key: char, desc: S, cmd: T) -> Self {
    Self {
      key,
      // TODO: We might be making unnecessary copies of strings here.
      desc: String::from(desc.as_ref()),
      cmd,
      hidden: false
    }
  }

  pub fn hidden(mut self) -> Self {
    self.hidden = true;
    self
  }
}

pub trait CommandProcessor<T: Copy> {
  fn from_char(c: char) -> Option<T> {
    // A HashMap here would obviously be more efficient, but
    // since we're not going to be called very often, it's probably
    // not that big a deal.
    for info in Self::get_command_info().iter() {
      if c == info.key { return Some(info.cmd); }
    }
    None
  }

  fn get_command_info() -> Vec<CommandInfo<T>>;

  fn show_help() {
    for info in Self::get_command_info().iter().filter(|i| !i.hidden) {
      println!("  {} - {}", info.key, info.desc);
    }
  }

  fn get_from_input(input: String) -> Option<T> {
    match input.chars().next() {
      Some(k) => {
        let k = k.to_ascii_lowercase();
        if k == 'h' || k == '?' {
          println!("Here's what I understand right now:\n");
          Self::show_help();
          println!("");
          return None;
        } else if let Some(cmd) = Self::from_char(k) {
          return Some(cmd);
        }
      },
      None => {}
    }
    println!("I have no idea what you're talking about.");
    return None;
  }
}
<|endoftext|>


<commit_before>extern crate collections;

use self::collections::hashmap::HashMap;

struct Command {
	action: ~str,
	flags: ~MutableMap<~str, ~str>
}

pub fn ParseCommand(line: ~str) -> Command {
	let words = line.split(' ').collect();
}<commit_msg>This is still a work in progress<commit_after>/*extern crate collections;

use self::collections::hashmap::HashMap;

struct Command {
	action: ~str,
	flags: ~MutableMap<~str, ~str>
}

pub fn ParseCommand(line: ~str) -> Command {
	let words = str::split_char(line, ' '); //line.split(' ').collect();
}*/
<|endoftext|>