<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 viewing key.
pub struct ViewingKey {
    pub sk_enc: [u8; 32],
    pub pk_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.pk_enc);
            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 mut a_pk = [0u8; 32];
        pseudorandom_function_a_pk(&mut a_pk, &self.a_sk);

        PaymentAddress {
            a_pk: a_pk,
            pk_enc: self.viewing_key().pk_enc,
        }
    }

    /// Computes a viewing key for this spending key.
    pub fn viewing_key(&self) -> ViewingKey {
        let mut sk_enc = [0u8; 32];
        pseudorandom_function_sk_enc(&mut sk_enc, &self.a_sk);
        clamp_curve25519(&mut sk_enc);
        let pk = &Scalar(sk_enc) * &ED25519_BASEPOINT;
        let pk_enc = pk.compress_montgomery().unwrap().to_bytes();

        ViewingKey {
            sk_enc: sk_enc,
            pk_enc: pk_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",
		"ZiVKW6e4pHW9jRSa6v1BSYQY9FYeDpF2dS8h61DQJzGJe8WdyErzb57W4rALT235UKmEcu2wNnrWNc2V9mjNbqPHTDYpEXdqa",
	),
	(
		"SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
		"zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
		"ZiVKkUCm1oqVfyN7e1TY1Xq6CPZgdG18CAirp1nAQWxWcDufxUUkmw11S1fRG23AVWv36rud1Fp3zDEUXEfLfbGyFU6XGTkKM",
	),
	(
		"SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
		"zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
		"ZiVKhFPyGkDJjTXUuSi9yWBSnkJF8jLcnE3wSNLnexSgwGeoR8N7ZHRfoyyBMRiM1nbNVDPmcTQPYPBqmS1MDDh1B9fSTauDD",
	),
	(
		"SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
		"zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
		"ZiVKnYSZMncaac1SmjB9nfbY3xPdwwiJCVP9inpy59pxVoAWAjSCHTn4rcbQJznWajZALLiFHkKE9XCZQAwSuSqSVbPKtjr5v",
	),
	(
		"SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
		"zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
		"ZiVKk7ZNUE2sjkNkmaJ5PwBEWm287XFZWeXvoAYcWacCtmuEpAGRDqSfvQstB6erm5bt8uptG3bSEjNipW63KRUv6TBbXrwk6"),
	(
		"SKxny894fJe2rmZjeuoE6GVfNkWoXfPp8337VrLLNWG56FfQtuS1",
		"zcbxovDeXGJJikZH5wQkcQvYx1gzsRt9mR5UnQir6NY8hhPHdgK7z7dE1vfa55Bq3JHJu7isfuWQGYrvMbLnud74z2vS4tS",
		"ZiVKZCknLZuWKZyyGoKDA4h2NRQeYkE5aLQFPPKAh7LYFjSKLhm4BAbk8ELyuwukmRxv29Cws8vaQNYLmq6unwvjZ1AgA5tZM",
	),
	(
		"SKxoo5QkFQgTbdc6EWRKyHPMdmtNDJhqudrAVhen9b4kjCwN6CeV",
		"zcRYvLiURno1LhXq95e8avXFcH2fKKToSFfhqaVKTy8mGH7i6SJbfuWcm4h9rEA6DvswrbxDhFGDQgpdDYV8zwUoHvwNvFX",
		"ZiVKkUCm1oqVfyN7e1TY1Xq6CPZgdG18CAirp1nAQWxWcDufxUUkmw11S1fRG23AVWv36rud1Fp3zDEUXEfLfbGyFU6XGTkKM",
	),
	(
		"SKxsVGKsCESoVb3Gfm762psjRtGHmjmv7HVjHckud5MnESfktUuG",
		"zcWGguu2UPfNhh1ygWW9Joo3osvncsuehtz5ewvXd78vFDdnDCRNG6QeKSZpwZmYmkfEutPVf8HzCfBytqXWsEcF2iBAM1e",
		"ZiVKhFPyGkDJjTXUuSi9yWBSnkJF8jLcnE3wSNLnexSgwGeoR8N7ZHRfoyyBMRiM1nbNVDPmcTQPYPBqmS1MDDh1B9fSTauDD",
	),
	(
		"SKxp72QGQ2qtovHSoVnPp8jRFQpHBhG1xF8s27iRFjPXXkYMQUA6",
		"zcWZomPYMEjJ49S4UHcvTnhjYqogfdYJuEDMURDpbkrz94bkzdTdJEZKWkkpQ8nK62eyLkZCvLZDFtLC2Cq5BmEK3WCKGMN",
		"ZiVKnYSZMncaac1SmjB9nfbY3xPdwwiJCVP9inpy59pxVoAWAjSCHTn4rcbQJznWajZALLiFHkKE9XCZQAwSuSqSVbPKtjr5v",
	),
	(
		"SKxpmLdykLu3xxSXtw1EA7iLJnXu8hFh8hhmW1B2J2194ijh5CR4",
		"zcgjj3fJF59QGBufopx3F51jCjUpXbgEzec7YQT6jRt4Ebu5EV3AW4jHPN6ZdXhmygBvQDRJrXoZLa3Lkh5GqnsFUzt7Qok",
		"ZiVKk7ZNUE2sjkNkmaJ5PwBEWm287XFZWeXvoAYcWacCtmuEpAGRDqSfvQstB6erm5bt8uptG3bSEjNipW63KRUv6TBbXrwk6",
	),
    ];

    #[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>Viewing keys should use a_pk not pk_enc<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 viewing key.
pub struct ViewingKey {
    pub sk_enc: [u8; 32],
    pub a_pk: [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 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);
        let pk = &Scalar(sk_enc) * &ED25519_BASEPOINT;
        let pk_enc = pk.compress_montgomery().unwrap().to_bytes();

        PaymentAddress {
            a_pk: a_pk,
            pk_enc: pk_enc,
        }
    }

    /// Computes a viewing key for this spending key.
    pub fn viewing_key(&self) -> ViewingKey {
        let mut sk_enc = [0u8; 32];
        let mut a_pk = [0u8; 32];
        pseudorandom_function_sk_enc(&mut sk_enc, &self.a_sk);
        clamp_curve25519(&mut sk_enc);
        pseudorandom_function_a_pk(&mut a_pk, &self.a_sk);

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

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>//! Determine how well a cockroach has performed in a environment
//!
//! A score can be given to a robot cockroach brain program.
//!
//! # Examples
//! To give an example of how to use a scorecard
//!
//! ```rust
//! #[macro_use]
//! extern crate nadezhda;
//!
//! use nadezhda::grammar::Program;
//! use nadezhda::environment::Environment;
//! use nadezhda::fitness::{Config, Score, ScoreCard};
//!
//! fn main() {
//!     let program = forward!(forward!(stop!()));
//!     let environment = Environment::new(3);
//!     let config = Config {
//!         program_length_weight: 1,
//!         food_distance_weight: 1,
//!     };
//!
//!     let score = config.score(program, environment);
//!
//!     assert_eq!(score, 4 as Score);
//! }
//! ```

use super::grammar::Program;
use super::environment::Environment;
use super::evaluate::Evaluator;

/// A score is just a number
pub type Score = i32;

/// A scorecard can keep score a program on an environment
pub trait ScoreCard {
    /// return a score
    fn score(&self, program: Program, environment: Environment) -> Score;
}

trait Length {
    fn length(&self) -> i32;
}

impl Length for Program {
    fn length(&self) -> i32 {
        1 + match *self {
            Program::Forward(ref contained_program) => {
                contained_program.length()
            },
            Program::Backward(ref contained_program) => {
                contained_program.length()
            },
            Program::Stop => 0
        }
    }
}

/// Configuration for a default implementation for a ScoreCard
pub struct Config {
    /// The weight with which the program length adds to the score
    pub program_length_weight: i32,
    /// The weight with which the distance to the source of food adds to the score
    pub food_distance_weight: i32,
}

impl ScoreCard for Config {
    fn score (&self, program: Program, environment: Environment) -> Score {
        let program_length = program.length();

        let final_environment = program.evaluate_on(environment);
        let food_distance = (final_environment.cockroach_location - final_environment.food_location).abs();

        self.program_length_weight * program_length + self.food_distance_weight * food_distance
    }
}

#[cfg(test)]
mod test {
    use super::super::grammar::Program;
    use super::super::environment::Environment;
    use super::*;

    #[test]
    fn should_determine_score() {
        let program = Program::Stop;
        let environment = Environment::new(5);
        let config = Config {
            program_length_weight: 1,
            food_distance_weight: 1,
        };

        let score = config.score(program, environment);

        assert_eq!(score, 6 as Score);
    }
}
<commit_msg> The score depends quadratically on the distance<commit_after>//! Determine how well a cockroach has performed in a environment
//!
//! A score can be given to a robot cockroach brain program.
//!
//! # Examples
//! To give an example of how to use a scorecard
//!
//! ```rust
//! #[macro_use]
//! extern crate nadezhda;
//!
//! use nadezhda::grammar::Program;
//! use nadezhda::environment::Environment;
//! use nadezhda::fitness::{Config, Score, ScoreCard};
//!
//! fn main() {
//!     let program = forward!(forward!(stop!()));
//!     let environment = Environment::new(3);
//!     let config = Config {
//!         program_length_weight: 1,
//!         food_distance_weight: 1,
//!     };
//!
//!     let score = config.score(program, environment);
//!
//!     assert_eq!(score, 4 as Score);
//! }
//! ```

use super::grammar::Program;
use super::environment::Environment;
use super::evaluate::Evaluator;

/// A score is just a number
pub type Score = i32;

/// A scorecard can keep score a program on an environment
pub trait ScoreCard {
    /// return a score
    fn score(&self, program: Program, environment: Environment) -> Score;
}

trait Length {
    fn length(&self) -> i32;
}

impl Length for Program {
    fn length(&self) -> i32 {
        1 + match *self {
            Program::Forward(ref contained_program) => {
                contained_program.length()
            },
            Program::Backward(ref contained_program) => {
                contained_program.length()
            },
            Program::Stop => 0
        }
    }
}

/// Configuration for a default implementation for a ScoreCard
pub struct Config {
    /// The weight with which the program length adds to the score
    pub program_length_weight: i32,
    /// The weight with which the distance to the source of food adds to the score
    pub food_distance_weight: i32,
}

impl ScoreCard for Config {
    fn score (&self, program: Program, environment: Environment) -> Score {
        let program_length = program.length();

        let final_environment = program.evaluate_on(environment);
        let food_distance = (final_environment.cockroach_location - final_environment.food_location).abs().pow(2);

        self.program_length_weight * program_length + self.food_distance_weight * food_distance
    }
}

#[cfg(test)]
mod test {
    use super::super::grammar::Program;
    use super::super::environment::Environment;
    use super::*;

    #[test]
    fn should_determine_score() {
        let program = Program::Stop;
        let environment = Environment::new(5);
        let config = Config {
            program_length_weight: 1,
            food_distance_weight: 1,
        };

        let score = config.score(program, environment);

        assert_eq!(score, 26 as Score);
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::string::*;

#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Soniton {
    pub bits: u8,
    pub be: bool,
    pub packed: bool,
    pub planar: bool,
    pub float: bool,
    pub signed: bool,
}

// TODO: make it a trait for usize?
fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

fn round_to_byte(v: usize) -> usize {
    (v + 7) >> 3
}

impl Soniton {
    pub fn new(bits: u8, be: bool, packed: bool, planar: bool, float: bool, signed: bool) -> Self {
        Soniton {
            bits,
            be,
            packed,
            planar,
            float,
            signed,
        }
    }

    pub fn get_audio_size(self, length: usize, alignment: usize) -> usize {
        let s = if self.packed {
            round_to_byte(length * (self.bits as usize))
        } else {
            length * round_to_byte(self.bits as usize)
        };

        align(s, alignment)
    }
}

impl fmt::Display for Soniton {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let fmt = if self.float {
            "float"
        } else if self.signed {
            "int"
        } else {
            "uint"
        };
        let end = if self.be { "BE" } else { "LE" };
        write!(
            f,
            "({} bps, {} planar: {} packed: {} {})",
            self.bits, end, self.packed, self.planar, fmt
        )
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ChannelType {
    C,
    L,
    R,
    Cs,
    Ls,
    Rs,
    Lss,
    Rss,
    LFE,
    Lc,
    Rc,
    Lh,
    Rh,
    Ch,
    LFE2,
    Lw,
    Rw,
    Ov,
    Lhs,
    Rhs,
    Chs,
    Ll,
    Rl,
    Cl,
    Lt,
    Rt,
    Lo,
    Ro,
}

impl ChannelType {
    pub fn is_center(self) -> bool {
        match self {
            ChannelType::C => true,
            ChannelType::Ch => true,
            ChannelType::Cl => true,
            ChannelType::Ov => true,
            ChannelType::LFE => true,
            ChannelType::LFE2 => true,
            ChannelType::Cs => true,
            ChannelType::Chs => true,
            _ => false,
        }
    }
    pub fn is_left(self) -> bool {
        match self {
            ChannelType::L => true,
            ChannelType::Ls => true,
            ChannelType::Lss => true,
            ChannelType::Lc => true,
            ChannelType::Lh => true,
            ChannelType::Lw => true,
            ChannelType::Lhs => true,
            ChannelType::Ll => true,
            ChannelType::Lt => true,
            ChannelType::Lo => true,
            _ => false,
        }
    }
    pub fn is_right(self) -> bool {
        match self {
            ChannelType::R => true,
            ChannelType::Rs => true,
            ChannelType::Rss => true,
            ChannelType::Rc => true,
            ChannelType::Rh => true,
            ChannelType::Rw => true,
            ChannelType::Rhs => true,
            ChannelType::Rl => true,
            ChannelType::Rt => true,
            ChannelType::Ro => true,
            _ => false,
        }
    }
}

impl fmt::Display for ChannelType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match *self {
            ChannelType::C => "C".to_string(),
            ChannelType::L => "L".to_string(),
            ChannelType::R => "R".to_string(),
            ChannelType::Cs => "Cs".to_string(),
            ChannelType::Ls => "Ls".to_string(),
            ChannelType::Rs => "Rs".to_string(),
            ChannelType::Lss => "Lss".to_string(),
            ChannelType::Rss => "Rss".to_string(),
            ChannelType::LFE => "LFE".to_string(),
            ChannelType::Lc => "Lc".to_string(),
            ChannelType::Rc => "Rc".to_string(),
            ChannelType::Lh => "Lh".to_string(),
            ChannelType::Rh => "Rh".to_string(),
            ChannelType::Ch => "Ch".to_string(),
            ChannelType::LFE2 => "LFE2".to_string(),
            ChannelType::Lw => "Lw".to_string(),
            ChannelType::Rw => "Rw".to_string(),
            ChannelType::Ov => "Ov".to_string(),
            ChannelType::Lhs => "Lhs".to_string(),
            ChannelType::Rhs => "Rhs".to_string(),
            ChannelType::Chs => "Chs".to_string(),
            ChannelType::Ll => "Ll".to_string(),
            ChannelType::Rl => "Rl".to_string(),
            ChannelType::Cl => "Cl".to_string(),
            ChannelType::Lt => "Lt".to_string(),
            ChannelType::Rt => "Rt".to_string(),
            ChannelType::Lo => "Lo".to_string(),
            ChannelType::Ro => "Ro".to_string(),
        };
        write!(f, "{}", name)
    }
}

#[derive(Clone, Debug, PartialEq, Default)]
pub struct ChannelMap {
    ids: Vec<ChannelType>,
}

impl ChannelMap {
    pub fn new() -> Self {
        ChannelMap { ids: Vec::new() }
    }
    pub fn add_channel(&mut self, ch: ChannelType) {
        self.ids.push(ch);
    }
    pub fn add_channels(&mut self, chs: &[ChannelType]) {
        for ch in chs {
            self.ids.push(*ch);
        }
    }
    pub fn len(&self) -> usize {
        self.ids.len()
    }
    pub fn get_channel(&self, idx: usize) -> ChannelType {
        self.ids[idx]
    }
    pub fn find_channel_id(&self, t: ChannelType) -> Option<u8> {
        for i in 0..self.ids.len() {
            if self.ids[i] as i32 == t as i32 {
                return Some(i as u8);
            }
        }
        None
    }
    pub fn default_map(count: usize) -> Self {
        use self::ChannelType::*;
        let ids = match count {
            1 => vec![C],
            2 => vec![R, L],
            _ => unimplemented!(),
        };

        ChannelMap { ids }
    }
}

pub mod formats {
    use super::*;

    pub const U8: Soniton = Soniton {
        bits: 8,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: false,
    };
    pub const S16: Soniton = Soniton {
        bits: 16,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: true,
    };
    pub const S32: Soniton = Soniton {
        bits: 32,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: true,
    };
    pub const F32: Soniton = Soniton {
        bits: 32,
        be: false,
        packed: false,
        planar: false,
        float: true,
        signed: true,
    };
    pub const F64: Soniton = Soniton {
        bits: 64,
        be: false,
        packed: false,
        planar: false,
        float: true,
        signed: true,
    };
}

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

    #[test]
    fn fmt() {
        println!("{}", formats::S16);
        println!("{}", formats::U8);
        println!("{}", formats::F32);
    }
}
<commit_msg>Add some documentation for the audio_sample module<commit_after>use std::fmt;
use std::string::*;

/// Parameters describing the representation of an audio sample
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Soniton {
    /// Number of bits in each sample.
    pub bits: u8,
    /// True if the sample data is in big endian layout, false if little endian.
    pub be: bool,
    /// For data which isn't a multiple of 8 this is true if there is no padding between samples,
    /// false if the data is padded.
    pub packed: bool,
    /// If true data is stored in a planar layout with channels occurring in sequence i.e. C1 C1
    /// C1... C2 C2 C2 . If false data is interleaved i.e. C1 C2 C1 C2.
    pub planar: bool,
    /// True if the sample data is a floating point data type, false otherwise.
    pub float: bool,
    /// True if the sample data is signed, false otherwise.
    pub signed: bool,
}

// TODO: make it a trait for usize?
/// Given a length in bytes `v` return the size when aligned to `a` bytes
fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

/// Round the size in bits to the number of bytes required
fn round_to_byte(v: usize) -> usize {
    (v + 7) >> 3
}

impl Soniton {
    /// Creates a new representation of a sample with the provided parameters
    pub fn new(bits: u8, be: bool, packed: bool, planar: bool, float: bool, signed: bool) -> Self {
        Soniton {
            bits,
            be,
            packed,
            planar,
            float,
            signed,
        }
    }

    /// For a stream of the given length and alignment return the size of the buffer containing the
    /// data
    pub fn get_audio_size(self, length: usize, alignment: usize) -> usize {
        let s = if self.packed {
            round_to_byte(length * (self.bits as usize))
        } else {
            length * round_to_byte(self.bits as usize)
        };

        align(s, alignment)
    }
}

impl fmt::Display for Soniton {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let fmt = if self.float {
            "float"
        } else if self.signed {
            "int"
        } else {
            "uint"
        };
        let end = if self.be { "BE" } else { "LE" };
        write!(
            f,
            "({} bps, {} planar: {} packed: {} {})",
            self.bits, end, self.packed, self.planar, fmt
        )
    }
}

/// Enum representing the different types of audio channels
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ChannelType {
    C,
    L,
    R,
    Cs,
    Ls,
    Rs,
    Lss,
    Rss,
    LFE,
    Lc,
    Rc,
    Lh,
    Rh,
    Ch,
    LFE2,
    Lw,
    Rw,
    Ov,
    Lhs,
    Rhs,
    Chs,
    Ll,
    Rl,
    Cl,
    Lt,
    Rt,
    Lo,
    Ro,
}

impl ChannelType {
    /// Returns true if the channel is centered
    pub fn is_center(self) -> bool {
        match self {
            ChannelType::C => true,
            ChannelType::Ch => true,
            ChannelType::Cl => true,
            ChannelType::Ov => true,
            ChannelType::LFE => true,
            ChannelType::LFE2 => true,
            ChannelType::Cs => true,
            ChannelType::Chs => true,
            _ => false,
        }
    }

    /// Returns true if the channel is a left audio channel, false otherwise
    pub fn is_left(self) -> bool {
        match self {
            ChannelType::L => true,
            ChannelType::Ls => true,
            ChannelType::Lss => true,
            ChannelType::Lc => true,
            ChannelType::Lh => true,
            ChannelType::Lw => true,
            ChannelType::Lhs => true,
            ChannelType::Ll => true,
            ChannelType::Lt => true,
            ChannelType::Lo => true,
            _ => false,
        }
    }

    /// Returns true if the channel is a right audio channel, false otherwise
    pub fn is_right(self) -> bool {
        match self {
            ChannelType::R => true,
            ChannelType::Rs => true,
            ChannelType::Rss => true,
            ChannelType::Rc => true,
            ChannelType::Rh => true,
            ChannelType::Rw => true,
            ChannelType::Rhs => true,
            ChannelType::Rl => true,
            ChannelType::Rt => true,
            ChannelType::Ro => true,
            _ => false,
        }
    }
}

impl fmt::Display for ChannelType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match *self {
            ChannelType::C => "C".to_string(),
            ChannelType::L => "L".to_string(),
            ChannelType::R => "R".to_string(),
            ChannelType::Cs => "Cs".to_string(),
            ChannelType::Ls => "Ls".to_string(),
            ChannelType::Rs => "Rs".to_string(),
            ChannelType::Lss => "Lss".to_string(),
            ChannelType::Rss => "Rss".to_string(),
            ChannelType::LFE => "LFE".to_string(),
            ChannelType::Lc => "Lc".to_string(),
            ChannelType::Rc => "Rc".to_string(),
            ChannelType::Lh => "Lh".to_string(),
            ChannelType::Rh => "Rh".to_string(),
            ChannelType::Ch => "Ch".to_string(),
            ChannelType::LFE2 => "LFE2".to_string(),
            ChannelType::Lw => "Lw".to_string(),
            ChannelType::Rw => "Rw".to_string(),
            ChannelType::Ov => "Ov".to_string(),
            ChannelType::Lhs => "Lhs".to_string(),
            ChannelType::Rhs => "Rhs".to_string(),
            ChannelType::Chs => "Chs".to_string(),
            ChannelType::Ll => "Ll".to_string(),
            ChannelType::Rl => "Rl".to_string(),
            ChannelType::Cl => "Cl".to_string(),
            ChannelType::Lt => "Lt".to_string(),
            ChannelType::Rt => "Rt".to_string(),
            ChannelType::Lo => "Lo".to_string(),
            ChannelType::Ro => "Ro".to_string(),
        };
        write!(f, "{}", name)
    }
}

/// Stores the audio channels in the file in order of occurence
#[derive(Clone, Debug, PartialEq, Default)]
pub struct ChannelMap {
    ids: Vec<ChannelType>,
}

impl ChannelMap {
    /// Creates a new channel map
    pub fn new() -> Self {
        ChannelMap { ids: Vec::new() }
    }

    /// Adds a single channel to the map
    pub fn add_channel(&mut self, ch: ChannelType) {
        self.ids.push(ch);
    }

    /// Adds number of channels to the map in order of occurrence
    pub fn add_channels(&mut self, chs: &[ChannelType]) {
        for ch in chs {
            self.ids.push(*ch);
        }
    }

    /// Returns the number of channels in the stream
    pub fn len(&self) -> usize {
        self.ids.len()
    }

    /// Gets the channel type for the channel at the given index
    pub fn get_channel(&self, idx: usize) -> ChannelType {
        self.ids[idx]
    }

    /// Finds the first channel index containing a channel of the given type
    pub fn find_channel_id(&self, t: ChannelType) -> Option<u8> {
        for i in 0..self.ids.len() {
            if self.ids[i] as i32 == t as i32 {
                return Some(i as u8);
            }
        }
        None
    }

    /// Creates a default channel map. For 1 channel this is a single centred channel for 2
    /// channels a Right channel and a Left channel. No implementation currently exists for other
    /// counts.
    pub fn default_map(count: usize) -> Self {
        use self::ChannelType::*;
        let ids = match count {
            1 => vec![C],
            2 => vec![R, L],
            _ => unimplemented!(),
        };

        ChannelMap { ids }
    }
}

/// A set of default constant channels for general use
pub mod formats {
    use super::*;

    pub const U8: Soniton = Soniton {
        bits: 8,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: false,
    };
    pub const S16: Soniton = Soniton {
        bits: 16,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: true,
    };
    pub const S32: Soniton = Soniton {
        bits: 32,
        be: false,
        packed: false,
        planar: false,
        float: false,
        signed: true,
    };
    pub const F32: Soniton = Soniton {
        bits: 32,
        be: false,
        packed: false,
        planar: false,
        float: true,
        signed: true,
    };
    pub const F64: Soniton = Soniton {
        bits: 64,
        be: false,
        packed: false,
        planar: false,
        float: true,
        signed: true,
    };
}

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

    #[test]
    fn fmt() {
        println!("{}", formats::S16);
        println!("{}", formats::U8);
        println!("{}", formats::F32);
    }
}
<|endoftext|>


<commit_before>////////////////
// INTERFACES //
////////////////

extern crate byteorder;
use self::byteorder::{ByteOrder, BigEndian, LittleEndian};

use errors::*;

/// Address type/size.
/// Just a simple alias to usize for easier-to-read code. (In my opinion, obviously.)
/// Originally, I wanted to make it use generics,
/// but that got messy really quick.
pub type Addr = usize;

/// Byte.
/// Convenience alias.
pub type Byte = u8;

/// Simple trait for a finite memory block.
pub trait MemoryBlock {
    /// Get the block's accessible size.
    /// Simply the highest address,
    /// NOT the number of addresses.
    fn get_size(&self) -> Addr;

    /// Set a byte at address.
    fn set(&mut self, Addr, Byte) -> Result<(), Error>;

    /// Get a byte at address.
    /// Returns `Ok(X)` on success, where X will be the byte.
    fn get(&self, Addr) -> Result<Byte, Error>;

    /// Delete data at address.
    /// May allow the block to efficiently delete it, marking it as unused.
    /// This could allow the block to do wear leveling, for example.
    fn delete(&mut self, from: Addr, to: Addr) -> Result<(), Error> {
        if from == to {
            self.set(from, 0).chain_err(|| format!("failure to delete {:#X}", from))?;
        }
        for i in from..to {
            self.set(i, 0).chain_err(|| format!("failure to delete byte {} in {:#X}-{:#X}", i-from, from, to))?;
        }
        self.set(to, 0).chain_err(|| format!("failure to delete {:#X}", to))
    }

    /// Flush writes out.
    /// In case it does any form of caching, calling this method
    /// ensures it has written all the data it needs to write.
    /// Or, well, it fails otherwise.
    fn flush(&mut self) -> Result<(), Error> {
        Ok(())
    }
}

/// Subtrait for 32 bit big-endian interface
pub trait MemoryBlock32be: MemoryBlock {
    /// Set a 32 bit value (4 bytes) in Big Endian format.
    fn set32be(&mut self, addr: Addr, data: u32) -> Result<(), Error> {
        let mut bytes = [0; 4];
        BigEndian::write_u32(&mut bytes, data);
        for i in 0..4 {
            self.set(addr + i, bytes[i]).chain_err(|| format!("failure to set32be {:#X} in byte {}", addr, i))?;
        }
        Ok(())
    }
    /// Get a 32 bit value (4 bytes) stored in Big Endian format.
    fn get32be(&mut self, addr: Addr) -> Result<u32, Error> {
        let mut bytes = [0; 4];
        for i in 0..4 {
            bytes[i] = self.get(addr + i).chain_err(|| format!("failure to get32be {:#X} in byte {}", addr, i))?;
        }
        Ok(BigEndian::read_u32(&bytes))
    }
}
/// Subtrait for 32 bit little-endian interface
pub trait MemoryBlock32le: MemoryBlock {
    /// Set a 32 bit value (4 bytes) in Little Endian format.
    fn set32le(&mut self, addr: Addr, data: u32) -> Result<(), Error> {
        let mut bytes = [0; 4];
        LittleEndian::write_u32(&mut bytes, data);
        for i in 0..4 {
            self.set(addr + i, bytes[i]).chain_err(|| format!("failure to set32le {:#X} in byte {}", addr, i))?;
        }
        Ok(())
    }
    /// Get a 32 bit value (4 bytes) stored in Little Endian format.
    fn get32le(&mut self, addr: Addr) -> Result<u32, Error> {
        let mut bytes = [0; 4];
        for i in 0..4 {
            bytes[i] = self.get(addr + i).chain_err(|| format!("failure to get32le {:#X} in byte {}", addr, i))?;
        }
        Ok(LittleEndian::read_u32(&bytes))
    }
}

/// Simple `new(size)` trait for backends.
// For convenience, mostly.
pub trait MemoryCreator<T> {
    /// Returns a `MemoryBlock` of size `size`.
    /// `size` being the highest address,
    /// NOT number of addresses.
    fn new(size: Addr) -> T;
}

/// Trait for a new function that wraps
/// another memory block, returning one itself.
///
/// Useful for invalid access checkers, debuggers, etc....
pub trait MemoryMiddlewareCreator<T> {
    fn new<M: MemoryBlock>(Box<M>) -> T;
}
<commit_msg>Oops, this is not supposed to be mutable.<commit_after>////////////////
// INTERFACES //
////////////////

extern crate byteorder;
use self::byteorder::{ByteOrder, BigEndian, LittleEndian};

use errors::*;

/// Address type/size.
/// Just a simple alias to usize for easier-to-read code. (In my opinion, obviously.)
/// Originally, I wanted to make it use generics,
/// but that got messy really quick.
pub type Addr = usize;

/// Byte.
/// Convenience alias.
pub type Byte = u8;

/// Simple trait for a finite memory block.
pub trait MemoryBlock {
    /// Get the block's accessible size.
    /// Simply the highest address,
    /// NOT the number of addresses.
    fn get_size(&self) -> Addr;

    /// Set a byte at address.
    fn set(&mut self, Addr, Byte) -> Result<(), Error>;

    /// Get a byte at address.
    /// Returns `Ok(X)` on success, where X will be the byte.
    fn get(&self, Addr) -> Result<Byte, Error>;

    /// Delete data at address.
    /// May allow the block to efficiently delete it, marking it as unused.
    /// This could allow the block to do wear leveling, for example.
    fn delete(&mut self, from: Addr, to: Addr) -> Result<(), Error> {
        if from == to {
            self.set(from, 0).chain_err(|| format!("failure to delete {:#X}", from))?;
        }
        for i in from..to {
            self.set(i, 0).chain_err(|| format!("failure to delete byte {} in {:#X}-{:#X}", i-from, from, to))?;
        }
        self.set(to, 0).chain_err(|| format!("failure to delete {:#X}", to))
    }

    /// Flush writes out.
    /// In case it does any form of caching, calling this method
    /// ensures it has written all the data it needs to write.
    /// Or, well, it fails otherwise.
    fn flush(&mut self) -> Result<(), Error> {
        Ok(())
    }
}

/// Subtrait for 32 bit big-endian interface
pub trait MemoryBlock32be: MemoryBlock {
    /// Set a 32 bit value (4 bytes) in Big Endian format.
    fn set32be(&mut self, addr: Addr, data: u32) -> Result<(), Error> {
        let mut bytes = [0; 4];
        BigEndian::write_u32(&mut bytes, data);
        for i in 0..4 {
            self.set(addr + i, bytes[i]).chain_err(|| format!("failure to set32be {:#X} in byte {}", addr, i))?;
        }
        Ok(())
    }
    /// Get a 32 bit value (4 bytes) stored in Big Endian format.
    fn get32be(&self, addr: Addr) -> Result<u32, Error> {
        let mut bytes = [0; 4];
        for i in 0..4 {
            bytes[i] = self.get(addr + i).chain_err(|| format!("failure to get32be {:#X} in byte {}", addr, i))?;
        }
        Ok(BigEndian::read_u32(&bytes))
    }
}
/// Subtrait for 32 bit little-endian interface
pub trait MemoryBlock32le: MemoryBlock {
    /// Set a 32 bit value (4 bytes) in Little Endian format.
    fn set32le(&mut self, addr: Addr, data: u32) -> Result<(), Error> {
        let mut bytes = [0; 4];
        LittleEndian::write_u32(&mut bytes, data);
        for i in 0..4 {
            self.set(addr + i, bytes[i]).chain_err(|| format!("failure to set32le {:#X} in byte {}", addr, i))?;
        }
        Ok(())
    }
    /// Get a 32 bit value (4 bytes) stored in Little Endian format.
    fn get32le(&self, addr: Addr) -> Result<u32, Error> {
        let mut bytes = [0; 4];
        for i in 0..4 {
            bytes[i] = self.get(addr + i).chain_err(|| format!("failure to get32le {:#X} in byte {}", addr, i))?;
        }
        Ok(LittleEndian::read_u32(&bytes))
    }
}

/// Simple `new(size)` trait for backends.
// For convenience, mostly.
pub trait MemoryCreator<T> {
    /// Returns a `MemoryBlock` of size `size`.
    /// `size` being the highest address,
    /// NOT number of addresses.
    fn new(size: Addr) -> T;
}

/// Trait for a new function that wraps
/// another memory block, returning one itself.
///
/// Useful for invalid access checkers, debuggers, etc....
pub trait MemoryMiddlewareCreator<T> {
    fn new<M: MemoryBlock>(Box<M>) -> T;
}
<|endoftext|>


<commit_before><commit_msg>Rename to_str to to_string<commit_after><|endoftext|>


<commit_before>extern crate xml;
extern crate regex;

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::Read;
use std::io::BufReader;

use xml::reader::{EventReader, XmlEvent};
use xml::attribute::OwnedAttribute;

#[derive(Debug, Clone, Default)]
struct LineAttributes {
    top: i32,
    left: i32,
    height: i32,
    page: u32,
}

fn read_attributes(attr_list: &Vec<OwnedAttribute>) -> LineAttributes {
    let mut attributes : LineAttributes = Default::default();

    for attr in attr_list {
        match attr.name.local_name.as_ref() {
            "top" => attributes.top = attr.value.parse().unwrap(),
            "left" => attributes.left = attr.value.parse().unwrap(),
            "height" => attributes.height = attr.value.parse().unwrap(),
            _ => {}
        }
    }

    attributes
}

#[derive(Debug, Clone, Default)]
struct ScriptProperties {
    direction_position: i32,
    dialog_position: i32,
    speaker_direction_position: i32,
    speaker_position: i32,
    intra_paragraph_line_height: i32,
}

fn read_and_analyze_script(reader: Box<Read>) -> (ScriptProperties, Vec<(LineAttributes, String)>) {
    let mut script_properties: ScriptProperties = Default::default();
    let mut lines: Vec<(LineAttributes, String)> = Vec::new();

    // these maps are used to heuristically determine the ScriptProperties
    let mut position_uses: HashMap<i32, u32> = HashMap::new();
    let mut line_height_uses: HashMap<i32, u32> = HashMap::new();

    // states for the streaming xml parsing
    let mut current_line_attributes: LineAttributes = Default::default();
    let mut current_text_buffer = String::new();
    let mut current_page_number = 0;
    let mut last_line_height = 0;

    let parser = EventReader::new(reader);
    for event in parser {
        match event {
            Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                match name.local_name.as_ref() {
                    "text" => {
                        current_line_attributes = read_attributes(&attributes);

                        current_text_buffer.clear();

                        // increase the count for the current left margin
                        *position_uses.entry(current_line_attributes.left)
                                      .or_insert(0) += 1;

                        // increase the count for the current top position diff
                        *line_height_uses.entry(current_line_attributes.top - last_line_height)
                                         .or_insert(0) += 1;
                        last_line_height = current_line_attributes.top;
                    }
                    "page" => {
                        last_line_height = 0;

                        for attr in attributes {
                            if "number" == attr.name.local_name {
                                current_page_number = attr.value.parse().unwrap();
                                break;
                            }
                        }
                    }
                    _ => {}
                }
            }
            Ok(XmlEvent::Characters(text)) => {
                if current_text_buffer.len() > 0 {
                    current_text_buffer.push(' ');
                }
                current_text_buffer.push_str(text.trim());
            }
            Ok(XmlEvent::EndElement { name, .. }) => {
                if name.local_name == "text" {
                    current_line_attributes.page = current_page_number;
                    lines.push((current_line_attributes.clone(), current_text_buffer.clone()));
                }
            }
            Ok(_) => {},
            Err(e) => panic!("Error parsing xml: {}", e),
        }
    }

    // the position_uses map should at least have 3 different entries
    if position_uses.len() < 3 {
        panic!("Script uses strange layout");
    }

    // copy the position uses map into a vector and sort by value desc
    let mut position_uses_vec = position_uses.into_iter().collect::<Vec<_>>();
    position_uses_vec.sort_by(|&a, &b| b.1.cmp(&a.1));

    // sort the top 3 entries by position again (right-exclusive range)
    position_uses_vec[0..3].sort_by(|&a, &b| b.0.cmp(&a.0));

    script_properties.speaker_position = position_uses_vec[0].0;
    script_properties.dialog_position = position_uses_vec[1].0;
    script_properties.direction_position = position_uses_vec[2].0;

    // take the next value between speaker and dialog as speaker_direction
    for &value in position_uses_vec.iter().skip(3) {
        let (position, _) = value;
        if position > script_properties.dialog_position &&
           position < script_properties.speaker_position {
            script_properties.speaker_direction_position = position;
            break;
        }
    }

    // use to most used line diff value to determine the specific sections
    let mut last_uses: u32 = 0;
    for (&line_height, &uses) in line_height_uses.iter() {
        if uses > last_uses {
            script_properties.intra_paragraph_line_height = line_height;
            last_uses = uses;
        }
    }

    (script_properties, lines)
}

#[derive(Debug, Clone)]
enum DialogPart {
    Dialog(String),
    Direction(String),
}

#[derive(Debug, Clone)]
enum ScenePart {
    Direction {
        direction: String,
        page: u32,
    },
    Dialog {
        speaker: String,
        dialog: Vec<DialogPart>,
        page: u32,
    }
}

#[derive(Debug, Clone)]
enum ScriptPart {
    Separator,
    ScenePart(ScenePart),
    LocationChange(String),
    SceneChange,
}

fn is_location_change(line: &str) -> bool {
    line.starts_with("INT.") || line.starts_with("EXT.")
}

fn is_scene_change(line: &str) -> bool {
    line.starts_with("CUT TO")
}

fn extract_script_parts(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>)
    -> Vec<ScriptPart> {
    let mut script_parts: Vec<ScriptPart> = Vec::new();
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // check if a new section starts
        if attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            // used to separate two consecutive script parts of the
            // same type. this is just a implementation detail of
            // the parsing. the Separator can be ignored later.
            script_parts.push(ScriptPart::Separator);
        }

        if attributes.left == properties.direction_position {
            if is_location_change(line) {
                script_parts.push(ScriptPart::LocationChange(line.clone()));
            } else {
                // ensure the last script part is a direction
                if let Some(&ScriptPart::ScenePart(ScenePart::Direction{..})) = script_parts.last() {
                } else {
                    script_parts.push(ScriptPart::ScenePart(
                        ScenePart::Direction{
                            direction: String::new(),
                            page: attributes.page,
                        }));
                }

                if let Some(&mut ScriptPart::ScenePart(ScenePart::Direction{ref mut direction, ..})) = script_parts.last_mut() {
                    if direction.len() > 0 {
                        direction.push(' ');
                    }
                    direction.push_str(line);
                }
            }
        } else if attributes.left == properties.speaker_position ||
                  attributes.left == properties.speaker_direction_position ||
                  attributes.left == properties.dialog_position {
            // Ensure the last script part is a dialog
            if let Some(&ScriptPart::ScenePart(ScenePart::Dialog{..})) = script_parts.last() {
            } else {
                script_parts.push(ScriptPart::ScenePart(
                    ScenePart::Dialog{
                        speaker: String::new(),
                        dialog: Vec::new(),
                        page: attributes.page,
                    }));
            }

            // get the dialog, should never fail (see above)
            if let Some(&mut ScriptPart::ScenePart(ScenePart::Dialog{ref mut speaker, ref mut dialog, ..})) = script_parts.last_mut() {
                if attributes.left == properties.speaker_position {
                    // there is only one speaker per dialog
                    speaker.push_str(line);
                } else if attributes.left == properties.speaker_direction_position {
                    if let Some(&DialogPart::Direction(_)) = dialog.last() {
                    } else {
                        dialog.push(DialogPart::Direction(String::new()));
                    }

                    if let Some(&mut DialogPart::Direction(ref mut direction)) = dialog.last_mut() {
                        if direction.len() > 0 {
                            direction.push(' ');
                        }
                        direction.push_str(line);
                    }
                } else if attributes.left == properties.dialog_position {
                    if let Some(&DialogPart::Dialog(_)) = dialog.last() {
                    } else {
                        dialog.push(DialogPart::Dialog(String::new()));
                    }

                    if let Some(&mut DialogPart::Dialog(ref mut dialog)) = dialog.last_mut() {
                        if dialog.len() > 0 {
                            dialog.push(' ');
                        }
                        dialog.push_str(line);
                    }
                }
            }
        } else {
            if is_scene_change(line) {
                script_parts.push(ScriptPart::SceneChange);
            }
        }

        last_top_position = attributes.top;
    }

    script_parts
}

#[derive(Clone, Debug)]
enum LocationType {
    Undefined,
    Internal,
    External,
    InternalExternal,
}

impl Default for LocationType {
    fn default() -> LocationType { LocationType::Undefined }
}

fn extract_location(name: &str) -> Location {
    let pattern = regex::Regex::new(r"(?:(?P<kind>INT\.|EXT\.|INT\./EXT\.)\s+)?(?P<location>.+)").unwrap();
    let mut location: Location = Default::default();

    if let Some(captures) = pattern.captures(name) {
        location.name = captures.name("location").unwrap().to_string();

        if let Some(location_kind) = captures.name("kind") {
            location.kind = match location_kind {
                "INT." => LocationType::Internal,
                "EXT." => LocationType::External,
                "INT./EXT." => LocationType::InternalExternal,
                _ => LocationType::Undefined,
            };
        }
    } else {
        location.name = name.to_string();
    }

    location
}

#[derive(Default, Clone, Debug)]
struct Location {
    kind: LocationType,
    name: String,
    parts: Vec<ScenePart>,
}

type Scene = Vec<Location>;

fn extract_scenes(script_parts: &Vec<ScriptPart>) -> Vec<Scene> {
    let mut scenes = Vec::new();

    // scene with default (empty) location
    let default_scene: Scene = vec![Default::default()];
    scenes.push(default_scene.clone());

    for script_part in script_parts.iter() {
        use ScriptPart::*;
        match script_part {
            &SceneChange => {
                // unwrap is save, see default_scene
                if scenes.last().unwrap().len() > 0 {
                    scenes.push(default_scene.clone());
                }
            }
            &LocationChange(ref location) => {
                // unwraps are safe, see default_scene
                let mut current_scene = scenes.last_mut().unwrap();

                if current_scene.last().unwrap().parts.len() == 0 {
                    current_scene.pop();
                }
                current_scene.push(extract_location(location));
            }
            &ScenePart(ref scene_part) => {
                let scene_parts = &mut scenes.last_mut().unwrap().last_mut().unwrap().parts;
                scene_parts.push(scene_part.clone());
            }
            &Separator => {} //ignore
        }
    }

    scenes
}

#[allow(dead_code)]
fn condense_script_demo(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>) {
    let mut last_left_position = 0;
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // start new section
        if last_left_position != attributes.left ||
           attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            if ! (last_left_position == properties.speaker_direction_position ||
                  last_left_position == properties.speaker_position) {
                // empty line after last section
                println!("");
            }
        }

        match attributes.left {
            k if (k == properties.direction_position)
                => println!("{}", line),
            k if (k == properties.dialog_position)
                => println!("          {}", line),
            k if (k == properties.speaker_direction_position)
                => println!("               {}", line),
            k if (k == properties.speaker_position)
                => println!("                    {}", line),
            _   => println!("------- {}", line),
        }

        last_left_position = attributes.left;
        last_top_position = attributes.top;
    }
}

fn main() {
    let xml_file_name = env::args().skip(1).next().expect("No input file given");

    let file_reader = File::open(xml_file_name).expect("Could not open file");
    let buffered_file_reader = Box::new(BufReader::new(file_reader));

    let (properties, lines) = read_and_analyze_script(buffered_file_reader);

    //condense_script_demo(properties, &lines);
    println!("{:#?}", extract_scenes(&extract_script_parts(properties, &lines)));
}
<commit_msg>Implement output to xml<commit_after>extern crate xml;
extern crate regex;

use std::collections::HashMap;
use std::env;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::io::BufReader;

use xml::{EventReader, EventWriter};
use xml::reader::XmlEvent;
use xml::writer::Result as XmlResult;
use xml::attribute::OwnedAttribute;

#[derive(Debug, Clone, Default)]
struct LineAttributes {
    top: i32,
    left: i32,
    height: i32,
    page: u32,
}

fn read_attributes(attr_list: &Vec<OwnedAttribute>) -> LineAttributes {
    let mut attributes : LineAttributes = Default::default();

    for attr in attr_list {
        match attr.name.local_name.as_ref() {
            "top" => attributes.top = attr.value.parse().unwrap(),
            "left" => attributes.left = attr.value.parse().unwrap(),
            "height" => attributes.height = attr.value.parse().unwrap(),
            _ => {}
        }
    }

    attributes
}

#[derive(Debug, Clone, Default)]
struct ScriptProperties {
    direction_position: i32,
    dialog_position: i32,
    speaker_direction_position: i32,
    speaker_position: i32,
    intra_paragraph_line_height: i32,
}

fn read_and_analyze_script(reader: Box<Read>) -> (ScriptProperties, Vec<(LineAttributes, String)>) {
    let mut script_properties: ScriptProperties = Default::default();
    let mut lines: Vec<(LineAttributes, String)> = Vec::new();

    // these maps are used to heuristically determine the ScriptProperties
    let mut position_uses: HashMap<i32, u32> = HashMap::new();
    let mut line_height_uses: HashMap<i32, u32> = HashMap::new();

    // states for the streaming xml parsing
    let mut current_line_attributes: LineAttributes = Default::default();
    let mut current_text_buffer = String::new();
    let mut current_page_number = 0;
    let mut last_line_height = 0;

    let parser = EventReader::new(reader);
    for event in parser {
        match event {
            Ok(XmlEvent::StartElement { name, attributes, .. }) => {
                match name.local_name.as_ref() {
                    "text" => {
                        current_line_attributes = read_attributes(&attributes);

                        current_text_buffer.clear();

                        // increase the count for the current left margin
                        *position_uses.entry(current_line_attributes.left)
                                      .or_insert(0) += 1;

                        // increase the count for the current top position diff
                        *line_height_uses.entry(current_line_attributes.top - last_line_height)
                                         .or_insert(0) += 1;
                        last_line_height = current_line_attributes.top;
                    }
                    "page" => {
                        last_line_height = 0;

                        for attr in attributes {
                            if "number" == attr.name.local_name {
                                current_page_number = attr.value.parse().unwrap();
                                break;
                            }
                        }
                    }
                    _ => {}
                }
            }
            Ok(XmlEvent::Characters(text)) => {
                if current_text_buffer.len() > 0 {
                    current_text_buffer.push(' ');
                }
                current_text_buffer.push_str(text.trim());
            }
            Ok(XmlEvent::EndElement { name, .. }) => {
                if name.local_name == "text" {
                    current_line_attributes.page = current_page_number;
                    lines.push((current_line_attributes.clone(), current_text_buffer.clone()));
                }
            }
            Ok(_) => {},
            Err(e) => panic!("Error parsing xml: {}", e),
        }
    }

    // the position_uses map should at least have 3 different entries
    if position_uses.len() < 3 {
        panic!("Script uses strange layout");
    }

    // copy the position uses map into a vector and sort by value desc
    let mut position_uses_vec = position_uses.into_iter().collect::<Vec<_>>();
    position_uses_vec.sort_by(|&a, &b| b.1.cmp(&a.1));

    // sort the top 3 entries by position again (right-exclusive range)
    position_uses_vec[0..3].sort_by(|&a, &b| b.0.cmp(&a.0));

    script_properties.speaker_position = position_uses_vec[0].0;
    script_properties.dialog_position = position_uses_vec[1].0;
    script_properties.direction_position = position_uses_vec[2].0;

    // take the next value between speaker and dialog as speaker_direction
    for &value in position_uses_vec.iter().skip(3) {
        let (position, _) = value;
        if position > script_properties.dialog_position &&
           position < script_properties.speaker_position {
            script_properties.speaker_direction_position = position;
            break;
        }
    }

    // use to most used line diff value to determine the specific sections
    let mut last_uses: u32 = 0;
    for (&line_height, &uses) in line_height_uses.iter() {
        if uses > last_uses {
            script_properties.intra_paragraph_line_height = line_height;
            last_uses = uses;
        }
    }

    (script_properties, lines)
}

#[derive(Debug, Clone)]
enum DialogPart {
    Dialog(String),
    Direction(String),
}

#[derive(Debug, Clone)]
enum ScenePart {
    Direction {
        direction: String,
        page: u32,
    },
    Dialog {
        speaker: String,
        dialog: Vec<DialogPart>,
        page: u32,
    }
}

#[derive(Debug, Clone)]
enum ScriptPart {
    Separator,
    ScenePart(ScenePart),
    LocationChange(String),
    SceneChange,
}

fn is_location_change(line: &str) -> bool {
    line.starts_with("INT.") || line.starts_with("EXT.")
}

fn is_scene_change(line: &str) -> bool {
    line.starts_with("CUT TO")
}

fn extract_script_parts(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>)
    -> Vec<ScriptPart> {
    let mut script_parts: Vec<ScriptPart> = Vec::new();
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // check if a new section starts
        if attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            // used to separate two consecutive script parts of the
            // same type. this is just a implementation detail of
            // the parsing. the Separator can be ignored later.
            script_parts.push(ScriptPart::Separator);
        }

        if attributes.left == properties.direction_position {
            if is_location_change(line) {
                script_parts.push(ScriptPart::LocationChange(line.clone()));
            } else {
                // ensure the last script part is a direction
                if let Some(&ScriptPart::ScenePart(ScenePart::Direction{..})) = script_parts.last() {
                } else {
                    script_parts.push(ScriptPart::ScenePart(
                        ScenePart::Direction{
                            direction: String::new(),
                            page: attributes.page,
                        }));
                }

                if let Some(&mut ScriptPart::ScenePart(ScenePart::Direction{ref mut direction, ..})) = script_parts.last_mut() {
                    if direction.len() > 0 {
                        direction.push(' ');
                    }
                    direction.push_str(line);
                }
            }
        } else if attributes.left == properties.speaker_position ||
                  attributes.left == properties.speaker_direction_position ||
                  attributes.left == properties.dialog_position {
            // Ensure the last script part is a dialog
            if let Some(&ScriptPart::ScenePart(ScenePart::Dialog{..})) = script_parts.last() {
            } else {
                script_parts.push(ScriptPart::ScenePart(
                    ScenePart::Dialog{
                        speaker: String::new(),
                        dialog: Vec::new(),
                        page: attributes.page,
                    }));
            }

            // get the dialog, should never fail (see above)
            if let Some(&mut ScriptPart::ScenePart(ScenePart::Dialog{ref mut speaker, ref mut dialog, ..})) = script_parts.last_mut() {
                if attributes.left == properties.speaker_position {
                    // there is only one speaker per dialog
                    speaker.push_str(line);
                } else if attributes.left == properties.speaker_direction_position {
                    if let Some(&DialogPart::Direction(_)) = dialog.last() {
                    } else {
                        dialog.push(DialogPart::Direction(String::new()));
                    }

                    if let Some(&mut DialogPart::Direction(ref mut direction)) = dialog.last_mut() {
                        if direction.len() > 0 {
                            direction.push(' ');
                        }
                        direction.push_str(line);
                    }
                } else if attributes.left == properties.dialog_position {
                    if let Some(&DialogPart::Dialog(_)) = dialog.last() {
                    } else {
                        dialog.push(DialogPart::Dialog(String::new()));
                    }

                    if let Some(&mut DialogPart::Dialog(ref mut dialog)) = dialog.last_mut() {
                        if dialog.len() > 0 {
                            dialog.push(' ');
                        }
                        dialog.push_str(line);
                    }
                }
            }
        } else {
            if is_scene_change(line) {
                script_parts.push(ScriptPart::SceneChange);
            }
        }

        last_top_position = attributes.top;
    }

    script_parts
}

#[derive(Clone, Debug)]
enum LocationType {
    Undefined,
    Internal,
    External,
    InternalExternal,
}

impl Default for LocationType {
    fn default() -> LocationType { LocationType::Undefined }
}

impl std::ops::Deref for LocationType {
    type Target = str;
    fn deref(&self) -> &str {
        match *self {
            LocationType::Undefined => "",
            LocationType::Internal => "internal",
            LocationType::External => "external",
            LocationType::InternalExternal => "internal,external",
        }
    }
}

fn extract_location(name: &str) -> Location {
    let pattern = regex::Regex::new(r"(?:(?P<kind>INT\.|EXT\.|INT\./EXT\.)\s+)?(?P<location>.+)").unwrap();
    let mut location: Location = Default::default();

    if let Some(captures) = pattern.captures(name) {
        location.name = captures.name("location").unwrap().to_string();

        if let Some(location_kind) = captures.name("kind") {
            location.kind = match location_kind {
                "INT." => LocationType::Internal,
                "EXT." => LocationType::External,
                "INT./EXT." => LocationType::InternalExternal,
                _ => LocationType::Undefined,
            };
        }
    } else {
        location.name = name.to_string();
    }

    location
}

#[derive(Default, Clone, Debug)]
struct Location {
    kind: LocationType,
    name: String,
    parts: Vec<ScenePart>,
}

type Scene = Vec<Location>;

fn extract_scenes(script_parts: &Vec<ScriptPart>) -> Vec<Scene> {
    let mut scenes = Vec::new();

    // scene with default (empty) location
    let default_scene: Scene = vec![Default::default()];
    scenes.push(default_scene.clone());

    for script_part in script_parts.iter() {
        use ScriptPart::*;
        match script_part {
            &SceneChange => {
                // unwrap is save, see default_scene
                if scenes.last().unwrap().len() > 0 {
                    scenes.push(default_scene.clone());
                }
            }
            &LocationChange(ref location) => {
                // unwraps are safe, see default_scene
                let mut current_scene = scenes.last_mut().unwrap();

                if current_scene.last().unwrap().parts.len() == 0 {
                    current_scene.pop();
                }
                current_scene.push(extract_location(location));
            }
            &ScenePart(ref scene_part) => {
                let scene_parts = &mut scenes.last_mut().unwrap().last_mut().unwrap().parts;
                scene_parts.push(scene_part.clone());
            }
            &Separator => {} //ignore
        }
    }

    scenes
}

fn format_scene_parts<W: Write>(scene_parts: &Vec<ScenePart>, writer: &mut EventWriter<W>) -> XmlResult<()> {
    use xml::writer::XmlEvent;

    for part in scene_parts.iter() {
        match part {
            &ScenePart::Direction { ref direction, ref page } => {
                try!(writer.write(XmlEvent::start_element("direction")
                                           .attr("page", page.to_string().as_ref())));

                try!(writer.write(XmlEvent::characters(direction)));

                try!(writer.write(XmlEvent::end_element()));
            }
            &ScenePart::Dialog { ref speaker, ref dialog, ref page } => {
                try!(writer.write(XmlEvent::start_element("dialog")
                                           .attr("speaker", speaker)
                                           .attr("page", page.to_string().as_ref())));

                for (i, dialog_part) in dialog.iter().enumerate() {
                    match dialog_part {
                        &DialogPart::Dialog(ref dialog) => {
                            try!(writer.write(XmlEvent::characters(dialog)));
                        }
                        &DialogPart::Direction(ref direction) => {
                            try!(writer.write(XmlEvent::start_element("direction")));
                            try!(writer.write(XmlEvent::characters(direction)));
                            try!(writer.write(XmlEvent::end_element()));
                        }
                    }

                    if i + 1 != dialog.len() {
                        try!(writer.write(XmlEvent::characters(" ")));
                    }
                }

                try!(writer.write(XmlEvent::end_element()));
            }
        }
    }

    Ok(())
}

fn format_scenes<W: Write>(scenes: &Vec<Scene>, output: &mut W) -> XmlResult<()> {
    use xml::writer::XmlEvent;

    let mut writer = xml::EmitterConfig::new().perform_indent(true).create_writer(output);

    try!(writer.write(XmlEvent::start_element("script")));

    for scene in scenes.iter() {
        try!(writer.write(XmlEvent::start_element("scene")));

        for location in scene.iter() {
            let mut location_event = XmlEvent::start_element("location");
            if location.name.len() > 0 {
                location_event = location_event.attr("name", &location.name);
            }
            match location.kind {
                LocationType::Undefined => {}
                _ => {
                    location_event = location_event.attr("kind", &location.kind);
                }
            }
            try!(writer.write(location_event));

            try!(format_scene_parts(&location.parts, &mut writer));

            try!(writer.write(XmlEvent::end_element()));
        }

        try!(writer.write(XmlEvent::end_element()));
    }

    try!(writer.write(XmlEvent::end_element()));

    Ok(())
}

#[allow(dead_code)]
fn condense_script_demo(properties: ScriptProperties, lines: &Vec<(LineAttributes, String)>) {
    let mut last_left_position = 0;
    let mut last_top_position = 0;

    for line in lines.iter() {
        let &(ref attributes, ref line) = line;

        if line.len() == 0 {
            continue;
        }

        // start new section
        if last_left_position != attributes.left ||
           attributes.top - last_top_position > 18 ||
           attributes.top - last_top_position < 0 {
            if ! (last_left_position == properties.speaker_direction_position ||
                  last_left_position == properties.speaker_position) {
                // empty line after last section
                println!("");
            }
        }

        match attributes.left {
            k if (k == properties.direction_position)
                => println!("{}", line),
            k if (k == properties.dialog_position)
                => println!("          {}", line),
            k if (k == properties.speaker_direction_position)
                => println!("               {}", line),
            k if (k == properties.speaker_position)
                => println!("                    {}", line),
            _   => println!("------- {}", line),
        }

        last_left_position = attributes.left;
        last_top_position = attributes.top;
    }
}

fn main() {
    let xml_file_name = env::args().skip(1).next().expect("No input file given");

    let file_reader = File::open(xml_file_name).expect("Could not open file");
    let buffered_file_reader = Box::new(BufReader::new(file_reader));

    let (properties, lines) = read_and_analyze_script(buffered_file_reader);

    //condense_script_demo(properties, &lines);
    format_scenes(&extract_scenes(&extract_script_parts(properties, &lines)), &mut std::io::stdout()).unwrap();
}
<|endoftext|>


<commit_before>extern crate rs6502;
extern crate sdl2;

mod input;
mod ship;
mod timer;

use std::io::{self, BufRead, Write};
use std::path::Path;
use std::sync::mpsc::channel;
use std::time::Duration;

use rs6502::{Assembler, Cpu, Disassembler};

use sdl2::event::Event;
use sdl2::keyboard::Keycode;
use sdl2::image::LoadTexture;
use sdl2::Sdl;

const FPS: u32 = 1000 / 60;

fn main() {
    let (tx, rx) = channel();

    input::handle(tx);

    let mut assembler = Assembler::new();
    let bytecode = assembler.assemble_file("level.asm").unwrap();
    let mut cpu = init_cpu(&bytecode);

    let sdl_context = sdl2::init().unwrap();
    let video_subsystem = sdl_context.video().unwrap();

    let window = video_subsystem.window("Hakka", 1280, 400)
        .position_centered()
        .build()
        .unwrap();

    let mut renderer = window.renderer()
        .accelerated()
        .build()
        .unwrap();

    let ship_texture = renderer.load_texture(Path::new("ship.png")).unwrap();

    let mut events = sdl_context.event_pump().unwrap();

    let mut ship = ship::Ship::new(ship_texture);

    let mut last_fps = 0;
    let mut monitor_enabled = false;
    let mut monitor_last = 0;

    'running: loop {
        cpu.step();

        if let Ok(input) = rx.try_recv() {
            let input = input.trim();
            if input == "exit" {
                break 'running;
            }

            if input == "list" {
                std::io::stdout().write(b"\n");
                std::io::stdout().write(b"-- Disassembly --\n");

                let mut disassembler = Disassembler::with_offset(0xC000);
                let asm = disassembler.disassemble(&bytecode[..]);

                std::io::stdout().write(asm.as_bytes());
                std::io::stdout().write(b"\nHAKKA> ");
                std::io::stdout().flush();
            }

            if input == "monitor" {
                monitor_enabled = !monitor_enabled;
            }

            if input.as_bytes().starts_with(b"memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else {
                    let dst = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
                    let src = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

                    cpu.memory[dst as usize] = src;
                }
            }
        }

        ship.process(&cpu.memory[..]);

        for event in events.poll_iter() {
            match event {
                Event::Quit { .. } => break 'running,
                Event::KeyDown { keycode: Option::Some(Keycode::Right), .. } => {
                    cpu.memory[0x04] = 39
                }
                Event::KeyDown { keycode: Option::Some(Keycode::Left), .. } => {
                    cpu.memory[0x04] = 37
                }
                Event::KeyUp { keycode: Option::Some(Keycode::Right), .. } => cpu.memory[0x04] = 0,
                Event::KeyUp { keycode: Option::Some(Keycode::Left), .. } => cpu.memory[0x04] = 0,
                Event::KeyDown { keycode: Option::Some(Keycode::Escape), .. } => break 'running,
                _ => (),
            }
        }

        let now = sdl_context.timer().unwrap().ticks();
        let delta = now - last_fps;
        if delta < FPS {
            sdl_context.timer().unwrap().delay(FPS - delta);
        } else {
            if !cpu.flags.interrupt_disabled {
                // Render stuff here
                renderer.clear();
                ship.render(&mut renderer);
                renderer.present();
                last_fps = now;
            }
        }

        let delta = now - monitor_last;
        if delta > 1000 && monitor_enabled {
            for b in &cpu.memory[0x00..0xA] {
                print!("{:02X} ", *b);
            }
            println!("");
            monitor_last = now;
        }

        // thread::sleep(Duration::from_millis(10));
    }
}

fn init_cpu(bytecode: &[u8]) -> Cpu {
    let mut cpu = Cpu::new();
    cpu.load(&bytecode[..], None);
    cpu.flags.interrupt_disabled = false;

    cpu.memory[0x02] = 0x80;
    cpu.memory[0x03] = 0x00;
    cpu.memory[0x05] = 0x19;
    cpu.memory[0x06] = 0x00;

    cpu
}<commit_msg>Move assembling to its own function<commit_after>extern crate rs6502;
extern crate sdl2;

mod input;
mod ship;

use std::io::{self, BufRead, Write};
use std::path::Path;
use std::sync::mpsc::channel;
use std::time::Duration;

use rs6502::{Assembler, Cpu, Disassembler};

use sdl2::event::Event;
use sdl2::keyboard::Keycode;
use sdl2::image::LoadTexture;
use sdl2::Sdl;

const FPS: u32 = 1000 / 60;

fn main() {
    let (tx, rx) = channel();

    input::handle(tx);

    let bytecode = assemble("level.asm");
    let mut cpu = init_cpu(&bytecode);

    let sdl_context = sdl2::init().unwrap();
    let video_subsystem = sdl_context.video().unwrap();

    let window = video_subsystem.window("Hakka", 1280, 400)
        .position_centered()
        .build()
        .unwrap();

    let mut renderer = window.renderer()
        .accelerated()
        .build()
        .unwrap();

    let ship_texture = renderer.load_texture(Path::new("ship.png")).unwrap();

    let mut events = sdl_context.event_pump().unwrap();

    let mut ship = ship::Ship::new(ship_texture);

    let mut last_fps = 0;
    let mut monitor_enabled = false;
    let mut monitor_last = 0;

    'running: loop {
        cpu.step();

        if let Ok(input) = rx.try_recv() {
            let input = input.trim();
            if input == "exit" {
                break 'running;
            }

            if input == "list" {
                std::io::stdout().write(b"\n");
                std::io::stdout().write(b"-- Disassembly --\n");

                let mut disassembler = Disassembler::with_offset(0xC000);
                let asm = disassembler.disassemble(&bytecode[..]);

                std::io::stdout().write(asm.as_bytes());
                std::io::stdout().write(b"\nHAKKA> ");
                std::io::stdout().flush();
            }

            if input == "monitor" {
                monitor_enabled = !monitor_enabled;
            }

            if input.as_bytes().starts_with(b"memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else {
                    let dst = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
                    let src = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

                    cpu.memory[dst as usize] = src;
                }
            }
        }

        ship.process(&cpu.memory[..]);

        for event in events.poll_iter() {
            match event {
                Event::Quit { .. } => break 'running,
                Event::KeyDown { keycode: Option::Some(Keycode::Right), .. } => {
                    cpu.memory[0x04] = 39
                }
                Event::KeyDown { keycode: Option::Some(Keycode::Left), .. } => {
                    cpu.memory[0x04] = 37
                }
                Event::KeyUp { keycode: Option::Some(Keycode::Right), .. } => cpu.memory[0x04] = 0,
                Event::KeyUp { keycode: Option::Some(Keycode::Left), .. } => cpu.memory[0x04] = 0,
                Event::KeyDown { keycode: Option::Some(Keycode::Escape), .. } => break 'running,
                _ => (),
            }
        }

        let now = sdl_context.timer().unwrap().ticks();
        let delta = now - last_fps;
        if delta < FPS {
            sdl_context.timer().unwrap().delay(FPS - delta);
        } else {
            if !cpu.flags.interrupt_disabled {
                // Render stuff here
                renderer.clear();
                ship.render(&mut renderer);
                renderer.present();
                last_fps = now;
            }
        }

        let delta = now - monitor_last;
        if delta > 1000 && monitor_enabled {
            for b in &cpu.memory[0x00..0xA] {
                print!("{:02X} ", *b);
            }
            println!("");
            monitor_last = now;
        }

        // thread::sleep(Duration::from_millis(10));
    }
}

fn assemble<P>(path: P) -> Vec<u8>
    where P: AsRef<Path>
{
    let mut assembler = Assembler::new();
    assembler.assemble_file("level.asm").unwrap()
}

fn init_cpu(bytecode: &[u8]) -> Cpu {
    let mut cpu = Cpu::new();
    cpu.load(&bytecode[..], None);
    cpu.flags.interrupt_disabled = false;

    cpu.memory[0x02] = 0x80;
    cpu.memory[0x03] = 0x00;
    cpu.memory[0x05] = 0x19;
    cpu.memory[0x06] = 0x00;

    cpu
}<|endoftext|>