<commit_before>use curve25519_dalek::constants::ED25519_BASEPOINT;
use curve25519_dalek::scalar::Scalar;
use sha256::sha256_compress;
use std::fmt;
use util::bs58_encode_check;

const PAYMENT_ADDRESS_PREFIX_LENGTH: usize = 2;
const SPENDING_KEY_PREFIX_LENGTH: usize = 2;
const INVIEWING_KEY_PREFIX_LENGTH: usize = 3;
pub const PAYMENT_ADDRESS_PREFIX: [u8; PAYMENT_ADDRESS_PREFIX_LENGTH] = [0x16, 0x9a];
pub const SPENDING_KEY_PREFIX: [u8; SPENDING_KEY_PREFIX_LENGTH] = [0xab, 0x36];
pub const INVIEWING_KEY_PREFIX: [u8; INVIEWING_KEY_PREFIX_LENGTH] = [0xa8, 0xab, 0xd3];

/// A Zcash spending key.
pub struct SpendingKey {
    a_sk: [u8; 32], // 252 bits; first 4 bits are always 0
}

/// A Zcash incoming viewing key.
pub struct ViewingKey {
    pub a_pk: [u8; 32],
    pub sk_enc: [u8; 32],
}

/// A Zcash payment address.
pub struct PaymentAddress {
    pub a_pk: [u8; 32],
    pub pk_enc: [u8; 32],
}

impl fmt::Display for SpendingKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; SPENDING_KEY_PREFIX_LENGTH + 32];
        {
            let (prefix, rest) = data.split_at_mut(SPENDING_KEY_PREFIX_LENGTH);
            prefix.copy_from_slice(&SPENDING_KEY_PREFIX);
            rest.copy_from_slice(&self.a_sk);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl fmt::Display for ViewingKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; INVIEWING_KEY_PREFIX_LENGTH + 32 * 2];
        {
            let (prefix, rest) = data.split_at_mut(INVIEWING_KEY_PREFIX_LENGTH);
            prefix.copy_from_slice(&INVIEWING_KEY_PREFIX);
            rest[..32].copy_from_slice(&self.a_pk);
            rest[32..].copy_from_slice(&self.sk_enc);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl fmt::Display for PaymentAddress {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; PAYMENT_ADDRESS_PREFIX_LENGTH + 32 * 2];
        {
            let (prefix, rest) = data.split_at_mut(PAYMENT_ADDRESS_PREFIX_LENGTH);
            prefix.copy_from_slice(&PAYMENT_ADDRESS_PREFIX);
            rest[..32].copy_from_slice(&self.a_pk);
            rest[32..].copy_from_slice(&self.pk_enc);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl SpendingKey {
    /// Creates a new spending key instance for a given 252-bit *a_sk*.
    /// *a_sk* is represented by 32 bytes; the first four bits of the first byte are zeroed.
    pub fn new(a_sk_u256: [u8; 32]) -> SpendingKey {
        let mut a_sk_u252 = [0u8; 32];
        a_sk_u252.copy_from_slice(&a_sk_u256);
        a_sk_u252[0] &= 0x0f;

        SpendingKey { a_sk: a_sk_u252 }
    }

    /// Calculates the payment address for this spending key
    pub fn address(&self) -> PaymentAddress {
        let viewing_key = self.viewing_key();
        let pk = &Scalar(viewing_key.sk_enc) * &ED25519_BASEPOINT;
        let pk_enc = pk.compress_montgomery().unwrap().to_bytes();

        PaymentAddress {
            a_pk: viewing_key.a_pk,
            pk_enc: pk_enc,
        }
    }

    /// Computes an incoming viewing key for this spending key.
    pub fn viewing_key(&self) -> ViewingKey {
        let mut a_pk = [0u8; 32];
        let mut sk_enc = [0u8; 32];

        pseudorandom_function_a_pk(&mut a_pk, &self.a_sk);
        pseudorandom_function_sk_enc(&mut sk_enc, &self.a_sk);
        clamp_curve25519(&mut sk_enc);

        ViewingKey {
            a_pk: a_pk,
            sk_enc: sk_enc,
        }
    }
}

pub fn pseudorandom_function_a_pk(a_pk: &mut [u8; 32], a_sk: &[u8; 32]) {
    pseudorandom_function(a_pk, a_sk, 0)
}

pub fn pseudorandom_function_sk_enc(sk_enc: &mut [u8; 32], a_sk: &[u8; 32]) {
    pseudorandom_function(sk_enc, a_sk, 1)
}

fn pseudorandom_function(dst: &mut [u8; 32], input: &[u8; 32], t: u8) {
    let mut data = [0u8; 64];
    data[..32].copy_from_slice(input);
    data[0] = 0xc0 | (data[0] & 0x0f);
    data[32] = t;
    sha256_compress(dst, &data);
}

fn clamp_curve25519(key: &mut [u8; 32]) {
    key[0] &= 248;
    key[31] &= 127;
    key[31] |= 64;
}

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

    static ENCODED: [(&str, &str, &str); 10] = [
        (
            "SKxt8pwrQipUL5KgZUcBAqyLj9R1YwMuRRR3ijGMCwCCqchmi8ut",
            "zcJLC7a3aRJohMNCVjSZQ8jFuofhAHJNAY4aX5soDkYfgNejzKnEZbucJmVibLWCwK8dyyfDhNhf3foXDDTouweC382LcX5",
            "ZiVKYQyUcyAJLKwcosSeDxkGRhygFdAPWsr3m8UgjC5X85yqNyLTtJJJYNH83Wf2AQKU6TZsd65MXBZLFj6eSCAFcnCFuVCFS",
        ),
        (
            "SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
            "zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
            "ZiVKfdhhmQ1fpXaxyW5zRXw4Dhg9cbKRgK7mNFoBLiKjiBZiHJYJTpV2gNMDMPY9sRC96vnKZcnTMSi65SKPyL4WNQNm9PT5H",
        ),
        (
            "SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
            "zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
            "ZiVKkMUGwx4GgtwxTedRHYewVVskWicz8APQgdcYmvUsiLYgSh3cLAa8TwiR3shyNngGbLiUbYMkZ8F1giXmmcED98rDMwNSG",
        ),
        (
            "SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
            "zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
            "ZiVKkeb8STw7kpJQsjRCQKovQBciPcfjkpajuuS25DTXSQSVasnq4BkyaMLBBxAkZ8fv6f18woWgaA8W7kGvYp1C1ESaWGjwV",
        ),
        (
            "SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
            "zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
            "ZiVKvpWQiDpxAvWTMLkjjSbCiBGc4kXhtkgAJfW1JVbCTUY4YaAVvVZzCz6wspG9qttciRFLEXm3HLQAmssFbUp9uPEkP3uu5"),
        (
            "SKxny894fJe2rmZjeuoE6GVfNkWoXfPp8337VrLLNWG56FfQtuS1",
            "zcbxovDeXGJJikZH5wQkcQvYx1gzsRt9mR5UnQir6NY8hhPHdgK7z7dE1vfa55Bq3JHJu7isfuWQGYrvMbLnud74z2vS4tS",
            "ZiVKr3bHGa79Kpy1zx2rC9xYd11tGvsY6fSvn2k1aEx97Z19WcMMW7KeSgE6VhJrAZ4RE6AmDU3oKqoXQiZfuHqzEBhW9o5UT",
        ),
        (
            "SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
            "zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
            "ZiVKfdhhmQ1fpXaxyW5zRXw4Dhg9cbKRgK7mNFoBLiKjiBZiHJYJTpV2gNMDMPY9sRC96vnKZcnTMSi65SKPyL4WNQNm9PT5H",
        ),
        (
            "SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
            "zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
            "ZiVKkMUGwx4GgtwxTedRHYewVVskWicz8APQgdcYmvUsiLYgSh3cLAa8TwiR3shyNngGbLiUbYMkZ8F1giXmmcED98rDMwNSG",
        ),
        (
            "SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
            "zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
            "ZiVKkeb8STw7kpJQsjRCQKovQBciPcfjkpajuuS25DTXSQSVasnq4BkyaMLBBxAkZ8fv6f18woWgaA8W7kGvYp1C1ESaWGjwV",
        ),
        (
            "SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
            "zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
            "ZiVKvpWQiDpxAvWTMLkjjSbCiBGc4kXhtkgAJfW1JVbCTUY4YaAVvVZzCz6wspG9qttciRFLEXm3HLQAmssFbUp9uPEkP3uu5",
        ),
    ];

    #[test]
    fn spending_keys_to_payment_addresses() {
        for &(spending_key_encoded, payment_address_encoded, _) in ENCODED.iter() {
            let mut a_sk = [0u8; 32];
            a_sk.copy_from_slice(
                &bs58::decode(spending_key_encoded).into_vec().unwrap()[2..34],
            );
            let spending_key = SpendingKey::new(a_sk);
            assert_eq!(spending_key_encoded, spending_key.to_string());
            let payment_address = spending_key.address();
            assert_eq!(payment_address_encoded, payment_address.to_string());
        }
    }

    #[test]
    fn spending_keys_to_viewing_keys() {
        for &(spending_key_encoded, _,viewing_key_encoded) in ENCODED.iter() {
            let mut a_sk = [0u8; 32];
            a_sk.copy_from_slice(
                &bs58::decode(spending_key_encoded).into_vec().unwrap()[2..34],
            );
            let spending_key = SpendingKey::new(a_sk);
            assert_eq!(spending_key_encoded, spending_key.to_string());
            let viewing_key = spending_key.viewing_key();
            assert_eq!(viewing_key_encoded, viewing_key.to_string());
        }
    }
}
<commit_msg>Whitespace.<commit_after>use curve25519_dalek::constants::ED25519_BASEPOINT;
use curve25519_dalek::scalar::Scalar;
use sha256::sha256_compress;
use std::fmt;
use util::bs58_encode_check;

const PAYMENT_ADDRESS_PREFIX_LENGTH: usize = 2;
const SPENDING_KEY_PREFIX_LENGTH: usize = 2;
const INVIEWING_KEY_PREFIX_LENGTH: usize = 3;
pub const PAYMENT_ADDRESS_PREFIX: [u8; PAYMENT_ADDRESS_PREFIX_LENGTH] = [0x16, 0x9a];
pub const SPENDING_KEY_PREFIX: [u8; SPENDING_KEY_PREFIX_LENGTH] = [0xab, 0x36];
pub const INVIEWING_KEY_PREFIX: [u8; INVIEWING_KEY_PREFIX_LENGTH] = [0xa8, 0xab, 0xd3];

/// A Zcash spending key.
pub struct SpendingKey {
    a_sk: [u8; 32], // 252 bits; first 4 bits are always 0
}

/// A Zcash incoming viewing key.
pub struct ViewingKey {
    pub a_pk: [u8; 32],
    pub sk_enc: [u8; 32],
}

/// A Zcash payment address.
pub struct PaymentAddress {
    pub a_pk: [u8; 32],
    pub pk_enc: [u8; 32],
}

impl fmt::Display for SpendingKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; SPENDING_KEY_PREFIX_LENGTH + 32];
        {
            let (prefix, rest) = data.split_at_mut(SPENDING_KEY_PREFIX_LENGTH);
            prefix.copy_from_slice(&SPENDING_KEY_PREFIX);
            rest.copy_from_slice(&self.a_sk);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl fmt::Display for ViewingKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; INVIEWING_KEY_PREFIX_LENGTH + 32 * 2];
        {
            let (prefix, rest) = data.split_at_mut(INVIEWING_KEY_PREFIX_LENGTH);
            prefix.copy_from_slice(&INVIEWING_KEY_PREFIX);
            rest[..32].copy_from_slice(&self.a_pk);
            rest[32..].copy_from_slice(&self.sk_enc);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl fmt::Display for PaymentAddress {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut data = [0u8; PAYMENT_ADDRESS_PREFIX_LENGTH + 32 * 2];
        {
            let (prefix, rest) = data.split_at_mut(PAYMENT_ADDRESS_PREFIX_LENGTH);
            prefix.copy_from_slice(&PAYMENT_ADDRESS_PREFIX);
            rest[..32].copy_from_slice(&self.a_pk);
            rest[32..].copy_from_slice(&self.pk_enc);
        }

        f.write_str(bs58_encode_check(&data).as_str())
    }
}

impl SpendingKey {
    /// Creates a new spending key instance for a given 252-bit *a_sk*.
    /// *a_sk* is represented by 32 bytes; the first four bits of the first byte are zeroed.
    pub fn new(a_sk_u256: [u8; 32]) -> SpendingKey {
        let mut a_sk_u252 = [0u8; 32];
        a_sk_u252.copy_from_slice(&a_sk_u256);
        a_sk_u252[0] &= 0x0f;

        SpendingKey { a_sk: a_sk_u252 }
    }

    /// Calculates the payment address for this spending key
    pub fn address(&self) -> PaymentAddress {
        let viewing_key = self.viewing_key();
        let pk = &Scalar(viewing_key.sk_enc) * &ED25519_BASEPOINT;
        let pk_enc = pk.compress_montgomery().unwrap().to_bytes();

        PaymentAddress {
            a_pk: viewing_key.a_pk,
            pk_enc: pk_enc,
        }
    }

    /// Computes an incoming viewing key for this spending key.
    pub fn viewing_key(&self) -> ViewingKey {
        let mut a_pk = [0u8; 32];
        let mut sk_enc = [0u8; 32];

        pseudorandom_function_a_pk(&mut a_pk, &self.a_sk);
        pseudorandom_function_sk_enc(&mut sk_enc, &self.a_sk);
        clamp_curve25519(&mut sk_enc);

        ViewingKey {
            a_pk: a_pk,
            sk_enc: sk_enc,
        }
    }
}

pub fn pseudorandom_function_a_pk(a_pk: &mut [u8; 32], a_sk: &[u8; 32]) {
    pseudorandom_function(a_pk, a_sk, 0)
}

pub fn pseudorandom_function_sk_enc(sk_enc: &mut [u8; 32], a_sk: &[u8; 32]) {
    pseudorandom_function(sk_enc, a_sk, 1)
}

fn pseudorandom_function(dst: &mut [u8; 32], input: &[u8; 32], t: u8) {
    let mut data = [0u8; 64];
    data[..32].copy_from_slice(input);
    data[0] = 0xc0 | (data[0] & 0x0f);
    data[32] = t;
    sha256_compress(dst, &data);
}

fn clamp_curve25519(key: &mut [u8; 32]) {
    key[0] &= 248;
    key[31] &= 127;
    key[31] |= 64;
}

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

    static ENCODED: [(&str, &str, &str); 10] = [
        (
            "SKxt8pwrQipUL5KgZUcBAqyLj9R1YwMuRRR3ijGMCwCCqchmi8ut",
            "zcJLC7a3aRJohMNCVjSZQ8jFuofhAHJNAY4aX5soDkYfgNejzKnEZbucJmVibLWCwK8dyyfDhNhf3foXDDTouweC382LcX5",
            "ZiVKYQyUcyAJLKwcosSeDxkGRhygFdAPWsr3m8UgjC5X85yqNyLTtJJJYNH83Wf2AQKU6TZsd65MXBZLFj6eSCAFcnCFuVCFS",
        ),
        (
            "SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
            "zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
            "ZiVKfdhhmQ1fpXaxyW5zRXw4Dhg9cbKRgK7mNFoBLiKjiBZiHJYJTpV2gNMDMPY9sRC96vnKZcnTMSi65SKPyL4WNQNm9PT5H",
        ),
        (
            "SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
            "zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
            "ZiVKkMUGwx4GgtwxTedRHYewVVskWicz8APQgdcYmvUsiLYgSh3cLAa8TwiR3shyNngGbLiUbYMkZ8F1giXmmcED98rDMwNSG",
        ),
        (
            "SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
            "zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
            "ZiVKkeb8STw7kpJQsjRCQKovQBciPcfjkpajuuS25DTXSQSVasnq4BkyaMLBBxAkZ8fv6f18woWgaA8W7kGvYp1C1ESaWGjwV",
        ),
        (
            "SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
            "zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
            "ZiVKvpWQiDpxAvWTMLkjjSbCiBGc4kXhtkgAJfW1JVbCTUY4YaAVvVZzCz6wspG9qttciRFLEXm3HLQAmssFbUp9uPEkP3uu5"),
        (
            "SKxny894fJe2rmZjeuoE6GVfNkWoXfPp8337VrLLNWG56FfQtuS1",
            "zcbxovDeXGJJikZH5wQkcQvYx1gzsRt9mR5UnQir6NY8hhPHdgK7z7dE1vfa55Bq3JHJu7isfuWQGYrvMbLnud74z2vS4tS",
            "ZiVKr3bHGa79Kpy1zx2rC9xYd11tGvsY6fSvn2k1aEx97Z19WcMMW7KeSgE6VhJrAZ4RE6AmDU3oKqoXQiZfuHqzEBhW9o5UT",
        ),
        (
            "SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
            "zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
            "ZiVKfdhhmQ1fpXaxyW5zRXw4Dhg9cbKRgK7mNFoBLiKjiBZiHJYJTpV2gNMDMPY9sRC96vnKZcnTMSi65SKPyL4WNQNm9PT5H",
        ),
        (
            "SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
            "zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
            "ZiVKkMUGwx4GgtwxTedRHYewVVskWicz8APQgdcYmvUsiLYgSh3cLAa8TwiR3shyNngGbLiUbYMkZ8F1giXmmcED98rDMwNSG",
        ),
        (
            "SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
            "zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
            "ZiVKkeb8STw7kpJQsjRCQKovQBciPcfjkpajuuS25DTXSQSVasnq4BkyaMLBBxAkZ8fv6f18woWgaA8W7kGvYp1C1ESaWGjwV",
        ),
        (
            "SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
            "zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
            "ZiVKvpWQiDpxAvWTMLkjjSbCiBGc4kXhtkgAJfW1JVbCTUY4YaAVvVZzCz6wspG9qttciRFLEXm3HLQAmssFbUp9uPEkP3uu5",
        ),
    ];

    #[test]
    fn spending_keys_to_payment_addresses() {
        for &(spending_key_encoded, payment_address_encoded, _) in ENCODED.iter() {
            let mut a_sk = [0u8; 32];
            a_sk.copy_from_slice(
                &bs58::decode(spending_key_encoded).into_vec().unwrap()[2..34],
            );
            let spending_key = SpendingKey::new(a_sk);
            assert_eq!(spending_key_encoded, spending_key.to_string());
            let payment_address = spending_key.address();
            assert_eq!(payment_address_encoded, payment_address.to_string());
        }
    }

    #[test]
    fn spending_keys_to_viewing_keys() {
        for &(spending_key_encoded, _, viewing_key_encoded) in ENCODED.iter() {
            let mut a_sk = [0u8; 32];
            a_sk.copy_from_slice(
                &bs58::decode(spending_key_encoded).into_vec().unwrap()[2..34],
            );
            let spending_key = SpendingKey::new(a_sk);
            assert_eq!(spending_key_encoded, spending_key.to_string());
            let viewing_key = spending_key.viewing_key();
            assert_eq!(viewing_key_encoded, viewing_key.to_string());
        }
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;

use bv::BitSlice;
use byteorder::{BigEndian, ByteOrder};

use crate::lz77;

#[derive(Debug, Copy, Clone)]
pub struct Colour {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

#[derive(Debug)]
pub struct PivImage {
    pub palette: Vec<Colour>,
    pixels: Vec<usize>,
}

impl PivImage {
    pub fn from_reader<T: Read>(reader: &mut T) -> Result<PivImage, io::Error> {
        let mut data: Vec<u8> = Vec::new();
        reader.read_to_end(&mut data)?;
        let header = PivImage::read_header(&data[..6]);

        let palette: Vec<Colour> =
            read_palette(header.bit_depth, &data[6..6 + (header.bit_depth * 2)]);

        let extracted = lz77::decompress(header.file_length, &data[6 + (header.bit_depth * 2)..])?;
        let pixels = PivImage::combine_bit_planes(&extracted);
        Ok(PivImage {
            palette: palette,
            pixels: pixels,
        })
    }

    pub fn to_rgba8(&self) -> Vec<u8> {
        let width = 320;
        let height = 240;

        let mut pixels: Vec<u8> = Vec::with_capacity(width * height * 4);
        for pel in self.pixels.iter() {
            let colour = &self.palette[*pel];
            pixels.extend([colour.r, colour.g, colour.b, colour.a].iter())
        }
        pixels
    }

    fn read_header(data: &[u8]) -> Header {
        let file_type = BigEndian::read_u16(&data[..2]);
        Header {
            file_length: BigEndian::read_u16(&data[4..6]),
            bit_depth: 1usize.wrapping_shl(file_type as u32),
        }
    }

    fn combine_bit_planes(data: &[u8]) -> Vec<usize> {
        let planes: Vec<BitSlice<u8>> =
            data.chunks(8000).map(|p| BitSlice::from_slice(p)).collect();

        let mut pixels: Vec<usize> = Vec::with_capacity(64000);
        for i in (0..64000).map(|x| 7 - (x % 8) + x / 8 * 8) {
            let mut sum = 0;
            for (j, plane) in planes.iter().enumerate() {
                sum += (plane[i] as usize) << j;
            }

            pixels.push(sum);
        }
        pixels
    }
}

#[derive(Debug, Copy, Clone)]
struct Header {
    file_length: u16,
    bit_depth: usize,
}

pub fn read_palette(bit_depth: usize, data: &[u8]) -> Vec<Colour> {
    let mut palette = vec![0; bit_depth];
    BigEndian::read_u16_into(&data, &mut palette);
    let palette: Vec<u16> = palette.iter().map(|pel| pel & 0x7fff).collect();

    palette
        .iter()
        .map(|pel| {
            let mut pel_bytes = [0u8; 2];
            BigEndian::write_u16(&mut pel_bytes, *pel);
            let mut alpha = 0u8;
            if *pel != 0u16 {
                alpha = 255;
            }
            Colour {
                r: (pel_bytes[0]) << 4,
                g: (((pel_bytes[1]) & 0xf0) >> 2) << 2,
                b: ((pel_bytes[1]) & 0x0f) << 4,
                a: alpha,
            }
        })
        .collect()
}
<commit_msg>Rust: only make first colour in a palette's alpha 0<commit_after>use std::io;
use std::io::prelude::*;

use bv::BitSlice;
use byteorder::{BigEndian, ByteOrder};

use crate::lz77;

#[derive(Debug, Copy, Clone)]
pub struct Colour {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

#[derive(Debug)]
pub struct PivImage {
    pub palette: Vec<Colour>,
    pixels: Vec<usize>,
}

impl PivImage {
    pub fn from_reader<T: Read>(reader: &mut T) -> Result<PivImage, io::Error> {
        let mut data: Vec<u8> = Vec::new();
        reader.read_to_end(&mut data)?;
        let header = PivImage::read_header(&data[..6]);

        let palette: Vec<Colour> =
            read_palette(header.bit_depth, &data[6..6 + (header.bit_depth * 2)]);

        let extracted = lz77::decompress(header.file_length, &data[6 + (header.bit_depth * 2)..])?;
        let pixels = PivImage::combine_bit_planes(&extracted);
        Ok(PivImage {
            palette: palette,
            pixels: pixels,
        })
    }

    pub fn to_rgba8(&self) -> Vec<u8> {
        let width = 320;
        let height = 240;

        let mut pixels: Vec<u8> = Vec::with_capacity(width * height * 4);
        for pel in self.pixels.iter() {
            let colour = &self.palette[*pel];
            pixels.extend([colour.r, colour.g, colour.b, colour.a].iter())
        }
        pixels
    }

    fn read_header(data: &[u8]) -> Header {
        let file_type = BigEndian::read_u16(&data[..2]);
        Header {
            file_length: BigEndian::read_u16(&data[4..6]),
            bit_depth: 1usize.wrapping_shl(file_type as u32),
        }
    }

    fn combine_bit_planes(data: &[u8]) -> Vec<usize> {
        let planes: Vec<BitSlice<u8>> =
            data.chunks(8000).map(|p| BitSlice::from_slice(p)).collect();

        let mut pixels: Vec<usize> = Vec::with_capacity(64000);
        for i in (0..64000).map(|x| 7 - (x % 8) + x / 8 * 8) {
            let mut sum = 0;
            for (j, plane) in planes.iter().enumerate() {
                sum += (plane[i] as usize) << j;
            }

            pixels.push(sum);
        }
        pixels
    }
}

#[derive(Debug, Copy, Clone)]
struct Header {
    file_length: u16,
    bit_depth: usize,
}

pub fn read_palette(bit_depth: usize, data: &[u8]) -> Vec<Colour> {
    let mut palette = vec![0; bit_depth];
    BigEndian::read_u16_into(&data, &mut palette);
    let palette: Vec<u16> = palette.iter().map(|pel| pel & 0x7fff).collect();

    let mut palette: Vec<Colour> = palette
        .iter()
        .map(|pel| {
            let mut pel_bytes = [0u8; 2];
            BigEndian::write_u16(&mut pel_bytes, *pel);
            Colour {
                r: (pel_bytes[0]) << 4,
                g: (((pel_bytes[1]) & 0xf0) >> 2) << 2,
                b: ((pel_bytes[1]) & 0x0f) << 4,
                a: 255,
            }
        })
        .collect();
    if let Some(first) = palette.get_mut(0) {
        first.a = 0;
    }
    palette
}
<|endoftext|>


<commit_before>use arena::{TypedArena, DroplessArena};
use std::mem;
use std::ptr;
use std::slice;
use std::cell::RefCell;
use std::marker::PhantomData;
use smallvec::SmallVec;

#[macro_export]
macro_rules! arena_types {
    ($macro:path, $args:tt, $tcx:lifetime) => (
        $macro!($args, [
            [] vtable_method: Option<(
                rustc::hir::def_id::DefId,
                rustc::ty::subst::SubstsRef<$tcx>
            )>,
            [few, decode] mir_keys: rustc::util::nodemap::DefIdSet,
            [decode] specialization_graph: rustc::traits::specialization_graph::Graph,
            [] region_scope_tree: rustc::middle::region::ScopeTree,
            [] item_local_set: rustc::util::nodemap::ItemLocalSet,
            [decode] mir_const_qualif: rustc_data_structures::bit_set::BitSet<rustc::mir::Local>,
            [] trait_impls_of: rustc::ty::trait_def::TraitImpls,
            [] dropck_outlives:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        rustc::traits::query::dropck_outlives::DropckOutlivesResult<'tcx>
                    >
                >,
            [] normalize_projection_ty:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        rustc::traits::query::normalize::NormalizationResult<'tcx>
                    >
                >,
            [] implied_outlives_bounds:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        Vec<rustc::traits::query::outlives_bounds::OutlivesBound<'tcx>>
                    >
                >,
            [] type_op_subtype:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, ()>
                >,
            [] type_op_normalize_poly_fn_sig:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::PolyFnSig<'tcx>>
                >,
            [] type_op_normalize_fn_sig:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::FnSig<'tcx>>
                >,
            [] type_op_normalize_predicate:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::Predicate<'tcx>>
                >,
            [] type_op_normalize_ty:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::Ty<'tcx>>
                >,
            [few] crate_inherent_impls: rustc::ty::CrateInherentImpls,
            [decode] borrowck: rustc::middle::borrowck::BorrowCheckResult,
            [few] upstream_monomorphizations:
                rustc::util::nodemap::DefIdMap<
                    rustc_data_structures::fx::FxHashMap<
                        rustc::ty::subst::SubstsRef<'tcx>,
                        rustc::hir::def_id::CrateNum
                    >
                >,
            [few] resolve_lifetimes: rustc::middle::resolve_lifetime::ResolveLifetimes,
            [decode] generic_predicates: rustc::ty::GenericPredicates<'tcx>,
        ], $tcx);
    )
}

macro_rules! arena_for_type {
    ([][$ty:ty]) => {
        TypedArena<$ty>
    };
    ([few $(, $attrs:ident)*][$ty:ty]) => {
        PhantomData<$ty>
    };
    ([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
        arena_for_type!([$($attrs),*]$args)
    };
}

macro_rules! declare_arena {
    ([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
        #[derive(Default)]
        pub struct Arena<$tcx> {
            dropless: DroplessArena,
            drop: DropArena,
            $($name: arena_for_type!($a[$ty]),)*
        }
    }
}

macro_rules! which_arena_for_type {
    ([][$arena:expr]) => {
        Some($arena)
    };
    ([few$(, $attrs:ident)*][$arena:expr]) => {
        None
    };
    ([$ignore:ident$(, $attrs:ident)*]$args:tt) => {
        which_arena_for_type!([$($attrs),*]$args)
    };
}

macro_rules! impl_arena_allocatable {
    ([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
        $(
            impl ArenaAllocatable for $ty {}
            unsafe impl<$tcx> ArenaField<$tcx> for $ty {
                #[inline]
                fn arena<'a>(_arena: &'a Arena<$tcx>) -> Option<&'a TypedArena<Self>> {
                    which_arena_for_type!($a[&_arena.$name])
                }
            }
        )*
    }
}

arena_types!(declare_arena, [], 'tcx);

arena_types!(impl_arena_allocatable, [], 'tcx);

pub trait ArenaAllocatable {}

impl<T: Copy> ArenaAllocatable for T {}

pub unsafe trait ArenaField<'tcx>: Sized {
    /// Returns a specific arena to allocate from.
    /// If None is returned, the DropArena will be used.
    fn arena<'a>(arena: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>>;
}

unsafe impl<'tcx, T> ArenaField<'tcx> for T {
    #[inline]
    default fn arena<'a>(_: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>> {
        panic!()
    }
}

impl<'tcx> Arena<'tcx> {
    #[inline]
    pub fn alloc<T: ArenaAllocatable>(&self, value: T) -> &mut T {
        if !mem::needs_drop::<T>() {
            return self.dropless.alloc(value);
        }
        match <T as ArenaField<'tcx>>::arena(self) {
            Some(arena) => arena.alloc(value),
            None => unsafe { self.drop.alloc(value) },
        }
    }

    #[inline]
    pub fn alloc_slice<T: Copy>(&self, value: &[T]) -> &mut [T] {
        if value.len() == 0 {
            return &mut []
        }
        self.dropless.alloc_slice(value)
    }

    pub fn alloc_from_iter<
        T: ArenaAllocatable,
        I: IntoIterator<Item = T>
    >(
        &'a self,
        iter: I
    ) -> &'a mut [T] {
        if !mem::needs_drop::<T>() {
            return self.dropless.alloc_from_iter(iter);
        }
        match <T as ArenaField<'tcx>>::arena(self) {
            Some(arena) => arena.alloc_from_iter(iter),
            None => unsafe { self.drop.alloc_from_iter(iter) },
        }
    }
}

/// Calls the destructor for an object when dropped.
struct DropType {
    drop_fn: unsafe fn(*mut u8),
    obj: *mut u8,
}

unsafe fn drop_for_type<T>(to_drop: *mut u8) {
    std::ptr::drop_in_place(to_drop as *mut T)
}

impl Drop for DropType {
    fn drop(&mut self) {
        unsafe {
            (self.drop_fn)(self.obj)
        }
    }
}

/// An arena which can be used to allocate any type.
/// Allocating in this arena is unsafe since the type system
/// doesn't know which types it contains. In order to
/// allocate safely, you must store a PhantomData<T>
/// alongside this arena for each type T you allocate.
#[derive(Default)]
struct DropArena {
    /// A list of destructors to run when the arena drops.
    /// Ordered so `destructors` gets dropped before the arena
    /// since its destructor can reference memory in the arena.
    destructors: RefCell<Vec<DropType>>,
    arena: DroplessArena,
}

impl DropArena {
    #[inline]
    unsafe fn alloc<T>(&self, object: T) -> &mut T {
        let mem = self.arena.alloc_raw(
            mem::size_of::<T>(),
            mem::align_of::<T>()
        ) as *mut _ as *mut T;
        // Write into uninitialized memory.
        ptr::write(mem, object);
        let result = &mut *mem;
        // Record the destructor after doing the allocation as that may panic
        // and would cause `object`'s destuctor to run twice if it was recorded before
        self.destructors.borrow_mut().push(DropType {
            drop_fn: drop_for_type::<T>,
            obj: result as *mut T as *mut u8,
        });
        result
    }

    #[inline]
    unsafe fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
        let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
        if vec.is_empty() {
            return &mut [];
        }
        let len = vec.len();

        let start_ptr = self.arena.alloc_raw(
            len.checked_mul(mem::size_of::<T>()).unwrap(),
            mem::align_of::<T>()
        ) as *mut _ as *mut T;

        let mut destructors = self.destructors.borrow_mut();
        // Reserve space for the destructors so we can't panic while adding them
        destructors.reserve(len);

        // Move the content to the arena by copying it and then forgetting
        // the content of the SmallVec
        vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
        mem::forget(vec.drain());

        // Record the destructors after doing the allocation as that may panic
        // and would cause `object`'s destuctor to run twice if it was recorded before
        for i in 0..len {
            destructors.push(DropType {
                drop_fn: drop_for_type::<T>,
                obj: start_ptr.offset(i as isize) as *mut u8,
            });
        }

        slice::from_raw_parts_mut(start_ptr, len)
    }
}
<commit_msg>Add a comment for arena_types!<commit_after>use arena::{TypedArena, DroplessArena};
use std::mem;
use std::ptr;
use std::slice;
use std::cell::RefCell;
use std::marker::PhantomData;
use smallvec::SmallVec;

/// This declares a list of types which can be allocated by `Arena`.
///
/// The `few` modifier will cause allocation to use the shared arena and recording the destructor.
/// This is faster and more memory efficient if there's only a few allocations of the type.
/// Leaving `few` out will cause the type to get its own dedicated `TypedArena` which is
/// faster and more memory efficient if there is lots of allocations.
///
/// Specifying the `decode` modifier will add decode impls for &T and &[T] where T is the type
/// listed. These impls will appear in the implement_ty_decoder! macro.
#[macro_export]
macro_rules! arena_types {
    ($macro:path, $args:tt, $tcx:lifetime) => (
        $macro!($args, [
            [] vtable_method: Option<(
                rustc::hir::def_id::DefId,
                rustc::ty::subst::SubstsRef<$tcx>
            )>,
            [few, decode] mir_keys: rustc::util::nodemap::DefIdSet,
            [decode] specialization_graph: rustc::traits::specialization_graph::Graph,
            [] region_scope_tree: rustc::middle::region::ScopeTree,
            [] item_local_set: rustc::util::nodemap::ItemLocalSet,
            [decode] mir_const_qualif: rustc_data_structures::bit_set::BitSet<rustc::mir::Local>,
            [] trait_impls_of: rustc::ty::trait_def::TraitImpls,
            [] dropck_outlives:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        rustc::traits::query::dropck_outlives::DropckOutlivesResult<'tcx>
                    >
                >,
            [] normalize_projection_ty:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        rustc::traits::query::normalize::NormalizationResult<'tcx>
                    >
                >,
            [] implied_outlives_bounds:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx,
                        Vec<rustc::traits::query::outlives_bounds::OutlivesBound<'tcx>>
                    >
                >,
            [] type_op_subtype:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, ()>
                >,
            [] type_op_normalize_poly_fn_sig:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::PolyFnSig<'tcx>>
                >,
            [] type_op_normalize_fn_sig:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::FnSig<'tcx>>
                >,
            [] type_op_normalize_predicate:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::Predicate<'tcx>>
                >,
            [] type_op_normalize_ty:
                rustc::infer::canonical::Canonical<'tcx,
                    rustc::infer::canonical::QueryResponse<'tcx, rustc::ty::Ty<'tcx>>
                >,
            [few] crate_inherent_impls: rustc::ty::CrateInherentImpls,
            [decode] borrowck: rustc::middle::borrowck::BorrowCheckResult,
            [few] upstream_monomorphizations:
                rustc::util::nodemap::DefIdMap<
                    rustc_data_structures::fx::FxHashMap<
                        rustc::ty::subst::SubstsRef<'tcx>,
                        rustc::hir::def_id::CrateNum
                    >
                >,
            [few] resolve_lifetimes: rustc::middle::resolve_lifetime::ResolveLifetimes,
            [decode] generic_predicates: rustc::ty::GenericPredicates<'tcx>,
        ], $tcx);
    )
}

macro_rules! arena_for_type {
    ([][$ty:ty]) => {
        TypedArena<$ty>
    };
    ([few $(, $attrs:ident)*][$ty:ty]) => {
        PhantomData<$ty>
    };
    ([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
        arena_for_type!([$($attrs),*]$args)
    };
}

macro_rules! declare_arena {
    ([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
        #[derive(Default)]
        pub struct Arena<$tcx> {
            dropless: DroplessArena,
            drop: DropArena,
            $($name: arena_for_type!($a[$ty]),)*
        }
    }
}

macro_rules! which_arena_for_type {
    ([][$arena:expr]) => {
        Some($arena)
    };
    ([few$(, $attrs:ident)*][$arena:expr]) => {
        None
    };
    ([$ignore:ident$(, $attrs:ident)*]$args:tt) => {
        which_arena_for_type!([$($attrs),*]$args)
    };
}

macro_rules! impl_arena_allocatable {
    ([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
        $(
            impl ArenaAllocatable for $ty {}
            unsafe impl<$tcx> ArenaField<$tcx> for $ty {
                #[inline]
                fn arena<'a>(_arena: &'a Arena<$tcx>) -> Option<&'a TypedArena<Self>> {
                    which_arena_for_type!($a[&_arena.$name])
                }
            }
        )*
    }
}

arena_types!(declare_arena, [], 'tcx);

arena_types!(impl_arena_allocatable, [], 'tcx);

pub trait ArenaAllocatable {}

impl<T: Copy> ArenaAllocatable for T {}

pub unsafe trait ArenaField<'tcx>: Sized {
    /// Returns a specific arena to allocate from.
    /// If None is returned, the DropArena will be used.
    fn arena<'a>(arena: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>>;
}

unsafe impl<'tcx, T> ArenaField<'tcx> for T {
    #[inline]
    default fn arena<'a>(_: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>> {
        panic!()
    }
}

impl<'tcx> Arena<'tcx> {
    #[inline]
    pub fn alloc<T: ArenaAllocatable>(&self, value: T) -> &mut T {
        if !mem::needs_drop::<T>() {
            return self.dropless.alloc(value);
        }
        match <T as ArenaField<'tcx>>::arena(self) {
            Some(arena) => arena.alloc(value),
            None => unsafe { self.drop.alloc(value) },
        }
    }

    #[inline]
    pub fn alloc_slice<T: Copy>(&self, value: &[T]) -> &mut [T] {
        if value.len() == 0 {
            return &mut []
        }
        self.dropless.alloc_slice(value)
    }

    pub fn alloc_from_iter<
        T: ArenaAllocatable,
        I: IntoIterator<Item = T>
    >(
        &'a self,
        iter: I
    ) -> &'a mut [T] {
        if !mem::needs_drop::<T>() {
            return self.dropless.alloc_from_iter(iter);
        }
        match <T as ArenaField<'tcx>>::arena(self) {
            Some(arena) => arena.alloc_from_iter(iter),
            None => unsafe { self.drop.alloc_from_iter(iter) },
        }
    }
}

/// Calls the destructor for an object when dropped.
struct DropType {
    drop_fn: unsafe fn(*mut u8),
    obj: *mut u8,
}

unsafe fn drop_for_type<T>(to_drop: *mut u8) {
    std::ptr::drop_in_place(to_drop as *mut T)
}

impl Drop for DropType {
    fn drop(&mut self) {
        unsafe {
            (self.drop_fn)(self.obj)
        }
    }
}

/// An arena which can be used to allocate any type.
/// Allocating in this arena is unsafe since the type system
/// doesn't know which types it contains. In order to
/// allocate safely, you must store a PhantomData<T>
/// alongside this arena for each type T you allocate.
#[derive(Default)]
struct DropArena {
    /// A list of destructors to run when the arena drops.
    /// Ordered so `destructors` gets dropped before the arena
    /// since its destructor can reference memory in the arena.
    destructors: RefCell<Vec<DropType>>,
    arena: DroplessArena,
}

impl DropArena {
    #[inline]
    unsafe fn alloc<T>(&self, object: T) -> &mut T {
        let mem = self.arena.alloc_raw(
            mem::size_of::<T>(),
            mem::align_of::<T>()
        ) as *mut _ as *mut T;
        // Write into uninitialized memory.
        ptr::write(mem, object);
        let result = &mut *mem;
        // Record the destructor after doing the allocation as that may panic
        // and would cause `object`'s destuctor to run twice if it was recorded before
        self.destructors.borrow_mut().push(DropType {
            drop_fn: drop_for_type::<T>,
            obj: result as *mut T as *mut u8,
        });
        result
    }

    #[inline]
    unsafe fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
        let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
        if vec.is_empty() {
            return &mut [];
        }
        let len = vec.len();

        let start_ptr = self.arena.alloc_raw(
            len.checked_mul(mem::size_of::<T>()).unwrap(),
            mem::align_of::<T>()
        ) as *mut _ as *mut T;

        let mut destructors = self.destructors.borrow_mut();
        // Reserve space for the destructors so we can't panic while adding them
        destructors.reserve(len);

        // Move the content to the arena by copying it and then forgetting
        // the content of the SmallVec
        vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
        mem::forget(vec.drain());

        // Record the destructors after doing the allocation as that may panic
        // and would cause `object`'s destuctor to run twice if it was recorded before
        for i in 0..len {
            destructors.push(DropType {
                drop_fn: drop_for_type::<T>,
                obj: start_ptr.offset(i as isize) as *mut u8,
            });
        }

        slice::from_raw_parts_mut(start_ptr, len)
    }
}
<|endoftext|>


<commit_before>use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ops::Deref;
use std::ptr;

use errors::*;

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

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

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

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

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

pub struct OpenContext {
    context: Context,
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        Ok(OpenContext { context: self })
    }
}

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

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

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

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

        Ok(Self { ptr })
    }
}

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

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

        let mut ptr = ptr::null_mut();

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

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

        Ok(Self { ptr })
    }
}

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

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

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

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

    if codec.is_null() {
        Ok(None)
    } else {
        Ok(Some(Codec { ptr: codec }))
    }
}
<commit_msg>Allow dead code in avcodec<commit_after>#![allow(dead_code)]
use ffmpeg_sys;
use libc::*;
use std::cmp;
use std::ffi::{ CStr, CString };
use std::ops::Deref;
use std::ptr;

use errors::*;

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

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

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

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

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

pub struct OpenContext {
    context: Context,
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        Ok(OpenContext { context: self })
    }
}

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

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

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

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

        Ok(Self { ptr })
    }
}

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

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

        let mut ptr = ptr::null_mut();

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

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

        Ok(Self { ptr })
    }
}

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

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

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

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

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


<commit_before>use crate::io::ErrorKind;
use std::io::{self, BufWriter, Write};
use std::process::exit;
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::thread::{self, JoinHandle};

use crate::{constants::*, refiner};
use regex::Regex;

const HUNK_HEADER: &str = "\x1b[36m"; // Cyan

lazy_static! {
    static ref STATIC_HEADER_PREFIXES: Vec<(&'static str, &'static str)> = vec![
        ("diff ", BOLD),
        ("index ", BOLD),
        ("--- ", BOLD),
        ("+++ ", BOLD),
        ("@@ ", HUNK_HEADER),
    ];
    static ref ANSI_COLOR_REGEX: Regex = Regex::new("\x1b[^m]*m").unwrap();
}

#[must_use]
fn get_fixed_highlight(line: &str) -> &str {
    for static_header_prefix in STATIC_HEADER_PREFIXES.iter() {
        let prefix = static_header_prefix.0;
        if line.starts_with(prefix) {
            return static_header_prefix.1;
        }
    }

    return "";
}

fn print<W: io::Write + Send>(stream: &mut BufWriter<W>, text: &str) {
    if let Err(error) = stream.write_all(text.as_bytes()) {
        if error.kind() == ErrorKind::BrokenPipe {
            // This is fine, somebody probably just quit their pager before it
            // was done reading our output.
            exit(0);
        }

        panic!("Error writing diff to pager: {:?}", error);
    }
}

struct StringFuture {
    // This field is only valid if we're done with the result_receiver (next
    // field)
    result: String,

    // If available, get() will await a result on this receiver, then populate
    // the result field and return it
    result_receiver: Option<Receiver<String>>,
}

impl StringFuture {
    pub fn from_string(result: String) -> StringFuture {
        return StringFuture {
            result,
            result_receiver: None,
        };
    }

    pub fn from_oldnew(old_text: String, new_text: String) -> StringFuture {
        // Create a String channel
        let (sender, receiver): (SyncSender<String>, Receiver<String>) = sync_channel(1);

        // Start diffing in a thread
        thread::spawn(move || {
            let mut result = String::new();
            for line in refiner::format(&old_text, &new_text) {
                result.push_str(&line);
                result.push('\n');
            }

            // Done, channel the result!
            sender.send(result).unwrap();
        });

        return StringFuture {
            result: "".to_string(),
            result_receiver: Some(receiver),
        };
    }

    pub fn is_empty(&mut self) -> bool {
        return self.get().is_empty();
    }

    pub fn get(&mut self) -> &str {
        // If the result is still pending...
        if let Some(receiver) = &self.result_receiver {
            // ... wait for it
            self.result = receiver.recv().unwrap();
            self.result_receiver = None;
        }

        return &self.result;
    }
}

pub struct LineCollector {
    old_text: String,
    new_text: String,
    plain_text: String,
    consumer_thread: Option<JoinHandle<()>>,

    // FIXME: I'd rather have had a SyncSender of some trait here. That would
    // enable us to have two separate result implementations, one which just
    // returns a string and another that does a background computation first.
    // But I failed to figure out how when I tried, more Googling needed!
    queue_putter: SyncSender<StringFuture>,
}

impl Drop for LineCollector {
    fn drop(&mut self) {
        // Flush any outstanding lines. This can be done in any order, at most
        // one of them is going to do anything anyway.
        self.drain_oldnew();
        self.drain_plain();

        // Tell the consumer thread to drain and quit. Sending an empty string
        // like this is the secret handshake for requesting a shutdown.
        self.queue_putter
            .send(StringFuture::from_string("".to_string()))
            .unwrap();

        // Wait for the consumer thread to finish
        // https://stackoverflow.com/q/57670145/473672
        self.consumer_thread.take().map(JoinHandle::join);
    }
}

impl LineCollector {
    pub fn new<W: io::Write + Send + 'static>(output: W) -> LineCollector {
        // The queue will be bounded to 2x the number of logical CPUs.
        //
        // 1x is for the entries that need CPU time for diffing.
        //
        // The other 1x is for the entries that just contain text to print and
        // won't need any background processing time.
        let queue_size = num_cpus::get() * 2;

        // Allocate a queue where we can push our futures to the consumer thread
        let (queue_putter, queue_getter): (SyncSender<StringFuture>, Receiver<StringFuture>) =
            sync_channel(queue_size);

        // This thread takes futures and prints their results
        let consumer = thread::spawn(move || {
            let mut output = BufWriter::new(output);

            loop {
                if let Ok(mut print_me) = queue_getter.recv() {
                    if print_me.is_empty() {
                        // Secret handshake received, done!
                        break;
                    }
                    print(&mut output, print_me.get());
                }
            }
        });

        return LineCollector {
            old_text: String::from(""),
            new_text: String::from(""),
            plain_text: String::from(""),
            consumer_thread: Some(consumer),
            queue_putter,
        };
    }

    fn drain_oldnew(&mut self) {
        if self.old_text.is_empty() && self.new_text.is_empty() {
            return;
        }

        self.queue_putter
            .send(StringFuture::from_oldnew(
                self.old_text.clone(),
                self.new_text.clone(),
            ))
            .unwrap();

        self.old_text.clear();
        self.new_text.clear();
    }

    fn drain_plain(&mut self) {
        if self.plain_text.is_empty() {
            return;
        }

        // Create an already-resolved future returning this text, then store
        // that future in our queue.
        self.queue_putter
            .send(StringFuture::from_string(String::from(&self.plain_text)))
            .unwrap();

        self.plain_text.clear();
    }

    fn consume_plain_line(&mut self, line: &str) {
        self.drain_oldnew();
        self.plain_text.push_str(line);
        self.plain_text.push('\n');
    }

    fn consume_old_line(&mut self, line: &str) {
        self.drain_plain();
        self.old_text.push_str(&line[1..]);
        self.old_text.push('\n');
    }

    fn consume_new_line(&mut self, line: &str) {
        self.drain_plain();
        self.new_text.push_str(&line[1..]);
        self.new_text.push('\n');
    }

    fn consume_no_eof_newline_marker(&mut self) {
        if !self.new_text.is_empty() {
            // New section comes after old, so if we get in here it's a new
            // section that doesn't end in a newline. Remove its trailing
            // newline.
            assert!(self.new_text.pop().unwrap() == '\n');
            return;
        }

        if !self.old_text.is_empty() {
            // Old text doesn't end in a newline, remove its trailing newline
            assert!(self.old_text.pop().unwrap() == '\n');
            return;
        }

        // It's a piece of unchanged text that doesn't end in a newline, just
        // consume the colorized marker as plain text
        self.consume_plain_line(&format!(
            "{}{}{}",
            NO_EOF_NEWLINE_COLOR, &NO_EOF_NEWLINE_MARKER, NORMAL
        ))
    }

    pub fn consume_line(&mut self, line: String) {
        // Strip out incoming ANSI formatting. This enables us to highlight
        // already-colored input.
        let line = ANSI_COLOR_REGEX.replace_all(&line, "");

        let fixed_highlight = get_fixed_highlight(&line);
        if !fixed_highlight.is_empty() {
            self.consume_plain_line(&format!("{}{}{}", fixed_highlight, line, NORMAL));
            return;
        }

        if line.is_empty() {
            self.consume_plain_line("");
            return;
        }

        if line.starts_with('-') {
            self.consume_old_line(&line);
            return;
        }

        if line.starts_with('+') {
            self.consume_new_line(&line);
            return;
        }

        if line == NO_EOF_NEWLINE_MARKER {
            self.consume_no_eof_newline_marker();
            return;
        }

        self.consume_plain_line(&line);
    }
}
<commit_msg>Improve docs<commit_after>use crate::io::ErrorKind;
use std::io::{self, BufWriter, Write};
use std::process::exit;
use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
use std::thread::{self, JoinHandle};

use crate::{constants::*, refiner};
use regex::Regex;

const HUNK_HEADER: &str = "\x1b[36m"; // Cyan

lazy_static! {
    static ref STATIC_HEADER_PREFIXES: Vec<(&'static str, &'static str)> = vec![
        ("diff ", BOLD),
        ("index ", BOLD),
        ("--- ", BOLD),
        ("+++ ", BOLD),
        ("@@ ", HUNK_HEADER),
    ];
    static ref ANSI_COLOR_REGEX: Regex = Regex::new("\x1b[^m]*m").unwrap();
}

#[must_use]
fn get_fixed_highlight(line: &str) -> &str {
    for static_header_prefix in STATIC_HEADER_PREFIXES.iter() {
        let prefix = static_header_prefix.0;
        if line.starts_with(prefix) {
            return static_header_prefix.1;
        }
    }

    return "";
}

fn print<W: io::Write + Send>(stream: &mut BufWriter<W>, text: &str) {
    if let Err(error) = stream.write_all(text.as_bytes()) {
        if error.kind() == ErrorKind::BrokenPipe {
            // This is fine, somebody probably just quit their pager before it
            // was done reading our output.
            exit(0);
        }

        panic!("Error writing diff to pager: {:?}", error);
    }
}

/**
A StringFuture can perform diffing in a background thread.

Doing get() on a future that isn't done yet will block until the result is
available.
*/
struct StringFuture {
    // This field is only valid if we're done with the result_receiver (next
    // field)
    result: String,

    // If available, get() will await a result on this receiver, then populate
    // the result field and return it
    result_receiver: Option<Receiver<String>>,
}

impl StringFuture {
    /// Create an already-finished future
    pub fn from_string(result: String) -> StringFuture {
        return StringFuture {
            result,
            result_receiver: None,
        };
    }

    /// Call get() to get the result of this diff
    pub fn from_oldnew(old_text: String, new_text: String) -> StringFuture {
        // Create a String channel
        let (sender, receiver): (SyncSender<String>, Receiver<String>) = sync_channel(1);

        // Start diffing in a thread
        thread::spawn(move || {
            let mut result = String::new();
            for line in refiner::format(&old_text, &new_text) {
                result.push_str(&line);
                result.push('\n');
            }

            // Done, channel the result!
            sender.send(result).unwrap();
        });

        return StringFuture {
            result: "".to_string(),
            result_receiver: Some(receiver),
        };
    }

    pub fn is_empty(&mut self) -> bool {
        return self.get().is_empty();
    }

    pub fn get(&mut self) -> &str {
        // If the result is still pending...
        if let Some(receiver) = &self.result_receiver {
            // ... wait for it
            self.result = receiver.recv().unwrap();
            self.result_receiver = None;
        }

        return &self.result;
    }
}

/**
The way this thing works from the outside is that you initialize it with an
output stream, you pass it one line of input at a time, and it writes
formatted lines to the output stream.

From the inside, it will collect blocks of either diff lines or not-diff-lines.

The not-diff-lines blocks will be enqueued for printing by the printing thread.

The diff lines blocks will also be enqueued for printing, but the actual diffing
will happen in background threads.
*/
pub struct LineCollector {
    old_text: String,
    new_text: String,
    plain_text: String,
    consumer_thread: Option<JoinHandle<()>>,

    // FIXME: I'd rather have had a SyncSender of some trait here. That would
    // enable us to have two separate result implementations, one which just
    // returns a string and another that does a background computation first.
    // But I failed to figure out how when I tried, more Googling needed!
    queue_putter: SyncSender<StringFuture>,
}

impl Drop for LineCollector {
    fn drop(&mut self) {
        // Flush any outstanding lines. This can be done in any order, at most
        // one of them is going to do anything anyway.
        self.drain_oldnew();
        self.drain_plain();

        // Tell the consumer thread to drain and quit. Sending an empty string
        // like this is the secret handshake for requesting a shutdown.
        self.queue_putter
            .send(StringFuture::from_string("".to_string()))
            .unwrap();

        // Wait for the consumer thread to finish
        // https://stackoverflow.com/q/57670145/473672
        self.consumer_thread.take().map(JoinHandle::join);
    }
}

impl LineCollector {
    pub fn new<W: io::Write + Send + 'static>(output: W) -> LineCollector {
        // The queue will be bounded to 2x the number of logical CPUs.
        //
        // 1x is for the entries that need CPU time for diffing.
        //
        // The other 1x is for the entries that just contain text to print and
        // won't need any background processing time.
        let queue_size = num_cpus::get() * 2;

        // Allocate a queue where we can push our futures to the consumer thread
        let (queue_putter, queue_getter): (SyncSender<StringFuture>, Receiver<StringFuture>) =
            sync_channel(queue_size);

        // This thread takes futures and prints their results
        let consumer = thread::spawn(move || {
            let mut output = BufWriter::new(output);

            loop {
                if let Ok(mut print_me) = queue_getter.recv() {
                    if print_me.is_empty() {
                        // Secret handshake received, done!
                        break;
                    }
                    print(&mut output, print_me.get());
                }
            }
        });

        return LineCollector {
            old_text: String::from(""),
            new_text: String::from(""),
            plain_text: String::from(""),
            consumer_thread: Some(consumer),
            queue_putter,
        };
    }

    fn drain_oldnew(&mut self) {
        if self.old_text.is_empty() && self.new_text.is_empty() {
            return;
        }

        self.queue_putter
            .send(StringFuture::from_oldnew(
                self.old_text.clone(),
                self.new_text.clone(),
            ))
            .unwrap();

        self.old_text.clear();
        self.new_text.clear();
    }

    fn drain_plain(&mut self) {
        if self.plain_text.is_empty() {
            return;
        }

        // Enqueue an already-resolved future
        self.queue_putter
            .send(StringFuture::from_string(String::from(&self.plain_text)))
            .unwrap();

        self.plain_text.clear();
    }

    fn consume_plain_line(&mut self, line: &str) {
        self.drain_oldnew();
        self.plain_text.push_str(line);
        self.plain_text.push('\n');
    }

    fn consume_old_line(&mut self, line: &str) {
        self.drain_plain();
        self.old_text.push_str(&line[1..]);
        self.old_text.push('\n');
    }

    fn consume_new_line(&mut self, line: &str) {
        self.drain_plain();
        self.new_text.push_str(&line[1..]);
        self.new_text.push('\n');
    }

    fn consume_no_eof_newline_marker(&mut self) {
        if !self.new_text.is_empty() {
            // New section comes after old, so if we get in here it's a new
            // section that doesn't end in a newline. Remove its trailing
            // newline.
            assert!(self.new_text.pop().unwrap() == '\n');
            return;
        }

        if !self.old_text.is_empty() {
            // Old text doesn't end in a newline, remove its trailing newline
            assert!(self.old_text.pop().unwrap() == '\n');
            return;
        }

        // It's a piece of unchanged text that doesn't end in a newline, just
        // consume the colorized marker as plain text
        self.consume_plain_line(&format!(
            "{}{}{}",
            NO_EOF_NEWLINE_COLOR, &NO_EOF_NEWLINE_MARKER, NORMAL
        ))
    }

    pub fn consume_line(&mut self, line: String) {
        // Strip out incoming ANSI formatting. This enables us to highlight
        // already-colored input.
        let line = ANSI_COLOR_REGEX.replace_all(&line, "");

        let fixed_highlight = get_fixed_highlight(&line);
        if !fixed_highlight.is_empty() {
            self.consume_plain_line(&format!("{}{}{}", fixed_highlight, line, NORMAL));
            return;
        }

        if line.is_empty() {
            self.consume_plain_line("");
            return;
        }

        if line.starts_with('-') {
            self.consume_old_line(&line);
            return;
        }

        if line.starts_with('+') {
            self.consume_new_line(&line);
            return;
        }

        if line == NO_EOF_NEWLINE_MARKER {
            self.consume_no_eof_newline_marker();
            return;
        }

        self.consume_plain_line(&line);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;
use std::str;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait BAMRead {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct BAMReader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMReader {
    /// Create a new BAMReader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: &P) -> Self {
        let f = bgzf_open(path, b"r");
        let header = unsafe { htslib::bam_hdr_read(f) };
        BAMReader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl BAMRead for BAMReader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for BAMReader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// A BAM writer.
pub struct BAMWriter {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl BAMWriter {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");
        let header_record = unsafe {
            let header_string = header.to_bytes();
            println!("{}", str::from_utf8(&header_string).unwrap());
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }
        BAMWriter { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }
        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }
        BAMWriter { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for BAMWriter {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: BAMReader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use super::record::*;
    use super::header::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = BAMReader::new(&"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));

        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }

    #[test]
    fn test_write() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bampath = tmp.path().join("test.bam");
        {
            let mut bam = BAMWriter::new(
                &bampath,
                Header::new().push_record(
                    HeaderRecord::new(b"SQ").push_tag(b"SN", &"chr1")
                                            .push_tag(b"LN", &15072423)
                )
            );

            let mut rec = record::Record::new();
            rec.set(qname, &cigar, seq, qual);
            rec.push_aux(b"NM", &Aux::Integer(15));

            bam.write(&mut rec).ok().expect("Failed to write record.");
        }
        {
            let bam = BAMReader::new(&bampath);

            let mut rec = record::Record::new();
            bam.read(&mut rec).ok().expect("Failed to read record.");

            assert_eq!(rec.qname(), qname);
            assert_eq!(rec.cigar(), cigar);
            assert_eq!(rec.seq().as_bytes(), seq);
            assert_eq!(rec.qual(), qual);
            assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<commit_msg>Refactoring.<commit_after>// Copyright 2014 Christopher Schröder, Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

pub mod record;
pub mod header;

use std::ffi;
use std::path;
use std::ffi::AsOsStr;
use std::os::unix::prelude::OsStrExt;
use std::str;

use htslib;

/// A trait for a BAM reader with a read method.
pub trait Read {
    /// Read next BAM record into given record.
    /// Use this method in combination with a single allocated record to avoid the reallocations
    /// occurring with the iterator.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to be filled
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError>;
}

/// A BAM reader.
pub struct Reader {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl Reader {
    /// Create a new Reader.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    pub fn new<P: path::AsPath>(path: &P) -> Self {
        let f = bgzf_open(path, b"r");
        let header = unsafe { htslib::bam_hdr_read(f) };
        Reader { f : f, header : header }
    }

    /// Iterator over the records of the BAM file.
    pub fn records(self) -> Records {
        Records { bam: self }
    }
}

impl Read for Reader {
    fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::bam_read1(self.f, &mut record.inner) } {
            -1 => Err(ReadError::EOF),
            -2 => Err(ReadError::Truncated),
            -4 => Err(ReadError::Invalid),
            _  => Ok(())
        }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// A BAM writer.
pub struct Writer {
    f: *mut htslib::Struct_BGZF,
    header: *mut htslib::bam_hdr_t,
}

impl Writer {
    /// Create a new BAM file.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `header` - header definition to use
    pub fn new<P: path::AsPath>(path: &P, header: &header::Header) -> Self {
        let f = bgzf_open(path, b"w");

        let header_record = unsafe {
            let header_string = header.to_bytes();
            println!("{}", str::from_utf8(&header_string).unwrap());
            htslib::sam_hdr_parse(
                header_string.len() as i32,
                ffi::CString::new(header_string).unwrap().as_ptr()
            )
        };
        unsafe { htslib::bam_hdr_write(f, header_record); }

        Writer { f: f, header: header_record }
    }

    /// Create a new BAM file from template.
    ///
    /// # Arguments
    ///
    /// * `path` - the path. Use "-" for stdin.
    /// * `template` - the template BAM. Use "-" for stdin.
    pub fn with_template<P: path::AsPath, T: path::AsPath>(path: &P, template: &T) -> Self {
        let t = bgzf_open(template, b"r");
        let header = unsafe { htslib::bam_hdr_read(t) };
        unsafe { htslib::bgzf_close(t); }

        let f = bgzf_open(path, b"w");
        unsafe { htslib::bam_hdr_write(f, header); }

        Writer { f: f, header: header }
    }

    /// Write record to BAM.
    ///
    /// # Arguments
    ///
    /// * `record` - the record to write
    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::bam_write1(self.f, &record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for Writer {
    fn drop(&mut self) {
        unsafe {
            htslib::bam_hdr_destroy(self.header);
            htslib::bgzf_close(self.f);
        }
    }
}

/// Iterator over the records of a BAM.
pub struct Records {
    bam: Reader
}

impl Iterator for Records {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.bam.read(&mut record) {
            Err(ReadError::EOF) => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

pub enum ReadError {
    Truncated,
    Invalid,
    EOF,
}

/// Wrapper for opening a BAM file.
fn bgzf_open<P: path::AsPath>(path: &P, mode: &[u8]) -> *mut htslib::Struct_BGZF {
    unsafe {
        htslib::bgzf_open(
            path.as_path().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use super::record::*;
    use super::header::*;
    use std::str;

    #[test]
    fn test_read() {
        let names = [b"I", b"II.14978392", b"III", b"IV", b"V", b"VI"];
        let flags = [16u16, 16u16, 16u16, 16u16, 16u16, 2048u16];
        let seqs = [
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAAGCCTAA",
            b"ACTAAGCCTAAGCCTAAGCCTAAGCCAATTATCGATTTCTGAAAAAATTATCGAATTTTCTAGAAATTTTGCAAATTTTTTCATAAAATTATCGATTTTA",
        ];
        let cigars = [
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)],
            [Cigar::Match(27), Cigar::Del(100000), Cigar::Match(73)],
        ];

        let bam = Reader::new(&"test.bam");

        for (i, record) in bam.records().enumerate() {
            let rec = record.ok().expect("Expected valid record");
            println!("{}", str::from_utf8(rec.qname()).ok().unwrap());
            assert_eq!(rec.qname(), names[i]);
            assert_eq!(rec.flags(), flags[i]);
            assert_eq!(rec.seq().as_bytes(), seqs[i]);
            assert_eq!(rec.cigar(), cigars[i]);
        }
    }

    #[test]
    fn test_set_record() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let mut rec = record::Record::new();
        rec.set_reverse();
        rec.set(qname, &cigar, seq, qual);
        rec.push_aux(b"NM", &Aux::Integer(15));

        assert_eq!(rec.qname(), qname);
        assert_eq!(rec.cigar(), cigar);
        assert_eq!(rec.seq().as_bytes(), seq);
        assert_eq!(rec.qual(), qual);
        assert!(rec.is_reverse());
        assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
    }

    #[test]
    fn test_write() {
        let qname = b"I";
        let cigar = [Cigar::Match(27), Cigar::Del(1), Cigar::Match(73)];
        let seq =  b"CCTAGCCCTAACCCTAACCCTAACCCTAGCCTAAGCCTAAGCCTAAGCCTAAGC";
        let qual = b"JJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJJ";

        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bampath = tmp.path().join("test.bam");
        {
            let mut bam = Writer::new(
                &bampath,
                Header::new().push_record(
                    HeaderRecord::new(b"SQ").push_tag(b"SN", &"chr1")
                                            .push_tag(b"LN", &15072423)
                )
            );

            let mut rec = record::Record::new();
            rec.set(qname, &cigar, seq, qual);
            rec.push_aux(b"NM", &Aux::Integer(15));

            bam.write(&mut rec).ok().expect("Failed to write record.");
        }
        {
            let bam = Reader::new(&bampath);

            let mut rec = record::Record::new();
            bam.read(&mut rec).ok().expect("Failed to read record.");

            assert_eq!(rec.qname(), qname);
            assert_eq!(rec.cigar(), cigar);
            assert_eq!(rec.seq().as_bytes(), seq);
            assert_eq!(rec.qual(), qual);
            assert_eq!(rec.aux(b"NM").unwrap(), Aux::Integer(15));
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<|endoftext|>