/*
 * ChaCha20 implementation based on RFC8439
 * ChaCha20 is a stream cipher developed independently by Daniel J. Bernstein.
 * To use it, the `chacha20` function should be called with appropriate
 * parameters and the output of the function should be XORed with plain text.
 */

macro_rules! quarter_round {
    ($a:expr,$b:expr,$c:expr,$d:expr) => {
        $a = $a.wrapping_add($b);
        $d = ($d ^ $a).rotate_left(16);
        $c = $c.wrapping_add($d);
        $b = ($b ^ $c).rotate_left(12);
        $a = $a.wrapping_add($b);
        $d = ($d ^ $a).rotate_left(8);
        $c = $c.wrapping_add($d);
        $b = ($b ^ $c).rotate_left(7);
    };
}

#[allow(dead_code)]
// "expand 32-byte k", written in little-endian order
pub const C: [u32; 4] = [0x61707865, 0x3320646e, 0x79622d32, 0x6b206574];

/**
 * `chacha20` function takes as input an array of 16 32-bit integers (512 bits)
 * of which 128 bits is the constant 'expand 32-byte k', 256 bits is the key,
 * and 128 bits are nonce and counter. According to RFC8439, the nonce should
 * be 96 bits long, which leaves 32 bits for the counter. Given that the block
 * length is 512 bits, this leaves enough counter values to encrypt 256GB of
 * data.
 *
 * The 16 input numbers can be thought of as the elements of a 4x4 matrix like
 * the one bellow, on which we do the main operations of the cipher.
 *
 * +----+----+----+----+
 * | 00 | 01 | 02 | 03 |
 * +----+----+----+----+
 * | 04 | 05 | 06 | 07 |
 * +----+----+----+----+
 * | 08 | 09 | 10 | 11 |
 * +----+----+----+----+
 * | 12 | 13 | 14 | 15 |
 * +----+----+----+----+
 *
 * As per the diagram bellow, input[0, 1, 2, 3] are the constants mentioned
 * above, input[4..=11] is filled with the key, and input[6..=9] should be
 * filled with nonce and counter values. The output of the function is stored
 * in `output` variable and can be XORed with the plain text to produce the
 * cipher text.
 *
 * +------+------+------+------+
 * |      |      |      |      |
 * | C[0] | C[1] | C[2] | C[3] |
 * |      |      |      |      |
 * +------+------+------+------+
 * |      |      |      |      |
 * | key0 | key1 | key2 | key3 |
 * |      |      |      |      |
 * +------+------+------+------+
 * |      |      |      |      |
 * | key4 | key5 | key6 | key7 |
 * |      |      |      |      |
 * +------+------+------+------+
 * |      |      |      |      |
 * | ctr0 | no.0 | no.1 | no.2 |
 * |      |      |      |      |
 * +------+------+------+------+
 *
 * Note that the constants, the key, and the nonce should be written in
 * little-endian order, meaning that for example if the key is 01:02:03:04
 * (in hex), it corresponds to the integer 0x04030201. It is important to
 * know that the hex value of the counter is meaningless, and only its integer
 * value matters, and it should start with (for example) 0x00000000, and then
 * 0x00000001 and so on until 0xffffffff. Keep in mind that as soon as we get
 * from bytes to words, we stop caring about their representation in memory,
 * and we only need the math to be correct.
 *
 * The output of the function can be used without any change, as long as the
 * plain text has the same endianness. For example if the plain text is
 * "hello world", and the first word of the output is 0x01020304, then the
 * first byte of plain text ('h') should be XORed with the least-significant
 * byte of 0x01020304, which is 0x04.
*/
pub fn chacha20(input: &[u32; 16], output: &mut [u32; 16]) {
    output.copy_from_slice(&input[..]);
    for _ in 0..10 {
        // Odd round (column round)
        quarter_round!(output[0], output[4], output[8], output[12]); //  column 1
        quarter_round!(output[1], output[5], output[9], output[13]); //  column 2
        quarter_round!(output[2], output[6], output[10], output[14]); // column 3
        quarter_round!(output[3], output[7], output[11], output[15]); // column 4

        // Even round (diagonal round)
        quarter_round!(output[0], output[5], output[10], output[15]); // diag 1
        quarter_round!(output[1], output[6], output[11], output[12]); // diag 2
        quarter_round!(output[2], output[7], output[8], output[13]); //  diag 3
        quarter_round!(output[3], output[4], output[9], output[14]); //  diag 4
    }
    for (a, &b) in output.iter_mut().zip(input.iter()) {
        *a = a.wrapping_add(b);
    }
}

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

    fn output_hex(inp: &[u32; 16]) -> String {
        let mut res = String::new();
        res.reserve(512 / 4);
        for &x in inp {
            write!(&mut res, "{x:08x}").unwrap();
        }
        res
    }

    #[test]
    // test vector 1
    fn basic_tv1() {
        let mut inp = [0u32; 16];
        let mut out = [0u32; 16];
        inp[0] = C[0];
        inp[1] = C[1];
        inp[2] = C[2];
        inp[3] = C[3];
        inp[4] = 0x03020100; // The key is 00:01:02:..:1f (hex)
        inp[5] = 0x07060504;
        inp[6] = 0x0b0a0908;
        inp[7] = 0x0f0e0d0c;
        inp[8] = 0x13121110;
        inp[9] = 0x17161514;
        inp[10] = 0x1b1a1918;
        inp[11] = 0x1f1e1d1c;
        inp[12] = 0x00000001; // The value of counter is 1 (an integer). Nonce:
        inp[13] = 0x09000000; // 00:00:00:09
        inp[14] = 0x4a000000; // 00:00:00:4a
        inp[15] = 0x00000000; // 00:00:00:00
        chacha20(&inp, &mut out);
        assert_eq!(
            output_hex(&out),
            concat!(
                "e4e7f11015593bd11fdd0f50c47120a3c7f4d1c70368c0339aaa22044e6cd4c3",
                "466482d209aa9f0705d7c214a2028bd9d19c12b5b94e16dee883d0cb4e3c50a2"
            )
        );
    }
}