<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::{
        epoch_ending::backup::{EpochEndingBackupController, EpochEndingBackupOpt},
        state_snapshot::backup::{StateSnapshotBackupController, StateSnapshotBackupOpt},
        transaction::backup::{TransactionBackupController, TransactionBackupOpt},
    },
    metadata,
    metadata::cache::MetadataCacheOpt,
    storage::BackupStorage,
    utils::{backup_service_client::BackupServiceClient, GlobalBackupOpt},
};
use anyhow::{anyhow, ensure, Result};
use futures::{stream, Future, StreamExt};
use libra_types::transaction::Version;
use libradb::backup::backup_handler::DbState;
use std::{fmt::Debug, sync::Arc};
use structopt::StructOpt;
use tokio::{
    sync::watch,
    time::{interval, Duration},
};

#[derive(StructOpt)]
pub struct BackupCoordinatorOpt {
    #[structopt(flatten)]
    pub metadata_cache_opt: MetadataCacheOpt,
    // Assuming epoch doesn't bump frequently, we default to backing epoch ending info up as soon
    // as possible.
    #[structopt(long, default_value = "1")]
    pub epoch_ending_batch_size: usize,
    // We replay transactions on top of a state snapshot at about 2000 tps, having a state snapshot
    // every 2000 * 3600 * 2 = 14.4 Mil versions will guarantee that to achieve any state in
    // history, it won't take us more than two hours in transaction replaying. Defaulting to 10 Mil
    // here to make it less than two, and easier for eyes.
    #[structopt(long, default_value = "10000000")]
    pub state_snapshot_interval: usize,
    // Assuming the network runs at 100 tps, it's 100 * 3600 = 360k transactions per hour, we don't
    // want the backups to lag behind too much. Defaulting to 100k here in case the network is way
    // slower than expected.
    #[structopt(long, default_value = "100000")]
    pub transaction_batch_size: usize,
}

impl BackupCoordinatorOpt {
    fn validate(&self) -> Result<()> {
        ensure!(
            self.epoch_ending_batch_size > 0
                && self.state_snapshot_interval > 0
                && self.transaction_batch_size > 0,
            "Backup interval and batch sizes must be greater than 0"
        );
        Ok(())
    }
}

pub struct BackupCoordinator {
    client: Arc<BackupServiceClient>,
    storage: Arc<dyn BackupStorage>,
    global_opt: GlobalBackupOpt,
    metadata_cache_opt: MetadataCacheOpt,
    epoch_ending_batch_size: usize,
    state_snapshot_interval: usize,
    transaction_batch_size: usize,
}

impl BackupCoordinator {
    pub fn new(
        opt: BackupCoordinatorOpt,
        global_opt: GlobalBackupOpt,
        client: Arc<BackupServiceClient>,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        opt.validate().unwrap();
        Self {
            client,
            storage,
            global_opt,
            metadata_cache_opt: opt.metadata_cache_opt,
            epoch_ending_batch_size: opt.epoch_ending_batch_size,
            state_snapshot_interval: opt.state_snapshot_interval,
            transaction_batch_size: opt.transaction_batch_size,
        }
    }
    pub async fn run(&self) -> Result<()> {
        // Connect to both the local Libra node and the backup storage.
        let backup_state =
            metadata::cache::sync_and_load(&self.metadata_cache_opt, Arc::clone(&self.storage))
                .await?
                .get_storage_state();
        let db_state = self
            .client
            .get_db_state()
            .await?
            .ok_or_else(|| anyhow!("DB not bootstrapped."))?;
        let (tx, rx) = watch::channel(db_state);

        // Schedule work streams.
        let watch_db_state = interval(Duration::from_secs(1))
            .then(|_| self.try_refresh_db_state(&tx))
            .boxed_local();

        let backup_epoch_endings = self
            .backup_work_stream(
                backup_state.latest_epoch_ending_epoch,
                &rx,
                Self::backup_epoch_endings,
            )
            .boxed_local();
        let backup_state_snapshots = self
            .backup_work_stream(
                backup_state.latest_state_snapshot_version,
                &rx,
                Self::backup_state_snapshot,
            )
            .boxed_local();
        let backup_transactions = self
            .backup_work_stream(
                backup_state.latest_transaction_version,
                &rx,
                Self::backup_transactions,
            )
            .boxed_local();

        println!("Backup coordinator started.");
        let mut all_work = stream::select_all(vec![
            watch_db_state,
            backup_epoch_endings,
            backup_state_snapshots,
            backup_transactions,
        ]);

        println!("Backup coordinator started.");
        loop {
            all_work
                .next()
                .await
                .ok_or_else(|| anyhow!("Must be a bug: we never returned None."))?
        }
    }
}

impl BackupCoordinator {
    async fn try_refresh_db_state(&self, db_state_broadcast: &watch::Sender<DbState>) {
        match self.client.get_db_state().await {
            Ok(s) => db_state_broadcast
                .broadcast(s.expect("Db should have been bootstrapped."))
                .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
                .unwrap(),
            Err(e) => println!("Failed pulling DbState from local Libra node: {}", e),
        };
    }

    async fn backup_epoch_endings(
        &self,
        last_epoch_ending_epoch_in_backup: Option<u64>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_epoch_ending_epoch_in_backup,
            self.epoch_ending_batch_size,
        );

        // <= because `db_state.epoch` hasn't ended yet
        if db_state.epoch <= last {
            // wait for the next db_state update
            return Ok(last_epoch_ending_epoch_in_backup);
        }

        EpochEndingBackupController::new(
            EpochEndingBackupOpt {
                start_epoch: first,
                end_epoch: last + 1,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(last))
    }

    async fn backup_state_snapshot(
        &self,
        last_snapshot_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<Version>> {
        let next_snapshot_version = get_next_snapshot(
            last_snapshot_version_in_backup,
            self.state_snapshot_interval,
        );

        if db_state.committed_version < next_snapshot_version {
            // wait for the next db_state update
            return Ok(last_snapshot_version_in_backup);
        }

        StateSnapshotBackupController::new(
            StateSnapshotBackupOpt {
                version: next_snapshot_version,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(next_snapshot_version))
    }

    async fn backup_transactions(
        &self,
        last_transaction_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (start_version, end_version) = get_batch_range(
            last_transaction_version_in_backup,
            self.transaction_batch_size,
        );

        if db_state.committed_version < end_version {
            // wait for the next db_state update
            return Ok(last_transaction_version_in_backup);
        }

        TransactionBackupController::new(
            TransactionBackupOpt {
                start_version,
                num_transactions: (end_version + 1 - start_version) as usize,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(end_version))
    }

    fn backup_work_stream<'a, S, W, Fut>(
        &'a self,
        initial_state: S,
        db_state_rx: &'a watch::Receiver<DbState>,
        worker: W,
    ) -> impl StreamExt<Item = ()> + 'a
    where
        S: Copy + Debug + 'a,
        W: Worker<'a, S, Fut> + Copy + 'static,
        Fut: Future<Output = Result<S>> + 'a,
    {
        stream::unfold(
            (initial_state, db_state_rx.clone()),
            move |(s, mut rx)| async move {
                let db_state = rx
                    .recv()
                    .await
                    .ok_or_else(|| anyhow!("The broadcaster has been dropped."))
                    .unwrap();
                let next_state = worker(self, s, db_state).await.unwrap_or_else(|e| {
                    println!("backup failed: {}. Keep trying with state {:?}.", e, s);
                    s
                });
                Some(((), (next_state, rx)))
            },
        )
    }
}

trait Worker<'a, S, Fut: Future<Output = Result<S>> + 'a>:
    Fn(&'a BackupCoordinator, S, DbState) -> Fut
{
}

impl<'a, T, S, Fut> Worker<'a, S, Fut> for T
where
    T: Fn(&'a BackupCoordinator, S, DbState) -> Fut,
    Fut: Future<Output = Result<S>> + 'a,
{
}

fn get_batch_range(last_in_backup: Option<u64>, batch_size: usize) -> (u64, u64) {
    // say, 5 is already in backup, and we target batches of size 10, we will return (6, 4) in this
    // case, so 6, 7, 8, 9 will be in this batch, and next time the backup worker will pass in 9,
    // and we will return (10, 10)
    let start = last_in_backup.map_or(0, |n| n + 1);
    let next_batch_start = (start / batch_size as u64 + 1) * batch_size as u64;

    (start, next_batch_start - 1)
}

fn get_next_snapshot(last_in_backup: Option<u64>, interval: usize) -> u64 {
    match last_in_backup {
        Some(last) => (last / interval as u64 + 1) * interval as u64,
        None => 0,
    }
}
<commit_msg>[backup] fix documentation of get_batch_range, add test<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::{
        epoch_ending::backup::{EpochEndingBackupController, EpochEndingBackupOpt},
        state_snapshot::backup::{StateSnapshotBackupController, StateSnapshotBackupOpt},
        transaction::backup::{TransactionBackupController, TransactionBackupOpt},
    },
    metadata,
    metadata::cache::MetadataCacheOpt,
    storage::BackupStorage,
    utils::{backup_service_client::BackupServiceClient, GlobalBackupOpt},
};
use anyhow::{anyhow, ensure, Result};
use futures::{stream, Future, StreamExt};
use libra_types::transaction::Version;
use libradb::backup::backup_handler::DbState;
use std::{fmt::Debug, sync::Arc};
use structopt::StructOpt;
use tokio::{
    sync::watch,
    time::{interval, Duration},
};

#[derive(StructOpt)]
pub struct BackupCoordinatorOpt {
    #[structopt(flatten)]
    pub metadata_cache_opt: MetadataCacheOpt,
    // Assuming epoch doesn't bump frequently, we default to backing epoch ending info up as soon
    // as possible.
    #[structopt(long, default_value = "1")]
    pub epoch_ending_batch_size: usize,
    // We replay transactions on top of a state snapshot at about 2000 tps, having a state snapshot
    // every 2000 * 3600 * 2 = 14.4 Mil versions will guarantee that to achieve any state in
    // history, it won't take us more than two hours in transaction replaying. Defaulting to 10 Mil
    // here to make it less than two, and easier for eyes.
    #[structopt(long, default_value = "10000000")]
    pub state_snapshot_interval: usize,
    // Assuming the network runs at 100 tps, it's 100 * 3600 = 360k transactions per hour, we don't
    // want the backups to lag behind too much. Defaulting to 100k here in case the network is way
    // slower than expected.
    #[structopt(long, default_value = "100000")]
    pub transaction_batch_size: usize,
}

impl BackupCoordinatorOpt {
    fn validate(&self) -> Result<()> {
        ensure!(
            self.epoch_ending_batch_size > 0
                && self.state_snapshot_interval > 0
                && self.transaction_batch_size > 0,
            "Backup interval and batch sizes must be greater than 0"
        );
        Ok(())
    }
}

pub struct BackupCoordinator {
    client: Arc<BackupServiceClient>,
    storage: Arc<dyn BackupStorage>,
    global_opt: GlobalBackupOpt,
    metadata_cache_opt: MetadataCacheOpt,
    epoch_ending_batch_size: usize,
    state_snapshot_interval: usize,
    transaction_batch_size: usize,
}

impl BackupCoordinator {
    pub fn new(
        opt: BackupCoordinatorOpt,
        global_opt: GlobalBackupOpt,
        client: Arc<BackupServiceClient>,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        opt.validate().unwrap();
        Self {
            client,
            storage,
            global_opt,
            metadata_cache_opt: opt.metadata_cache_opt,
            epoch_ending_batch_size: opt.epoch_ending_batch_size,
            state_snapshot_interval: opt.state_snapshot_interval,
            transaction_batch_size: opt.transaction_batch_size,
        }
    }
    pub async fn run(&self) -> Result<()> {
        // Connect to both the local Libra node and the backup storage.
        let backup_state =
            metadata::cache::sync_and_load(&self.metadata_cache_opt, Arc::clone(&self.storage))
                .await?
                .get_storage_state();
        let db_state = self
            .client
            .get_db_state()
            .await?
            .ok_or_else(|| anyhow!("DB not bootstrapped."))?;
        let (tx, rx) = watch::channel(db_state);

        // Schedule work streams.
        let watch_db_state = interval(Duration::from_secs(1))
            .then(|_| self.try_refresh_db_state(&tx))
            .boxed_local();

        let backup_epoch_endings = self
            .backup_work_stream(
                backup_state.latest_epoch_ending_epoch,
                &rx,
                Self::backup_epoch_endings,
            )
            .boxed_local();
        let backup_state_snapshots = self
            .backup_work_stream(
                backup_state.latest_state_snapshot_version,
                &rx,
                Self::backup_state_snapshot,
            )
            .boxed_local();
        let backup_transactions = self
            .backup_work_stream(
                backup_state.latest_transaction_version,
                &rx,
                Self::backup_transactions,
            )
            .boxed_local();

        println!("Backup coordinator started.");
        let mut all_work = stream::select_all(vec![
            watch_db_state,
            backup_epoch_endings,
            backup_state_snapshots,
            backup_transactions,
        ]);

        println!("Backup coordinator started.");
        loop {
            all_work
                .next()
                .await
                .ok_or_else(|| anyhow!("Must be a bug: we never returned None."))?
        }
    }
}

impl BackupCoordinator {
    async fn try_refresh_db_state(&self, db_state_broadcast: &watch::Sender<DbState>) {
        match self.client.get_db_state().await {
            Ok(s) => db_state_broadcast
                .broadcast(s.expect("Db should have been bootstrapped."))
                .map_err(|e| anyhow!("Receivers should not be cancelled: {}", e))
                .unwrap(),
            Err(e) => println!("Failed pulling DbState from local Libra node: {}", e),
        };
    }

    async fn backup_epoch_endings(
        &self,
        last_epoch_ending_epoch_in_backup: Option<u64>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_epoch_ending_epoch_in_backup,
            self.epoch_ending_batch_size,
        );

        // <= because `db_state.epoch` hasn't ended yet
        if db_state.epoch <= last {
            // wait for the next db_state update
            return Ok(last_epoch_ending_epoch_in_backup);
        }

        EpochEndingBackupController::new(
            EpochEndingBackupOpt {
                start_epoch: first,
                end_epoch: last + 1,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(last))
    }

    async fn backup_state_snapshot(
        &self,
        last_snapshot_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<Version>> {
        let next_snapshot_version = get_next_snapshot(
            last_snapshot_version_in_backup,
            self.state_snapshot_interval,
        );

        if db_state.committed_version < next_snapshot_version {
            // wait for the next db_state update
            return Ok(last_snapshot_version_in_backup);
        }

        StateSnapshotBackupController::new(
            StateSnapshotBackupOpt {
                version: next_snapshot_version,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(next_snapshot_version))
    }

    async fn backup_transactions(
        &self,
        last_transaction_version_in_backup: Option<Version>,
        db_state: DbState,
    ) -> Result<Option<u64>> {
        let (first, last) = get_batch_range(
            last_transaction_version_in_backup,
            self.transaction_batch_size,
        );

        if db_state.committed_version < last {
            // wait for the next db_state update
            return Ok(last_transaction_version_in_backup);
        }

        TransactionBackupController::new(
            TransactionBackupOpt {
                start_version: first,
                num_transactions: (last + 1 - first) as usize,
            },
            self.global_opt.clone(),
            Arc::clone(&self.client),
            Arc::clone(&self.storage),
        )
        .run()
        .await?;

        Ok(Some(last))
    }

    fn backup_work_stream<'a, S, W, Fut>(
        &'a self,
        initial_state: S,
        db_state_rx: &'a watch::Receiver<DbState>,
        worker: W,
    ) -> impl StreamExt<Item = ()> + 'a
    where
        S: Copy + Debug + 'a,
        W: Worker<'a, S, Fut> + Copy + 'static,
        Fut: Future<Output = Result<S>> + 'a,
    {
        stream::unfold(
            (initial_state, db_state_rx.clone()),
            move |(s, mut rx)| async move {
                let db_state = rx
                    .recv()
                    .await
                    .ok_or_else(|| anyhow!("The broadcaster has been dropped."))
                    .unwrap();
                let next_state = worker(self, s, db_state).await.unwrap_or_else(|e| {
                    println!("backup failed: {}. Keep trying with state {:?}.", e, s);
                    s
                });
                Some(((), (next_state, rx)))
            },
        )
    }
}

trait Worker<'a, S, Fut: Future<Output = Result<S>> + 'a>:
    Fn(&'a BackupCoordinator, S, DbState) -> Fut
{
}

impl<'a, T, S, Fut> Worker<'a, S, Fut> for T
where
    T: Fn(&'a BackupCoordinator, S, DbState) -> Fut,
    Fut: Future<Output = Result<S>> + 'a,
{
}

fn get_batch_range(last_in_backup: Option<u64>, batch_size: usize) -> (u64, u64) {
    // say, 5 is already in backup, and we target batches of size 10, we will return (6, 9) in this
    // case, so 6, 7, 8, 9 will be in this batch, and next time the backup worker will pass in 9,
    // and we will return (10, 19)
    let start = last_in_backup.map_or(0, |n| n + 1);
    let next_batch_start = (start / batch_size as u64 + 1) * batch_size as u64;

    (start, next_batch_start - 1)
}

fn get_next_snapshot(last_in_backup: Option<u64>, interval: usize) -> u64 {
    match last_in_backup {
        Some(last) => (last / interval as u64 + 1) * interval as u64,
        None => 0,
    }
}

#[cfg(test)]
mod tests {
    use crate::coordinators::backup::get_batch_range;

    #[test]
    fn test_get_batch_range() {
        assert_eq!(get_batch_range(None, 100), (0, 99));
        assert_eq!(get_batch_range(Some(99), 50), (100, 149));
        assert_eq!(get_batch_range(Some(149), 100), (150, 199));
        assert_eq!(get_batch_range(Some(199), 100), (200, 299));
    }
}
<|endoftext|>


<commit_before>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use amq_protocol::frame::{AMQPFrame, gen_frame, parse_frame};

use nom::Offset;
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::io;
use std::iter::repeat;
use futures::{Async,AsyncSink,Poll,Sink,StartSend,Stream,Future,future};
use tokio_codec::{Decoder,Encoder,Framed};
use tokio_io::{AsyncRead,AsyncWrite};

use channel::BasicProperties;
use client::ConnectionOptions;
use error::Error;

#[derive(Fail, Debug)]
pub enum CodecError {
  #[fail(display = "IO Error: {}", _0)]
  IoError(io::Error),
  #[fail(display = "Couldn't parse incoming frame")]
  ParseError,
  #[fail(display = "Couldn't generate outcoming frame: {:?}", _0)]
  GenerationError(GenError),
}

impl From<io::Error> for CodecError {
  fn from(err: io::Error) -> Self {
    CodecError::IoError(err)
  }
}

/// During my testing, it appeared to be the "best" value.
/// To fine-tune it, use a queue with a very large amount of messages, consume it and compare the
/// delivery rate with the ack one.
/// Decreasing this number will increase the ack rate slightly, but decrease the delivery rate by a
/// lot.
/// Increasing this number will increate the delivery rate slightly, but the ack rate will be very
/// close to 0.
const POLL_RECV_LIMIT: u32 = 128;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = CodecError;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<AMQPFrame>, CodecError> {
        let (consumed, f) = match parse_frame(buf) {
          Err(e) => {
            if e.is_incomplete() {
              return Ok(None);
            } else {
              return Err(CodecError::ParseError);
            }
          },
          Ok((i, frame)) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = CodecError;

    fn encode(&mut self, frame: AMQPFrame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        // If the buffer starts running out of capacity (the threshold is 1/4 of a frame), we
        // reserve more bytes upfront to avoid putting too much strain on the allocator.
        if buf.remaining_mut() < frame_max / 4 {
          trace!("encoder; reserve={}", frame_max * 2);
          buf.reserve(frame_max * 2);
        }

        let gen_res = gen_frame((buf, offset), &frame).map(|tup| tup.1);

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(CodecError::GenerationError(e));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream:  Framed<T,AMQPCodec>,
  pub conn:  Connection,
  heartbeat: Option<AMQPFrame>,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: ConnectionOptions) -> impl Future<Item = AMQPTransport<T>, Error = Error> + Send + 'static {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);

    future::result(conn.connect())
      .map_err(|e| Error::new(format!("Failed to connect: {:?}", e)))
      .and_then(|_| {
        let codec = AMQPCodec {
          frame_max: conn.configuration.frame_max,
        };
        let t = AMQPTransport {
          upstream:  codec.framed(stream),
          conn:      conn,
          heartbeat: Some(AMQPFrame::Heartbeat(0)),
        };

        AMQPTransportConnector {
          transport: Some(t),
        }
    })
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: AMQPFrame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Preemptively send an heartbeat frame
  pub fn send_heartbeat(&mut self) -> Poll<(), Error> {
    if let Some(frame) = self.heartbeat.take() {
      self.conn.frame_queue.push_front(frame);
    }
    self.poll_send().map(|r| r.map(|r| {
      // poll_send succeeded, reinitialize self.heartbeat so that we sned a new frame on the next
      // send_heartbeat call
      self.heartbeat = Some(AMQPFrame::Heartbeat(0));
      r
    }))
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  fn poll_recv(&mut self) -> Poll<(), Error> {
    for _ in 0..POLL_RECV_LIMIT {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            return Err(Error::new(format!("failed to handle frame: {:?}", e)));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(Error::new(e.to_string()));
        },
      };
    }
    Ok(Async::NotReady)
  }

  /// Poll the network to send outcoming frames.
  fn poll_send(&mut self) -> Poll<(), Error> {
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.start_send(frame) {
        Ok(AsyncSink::Ready) => {
          trace!("transport poll_send; status=Ready");
        },
        Ok(AsyncSink::NotReady(frame)) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          return Ok(Async::NotReady);
        },
        Err(e) => {
          return Err(Error::new(e.to_string()));
        }
      }
    }
    self.poll_complete()
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<()>, Error> {
      trace!("transport poll");
      if let Async::Ready(()) = self.poll_recv()? {
        trace!("poll transport; status=Ready");
        return Err(Error::new("The connection was closed by the remote peer"));
      }
      self.poll_send().map(|r| r.map(Some))
    }
}

impl <T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Send,
          T: 'static {
    type SinkItem = AMQPFrame;
    type SinkError = Error;

    fn start_send(&mut self, frame: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        trace!("transport start_send; frame={:?}", frame);
        self.upstream.start_send(frame)
          .map_err(|e| Error::new(e.to_string()))
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        trace!("transport poll_complete");
        self.upstream.poll_complete()
          .map_err(|e| Error::new(e.to_string()))
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    transport.poll()?;

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! lock_transport (
    ($t: expr) => ({
        match $t.lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(Error::new("Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

#[cfg(test)]
mod tests {
  extern crate env_logger;

  use super::*;

  #[test]
  fn encode_multiple_frames() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_nested_frame() {
    use amq_protocol::frame::AMQPContentHeader;

    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let frame = AMQPFrame::Header(0, 10, Box::new(AMQPContentHeader {
      class_id: 10,
      weight: 0,
      body_size: 64,
      properties: BasicProperties::default()
    }));
    let r = codec.encode(frame, &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(22, buffer.len());
  }

  #[test]
  fn encode_initial_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= frame_max);
    assert_eq!(8, buffer.len());
  }

  #[test]
  fn encode_anticipation_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());
    assert_eq!(8, buffer.len());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it shouldn't trigger buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it should trigger a buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 4, buffer.capacity());
  }
}
<commit_msg>Display the cause for a ParseError in AMQPCodec<commit_after>/// low level wrapper for the state machine, encoding and decoding from lapin-async
use lapin_async::connection::*;
use amq_protocol::frame::{AMQPFrame, gen_frame, parse_frame};

use nom::Offset;
use cookie_factory::GenError;
use bytes::{BufMut, BytesMut};
use std::cmp;
use std::io;
use std::iter::repeat;
use futures::{Async,AsyncSink,Poll,Sink,StartSend,Stream,Future,future};
use tokio_codec::{Decoder,Encoder,Framed};
use tokio_io::{AsyncRead,AsyncWrite};

use channel::BasicProperties;
use client::ConnectionOptions;
use error::Error;

#[derive(Fail, Debug)]
pub enum CodecError {
  #[fail(display = "IO Error: {}", _0)]
  IoError(io::Error),
  #[fail(display = "Couldn't parse incoming frame: {}", _0)]
  ParseError(String),
  #[fail(display = "Couldn't generate outcoming frame: {:?}", _0)]
  GenerationError(GenError),
}

impl From<io::Error> for CodecError {
  fn from(err: io::Error) -> Self {
    CodecError::IoError(err)
  }
}

/// During my testing, it appeared to be the "best" value.
/// To fine-tune it, use a queue with a very large amount of messages, consume it and compare the
/// delivery rate with the ack one.
/// Decreasing this number will increase the ack rate slightly, but decrease the delivery rate by a
/// lot.
/// Increasing this number will increate the delivery rate slightly, but the ack rate will be very
/// close to 0.
const POLL_RECV_LIMIT: u32 = 128;

/// implements tokio-io's Decoder and Encoder
pub struct AMQPCodec {
    pub frame_max: u32,
}

impl Decoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = CodecError;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<AMQPFrame>, CodecError> {
        let (consumed, f) = match parse_frame(buf) {
          Err(e) => {
            if e.is_incomplete() {
              return Ok(None);
            } else {
              return Err(CodecError::ParseError(format!("{:?}", e)));
            }
          },
          Ok((i, frame)) => {
            (buf.offset(i), frame)
          }
        };

        trace!("amqp decoder; frame={:?}", f);

        buf.split_to(consumed);

        Ok(Some(f))
    }
}

impl Encoder for AMQPCodec {
    type Item = AMQPFrame;
    type Error = CodecError;

    fn encode(&mut self, frame: AMQPFrame, buf: &mut BytesMut) -> Result<(), Self::Error> {
      let frame_max = cmp::max(self.frame_max, 8192) as usize;
      trace!("encoder; frame={:?}", frame);
      let offset = buf.len();
      loop {
        // If the buffer starts running out of capacity (the threshold is 1/4 of a frame), we
        // reserve more bytes upfront to avoid putting too much strain on the allocator.
        if buf.remaining_mut() < frame_max / 4 {
          trace!("encoder; reserve={}", frame_max * 2);
          buf.reserve(frame_max * 2);
        }

        let gen_res = gen_frame((buf, offset), &frame).map(|tup| tup.1);

        match gen_res {
          Ok(sz) => {
            trace!("encoder; frame_size={}", sz - offset);
            return Ok(());
          },
          Err(GenError::BufferTooSmall(sz)) => {
            // BufferTooSmall error variant returns the index the next write would have
            // occured if there was enough space in the buffer. Thus we subtract the
            // buffer's length to know how much bytes we sould make available.
            let length = buf.len();
            trace!("encoder; sz={} length={} extend={}", sz, length, sz - length);
            buf.extend(repeat(0).take(sz - length));
          },
          Err(e) => {
            error!("error generating frame: {:?}", e);
            return Err(CodecError::GenerationError(e));
          }
        }
      }
    }
}

/// Wrappers over a `Framed` stream using `AMQPCodec` and lapin-async's `Connection`
pub struct AMQPTransport<T> {
  upstream:  Framed<T,AMQPCodec>,
  pub conn:  Connection,
  heartbeat: Option<AMQPFrame>,
}

impl<T> AMQPTransport<T>
   where T: AsyncRead+AsyncWrite,
         T: Send,
         T: 'static               {

  /// starts the connection process
  ///
  /// returns a future of a `AMQPTransport` that is connected
  pub fn connect(stream: T, options: ConnectionOptions) -> impl Future<Item = AMQPTransport<T>, Error = Error> + Send + 'static {
    let mut conn = Connection::new();
    conn.set_credentials(&options.username, &options.password);
    conn.set_vhost(&options.vhost);
    conn.set_frame_max(options.frame_max);
    conn.set_heartbeat(options.heartbeat);

    future::result(conn.connect())
      .map_err(|e| Error::new(format!("Failed to connect: {:?}", e)))
      .and_then(|_| {
        let codec = AMQPCodec {
          frame_max: conn.configuration.frame_max,
        };
        let t = AMQPTransport {
          upstream:  codec.framed(stream),
          conn:      conn,
          heartbeat: Some(AMQPFrame::Heartbeat(0)),
        };

        AMQPTransportConnector {
          transport: Some(t),
        }
    })
  }

  /// Send a frame to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frame to a queue, to actually send the frame you have to
  /// call either `poll` or `poll_send`.
  pub fn send_frame(&mut self, frame: AMQPFrame) {
    self.conn.frame_queue.push_back(frame);
  }

  /// Send content frames to the broker.
  ///
  /// # Notes
  ///
  /// This function only appends the frames to a queue, to actually send the frames you have to
  /// call either `poll` or `poll_send`.
  pub fn send_content_frames(&mut self, channel_id: u16, payload: &[u8], properties: BasicProperties) {
    self.conn.send_content_frames(channel_id, 60, payload, properties);
  }

  /// Preemptively send an heartbeat frame
  pub fn send_heartbeat(&mut self) -> Poll<(), Error> {
    if let Some(frame) = self.heartbeat.take() {
      self.conn.frame_queue.push_front(frame);
    }
    self.poll_send().map(|r| r.map(|r| {
      // poll_send succeeded, reinitialize self.heartbeat so that we sned a new frame on the next
      // send_heartbeat call
      self.heartbeat = Some(AMQPFrame::Heartbeat(0));
      r
    }))
  }

  /// Poll the network to receive & handle incoming frames.
  ///
  /// # Return value
  ///
  /// This function will always return `Ok(Async::NotReady)` except in two cases:
  ///
  /// * In case of error, it will return `Err(e)`
  /// * If the socket was closed, it will return `Ok(Async::Ready(()))`
  fn poll_recv(&mut self) -> Poll<(), Error> {
    for _ in 0..POLL_RECV_LIMIT {
      match self.upstream.poll() {
        Ok(Async::Ready(Some(frame))) => {
          trace!("transport poll_recv; frame={:?}", frame);
          if let Err(e) = self.conn.handle_frame(frame) {
            return Err(Error::new(format!("failed to handle frame: {:?}", e)));
          }
        },
        Ok(Async::Ready(None)) => {
          trace!("transport poll_recv; status=Ready(None)");
          return Ok(Async::Ready(()));
        },
        Ok(Async::NotReady) => {
          trace!("transport poll_recv; status=NotReady");
          return Ok(Async::NotReady);
        },
        Err(e) => {
          error!("transport poll_recv; status=Err({:?})", e);
          return Err(Error::new(e.to_string()));
        },
      };
    }
    Ok(Async::NotReady)
  }

  /// Poll the network to send outcoming frames.
  fn poll_send(&mut self) -> Poll<(), Error> {
    while let Some(frame) = self.conn.next_frame() {
      trace!("transport poll_send; frame={:?}", frame);
      match self.start_send(frame) {
        Ok(AsyncSink::Ready) => {
          trace!("transport poll_send; status=Ready");
        },
        Ok(AsyncSink::NotReady(frame)) => {
          trace!("transport poll_send; status=NotReady");
          self.conn.frame_queue.push_front(frame);
          return Ok(Async::NotReady);
        },
        Err(e) => {
          return Err(Error::new(e.to_string()));
        }
      }
    }
    self.poll_complete()
  }
}

impl<T> Stream for AMQPTransport<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<()>, Error> {
      trace!("transport poll");
      if let Async::Ready(()) = self.poll_recv()? {
        trace!("poll transport; status=Ready");
        return Err(Error::new("The connection was closed by the remote peer"));
      }
      self.poll_send().map(|r| r.map(Some))
    }
}

impl <T> Sink for AMQPTransport<T>
    where T: AsyncWrite,
          T: Send,
          T: 'static {
    type SinkItem = AMQPFrame;
    type SinkError = Error;

    fn start_send(&mut self, frame: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        trace!("transport start_send; frame={:?}", frame);
        self.upstream.start_send(frame)
          .map_err(|e| Error::new(e.to_string()))
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        trace!("transport poll_complete");
        self.upstream.poll_complete()
          .map_err(|e| Error::new(e.to_string()))
    }
}

/// implements a future of `AMQPTransport`
///
/// this structure is used to perform the AMQP handshake and provide
/// a connected transport afterwards
pub struct AMQPTransportConnector<T> {
  pub transport: Option<AMQPTransport<T>>,
}

impl<T> Future for AMQPTransportConnector<T>
    where T: AsyncRead + AsyncWrite,
          T: Send,
          T: 'static {

  type Item  = AMQPTransport<T>;
  type Error = Error;

  fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
    trace!("connector poll; has_transport={:?}", !self.transport.is_none());
    let mut transport = self.transport.take().unwrap();

    transport.poll()?;

    trace!("connector poll; state=ConnectionState::{:?}", transport.conn.state);
    if transport.conn.state == ConnectionState::Connected {
      return Ok(Async::Ready(transport))
    }

    self.transport = Some(transport);
    Ok(Async::NotReady)
  }
}

#[macro_export]
macro_rules! lock_transport (
    ($t: expr) => ({
        match $t.lock() {
            Ok(t) => t,
            Err(_) => if $t.is_poisoned() {
                return Err(Error::new("Transport mutex is poisoned"))
            } else {
                task::current().notify();
                return Ok(Async::NotReady)
            }
        }
    });
);

#[cfg(test)]
mod tests {
  extern crate env_logger;

  use super::*;

  #[test]
  fn encode_multiple_frames() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(8, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(16, buffer.len());
    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(24, buffer.len());
  }

  #[test]
  fn encode_nested_frame() {
    use amq_protocol::frame::AMQPContentHeader;

    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let mut buffer = BytesMut::with_capacity(8192);
    let frame = AMQPFrame::Header(0, 10, Box::new(AMQPContentHeader {
      class_id: 10,
      weight: 0,
      body_size: 64,
      properties: BasicProperties::default()
    }));
    let r = codec.encode(frame, &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(22, buffer.len());
  }

  #[test]
  fn encode_initial_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(true, buffer.capacity() >= frame_max);
    assert_eq!(8, buffer.len());
  }

  #[test]
  fn encode_anticipation_extend_buffer() {
    let _ = env_logger::try_init();

    let mut codec = AMQPCodec { frame_max: 8192 };
    let frame_max = codec.frame_max as usize;
    let mut buffer = BytesMut::new();

    let r = codec.encode(AMQPFrame::Heartbeat(0), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());
    assert_eq!(8, buffer.len());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it shouldn't trigger buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 2, buffer.capacity());

    let payload = repeat(0u8)
      // Use 80% of the remaining space (it should trigger a buffer capacity expansion)
      .take(((buffer.capacity() as f64 - buffer.len() as f64) * 0.8) as usize)
      .collect::<Vec<u8>>();
    let r = codec.encode(AMQPFrame::Body(1, payload), &mut buffer);
    assert_eq!(false, r.is_err());
    assert_eq!(frame_max * 4, buffer.capacity());
  }
}
<|endoftext|>


<commit_before>use std::io;

use tag;
use err;
use serial;
use byte::{read_byte, write_byte};

use std::cmp;

macro_rules! asn1_cereal_int {
  ($rs_type:ty, $size:expr) => (
    impl serial::traits::Asn1Info for $rs_type {
      fn asn1_tag() -> tag::Tag {
        tag::Tag {
          class: tag::Class::Universal,
          tagnum: 2u8.into(),
          constructed: false,
        }
      }

      fn asn1_type() -> tag::Type {
        tag::Type::from("INTEGER")
      }
    }

    impl serial::traits::Asn1Serialize for $rs_type {
      fn serialize_imp<W: io::Write>(&self, writer: &mut W) -> Result<(), err::EncodeError> {
        let mut started = false;
        // Loop through bytes in int backwards, start writing when first non-zero byte is encounted.
        for offset in (0..8).rev() {
          let shifted: $rs_type = self >> (offset * 8);
          let byte: u8 = (shifted & 0xff) as u8;
          if !started {
            if byte == 0 {
              continue;
            }
            started = true;
          }
          try!(write_byte(writer, byte));
        }

        // If we haven't written anything, write a zero byte.
        if !started {
          try!(write_byte(writer, 0 as u8));
        }
        Ok(())
      }
    }

    impl serial::traits::Asn1Deserialize for $rs_type {
      fn deserialize_imp<I: Iterator<Item=io::Result<u8>>>(reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
        let len_num = try!(len.as_num().ok_or(err::DecodeError::PrimIndef));

        let mut int: $rs_type = 0;
        let size = cmp::min($size, len_num);

        for _ in 0..size {
          let byte = try!(read_byte(reader));
          int = (int << 8) + (byte as $rs_type);
        }

        Ok(int)
      }
    }
  )
}

// Can't bitshift a u8 by 8 bits, need to update deserialize_imp.
// asn1_cereal_int!(u8, 1);
asn1_cereal_int!(u16, 2);
asn1_cereal_int!(u32, 4);
asn1_cereal_int!(u64, 8);
<commit_msg>Rename 'int' macro to 'uint'<commit_after>use std::io;

use tag;
use err;
use serial;
use byte::{read_byte, write_byte};

use std::cmp;

macro_rules! asn1_cereal_uint {
  ($rs_type:ty, $size:expr) => (
    impl serial::traits::Asn1Info for $rs_type {
      fn asn1_tag() -> tag::Tag {
        tag::Tag {
          class: tag::Class::Universal,
          tagnum: 2u8.into(),
          constructed: false,
        }
      }

      fn asn1_type() -> tag::Type {
        tag::Type::from("INTEGER")
      }
    }

    impl serial::traits::Asn1Serialize for $rs_type {
      fn serialize_imp<W: io::Write>(&self, writer: &mut W) -> Result<(), err::EncodeError> {
        let mut started = false;
        // Loop through bytes in int backwards, start writing when first non-zero byte is encounted.
        for offset in (0..8).rev() {
          let shifted: $rs_type = self >> (offset * 8);
          let byte: u8 = (shifted & 0xff) as u8;
          if !started {
            if byte == 0 {
              continue;
            }
            started = true;
          }
          try!(write_byte(writer, byte));
        }

        // If we haven't written anything, write a zero byte.
        if !started {
          try!(write_byte(writer, 0 as u8));
        }
        Ok(())
      }
    }

    impl serial::traits::Asn1Deserialize for $rs_type {
      fn deserialize_imp<I: Iterator<Item=io::Result<u8>>>(reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
        let len_num = try!(len.as_num().ok_or(err::DecodeError::PrimIndef));

        let mut int: $rs_type = 0;
        let size = cmp::min($size, len_num);

        for _ in 0..size {
          let byte = try!(read_byte(reader));
          int = (int << 8) + (byte as $rs_type);
        }

        Ok(int)
      }
    }
  )
}

// Can't bitshift a u8 by 8 bits, need to update deserialize_imp.
// asn1_cereal_int!(u8, 1);
asn1_cereal_uint!(u16, 2);
asn1_cereal_uint!(u32, 4);
asn1_cereal_uint!(u64, 8);
<|endoftext|>


<commit_before>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint, RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float, FromPrimitive};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::clone::Clone;
use std::ops::{Shr, BitAnd};

static pow_offsets : &'static[uint] = &[0u, 4u, 6u, 10u, 12u, 16u];

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Miller-Rabin, probabilistic
fn try_composite(a : &BigUint, d : &BigUint, n : &BigUint, s : &BigUint) -> bool {

    let one : BigUint = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if big(i) == *s {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn is_prime(n : &BigUint) -> bool {

    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_biguint_range(&two, n);
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn fermat_is_prime(n : &BigUint) -> bool {
    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    if n == &two {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&two, &(n-one), n) == one {
        true
    } else {
        false
    }
}

fn is_valid_pow(prime : &BigUint) -> bool {
    for offset in [big(0u), big(4u), big(6u), big(10u), big(12u), big(16u)].iter() {
        if !is_prime(&(*prime + *offset)) {
            return false;
        }
    }
    true
}

// Source: http://rosettacode.org/wiki/Sieve_of_Eratosthenes#Rust
fn simple_sieve(limit: uint) -> ~[uint] {

    #[inline]
    fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

    if limit < 2 {
        return ~[];
    }

    let mut primes = Vec::from_elem(limit + 1, true);

    for prime in iter::range_inclusive(2, int_sqrt(limit) + 1) {
        if *primes.get(prime) {
            for multiple in iter::range_step(prime * prime, limit + 1, prime) {
                *primes.get_mut(multiple) = false
            }
        }
    }
    iter::range_inclusive(2, limit).filter(|&n| *primes.get(n)).collect()

}

fn gen_prime(max_val : &BigUint, verbose : bool) -> uint {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }
    let primorial_max : uint = 29u;
    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);
    for i in simple_sieve(primorial_max).iter() {
        if *i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(*i), big(primorial));
        primorial = primorial * *i;
    }
    let mut count : uint = 0;
    for o in offsets.iter() {
        if is_valid_pow(o) {
            count = count + 1;
            if verbose { println!("prime: {}", o) }
        }
    }
    count
}

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

    #[test]
    fn gen_prime_is_correct() {
        let v = gen_prime(&big(1000000000), true);
        println!("{}", v);
        assert!(v == 81u);
    }
}

fn main() {
    println!("{}", gen_prime(&big(1000000000), false));
}<commit_msg>Fix some bugs<commit_after>extern crate num;
extern crate rand;
extern crate bignum = "bignum#0.1.1-pre";

use bignum::{BigUint, ToBigUint, RandBigInt};
use num::{div_mod_floor, Integer};
use std::num::{Zero, One, pow, Float, FromPrimitive};
use rand::{task_rng, Rng};
use rand::distributions::range::SampleRange;
use std::vec::Vec;
use std::iter;
use std::clone::Clone;
use std::ops::{Shr, BitAnd};

static pow_offsets : &'static[uint] = &[0u, 4u, 6u, 10u, 12u, 16u];

#[inline]
fn big<T: ToBigUint>(i : T) -> BigUint {
    i.to_biguint().unwrap()
}

// Source: https://gist.github.com/jsanders/8739134#file-generate_big_primes-rs-L35
fn mod_exp<T: Integer + Clone + Shr<uint,T> + BitAnd<T,T>>
    (base: &T, exponent: &T, modulus: &T) -> T {

    let (zero, one): (T, T) = (Zero::zero(), One::one());
    let mut result = one.clone();
    let mut baseAcc = base.clone();
    let mut exponentAcc = exponent.clone();

    while exponentAcc > zero {
        // Accumulate current base if current exponent bit is 1
        if (exponentAcc & one) == one {
            result = result.mul(&baseAcc);
            result = result.rem(modulus);
        }
        // Get next base by squaring
        baseAcc = baseAcc * baseAcc;
        baseAcc = baseAcc % *modulus;

        // Get next bit of exponent
        exponentAcc = exponentAcc.shr(&1);
    }
    result

}

// Rabin-Miller primality test, probabilistic
fn try_composite(a : &BigUint, d : &BigUint, n : &BigUint, s : &BigUint) -> bool {

    let one : BigUint = One::one();
    if mod_exp(a, d, n) == one {
        return false
    }
    let mut i = 0u;
    loop {
        if big(i) == *s {
            break
        }
        if mod_exp(a, &(pow(one+one, i) * *d), n) == *n-one {
            return false
        }
        i += 1;
    }
    true

}
fn is_prime(n : &BigUint) -> bool {

    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    assert!(n >= &two);

    let mut s = zero.clone();
    let mut d = *n-one;
    let num_trials = 5;
    if n == &two {
        return true
    } else if n == &one || n.is_even() {
        return false
    } else {
        loop {
            let (quotient, remainder) = div_mod_floor(d.clone(), two.clone());
            if remainder == one { break }
            s = s + one;
            d = quotient;
        }
    }

    let mut rng = task_rng();
    for _ in range(0, num_trials) {
        let a = rng.gen_biguint_range(&two, n);
        if try_composite(&a, &d, n, &s) {
            return false
        }
    }
    true

}

fn fermat_is_prime(n : &BigUint) -> bool {
    let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
    let two = one + one;

    if n == &two {
        true
    } else if n.is_even() {
        false
    } else if mod_exp(&two, &(n-one), n) == one {
        true
    } else {
        false
    }
}

fn is_valid_pow(prime : &BigUint) -> bool {
    for offset in [big(0u), big(4u), big(6u), big(10u), big(12u), big(16u)].iter() {
        if !is_prime(&(*prime + *offset)) {
            return false;
        }
    }
    true
}

// Source: http://rosettacode.org/wiki/Sieve_of_Eratosthenes#Rust
fn simple_sieve(limit: uint) -> ~[uint] {

    #[inline]
    fn int_sqrt(n: uint) -> uint { (n as f64).sqrt() as uint }

    if limit < 2 {
        return ~[];
    }

    let mut primes = Vec::from_elem(limit + 1, true);

    for prime in iter::range_inclusive(2, int_sqrt(limit) + 1) {
        if *primes.get(prime) {
            for multiple in iter::range_step(prime * prime, limit + 1, prime) {
                *primes.get_mut(multiple) = false
            }
        }
    }
    iter::range_inclusive(2, limit).filter(|&n| *primes.get(n)).collect()

}

fn gen_prime(max_val : &BigUint, verbose : bool) -> uint {

    fn candidate_killed_by(candidate : &BigUint, prime : &BigUint) -> bool {
        let zero : BigUint = Zero::zero();
        for &offset in pow_offsets.iter() {
            let o : BigUint = FromPrimitive::from_uint(offset).unwrap();
            if (*candidate + o) % *prime == zero {
                return true;
            }
        }
        false
    }

    fn add_next_prime (max_val : &BigUint, offsets : Vec<BigUint>, prime : BigUint, primorial : BigUint) -> Vec<BigUint> {
        let (zero, one): (BigUint, BigUint) = (Zero::zero(), One::one());
        let mut base : BigUint = Zero::zero();
        let mut new_offsets : Vec<BigUint> = Vec::new();
        let mut counter = zero.clone();
        while counter < prime {
            if base > *max_val {
                break
            }
            for o in offsets.iter() {
                    let val = base + *o;
                    if val > *max_val {
                        break
                    }
                    if !candidate_killed_by(&val, &prime) {
                        new_offsets.push(val);
                    }
            }
            base = base + primorial;
            counter = counter + one;
        }
        return new_offsets
    }
    let primorial_max : uint = 29u;
    let primorial_start : uint = 7u;
    let mut primorial : uint = 210u;  // 2*3*5*7
    let mut offsets : Vec<BigUint> = Vec::from_slice([big(97u)]);
    for i in simple_sieve(primorial_max).iter() {
        if *i <= primorial_start {
            continue
        }
        offsets = add_next_prime(max_val, offsets, big(*i), big(primorial));
        primorial = primorial * *i;
    }
    let mut count : uint = 0;
    for o in offsets.iter() {
        if is_valid_pow(o) {
            count = count + 1;
            if verbose { println!("prime: {}", o) }
        }
    }
    count
}

#[cfg(test)]
mod test_primes {
    use super::{big, gen_prime};

    #[test]
    fn gen_prime_is_correct() {
        let v = gen_prime(&big(1000000000), true);
        println!("{}", v);
        assert!(v == 81u);
    }
}

fn main() {
    println!("{}", gen_prime(&big(1000000000), false));
}<|endoftext|>


<commit_before>/**
 * > Every morning in the early part of October 1843, on my coming down to
 *   breakfast, your brother William Edward and yourself used to ask me: "Well,
 *   Papa, can you multiply triples?" Whereto I was always obliged to reply,
 *   with a sad shake of the head, "No, I can only add and subtract them."
 *
 *   Sir William Hamilton
 */

use core::cast::transmute;
use core::cmp::{Eq, Ord};
use core::ptr::to_unsafe_ptr;
use core::sys::size_of;
use core::vec::raw::buf_as_slice;

use std::cmp::FuzzyEq;

use angle::Angle;
use funs::common::*;
use funs::exponential::*;
use funs::triganomic::*;
use mat::{Mat3, Mat4};
use num::types::{Float, Number};
use vec::Vec3;

/**
 * The base quaternion trait
 *
 * # Type parameters
 *
 * * `T` - The type of the components. Should be a floating point type.
 * * `V3` - The 3-dimensional vector type that will containin the imaginary
 *          components of the quaternion.
 */
pub trait Quaternion<T,V3>: Index<uint, T> Eq Neg<self> {
    /**
     * # Return value
     *
     * The multiplicative identity, ie: `q = 1 + 0i + 0j + 0i`
     */
    static pure fn identity() -> self;

    /**
     * # Return value
     *
     * The additive identity, ie: `q = 0 + 0i + 0j + 0i`
     */
    static pure fn zero() -> self;

    /**
     * # Return value
     *
     * The result of multiplying the quaternion a scalar
     */
    pure fn mul_t(&self, value: T) -> self;

    /**
     * # Return value
     *
     * The result of dividing the quaternion a scalar
     */
    pure fn div_t(&self, value: T) -> self;

    /**
     * # Return value
     *
     * The result of multiplying the quaternion by a vector
     */
    pure fn mul_v(&self, vec: &V3) -> V3;

    /**
     * # Return value
     *
     * The sum of this quaternion and `other`
     */
    pure fn add_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The sum of this quaternion and `other`
     */
    pure fn sub_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The the result of multipliplying the quaternion by `other`
     */
    pure fn mul_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The dot product of the quaternion and `other`
     */
    pure fn dot(&self, other: &self) -> T;

    /**
     * # Return value
     *
     * The conjugate of the quaternion
     */
    pure fn conjugate(&self) -> self;

    /**
     * # Return value
     *
     * The multiplicative inverse of the quaternion
     */
    pure fn inverse(&self) -> self;

    /**
     * # Return value
     *
     * The squared magnitude of the quaternion. This is useful for
     * magnitude comparisons where the exact magnitude does not need to be
     * calculated.
     */
    pure fn magnitude2(&self) -> T;

    /**
     * # Return value
     *
     * The magnitude of the quaternion
     *
     * # Performance notes
     *
     * For instances where the exact magnitude of the quaternion does not need
     * to be known, for example for quaternion-quaternion magnitude comparisons,
     * it is advisable to use the `magnitude2` method instead.
     */
    pure fn magnitude(&self) -> T;

    /**
     * # Return value
     *
     * The normalized quaternion
     */
    pure fn normalize(&self) -> self;

    /**
     * Normalised linear interpolation
     *
     * # Return value
     *
     * The intoperlated quaternion
     */
    pure fn nlerp(&self, other: &self, amount: T) -> self;

    /**
     * Perform a spherical linear interpolation between the quaternion and
     * `other`.
     *
     * # Return value
     *
     * The intoperlated quaternion
     *
     * # Performance notes
     *
     * This is more accurate than `nlerp` but is also more
     * computationally intensive.
     */
    pure fn slerp(&self, other: &self, amount: T) -> self;

    /**
     * Convert the quaternion to a 3 x 3 rotation matrix
     */
    pure fn to_mat3(&self) -> Mat3<T>;

    /**
     * Convert the quaternion to a 4 x 4 transformation matrix
     */
    pure fn to_mat4(&self) -> Mat4<T>;

    /**
     * # Return value
     *
     * A pointer to the first component of the quaternion
     */
    pure fn to_ptr(&self) -> *T;
}

pub trait ToQuat<T> {
    /**
     * Convert `self` to a quaternion
     */
    pure fn to_Quat() -> Quat<T>;
}

/**
 * A quaternion in scalar/vector form
 *
 * # Type parameters
 *
 * * `T` - The type of the components. Should be a floating point type.
 *
 * # Fields
 *
 * * `s` - the scalar component
 * * `v` - a vector containing the three imaginary components
 */
pub struct Quat<T> { s: T, v: Vec3<T> }

pub impl<T:Copy Float> Quat<T> {
    /**
     * Construct the quaternion from one scalar component and three
     * imaginary components
     *
     * # Arguments
     *
     * * `w`  - the scalar component
     * * `xi` - the fist imaginary component
     * * `yj` - the second imaginary component
     * * `zk` - the third imaginary component
     */
    #[inline(always)]
    static pure fn new(w: T, xi: T, yj: T, zk: T) -> Quat<T> {
        Quat::from_sv(move w, move Vec3::new(move xi, move yj, move zk))
    }

    /**
     * Construct the quaternion from a scalar and a vector
     *
     * # Arguments
     *
     * * `s` - the scalar component
     * * `v` - a vector containing the three imaginary components
     */
    #[inline(always)]
    static pure fn from_sv(s: T, v: Vec3<T>) -> Quat<T> {
        Quat { s: move s, v: move v }
    }

    #[inline(always)]
    static pure fn from_axis_angle<A:Angle<T>>(axis: &Vec3<T>, theta: A) -> Quat<T> {
        let half = theta.to_radians() / Number::from(2);
        Quat::from_sv(cos(&half), axis.mul_t(sin(&half)))
    }
}

pub impl<T:Copy> Quat<T>: Index<uint, T> {
    #[inline(always)]
    pure fn index(&self, i: uint) -> T {
        unsafe { do buf_as_slice(
            transmute::<*Quat<T>, *T>(
                to_unsafe_ptr(self)), 4) |slice| { slice[i] }
        }
    }
}

pub impl<T:Copy Float Exp Extent InvTrig> Quat<T>: Quaternion<T, Vec3<T>> {
    #[inline(always)]
    static pure fn identity() -> Quat<T> {
        Quat::new(Number::from(1),
                  Number::from(0),
                  Number::from(0),
                  Number::from(0))
    }

    #[inline(always)]
    static pure fn zero() -> Quat<T> {
        Quat::new(Number::from(0),
                  Number::from(0),
                  Number::from(0),
                  Number::from(0))
    }

    #[inline(always)]
    pure fn mul_t(&self, value: T) -> Quat<T> {
        Quat::new(self[0] * value,
                  self[1] * value,
                  self[2] * value,
                  self[3] * value)
    }

    #[inline(always)]
    pure fn div_t(&self, value: T) -> Quat<T> {
        Quat::new(self[0] / value,
                  self[1] / value,
                  self[2] / value,
                  self[3] / value)
    }

    #[inline(always)]
    pure fn mul_v(&self, vec: &Vec3<T>) -> Vec3<T>  {
        let tmp = self.v.cross(vec).add_v(&vec.mul_t(self.s));
        self.v.cross(&tmp).mul_t(Number::from(2)).add_v(vec)
    }

    #[inline(always)]
    pure fn add_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self[0] + other[0],
                  self[1] + other[1],
                  self[2] + other[2],
                  self[3] + other[3])
    }

    #[inline(always)]
    pure fn sub_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self[0] - other[0],
                  self[1] - other[1],
                  self[2] - other[2],
                  self[3] - other[3])
    }

    #[inline(always)]
    pure fn mul_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self.s * other.s   - self.v.x * other.v.x - self.v.y * other.v.y - self.v.z * other.v.z,
                  self.s * other.v.x + self.v.x * other.s   + self.v.y * other.v.z - self.v.z * other.v.y,
                  self.s * other.v.y + self.v.y * other.s   + self.v.z * other.v.x - self.v.x * other.v.z,
                  self.s * other.v.z + self.v.z * other.s   + self.v.x * other.v.y - self.v.y * other.v.x)
    }

    #[inline(always)]
    pure fn dot(&self, other: &Quat<T>) -> T {
        self.s * other.s + self.v.dot(&other.v)
    }

    #[inline(always)]
    pure fn conjugate(&self) -> Quat<T> {
        Quat::from_sv(self.s, -self.v)
    }

    #[inline(always)]
    pure fn inverse(&self) -> Quat<T> {
        self.conjugate().div_t(self.magnitude2())
    }

    #[inline(always)]
    pure fn magnitude2(&self) -> T {
        self.s * self.s + self.v.length2()
    }

    #[inline(always)]
    pure fn magnitude(&self) -> T {
        self.magnitude2().sqrt()
    }

    #[inline(always)]
    pure fn normalize(&self) -> Quat<T> {
        let mut n: T = Number::from(1);
        n /= self.magnitude();
        return self.mul_t(n);
    }

    #[inline(always)]
    pure fn nlerp(&self, other: &Quat<T>, amount: T) -> Quat<T> {
        let _1: T = Number::from(1);
        self.mul_t(_1 - amount).add_q(&other.mul_t(amount)).normalize()
    }

    /**
     * Spherical Linear Intoperlation
     *
     * Both quaternions should be normalized first, or else strange things will
     * will happen...
     *
     * # Performance notes
     *
     * The `acos` operation used in `slerp` is an expensive operation, so unless
     * your quarternions a far away from each other it's generally more advisable
     * to use `nlerp` when you know your rotations are going to be small.
     *
     * - [Understanding Slerp, Then Not Using It]
     *   (http://number-none.com/product/Understanding%20Slerp,%20Then%20Not%20Using%20It/)
     * - [Arcsynthesis OpenGL tutorial]
     *   (http://www.arcsynthesis.org/gltut/Positioning/Tut08%20Interpolation.html)
     */
    #[inline(always)]
    pure fn slerp(&self, other: &Quat<T>, amount: T) -> Quat<T> {
        let dot = self.dot(other);

        let dot_threshold = Number::from(0.9995);

        if dot > dot_threshold {
            return self.nlerp(other, amount);                   // if quaternions are close together use `nlerp`
        } else {
            let robust_dot = dot.clamp(&-Number::from(1),
                                       &Number::from(1));       // stay within the domain of acos()

            let theta_0 = acos(&robust_dot);                    // the angle between the quaternions
            let theta = theta_0 * amount;                       // the fraction of theta specified by `amount`

            let q = other.sub_q(&self.mul_t(robust_dot))
                         .normalize();

            return self.mul_t(cos(&theta)).add_q(&q.mul_t(sin(&theta)));
        }
    }

    #[inline(always)]
    pure fn to_mat3(&self) -> Mat3<T> {
        let x2 = self.v.x + self.v.x;
        let y2 = self.v.y + self.v.y;
        let z2 = self.v.z + self.v.z;

        let xx2 = x2 * self.v.x;
        let xy2 = x2 * self.v.y;
        let xz2 = x2 * self.v.z;

        let yy2 = y2 * self.v.y;
        let yz2 = y2 * self.v.z;
        let zz2 = z2 * self.v.z;

        let sy2 = y2 * self.s;
        let sz2 = z2 * self.s;
        let sx2 = x2 * self.s;

        let _1: T = Number::from(1);

        Mat3::new(_1 - yy2 - zz2,      xy2 + sz2,      xz2 - sy2,
                       xy2 - sz2, _1 - xx2 - zz2,      yz2 + sx2,
                       xz2 + sy2,      yz2 - sx2, _1 - xx2 - yy2)
    }

    #[inline(always)]
    pure fn to_mat4(&self) -> Mat4<T> {
        self.to_mat3().to_mat4()
    }

    #[inline(always)]
    pure fn to_ptr(&self) -> *T {
        unsafe {
            transmute::<*Quat<T>, *T>(
                to_unsafe_ptr(self)
            )
        }
    }
}

pub impl<T:Copy Float> Quat<T>: Neg<Quat<T>> {
    #[inline(always)]
    pure fn neg(&self) -> Quat<T> {
        Quat::new(-self[0], -self[1], -self[2], -self[3])
    }
}

pub impl<T:Copy Eq> Quat<T>: Eq {
    #[inline(always)]
    pure fn eq(&self, other: &Quat<T>) -> bool {
        self[0] == other[0] &&
        self[1] == other[1] &&
        self[2] == other[2] &&
        self[3] == other[3]
    }

    #[inline(always)]
    pure fn ne(&self, other: &Quat<T>) -> bool {
        !(self == other)
    }
}

pub impl<T:Copy FuzzyEq> Quat<T>: FuzzyEq {
    #[inline(always)]
    pure fn fuzzy_eq(other: &Quat<T>) -> bool {
        self[0].fuzzy_eq(&other[0]) &&
        self[1].fuzzy_eq(&other[1]) &&
        self[2].fuzzy_eq(&other[2]) &&
        self[3].fuzzy_eq(&other[3])
    }
}
<commit_msg>Add identity and zero static methods to Quat's impl<commit_after>/**
 * > Every morning in the early part of October 1843, on my coming down to
 *   breakfast, your brother William Edward and yourself used to ask me: "Well,
 *   Papa, can you multiply triples?" Whereto I was always obliged to reply,
 *   with a sad shake of the head, "No, I can only add and subtract them."
 *
 *   Sir William Hamilton
 */

use core::cast::transmute;
use core::cmp::{Eq, Ord};
use core::ptr::to_unsafe_ptr;
use core::sys::size_of;
use core::vec::raw::buf_as_slice;

use std::cmp::FuzzyEq;

use angle::Angle;
use funs::common::*;
use funs::exponential::*;
use funs::triganomic::*;
use mat::{Mat3, Mat4};
use num::types::{Float, Number};
use vec::Vec3;

/**
 * The base quaternion trait
 *
 * # Type parameters
 *
 * * `T` - The type of the components. Should be a floating point type.
 * * `V3` - The 3-dimensional vector type that will containin the imaginary
 *          components of the quaternion.
 */
pub trait Quaternion<T,V3>: Index<uint, T> Eq Neg<self> {
    /**
     * # Return value
     *
     * The multiplicative identity, ie: `q = 1 + 0i + 0j + 0i`
     */
    static pure fn identity() -> self;

    /**
     * # Return value
     *
     * The additive identity, ie: `q = 0 + 0i + 0j + 0i`
     */
    static pure fn zero() -> self;

    /**
     * # Return value
     *
     * The result of multiplying the quaternion a scalar
     */
    pure fn mul_t(&self, value: T) -> self;

    /**
     * # Return value
     *
     * The result of dividing the quaternion a scalar
     */
    pure fn div_t(&self, value: T) -> self;

    /**
     * # Return value
     *
     * The result of multiplying the quaternion by a vector
     */
    pure fn mul_v(&self, vec: &V3) -> V3;

    /**
     * # Return value
     *
     * The sum of this quaternion and `other`
     */
    pure fn add_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The sum of this quaternion and `other`
     */
    pure fn sub_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The the result of multipliplying the quaternion by `other`
     */
    pure fn mul_q(&self, other: &self) -> self;

    /**
     * # Return value
     *
     * The dot product of the quaternion and `other`
     */
    pure fn dot(&self, other: &self) -> T;

    /**
     * # Return value
     *
     * The conjugate of the quaternion
     */
    pure fn conjugate(&self) -> self;

    /**
     * # Return value
     *
     * The multiplicative inverse of the quaternion
     */
    pure fn inverse(&self) -> self;

    /**
     * # Return value
     *
     * The squared magnitude of the quaternion. This is useful for
     * magnitude comparisons where the exact magnitude does not need to be
     * calculated.
     */
    pure fn magnitude2(&self) -> T;

    /**
     * # Return value
     *
     * The magnitude of the quaternion
     *
     * # Performance notes
     *
     * For instances where the exact magnitude of the quaternion does not need
     * to be known, for example for quaternion-quaternion magnitude comparisons,
     * it is advisable to use the `magnitude2` method instead.
     */
    pure fn magnitude(&self) -> T;

    /**
     * # Return value
     *
     * The normalized quaternion
     */
    pure fn normalize(&self) -> self;

    /**
     * Normalised linear interpolation
     *
     * # Return value
     *
     * The intoperlated quaternion
     */
    pure fn nlerp(&self, other: &self, amount: T) -> self;

    /**
     * Perform a spherical linear interpolation between the quaternion and
     * `other`.
     *
     * # Return value
     *
     * The intoperlated quaternion
     *
     * # Performance notes
     *
     * This is more accurate than `nlerp` but is also more
     * computationally intensive.
     */
    pure fn slerp(&self, other: &self, amount: T) -> self;

    /**
     * Convert the quaternion to a 3 x 3 rotation matrix
     */
    pure fn to_mat3(&self) -> Mat3<T>;

    /**
     * Convert the quaternion to a 4 x 4 transformation matrix
     */
    pure fn to_mat4(&self) -> Mat4<T>;

    /**
     * # Return value
     *
     * A pointer to the first component of the quaternion
     */
    pure fn to_ptr(&self) -> *T;
}

pub trait ToQuat<T> {
    /**
     * Convert `self` to a quaternion
     */
    pure fn to_Quat() -> Quat<T>;
}

/**
 * A quaternion in scalar/vector form
 *
 * # Type parameters
 *
 * * `T` - The type of the components. Should be a floating point type.
 *
 * # Fields
 *
 * * `s` - the scalar component
 * * `v` - a vector containing the three imaginary components
 */
pub struct Quat<T> { s: T, v: Vec3<T> }

pub impl<T:Copy Float> Quat<T> {
    /**
     * Construct the quaternion from one scalar component and three
     * imaginary components
     *
     * # Arguments
     *
     * * `w`  - the scalar component
     * * `xi` - the fist imaginary component
     * * `yj` - the second imaginary component
     * * `zk` - the third imaginary component
     */
    #[inline(always)]
    static pure fn new(w: T, xi: T, yj: T, zk: T) -> Quat<T> {
        Quat::from_sv(move w, move Vec3::new(move xi, move yj, move zk))
    }

    /**
     * Construct the quaternion from a scalar and a vector
     *
     * # Arguments
     *
     * * `s` - the scalar component
     * * `v` - a vector containing the three imaginary components
     */
    #[inline(always)]
    static pure fn from_sv(s: T, v: Vec3<T>) -> Quat<T> {
        Quat { s: move s, v: move v }
    }

    #[inline(always)]
    static pure fn from_axis_angle<A:Angle<T>>(axis: &Vec3<T>, theta: A) -> Quat<T> {
        let half = theta.to_radians() / Number::from(2);
        Quat::from_sv(cos(&half), axis.mul_t(sin(&half)))
    }

    /**
     * # Return value
     *
     * The multiplicative identity, ie: `q = 1 + 0i + 0j + 0i`
     */
    #[inline(always)]
    static pure fn identity() -> Quat<T> {
        Quat::new(Number::one(),
                  Number::zero(),
                  Number::zero(),
                  Number::zero())
    }

    /**
     * # Return value
     *
     * The additive identity, ie: `q = 0 + 0i + 0j + 0i`
     */
    #[inline(always)]
    static pure fn zero() -> Quat<T> {
        Quat::new(Number::zero(),
                  Number::zero(),
                  Number::zero(),
                  Number::zero())
    }
}

pub impl<T:Copy> Quat<T>: Index<uint, T> {
    #[inline(always)]
    pure fn index(&self, i: uint) -> T {
        unsafe { do buf_as_slice(
            transmute::<*Quat<T>, *T>(
                to_unsafe_ptr(self)), 4) |slice| { slice[i] }
        }
    }
}

pub impl<T:Copy Float Exp Extent InvTrig> Quat<T>: Quaternion<T, Vec3<T>> {
    #[inline(always)]
    static pure fn identity() -> Quat<T> {
        Quat::new(Number::from(1),
                  Number::from(0),
                  Number::from(0),
                  Number::from(0))
    }

    #[inline(always)]
    static pure fn zero() -> Quat<T> {
        Quat::new(Number::from(0),
                  Number::from(0),
                  Number::from(0),
                  Number::from(0))
    }

    #[inline(always)]
    pure fn mul_t(&self, value: T) -> Quat<T> {
        Quat::new(self[0] * value,
                  self[1] * value,
                  self[2] * value,
                  self[3] * value)
    }

    #[inline(always)]
    pure fn div_t(&self, value: T) -> Quat<T> {
        Quat::new(self[0] / value,
                  self[1] / value,
                  self[2] / value,
                  self[3] / value)
    }

    #[inline(always)]
    pure fn mul_v(&self, vec: &Vec3<T>) -> Vec3<T>  {
        let tmp = self.v.cross(vec).add_v(&vec.mul_t(self.s));
        self.v.cross(&tmp).mul_t(Number::from(2)).add_v(vec)
    }

    #[inline(always)]
    pure fn add_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self[0] + other[0],
                  self[1] + other[1],
                  self[2] + other[2],
                  self[3] + other[3])
    }

    #[inline(always)]
    pure fn sub_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self[0] - other[0],
                  self[1] - other[1],
                  self[2] - other[2],
                  self[3] - other[3])
    }

    #[inline(always)]
    pure fn mul_q(&self, other: &Quat<T>) -> Quat<T> {
        Quat::new(self.s * other.s   - self.v.x * other.v.x - self.v.y * other.v.y - self.v.z * other.v.z,
                  self.s * other.v.x + self.v.x * other.s   + self.v.y * other.v.z - self.v.z * other.v.y,
                  self.s * other.v.y + self.v.y * other.s   + self.v.z * other.v.x - self.v.x * other.v.z,
                  self.s * other.v.z + self.v.z * other.s   + self.v.x * other.v.y - self.v.y * other.v.x)
    }

    #[inline(always)]
    pure fn dot(&self, other: &Quat<T>) -> T {
        self.s * other.s + self.v.dot(&other.v)
    }

    #[inline(always)]
    pure fn conjugate(&self) -> Quat<T> {
        Quat::from_sv(self.s, -self.v)
    }

    #[inline(always)]
    pure fn inverse(&self) -> Quat<T> {
        self.conjugate().div_t(self.magnitude2())
    }

    #[inline(always)]
    pure fn magnitude2(&self) -> T {
        self.s * self.s + self.v.length2()
    }

    #[inline(always)]
    pure fn magnitude(&self) -> T {
        self.magnitude2().sqrt()
    }

    #[inline(always)]
    pure fn normalize(&self) -> Quat<T> {
        let mut n: T = Number::from(1);
        n /= self.magnitude();
        return self.mul_t(n);
    }

    #[inline(always)]
    pure fn nlerp(&self, other: &Quat<T>, amount: T) -> Quat<T> {
        let _1: T = Number::from(1);
        self.mul_t(_1 - amount).add_q(&other.mul_t(amount)).normalize()
    }

    /**
     * Spherical Linear Intoperlation
     *
     * Both quaternions should be normalized first, or else strange things will
     * will happen...
     *
     * # Performance notes
     *
     * The `acos` operation used in `slerp` is an expensive operation, so unless
     * your quarternions a far away from each other it's generally more advisable
     * to use `nlerp` when you know your rotations are going to be small.
     *
     * - [Understanding Slerp, Then Not Using It]
     *   (http://number-none.com/product/Understanding%20Slerp,%20Then%20Not%20Using%20It/)
     * - [Arcsynthesis OpenGL tutorial]
     *   (http://www.arcsynthesis.org/gltut/Positioning/Tut08%20Interpolation.html)
     */
    #[inline(always)]
    pure fn slerp(&self, other: &Quat<T>, amount: T) -> Quat<T> {
        let dot = self.dot(other);

        let dot_threshold = Number::from(0.9995);

        if dot > dot_threshold {
            return self.nlerp(other, amount);                   // if quaternions are close together use `nlerp`
        } else {
            let robust_dot = dot.clamp(&-Number::from(1),
                                       &Number::from(1));       // stay within the domain of acos()

            let theta_0 = acos(&robust_dot);                    // the angle between the quaternions
            let theta = theta_0 * amount;                       // the fraction of theta specified by `amount`

            let q = other.sub_q(&self.mul_t(robust_dot))
                         .normalize();

            return self.mul_t(cos(&theta)).add_q(&q.mul_t(sin(&theta)));
        }
    }

    #[inline(always)]
    pure fn to_mat3(&self) -> Mat3<T> {
        let x2 = self.v.x + self.v.x;
        let y2 = self.v.y + self.v.y;
        let z2 = self.v.z + self.v.z;

        let xx2 = x2 * self.v.x;
        let xy2 = x2 * self.v.y;
        let xz2 = x2 * self.v.z;

        let yy2 = y2 * self.v.y;
        let yz2 = y2 * self.v.z;
        let zz2 = z2 * self.v.z;

        let sy2 = y2 * self.s;
        let sz2 = z2 * self.s;
        let sx2 = x2 * self.s;

        let _1: T = Number::from(1);

        Mat3::new(_1 - yy2 - zz2,      xy2 + sz2,      xz2 - sy2,
                       xy2 - sz2, _1 - xx2 - zz2,      yz2 + sx2,
                       xz2 + sy2,      yz2 - sx2, _1 - xx2 - yy2)
    }

    #[inline(always)]
    pure fn to_mat4(&self) -> Mat4<T> {
        self.to_mat3().to_mat4()
    }

    #[inline(always)]
    pure fn to_ptr(&self) -> *T {
        unsafe {
            transmute::<*Quat<T>, *T>(
                to_unsafe_ptr(self)
            )
        }
    }
}

pub impl<T:Copy Float> Quat<T>: Neg<Quat<T>> {
    #[inline(always)]
    pure fn neg(&self) -> Quat<T> {
        Quat::new(-self[0], -self[1], -self[2], -self[3])
    }
}

pub impl<T:Copy Eq> Quat<T>: Eq {
    #[inline(always)]
    pure fn eq(&self, other: &Quat<T>) -> bool {
        self[0] == other[0] &&
        self[1] == other[1] &&
        self[2] == other[2] &&
        self[3] == other[3]
    }

    #[inline(always)]
    pure fn ne(&self, other: &Quat<T>) -> bool {
        !(self == other)
    }
}

pub impl<T:Copy FuzzyEq> Quat<T>: FuzzyEq {
    #[inline(always)]
    pure fn fuzzy_eq(other: &Quat<T>) -> bool {
        self[0].fuzzy_eq(&other[0]) &&
        self[1].fuzzy_eq(&other[1]) &&
        self[2].fuzzy_eq(&other[2]) &&
        self[3].fuzzy_eq(&other[3])
    }
}
<|endoftext|>