<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>Replace tabs with spaces.<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 request::{Request, ToRequest};
use service::Service;
use application::Application;
use user::User;

#[derive(Debug)]
pub enum Type {
    Authorize,
    AuthRep,
    Report
}

impl Type {
    pub fn method(&self) -> String {
        use self::Type::*;
        match *self {
            Report => "POST".to_owned(),
            AuthRep | Authorize => "GET".to_owned(),
        }
    }
}

#[derive(Debug)]
pub struct Info<'service, 'app, 'user> {
    kind: Type,
    service: &'service Service,
    application: &'app Application,
    user: Option<&'user User>,
}

// TODO: define endpoints for oauth
const AUTH_ENDPOINT: &str = "/transactions/authorize.xml";
const AUTHREP_ENDPOINT: &str = "/transactions/authrep.xml";
const REPORT_ENDPOINT: &str = "/transactions.xml";

impl<'service, 'app, 'user> Info<'service, 'app, 'user> {
    pub fn new(kind: Type, service: &'service Service, application: &'app Application, user: Option<&'user User>) -> Self {
        Self { kind, service, application, user }
    }

    fn endpoint(&self) -> &str {
        use self::Type::*;

        match self.kind {
            Authorize => AUTH_ENDPOINT,
            AuthRep => AUTHREP_ENDPOINT,
            Report => REPORT_ENDPOINT,
        }
    }

    fn path(&self) -> String {
        use request::ToParams;

        let mut params: Vec<(&str, &str)> = Vec::new();
        params.extend(self.service.to_params());
        params.extend(self.application.to_params());

        if let Some(user_params) = self.user.as_ref() {
            params.extend(user_params.to_params());
        }

        let params = params.into_iter()
            .map(|(param, value)| param.to_owned() + "=" + value)
            .collect::<Vec<String>>()
            .join("&");

        self.endpoint().to_owned() + "?" + params.as_str()
    }
}

impl<'service, 'app, 'user> ToRequest for Info<'service, 'app, 'user> {
    fn to_request(&self) -> Request {
        Request::new(self.kind.method(), self.path(), None)
    }
}
<commit_msg>apicall: add OAuth endpoint paths<commit_after>use request::{Request, ToRequest};
use service::Service;
use application::Application;
use user::User;

#[derive(Debug)]
pub enum Type {
    Authorize,
    AuthRep,
    Report
}

impl Type {
    pub fn method(&self) -> String {
        use self::Type::*;
        match *self {
            Report => "POST".to_owned(),
            AuthRep | Authorize => "GET".to_owned(),
        }
    }
}

#[derive(Debug)]
pub struct Info<'service, 'app, 'user> {
    kind: Type,
    service: &'service Service,
    application: &'app Application,
    user: Option<&'user User>,
}

const AUTHORIZE_ENDPOINT: &str = "/transactions/authorize.xml";
const AUTHREP_ENDPOINT: &str = "/transactions/authrep.xml";
const REPORT_ENDPOINT: &str = "/transactions.xml";
const OAUTH_AUTHORIZE_ENDPOINT: &str = "/transactions/oauth_authorize.xml";
const OAUTH_AUTHREP_ENDPOINT: &str = "/transactions/oauth_authrep.xml";

impl<'service, 'app, 'user> Info<'service, 'app, 'user> {
    pub fn new(kind: Type, service: &'service Service, application: &'app Application, user: Option<&'user User>) -> Self {
        Self { kind, service, application, user }
    }

    fn endpoint(&self) -> &str {
        use self::Type::*;
        use self::Application::*;

        match (&self.kind, self.application) {
            (&Authorize, &OAuthToken(_)) => OAUTH_AUTHORIZE_ENDPOINT,
            (&Authorize, _) => AUTHORIZE_ENDPOINT,
            (&AuthRep, &OAuthToken(_)) => OAUTH_AUTHREP_ENDPOINT,
            (&AuthRep, _) => AUTHREP_ENDPOINT,
            (&Report, _) => REPORT_ENDPOINT,
        }
    }

    fn path(&self) -> String {
        use request::ToParams;

        let mut params: Vec<(&str, &str)> = Vec::new();
        params.extend(self.service.to_params());
        params.extend(self.application.to_params());

        if let Some(user_params) = self.user.as_ref() {
            params.extend(user_params.to_params());
        }

        let params = params.into_iter()
            .map(|(param, value)| param.to_owned() + "=" + value)
            .collect::<Vec<String>>()
            .join("&");

        self.endpoint().to_owned() + "?" + params.as_str()
    }
}

impl<'service, 'app, 'user> ToRequest for Info<'service, 'app, 'user> {
    fn to_request(&self) -> Request {
        Request::new(self.kind.method(), self.path(), None)
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::BufReader;
use std::io::prelude::*;
use std::env;
use std::process::exit;

fn usage() {
  println!("Usage: cat input.csv | csv_to_json");
}

fn main() {
  let args : Vec<String> = env::args().collect();

  if args.len() > 1 {
    usage();
    exit(1);
  }

  let mut reader = BufReader::new(io::stdin());
  let split_char = ',';

  let mut line = String::new();
  let _unused_size = reader.read_line(&mut line);

  let headers : Vec<&str> = line.trim().split(split_char).collect();

  let mut first = true;

  for line in reader.lines() {
    if first {
      print!("[");
      first = false;
    } else {
      println!(",");
    }

    // Opening brace (escaped).
    print!("{{");

    let line_str = line.unwrap();

    let mut first = true;
    let components : Vec<&str> = line_str.split(split_char).collect();
    for (header, component) in headers.iter().zip(components.iter()) {
      if !first {
        print!(",");
      } else {
        first = false;
      }
      print!("\"{}\": \"{}\"", header, component);
    }

    // Closing brace (escaped).
    print!("}}");
  }
  println!("]");
}
<commit_msg>Take a file on the command line directly.<commit_after>use std::io;
use std::io::BufReader;
use std::io::prelude::*;
use std::env;
use std::process::exit;
use std::fs::File;

fn usage() {
  println!("Usage: cat input.csv | csv_to_json");
}

fn main() {
  let stdin = io::stdin();
  let mut filereader;
  let mut reader = &mut stdin.lock() as &mut BufRead;
  let args : Vec<String> = env::args().collect();

  if args.len() > 2 {
    usage();
    exit(1);
  } else if args.len() == 2 {
    let file = File::open(args[1].clone());
    if !file.is_ok() {
      usage();
      exit(1);
    }
    filereader = BufReader::new(file.unwrap());
    reader = &mut filereader;
  }

  let split_char = ',';

  let mut line = String::new();
  let _unused_size = reader.read_line(&mut line);

  let headers : Vec<&str> = line.trim().split(split_char).collect();

  let mut first = true;

  for line in reader.lines() {
    if first {
      print!("[");
      first = false;
    } else {
      println!(",");
    }

    // Opening brace (escaped).
    print!("{{");

    let line_str = line.unwrap();

    let mut first = true;
    let components : Vec<&str> = line_str.split(split_char).collect();
    for (header, component) in headers.iter().zip(components.iter()) {
      if !first {
        print!(",");
      } else {
        first = false;
      }
      print!("\"{}\": \"{}\"", header, component);
    }

    // Closing brace (escaped).
    print!("}}");
  }
  println!("]");
}
<|endoftext|>


<commit_before>//! It's a tree!

use std::ptr;
use std::ops::{Deref, DerefMut};
use std::fmt::Debug;

#[derive(Debug, PartialEq, Eq)]
struct Node<T> {
    // We need a mut pointer so we can modify the parent
    parent: *mut Node<T>,
    val: T,
    children: Vec<Node<T>>
}

impl<T> Node<T> {
    /// Create a new node with the existing value.
    /// For root-style constructors.
    pub fn new(val: T) -> Node<T> {
        Node {
            parent: ptr::null_mut(),
            val: val,
            children: Vec::new()
        }
    }

    /// Add a new child node to this node, using a value
    pub fn new_child(&mut self, val: T) {
        let self_mut = self as *mut Node<T>;
        self.children.push(Node {
            parent: self_mut,
            val: val,
            children: Vec::new()
        })
    }

    /// Whether this node has a (currently-reachable) parent
    pub fn has_parent(&self) -> bool {
        self.parent.is_null()
    }

    /// Gets the parent of this node (if it exists)
    pub fn get_parent<'a>(&'a self) -> Option<&'a mut Node<T>> {
        if self.parent.is_null() {
            return None;
        }
        unsafe {
            return Some(&mut *self.parent);
        }
    }

    /// Borrow the children of this node.
    pub fn get_children<'a>(&'a self) -> &Vec<Node<T>> {
        &self.children
    }

    /// Mutably borrow the children of this mutable node
    pub fn get_mut_children(&mut self) -> &mut Vec<Node<T>> {
        &mut self.children
    }

    /// Remove a child at the given index
    pub fn remove_child_at(&mut self, index: usize) -> Node<T> {
        self.children.remove(index)
    }

    /// Whether this node is a parent of another node
    pub fn is_parent_of(&self, other: Node<T>) -> bool {
        self.parent == other.parent
    }
}

impl <T: PartialEq> Node<T> {
    pub fn remove_from_parent(&self) {
        if let Some(ref mut parent) = self.get_parent() {
            let children = parent.get_mut_children();
            if let Some(index) = children.iter().position(|c| c == self) {
                children.remove(index);
            }
        }
    }

    pub fn remove_child(&mut self, other: &Node<T>) -> Option<Node<T>> {
        if let Some(index) = self.children.iter().position(|c| c == other) {
            Some(self.children.remove(index))
        }
        else {
            None
        }
    }

    pub fn move_to(mut self, new_parent: &mut Node<T>) {
        self.remove_from_parent();
        self.parent = new_parent as *mut Node<T>;
        new_parent.children.push(self);
    }
}

impl<T> Drop for Node<T> {
    fn drop(&mut self) {
        trace!("Dropping a node!!!!");
        self.parent = ptr::null_mut();
    }
}

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

    fn test_add_child() {

    }
}
<commit_msg>Node drop nullifies child's parent pointer<commit_after>//! It's a tree!

use std::ptr;
use std::ops::{Deref, DerefMut};
use std::fmt::Debug;

#[derive(Debug, PartialEq, Eq)]
struct Node<T> {
    // We need a mut pointer so we can modify the parent
    parent: *mut Node<T>,
    val: T,
    children: Vec<Node<T>>
}

impl<T> Node<T> {
    /// Create a new node with the existing value.
    /// For root-style constructors.
    pub fn new(val: T) -> Node<T> {
        Node {
            parent: ptr::null_mut(),
            val: val,
            children: Vec::new()
        }
    }

    /// Add a new child node to this node, using a value
    pub fn new_child(&mut self, val: T) -> &mut Node<T> {
        let self_mut = self as *mut Node<T>;
        self.children.push(Node {
            parent: self_mut,
            val: val,
            children: Vec::new()
        });
        let last_ix = self.children.len() -1;
        &mut self.children[last_ix]
    }

    /// Whether this node has a (currently-reachable) parent
    pub fn has_parent(&self) -> bool {
        self.parent.is_null()
    }

    /// Gets the parent of this node (if it exists)
    pub fn get_parent<'a>(&'a self) -> Option<&'a mut Node<T>> {
        if self.parent.is_null() {
            return None;
        }
        unsafe {
            return Some(&mut *self.parent);
        }
    }

    /// Borrow the children of this node.
    pub fn get_children<'a>(&'a self) -> &Vec<Node<T>> {
        &self.children
    }

    /// Mutably borrow the children of this mutable node
    pub fn get_mut_children(&mut self) -> &mut Vec<Node<T>> {
        &mut self.children
    }

    /// Remove a child at the given index
    pub fn remove_child_at(&mut self, index: usize) -> Node<T> {
        self.children.remove(index)
    }

    /// Whether this node is a parent of another node
    pub fn is_parent_of(&self, other: Node<T>) -> bool {
        self.parent == other.parent
    }
}

impl <T: PartialEq> Node<T> {

    /// Remove a node from its parent.
    /// This method will mutate the parent if it exists.
    pub fn remove_from_parent(&self) {
        if let Some(ref mut parent) = self.get_parent() {
            let children = parent.get_mut_children();
            if let Some(index) = children.iter().position(|c| c == self) {
                children.remove(index);
            }
        }
    }

    /// Removes a child from self
    pub fn remove_child(&mut self, other: &Node<T>) -> Option<Node<T>> {
        if let Some(index) = self.children.iter().position(|c| c == other) {
            Some(self.children.remove(index))
        }
        else {
            None
        }
    }

    pub fn move_to(mut self, new_parent: &mut Node<T>) {
        self.remove_from_parent();
        self.parent = new_parent as *mut Node<T>;
        new_parent.children.push(self);
    }
}

impl<T> Drop for Node<T> {
    fn drop(&mut self) {
        println!("Dropping a node.");
        self.parent = ptr::null_mut();
        for child in &mut self.children {
            println!("> Unlinking a child");
            child.parent = ptr::null_mut();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Node;
    use log::LogLevelFilter;
    use env_logger::LogBuilder;

    /// Nodes can have children added to them
    #[test]
    fn test_add_child() {
        let mut root = Node::new(0);
        root.new_child(1);
        root.new_child(2); // This is okay
        {
            let mut third_child = root.new_child(3);
            //root.new_child(4); // Have to wait for 3rd child to drop
        }
        root.new_child(4); // Now this works yay standard borrowing
        assert_eq!(root.children.len(), 4);
        println!("Done with the test now.");
    }

    /// These operations will for example operate on the parent
    /// under an if let. `remove_from_parent` will not panic if the node
    /// already is parentless, for example.
    #[test]
    fn optional_operations() {

    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use to_str::{ToStr,ToStrConsume};
use str;
use cast;

#[cfg(test)]
pub struct Ascii { priv chr: u8 }

/// Datatype to hold one ascii character. It is 8 bit long.
#[cfg(notest)]
#[deriving(Clone, Eq, Ord)]
pub struct Ascii { priv chr: u8 }

pub impl Ascii {
    /// Converts a ascii character into a `u8`.
    #[inline(always)]
    fn to_byte(self) -> u8 {
        self.chr
    }

    /// Converts a ascii character into a `char`.
    #[inline(always)]
    fn to_char(self) -> char {
        self.chr as char
    }

    /// Convert to lowercase.
    #[inline(always)]
    fn to_lower(self) -> Ascii {
        if self.chr >= 65 && self.chr <= 90 {
            Ascii{chr: self.chr | 0x20 }
        } else {
            self
        }
    }

    /// Convert to uppercase.
    #[inline(always)]
    fn to_upper(self) -> Ascii {
        if self.chr >= 97 && self.chr <= 122 {
            Ascii{chr: self.chr & !0x20 }
        } else {
            self
        }
    }

    // Compares two ascii characters of equality, ignoring case.
    #[inline(always)]
    fn eq_ignore_case(self, other: Ascii) -> bool {
        self.to_lower().chr == other.to_lower().chr
    }
}

impl ToStr for Ascii {
    #[inline(always)]
    fn to_str(&self) -> ~str { str::from_bytes(['\'' as u8, self.chr, '\'' as u8]) }
}

/// Trait for converting into an ascii type.
pub trait AsciiCast<T> {
    /// Convert to an ascii type
    fn to_ascii(&self) -> T;

    /// Check if convertible to ascii
    fn is_ascii(&self) -> bool;
}

impl<'self> AsciiCast<&'self[Ascii]> for &'self [u8] {
    #[inline(always)]
    fn to_ascii(&self) -> &'self[Ascii] {
        assert!(self.is_ascii());

        unsafe{ cast::transmute(*self) }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        for self.each |b| {
            if !b.is_ascii() { return false; }
        }
        true
    }
}

impl<'self> AsciiCast<&'self[Ascii]> for &'self str {
    #[inline(always)]
    fn to_ascii(&self) -> &'self[Ascii] {
        assert!(self.is_ascii());

        let (p,len): (*u8, uint) = unsafe{ cast::transmute(*self) };
        unsafe{ cast::transmute((p, len - 1))}
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        for self.each |b| {
            if !b.is_ascii() { return false; }
        }
        true
    }
}

impl AsciiCast<Ascii> for u8 {
    #[inline(always)]
    fn to_ascii(&self) -> Ascii {
        assert!(self.is_ascii());
        Ascii{ chr: *self }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        *self & 128 == 0u8
    }
}

impl AsciiCast<Ascii> for char {

    #[inline(always)]
    fn to_ascii(&self) -> Ascii {
        assert!(self.is_ascii());
        Ascii{ chr: *self as u8 }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        *self - ('\x7F' & *self) == '\x00'
    }
}

/// Trait for copyless casting to an ascii vector.
pub trait OwnedAsciiCast {
    /// Take ownership and cast to an ascii vector without trailing zero element.
    fn to_ascii_consume(self) -> ~[Ascii];
}

impl OwnedAsciiCast for ~[u8] {
    #[inline(always)]
    fn to_ascii_consume(self) -> ~[Ascii] {
        assert!(self.is_ascii());

        unsafe {cast::transmute(self)}
    }
}

impl OwnedAsciiCast for ~str {
    #[inline(always)]
    fn to_ascii_consume(self) -> ~[Ascii] {
        let mut s = self;
        unsafe {
            str::raw::pop_byte(&mut s);
            cast::transmute(s)
        }
    }
}

/// Trait for converting an ascii type to a string. Needed to convert `&[Ascii]` to `~str`
pub trait ToStrAscii {
    /// Convert to a string.
    fn to_str_ascii(&self) -> ~str;
}

impl<'self> ToStrAscii for &'self [Ascii] {
    #[inline(always)]
    fn to_str_ascii(&self) -> ~str {
        let mut cpy = self.to_owned();
        cpy.push(0u8.to_ascii());
        unsafe {cast::transmute(cpy)}
    }
}

impl ToStrConsume for ~[Ascii] {
    #[inline(always)]
    fn to_str_consume(self) -> ~str {
        let mut cpy = self;
        cpy.push(0u8.to_ascii());
        unsafe {cast::transmute(cpy)}
    }
}

// NOTE: Remove stage0 marker after snapshot
#[cfg(and(test, not(stage0)))]
mod tests {
    use super::*;
    use to_str::{ToStr,ToStrConsume};
    use str;
    use cast;

    macro_rules! v2ascii (
        ( [$($e:expr),*]) => ( [$(Ascii{chr:$e}),*]);
        (~[$($e:expr),*]) => (~[$(Ascii{chr:$e}),*]);
    )

    #[test]
    fn test_ascii() {
        assert_eq!(65u8.to_ascii().to_byte(), 65u8);
        assert_eq!(65u8.to_ascii().to_char(), 'A');
        assert_eq!('A'.to_ascii().to_char(), 'A');
        assert_eq!('A'.to_ascii().to_byte(), 65u8);

        assert_eq!('A'.to_ascii().to_lower().to_char, 'a');
        assert_eq!('Z'.to_ascii().to_lower().to_char, 'z');
        assert_eq!('a'.to_ascii().to_upper().to_char, 'A');
        assert_eq!('z'.to_ascii().to_upper().to_char, 'Z');

        assert_eq!('@'.to_ascii().to_lower().to_char, '@');
        assert_eq!('['.to_ascii().to_lower().to_char, '[');
        assert_eq!('`'.to_ascii().to_upper().to_char, '`');
        assert_eq!('{'.to_ascii().to_upper().to_char, '{');
    }

    #[test]
    fn test_ascii_vec() {
        assert_eq!((&[40u8, 32u8, 59u8]).to_ascii(), v2ascii!([40, 32, 59]));
        assert_eq!("( ;".to_ascii(),                 v2ascii!([40, 32, 59]));
        // FIXME: #5475 borrowchk error, owned vectors do not live long enough
        // if chained-from directly
        let v = ~[40u8, 32u8, 59u8]; assert_eq!(v.to_ascii(), v2ascii!([40, 32, 59]));
        let v = ~"( ;";              assert_eq!(v.to_ascii(), v2ascii!([40, 32, 59]));
    }

    #[test]
    fn test_owned_ascii_vec() {
        // FIXME: #4318 Compiler crashes on moving self
        //assert_eq!(~"( ;".to_ascii_consume(), v2ascii!(~[40, 32, 59]));
        //assert_eq!(~[40u8, 32u8, 59u8].to_ascii_consume(), v2ascii!(~[40, 32, 59]));
        //assert_eq!(~"( ;".to_ascii_consume_with_null(), v2ascii!(~[40, 32, 59, 0]));
        //assert_eq!(~[40u8, 32u8, 59u8].to_ascii_consume_with_null(),
        //           v2ascii!(~[40, 32, 59, 0]));
    }

    #[test]
    fn test_ascii_to_str() { assert_eq!(v2ascii!([40, 32, 59]).to_str_ascii(), ~"( ;"); }

    #[test]
    fn test_ascii_to_str_consume() {
        // FIXME: #4318 Compiler crashes on moving self
        //assert_eq!(v2ascii!(~[40, 32, 59]).to_str_consume(), ~"( ;");
    }

    #[test] #[should_fail]
    fn test_ascii_vec_fail_u8_slice()  { (&[127u8, 128u8, 255u8]).to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_vec_fail_str_slice() { "zoä华".to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_fail_u8_slice() { 255u8.to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_fail_char_slice() { 'λ'.to_ascii(); }
}

<commit_msg>Corrected history<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use to_str::{ToStr,ToStrConsume};
use str;
use cast;

#[cfg(test)]
pub struct Ascii { priv chr: u8 }

/// Datatype to hold one ascii character. It is 8 bit long.
#[cfg(notest)]
#[deriving(Clone, Eq, Ord)]
pub struct Ascii { priv chr: u8 }

pub impl Ascii {
    /// Converts a ascii character into a `u8`.
    #[inline(always)]
    fn to_byte(self) -> u8 {
        self.chr
    }

    /// Converts a ascii character into a `char`.
    #[inline(always)]
    fn to_char(self) -> char {
        self.chr as char
    }

    /// Convert to lowercase.
    #[inline(always)]
    fn to_lower(self) -> Ascii {
        if self.chr >= 65 && self.chr <= 90 {
            Ascii{chr: self.chr | 0x20 }
        } else {
            self
        }
    }

    /// Convert to uppercase.
    #[inline(always)]
    fn to_upper(self) -> Ascii {
        if self.chr >= 97 && self.chr <= 122 {
            Ascii{chr: self.chr & !0x20 }
        } else {
            self
        }
    }

    // Compares two ascii characters of equality, ignoring case.
    #[inline(always)]
    fn eq_ignore_case(self, other: Ascii) -> bool {
        self.to_lower().chr == other.to_lower().chr
    }
}

impl ToStr for Ascii {
    #[inline(always)]
    fn to_str(&self) -> ~str { str::from_bytes(['\'' as u8, self.chr, '\'' as u8]) }
}

/// Trait for converting into an ascii type.
pub trait AsciiCast<T> {
    /// Convert to an ascii type
    fn to_ascii(&self) -> T;

    /// Check if convertible to ascii
    fn is_ascii(&self) -> bool;
}

impl<'self> AsciiCast<&'self[Ascii]> for &'self [u8] {
    #[inline(always)]
    fn to_ascii(&self) -> &'self[Ascii] {
        assert!(self.is_ascii());

        unsafe{ cast::transmute(*self) }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        for self.each |b| {
            if !b.is_ascii() { return false; }
        }
        true
    }
}

impl<'self> AsciiCast<&'self[Ascii]> for &'self str {
    #[inline(always)]
    fn to_ascii(&self) -> &'self[Ascii] {
        assert!(self.is_ascii());

        let (p,len): (*u8, uint) = unsafe{ cast::transmute(*self) };
        unsafe{ cast::transmute((p, len - 1))}
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        for self.each |b| {
            if !b.is_ascii() { return false; }
        }
        true
    }
}

impl AsciiCast<Ascii> for u8 {
    #[inline(always)]
    fn to_ascii(&self) -> Ascii {
        assert!(self.is_ascii());
        Ascii{ chr: *self }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        *self & 128 == 0u8
    }
}

impl AsciiCast<Ascii> for char {

    #[inline(always)]
    fn to_ascii(&self) -> Ascii {
        assert!(self.is_ascii());
        Ascii{ chr: *self as u8 }
    }

    #[inline(always)]
    fn is_ascii(&self) -> bool {
        *self - ('\x7F' & *self) == '\x00'
    }
}

/// Trait for copyless casting to an ascii vector.
pub trait OwnedAsciiCast {
    /// Take ownership and cast to an ascii vector without trailing zero element.
    fn to_ascii_consume(self) -> ~[Ascii];
}

impl OwnedAsciiCast for ~[u8] {
    #[inline(always)]
    fn to_ascii_consume(self) -> ~[Ascii] {
        assert!(self.is_ascii());

        unsafe {cast::transmute(self)}
    }
}

impl OwnedAsciiCast for ~str {
    #[inline(always)]
    fn to_ascii_consume(self) -> ~[Ascii] {
        let mut s = self;
        unsafe {
            str::raw::pop_byte(&mut s);
            cast::transmute(s)
        }
    }
}

/// Trait for converting an ascii type to a string. Needed to convert `&[Ascii]` to `~str`
pub trait ToStrAscii {
    /// Convert to a string.
    fn to_str_ascii(&self) -> ~str;
}

impl<'self> ToStrAscii for &'self [Ascii] {
    #[inline(always)]
    fn to_str_ascii(&self) -> ~str {
        let mut cpy = self.to_owned();
        cpy.push(0u8.to_ascii());
        unsafe {cast::transmute(cpy)}
    }
}

impl ToStrConsume for ~[Ascii] {
    #[inline(always)]
    fn to_str_consume(self) -> ~str {
        let mut cpy = self;
        cpy.push(0u8.to_ascii());
        unsafe {cast::transmute(cpy)}
    }
}

// NOTE: Remove stage0 marker after snapshot
#[cfg(and(test, not(stage0)))]
mod tests {
    use super::*;
    use to_str::{ToStr,ToStrConsume};
    use str;
    use cast;

    macro_rules! v2ascii (
        ( [$($e:expr),*]) => ( [$(Ascii{chr:$e}),*]);
        (~[$($e:expr),*]) => (~[$(Ascii{chr:$e}),*]);
    )

    #[test]
    fn test_ascii() {
        assert_eq!(65u8.to_ascii().to_byte(), 65u8);
        assert_eq!(65u8.to_ascii().to_char(), 'A');
        assert_eq!('A'.to_ascii().to_char(), 'A');
        assert_eq!('A'.to_ascii().to_byte(), 65u8);

        assert_eq!('A'.to_ascii().to_lower().to_char, 'a');
        assert_eq!('Z'.to_ascii().to_lower().to_char, 'z');
        assert_eq!('a'.to_ascii().to_upper().to_char, 'A');
        assert_eq!('z'.to_ascii().to_upper().to_char, 'Z');

        assert_eq!('@'.to_ascii().to_lower().to_char, '@');
        assert_eq!('['.to_ascii().to_lower().to_char, '[');
        assert_eq!('`'.to_ascii().to_upper().to_char, '`');
        assert_eq!('{'.to_ascii().to_upper().to_char, '{');
    }

    #[test]
    fn test_ascii_vec() {
        assert_eq!((&[40u8, 32u8, 59u8]).to_ascii(), v2ascii!([40, 32, 59]));
        assert_eq!("( ;".to_ascii(),                 v2ascii!([40, 32, 59]));
        // FIXME: #5475 borrowchk error, owned vectors do not live long enough
        // if chained-from directly
        let v = ~[40u8, 32u8, 59u8]; assert_eq!(v.to_ascii(), v2ascii!([40, 32, 59]));
        let v = ~"( ;";              assert_eq!(v.to_ascii(), v2ascii!([40, 32, 59]));
    }

    #[test]
    fn test_owned_ascii_vec() {
        // FIXME: #4318 Compiler crashes on moving self
        //assert_eq!(~"( ;".to_ascii_consume(), v2ascii!(~[40, 32, 59]));
        //assert_eq!(~[40u8, 32u8, 59u8].to_ascii_consume(), v2ascii!(~[40, 32, 59]));
        //assert_eq!(~"( ;".to_ascii_consume_with_null(), v2ascii!(~[40, 32, 59, 0]));
        //assert_eq!(~[40u8, 32u8, 59u8].to_ascii_consume_with_null(),
        //           v2ascii!(~[40, 32, 59, 0]));
    }

    #[test]
    fn test_ascii_to_str() { assert_eq!(v2ascii!([40, 32, 59]).to_str_ascii(), ~"( ;"); }

    #[test]
    fn test_ascii_to_str_consume() {
        // FIXME: #4318 Compiler crashes on moving self
        //assert_eq!(v2ascii!(~[40, 32, 59]).to_str_consume(), ~"( ;");
    }

    #[test] #[should_fail]
    fn test_ascii_vec_fail_u8_slice()  { (&[127u8, 128u8, 255u8]).to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_vec_fail_str_slice() { "zoä华".to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_fail_u8_slice() { 255u8.to_ascii(); }

    #[test] #[should_fail]
    fn test_ascii_fail_char_slice() { 'λ'.to_ascii(); }
}

<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker, Configuration, ConnectionStatus,
    Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.critical_error(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.handle_frames()?;
            if self.should_heartbeat() {
                self.channels.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()?;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>IoLoop: use new Heartbeat helper<commit_after>use crate::{
    buffer::Buffer, channels::Channels, connection_status::ConnectionState, frames::Frames,
    heartbeat::Heartbeat, internal_rpc::InternalRPC, thread::ThreadHandle, waker::Waker,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker as MioWaker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    connection_io_loop_handle: ThreadHandle,
    waker: Waker,
    socket: T,
    status: Status,
    poll: Poll,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        connection_io_loop_handle: ThreadHandle,
        waker: Waker,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        waker.set_waker(MioWaker::new(poll.registry(), WAKER)?);
        let heartbeat = Heartbeat::new(channels.clone());
        let frame_size = std::cmp::max(8192, configuration.frame_max() as usize);
        let mut inner = Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            connection_io_loop_handle,
            waker,
            socket,
            status: Status::Initial,
            poll,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            serialized_frames: VecDeque::default(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) {
        if self.status != Status::Setup && self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
            }
            self.status = Status::Setup;
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status == Status::Initial
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.critical_error(err)?;
                        }
                    }
                    Ok(())
                })?,
        );
        waker.wake()
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.heartbeat.poll_timeout()?)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup();
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection_status.closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.handle_frames()?;
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection_status.errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.socket.flush()?;

        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }

            self.socket.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer.read_from(&mut self.socket).map(|sz| {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                })?;
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            self.send_continue()?;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>use collections::slice::{mod, AsSlice, SlicePrelude};
use collections::vec::{mod, Vec};
use core::clone::Clone;
use core::fmt;
use core::mem;
use core::iter::{mod, AdditiveIterator, Iterator};
use core::option::{mod, Some, None, Option};
use core::result::{Ok, Err};

use iobuf::Iobuf;

/// A span over potentially many Iobufs. This is useful as a "string" type where
/// the contents of the string can come from multiple IObufs, and you want to
/// avoid copying the buffer contents unnecessarily.
///
/// As an optimization, pushing an Iobuf that points to data immediately after
/// the range represented by the last Iobuf pushed will result in just expanding
/// the held Iobuf's range. This prevents allocating lots of unnecessary
/// intermediate buffers, while still maintaining the illusion of "pushing lots
/// of buffers" while incrementally parsing.
///
/// A `BufSpan` is internally represented as either an `Iobuf` or a `Vec<Iobuf>`,
/// depending on how many different buffers were used.
pub enum BufSpan<Buf> {
  Empty,
  One (Buf),
  Many(Vec<Buf>),
}

impl<Buf: Iobuf> Clone for BufSpan<Buf> {
  #[inline]
  fn clone(&self) -> BufSpan<Buf> {
    match *self {
      Empty       => Empty,
      One(ref b)  => One ((*b).clone()),
      Many(ref v) => Many((*v).clone()),
    }
  }
}

impl<Buf: Iobuf> fmt::Show for BufSpan<Buf> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    let mut first_time = true;

    for b in self.iter() {
      if !first_time {
        try!(write!(f, "\n"));
      }

      try!(b.fmt(f));

      first_time = false;
    }

    Ok(())
  }
}

impl<Buf: Iobuf> BufSpan<Buf> {
  /// Creates a new, empty `Bufspan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s: BufSpan<ROIobuf<'static>> = BufSpan::new();
  /// assert!(s.is_empty());
  /// ```
  #[inline]
  pub fn new() -> BufSpan<Buf> {
    Empty
  }

  /// Creates a new `BufSpan` from an Iobuf.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s = BufSpan::from_buf(ROIobuf::from_slice(b"hello"));
  /// assert_eq!(s.iter().count(), 1);
  /// assert_eq!(s.count_bytes(), 5);
  /// ```
  #[inline]
  pub fn from_buf(b: Buf) -> BufSpan<Buf> {
    One(b)
  }

  /// Returns `true` iff the span is over an empty range.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str(""));
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str("hello, world!"));
  /// assert!(!s.is_empty());
  /// ```
  #[inline]
  pub fn is_empty(&self) -> bool {
    match *self {
      Empty => true,
      _     => false,
    }
  }

  /// The fast path during pushing -- either fills in the first buffer, or
  /// extends an existing one.
  ///
  /// Returns `None` if the fast path was taken and nothing more needs to be
  /// done. Returns `Some` if we need to do a slow push.
  fn try_to_extend(&mut self, b: Buf) -> Option<Buf> {
    let len = b.len();
    if len == 0 { return None; }

    match *self {
      Empty => {},
      One(ref mut b0) => {
        unsafe {
          if b0.is_extended_by(&b) {
            b0.unsafe_extend(len);
            return None;
          } else {
            return Some(b); // Upgrade the `One` into a `Many` in the slow path.
          }
        }
      }
      Many(ref mut v) => {
        // I wish this wouldn't unwind, and just abort... :(
        let last = v.last_mut().unwrap();
        unsafe {
          if last.is_extended_by(&b) {
            last.unsafe_extend(len);
            return None;
          } else {
            return Some(b);
          }
        }
      }
    }

    // Handle this case in the fast path, instead of leaving it for slow_push
    // to clean up.
    *self = One(b);
    None
  }

  /// Appends a buffer to a `BufSpan`. If the buffer is an extension of the
  /// previously pushed buffer, the range will be extended. Otherwise, the new
  /// non-extension buffer will be added to the end of a vector.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// s.push(ROIobuf::from_str("he"));
  /// s.push(ROIobuf::from_str("llo"));
  ///
  /// assert_eq!(s.count_bytes() as uint, "hello".len());
  /// assert_eq!(s.iter().count(), 2);
  ///
  /// let mut b0 = ROIobuf::from_str(" world");
  /// let mut b1 = b0.clone();
  ///
  /// b0.resize(2).unwrap();
  /// b1.advance(2).unwrap();
  ///
  /// s.push(b0);
  /// s.push(b1);
  ///
  /// // b0 and b1 are immediately after each other, and from the same buffer,
  /// // so get merged into one Iobuf.
  /// assert_eq!(s.count_bytes() as uint, "hello world".len());
  /// assert_eq!(s.iter().count(), 3);
  /// ```
  pub fn push(&mut self, b: Buf) {
    match self.try_to_extend(b) {
      None    => {},
      Some(b) => self.slow_push(b),
    }
  }

  /// The slow path during a push. This is only taken if a `BufSpan` must span
  /// multiple backing buffers.
  #[cold]
  fn slow_push(&mut self, b: Buf) {
    let this = mem::replace(self, Empty);
    *self =
      match this {
        Empty   => One(b),
        One(b0) => {
          let mut v = Vec::with_capacity(2);
          v.push(b0);
          v.push(b);
          Many(v)
        },
        Many(mut bs) => { bs.push(b); Many(bs) }
      };
  }

  /// Returns an iterator over references to the buffers inside the `BufSpan`.
  #[inline]
  pub fn iter<'a>(&'a self) -> SpanIter<'a, Buf> {
    match *self {
      Empty       => Opt(None.into_iter()),
      One (ref b) => Opt(Some(b).into_iter()),
      Many(ref v) => Lot(v.as_slice().iter()),
    }
  }

  /// Returns a moving iterator over the buffers inside the `BufSpan`.
  #[inline]
  pub fn into_iter(self) -> SpanMoveIter<Buf> {
    match self {
      Empty   => MoveOpt(None.into_iter()),
      One (b) => MoveOpt(Some(b).into_iter()),
      Many(v) => MoveLot(v.into_iter()),
    }
  }

  /// Returns an iterator over the bytes in the `BufSpan`.
  #[inline]
  pub fn iter_bytes<'a>(&'a self) -> ByteIter<'a, Buf> {
    self.iter()
        .flat_map(|buf| unsafe { buf.as_window_slice().iter() })
        .map(|&b| b)
  }

  /// Returns `true` iff the bytes in this `BufSpan` are the same as the bytes
  /// in the other `BufSpan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf, RWIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  /// let b = BufSpan::from_buf(RWIobuf::from_string("hello".into_string()));
  ///
  /// assert!(a.byte_equal(&b));
  ///
  /// let mut c = BufSpan::from_buf(ROIobuf::from_str("hel"));
  /// c.push(ROIobuf::from_str("lo"));
  ///
  /// assert!(a.byte_equal(&c)); assert!(c.byte_equal(&a));
  ///
  /// let d = BufSpan::from_buf(ROIobuf::from_str("helo"));
  /// assert!(!a.byte_equal(&d));
  /// ```
  #[inline]
  pub fn byte_equal<Buf2: Iobuf>(&self, other: &BufSpan<Buf2>) -> bool {
    self.count_bytes() == other.count_bytes()
    && self.iter_bytes().zip(other.iter_bytes()).all(|(a, b)| a == b)
  }

  /// A more efficient version of byte_equal, specialized to work exclusively on
  /// slices.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  ///
  /// assert!(a.byte_equal_slice(b"hello"));
  /// assert!(!a.byte_equal_slice(b"helo"));
  /// ```
  #[inline]
  pub fn byte_equal_slice(&self, other: &[u8]) -> bool {
    self.count_bytes() as uint == other.len()
    && self.iter_bytes().zip(other.iter()).all(|(a, &b)| a == b)
  }

  /// Counts the number of bytes this `BufSpan` is over. This is
  /// `O(self.iter().len())`.
  #[inline]
  pub fn count_bytes(&self) -> u32 {
    // `self.iter().map(|b| b.len()).sum()` would be shorter, but I like to
    // specialize for the much more common case of empty or singular `BufSpan`s.
    match *self {
      Empty       => 0,
      One (ref b) => b.len(),
      Many(ref v) => v.iter().map(|b| b.len()).sum(),
    }
  }
}

/// An iterator over the bytes in a `BufSpan`.
pub type ByteIter<'a, Buf> =
  iter::Map<'static, &'a u8, u8,
    iter::FlatMap<'static, &'a Buf,
      SpanIter<'a, Buf>,
      slice::Items<'a, u8>>>;

/// An iterator over references to buffers inside a `BufSpan`.
pub enum SpanIter<'a, Buf: 'a> {
  Opt(option::Item<&'a Buf>),
  Lot(slice::Items<'a, Buf>),
}

impl<'a, Buf: Iobuf> Iterator<&'a Buf> for SpanIter<'a, Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<&'a Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      Opt(ref mut iter) => iter.next(),
      Lot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      Opt(ref iter) => iter.size_hint(),
      Lot(ref iter) => iter.size_hint(),
    }
  }
}

/// A moving iterator over buffers inside a `BufSpan`.
pub enum SpanMoveIter<Buf> {
  MoveOpt(option::Item<Buf>),
  MoveLot(vec::MoveItems<Buf>),
}

impl<Buf: Iobuf> Iterator<Buf> for SpanMoveIter<Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      MoveOpt(ref mut iter) => iter.next(),
      MoveLot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      MoveOpt(ref iter) => iter.size_hint(),
      MoveLot(ref iter) => iter.size_hint(),
    }
  }
}
<commit_msg>extendzor!<commit_after>use collections::slice::{mod, AsSlice, SlicePrelude};
use collections::vec::{mod, Vec};
use core::clone::Clone;
use core::fmt;
use core::mem;
use core::iter::{mod, Extend, AdditiveIterator, Iterator};
use core::option::{mod, Some, None, Option};
use core::result::{Ok, Err};

use iobuf::Iobuf;

/// A span over potentially many Iobufs. This is useful as a "string" type where
/// the contents of the string can come from multiple IObufs, and you want to
/// avoid copying the buffer contents unnecessarily.
///
/// As an optimization, pushing an Iobuf that points to data immediately after
/// the range represented by the last Iobuf pushed will result in just expanding
/// the held Iobuf's range. This prevents allocating lots of unnecessary
/// intermediate buffers, while still maintaining the illusion of "pushing lots
/// of buffers" while incrementally parsing.
///
/// A `BufSpan` is internally represented as either an `Iobuf` or a `Vec<Iobuf>`,
/// depending on how many different buffers were used.
pub enum BufSpan<Buf> {
  Empty,
  One (Buf),
  Many(Vec<Buf>),
}

impl<Buf: Iobuf> Clone for BufSpan<Buf> {
  #[inline]
  fn clone(&self) -> BufSpan<Buf> {
    match *self {
      Empty       => Empty,
      One(ref b)  => One ((*b).clone()),
      Many(ref v) => Many((*v).clone()),
    }
  }
}

impl<Buf: Iobuf> fmt::Show for BufSpan<Buf> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    let mut first_time = true;

    for b in self.iter() {
      if !first_time {
        try!(write!(f, "\n"));
      }

      try!(b.fmt(f));

      first_time = false;
    }

    Ok(())
  }
}

impl<Buf: Iobuf> Extend<Buf> for BufSpan<Buf> {
  fn extend<T: Iterator<Buf>>(&mut self, mut iterator: T) {
    for x in iterator {
      self.push(x);
    }
  }
}

impl<Buf: Iobuf> BufSpan<Buf> {
  /// Creates a new, empty `Bufspan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s: BufSpan<ROIobuf<'static>> = BufSpan::new();
  /// assert!(s.is_empty());
  /// ```
  #[inline]
  pub fn new() -> BufSpan<Buf> {
    Empty
  }

  /// Creates a new `BufSpan` from an Iobuf.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let s = BufSpan::from_buf(ROIobuf::from_slice(b"hello"));
  /// assert_eq!(s.iter().count(), 1);
  /// assert_eq!(s.count_bytes(), 5);
  /// ```
  #[inline]
  pub fn from_buf(b: Buf) -> BufSpan<Buf> {
    One(b)
  }

  /// Returns `true` iff the span is over an empty range.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str(""));
  /// assert!(s.is_empty());
  ///
  /// s.push(ROIobuf::from_str("hello, world!"));
  /// assert!(!s.is_empty());
  /// ```
  #[inline]
  pub fn is_empty(&self) -> bool {
    match *self {
      Empty => true,
      _     => false,
    }
  }

  /// The fast path during pushing -- either fills in the first buffer, or
  /// extends an existing one.
  ///
  /// Returns `None` if the fast path was taken and nothing more needs to be
  /// done. Returns `Some` if we need to do a slow push.
  fn try_to_extend(&mut self, b: Buf) -> Option<Buf> {
    let len = b.len();
    if len == 0 { return None; }

    match *self {
      Empty => {},
      One(ref mut b0) => {
        unsafe {
          if b0.is_extended_by(&b) {
            b0.unsafe_extend(len);
            return None;
          } else {
            return Some(b); // Upgrade the `One` into a `Many` in the slow path.
          }
        }
      }
      Many(ref mut v) => {
        // I wish this wouldn't unwind, and just abort... :(
        let last = v.last_mut().unwrap();
        unsafe {
          if last.is_extended_by(&b) {
            last.unsafe_extend(len);
            return None;
          } else {
            return Some(b);
          }
        }
      }
    }

    // Handle this case in the fast path, instead of leaving it for slow_push
    // to clean up.
    *self = One(b);
    None
  }

  /// Appends a buffer to a `BufSpan`. If the buffer is an extension of the
  /// previously pushed buffer, the range will be extended. Otherwise, the new
  /// non-extension buffer will be added to the end of a vector.
  ///
  /// ```
  /// use iobuf::{BufSpan, Iobuf, ROIobuf};
  ///
  /// let mut s = BufSpan::new();
  ///
  /// s.push(ROIobuf::from_str("he"));
  /// s.push(ROIobuf::from_str("llo"));
  ///
  /// assert_eq!(s.count_bytes() as uint, "hello".len());
  /// assert_eq!(s.iter().count(), 2);
  ///
  /// let mut b0 = ROIobuf::from_str(" world");
  /// let mut b1 = b0.clone();
  ///
  /// b0.resize(2).unwrap();
  /// b1.advance(2).unwrap();
  ///
  /// s.push(b0);
  /// s.push(b1);
  ///
  /// // b0 and b1 are immediately after each other, and from the same buffer,
  /// // so get merged into one Iobuf.
  /// assert_eq!(s.count_bytes() as uint, "hello world".len());
  /// assert_eq!(s.iter().count(), 3);
  /// ```
  pub fn push(&mut self, b: Buf) {
    match self.try_to_extend(b) {
      None    => {},
      Some(b) => self.slow_push(b),
    }
  }

  /// The slow path during a push. This is only taken if a `BufSpan` must span
  /// multiple backing buffers.
  #[cold]
  fn slow_push(&mut self, b: Buf) {
    let this = mem::replace(self, Empty);
    *self =
      match this {
        Empty   => One(b),
        One(b0) => {
          let mut v = Vec::with_capacity(2);
          v.push(b0);
          v.push(b);
          Many(v)
        },
        Many(mut bs) => { bs.push(b); Many(bs) }
      };
  }

  /// Returns an iterator over references to the buffers inside the `BufSpan`.
  #[inline]
  pub fn iter<'a>(&'a self) -> SpanIter<'a, Buf> {
    match *self {
      Empty       => Opt(None.into_iter()),
      One (ref b) => Opt(Some(b).into_iter()),
      Many(ref v) => Lot(v.as_slice().iter()),
    }
  }

  /// Returns a moving iterator over the buffers inside the `BufSpan`.
  #[inline]
  pub fn into_iter(self) -> SpanMoveIter<Buf> {
    match self {
      Empty   => MoveOpt(None.into_iter()),
      One (b) => MoveOpt(Some(b).into_iter()),
      Many(v) => MoveLot(v.into_iter()),
    }
  }

  /// Returns an iterator over the bytes in the `BufSpan`.
  #[inline]
  pub fn iter_bytes<'a>(&'a self) -> ByteIter<'a, Buf> {
    self.iter()
        .flat_map(|buf| unsafe { buf.as_window_slice().iter() })
        .map(|&b| b)
  }

  /// Returns `true` iff the bytes in this `BufSpan` are the same as the bytes
  /// in the other `BufSpan`.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf, RWIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  /// let b = BufSpan::from_buf(RWIobuf::from_string("hello".into_string()));
  ///
  /// assert!(a.byte_equal(&b));
  ///
  /// let mut c = BufSpan::from_buf(ROIobuf::from_str("hel"));
  /// c.push(ROIobuf::from_str("lo"));
  ///
  /// assert!(a.byte_equal(&c)); assert!(c.byte_equal(&a));
  ///
  /// let d = BufSpan::from_buf(ROIobuf::from_str("helo"));
  /// assert!(!a.byte_equal(&d));
  /// ```
  #[inline]
  pub fn byte_equal<Buf2: Iobuf>(&self, other: &BufSpan<Buf2>) -> bool {
    self.count_bytes() == other.count_bytes()
    && self.iter_bytes().zip(other.iter_bytes()).all(|(a, b)| a == b)
  }

  /// A more efficient version of byte_equal, specialized to work exclusively on
  /// slices.
  ///
  /// ```
  /// use iobuf::{BufSpan, ROIobuf};
  ///
  /// let a = BufSpan::from_buf(ROIobuf::from_str("hello"));
  ///
  /// assert!(a.byte_equal_slice(b"hello"));
  /// assert!(!a.byte_equal_slice(b"helo"));
  /// ```
  #[inline]
  pub fn byte_equal_slice(&self, other: &[u8]) -> bool {
    self.count_bytes() as uint == other.len()
    && self.iter_bytes().zip(other.iter()).all(|(a, &b)| a == b)
  }

  /// Counts the number of bytes this `BufSpan` is over. This is
  /// `O(self.iter().len())`.
  #[inline]
  pub fn count_bytes(&self) -> u32 {
    // `self.iter().map(|b| b.len()).sum()` would be shorter, but I like to
    // specialize for the much more common case of empty or singular `BufSpan`s.
    match *self {
      Empty       => 0,
      One (ref b) => b.len(),
      Many(ref v) => v.iter().map(|b| b.len()).sum(),
    }
  }

  /// Extends this span to include the range denoted by another span.
  #[inline]
  pub fn append(&mut self, other: BufSpan<Buf>) {
    self.extend(other.into_iter())
  }
}

/// An iterator over the bytes in a `BufSpan`.
pub type ByteIter<'a, Buf> =
  iter::Map<'static, &'a u8, u8,
    iter::FlatMap<'static, &'a Buf,
      SpanIter<'a, Buf>,
      slice::Items<'a, u8>>>;

/// An iterator over references to buffers inside a `BufSpan`.
pub enum SpanIter<'a, Buf: 'a> {
  Opt(option::Item<&'a Buf>),
  Lot(slice::Items<'a, Buf>),
}

impl<'a, Buf: Iobuf> Iterator<&'a Buf> for SpanIter<'a, Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<&'a Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      Opt(ref mut iter) => iter.next(),
      Lot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      Opt(ref iter) => iter.size_hint(),
      Lot(ref iter) => iter.size_hint(),
    }
  }
}

/// A moving iterator over buffers inside a `BufSpan`.
pub enum SpanMoveIter<Buf> {
  MoveOpt(option::Item<Buf>),
  MoveLot(vec::MoveItems<Buf>),
}

impl<Buf: Iobuf> Iterator<Buf> for SpanMoveIter<Buf> {
  #[inline(always)]
  fn next(&mut self) -> Option<Buf> {
    // I'm couting on this match getting lifted out of the loop with
    // loop-invariant code motion.
    match *self {
      MoveOpt(ref mut iter) => iter.next(),
      MoveLot(ref mut iter) => iter.next(),
    }
  }

  #[inline(always)]
  fn size_hint(&self) -> (uint, Option<uint>) {
    match *self {
      MoveOpt(ref iter) => iter.size_hint(),
      MoveLot(ref iter) => iter.size_hint(),
    }
  }
}
<|endoftext|>