<commit_before>/// Full runtime loom tests. These are heavy tests and take significant time to
/// run on CI.
///
/// Use `LOOM_MAX_PREEMPTIONS=1` to do a "quick" run as a smoke test.
///
/// In order to speed up the C
use crate::future::poll_fn;
use crate::runtime::tests::loom_oneshot as oneshot;
use crate::runtime::{self, Runtime};
use crate::{spawn, task};
use tokio_test::assert_ok;

use loom::sync::atomic::{AtomicBool, AtomicUsize};
use loom::sync::{Arc, Mutex};

use pin_project_lite::pin_project;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};
use std::task::{Context, Poll};

/// Tests are divided into groups to make the runs faster on CI.
mod group_a {
    use super::*;

    #[test]
    fn racy_shutdown() {
        loom::model(|| {
            let pool = mk_pool(1);

            // here's the case we want to exercise:
            //
            // a worker that still has tasks in its local queue gets sent to the blocking pool (due to
            // block_in_place). the blocking pool is shut down, so drops the worker. the worker's
            // shutdown method never gets run.
            //
            // we do this by spawning two tasks on one worker, the first of which does block_in_place,
            // and then immediately drop the pool.

            pool.spawn(track(async {
                crate::task::block_in_place(|| {});
            }));
            pool.spawn(track(async {}));
            drop(pool);
        });
    }

    #[test]
    fn pool_multi_spawn() {
        loom::model(|| {
            let pool = mk_pool(2);
            let c1 = Arc::new(AtomicUsize::new(0));

            let (tx, rx) = oneshot::channel();
            let tx1 = Arc::new(Mutex::new(Some(tx)));

            // Spawn a task
            let c2 = c1.clone();
            let tx2 = tx1.clone();
            pool.spawn(track(async move {
                spawn(track(async move {
                    if 1 == c1.fetch_add(1, Relaxed) {
                        tx1.lock().unwrap().take().unwrap().send(());
                    }
                }));
            }));

            // Spawn a second task
            pool.spawn(track(async move {
                spawn(track(async move {
                    if 1 == c2.fetch_add(1, Relaxed) {
                        tx2.lock().unwrap().take().unwrap().send(());
                    }
                }));
            }));

            rx.recv();
        });
    }

    fn only_blocking_inner(first_pending: bool) {
        loom::model(move || {
            let pool = mk_pool(1);
            let (block_tx, block_rx) = oneshot::channel();

            pool.spawn(track(async move {
                crate::task::block_in_place(move || {
                    block_tx.send(());
                });
                if first_pending {
                    task::yield_now().await
                }
            }));

            block_rx.recv();
            drop(pool);
        });
    }

    #[test]
    fn only_blocking_without_pending() {
        only_blocking_inner(false)
    }

    #[test]
    fn only_blocking_with_pending() {
        only_blocking_inner(true)
    }
}

mod group_b {
    use super::*;

    fn blocking_and_regular_inner(first_pending: bool) {
        const NUM: usize = 3;
        loom::model(move || {
            let pool = mk_pool(1);
            let cnt = Arc::new(AtomicUsize::new(0));

            let (block_tx, block_rx) = oneshot::channel();
            let (done_tx, done_rx) = oneshot::channel();
            let done_tx = Arc::new(Mutex::new(Some(done_tx)));

            pool.spawn(track(async move {
                crate::task::block_in_place(move || {
                    block_tx.send(());
                });
                if first_pending {
                    task::yield_now().await
                }
            }));

            for _ in 0..NUM {
                let cnt = cnt.clone();
                let done_tx = done_tx.clone();

                pool.spawn(track(async move {
                    if NUM == cnt.fetch_add(1, Relaxed) + 1 {
                        done_tx.lock().unwrap().take().unwrap().send(());
                    }
                }));
            }

            done_rx.recv();
            block_rx.recv();

            drop(pool);
        });
    }

    #[test]
    fn blocking_and_regular() {
        blocking_and_regular_inner(false);
    }

    #[test]
    fn blocking_and_regular_with_pending() {
        blocking_and_regular_inner(true);
    }

    #[test]
    fn pool_shutdown() {
        loom::model(|| {
            let pool = mk_pool(2);

            pool.spawn(track(async move {
                gated2(true).await;
            }));

            pool.spawn(track(async move {
                gated2(false).await;
            }));

            drop(pool);
        });
    }

    #[test]
    fn join_output() {
        loom::model(|| {
            let rt = mk_pool(1);

            rt.block_on(async {
                let t = crate::spawn(track(async { "hello" }));

                let out = assert_ok!(t.await);
                assert_eq!("hello", out.into_inner());
            });
        });
    }

    #[test]
    fn poll_drop_handle_then_drop() {
        loom::model(|| {
            let rt = mk_pool(1);

            rt.block_on(async move {
                let mut t = crate::spawn(track(async { "hello" }));

                poll_fn(|cx| {
                    let _ = Pin::new(&mut t).poll(cx);
                    Poll::Ready(())
                })
                .await;
            });
        })
    }

    #[test]
    fn complete_block_on_under_load() {
        loom::model(|| {
            let pool = mk_pool(1);

            pool.block_on(async {
                // Trigger a re-schedule
                crate::spawn(track(async {
                    for _ in 0..2 {
                        task::yield_now().await;
                    }
                }));

                gated2(true).await
            });
        });
    }
}

mod group_c {
    use super::*;

    #[test]
    fn shutdown_with_notification() {
        use crate::sync::oneshot;

        loom::model(|| {
            let rt = mk_pool(2);
            let (done_tx, done_rx) = oneshot::channel::<()>();

            rt.spawn(track(async move {
                let (tx, rx) = oneshot::channel::<()>();

                crate::spawn(async move {
                    crate::task::spawn_blocking(move || {
                        let _ = tx.send(());
                    });

                    let _ = done_rx.await;
                });

                let _ = rx.await;

                let _ = done_tx.send(());
            }));
        });
    }
}

mod group_d {
    use super::*;

    #[test]
    fn pool_multi_notify() {
        loom::model(|| {
            let pool = mk_pool(2);

            let c1 = Arc::new(AtomicUsize::new(0));

            let (done_tx, done_rx) = oneshot::channel();
            let done_tx1 = Arc::new(Mutex::new(Some(done_tx)));

            // Spawn a task
            let c2 = c1.clone();
            let done_tx2 = done_tx1.clone();
            pool.spawn(track(async move {
                gated().await;
                gated().await;

                if 1 == c1.fetch_add(1, Relaxed) {
                    done_tx1.lock().unwrap().take().unwrap().send(());
                }
            }));

            // Spawn a second task
            pool.spawn(track(async move {
                gated().await;
                gated().await;

                if 1 == c2.fetch_add(1, Relaxed) {
                    done_tx2.lock().unwrap().take().unwrap().send(());
                }
            }));

            done_rx.recv();
        });
    }
}

fn mk_pool(num_threads: usize) -> Runtime {
    runtime::Builder::new_multi_thread()
        .worker_threads(num_threads)
        .build()
        .unwrap()
}

fn gated() -> impl Future<Output = &'static str> {
    gated2(false)
}

fn gated2(thread: bool) -> impl Future<Output = &'static str> {
    use loom::thread;
    use std::sync::Arc;

    let gate = Arc::new(AtomicBool::new(false));
    let mut fired = false;

    poll_fn(move |cx| {
        if !fired {
            let gate = gate.clone();
            let waker = cx.waker().clone();

            if thread {
                thread::spawn(move || {
                    gate.store(true, SeqCst);
                    waker.wake_by_ref();
                });
            } else {
                spawn(track(async move {
                    gate.store(true, SeqCst);
                    waker.wake_by_ref();
                }));
            }

            fired = true;

            return Poll::Pending;
        }

        if gate.load(SeqCst) {
            Poll::Ready("hello world")
        } else {
            Poll::Pending
        }
    })
}

fn track<T: Future>(f: T) -> Track<T> {
    Track {
        inner: f,
        arc: Arc::new(()),
    }
}

pin_project! {
    struct Track<T> {
        #[pin]
        inner: T,
        // Arc is used to hook into loom's leak tracking.
        arc: Arc<()>,
    }
}

impl<T> Track<T> {
    fn into_inner(self) -> T {
        self.inner
    }
}

impl<T: Future> Future for Track<T> {
    type Output = Track<T::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.project();

        Poll::Ready(Track {
            inner: ready!(me.inner.poll(cx)),
            arc: me.arc.clone(),
        })
    }
}
<commit_msg>tests: simplify loom tests (#3995)<commit_after>/// Full runtime loom tests. These are heavy tests and take significant time to
/// run on CI.
///
/// Use `LOOM_MAX_PREEMPTIONS=1` to do a "quick" run as a smoke test.
///
/// In order to speed up the C
use crate::future::poll_fn;
use crate::runtime::tests::loom_oneshot as oneshot;
use crate::runtime::{self, Runtime};
use crate::{spawn, task};
use tokio_test::assert_ok;

use loom::sync::atomic::{AtomicBool, AtomicUsize};
use loom::sync::Arc;

use pin_project_lite::pin_project;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};
use std::task::{Context, Poll};

mod atomic_take {
    use loom::sync::atomic::AtomicBool;
    use std::mem::MaybeUninit;
    use std::sync::atomic::Ordering::SeqCst;

    pub(super) struct AtomicTake<T> {
        inner: MaybeUninit<T>,
        taken: AtomicBool,
    }

    impl<T> AtomicTake<T> {
        pub(super) fn new(value: T) -> Self {
            Self {
                inner: MaybeUninit::new(value),
                taken: AtomicBool::new(false),
            }
        }

        pub(super) fn take(&self) -> Option<T> {
            // safety: Only one thread will see the boolean change from false
            // to true, so that thread is able to take the value.
            match self.taken.fetch_or(true, SeqCst) {
                false => unsafe { Some(std::ptr::read(self.inner.as_ptr())) },
                true => None,
            }
        }
    }

    impl<T> Drop for AtomicTake<T> {
        fn drop(&mut self) {
            drop(self.take());
        }
    }
}

#[derive(Clone)]
struct AtomicOneshot<T> {
    value: std::sync::Arc<atomic_take::AtomicTake<oneshot::Sender<T>>>,
}
impl<T> AtomicOneshot<T> {
    fn new(sender: oneshot::Sender<T>) -> Self {
        Self {
            value: std::sync::Arc::new(atomic_take::AtomicTake::new(sender)),
        }
    }

    fn assert_send(&self, value: T) {
        self.value.take().unwrap().send(value);
    }
}

/// Tests are divided into groups to make the runs faster on CI.
mod group_a {
    use super::*;

    #[test]
    fn racy_shutdown() {
        loom::model(|| {
            let pool = mk_pool(1);

            // here's the case we want to exercise:
            //
            // a worker that still has tasks in its local queue gets sent to the blocking pool (due to
            // block_in_place). the blocking pool is shut down, so drops the worker. the worker's
            // shutdown method never gets run.
            //
            // we do this by spawning two tasks on one worker, the first of which does block_in_place,
            // and then immediately drop the pool.

            pool.spawn(track(async {
                crate::task::block_in_place(|| {});
            }));
            pool.spawn(track(async {}));
            drop(pool);
        });
    }

    #[test]
    fn pool_multi_spawn() {
        loom::model(|| {
            let pool = mk_pool(2);
            let c1 = Arc::new(AtomicUsize::new(0));

            let (tx, rx) = oneshot::channel();
            let tx1 = AtomicOneshot::new(tx);

            // Spawn a task
            let c2 = c1.clone();
            let tx2 = tx1.clone();
            pool.spawn(track(async move {
                spawn(track(async move {
                    if 1 == c1.fetch_add(1, Relaxed) {
                        tx1.assert_send(());
                    }
                }));
            }));

            // Spawn a second task
            pool.spawn(track(async move {
                spawn(track(async move {
                    if 1 == c2.fetch_add(1, Relaxed) {
                        tx2.assert_send(());
                    }
                }));
            }));

            rx.recv();
        });
    }

    fn only_blocking_inner(first_pending: bool) {
        loom::model(move || {
            let pool = mk_pool(1);
            let (block_tx, block_rx) = oneshot::channel();

            pool.spawn(track(async move {
                crate::task::block_in_place(move || {
                    block_tx.send(());
                });
                if first_pending {
                    task::yield_now().await
                }
            }));

            block_rx.recv();
            drop(pool);
        });
    }

    #[test]
    fn only_blocking_without_pending() {
        only_blocking_inner(false)
    }

    #[test]
    fn only_blocking_with_pending() {
        only_blocking_inner(true)
    }
}

mod group_b {
    use super::*;

    fn blocking_and_regular_inner(first_pending: bool) {
        const NUM: usize = 3;
        loom::model(move || {
            let pool = mk_pool(1);
            let cnt = Arc::new(AtomicUsize::new(0));

            let (block_tx, block_rx) = oneshot::channel();
            let (done_tx, done_rx) = oneshot::channel();
            let done_tx = AtomicOneshot::new(done_tx);

            pool.spawn(track(async move {
                crate::task::block_in_place(move || {
                    block_tx.send(());
                });
                if first_pending {
                    task::yield_now().await
                }
            }));

            for _ in 0..NUM {
                let cnt = cnt.clone();
                let done_tx = done_tx.clone();

                pool.spawn(track(async move {
                    if NUM == cnt.fetch_add(1, Relaxed) + 1 {
                        done_tx.assert_send(());
                    }
                }));
            }

            done_rx.recv();
            block_rx.recv();

            drop(pool);
        });
    }

    #[test]
    fn blocking_and_regular() {
        blocking_and_regular_inner(false);
    }

    #[test]
    fn blocking_and_regular_with_pending() {
        blocking_and_regular_inner(true);
    }

    #[test]
    fn pool_shutdown() {
        loom::model(|| {
            let pool = mk_pool(2);

            pool.spawn(track(async move {
                gated2(true).await;
            }));

            pool.spawn(track(async move {
                gated2(false).await;
            }));

            drop(pool);
        });
    }

    #[test]
    fn join_output() {
        loom::model(|| {
            let rt = mk_pool(1);

            rt.block_on(async {
                let t = crate::spawn(track(async { "hello" }));

                let out = assert_ok!(t.await);
                assert_eq!("hello", out.into_inner());
            });
        });
    }

    #[test]
    fn poll_drop_handle_then_drop() {
        loom::model(|| {
            let rt = mk_pool(1);

            rt.block_on(async move {
                let mut t = crate::spawn(track(async { "hello" }));

                poll_fn(|cx| {
                    let _ = Pin::new(&mut t).poll(cx);
                    Poll::Ready(())
                })
                .await;
            });
        })
    }

    #[test]
    fn complete_block_on_under_load() {
        loom::model(|| {
            let pool = mk_pool(1);

            pool.block_on(async {
                // Trigger a re-schedule
                crate::spawn(track(async {
                    for _ in 0..2 {
                        task::yield_now().await;
                    }
                }));

                gated2(true).await
            });
        });
    }
}

mod group_c {
    use super::*;

    #[test]
    fn shutdown_with_notification() {
        use crate::sync::oneshot;

        loom::model(|| {
            let rt = mk_pool(2);
            let (done_tx, done_rx) = oneshot::channel::<()>();

            rt.spawn(track(async move {
                let (tx, rx) = oneshot::channel::<()>();

                crate::spawn(async move {
                    crate::task::spawn_blocking(move || {
                        let _ = tx.send(());
                    });

                    let _ = done_rx.await;
                });

                let _ = rx.await;

                let _ = done_tx.send(());
            }));
        });
    }
}

mod group_d {
    use super::*;

    #[test]
    fn pool_multi_notify() {
        loom::model(|| {
            let pool = mk_pool(2);

            let c1 = Arc::new(AtomicUsize::new(0));

            let (done_tx, done_rx) = oneshot::channel();
            let done_tx1 = AtomicOneshot::new(done_tx);
            let done_tx2 = done_tx1.clone();

            // Spawn a task
            let c2 = c1.clone();
            pool.spawn(track(async move {
                gated().await;
                gated().await;

                if 1 == c1.fetch_add(1, Relaxed) {
                    done_tx1.assert_send(());
                }
            }));

            // Spawn a second task
            pool.spawn(track(async move {
                gated().await;
                gated().await;

                if 1 == c2.fetch_add(1, Relaxed) {
                    done_tx2.assert_send(());
                }
            }));

            done_rx.recv();
        });
    }
}

fn mk_pool(num_threads: usize) -> Runtime {
    runtime::Builder::new_multi_thread()
        .worker_threads(num_threads)
        .build()
        .unwrap()
}

fn gated() -> impl Future<Output = &'static str> {
    gated2(false)
}

fn gated2(thread: bool) -> impl Future<Output = &'static str> {
    use loom::thread;
    use std::sync::Arc;

    let gate = Arc::new(AtomicBool::new(false));
    let mut fired = false;

    poll_fn(move |cx| {
        if !fired {
            let gate = gate.clone();
            let waker = cx.waker().clone();

            if thread {
                thread::spawn(move || {
                    gate.store(true, SeqCst);
                    waker.wake_by_ref();
                });
            } else {
                spawn(track(async move {
                    gate.store(true, SeqCst);
                    waker.wake_by_ref();
                }));
            }

            fired = true;

            return Poll::Pending;
        }

        if gate.load(SeqCst) {
            Poll::Ready("hello world")
        } else {
            Poll::Pending
        }
    })
}

fn track<T: Future>(f: T) -> Track<T> {
    Track {
        inner: f,
        arc: Arc::new(()),
    }
}

pin_project! {
    struct Track<T> {
        #[pin]
        inner: T,
        // Arc is used to hook into loom's leak tracking.
        arc: Arc<()>,
    }
}

impl<T> Track<T> {
    fn into_inner(self) -> T {
        self.inner
    }
}

impl<T: Future> Future for Track<T> {
    type Output = Track<T::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.project();

        Poll::Ready(Track {
            inner: ready!(me.inner.poll(cx)),
            arc: me.arc.clone(),
        })
    }
}
<|endoftext|>


<commit_before>use routing;
use rpc;
use bus;

use rpc::Rpc;
use time;
use {SubotaiError, SubotaiResult};

use hash::Hash;
use std::net;
use std::sync::Arc;
use std::sync;
use node::*;

/// Node resources for synchronous operations.
///
/// All methods on this module are synchronous, and will wait for any
/// remote nodes queried to reply to the RPCs sent, up to the timeout
/// defined at `node::NETWORK_TIMEOUT_S`. The node layer above is in
/// charge of parallelizing those operations by spawning threads when
/// adequate.
pub struct Resources {
   pub id       : Hash,
   pub table    : routing::Table,
   pub outbound : net::UdpSocket,
   pub inbound  : net::UdpSocket,
   pub state    : sync::Mutex<State>,
   pub updates  : sync::Mutex<bus::Bus<Update>>,
}

#[derive(Clone)]
pub enum Update {
   RpcReceived(Rpc),
   Tick,
}

impl Resources {
   pub fn local_info(&self) -> routing::NodeInfo {
      routing::NodeInfo {
         id      : self.id.clone(),
         address : self.inbound.local_addr().unwrap(),
      }
   }

   pub fn ping(&self, id: Hash) -> SubotaiResult<()> {
      let node = match self.table.specific_node(&id) {
         None => Some(try!(self.find_node(&id))),
         Some(node) => Some(node),
      };

      if let Some(node) = node {
         let rpc = Rpc::ping(self.id.clone(), self.inbound.local_addr().unwrap().port());
         let packet = rpc.serialize();
         let responses = self.receptions().during(time::Duration::seconds(NETWORK_TIMEOUT_S))
            .rpc(receptions::RpcFilter::PingResponse).from(id.clone()).take(1);

         try!(self.outbound.send_to(&packet, node.address));

         for _ in responses {
            return Ok(());
         }
      }
      Err(SubotaiError::NoResponse)
   }

   /// Sends a ping and doesn't wait for a response. Used by the maintenance thread.
   pub fn ping_and_forget(&self, id: Hash) -> SubotaiResult<()> {
      let node = match self.table.specific_node(&id) {
         None => Some(try!(self.find_node(&id))),
         Some(node) => Some(node),
      };

      if let Some(node) = node {
         let rpc = Rpc::ping(self.id.clone(), self.inbound.local_addr().unwrap().port());
         let packet = rpc.serialize();
         try!(self.outbound.send_to(&packet, node.address));
      }
      Ok(())
   }

   pub fn receptions(&self) -> receptions::Receptions {
      receptions::Receptions::new(self)
   }

   /// Attempts to find a node through the network.
   pub fn find_node(&self, id_to_find: &Hash) -> SubotaiResult<routing::NodeInfo> {
      // If the node is already present in our table, we are done early.
      if let Some(node) = self.table.specific_node(id_to_find) {
         return Ok(node);
      }

      let mut queried_ids = Vec::<Hash>::with_capacity(routing::K);
      let all_receptions = self.receptions();

      while queried_ids.len() < routing::K {
         // We query the 'ALPHA' nodes closest to the target we haven't yet queried.
         let nodes_to_query: Vec<routing::NodeInfo> = self.table.closest_nodes_to(id_to_find)
            .filter(|ref info| !queried_ids.contains(&info.id))
            .take(routing::ALPHA)
            .collect();

         // We wait for the response from the same number of nodes, minus the 'IMPATIENCE' factor.
         let responses = self.receptions()
            .during(time::Duration::seconds(NETWORK_TIMEOUT_S))
            .filter(|ref rpc| rpc.is_finding_node(id_to_find))
            .take(usize::saturating_sub(nodes_to_query.len(), routing::IMPATIENCE));

         // We compose the RPCs and send the UDP packets.
         try!(self.lookup_wave(id_to_find, &nodes_to_query, &mut queried_ids));

         for response in responses {
            println!("<-- Response from {}", response.sender_id);
            if response.found_node(id_to_find) {
               return self.table.specific_node(id_to_find).ok_or(SubotaiError::NodeNotFound);
            }
         }
      }

      // One last wait until success or timeout, to compensate for impatience
      // (It could be that the nodes we ignored earlier come back with the response)
      all_receptions.during(time::Duration::seconds(NETWORK_TIMEOUT_S))
         .filter(|ref rpc| rpc.found_node(id_to_find))
         .take(1)
         .count();
      self.table.specific_node(id_to_find).ok_or(SubotaiError::NodeNotFound)
   }

   pub fn bootstrap(&self, seed: routing::NodeInfo, network_size: Option<usize>) -> SubotaiResult<()>  {
      self.table.insert_node(seed);

      // Timeout for the entire operation.
      let total_timeout = time::Duration::seconds(2 * NETWORK_TIMEOUT_S);
      let deadline = time::SteadyTime::now() + total_timeout;
      let mut responses = self.receptions()
         .rpc(receptions::RpcFilter::BootstrapResponse)
         .during(total_timeout);

      let expected_length = match network_size {
         Some(size) => size,
         None => routing::K,
      };

      let mut queried_ids = Vec::<Hash>::with_capacity(routing::K);
      while queried_ids.len() < expected_length {
         let nodes_to_query: Vec<routing::NodeInfo> = self.table.all_nodes()
            .filter(|ref node_info| !queried_ids.contains(&node_info.id))
            .collect();

         try!(self.bootstrap_wave(&nodes_to_query, &mut queried_ids));
         responses.next();

         if time::SteadyTime::now() >= deadline {
            return Err(SubotaiError::UnresponsiveNetwork);
         }
      }
      Ok(())
   }

   fn bootstrap_wave(&self, nodes_to_query: &[routing::NodeInfo], queried: &mut Vec<Hash>) -> SubotaiResult<()> {
      let rpc = Rpc::bootstrap(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      for node in nodes_to_query {
         try!(self.outbound.send_to(&packet, node.address));
         queried.push(node.id.clone());
      }
      Ok(())
   }

   fn lookup_wave(&self, id_to_find: &Hash, nodes_to_query: &[routing::NodeInfo], queried: &mut Vec<Hash>) -> SubotaiResult<()> {
      let rpc = Rpc::find_node(
         self.id.clone(),
         self.inbound.local_addr().unwrap().port(),
         id_to_find.clone(),
         routing::ALPHA,
      );
      let packet = rpc.serialize();
      for node in nodes_to_query {
         println!("--> Sending to {}", node.id);
         try!(self.outbound.send_to(&packet, node.address));
         queried.push(node.id.clone());
      }

      Ok(())
   }

   pub fn process_incoming_rpc(&self, rpc: Rpc, mut source: net::SocketAddr) -> SubotaiResult<()>{
      source.set_port(rpc.reply_port);
      let sender = routing::NodeInfo {
         id      : rpc.sender_id.clone(),
         address : source,
      };

      let result = match rpc.kind {
         rpc::Kind::Ping                           => self.handle_ping(sender),
         rpc::Kind::PingResponse                   => self.handle_ping_response(sender),
         rpc::Kind::FindNode(ref payload)          => self.handle_find_node(payload.clone(), sender),
         rpc::Kind::FindNodeResponse(ref payload)  => self.handle_find_node_response(payload.clone(), sender),
         rpc::Kind::Bootstrap                      => self.handle_bootstrap(sender),
         rpc::Kind::BootstrapResponse(ref payload) => self.handle_bootstrap_response(payload.clone(),sender),
         _ => unimplemented!(),
      };

      self.updates.lock().unwrap().broadcast(Update::RpcReceived(rpc));
      result
   }

   fn handle_ping(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let rpc = Rpc::ping_response(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_bootstrap(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let closest_to_sender: Vec<_> = self.table.closest_nodes_to(&sender.id)
         .filter(|ref info| &info.id != &sender.id) // We don't want to reply with the sender itself
         .take(routing::K)
         .collect();

      let rpc = Rpc::bootstrap_response(self.id.clone(), self.inbound.local_addr().unwrap().port(), closest_to_sender);
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_bootstrap_response(&self, payload: Arc<rpc::BootstrapResponsePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      for node in &payload.nodes {
         self.table.insert_node(node.clone());
      }
      Ok(())
   }

   fn handle_ping_response(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender);
      Ok(())
   }

   fn handle_find_node(&self, payload: Arc<rpc::FindNodePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let lookup_results = self.table.lookup(&payload.id_to_find, payload.nodes_wanted, None);
      let rpc = Rpc::find_node_response(self.id.clone(),
                                        self.inbound.local_addr().unwrap().port(),
                                        payload.id_to_find.clone(),
                                        lookup_results);
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_find_node_response(&self, payload: Arc<rpc::FindNodeResponsePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender);
      match payload.result {
         routing::LookupResult::ClosestNodes(ref nodes) => for node in nodes { self.table.insert_node(node.clone()) },
         routing::LookupResult::Found(ref node) => self.table.insert_node(node.clone()),
         _ => (),
      }
      Ok(())
   }
}
<commit_msg>Simplified ping implementation<commit_after>use routing;
use rpc;
use bus;

use rpc::Rpc;
use time;
use {SubotaiError, SubotaiResult};

use hash::Hash;
use std::net;
use std::sync::Arc;
use std::sync;
use node::*;

/// Node resources for synchronous operations.
///
/// All methods on this module are synchronous, and will wait for any
/// remote nodes queried to reply to the RPCs sent, up to the timeout
/// defined at `node::NETWORK_TIMEOUT_S`. The node layer above is in
/// charge of parallelizing those operations by spawning threads when
/// adequate.
pub struct Resources {
   pub id       : Hash,
   pub table    : routing::Table,
   pub outbound : net::UdpSocket,
   pub inbound  : net::UdpSocket,
   pub state    : sync::Mutex<State>,
   pub updates  : sync::Mutex<bus::Bus<Update>>,
}

#[derive(Clone)]
pub enum Update {
   RpcReceived(Rpc),
   Tick,
}

impl Resources {
   pub fn local_info(&self) -> routing::NodeInfo {
      routing::NodeInfo {
         id      : self.id.clone(),
         address : self.inbound.local_addr().unwrap(),
      }
   }

   pub fn ping(&self, id: Hash) -> SubotaiResult<()> {
      let node = match self.table.specific_node(&id) {
         None => try!(self.find_node(&id)),
         Some(node) => node,
      };

      let rpc = Rpc::ping(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      let responses = self.receptions().during(time::Duration::seconds(NETWORK_TIMEOUT_S))
         .rpc(receptions::RpcFilter::PingResponse).from(id.clone()).take(1);

      try!(self.outbound.send_to(&packet, node.address));

      match responses.count() {
         1 => Ok(()),
         _ => Err(SubotaiError::NoResponse),
      }
   }

   /// Sends a ping and doesn't wait for a response. Used by the maintenance thread.
   pub fn ping_and_forget(&self, id: Hash) -> SubotaiResult<()> {
      let node = match self.table.specific_node(&id) {
         None => try!(self.find_node(&id)),
         Some(node) => node,
      };

      let rpc = Rpc::ping(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, node.address));

      Ok(())
   }

   pub fn receptions(&self) -> receptions::Receptions {
      receptions::Receptions::new(self)
   }

   /// Attempts to find a node through the network.
   pub fn find_node(&self, id_to_find: &Hash) -> SubotaiResult<routing::NodeInfo> {
      // If the node is already present in our table, we are done early.
      if let Some(node) = self.table.specific_node(id_to_find) {
         return Ok(node);
      }

      let mut queried_ids = Vec::<Hash>::with_capacity(routing::K);
      let all_receptions = self.receptions();

      while queried_ids.len() < routing::K {
         // We query the 'ALPHA' nodes closest to the target we haven't yet queried.
         let nodes_to_query: Vec<routing::NodeInfo> = self.table.closest_nodes_to(id_to_find)
            .filter(|ref info| !queried_ids.contains(&info.id))
            .take(routing::ALPHA)
            .collect();

         // We wait for the response from the same number of nodes, minus the 'IMPATIENCE' factor.
         let responses = self.receptions()
            .during(time::Duration::seconds(NETWORK_TIMEOUT_S))
            .filter(|ref rpc| rpc.is_finding_node(id_to_find))
            .take(usize::saturating_sub(nodes_to_query.len(), routing::IMPATIENCE));

         // We compose the RPCs and send the UDP packets.
         try!(self.lookup_wave(id_to_find, &nodes_to_query, &mut queried_ids));

         for response in responses {
            println!("<-- Response from {}", response.sender_id);
            if response.found_node(id_to_find) {
               return self.table.specific_node(id_to_find).ok_or(SubotaiError::NodeNotFound);
            }
         }
      }

      // One last wait until success or timeout, to compensate for impatience
      // (It could be that the nodes we ignored earlier come back with the response)
      all_receptions.during(time::Duration::seconds(NETWORK_TIMEOUT_S))
         .filter(|ref rpc| rpc.found_node(id_to_find))
         .take(1)
         .count();
      self.table.specific_node(id_to_find).ok_or(SubotaiError::NodeNotFound)
   }

   pub fn bootstrap(&self, seed: routing::NodeInfo, network_size: Option<usize>) -> SubotaiResult<()>  {
      self.table.insert_node(seed);

      // Timeout for the entire operation.
      let total_timeout = time::Duration::seconds(2 * NETWORK_TIMEOUT_S);
      let deadline = time::SteadyTime::now() + total_timeout;
      let mut responses = self.receptions()
         .rpc(receptions::RpcFilter::BootstrapResponse)
         .during(total_timeout);

      let expected_length = match network_size {
         Some(size) => size,
         None => routing::K,
      };

      let mut queried_ids = Vec::<Hash>::with_capacity(routing::K);
      while queried_ids.len() < expected_length {
         let nodes_to_query: Vec<routing::NodeInfo> = self.table.all_nodes()
            .filter(|ref node_info| !queried_ids.contains(&node_info.id))
            .collect();

         try!(self.bootstrap_wave(&nodes_to_query, &mut queried_ids));
         responses.next();

         if time::SteadyTime::now() >= deadline {
            return Err(SubotaiError::UnresponsiveNetwork);
         }
      }
      Ok(())
   }

   fn bootstrap_wave(&self, nodes_to_query: &[routing::NodeInfo], queried: &mut Vec<Hash>) -> SubotaiResult<()> {
      let rpc = Rpc::bootstrap(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      for node in nodes_to_query {
         try!(self.outbound.send_to(&packet, node.address));
         queried.push(node.id.clone());
      }
      Ok(())
   }

   fn lookup_wave(&self, id_to_find: &Hash, nodes_to_query: &[routing::NodeInfo], queried: &mut Vec<Hash>) -> SubotaiResult<()> {
      let rpc = Rpc::find_node(
         self.id.clone(),
         self.inbound.local_addr().unwrap().port(),
         id_to_find.clone(),
         routing::ALPHA,
      );
      let packet = rpc.serialize();
      for node in nodes_to_query {
         println!("--> Sending to {}", node.id);
         try!(self.outbound.send_to(&packet, node.address));
         queried.push(node.id.clone());
      }

      Ok(())
   }

   pub fn process_incoming_rpc(&self, rpc: Rpc, mut source: net::SocketAddr) -> SubotaiResult<()>{
      source.set_port(rpc.reply_port);
      let sender = routing::NodeInfo {
         id      : rpc.sender_id.clone(),
         address : source,
      };

      let result = match rpc.kind {
         rpc::Kind::Ping                           => self.handle_ping(sender),
         rpc::Kind::PingResponse                   => self.handle_ping_response(sender),
         rpc::Kind::FindNode(ref payload)          => self.handle_find_node(payload.clone(), sender),
         rpc::Kind::FindNodeResponse(ref payload)  => self.handle_find_node_response(payload.clone(), sender),
         rpc::Kind::Bootstrap                      => self.handle_bootstrap(sender),
         rpc::Kind::BootstrapResponse(ref payload) => self.handle_bootstrap_response(payload.clone(),sender),
         _ => unimplemented!(),
      };

      self.updates.lock().unwrap().broadcast(Update::RpcReceived(rpc));
      result
   }

   fn handle_ping(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let rpc = Rpc::ping_response(self.id.clone(), self.inbound.local_addr().unwrap().port());
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_bootstrap(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let closest_to_sender: Vec<_> = self.table.closest_nodes_to(&sender.id)
         .filter(|ref info| &info.id != &sender.id) // We don't want to reply with the sender itself
         .take(routing::K)
         .collect();

      let rpc = Rpc::bootstrap_response(self.id.clone(), self.inbound.local_addr().unwrap().port(), closest_to_sender);
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_bootstrap_response(&self, payload: Arc<rpc::BootstrapResponsePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      for node in &payload.nodes {
         self.table.insert_node(node.clone());
      }
      Ok(())
   }

   fn handle_ping_response(&self, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender);
      Ok(())
   }

   fn handle_find_node(&self, payload: Arc<rpc::FindNodePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender.clone());
      let lookup_results = self.table.lookup(&payload.id_to_find, payload.nodes_wanted, None);
      let rpc = Rpc::find_node_response(self.id.clone(),
                                        self.inbound.local_addr().unwrap().port(),
                                        payload.id_to_find.clone(),
                                        lookup_results);
      let packet = rpc.serialize();
      try!(self.outbound.send_to(&packet, sender.address));
      Ok(())
   }

   fn handle_find_node_response(&self, payload: Arc<rpc::FindNodeResponsePayload>, sender: routing::NodeInfo) -> SubotaiResult<()> {
      self.table.insert_node(sender);
      match payload.result {
         routing::LookupResult::ClosestNodes(ref nodes) => for node in nodes { self.table.insert_node(node.clone()) },
         routing::LookupResult::Found(ref node) => self.table.insert_node(node.clone()),
         _ => (),
      }
      Ok(())
   }
}
<|endoftext|>


<commit_before>use ast::*;

pub trait PrettyPrint {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> ();
    fn pretty_print(&self, identifier_map: &IdentifierMap) -> String {
        let mut pretty_printed = String::new();
        self.pretty_print_into(identifier_map, &mut pretty_printed);
        pretty_printed
    }
}

impl PrettyPrint for Expr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Expr::This(_, _) => pretty_printed.push_str("this"),
            Expr::Literal(ref l) => l.pretty_print_into(identifier_map, pretty_printed),
            Expr::Unary(ref u) => u.pretty_print_into(identifier_map, pretty_printed),
            Expr::Binary(ref b) => b.pretty_print_into(identifier_map, pretty_printed),
            Expr::Logic(ref b) => b.pretty_print_into(identifier_map, pretty_printed),
            Expr::Grouping(ref g) => g.pretty_print_into(identifier_map, pretty_printed),
            Expr::Identifier(ref _h, ref i) => i.pretty_print_into(identifier_map, pretty_printed),
            Expr::Assignment(ref a) => a.pretty_print_into(identifier_map, pretty_printed),
            Expr::Call(ref c) => c.pretty_print_into(identifier_map, pretty_printed),
            Expr::Get(ref g) => g.pretty_print_into(identifier_map, pretty_printed),
            Expr::Set(ref s) => s.pretty_print_into(identifier_map, pretty_printed),
        }
    }
}

impl PrettyPrint for UnaryOperator {
    fn pretty_print_into(&self,
                         _identifier_map: &IdentifierMap,
                         pretty_printed: &mut String)
                         -> () {
        match *self {
            UnaryOperator::Bang => pretty_printed.push_str("!"),
            UnaryOperator::Minus => pretty_printed.push_str("-"),
        }
    }
}

impl PrettyPrint for BinaryOperator {
    fn pretty_print_into(&self,
                         _identifier_map: &IdentifierMap,
                         pretty_printed: &mut String)
                         -> () {
        match *self {
            BinaryOperator::Minus => pretty_printed.push_str("-"),
            BinaryOperator::Plus => pretty_printed.push_str("+"),
            BinaryOperator::Slash => pretty_printed.push_str("/"),
            BinaryOperator::Star => pretty_printed.push_str("*"),
            BinaryOperator::Equal => pretty_printed.push_str("=="),
            BinaryOperator::NotEqual => pretty_printed.push_str("!="),
            BinaryOperator::Less => pretty_printed.push_str("<"),
            BinaryOperator::LessEqual => pretty_printed.push_str("<="),
            BinaryOperator::Greater => pretty_printed.push_str(">"),
            BinaryOperator::GreaterEqual => pretty_printed.push_str(">="),
        }
    }
}

impl PrettyPrint for LogicOperator {
    fn pretty_print_into(&self,
                         _identifier_map: &IdentifierMap,
                         pretty_printed: &mut String)
                         -> () {
        match *self {
            LogicOperator::Or => pretty_printed.push_str("or"),
            LogicOperator::And => pretty_printed.push_str("and"),
        }
    }
}

impl PrettyPrint for Literal {
    fn pretty_print_into(&self,
                         _identifier_map: &IdentifierMap,
                         pretty_printed: &mut String)
                         -> () {
        match *self {
            Literal::NilLiteral => pretty_printed.push_str("null"),
            Literal::BoolLiteral(ref b) => pretty_printed.push_str(&b.to_string()),
            Literal::StringLiteral(ref s) => pretty_printed.push_str(s),
            Literal::NumberLiteral(n) => pretty_printed.push_str(&n.to_string()),
        }
    }
}

impl PrettyPrint for Identifier {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str(identifier_map.lookup(self).unwrap())
    }
}

impl PrettyPrint for Target {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Target::Identifier(ref i) => i.pretty_print_into(identifier_map, pretty_printed),
        }
    }
}

impl PrettyPrint for Assignment {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.lvalue
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" = ");
        self.rvalue
            .pretty_print_into(identifier_map, pretty_printed);
    }
}

impl PrettyPrint for Grouping {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(group ");
        self.expr.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for UnaryExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for LogicExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.left.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for BinaryExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.left.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for Statement {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Statement::Print(ref e) => {
                pretty_printed.push_str("print ");
                e.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::Return(ref e) => {
                pretty_printed.push_str("return");
                if let &Some(ref e) = e {
                    pretty_printed.push_str(" ");
                    e.pretty_print_into(identifier_map, pretty_printed);
                }
                pretty_printed.push_str(";");
            }
            Statement::Expression(ref e) => {
                e.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::VariableDefinition(ref identifier) => {
                pretty_printed.push_str("var ");
                identifier.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::VariableDefinitionWithInitalizer(ref identifier, ref initializer) => {
                pretty_printed.push_str("var ");
                identifier.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" = ");
                initializer.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::Block(ref b) => {
                pretty_printed.push_str("{ ");
                for statement in &b.statements {
                    statement.pretty_print_into(identifier_map, pretty_printed);
                    pretty_printed.push_str(" ");
                }
                pretty_printed.push_str("}");
            }
            Statement::IfThen(ref c) => {
                pretty_printed.push_str("if ( ");
                c.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                c.then_branch
                    .pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::IfThenElse(ref c) => {
                pretty_printed.push_str("if ( ");
                c.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                c.then_branch
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" else ");
                c.else_branch
                    .pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::While(ref l) => {
                pretty_printed.push_str("while ( ");
                l.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                l.body.pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::FunctionDefinition(ref f) => {
                if let FunctionKind::Function = f.kind {
                    pretty_printed.push_str("fun ");
                }
                f.name.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" (");
                for argument in &f.arguments {
                    argument.pretty_print_into(identifier_map, pretty_printed);
                    pretty_printed.push_str(" ");
                }
                pretty_printed.push_str(") ");
                f.body.pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::Class(ref c) => {
                pretty_printed.push_str("class ");
                c.name.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" {");
                for method in &c.methods {
                    pretty_printed.push_str(" ");
                    Statement::FunctionDefinition(method.clone())
                        .pretty_print_into(identifier_map, pretty_printed);
                }
                pretty_printed.push_str(" }");
            }
        };
    }
}

impl PrettyPrint for Call {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.callee
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str("( ");
        for arg in self.arguments.iter() {
            arg.pretty_print_into(identifier_map, pretty_printed);
            pretty_printed.push_str(" ");
        }
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for Get {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.instance
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(".");
        self.property
            .pretty_print_into(identifier_map, pretty_printed);
    }
}

impl PrettyPrint for Set {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.instance
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(".");
        self.property
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" = ");
        self.value.pretty_print_into(identifier_map, pretty_printed);
    }
}

#[cfg(test)]
mod tests {
    use ast::*;
    use pretty_printer::PrettyPrint;

    #[test]
    fn literal() {
        let identifier_map = IdentifierMap::new();
        let string = String::from("abc");
        let expr = Expr::Literal(Literal::StringLiteral(string.clone()));
        assert_eq!(string, expr.pretty_print(&identifier_map));
    }

    #[test]
    fn complex_expression() {
        let identifier_map = IdentifierMap::new();
        let subexpr1 = UnaryExpr {
            operator: UnaryOperator::Minus,
            right: Expr::Literal(Literal::NumberLiteral(123f64)),
        };
        let subexpr2 = Grouping { expr: Expr::Literal(Literal::NumberLiteral(45.67f64)) };
        let binary_expr = BinaryExpr {
            left: Expr::Unary(Box::new(subexpr1)),
            operator: BinaryOperator::Star,
            right: Expr::Grouping(Box::new(subexpr2)),
        };
        let expr = Expr::Binary(Box::new(binary_expr));
        assert_eq!("(* (- 123) (group 45.67))", &expr.pretty_print(&identifier_map));
    }

    #[test]
    fn block() {
        let mut identifier_map = IdentifierMap::new();
        let mut handle_factory = VariableUseHandleFactory::new();
        let identifier = identifier_map.from_name(&"x");
        let statements = vec![
            Statement::VariableDefinitionWithInitalizer(
                identifier.clone(),
                Expr::Literal(Literal::BoolLiteral(true))
                ),
            Statement::Print(
                Expr::Identifier(handle_factory.next(), identifier.clone()))
                    ];
        let block = Statement::Block(Box::new(Block { statements: statements }));
        assert_eq!("{ var x = true; print x; }", &block.pretty_print(&identifier_map));
    }
}
<commit_msg>pretty print super expression<commit_after>use ast::*;

pub trait PrettyPrint {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> ();
    fn pretty_print(&self, identifier_map: &IdentifierMap) -> String {
        let mut pretty_printed = String::new();
        self.pretty_print_into(identifier_map, &mut pretty_printed);
        pretty_printed
    }
}

impl PrettyPrint for Expr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Expr::This(_, _) => pretty_printed.push_str("this"),
            Expr::Super(_, _, ref identifier) => {
                pretty_printed.push_str("super.");
                identifier.pretty_print_into(identifier_map, pretty_printed)
            }
            Expr::Literal(ref l) => l.pretty_print_into(identifier_map, pretty_printed),
            Expr::Unary(ref u) => u.pretty_print_into(identifier_map, pretty_printed),
            Expr::Binary(ref b) => b.pretty_print_into(identifier_map, pretty_printed),
            Expr::Logic(ref b) => b.pretty_print_into(identifier_map, pretty_printed),
            Expr::Grouping(ref g) => g.pretty_print_into(identifier_map, pretty_printed),
            Expr::Identifier(ref _h, ref i) => i.pretty_print_into(identifier_map, pretty_printed),
            Expr::Assignment(ref a) => a.pretty_print_into(identifier_map, pretty_printed),
            Expr::Call(ref c) => c.pretty_print_into(identifier_map, pretty_printed),
            Expr::Get(ref g) => g.pretty_print_into(identifier_map, pretty_printed),
            Expr::Set(ref s) => s.pretty_print_into(identifier_map, pretty_printed),
        }
    }
}

impl PrettyPrint for UnaryOperator {
    fn pretty_print_into(
        &self,
        _identifier_map: &IdentifierMap,
        pretty_printed: &mut String,
    ) -> () {
        match *self {
            UnaryOperator::Bang => pretty_printed.push_str("!"),
            UnaryOperator::Minus => pretty_printed.push_str("-"),
        }
    }
}

impl PrettyPrint for BinaryOperator {
    fn pretty_print_into(
        &self,
        _identifier_map: &IdentifierMap,
        pretty_printed: &mut String,
    ) -> () {
        match *self {
            BinaryOperator::Minus => pretty_printed.push_str("-"),
            BinaryOperator::Plus => pretty_printed.push_str("+"),
            BinaryOperator::Slash => pretty_printed.push_str("/"),
            BinaryOperator::Star => pretty_printed.push_str("*"),
            BinaryOperator::Equal => pretty_printed.push_str("=="),
            BinaryOperator::NotEqual => pretty_printed.push_str("!="),
            BinaryOperator::Less => pretty_printed.push_str("<"),
            BinaryOperator::LessEqual => pretty_printed.push_str("<="),
            BinaryOperator::Greater => pretty_printed.push_str(">"),
            BinaryOperator::GreaterEqual => pretty_printed.push_str(">="),
        }
    }
}

impl PrettyPrint for LogicOperator {
    fn pretty_print_into(
        &self,
        _identifier_map: &IdentifierMap,
        pretty_printed: &mut String,
    ) -> () {
        match *self {
            LogicOperator::Or => pretty_printed.push_str("or"),
            LogicOperator::And => pretty_printed.push_str("and"),
        }
    }
}

impl PrettyPrint for Literal {
    fn pretty_print_into(
        &self,
        _identifier_map: &IdentifierMap,
        pretty_printed: &mut String,
    ) -> () {
        match *self {
            Literal::NilLiteral => pretty_printed.push_str("null"),
            Literal::BoolLiteral(ref b) => pretty_printed.push_str(&b.to_string()),
            Literal::StringLiteral(ref s) => pretty_printed.push_str(s),
            Literal::NumberLiteral(n) => pretty_printed.push_str(&n.to_string()),
        }
    }
}

impl PrettyPrint for Identifier {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str(identifier_map.lookup(self).unwrap())
    }
}

impl PrettyPrint for Target {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Target::Identifier(ref i) => i.pretty_print_into(identifier_map, pretty_printed),
        }
    }
}

impl PrettyPrint for Assignment {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.lvalue
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" = ");
        self.rvalue
            .pretty_print_into(identifier_map, pretty_printed);
    }
}

impl PrettyPrint for Grouping {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(group ");
        self.expr.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for UnaryExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for LogicExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.left.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for BinaryExpr {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        pretty_printed.push_str("(");
        self.operator
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.left.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" ");
        self.right.pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for Statement {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        match *self {
            Statement::Print(ref e) => {
                pretty_printed.push_str("print ");
                e.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::Return(ref e) => {
                pretty_printed.push_str("return");
                if let &Some(ref e) = e {
                    pretty_printed.push_str(" ");
                    e.pretty_print_into(identifier_map, pretty_printed);
                }
                pretty_printed.push_str(";");
            }
            Statement::Expression(ref e) => {
                e.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::VariableDefinition(ref identifier) => {
                pretty_printed.push_str("var ");
                identifier.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::VariableDefinitionWithInitalizer(ref identifier, ref initializer) => {
                pretty_printed.push_str("var ");
                identifier.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" = ");
                initializer.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(";");
            }
            Statement::Block(ref b) => {
                pretty_printed.push_str("{ ");
                for statement in &b.statements {
                    statement.pretty_print_into(identifier_map, pretty_printed);
                    pretty_printed.push_str(" ");
                }
                pretty_printed.push_str("}");
            }
            Statement::IfThen(ref c) => {
                pretty_printed.push_str("if ( ");
                c.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                c.then_branch
                    .pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::IfThenElse(ref c) => {
                pretty_printed.push_str("if ( ");
                c.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                c.then_branch
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" else ");
                c.else_branch
                    .pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::While(ref l) => {
                pretty_printed.push_str("while ( ");
                l.condition
                    .pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" ) ");
                l.body.pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::FunctionDefinition(ref f) => {
                if let FunctionKind::Function = f.kind {
                    pretty_printed.push_str("fun ");
                }
                f.name.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" (");
                for argument in &f.arguments {
                    argument.pretty_print_into(identifier_map, pretty_printed);
                    pretty_printed.push_str(" ");
                }
                pretty_printed.push_str(") ");
                f.body.pretty_print_into(identifier_map, pretty_printed);
            }
            Statement::Class(ref c) => {
                pretty_printed.push_str("class ");
                c.name.pretty_print_into(identifier_map, pretty_printed);
                pretty_printed.push_str(" {");
                for method in &c.methods {
                    pretty_printed.push_str(" ");
                    Statement::FunctionDefinition(method.clone())
                        .pretty_print_into(identifier_map, pretty_printed);
                }
                pretty_printed.push_str(" }");
            }
        };
    }
}

impl PrettyPrint for Call {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.callee
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str("( ");
        for arg in self.arguments.iter() {
            arg.pretty_print_into(identifier_map, pretty_printed);
            pretty_printed.push_str(" ");
        }
        pretty_printed.push_str(")");
    }
}

impl PrettyPrint for Get {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.instance
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(".");
        self.property
            .pretty_print_into(identifier_map, pretty_printed);
    }
}

impl PrettyPrint for Set {
    fn pretty_print_into(&self, identifier_map: &IdentifierMap, pretty_printed: &mut String) -> () {
        self.instance
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(".");
        self.property
            .pretty_print_into(identifier_map, pretty_printed);
        pretty_printed.push_str(" = ");
        self.value.pretty_print_into(identifier_map, pretty_printed);
    }
}

#[cfg(test)]
mod tests {
    use ast::*;
    use pretty_printer::PrettyPrint;

    #[test]
    fn literal() {
        let identifier_map = IdentifierMap::new();
        let string = String::from("abc");
        let expr = Expr::Literal(Literal::StringLiteral(string.clone()));
        assert_eq!(string, expr.pretty_print(&identifier_map));
    }

    #[test]
    fn complex_expression() {
        let identifier_map = IdentifierMap::new();
        let subexpr1 = UnaryExpr {
            operator: UnaryOperator::Minus,
            right: Expr::Literal(Literal::NumberLiteral(123f64)),
        };
        let subexpr2 = Grouping {
            expr: Expr::Literal(Literal::NumberLiteral(45.67f64)),
        };
        let binary_expr = BinaryExpr {
            left: Expr::Unary(Box::new(subexpr1)),
            operator: BinaryOperator::Star,
            right: Expr::Grouping(Box::new(subexpr2)),
        };
        let expr = Expr::Binary(Box::new(binary_expr));
        assert_eq!(
            "(* (- 123) (group 45.67))",
            &expr.pretty_print(&identifier_map)
        );
    }

    #[test]
    fn block() {
        let mut identifier_map = IdentifierMap::new();
        let mut handle_factory = VariableUseHandleFactory::new();
        let identifier = identifier_map.from_name(&"x");
        let statements = vec![
            Statement::VariableDefinitionWithInitalizer(
                identifier.clone(),
                Expr::Literal(Literal::BoolLiteral(true)),
            ),
            Statement::Print(Expr::Identifier(handle_factory.next(), identifier.clone())),
        ];
        let block = Statement::Block(Box::new(Block {
            statements: statements,
        }));
        assert_eq!(
            "{ var x = true; print x; }",
            &block.pretty_print(&identifier_map)
        );
    }
}
<|endoftext|>


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

use core::prelude::*;

use util::smallcharset::SmallCharSet;

use core::str::CharRange;
use collections::string::String;
use collections::{MutableSeq, Deque};
use collections::dlist::DList;

struct Buffer {
    /// Byte position within the buffer.
    pub pos: uint,
    /// The buffer.
    pub buf: String,
}

/// Result from `pop_except_from`.
#[deriving(PartialEq, Eq, Show)]
pub enum SetResult {
    FromSet(char),
    NotFromSet(String),
}

/// A queue of owned string buffers, which supports incrementally
/// consuming characters.
pub struct BufferQueue {
    /// Buffers to process.
    buffers: DList<Buffer>,

    /// Number of available characters.
    available: uint,
}

impl BufferQueue {
    /// Create an empty BufferQueue.
    pub fn new() -> BufferQueue {
        BufferQueue {
            buffers: DList::new(),
            available: 0,
        }
    }

    /// Add a buffer to the beginning of the queue.
    pub fn push_front(&mut self, buf: String) {
        if buf.len() == 0 {
            return;
        }
        self.account_new(buf.as_slice());
        self.buffers.push_front(Buffer {
            pos: 0,
            buf: buf,
        });
    }

    /// Add a buffer to the end of the queue.
    /// 'pos' can be non-zero to remove that many bytes
    /// from the beginning.
    pub fn push_back(&mut self, buf: String, pos: uint) {
        if pos >= buf.len() {
            return;
        }
        self.account_new(buf.as_slice().slice_from(pos));
        self.buffers.push(Buffer {
            pos: pos,
            buf: buf,
        });
    }

    /// Do we have at least n characters available?
    pub fn has(&self, n: uint) -> bool {
        self.available >= n
    }

    /// Get multiple characters, if that many are available.
    pub fn pop_front(&mut self, n: uint) -> Option<String> {
        if !self.has(n) {
            return None;
        }
        // FIXME: this is probably pretty inefficient
        Some(self.by_ref().take(n).collect())
    }

    /// Look at the next available character, if any.
    pub fn peek(&mut self) -> Option<char> {
        match self.buffers.front() {
            Some(&Buffer { pos, ref buf }) => Some(buf.as_slice().char_at(pos)),
            None => None,
        }
    }

    /// Pops and returns either a single character from the given set, or
    /// a `String` of characters none of which are in the set.  The set
    /// is represented as a bitmask and so can only contain the first 64
    /// ASCII characters.
    pub fn pop_except_from(&mut self, set: SmallCharSet) -> Option<SetResult> {
        let (result, now_empty) = match self.buffers.front_mut() {
            Some(&Buffer { ref mut pos, ref buf }) => {
                let n = set.nonmember_prefix_len(buf.as_slice().slice_from(*pos));
                if n > 0 {
                    let new_pos = *pos + n;
                    let out = String::from_str(buf.as_slice().slice(*pos, new_pos));
                    *pos = new_pos;
                    self.available -= n;
                    (Some(NotFromSet(out)), new_pos >= buf.len())
                } else {
                    let CharRange { ch, next } = buf.as_slice().char_range_at(*pos);
                    *pos = next;
                    self.available -= 1;
                    (Some(FromSet(ch)), next >= buf.len())
                }
            }
            _ => (None, false),
        };

        // Unborrow self for this part.
        if now_empty {
            self.buffers.pop_front();
        }

        result
    }

    fn account_new(&mut self, buf: &str) {
        // FIXME: We could pass through length from the initial [u8] -> String
        // conversion, which already must re-encode or at least scan for UTF-8
        // validity.
        self.available += buf.char_len();
    }
}

impl Iterator<char> for BufferQueue {
    /// Get the next character, if one is available.
    ///
    /// Because more data can arrive at any time, this can return Some(c) after
    /// it returns None.  That is allowed by the Iterator protocol, but it's
    /// unusual!
    fn next(&mut self) -> Option<char> {
        let (result, now_empty) = match self.buffers.front_mut() {
            None => (None, false),
            Some(&Buffer { ref mut pos, ref buf }) => {
                let CharRange { ch, next } = buf.as_slice().char_range_at(*pos);
                *pos = next;
                self.available -= 1;
                (Some(ch), next >= buf.len())
            }
        };

        if now_empty {
            self.buffers.pop_front();
        }

        result
    }
}

#[cfg(test)]
#[allow(non_snake_case)]
mod test {
    use core::prelude::*;
    use collections::string::String;
    use super::{BufferQueue, FromSet, NotFromSet};

    #[test]
    fn smoke_test() {
        let mut bq = BufferQueue::new();
        assert_eq!(bq.has(1), false);
        assert_eq!(bq.peek(), None);
        assert_eq!(bq.next(), None);

        bq.push_back(String::from_str("abc"), 0);
        assert_eq!(bq.has(1), true);
        assert_eq!(bq.has(3), true);
        assert_eq!(bq.has(4), false);

        assert_eq!(bq.peek(), Some('a'));
        assert_eq!(bq.next(), Some('a'));
        assert_eq!(bq.peek(), Some('b'));
        assert_eq!(bq.peek(), Some('b'));
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.peek(), Some('c'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.peek(), None);
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_pop_front() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 0);

        assert_eq!(bq.pop_front(2), Some(String::from_str("ab")));
        assert_eq!(bq.peek(), Some('c'));
        assert_eq!(bq.pop_front(2), None);
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_unconsume() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 0);
        assert_eq!(bq.next(), Some('a'));

        bq.push_front(String::from_str("xy"));
        assert_eq!(bq.next(), Some('x'));
        assert_eq!(bq.next(), Some('y'));
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_pop_except_set() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc&def"), 0);
        let pop = || bq.pop_except_from(small_char_set!('&'));
        assert_eq!(pop(), Some(NotFromSet(String::from_str("abc"))));
        assert_eq!(pop(), Some(FromSet('&')));
        assert_eq!(pop(), Some(NotFromSet(String::from_str("def"))));
        assert_eq!(pop(), None);
    }

    #[test]
    fn can_push_truncated() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 1);
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }
}
<commit_msg>Add a regression test for #37<commit_after>// Copyright 2014 The html5ever Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use core::prelude::*;

use util::smallcharset::SmallCharSet;

use core::str::CharRange;
use collections::string::String;
use collections::{MutableSeq, Deque};
use collections::dlist::DList;

struct Buffer {
    /// Byte position within the buffer.
    pub pos: uint,
    /// The buffer.
    pub buf: String,
}

/// Result from `pop_except_from`.
#[deriving(PartialEq, Eq, Show)]
pub enum SetResult {
    FromSet(char),
    NotFromSet(String),
}

/// A queue of owned string buffers, which supports incrementally
/// consuming characters.
pub struct BufferQueue {
    /// Buffers to process.
    buffers: DList<Buffer>,

    /// Number of available characters.
    available: uint,
}

impl BufferQueue {
    /// Create an empty BufferQueue.
    pub fn new() -> BufferQueue {
        BufferQueue {
            buffers: DList::new(),
            available: 0,
        }
    }

    /// Add a buffer to the beginning of the queue.
    pub fn push_front(&mut self, buf: String) {
        if buf.len() == 0 {
            return;
        }
        self.account_new(buf.as_slice());
        self.buffers.push_front(Buffer {
            pos: 0,
            buf: buf,
        });
    }

    /// Add a buffer to the end of the queue.
    /// 'pos' can be non-zero to remove that many bytes
    /// from the beginning.
    pub fn push_back(&mut self, buf: String, pos: uint) {
        if pos >= buf.len() {
            return;
        }
        self.account_new(buf.as_slice().slice_from(pos));
        self.buffers.push(Buffer {
            pos: pos,
            buf: buf,
        });
    }

    /// Do we have at least n characters available?
    pub fn has(&self, n: uint) -> bool {
        self.available >= n
    }

    /// Get multiple characters, if that many are available.
    pub fn pop_front(&mut self, n: uint) -> Option<String> {
        if !self.has(n) {
            return None;
        }
        // FIXME: this is probably pretty inefficient
        Some(self.by_ref().take(n).collect())
    }

    /// Look at the next available character, if any.
    pub fn peek(&mut self) -> Option<char> {
        match self.buffers.front() {
            Some(&Buffer { pos, ref buf }) => Some(buf.as_slice().char_at(pos)),
            None => None,
        }
    }

    /// Pops and returns either a single character from the given set, or
    /// a `String` of characters none of which are in the set.  The set
    /// is represented as a bitmask and so can only contain the first 64
    /// ASCII characters.
    pub fn pop_except_from(&mut self, set: SmallCharSet) -> Option<SetResult> {
        let (result, now_empty) = match self.buffers.front_mut() {
            Some(&Buffer { ref mut pos, ref buf }) => {
                let n = set.nonmember_prefix_len(buf.as_slice().slice_from(*pos));
                if n > 0 {
                    let new_pos = *pos + n;
                    let out = String::from_str(buf.as_slice().slice(*pos, new_pos));
                    *pos = new_pos;
                    self.available -= n;
                    (Some(NotFromSet(out)), new_pos >= buf.len())
                } else {
                    let CharRange { ch, next } = buf.as_slice().char_range_at(*pos);
                    *pos = next;
                    self.available -= 1;
                    (Some(FromSet(ch)), next >= buf.len())
                }
            }
            _ => (None, false),
        };

        // Unborrow self for this part.
        if now_empty {
            self.buffers.pop_front();
        }

        result
    }

    fn account_new(&mut self, buf: &str) {
        // FIXME: We could pass through length from the initial [u8] -> String
        // conversion, which already must re-encode or at least scan for UTF-8
        // validity.
        self.available += buf.char_len();
    }
}

impl Iterator<char> for BufferQueue {
    /// Get the next character, if one is available.
    ///
    /// Because more data can arrive at any time, this can return Some(c) after
    /// it returns None.  That is allowed by the Iterator protocol, but it's
    /// unusual!
    fn next(&mut self) -> Option<char> {
        let (result, now_empty) = match self.buffers.front_mut() {
            None => (None, false),
            Some(&Buffer { ref mut pos, ref buf }) => {
                let CharRange { ch, next } = buf.as_slice().char_range_at(*pos);
                *pos = next;
                self.available -= 1;
                (Some(ch), next >= buf.len())
            }
        };

        if now_empty {
            self.buffers.pop_front();
        }

        result
    }
}

#[cfg(test)]
#[allow(non_snake_case)]
mod test {
    use core::prelude::*;
    use collections::string::String;
    use super::{BufferQueue, FromSet, NotFromSet};

    #[test]
    fn smoke_test() {
        let mut bq = BufferQueue::new();
        assert_eq!(bq.has(1), false);
        assert_eq!(bq.peek(), None);
        assert_eq!(bq.next(), None);

        bq.push_back(String::from_str("abc"), 0);
        assert_eq!(bq.has(1), true);
        assert_eq!(bq.has(3), true);
        assert_eq!(bq.has(4), false);

        assert_eq!(bq.peek(), Some('a'));
        assert_eq!(bq.next(), Some('a'));
        assert_eq!(bq.peek(), Some('b'));
        assert_eq!(bq.peek(), Some('b'));
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.peek(), Some('c'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.peek(), None);
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_pop_front() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 0);

        assert_eq!(bq.pop_front(2), Some(String::from_str("ab")));
        assert_eq!(bq.peek(), Some('c'));
        assert_eq!(bq.pop_front(2), None);
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_unconsume() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 0);
        assert_eq!(bq.next(), Some('a'));

        bq.push_front(String::from_str("xy"));
        assert_eq!(bq.next(), Some('x'));
        assert_eq!(bq.next(), Some('y'));
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }

    #[test]
    fn can_pop_except_set() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc&def"), 0);
        let pop = || bq.pop_except_from(small_char_set!('&'));
        assert_eq!(pop(), Some(NotFromSet(String::from_str("abc"))));
        assert_eq!(pop(), Some(FromSet('&')));
        assert_eq!(pop(), Some(NotFromSet(String::from_str("def"))));
        assert_eq!(pop(), None);
    }

    #[test]
    fn can_push_truncated() {
        let mut bq = BufferQueue::new();
        bq.push_back(String::from_str("abc"), 1);
        assert!(!bq.has(3)); // regression test for #37
        assert_eq!(bq.next(), Some('b'));
        assert_eq!(bq.next(), Some('c'));
        assert_eq!(bq.next(), None);
    }
}
<|endoftext|>


<commit_before>use bindings::{
    jit_context_t,
    jit_context_create,
    jit_context_destroy,
    jit_context_build_start,
    jit_context_build_end,
    jit_function_t,
    jit_function_next
};
use function::Function;
use std::kinds::marker::ContravariantLifetime;
use util::NativeRef;
/// Holds all of the functions you have built and compiled. There can be multiple, but normally there is only one.
native_ref!(Context, _context, jit_context_t, ContravariantLifetime)
impl<'a> Context<'a> {
    /// Create a new JIT Context
    pub fn new() -> Context<'a> {
        unsafe {
            NativeRef::from_ptr(jit_context_create())
        }
    }
    /// Run a closure that can generate IR
    pub fn build<R>(&self, cb: || -> R) -> R {
        unsafe {
            jit_context_build_start(self.as_ptr());
        }
        let rv = cb();
        unsafe {
            jit_context_build_end(self.as_ptr());
        }
        rv
    }
    /// Iterate through all the functions in this context
    pub fn iter_funcs(&self) -> Functions<'a> {
        Functions::new(self)
    }
}
#[unsafe_destructor]
impl<'a> Drop for Context<'a> {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            jit_context_destroy(self.as_ptr());
        }
    }
}
/// Any JIT object which is in a context
pub trait InContext<'a> {
    /// Get the context for this object
    fn get_context(&self) -> Context<'a>;
}

/// An iterator over a context's functions
pub struct Functions<'a> {
    ctx: jit_context_t,
    last: jit_function_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Functions<'a> {
    fn new(ctx:&Context<'a>) -> Functions<'a> {
        unsafe {
            Functions {
                ctx: ctx.as_ptr(),
                last: RawPtr::null(),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Function<'a>> for Functions<'a> {
    fn next(&mut self) -> Option<Function> {
        unsafe {
            let native_next = jit_function_next(self.ctx, self.last);
            self.last = native_next;
            NativeRef::from_ptr_opt(native_next)
        }
    }
    fn size_hint(&self) -> (uint, Option<uint>) {
        unsafe {
            let mut size : uint = 0;
            let mut last = self.last;
            loop {
                last = jit_function_next(self.ctx, last);
                if last.is_null() {
                    break;
                } else {
                    size += 1;
                }
            }
            (size, Some(size))
        }
    }
}<commit_msg>Group unsafe blocks in `context.build`<commit_after>use bindings::{
    jit_context_t,
    jit_context_create,
    jit_context_destroy,
    jit_context_build_start,
    jit_context_build_end,
    jit_function_t,
    jit_function_next
};
use function::Function;
use std::kinds::marker::ContravariantLifetime;
use util::NativeRef;
/// Holds all of the functions you have built and compiled. There can be multiple, but normally there is only one.
native_ref!(Context, _context, jit_context_t, ContravariantLifetime)
impl<'a> Context<'a> {
    /// Create a new JIT Context
    pub fn new() -> Context<'a> {
        unsafe {
            NativeRef::from_ptr(jit_context_create())
        }
    }
    /// Run a closure that can generate IR
    pub fn build<R>(&self, cb: || -> R) -> R {
        unsafe {
            jit_context_build_start(self.as_ptr());
            let rv = cb();
            jit_context_build_end(self.as_ptr());
            rv
        }
    }
    /// Iterate through all the functions in this context
    pub fn iter_funcs(&self) -> Functions<'a> {
        Functions::new(self)
    }
}
#[unsafe_destructor]
impl<'a> Drop for Context<'a> {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            jit_context_destroy(self.as_ptr());
        }
    }
}
/// Any JIT object which is in a context
pub trait InContext<'a> {
    /// Get the context for this object
    fn get_context(&self) -> Context<'a>;
}

/// An iterator over a context's functions
pub struct Functions<'a> {
    ctx: jit_context_t,
    last: jit_function_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Functions<'a> {
    fn new(ctx:&Context<'a>) -> Functions<'a> {
        unsafe {
            Functions {
                ctx: ctx.as_ptr(),
                last: RawPtr::null(),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Function<'a>> for Functions<'a> {
    fn next(&mut self) -> Option<Function> {
        unsafe {
            let native_next = jit_function_next(self.ctx, self.last);
            self.last = native_next;
            NativeRef::from_ptr_opt(native_next)
        }
    }
    fn size_hint(&self) -> (uint, Option<uint>) {
        unsafe {
            let mut size : uint = 0;
            let mut last = self.last;
            loop {
                last = jit_function_next(self.ctx, last);
                if last.is_null() {
                    break;
                } else {
                    size += 1;
                }
            }
            (size, Some(size))
        }
    }
}<|endoftext|>