<commit_before>use super::utils::types::{ Sha256Digest, Bitfield, StateHash };
use super::utils::bls::{ Signature, AggregateSignature, Keypair, PublicKey };
use super::aggregate_vote::AggregateVote;
use super::rlp::{ RlpStream, Encodable } ;

pub struct Block {
    pub parent_hash: Sha256Digest,
    pub skip_count: u64,
    pub randao_reveal: Sha256Digest,
    pub attestation_bitfield: Bitfield,
    pub attestation_aggregate_sig: AggregateSignature,
    pub shard_aggregate_votes: Vec<AggregateVote>,
    pub main_chain_ref: Sha256Digest,
    pub state_hash: StateHash,
    pub sig: Option<Signature>
}
impl Block {
    pub fn new(parent_hash: Sha256Digest,
               randao_reveal: Sha256Digest,
               main_chain_ref: Sha256Digest,
               state_hash: StateHash) -> Block {
        Block {
            parent_hash: parent_hash,
            skip_count: 0,
            randao_reveal: randao_reveal,
            attestation_bitfield: Bitfield::new(),
            attestation_aggregate_sig: AggregateSignature::new(),
            shard_aggregate_votes: Vec::new(),
            main_chain_ref: main_chain_ref,
            state_hash: state_hash,
            sig: None
        }
    }

    /*
     * Take a Block and covert it into an array of u8 for BLS signing
     * or verfication. The `sig` field is purposefully omitted.
     */
    pub fn encode_to_signable_message(&self) -> [u8; 9140] {
        // Using biggest avg. block size from v2 spec
        let mut message: [u8; 9140] = [0; 9140];

        // Create the RLP vector
        let mut s = RlpStream::new();
        s.append(&self.parent_hash);
        s.append(&self.skip_count);
        s.append(&self.randao_reveal);
        s.append(&self.attestation_bitfield);
        // s.append(&self.attestation_aggregate_sig);   // TODO: RLP this
        s.append_list(&self.shard_aggregate_votes);
        s.append(&self.main_chain_ref);
        // TODO: state hash serialization is probably incorrect.
        s.append(&self.state_hash.crystallized_state);
        s.append(&self.state_hash.active_state);
        let rlp_vec = s.out();

        // Parse the RLP vector into an array compatible with the BLS signer
        let len = rlp_vec.len();
        message[..len].copy_from_slice(&rlp_vec[..len]);
        message
    }

    /*
     * Sign the block with the given keypair.
     */
    pub fn sig_sign(&mut self, keypair: &Keypair) {
        let message = self.encode_to_signable_message();
        self.sig = Some(keypair.sign(&message));
    }

    /*
     * Verify a block signature given some keypair.
     */
    pub fn sig_verify(&self, pub_key: &PublicKey) -> bool {
        let message = self.encode_to_signable_message();
        match &self.sig {
            None => false,
            Some(sig) => {
                pub_key.verify(&message, &sig)
            },
        }
    }
}

impl Encodable for Block {
    fn rlp_append(&self, s: &mut RlpStream) {
        s.append(&self.parent_hash);
        s.append(&self.skip_count);
        s.append(&self.randao_reveal);
        s.append(&self.attestation_bitfield);
        // s.append(&self.attestation_aggregate_sig);   // TODO: RLP this
        s.append_list(&self.shard_aggregate_votes);
        s.append(&self.main_chain_ref);
        // TODO: state hash serialization is probably incorrect.
        s.append(&self.state_hash.crystallized_state);
        s.append(&self.state_hash.active_state);
        // s.append(&self.sig);    // TODO: RLP this
    }
}

#[cfg(test)]
mod tests {
    use super::super::rlp;
    extern crate rand;

    use super::*;
    use self::rand::{ SeedableRng, XorShiftRng };

    #[test]
    fn test_signable_message_encoding() {
        let parent_hash = Sha256Digest::from([0; 32]);
        let randao_reveal = Sha256Digest::from([1; 32]);
        let main_chain_ref = Sha256Digest::from([2; 32]);
        let state_hash = StateHash::zero();
        let mut b = Block::new(parent_hash,
                           randao_reveal,
                           main_chain_ref,
                           state_hash);
        b.skip_count = 2;
        let output = b.encode_to_signable_message();
        // TODO: test this better
        assert_eq!(output[0], 160);
        assert_eq!(output[1..21], [0; 20]);
    }

    #[test]
    fn test_sign_and_verify() {
        let mut rng = XorShiftRng::from_seed([0xbc4f6d44, 0xd62f276c, 0xb963afd0, 0x5455863d]);
        let alice_keypair = Keypair::generate(&mut rng);
        let bob_keypair = Keypair::generate(&mut rng);
        let mut b = Block::new(Sha256Digest::random(),
                           Sha256Digest::random(),
                           Sha256Digest::random(),
                           StateHash::zero());

        // Both signatures fail before signing
        assert_eq!(b.sig_verify(&alice_keypair.public), false);
        assert_eq!(b.sig_verify(&bob_keypair.public), false);

        // Sign as Alice
        b.sig_sign(&alice_keypair);

        // Alice signature passes, bobs fails
        assert_eq!(b.sig_verify(&alice_keypair.public), true);
        assert_eq!(b.sig_verify(&bob_keypair.public), false);
    }

    #[test]
    fn test_rlp_serialization() {
        let b = Block {
            parent_hash: Sha256Digest::zero(),
            skip_count: 100,
            randao_reveal: Sha256Digest::zero(),
            attestation_bitfield: Bitfield::new(),
            attestation_aggregate_sig: AggregateSignature::new(),
            shard_aggregate_votes: Vec::new(),
            main_chain_ref: Sha256Digest::zero(),
            state_hash: StateHash::zero(),
            sig: None
        };
        let e = rlp::encode(&b);
        assert_eq!(e.len(), 168);
        assert_eq!(e[0], 160);
        assert_eq!(e[1..33], [0; 32]);
        assert_eq!(e[33], 100);
        assert_eq!(e[34], 160);
        assert_eq!(e[35..67], [0; 32]);
        assert_eq!(e[67], 128);
        assert_eq!(e[69], 160);
        assert_eq!(e[70..102], [0; 32]);
        /*
        assert_eq!(e[102], 248);
        assert_eq!(e[103], 64);
        assert_eq!(e[104..136], [128; 32]);
        assert_eq!(e[136..168], [128; 32]);
        */
    }
}
<commit_msg>Add zero() method for Block<commit_after>use super::utils::types::{ Sha256Digest, Bitfield, StateHash };
use super::utils::bls::{ Signature, AggregateSignature, Keypair, PublicKey };
use super::aggregate_vote::AggregateVote;
use super::rlp::{ RlpStream, Encodable } ;

pub struct Block {
    pub parent_hash: Sha256Digest,
    pub skip_count: u64,
    pub randao_reveal: Sha256Digest,
    pub attestation_bitfield: Bitfield,
    pub attestation_aggregate_sig: AggregateSignature,
    pub shard_aggregate_votes: Vec<AggregateVote>,
    pub main_chain_ref: Sha256Digest,
    pub state_hash: StateHash,
    pub sig: Option<Signature>
}
impl Block {
    pub fn new(parent_hash: Sha256Digest,
               randao_reveal: Sha256Digest,
               main_chain_ref: Sha256Digest,
               state_hash: StateHash) -> Block {
        Block {
            parent_hash: parent_hash,
            skip_count: 0,
            randao_reveal: randao_reveal,
            attestation_bitfield: Bitfield::new(),
            attestation_aggregate_sig: AggregateSignature::new(),
            shard_aggregate_votes: Vec::new(),
            main_chain_ref: main_chain_ref,
            state_hash: state_hash,
            sig: None
        }
    }

    pub fn zero() -> Block {
        Block {
            parent_hash: Sha256Digest::zero(),
            skip_count: 0,
            randao_reveal: Sha256Digest::zero(),
            attestation_bitfield: Bitfield::new(),
            attestation_aggregate_sig: AggregateSignature::new(),
            shard_aggregate_votes: vec![],
            main_chain_ref: Sha256Digest::zero(),
            state_hash: StateHash::zero(),
            sig: None
        }
    }

    /*
     * Take a Block and covert it into an array of u8 for BLS signing
     * or verfication. The `sig` field is purposefully omitted.
     */
    pub fn encode_to_signable_message(&self) -> [u8; 9140] {
        // Using biggest avg. block size from v2 spec
        let mut message: [u8; 9140] = [0; 9140];

        // Create the RLP vector
        let mut s = RlpStream::new();
        s.append(&self.parent_hash);
        s.append(&self.skip_count);
        s.append(&self.randao_reveal);
        s.append(&self.attestation_bitfield);
        // s.append(&self.attestation_aggregate_sig);   // TODO: RLP this
        s.append_list(&self.shard_aggregate_votes);
        s.append(&self.main_chain_ref);
        // TODO: state hash serialization is probably incorrect.
        s.append(&self.state_hash.crystallized_state);
        s.append(&self.state_hash.active_state);
        let rlp_vec = s.out();

        // Parse the RLP vector into an array compatible with the BLS signer
        let len = rlp_vec.len();
        message[..len].copy_from_slice(&rlp_vec[..len]);
        message
    }

    /*
     * Sign the block with the given keypair.
     */
    pub fn sig_sign(&mut self, keypair: &Keypair) {
        let message = self.encode_to_signable_message();
        self.sig = Some(keypair.sign(&message));
    }

    /*
     * Verify a block signature given some keypair.
     */
    pub fn sig_verify(&self, pub_key: &PublicKey) -> bool {
        let message = self.encode_to_signable_message();
        match &self.sig {
            None => false,
            Some(sig) => {
                pub_key.verify(&message, &sig)
            },
        }
    }
}

impl Encodable for Block {
    fn rlp_append(&self, s: &mut RlpStream) {
        s.append(&self.parent_hash);
        s.append(&self.skip_count);
        s.append(&self.randao_reveal);
        s.append(&self.attestation_bitfield);
        // s.append(&self.attestation_aggregate_sig);   // TODO: RLP this
        s.append_list(&self.shard_aggregate_votes);
        s.append(&self.main_chain_ref);
        // TODO: state hash serialization is probably incorrect.
        s.append(&self.state_hash.crystallized_state);
        s.append(&self.state_hash.active_state);
        // s.append(&self.sig);    // TODO: RLP this
    }
}

#[cfg(test)]
mod tests {
    use super::super::rlp;
    extern crate rand;

    use super::*;
    use self::rand::{ SeedableRng, XorShiftRng };

    #[test]
    fn test_signable_message_encoding() {
        let parent_hash = Sha256Digest::from([0; 32]);
        let randao_reveal = Sha256Digest::from([1; 32]);
        let main_chain_ref = Sha256Digest::from([2; 32]);
        let state_hash = StateHash::zero();
        let mut b = Block::new(parent_hash,
                           randao_reveal,
                           main_chain_ref,
                           state_hash);
        b.skip_count = 2;
        let output = b.encode_to_signable_message();
        // TODO: test this better
        assert_eq!(output[0], 160);
        assert_eq!(output[1..21], [0; 20]);
    }

    #[test]
    fn test_sign_and_verify() {
        let mut rng = XorShiftRng::from_seed([0xbc4f6d44, 0xd62f276c, 0xb963afd0, 0x5455863d]);
        let alice_keypair = Keypair::generate(&mut rng);
        let bob_keypair = Keypair::generate(&mut rng);
        let mut b = Block::new(Sha256Digest::random(),
                           Sha256Digest::random(),
                           Sha256Digest::random(),
                           StateHash::zero());

        // Both signatures fail before signing
        assert_eq!(b.sig_verify(&alice_keypair.public), false);
        assert_eq!(b.sig_verify(&bob_keypair.public), false);

        // Sign as Alice
        b.sig_sign(&alice_keypair);

        // Alice signature passes, bobs fails
        assert_eq!(b.sig_verify(&alice_keypair.public), true);
        assert_eq!(b.sig_verify(&bob_keypair.public), false);
    }

    #[test]
    fn test_rlp_serialization() {
        let b = Block {
            parent_hash: Sha256Digest::zero(),
            skip_count: 100,
            randao_reveal: Sha256Digest::zero(),
            attestation_bitfield: Bitfield::new(),
            attestation_aggregate_sig: AggregateSignature::new(),
            shard_aggregate_votes: Vec::new(),
            main_chain_ref: Sha256Digest::zero(),
            state_hash: StateHash::zero(),
            sig: None
        };
        let e = rlp::encode(&b);
        assert_eq!(e.len(), 168);
        assert_eq!(e[0], 160);
        assert_eq!(e[1..33], [0; 32]);
        assert_eq!(e[33], 100);
        assert_eq!(e[34], 160);
        assert_eq!(e[35..67], [0; 32]);
        assert_eq!(e[67], 128);
        assert_eq!(e[69], 160);
        assert_eq!(e[70..102], [0; 32]);
        /*
        assert_eq!(e[102], 248);
        assert_eq!(e[103], 64);
        assert_eq!(e[104..136], [128; 32]);
        assert_eq!(e[136..168], [128; 32]);
        */
    }
}
<|endoftext|>


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

//! Handling log-probabilities.

use std::mem;
use std::f64;
use std::iter;
use std::ops::{Add, Sub, Mul, Div};
use num::traits::cast;
use num::NumCast;

use itertools::linspace;
use itertools::Itertools;
use itertools::misc::ToFloat;

/// A factor to convert log-probabilities to PHRED-scale (phred = p * `LOG_TO_PHRED_FACTOR`).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * `PHRED_TO_LOG_FACTOR`).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
fn ln_1m_exp(p: f64) -> f64 {
    assert!(p <= 0.0);
    if p < -0.693 {
        (-p.exp()).ln_1p()
    } else {
        (-p.exp_m1()).ln()
    }
}

/// A newtype for probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct Prob(f64);
}

/// A newtype for log-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct LogProb(f64);
}

/// A newtype for PHRED-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct PHREDProb(f64);
}

/// Iterator returned by scans over logprobs.
pub type ScanIter<I> = iter::Scan<<I as IntoIterator>::IntoIter, LogProb, fn(&mut LogProb, LogProb) -> Option<LogProb>>;

impl LogProb {
    pub fn ln_zero() -> LogProb {
        LogProb(f64::NEG_INFINITY)
    }

    pub fn ln_one() -> LogProb {
        LogProb(0.0)
    }

    pub fn ln_one_minus_exp(&self) -> LogProb {
        LogProb(ln_1m_exp(**self))
    }

    pub fn ln_sum_exp(probs: &[LogProb]) -> LogProb {
        if probs.is_empty() {
            Self::ln_zero()
        } else {
            let mut pmax = probs[0];
            let mut imax = 0;
            for (i, &p) in probs.iter().enumerate().skip(1) {
                if p > pmax {
                    pmax = p;
                    imax = i;
                }
            }
            if pmax == Self::ln_zero() {
                Self::ln_zero()
            } else if *pmax == f64::INFINITY {
                LogProb(f64::INFINITY)
            } else {
                // TODO use sum() once it has been stabilized: .sum::<usize>()
                pmax + LogProb(
                    (probs.iter()
                          .enumerate()
                          .filter_map(|(i, p)| {
                              if i == imax {
                                  None
                              } else {
                                  Some((p - pmax).exp())
                              }
                          })
                          .fold(0.0, |s, e| s + e)
                     ).ln_1p()
                 )
            }
        }
    }

    pub fn ln_add_exp(self, other: LogProb) -> LogProb {
        let (mut p0, mut p1) = (self, other);
        if p1 > p0 {
            mem::swap(&mut p0, &mut p1);
        }
        if p0 == Self::ln_zero() {
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + LogProb((p1 - p0).exp().ln_1p())
        }
    }

    pub fn ln_sub_exp(self, other: LogProb) -> LogProb {
        let (p0, p1) = (self, other);
        assert!(p0 >= p1,
                "Subtraction would lead to negative probability, which is undefined in log space.");
        if relative_eq!(*p0, *p1) || p0 == Self::ln_zero() {
            // the first case leads to zero,
            // in the second case p0 and p1 are -inf, which is fine
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + (p1 - p0).ln_one_minus_exp()
        }
    }

    /// Calculate the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
    pub fn ln_cumsum_exp<I: IntoIterator<Item = LogProb>>(probs: I) -> ScanIter<I> {
        probs.into_iter().scan(Self::ln_zero(), Self::scan_ln_add_exp)
    }

    /// Integrate numerically stable over given log-space density in the interval [a, b]. Uses the trapezoidal rule with n grid points.
    pub fn ln_integrate_exp<T, D>(density: &D, a: T, b: T, n: usize) -> LogProb
        where T: NumCast + Copy + Add<Output=T> + Sub<Output=T> + Div<Output=T> + Mul<Output=T>, D: Fn(T) -> LogProb, usize: ToFloat<T>
    {
        let mut probs = linspace(a, b, n).dropping(1).dropping_back(1).map(|v| LogProb(*density(v) + 2.0f64.ln())).collect_vec();
        probs.push(density(a));
        probs.push(density(b));
        let width: f64 = cast(b - a).unwrap();

        LogProb(*Self::ln_sum_exp(&probs) + width.ln() - (2.0 * (n - 1) as f64).ln())
    }

    fn scan_ln_add_exp(s: &mut LogProb, p: LogProb) -> Option<LogProb> {
        *s = s.ln_add_exp(p);
        Some(*s)
    }
}

impl From<LogProb> for Prob {
    fn from(p: LogProb) -> Prob {
        Prob(p.exp())
    }
}

impl From<PHREDProb> for Prob {
    fn from(p: PHREDProb) -> Prob {
        Prob(10.0f64.powf(-*p / 10.0))
    }
}

impl From<Prob> for LogProb {
    fn from(p: Prob) -> LogProb {
        LogProb(p.ln())
    }
}

impl From<PHREDProb> for LogProb {
    fn from(p: PHREDProb) -> LogProb {
        LogProb(*p * PHRED_TO_LOG_FACTOR)
    }
}

impl From<LogProb> for PHREDProb {
    fn from(p: LogProb) -> PHREDProb {
        PHREDProb(*p * LOG_TO_PHRED_FACTOR)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64;
    use itertools::Itertools;

    #[test]
    fn test_sum() {
        let probs = [LogProb::ln_zero(), LogProb::ln_one(), LogProb::ln_zero()];
        assert_eq!(LogProb::ln_sum_exp(&probs), LogProb::ln_one());
    }

    #[test]
    fn test_empty_sum() {
        assert_eq!(LogProb::ln_sum_exp(&[]), LogProb::ln_zero());
    }

    #[test]
    fn test_cumsum() {
        let probs = vec![LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.001f64.ln())];
        assert_eq!(LogProb::ln_cumsum_exp(probs).collect_vec(),
                   [LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.011f64.ln())]);
    }

    #[test]
    fn test_sub() {
        assert_eq!(LogProb::ln_one().ln_sub_exp(LogProb::ln_one()), LogProb::ln_zero());
        assert_relative_eq!(*LogProb::ln_one().ln_sub_exp(LogProb(0.5f64.ln())), *LogProb(0.5f64.ln()));
    }

    #[test]
    fn test_one_minus() {
        assert_eq!(LogProb::ln_zero().ln_one_minus_exp(), LogProb::ln_one());
        assert_eq!(LogProb::ln_one().ln_one_minus_exp(), LogProb::ln_zero());
    }

    #[test]
    fn test_integrate() {
        let density = |_| LogProb(0.1f64.ln());
        let prob = LogProb::ln_integrate_exp(&density, 0.0, 10.0, 5);
        assert_relative_eq!(*prob, *LogProb::ln_one(), epsilon=0.0000001);
    }
}
<commit_msg>Add automatic deriving of Mul and Div traits for Prob.<commit_after>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Handling log-probabilities.

use std::mem;
use std::f64;
use std::iter;
use std::ops::{Add, Sub, Mul, Div};
use num::traits::cast;
use num::NumCast;

use itertools::linspace;
use itertools::Itertools;
use itertools::misc::ToFloat;

/// A factor to convert log-probabilities to PHRED-scale (phred = p * `LOG_TO_PHRED_FACTOR`).
const LOG_TO_PHRED_FACTOR: f64 = -4.3429448190325175; // -10 * 1 / ln(10)

/// A factor to convert PHRED-scale to log-probabilities (p = phred * `PHRED_TO_LOG_FACTOR`).
const PHRED_TO_LOG_FACTOR: f64 = -0.23025850929940456; // 1 / (-10 * log10(e))

/// Calculate log(1 - p) with p given in log space without loss of precision as described in
/// http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf.
fn ln_1m_exp(p: f64) -> f64 {
    assert!(p <= 0.0);
    if p < -0.693 {
        (-p.exp()).ln_1p()
    } else {
        (-p.exp_m1()).ln()
    }
}

/// A newtype for probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        NewtypeMul(*),
        NewtypeDiv(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct Prob(f64);
}

/// A newtype for log-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct LogProb(f64);
}

/// A newtype for PHRED-scale probabilities.
custom_derive! {
    #[derive(
        NewtypeFrom,
        NewtypeDeref,
        NewtypeAdd(*),
        NewtypeSub(*),
        PartialEq,
        PartialOrd,
        Copy,
        Clone,
        Debug
    )]
    pub struct PHREDProb(f64);
}

/// Iterator returned by scans over logprobs.
pub type ScanIter<I> = iter::Scan<<I as IntoIterator>::IntoIter, LogProb, fn(&mut LogProb, LogProb) -> Option<LogProb>>;

impl LogProb {
    pub fn ln_zero() -> LogProb {
        LogProb(f64::NEG_INFINITY)
    }

    pub fn ln_one() -> LogProb {
        LogProb(0.0)
    }

    pub fn ln_one_minus_exp(&self) -> LogProb {
        LogProb(ln_1m_exp(**self))
    }

    pub fn ln_sum_exp(probs: &[LogProb]) -> LogProb {
        if probs.is_empty() {
            Self::ln_zero()
        } else {
            let mut pmax = probs[0];
            let mut imax = 0;
            for (i, &p) in probs.iter().enumerate().skip(1) {
                if p > pmax {
                    pmax = p;
                    imax = i;
                }
            }
            if pmax == Self::ln_zero() {
                Self::ln_zero()
            } else if *pmax == f64::INFINITY {
                LogProb(f64::INFINITY)
            } else {
                // TODO use sum() once it has been stabilized: .sum::<usize>()
                pmax + LogProb(
                    (probs.iter()
                          .enumerate()
                          .filter_map(|(i, p)| {
                              if i == imax {
                                  None
                              } else {
                                  Some((p - pmax).exp())
                              }
                          })
                          .fold(0.0, |s, e| s + e)
                     ).ln_1p()
                 )
            }
        }
    }

    pub fn ln_add_exp(self, other: LogProb) -> LogProb {
        let (mut p0, mut p1) = (self, other);
        if p1 > p0 {
            mem::swap(&mut p0, &mut p1);
        }
        if p0 == Self::ln_zero() {
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + LogProb((p1 - p0).exp().ln_1p())
        }
    }

    pub fn ln_sub_exp(self, other: LogProb) -> LogProb {
        let (p0, p1) = (self, other);
        assert!(p0 >= p1,
                "Subtraction would lead to negative probability, which is undefined in log space.");
        if relative_eq!(*p0, *p1) || p0 == Self::ln_zero() {
            // the first case leads to zero,
            // in the second case p0 and p1 are -inf, which is fine
            Self::ln_zero()
        } else if *p0 == f64::INFINITY {
            LogProb(f64::INFINITY)
        } else {
            p0 + (p1 - p0).ln_one_minus_exp()
        }
    }

    /// Calculate the cumulative sum of the given probabilities in a numerically stable way (Durbin 1998).
    pub fn ln_cumsum_exp<I: IntoIterator<Item = LogProb>>(probs: I) -> ScanIter<I> {
        probs.into_iter().scan(Self::ln_zero(), Self::scan_ln_add_exp)
    }

    /// Integrate numerically stable over given log-space density in the interval [a, b]. Uses the trapezoidal rule with n grid points.
    pub fn ln_integrate_exp<T, D>(density: &D, a: T, b: T, n: usize) -> LogProb
        where T: NumCast + Copy + Add<Output=T> + Sub<Output=T> + Div<Output=T> + Mul<Output=T>, D: Fn(T) -> LogProb, usize: ToFloat<T>
    {
        let mut probs = linspace(a, b, n).dropping(1).dropping_back(1).map(|v| LogProb(*density(v) + 2.0f64.ln())).collect_vec();
        probs.push(density(a));
        probs.push(density(b));
        let width: f64 = cast(b - a).unwrap();

        LogProb(*Self::ln_sum_exp(&probs) + width.ln() - (2.0 * (n - 1) as f64).ln())
    }

    fn scan_ln_add_exp(s: &mut LogProb, p: LogProb) -> Option<LogProb> {
        *s = s.ln_add_exp(p);
        Some(*s)
    }
}

impl From<LogProb> for Prob {
    fn from(p: LogProb) -> Prob {
        Prob(p.exp())
    }
}

impl From<PHREDProb> for Prob {
    fn from(p: PHREDProb) -> Prob {
        Prob(10.0f64.powf(-*p / 10.0))
    }
}

impl From<Prob> for LogProb {
    fn from(p: Prob) -> LogProb {
        LogProb(p.ln())
    }
}

impl From<PHREDProb> for LogProb {
    fn from(p: PHREDProb) -> LogProb {
        LogProb(*p * PHRED_TO_LOG_FACTOR)
    }
}

impl From<LogProb> for PHREDProb {
    fn from(p: LogProb) -> PHREDProb {
        PHREDProb(*p * LOG_TO_PHRED_FACTOR)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64;
    use itertools::Itertools;

    #[test]
    fn test_sum() {
        let probs = [LogProb::ln_zero(), LogProb::ln_one(), LogProb::ln_zero()];
        assert_eq!(LogProb::ln_sum_exp(&probs), LogProb::ln_one());
    }

    #[test]
    fn test_empty_sum() {
        assert_eq!(LogProb::ln_sum_exp(&[]), LogProb::ln_zero());
    }

    #[test]
    fn test_cumsum() {
        let probs = vec![LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.001f64.ln())];
        assert_eq!(LogProb::ln_cumsum_exp(probs).collect_vec(),
                   [LogProb::ln_zero(), LogProb(0.01f64.ln()), LogProb(0.011f64.ln())]);
    }

    #[test]
    fn test_sub() {
        assert_eq!(LogProb::ln_one().ln_sub_exp(LogProb::ln_one()), LogProb::ln_zero());
        assert_relative_eq!(*LogProb::ln_one().ln_sub_exp(LogProb(0.5f64.ln())), *LogProb(0.5f64.ln()));
    }

    #[test]
    fn test_one_minus() {
        assert_eq!(LogProb::ln_zero().ln_one_minus_exp(), LogProb::ln_one());
        assert_eq!(LogProb::ln_one().ln_one_minus_exp(), LogProb::ln_zero());
    }

    #[test]
    fn test_integrate() {
        let density = |_| LogProb(0.1f64.ln());
        let prob = LogProb::ln_integrate_exp(&density, 0.0, 10.0, 5);
        assert_relative_eq!(*prob, *LogProb::ln_one(), epsilon=0.0000001);
    }
}
<|endoftext|>


<commit_before>/// Strongly borrowed from TypedArena
/// Differs in growth strategy (non-doubling)
/// And only for variable-sized strings

use std::borrow::BorrowFrom;
use std::cell::{Cell,RefCell};
use std::cmp::max;
use std::collections::DList;
use std::collections::hash_map::HashMap;
use std::collections::hash_map::Entry::{Occupied,Vacant};
use std::fmt;
use std::hash;
use std::mem;
use std::ptr;
use std::raw::Slice;
use std::rt::heap::{allocate, deallocate};
use std::str;
use xxhash::{XXHasher};

struct Chunk {
    start: *mut u8,
    capacity: uint,
}

impl Chunk {
    fn new(capacity: uint) -> Chunk {
        Chunk {
            start: unsafe { allocate(capacity, mem::min_align_of::<u8>()) },
            capacity: capacity,
        }
    }

    // Returns a pointer to the beginning of the allocated space.
    #[inline]
    fn start(&self) -> *const u8 {
        self.start as *const u8
    }

    // Returns a pointer to the end of the allocated space.
    #[inline]
    fn end(&self) -> *const u8 {
        unsafe { self.start().offset(self.capacity as int) }
    }
}

impl Drop for Chunk {
    fn drop(&mut self) {
        unsafe {
            deallocate(self.start, self.capacity, mem::min_align_of::<u8>());
        }
    }
}

#[derive(Copy)]
pub struct InternedString {
    slice: Slice<u8>,
}

impl InternedString {
    fn from_parts(data: *const u8, len: uint) -> InternedString {
        InternedString {
            slice: Slice {
                data: data,
                len: len,
            },
        }
    }

    pub fn from_str(s: &str) -> InternedString {
        let bytes: &[u8] = s.as_bytes();
        let slice: Slice<u8> = unsafe { mem::transmute(bytes) };
        InternedString {
            slice: slice,
        }
    }

    pub fn as_slice<'s>(&self) -> &'s str {
        unsafe {
            let bytes: &[u8] = mem::transmute(self.slice);
            str::from_utf8_unchecked(bytes)
        }
    }
}

impl fmt::Show for InternedString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl PartialEq for InternedString {
    fn eq(&self, other: &InternedString) -> bool {
        self.as_slice().eq(other.as_slice())
    }
}

impl Eq for InternedString {}

impl<S> hash::Hash<S> for InternedString
    where S: hash::Writer
{
    fn hash(&self, state: &mut S) {
        self.as_slice().hash(state)
    }
}

impl BorrowFrom<InternedString> for str {
    fn borrow_from(owned: &InternedString) -> &str {
        owned.as_slice()
    }
}

pub struct StringPool {
    start: Cell<*mut u8>,
    end: Cell<*const u8>,
    chunks: RefCell<DList<Chunk>>,
    index: RefCell<HashMap<InternedString, InternedString, XXHasher>>,
}

static CAPACITY: uint = 10240;

impl StringPool {
    pub fn new() -> StringPool {
        StringPool{
            start: Cell::new(ptr::null_mut()),
            end: Cell::new(ptr::null()),
            chunks: RefCell::new(DList::new()),
            index: RefCell::new(HashMap::with_hasher(XXHasher::new())),
        }
    }

    pub fn intern<'s>(&'s self, s: &str) -> &'s str {
        if s == "" { return ""; }

        let search_string = InternedString::from_str(s);

        let interned_str = match self.index.borrow_mut().entry(search_string) {
            Occupied(entry) => *entry.get(),
            Vacant(entry) => *entry.insert(self.do_intern(s)),
        };

        // The lifetime is really matched to us
        unsafe { mem::transmute(interned_str.as_slice()) }
    }

    fn do_intern(&self, s: &str) -> InternedString {
        self.ensure_capacity(s.len());
        self.store(s)
    }

    fn ensure_capacity(&self, str_len: uint) {
        let remaining = self.end.get() as uint - self.start.get() as uint;
        if remaining < str_len {
            self.allocate_another(max(CAPACITY, str_len))
        }
    }

    fn allocate_another(&self, capacity: uint) {
        let chunk = Chunk::new(capacity);
        self.start.set(chunk.start() as *mut u8);
        self.end.set(chunk.end());
        self.chunks.borrow_mut().push_front(chunk);
    }

    fn store(&self, s: &str) -> InternedString {
        let str_start = s.as_bytes().as_ptr();
        let str_len = s.len();

        unsafe {
            ptr::copy_nonoverlapping_memory(self.start.get(), str_start, str_len);

            // Rebuild the string from our buffer
            let interned_str = InternedString::from_parts(self.start.get() as *const u8, str_len);

            // Increase current pointer
            self.start.set(self.start.get().offset(str_len as int));

            interned_str
        }
    }
}

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

    #[test]
    fn keeps_the_same_string() {
        let s = StringPool::new();

        let interned = s.intern("hello");

        assert_eq!(interned, "hello");
    }

    #[test]
    fn does_not_reuse_the_pointer_of_the_input() {
        let s = StringPool::new();
        let input = "hello";

        let interned = s.intern(input);

        assert!(
            input.as_bytes().as_ptr() !=
            interned.as_bytes().as_ptr()
        );
    }

    #[test]
    fn reuses_the_pointer_for_repeated_input() {
        let s = StringPool::new();

        let interned1 = s.intern("world");
        let interned2 = s.intern("world");

        assert_eq!(
            interned1.as_bytes().as_ptr(),
            interned2.as_bytes().as_ptr()
        );
    }

    // #[test]
    // #[compile_failure]
    // fn string_cannot_outlive_the_pool() {
    //     let z = {
    //         let s = StringPool::new();
    //         s.intern("hello")
    //     };
    // }

    #[test]
    fn ignores_the_lifetime_of_the_input_string() {
        let s = StringPool::new();

        let interned = {
            let allocated_string = String::from_str("green");
            s.intern(allocated_string.as_slice())
        };

        // allocated_string is gone now, but we should be able to
        // access the result value until the storage goes away.

        assert_eq!(interned, "green");
    }

    #[test]
    fn can_be_dropped_immediately() {
        StringPool::new();
    }

    fn return_populated_storage() -> (StringPool, *const u8) {
        let s = StringPool::new();
        let ptr = {
            let interned = s.intern("hello");
            interned.as_ptr()
        };
        (s, ptr)
    }

    #[test]
    fn can_return_storage_populated_with_values() {
        let (s, ptr_val) = return_populated_storage();
        let interned = s.intern("hello");
        assert_eq!(interned.as_ptr(), ptr_val);
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;

    use super::StringPool;

    #[bench]
    fn single_string(b: &mut Bencher) {
        let s = StringPool::new();
        b.iter(|| s.intern("hello"));
        b.bytes = "hello".len() as u64;
    }

    #[bench]
    fn many_unique_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0u, 1000)
            .map(|i| format!("str{}str", i))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss.as_slice());
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }

    #[bench]
    fn many_repeated_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0u, 1000)
            .map(|i| format!("str{}str", i % 100))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss.as_slice());
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }
}
<commit_msg>Hash requires two traits now<commit_after>/// Strongly borrowed from TypedArena
/// Differs in growth strategy (non-doubling)
/// And only for variable-sized strings

use std::borrow::BorrowFrom;
use std::cell::{Cell,RefCell};
use std::cmp::max;
use std::collections::DList;
use std::collections::hash_map::HashMap;
use std::collections::hash_map::Entry::{Occupied,Vacant};
use std::fmt;
use std::hash;
use std::mem;
use std::ptr;
use std::raw::Slice;
use std::rt::heap::{allocate, deallocate};
use std::str;
use xxhash::{XXHasher};

struct Chunk {
    start: *mut u8,
    capacity: uint,
}

impl Chunk {
    fn new(capacity: uint) -> Chunk {
        Chunk {
            start: unsafe { allocate(capacity, mem::min_align_of::<u8>()) },
            capacity: capacity,
        }
    }

    // Returns a pointer to the beginning of the allocated space.
    #[inline]
    fn start(&self) -> *const u8 {
        self.start as *const u8
    }

    // Returns a pointer to the end of the allocated space.
    #[inline]
    fn end(&self) -> *const u8 {
        unsafe { self.start().offset(self.capacity as int) }
    }
}

impl Drop for Chunk {
    fn drop(&mut self) {
        unsafe {
            deallocate(self.start, self.capacity, mem::min_align_of::<u8>());
        }
    }
}

#[derive(Copy)]
pub struct InternedString {
    slice: Slice<u8>,
}

impl InternedString {
    fn from_parts(data: *const u8, len: uint) -> InternedString {
        InternedString {
            slice: Slice {
                data: data,
                len: len,
            },
        }
    }

    pub fn from_str(s: &str) -> InternedString {
        let bytes: &[u8] = s.as_bytes();
        let slice: Slice<u8> = unsafe { mem::transmute(bytes) };
        InternedString {
            slice: slice,
        }
    }

    pub fn as_slice<'s>(&self) -> &'s str {
        unsafe {
            let bytes: &[u8] = mem::transmute(self.slice);
            str::from_utf8_unchecked(bytes)
        }
    }
}

impl fmt::Show for InternedString {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_slice().fmt(f)
    }
}

impl PartialEq for InternedString {
    fn eq(&self, other: &InternedString) -> bool {
        self.as_slice().eq(other.as_slice())
    }
}

impl Eq for InternedString {}

impl<S> hash::Hash<S> for InternedString
    where S: hash::Writer + hash::Hasher
{
    fn hash(&self, state: &mut S) {
        self.as_slice().hash(state)
    }
}

impl BorrowFrom<InternedString> for str {
    fn borrow_from(owned: &InternedString) -> &str {
        owned.as_slice()
    }
}

pub struct StringPool {
    start: Cell<*mut u8>,
    end: Cell<*const u8>,
    chunks: RefCell<DList<Chunk>>,
    index: RefCell<HashMap<InternedString, InternedString, XXHasher>>,
}

static CAPACITY: uint = 10240;

impl StringPool {
    pub fn new() -> StringPool {
        StringPool{
            start: Cell::new(ptr::null_mut()),
            end: Cell::new(ptr::null()),
            chunks: RefCell::new(DList::new()),
            index: RefCell::new(HashMap::with_hasher(XXHasher::new())),
        }
    }

    pub fn intern<'s>(&'s self, s: &str) -> &'s str {
        if s == "" { return ""; }

        let search_string = InternedString::from_str(s);

        let interned_str = match self.index.borrow_mut().entry(search_string) {
            Occupied(entry) => *entry.get(),
            Vacant(entry) => *entry.insert(self.do_intern(s)),
        };

        // The lifetime is really matched to us
        unsafe { mem::transmute(interned_str.as_slice()) }
    }

    fn do_intern(&self, s: &str) -> InternedString {
        self.ensure_capacity(s.len());
        self.store(s)
    }

    fn ensure_capacity(&self, str_len: uint) {
        let remaining = self.end.get() as uint - self.start.get() as uint;
        if remaining < str_len {
            self.allocate_another(max(CAPACITY, str_len))
        }
    }

    fn allocate_another(&self, capacity: uint) {
        let chunk = Chunk::new(capacity);
        self.start.set(chunk.start() as *mut u8);
        self.end.set(chunk.end());
        self.chunks.borrow_mut().push_front(chunk);
    }

    fn store(&self, s: &str) -> InternedString {
        let str_start = s.as_bytes().as_ptr();
        let str_len = s.len();

        unsafe {
            ptr::copy_nonoverlapping_memory(self.start.get(), str_start, str_len);

            // Rebuild the string from our buffer
            let interned_str = InternedString::from_parts(self.start.get() as *const u8, str_len);

            // Increase current pointer
            self.start.set(self.start.get().offset(str_len as int));

            interned_str
        }
    }
}

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

    #[test]
    fn keeps_the_same_string() {
        let s = StringPool::new();

        let interned = s.intern("hello");

        assert_eq!(interned, "hello");
    }

    #[test]
    fn does_not_reuse_the_pointer_of_the_input() {
        let s = StringPool::new();
        let input = "hello";

        let interned = s.intern(input);

        assert!(
            input.as_bytes().as_ptr() !=
            interned.as_bytes().as_ptr()
        );
    }

    #[test]
    fn reuses_the_pointer_for_repeated_input() {
        let s = StringPool::new();

        let interned1 = s.intern("world");
        let interned2 = s.intern("world");

        assert_eq!(
            interned1.as_bytes().as_ptr(),
            interned2.as_bytes().as_ptr()
        );
    }

    // #[test]
    // #[compile_failure]
    // fn string_cannot_outlive_the_pool() {
    //     let z = {
    //         let s = StringPool::new();
    //         s.intern("hello")
    //     };
    // }

    #[test]
    fn ignores_the_lifetime_of_the_input_string() {
        let s = StringPool::new();

        let interned = {
            let allocated_string = String::from_str("green");
            s.intern(allocated_string.as_slice())
        };

        // allocated_string is gone now, but we should be able to
        // access the result value until the storage goes away.

        assert_eq!(interned, "green");
    }

    #[test]
    fn can_be_dropped_immediately() {
        StringPool::new();
    }

    fn return_populated_storage() -> (StringPool, *const u8) {
        let s = StringPool::new();
        let ptr = {
            let interned = s.intern("hello");
            interned.as_ptr()
        };
        (s, ptr)
    }

    #[test]
    fn can_return_storage_populated_with_values() {
        let (s, ptr_val) = return_populated_storage();
        let interned = s.intern("hello");
        assert_eq!(interned.as_ptr(), ptr_val);
    }
}

#[cfg(test)]
mod bench {
    use test::Bencher;

    use super::StringPool;

    #[bench]
    fn single_string(b: &mut Bencher) {
        let s = StringPool::new();
        b.iter(|| s.intern("hello"));
        b.bytes = "hello".len() as u64;
    }

    #[bench]
    fn many_unique_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0u, 1000)
            .map(|i| format!("str{}str", i))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss.as_slice());
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }

    #[bench]
    fn many_repeated_string(b: &mut Bencher) {
        let s = StringPool::new();

        let strings: Vec<String> =
            range(0u, 1000)
            .map(|i| format!("str{}str", i % 100))
            .collect();
        b.iter(|| {
            for ss in strings.iter() {
                s.intern(ss.as_slice());
            }
        });
        b.bytes = strings.iter().fold(0, |a, s| a + s.len()) as u64;
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::{Read, Seek};
use byteorder::{ReadBytesExt, LittleEndian};

use image::ImageError;
use image::ImageResult;
use image::ImageDecoder;
use image::DecodingResult;
use color::ColorType;

enum ImageType {
    NoImageData = 0,
    /// Uncompressed images
    RawColorMap = 1,
    RawTrueColor = 2,
    RawGrayScale = 3,
    /// Run length encoded images
    RunColorMap = 9,
    RunTrueColor = 10,
    RunGrayScale = 11,
    Unknown,
}

impl ImageType {
    /// Create a new image type from a u8
    fn new(img_type: u8) -> ImageType {
        match img_type {
            0  => ImageType::NoImageData,

            1  => ImageType::RawColorMap,
            2  => ImageType::RawTrueColor,
            3  => ImageType::RawGrayScale,

            9  => ImageType::RunColorMap,
            10 => ImageType::RunTrueColor,
            11 => ImageType::RunGrayScale,

            _  => ImageType::Unknown,
        }
    }

    /// Check if the image format uses colors as opposed to gray scale
    fn is_color(&self) -> bool {
        match *self {
            ImageType::RawColorMap  |
            ImageType::RawTrueColor |
            ImageType::RunTrueColor |
            ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Does the image use a color map
    fn is_color_mapped(&self) -> bool {
        match *self {
            ImageType::RawColorMap |
            ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Is the image run length encoded
    fn is_encoded(&self) -> bool {
        match *self {
            ImageType::RunColorMap |
            ImageType::RunTrueColor |
            ImageType::RunGrayScale => true,
            _ => false,
        }
    }
}

/// Header used by TGA image files
#[derive(Debug)]
struct Header {
    id_length: u8,         // length of ID string
    map_type: u8,          // color map type
    image_type: u8,        // image type code
    map_origin: u16,       // starting index of map
    map_length: u16,      // length of map
    map_entry_size: u8,    // size of map entries in bits
    x_origin: u16,         // x-origin of image
    y_origin: u16,         // y-origin of image
    image_width: u16,      // width of image
    image_height: u16,     // height of image
    pixel_depth: u8,       // bits per pixel
    image_desc: u8,        // image descriptor
}

impl Header {
    /// Create a header with all valuse set to zero
    fn new() -> Header {
        Header {
            id_length: 0,
            map_type: 0,
            image_type: 0,
            map_origin: 0,
            map_length: 0,
            map_entry_size: 0,
            x_origin: 0,
            y_origin: 0,
            image_width: 0,
            image_height: 0,
            pixel_depth: 0,
            image_desc: 0,
        }
    }

    /// Load the header with values from the reader
    fn from_reader(r: &mut Read) -> ImageResult<Header> {
        Ok(Header {
            id_length:         try!(r.read_u8()),
            map_type:          try!(r.read_u8()),
            image_type:        try!(r.read_u8()),
            map_origin:        try!(r.read_u16::<LittleEndian>()),
            map_length:        try!(r.read_u16::<LittleEndian>()),
            map_entry_size:    try!(r.read_u8()),
            x_origin:          try!(r.read_u16::<LittleEndian>()),
            y_origin:          try!(r.read_u16::<LittleEndian>()),
            image_width:       try!(r.read_u16::<LittleEndian>()),
            image_height:      try!(r.read_u16::<LittleEndian>()),
            pixel_depth:       try!(r.read_u8()),
            image_desc:        try!(r.read_u8()),
        })
    }
}

struct ColorMap {
    /// sizes in bytes
    start_offset: usize,
    entry_size: usize,
    bytes: Vec<u8>,
}

impl ColorMap {
    pub fn from_reader(r: &mut Read,
                       start_offset: u16,
                       num_entries: u16,
                       bits_per_entry: u8)
        -> ImageResult<ColorMap> {
            let bytes_per_entry = (bits_per_entry as usize + 7) / 8;

            let mut bytes = vec![0; bytes_per_entry * num_entries as usize];
            try!(r.read_exact(&mut bytes));

            Ok(ColorMap {
                entry_size: bytes_per_entry,
                start_offset: start_offset as usize,
                bytes: bytes,
            })
        }

    /// Get one entry from the color map
    pub fn get(&self, index: usize) -> &[u8] {
        let entry = self.start_offset + self.entry_size * index;
        &self.bytes[entry..entry + self.entry_size]
    }
}

/// The representation of a TGA decoder
pub struct TGADecoder<R> {
    r: R,

    width: usize,
    height: usize,
    bytes_per_pixel: usize,
    has_loaded_metadata: bool,

    image_type: ImageType,
    color_type: ColorType,

    header: Header,
    color_map: Option<ColorMap>,
}

impl<R: Read + Seek> TGADecoder<R> {
    /// Create a new decoder that decodes from the stream `r`
    pub fn new(r: R) -> TGADecoder<R> {
        TGADecoder {
            r: r,

            width: 0,
            height: 0,
            bytes_per_pixel: 0,
            has_loaded_metadata: false,

            image_type: ImageType::Unknown,
            color_type: ColorType::Gray(1),

            header: Header::new(),
            color_map: None,
        }
    }

    fn read_header(&mut self) -> ImageResult<()> {
        self.header = try!(Header::from_reader(&mut self.r));
        self.image_type = ImageType::new(self.header.image_type);
        self.width = self.header.image_width as usize;
        self.height = self.header.image_height as usize;
        self.bytes_per_pixel = (self.header.pixel_depth as usize + 7) / 8;
        Ok(())
    }

    fn read_metadata(&mut self) -> ImageResult<()> {
        if !self.has_loaded_metadata {
            try!(self.read_header());
            try!(self.read_image_id());
            try!(self.read_color_map());
            try!(self.read_color_information());
            self.has_loaded_metadata = true;
        }
        Ok(())
    }

    /// Loads the color information for the decoder
    ///
    /// To keep things simple, we won't handle bit depths that aren't divisible
    /// by 8 and are less than 32.
    fn read_color_information(&mut self) -> ImageResult<()> {
        if self.header.pixel_depth % 8 != 0 {
            return Err(ImageError::UnsupportedError("\
                Bit depth must be divisible by 8".to_string()));
        }
        if self.header.pixel_depth > 32 {
            return Err(ImageError::UnsupportedError("\
                Bit depth must be less than 32".to_string()));
        }

        let num_alpha_bits = self.header.image_desc & 0b1111;

        let other_channel_bits = if self.header.map_type != 0 {
            self.header.map_entry_size
        } else {
            self.header.pixel_depth - num_alpha_bits
        };
        let color = self.image_type.is_color();

        match (num_alpha_bits, other_channel_bits, color) {
            // really, the encoding is BGR and BGRA, this is fixed
            // up with `TGADecoder::reverse_encoding`.
            (8, 24, true) => self.color_type = ColorType::RGBA(8),
            (0, 24, true) => self.color_type = ColorType::RGB(8),
            (8, 8, false) => self.color_type = ColorType::GrayA(8),
            (0, 8, false) => self.color_type = ColorType::Gray(8),
            _ => return Err(ImageError::UnsupportedError(format!("\
                    Color format not supported. Bit depth: {}, Alpha bits: {}",
                    other_channel_bits, num_alpha_bits).to_string())),
        }
        Ok(())
    }

    /// Read the image id field
    ///
    /// We're not interested in this field, so this function skips it if it
    /// is present
    fn read_image_id(&mut self) -> ImageResult<()> {
        try!(self.r.seek(io::SeekFrom::Current(self.header.id_length as i64)));
        Ok(())
    }

    fn read_color_map(&mut self) -> ImageResult<()> {
        if self.header.map_type == 1 {
            self.color_map = Some(try!(
                ColorMap::from_reader(&mut self.r,
                                      self.header.map_origin,
                                      self.header.map_length,
                                      self.header.map_entry_size)));
        }
        Ok(())
    }

    /// Expands indices into its mapped color
    fn expand_color_map(&mut self, pixel_data: Vec<u8>) -> Vec<u8> {
        #[inline]
        fn bytes_to_index(bytes: &[u8]) -> usize {
            let mut result = 0usize;
            for byte in bytes.iter() {
                result = result << 8 | *byte as usize;
            }
            result
        }

        let bytes_per_entry = (self.header.map_entry_size as usize + 7) / 8;
        let mut result = Vec::with_capacity(self.width * self.height *
                                            bytes_per_entry);

        let color_map = match self.color_map {
            Some(ref color_map) => color_map,
            None => unreachable!(),
        };

        for chunk in pixel_data.chunks(self.bytes_per_pixel) {
            let index = bytes_to_index(chunk);
            result.extend(color_map.get(index).iter().map(|&c| c));
        }

        result
    }

    fn read_image_data(&mut self) -> ImageResult<Vec<u8>> {
        // read the pixels from the data region
        let mut pixel_data = if self.image_type.is_encoded() {
            try!(self.read_encoded_data())
        } else {
            let num_raw_bytes = self.width * self.height * self.bytes_per_pixel;
            let mut buf = vec![0; num_raw_bytes];
            try!(self.r.by_ref().read_exact(&mut buf));
            buf
        };

        // expand the indices using the color map if necessary
        if self.image_type.is_color_mapped() {
            pixel_data = self.expand_color_map(pixel_data)
        }

        self.reverse_encoding(&mut pixel_data);
        Ok(pixel_data)
    }

    /// Reads a run length encoded packet
    fn read_encoded_data(&mut self) -> ImageResult<Vec<u8>> {
        let num_bytes = self.width * self.height * self.bytes_per_pixel;
        let mut pixel_data = Vec::with_capacity(num_bytes);

        while pixel_data.len() < num_bytes {
            let run_packet = try!(self.r.read_u8());
            // If the highest bit in `run_packet` is set, then we repeat pixels
            //
            // Note: the TGA format adds 1 to both counts because having a count
            // of 0 would be pointless.
            if (run_packet & 0x80) != 0 {
                // high bit set, so we will repeat the data
                let repeat_count = ((run_packet & !0x80) + 1) as usize;
                let mut data = Vec::with_capacity(self.bytes_per_pixel);
                try!(self.r.by_ref().take(self.bytes_per_pixel as u64).read_to_end(&mut data));
                for _ in 0usize..repeat_count {
                    pixel_data.extend(data.iter().map(|&c| c));
                }
            } else {
                // not set, so `run_packet+1` is the number of non-encoded pixels
                let num_raw_bytes = (run_packet + 1) as usize * self.bytes_per_pixel;
                try!(self.r.by_ref().take(num_raw_bytes as u64).read_to_end(&mut pixel_data));
            }
        }

        Ok(pixel_data)
    }

    /// Reverse from BGR encoding to RGB encoding
    ///
    /// TGA files are stored in the BGRA encoding. This function swaps
    /// the blue and red bytes in the `pixels` array.
    fn reverse_encoding(&mut self, pixels: &mut [u8]) {
        // We only need to reverse the encoding of color images
        match self.color_type {
            ColorType::RGB(8) => {
                for chunk in pixels.chunks_mut(self.bytes_per_pixel) {
                    let r = chunk[0];
                    chunk[0] = chunk[2];
                    chunk[2] = r;
                }
            }
            ColorType::RGBA(8) => {
                for chunk in pixels.chunks_mut(self.bytes_per_pixel) {
                    let r = chunk[0];
                    chunk[0] = chunk[2];
                    chunk[2] = r;
                }
            }
            _ => { }
        }
    }
}

impl<R: Read + Seek> ImageDecoder for TGADecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        try!(self.read_metadata());

        Ok((self.width as u32, self.height as u32))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        try!(self.read_metadata());

        Ok(self.color_type)
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        try!(self.read_metadata());

        Ok(self.bytes_per_pixel * 8 * self.width)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        try!(self.read_metadata());
        self.read_image_data().map(|v| DecodingResult::U8(v) )
    }
}
<commit_msg>Adds functionality to TGA encoder to flip image based on the screen origin bit, so some TGA images are not upside down when loaded<commit_after>use std::io;
use std::io::{Read, Seek};
use byteorder::{ReadBytesExt, LittleEndian};

use image::ImageError;
use image::ImageResult;
use image::ImageDecoder;
use image::DecodingResult;
use color::ColorType;

enum ImageType {
    NoImageData = 0,
    /// Uncompressed images
    RawColorMap = 1,
    RawTrueColor = 2,
    RawGrayScale = 3,
    /// Run length encoded images
    RunColorMap = 9,
    RunTrueColor = 10,
    RunGrayScale = 11,
    Unknown,
}

impl ImageType {
    /// Create a new image type from a u8
    fn new(img_type: u8) -> ImageType {
        match img_type {
            0  => ImageType::NoImageData,

            1  => ImageType::RawColorMap,
            2  => ImageType::RawTrueColor,
            3  => ImageType::RawGrayScale,

            9  => ImageType::RunColorMap,
            10 => ImageType::RunTrueColor,
            11 => ImageType::RunGrayScale,

            _  => ImageType::Unknown,
        }
    }

    /// Check if the image format uses colors as opposed to gray scale
    fn is_color(&self) -> bool {
        match *self {
            ImageType::RawColorMap  |
            ImageType::RawTrueColor |
            ImageType::RunTrueColor |
            ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Does the image use a color map
    fn is_color_mapped(&self) -> bool {
        match *self {
            ImageType::RawColorMap |
            ImageType::RunColorMap => true,
            _ => false,
        }
    }

    /// Is the image run length encoded
    fn is_encoded(&self) -> bool {
        match *self {
            ImageType::RunColorMap |
            ImageType::RunTrueColor |
            ImageType::RunGrayScale => true,
            _ => false,
        }
    }
}

/// Header used by TGA image files
#[derive(Debug)]
struct Header {
    id_length: u8,         // length of ID string
    map_type: u8,          // color map type
    image_type: u8,        // image type code
    map_origin: u16,       // starting index of map
    map_length: u16,      // length of map
    map_entry_size: u8,    // size of map entries in bits
    x_origin: u16,         // x-origin of image
    y_origin: u16,         // y-origin of image
    image_width: u16,      // width of image
    image_height: u16,     // height of image
    pixel_depth: u8,       // bits per pixel
    image_desc: u8,        // image descriptor
}

impl Header {
    /// Create a header with all valuse set to zero
    fn new() -> Header {
        Header {
            id_length: 0,
            map_type: 0,
            image_type: 0,
            map_origin: 0,
            map_length: 0,
            map_entry_size: 0,
            x_origin: 0,
            y_origin: 0,
            image_width: 0,
            image_height: 0,
            pixel_depth: 0,
            image_desc: 0,
        }
    }

    /// Load the header with values from the reader
    fn from_reader(r: &mut Read) -> ImageResult<Header> {
        Ok(Header {
            id_length:         try!(r.read_u8()),
            map_type:          try!(r.read_u8()),
            image_type:        try!(r.read_u8()),
            map_origin:        try!(r.read_u16::<LittleEndian>()),
            map_length:        try!(r.read_u16::<LittleEndian>()),
            map_entry_size:    try!(r.read_u8()),
            x_origin:          try!(r.read_u16::<LittleEndian>()),
            y_origin:          try!(r.read_u16::<LittleEndian>()),
            image_width:       try!(r.read_u16::<LittleEndian>()),
            image_height:      try!(r.read_u16::<LittleEndian>()),
            pixel_depth:       try!(r.read_u8()),
            image_desc:        try!(r.read_u8()),
        })
    }
}

struct ColorMap {
    /// sizes in bytes
    start_offset: usize,
    entry_size: usize,
    bytes: Vec<u8>,
}

impl ColorMap {
    pub fn from_reader(r: &mut Read,
                       start_offset: u16,
                       num_entries: u16,
                       bits_per_entry: u8)
        -> ImageResult<ColorMap> {
            let bytes_per_entry = (bits_per_entry as usize + 7) / 8;

            let mut bytes = vec![0; bytes_per_entry * num_entries as usize];
            try!(r.read_exact(&mut bytes));

            Ok(ColorMap {
                entry_size: bytes_per_entry,
                start_offset: start_offset as usize,
                bytes: bytes,
            })
        }

    /// Get one entry from the color map
    pub fn get(&self, index: usize) -> &[u8] {
        let entry = self.start_offset + self.entry_size * index;
        &self.bytes[entry..entry + self.entry_size]
    }
}

/// The representation of a TGA decoder
pub struct TGADecoder<R> {
    r: R,

    width: usize,
    height: usize,
    bytes_per_pixel: usize,
    has_loaded_metadata: bool,

    image_type: ImageType,
    color_type: ColorType,

    header: Header,
    color_map: Option<ColorMap>,
}

impl<R: Read + Seek> TGADecoder<R> {
    /// Create a new decoder that decodes from the stream `r`
    pub fn new(r: R) -> TGADecoder<R> {
        TGADecoder {
            r: r,

            width: 0,
            height: 0,
            bytes_per_pixel: 0,
            has_loaded_metadata: false,

            image_type: ImageType::Unknown,
            color_type: ColorType::Gray(1),

            header: Header::new(),
            color_map: None,
        }
    }

    fn read_header(&mut self) -> ImageResult<()> {
        self.header = try!(Header::from_reader(&mut self.r));
        self.image_type = ImageType::new(self.header.image_type);
        self.width = self.header.image_width as usize;
        self.height = self.header.image_height as usize;
        self.bytes_per_pixel = (self.header.pixel_depth as usize + 7) / 8;
        Ok(())
    }

    fn read_metadata(&mut self) -> ImageResult<()> {
        if !self.has_loaded_metadata {
            try!(self.read_header());
            try!(self.read_image_id());
            try!(self.read_color_map());
            try!(self.read_color_information());
            self.has_loaded_metadata = true;
        }
        Ok(())
    }

    /// Loads the color information for the decoder
    ///
    /// To keep things simple, we won't handle bit depths that aren't divisible
    /// by 8 and are less than 32.
    fn read_color_information(&mut self) -> ImageResult<()> {
        if self.header.pixel_depth % 8 != 0 {
            return Err(ImageError::UnsupportedError("\
                Bit depth must be divisible by 8".to_string()));
        }
        if self.header.pixel_depth > 32 {
            return Err(ImageError::UnsupportedError("\
                Bit depth must be less than 32".to_string()));
        }

        let num_alpha_bits = self.header.image_desc & 0b1111;

        let other_channel_bits = if self.header.map_type != 0 {
            self.header.map_entry_size
        } else {
            self.header.pixel_depth - num_alpha_bits
        };
        let color = self.image_type.is_color();

        match (num_alpha_bits, other_channel_bits, color) {
            // really, the encoding is BGR and BGRA, this is fixed
            // up with `TGADecoder::reverse_encoding`.
            (8, 24, true) => self.color_type = ColorType::RGBA(8),
            (0, 24, true) => self.color_type = ColorType::RGB(8),
            (8, 8, false) => self.color_type = ColorType::GrayA(8),
            (0, 8, false) => self.color_type = ColorType::Gray(8),
            _ => return Err(ImageError::UnsupportedError(format!("\
                    Color format not supported. Bit depth: {}, Alpha bits: {}",
                    other_channel_bits, num_alpha_bits).to_string())),
        }
        Ok(())
    }

    /// Read the image id field
    ///
    /// We're not interested in this field, so this function skips it if it
    /// is present
    fn read_image_id(&mut self) -> ImageResult<()> {
        try!(self.r.seek(io::SeekFrom::Current(self.header.id_length as i64)));
        Ok(())
    }

    fn read_color_map(&mut self) -> ImageResult<()> {
        if self.header.map_type == 1 {
            self.color_map = Some(try!(
                ColorMap::from_reader(&mut self.r,
                                      self.header.map_origin,
                                      self.header.map_length,
                                      self.header.map_entry_size)));
        }
        Ok(())
    }

    /// Expands indices into its mapped color
    fn expand_color_map(&mut self, pixel_data: Vec<u8>) -> Vec<u8> {
        #[inline]
        fn bytes_to_index(bytes: &[u8]) -> usize {
            let mut result = 0usize;
            for byte in bytes.iter() {
                result = result << 8 | *byte as usize;
            }
            result
        }

        let bytes_per_entry = (self.header.map_entry_size as usize + 7) / 8;
        let mut result = Vec::with_capacity(self.width * self.height *
                                            bytes_per_entry);

        let color_map = match self.color_map {
            Some(ref color_map) => color_map,
            None => unreachable!(),
        };

        for chunk in pixel_data.chunks(self.bytes_per_pixel) {
            let index = bytes_to_index(chunk);
            result.extend(color_map.get(index).iter().map(|&c| c));
        }

        result
    }

    fn read_image_data(&mut self) -> ImageResult<Vec<u8>> {
        // read the pixels from the data region
        let mut pixel_data = if self.image_type.is_encoded() {
            try!(self.read_encoded_data())
        } else {
            let num_raw_bytes = self.width * self.height * self.bytes_per_pixel;
            let mut buf = vec![0; num_raw_bytes];
            try!(self.r.by_ref().read_exact(&mut buf));
            buf
        };

        // expand the indices using the color map if necessary
        if self.image_type.is_color_mapped() {
            pixel_data = self.expand_color_map(pixel_data)
        }

        self.reverse_encoding(&mut pixel_data);
        Ok(pixel_data)
    }

    /// Reads a run length encoded packet
    fn read_encoded_data(&mut self) -> ImageResult<Vec<u8>> {
        let num_bytes = self.width * self.height * self.bytes_per_pixel;
        let mut pixel_data = Vec::with_capacity(num_bytes);

        while pixel_data.len() < num_bytes {
            let run_packet = try!(self.r.read_u8());
            // If the highest bit in `run_packet` is set, then we repeat pixels
            //
            // Note: the TGA format adds 1 to both counts because having a count
            // of 0 would be pointless.
            if (run_packet & 0x80) != 0 {
                // high bit set, so we will repeat the data
                let repeat_count = ((run_packet & !0x80) + 1) as usize;
                let mut data = Vec::with_capacity(self.bytes_per_pixel);
                try!(self.r.by_ref().take(self.bytes_per_pixel as u64).read_to_end(&mut data));
                for _ in 0usize..repeat_count {
                    pixel_data.extend(data.iter().map(|&c| c));
                }
            } else {
                // not set, so `run_packet+1` is the number of non-encoded pixels
                let num_raw_bytes = (run_packet + 1) as usize * self.bytes_per_pixel;
                try!(self.r.by_ref().take(num_raw_bytes as u64).read_to_end(&mut pixel_data));
            }
        }

        Ok(pixel_data)
    }

    /// Reverse from BGR encoding to RGB encoding
    ///
    /// TGA files are stored in the BGRA encoding. This function swaps
    /// the blue and red bytes in the `pixels` array.
    fn reverse_encoding(&mut self, pixels: &mut [u8]) {
        // We only need to reverse the encoding of color images
        match self.color_type {
            ColorType::RGB(8) => {
                for chunk in pixels.chunks_mut(self.bytes_per_pixel) {
                    let r = chunk[0];
                    chunk[0] = chunk[2];
                    chunk[2] = r;
                }
            }
            ColorType::RGBA(8) => {
                for chunk in pixels.chunks_mut(self.bytes_per_pixel) {
                    let r = chunk[0];
                    chunk[0] = chunk[2];
                    chunk[2] = r;
                }
            }
            _ => { }
        }
    }

    /// Flip the image vertically depending on the screen origin bit
    ///
    /// The bit in position 5 of the image descriptor byte is the screen origin bit.
    /// If it's 1, the origin is in the top left corner.
    /// If it's 0, the origin is in the bottom left corner.
    /// This function checks the bit, and if it's 0, flips the image vertically.
    fn flip_vertically(&mut self, pixels: &mut [u8]) {
        let screen_origin_bit = 0b100000 & self.header.image_desc != 0;

        if !screen_origin_bit {
            let num_bytes = pixels.len();

            // Make a copy of the current pixels that we can take values from when flipping.
            let mut pixels_copy = Vec::with_capacity(num_bytes);
            for i in 0..pixels.len() {
                pixels_copy.push(pixels[i]);
            }

            let width_bytes = num_bytes / self.height;

            // Flip the image vertically.
            for vertical_index in 0..self.height {
                let vertical_target = (self.height - vertical_index) * width_bytes - width_bytes;

                for horizontal_index in 0..width_bytes {
                    let source = vertical_index * width_bytes + horizontal_index;
                    let target = vertical_target + horizontal_index;

                    pixels[target] = pixels_copy[source];
                }
            }
        }
    }
}

impl<R: Read + Seek> ImageDecoder for TGADecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        try!(self.read_metadata());

        Ok((self.width as u32, self.height as u32))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        try!(self.read_metadata());

        Ok(self.color_type)
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        try!(self.read_metadata());

        Ok(self.bytes_per_pixel * 8 * self.width)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        try!(self.read_metadata());
        self.read_image_data().map(|v| DecodingResult::U8(v) )
    }
}
<|endoftext|>


<commit_before>
/// Implemented by contexts that can transform.
pub trait Transform2d<'a> {
    /// Translate x and y.
    fn trans(&'a self, x: f64, y: f64) -> Self;

    /// Translate x an y in local coordinates.
    fn trans_local(&'a self, x: f64, y: f64) -> Self;

    /// Rotates degrees.
    #[inline(always)]
    fn rot_deg(&'a self, angle: f64) -> Self {
        let pi: f64 = Float::pi();
        self.rot_rad(angle * pi / 180.0)
    }

    /// Rotates degrees.
    #[inline(always)]
    fn rot_deg_local(&'a self, angle: f64) -> Self {
        let pi: f64 = Float::pi();
        self.rot_rad_local(angle * pi / 180.0)
    }

    /// Rotate radians.
    fn rot_rad(&'a self, angle: f64) -> Self;

    /// Rotate radians in local coordinates.
    fn rot_rad_local(&'a self, angle: f64) -> Self;

    /// Scale.
    fn scale(&'a self, sx: f64, sy: f64) -> Self;

    /// Scales in local coordinates.
    fn scale_local(&'a self, sx: f64, sy: f64) -> Self;

    /// Scales in both directions.
    #[inline(always)]
    fn zoom(&'a self, s: f64) -> Self {
        self.scale(s, s)
    }

    /// Scales in both directions.
    #[inline(always)]
    fn zoom_local(&'a self, s: f64) -> Self {
        self.scale_local(s, s)
    }

    /// Flips vertically.
    #[inline(always)]
    fn flip_v(&'a self) -> Self {
        self.scale(1.0, -1.0)
    }

    /// Flips vertically in local coordinates.
    #[inline(always)]
    fn flip_v_local(&'a self) -> Self {
        self.scale_local(1.0, -1.0)
    }

    /// Flips horizontally.
    #[inline(always)]
    fn flip_h(&'a self) -> Self {
        self.scale(-1.0, 0.0)
    }

    /// Shear.
    fn shear(&'a self, sx: f64, sy: f64) -> Self;

    /// Shears in local coordinates.
    fn shear_local(&'a self, sx: f64, sy: f64) -> Self;
}
<commit_msg>Added `flip_h_local` in `Transform2d`<commit_after>
/// Implemented by contexts that can transform.
pub trait Transform2d<'a> {
    /// Translate x and y.
    fn trans(&'a self, x: f64, y: f64) -> Self;

    /// Translate x an y in local coordinates.
    fn trans_local(&'a self, x: f64, y: f64) -> Self;

    /// Rotates degrees.
    #[inline(always)]
    fn rot_deg(&'a self, angle: f64) -> Self {
        let pi: f64 = Float::pi();
        self.rot_rad(angle * pi / 180.0)
    }

    /// Rotates degrees.
    #[inline(always)]
    fn rot_deg_local(&'a self, angle: f64) -> Self {
        let pi: f64 = Float::pi();
        self.rot_rad_local(angle * pi / 180.0)
    }

    /// Rotate radians.
    fn rot_rad(&'a self, angle: f64) -> Self;

    /// Rotate radians in local coordinates.
    fn rot_rad_local(&'a self, angle: f64) -> Self;

    /// Scale.
    fn scale(&'a self, sx: f64, sy: f64) -> Self;

    /// Scales in local coordinates.
    fn scale_local(&'a self, sx: f64, sy: f64) -> Self;

    /// Scales in both directions.
    #[inline(always)]
    fn zoom(&'a self, s: f64) -> Self {
        self.scale(s, s)
    }

    /// Scales in both directions.
    #[inline(always)]
    fn zoom_local(&'a self, s: f64) -> Self {
        self.scale_local(s, s)
    }

    /// Flips vertically.
    #[inline(always)]
    fn flip_v(&'a self) -> Self {
        self.scale(1.0, -1.0)
    }

    /// Flips vertically in local coordinates.
    #[inline(always)]
    fn flip_v_local(&'a self) -> Self {
        self.scale_local(1.0, -1.0)
    }

    /// Flips horizontally.
    #[inline(always)]
    fn flip_h(&'a self) -> Self {
        self.scale(-1.0, 0.0)
    }

    /// Flips horizontally in local coordinates.
    #[inline(always)]
    fn flip_h_local(&'a self) -> Self {
        self.scale_local(-1.0, 0.0)
    }

    /// Shear.
    fn shear(&'a self, sx: f64, sy: f64) -> Self;

    /// Shears in local coordinates.
    fn shear_local(&'a self, sx: f64, sy: f64) -> Self;
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;

use ast::*;
use ast;
use interpreter::InterpreterError;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Type {
    Number,
    Bool,
    Any,
}

impl From<ast::Literal> for Type {
    fn from(from: ast::Literal) -> Self {
        match from {
            ast::Literal::Integer(_) => Type::Number,
            ast::Literal::Float(_) => Type::Number,
            ast::Literal::Bool(_) => Type::Bool,
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Type::Number => write!(f, "Number"),
            Type::Bool => write!(f, "Bool"),
            Type::Any => write!(f, "Any"),
        }
    }
}

#[derive(Clone)]
pub struct TypeEnvironment {
    pub symbol_tables: Vec<HashMap<String, Type>>,
}

impl TypeEnvironment {
    pub fn new() -> TypeEnvironment {
        TypeEnvironment { symbol_tables: Vec::new() }
    }

    pub fn start_scope(&mut self) {
        self.symbol_tables.push(HashMap::new());
    }

    pub fn end_scope(&mut self) {
        self.symbol_tables.pop();
    }

    pub fn declare(&mut self, variable: &Variable, typ: &Type) {
        match *variable {
            Variable::Identifier(_, ref id) => {
                self.symbol_tables.last_mut().unwrap().insert(id.clone(), *typ);
            }
        };
    }

    pub fn set(&mut self, identifier: &String, typ: Type) -> Result<(), InterpreterError> {
        for table in self.symbol_tables.iter_mut().rev() {
            // TODO: Entry API
            if table.contains_key(identifier) {
                table.insert(identifier.clone(), typ);
                return Ok(());
            }
        }
        Err(InterpreterError::UndeclaredAssignment(identifier.clone()))
    }

    pub fn get_type(&mut self, identifier: &String) -> Result<Type, InterpreterError> {
        for table in self.symbol_tables.iter().rev() {
            if let Some(typ) = table.get(identifier) {
                return Ok(*typ);
            }
        }
        Err(InterpreterError::ReferenceError(identifier.clone()))
    }

    pub fn get_all_keys(&self) -> HashSet<String> {
        let mut keys = HashSet::new();
        for table in self.symbol_tables.iter() {
            for key in table.keys() {
                keys.insert(key.clone());
            }
        }
        keys
    }
}

pub fn check_program(ast: &Vec<Statement>) -> Result<(), Vec<InterpreterError>> {
    let mut env = TypeEnvironment::new();
    env.start_scope();
    let result = check_statements(ast, &mut env);
    env.end_scope();
    result
}

pub fn check_statements(ast: &Vec<Statement>,
                        env: &mut TypeEnvironment)
                        -> Result<(), Vec<InterpreterError>> {
    let mut errors = Vec::new();
    for statement in ast.iter() {
        if let Err(mut e) = check_statement(statement, env) {
            errors.append(&mut e);
        }
    }
    if errors.len() == 0 {
        Ok(())
    } else {
        Err(errors)
    }
}

pub fn check_statement(s: &Statement,
                       env: &mut TypeEnvironment)
                       -> Result<(), Vec<InterpreterError>> {
    let mut errors = Vec::new();
    match *s {
        Statement::VariableDeclaration(ref variable, ref expr) => {
            let checked_type = match check_expr(expr, env) {
                Ok(t) => t,
                Err(mut e) => {
                    errors.append(&mut e);
                    Type::Any
                }
            };
            env.declare(variable, &checked_type);
            // println!("{:?} => {}", variable, checked_type);
        }
        Statement::Assignment(ref lhs_expr, ref expr) => {
            let checked_type = match check_expr(expr, env) {
                Ok(t) => t,
                Err(mut e) => {
                    errors.append(&mut e);
                    Type::Any
                }
            };
            match *lhs_expr {
                LhsExpr::Identifier(ref id) => {
                    if let Err(e) = env.set(id, checked_type) {
                        errors.push(e);
                    }
                }
            };
        }
        Statement::Block(ref statements) => {
            env.start_scope();
            if let Err(mut e) = check_statements(statements, env) {
                errors.append(&mut e);
            }
            env.end_scope();
        }
        Statement::Expression(ref expr) => {
            if let Err(mut e) = check_expr(expr, env) {
                errors.append(&mut e);
            }
        }
        Statement::IfThen(ref if_expr, ref then_block) => {
            if let Err(mut e) = check_expr(if_expr, env) {
                errors.append(&mut e);
            }
            if let Err(mut e) = check_statement(then_block, env) {
                errors.append(&mut e);
            }
        }
        Statement::IfThenElse(ref if_expr, ref then_block, ref else_block) => {
            let mut then_env = env.clone();
            let mut else_env = env.clone();
            if let Err(mut e) = check_expr(if_expr, env) {
                errors.append(&mut e);
            }
            if let Err(mut e) = check_statement(then_block, &mut then_env) {
                errors.append(&mut e);
            }
            if let Err(mut e) = check_statement(else_block, &mut else_env) {
                errors.append(&mut e);
            }

            for name in then_env.get_all_keys() {
                let then_type = then_env.get_type(&name).unwrap();
                if else_env.get_type(&name).unwrap() != then_type {
                    env.set(&name, Type::Any).unwrap();
                } else {
                    env.set(&name, then_type).unwrap();
                }
            }
        }
        Statement::Loop(ref block) => {
            if let Err(mut e) = check_statement(block, env) {
                errors.append(&mut e);
            }
        }
        Statement::Break => {},
        Statement::Empty => {}
    };
    if errors.len() == 0 {
        Ok(())
    } else {
        Err(errors)
    }
}

fn check_expr(expr: &Expr, env: &mut TypeEnvironment) -> Result<Type, Vec<InterpreterError>> {
    match *expr {
        Expr::Literal(ref x) => Ok(Type::from(x.clone())),
        Expr::Identifier(ref id) => {
            match env.get_type(&id) {
                Ok(t) => Ok(t),
                Err(e) => Err(vec![e]),
            }
        }
        Expr::UnaryExpression(ref op, ref expr) => {
            match check_expr(expr, env) {
                Ok(t) => {
                    match *op {
                        UnaryOp::Minus => {
                            match check_unary_minus_for_type(t) {
                                Ok(t) => Ok(t),
                                Err(e) => Err(vec![e]),
                            }
                        }
                    }
                }
                Err(e) => Err(e),
            }
        }
        Expr::UnaryLogicalExpression(ref op, ref expr) => {
            match check_expr(expr, env) {
                Ok(_) => {
                    match *op {
                        LogicalUnaryOp::Not => Ok(Type::Bool),
                    }
                }
                Err(e) => Err(e),
            }
        }
        Expr::BinaryExpression(ref expr1, ref op, ref expr2) => {
            let mut errors = Vec::new();
            let checked_type_1 = match check_expr(expr1, env) {
                Ok(t) => t,
                Err(mut e) => {
                    errors.append(&mut e);
                    Type::Any
                }
            };
            let checked_type_2 = match check_expr(expr2, env) {
                Ok(t) => t,
                Err(mut e) => {
                    errors.append(&mut e);
                    Type::Any
                }
            };
            use ast::BinaryOp::*;
            let result = match *op {
                ref op @ Add |
                ref op @ Sub |
                ref op @ Mul |
                ref op @ Div |
                ref op @ FloorDiv => {
                    check_binary_arithmetic_for_types(op.clone(), checked_type_1, checked_type_2)
                }
                ref op @ LessThan |
                ref op @ LessThanOrEqual |
                ref op @ GreaterThan |
                ref op @ GreaterThanOrEqual => {
                    check_binary_comparison_for_types(op.clone(), checked_type_1, checked_type_2)
                }
                StrictEquals => Ok(Type::Bool),
            };
            match result {
                Err(e) => {
                    errors.push(e);
                    Err(errors)
                }
                Ok(t) => {
                    if errors.len() == 0 {
                        Ok(t)
                    } else {
                        Err(errors)
                    }
                }
            }
        }
        Expr::BinaryLogicalExpression(ref expr1, ref op, ref expr2) => {
            let mut errors = Vec::new();
            match *op {
                LogicalBinaryOp::LogicalAnd |
                LogicalBinaryOp::LogicalOr => {
                    if let Err(mut e) = check_expr(expr1, env) {
                        errors.append(&mut e);
                    };
                    if let Err(mut e) = check_expr(expr2, env) {
                        errors.append(&mut e);
                    };
                    if errors.len() == 0 {
                        Ok(Type::Bool)
                    } else {
                        Err(errors)
                    }
                }
            }
        }
        Expr::FunctionCall(ref id, ref args) => {
            match id.as_ref() {
                "println" => {},
                _ => {
                    return Err(vec![InterpreterError::ReferenceError(id.clone())]);
                }
            };
            let mut errors = Vec::new();
            for arg in args.iter() {
                if let Err(mut e) = check_expr(arg, env) {
                    errors.append(&mut e);
                }
            }
            if errors.len() == 0 {
                Ok(Type::Any)
            } else {
                Err(errors)
            }
        }
    }
}

fn check_unary_minus_for_type(typ: Type) -> Result<Type, InterpreterError> {
    match typ {
        Type::Number => Ok(Type::Number),
        Type::Any => Ok(Type::Any),
        _ => Err(InterpreterError::UnaryTypeError(UnaryOp::Minus, typ)),
    }
}

fn check_binary_arithmetic_for_types(op: BinaryOp,
                                     t1: Type,
                                     t2: Type)
                                     -> Result<Type, InterpreterError> {
    match (t1, t2) {
        (Type::Number, Type::Number) => Ok(Type::Number),
        (Type::Any, _) => Ok(Type::Any),
        (_, Type::Any) => Ok(Type::Any),
        _ => Err(InterpreterError::BinaryTypeError(op, t1, t2)),
    }
}

fn check_binary_comparison_for_types(op: BinaryOp,
                                     t1: Type,
                                     t2: Type)
                                     -> Result<Type, InterpreterError> {
    match (t1, t2) {
        (Type::Number, Type::Number) => Ok(Type::Bool),
        (Type::Any, _) => Ok(Type::Any),
        (_, Type::Any) => Ok(Type::Any),
        _ => Err(InterpreterError::BinaryTypeError(op, t1, t2)),
    }
}
<commit_msg>Add MultipleTypesFromBranchWarning as a type checker issue<commit_after>use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;

use ast::*;
use ast;
use interpreter::InterpreterError;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Type {
    Number,
    Bool,
    Any,
}

impl From<ast::Literal> for Type {
    fn from(from: ast::Literal) -> Self {
        match from {
            ast::Literal::Integer(_) => Type::Number,
            ast::Literal::Float(_) => Type::Number,
            ast::Literal::Bool(_) => Type::Bool,
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Type::Number => write!(f, "Number"),
            Type::Bool => write!(f, "Bool"),
            Type::Any => write!(f, "Any"),
        }
    }
}

#[derive(Debug)]
pub enum TypeCheckerIssue {
    InterpreterError(InterpreterError),
    MultipleTypesFromBranchWarning,
}

impl From<InterpreterError> for TypeCheckerIssue {
    fn from(from: InterpreterError) -> Self {
        TypeCheckerIssue::InterpreterError(from)
    }
}

#[derive(Clone)]
pub struct TypeEnvironment {
    pub symbol_tables: Vec<HashMap<String, Type>>,
}

impl TypeEnvironment {
    pub fn new() -> TypeEnvironment {
        TypeEnvironment { symbol_tables: Vec::new() }
    }

    pub fn start_scope(&mut self) {
        self.symbol_tables.push(HashMap::new());
    }

    pub fn end_scope(&mut self) {
        self.symbol_tables.pop();
    }

    pub fn declare(&mut self, variable: &Variable, typ: &Type) {
        match *variable {
            Variable::Identifier(_, ref id) => {
                self.symbol_tables.last_mut().unwrap().insert(id.clone(), *typ);
            }
        };
    }

    pub fn set(&mut self, identifier: &String, typ: Type) -> Result<(), TypeCheckerIssue> {
        for table in self.symbol_tables.iter_mut().rev() {
            // TODO: Entry API
            if table.contains_key(identifier) {
                table.insert(identifier.clone(), typ);
                return Ok(());
            }
        }
        Err(InterpreterError::UndeclaredAssignment(identifier.clone()).into())
    }

    pub fn get_type(&mut self, identifier: &String) -> Result<Type, TypeCheckerIssue> {
        for table in self.symbol_tables.iter().rev() {
            if let Some(typ) = table.get(identifier) {
                return Ok(*typ);
            }
        }
        Err(InterpreterError::ReferenceError(identifier.clone()).into())
    }

    pub fn get_all_keys(&self) -> HashSet<String> {
        let mut keys = HashSet::new();
        for table in self.symbol_tables.iter() {
            for key in table.keys() {
                keys.insert(key.clone());
            }
        }
        keys
    }
}

pub fn check_program(ast: &Vec<Statement>) -> Result<(), Vec<TypeCheckerIssue>> {
    let mut env = TypeEnvironment::new();
    env.start_scope();
    let result = check_statements(ast, &mut env);
    env.end_scope();
    result
}

pub fn check_statements(ast: &Vec<Statement>,
                        env: &mut TypeEnvironment)
                        -> Result<(), Vec<TypeCheckerIssue>> {
    let mut issues = Vec::new();
    for statement in ast.iter() {
        if let Err(mut e) = check_statement(statement, env) {
            issues.append(&mut e);
        }
    }
    if issues.len() == 0 {
        Ok(())
    } else {
        Err(issues)
    }
}

pub fn check_statement(s: &Statement,
                       env: &mut TypeEnvironment)
                       -> Result<(), Vec<TypeCheckerIssue>> {
    let mut issues = Vec::new();
    match *s {
        Statement::VariableDeclaration(ref variable, ref expr) => {
            let checked_type = match check_expr(expr, env) {
                Ok(t) => t,
                Err(mut e) => {
                    issues.append(&mut e);
                    Type::Any
                }
            };
            env.declare(variable, &checked_type);
            // println!("{:?} => {}", variable, checked_type);
        }
        Statement::Assignment(ref lhs_expr, ref expr) => {
            let checked_type = match check_expr(expr, env) {
                Ok(t) => t,
                Err(mut e) => {
                    issues.append(&mut e);
                    Type::Any
                }
            };
            match *lhs_expr {
                LhsExpr::Identifier(ref id) => {
                    if let Err(e) = env.set(id, checked_type) {
                        issues.push(e);
                    }
                }
            };
        }
        Statement::Block(ref statements) => {
            env.start_scope();
            if let Err(mut e) = check_statements(statements, env) {
                issues.append(&mut e);
            }
            env.end_scope();
        }
        Statement::Expression(ref expr) => {
            if let Err(mut e) = check_expr(expr, env) {
                issues.append(&mut e);
            }
        }
        Statement::IfThen(ref if_expr, ref then_block) => {
            if let Err(mut e) = check_expr(if_expr, env) {
                issues.append(&mut e);
            }
            if let Err(mut e) = check_statement(then_block, env) {
                issues.append(&mut e);
            }
        }
        Statement::IfThenElse(ref if_expr, ref then_block, ref else_block) => {
            let mut then_env = env.clone();
            let mut else_env = env.clone();
            if let Err(mut e) = check_expr(if_expr, env) {
                issues.append(&mut e);
            }
            if let Err(mut e) = check_statement(then_block, &mut then_env) {
                issues.append(&mut e);
            }
            if let Err(mut e) = check_statement(else_block, &mut else_env) {
                issues.append(&mut e);
            }

            for name in then_env.get_all_keys() {
                let then_type = then_env.get_type(&name).unwrap();
                if else_env.get_type(&name).unwrap() != then_type {
                    issues.push(TypeCheckerIssue::MultipleTypesFromBranchWarning);
                    env.set(&name, Type::Any).unwrap();
                } else {
                    env.set(&name, then_type).unwrap();
                }
            }
        }
        Statement::Loop(ref block) => {
            if let Err(mut e) = check_statement(block, env) {
                issues.append(&mut e);
            }
        }
        Statement::Break => {},
        Statement::Empty => {}
    };
    if issues.len() == 0 {
        Ok(())
    } else {
        Err(issues)
    }
}

fn check_expr(expr: &Expr, env: &mut TypeEnvironment) -> Result<Type, Vec<TypeCheckerIssue>> {
    match *expr {
        Expr::Literal(ref x) => Ok(Type::from(x.clone())),
        Expr::Identifier(ref id) => {
            match env.get_type(&id) {
                Ok(t) => Ok(t),
                Err(e) => Err(vec![e]),
            }
        }
        Expr::UnaryExpression(ref op, ref expr) => {
            match check_expr(expr, env) {
                Ok(t) => {
                    match *op {
                        UnaryOp::Minus => {
                            match check_unary_minus_for_type(t) {
                                Ok(t) => Ok(t),
                                Err(e) => Err(vec![e]),
                            }
                        }
                    }
                }
                Err(e) => Err(e),
            }
        }
        Expr::UnaryLogicalExpression(ref op, ref expr) => {
            match check_expr(expr, env) {
                Ok(_) => {
                    match *op {
                        LogicalUnaryOp::Not => Ok(Type::Bool),
                    }
                }
                Err(e) => Err(e),
            }
        }
        Expr::BinaryExpression(ref expr1, ref op, ref expr2) => {
            let mut issues = Vec::new();
            let checked_type_1 = match check_expr(expr1, env) {
                Ok(t) => t,
                Err(mut e) => {
                    issues.append(&mut e);
                    Type::Any
                }
            };
            let checked_type_2 = match check_expr(expr2, env) {
                Ok(t) => t,
                Err(mut e) => {
                    issues.append(&mut e);
                    Type::Any
                }
            };
            use ast::BinaryOp::*;
            let result = match *op {
                ref op @ Add |
                ref op @ Sub |
                ref op @ Mul |
                ref op @ Div |
                ref op @ FloorDiv => {
                    check_binary_arithmetic_for_types(op.clone(), checked_type_1, checked_type_2)
                }
                ref op @ LessThan |
                ref op @ LessThanOrEqual |
                ref op @ GreaterThan |
                ref op @ GreaterThanOrEqual => {
                    check_binary_comparison_for_types(op.clone(), checked_type_1, checked_type_2)
                }
                StrictEquals => Ok(Type::Bool),
            };
            match result {
                Err(e) => {
                    issues.push(e);
                    Err(issues)
                }
                Ok(t) => {
                    if issues.len() == 0 {
                        Ok(t)
                    } else {
                        Err(issues)
                    }
                }
            }
        }
        Expr::BinaryLogicalExpression(ref expr1, ref op, ref expr2) => {
            let mut issues = Vec::new();
            match *op {
                LogicalBinaryOp::LogicalAnd |
                LogicalBinaryOp::LogicalOr => {
                    if let Err(mut e) = check_expr(expr1, env) {
                        issues.append(&mut e);
                    };
                    if let Err(mut e) = check_expr(expr2, env) {
                        issues.append(&mut e);
                    };
                    if issues.len() == 0 {
                        Ok(Type::Bool)
                    } else {
                        Err(issues)
                    }
                }
            }
        }
        Expr::FunctionCall(ref id, ref args) => {
            match id.as_ref() {
                "println" => {},
                _ => {
                    return Err(vec![InterpreterError::ReferenceError(id.clone()).into()]);
                }
            };
            let mut issues = Vec::new();
            for arg in args.iter() {
                if let Err(mut e) = check_expr(arg, env) {
                    issues.append(&mut e);
                }
            }
            if issues.len() == 0 {
                Ok(Type::Any)
            } else {
                Err(issues)
            }
        }
    }
}

fn check_unary_minus_for_type(typ: Type) -> Result<Type, TypeCheckerIssue> {
    match typ {
        Type::Number => Ok(Type::Number),
        Type::Any => Ok(Type::Any),
        _ => Err(InterpreterError::UnaryTypeError(UnaryOp::Minus, typ).into()),
    }
}

fn check_binary_arithmetic_for_types(op: BinaryOp,
                                     t1: Type,
                                     t2: Type)
                                     -> Result<Type, TypeCheckerIssue> {
    match (t1, t2) {
        (Type::Number, Type::Number) => Ok(Type::Number),
        (Type::Any, _) => Ok(Type::Any),
        (_, Type::Any) => Ok(Type::Any),
        _ => Err(InterpreterError::BinaryTypeError(op, t1, t2).into()),
    }
}

fn check_binary_comparison_for_types(op: BinaryOp,
                                     t1: Type,
                                     t2: Type)
                                     -> Result<Type, TypeCheckerIssue> {
    match (t1, t2) {
        (Type::Number, Type::Number) => Ok(Type::Bool),
        (Type::Any, _) => Ok(Type::Any),
        (_, Type::Any) => Ok(Type::Any),
        _ => Err(InterpreterError::BinaryTypeError(op, t1, t2).into()),
    }
}
<|endoftext|>