<commit_before>use token::Token;
use util::Slab;
use time::precise_time_ns;
use std::{usize, iter};
use std::cmp::max;

use self::TimerErrorKind::TimerOverflow;

const EMPTY: Token = Token(usize::MAX);
const NS_PER_MS: u64 = 1_000_000;

// Implements coarse-grained timeouts using an algorithm based on hashed timing
// wheels by Varghese & Lauck.
//
// TODO:
// * Handle the case when the timer falls more than an entire wheel behind. There
//   is no point to loop multiple times around the wheel in one go.
// * New type for tick, now() -> Tick
#[derive(Debug)]
pub struct Timer<T> {
    // Size of each tick in milliseconds
    tick_ms: u64,
    // Slab of timeout entries
    entries: Slab<Entry<T>>,
    // Timeout wheel. Each tick, the timer will look at the next slot for
    // timeouts that match the current tick.
    wheel: Vec<Token>,
    // Tick 0's time in milliseconds
    start: u64,
    // The current tick
    tick: u64,
    // The next entry to possibly timeout
    next: Token,
    // Masks the target tick to get the slot
    mask: u64,
}

#[derive(Copy, Clone)]
pub struct Timeout {
    // Reference into the timer entry slab
    token: Token,
    // Tick that it should matchup with
    tick: u64,
}

impl<T> Timer<T> {
    pub fn new(tick_ms: u64, mut slots: usize, mut capacity: usize) -> Timer<T> {
        slots = slots.next_power_of_two();
        capacity = capacity.next_power_of_two();

        Timer {
            tick_ms: tick_ms,
            entries: Slab::new(capacity),
            wheel: iter::repeat(EMPTY).take(slots).collect(),
            start: 0,
            tick: 0,
            next: EMPTY,
            mask: (slots as u64) - 1
        }
    }

    #[cfg(test)]
    pub fn count(&self) -> usize {
        self.entries.count()
    }

    // Number of ms remaining until the next tick
    pub fn next_tick_in_ms(&self) -> Option<u64> {
        if self.entries.count() == 0 {
            return None;
        }

        let now = self.now_ms();
        let nxt = self.start + (self.tick + 1) * self.tick_ms;

        if nxt <= now {
            return Some(0);
        }

        Some(nxt - now)
    }

    /*
     *
     * ===== Initialization =====
     *
     */

    // Sets the starting time of the timer using the current system time
    pub fn setup(&mut self) {
        let now = self.now_ms();
        self.set_start_ms(now);
    }

    fn set_start_ms(&mut self, start: u64) {
        assert!(!self.is_initialized(), "the timer has already started");
        self.start = start;
    }

    /*
     *
     * ===== Timeout create / cancel =====
     *
     */

    pub fn timeout_ms(&mut self, token: T, delay: u64) -> TimerResult<Timeout> {
        let at = self.now_ms() + max(0, delay);
        self.timeout_at_ms(token, at)
    }

    pub fn timeout_at_ms(&mut self, token: T, mut at: u64) -> TimerResult<Timeout> {
        // Make relative to start
        at -= self.start;
        // Calculate tick
        let mut tick = (at + self.tick_ms - 1) / self.tick_ms;

        // Always target at least 1 tick in the future
        if tick <= self.tick {
            tick = self.tick + 1;
        }

        self.insert(token, tick)
    }

    pub fn clear(&mut self, timeout: Timeout) -> bool {
        let links = match self.entries.get(timeout.token) {
            Some(e) => e.links,
            None => return false
        };

        // Sanity check
        if links.tick != timeout.tick {
            return false;
        }

        self.unlink(&links, timeout.token);
        self.entries.remove(timeout.token);
        true
    }

    fn insert(&mut self, token: T, tick: u64) -> TimerResult<Timeout> {
        // Get the slot for the requested tick
        let slot = (tick & self.mask) as usize;
        let curr = self.wheel[slot];

        // Insert the new entry
        let token = try!(
            self.entries.insert(Entry::new(token, tick, curr))
            .map_err(|_| TimerError::overflow()));

        if curr != EMPTY {
            // If there was a previous entry, set its prev pointer to the new
            // entry
            self.entries[curr].links.prev = token;
        }

        // Update the head slot
        self.wheel[slot] = token;

        trace!("inserted timout; slot={}; token={:?}", slot, token);

        // Return the new timeout
        Ok(Timeout {
            token: token,
            tick: tick
        })
    }

    fn unlink(&mut self, links: &EntryLinks, token: Token) {
       trace!("unlinking timeout; slot={}; token={:?}",
               self.slot_for(links.tick), token);

        if links.prev == EMPTY {
            let slot = self.slot_for(links.tick);
            self.wheel[slot] = links.next;
        } else {
            self.entries[links.prev].links.next = links.next;
        }

        if links.next != EMPTY {
            self.entries[links.next].links.prev = links.prev;

            if token == self.next {
                self.next = links.next;
            }
        } else if token == self.next {
            self.next = EMPTY;
        }
    }

    /*
     *
     * ===== Advance time =====
     *
     */

    pub fn now(&self) -> u64 {
        self.ms_to_tick(self.now_ms())
    }

    pub fn tick_to(&mut self, now: u64) -> Option<T> {
        trace!("tick_to; now={}; tick={}", now, self.tick);

        while self.tick <= now {
            let curr = self.next;

            trace!("ticking; curr={:?}", curr);

            if curr == EMPTY {
                self.tick += 1;
                self.next = self.wheel[self.slot_for(self.tick)];
            } else {
                let links = self.entries[curr].links;

                if links.tick <= self.tick {
                    trace!("triggering; token={:?}", curr);

                    // Unlink will also advance self.next
                    self.unlink(&links, curr);

                    // Remove and return the token
                    return self.entries.remove(curr)
                        .map(|e| e.token);
                } else {
                    self.next = links.next;
                }
            }
        }

        None
    }

    /*
     *
     * ===== Misc =====
     *
     */

    // Timers are initialized when either the current time has been advanced or a timeout has been set
    #[inline]
    fn is_initialized(&self) -> bool {
        self.tick > 0 || !self.entries.is_empty()
    }

    #[inline]
    fn slot_for(&self, tick: u64) -> usize {
        (self.mask & tick) as usize
    }

    // Convert a ms duration into a number of ticks, rounds up
    #[inline]
    fn ms_to_tick(&self, ms: u64) -> u64 {
        (ms - self.start) / self.tick_ms
    }

    #[inline]
    fn now_ms(&self) -> u64 {
        precise_time_ns() / NS_PER_MS
    }
}

// Doubly linked list of timer entries. Allows for efficient insertion /
// removal of timeouts.
struct Entry<T> {
    token: T,
    links: EntryLinks,
}

impl<T> Entry<T> {
    fn new(token: T, tick: u64, next: Token) -> Entry<T> {
        Entry {
            token: token,
            links: EntryLinks {
                tick: tick,
                prev: EMPTY,
                next: next,
            },
        }
    }
}

#[derive(Copy, Clone)]
struct EntryLinks {
    tick: u64,
    prev: Token,
    next: Token
}

pub type TimerResult<T> = Result<T, TimerError>;

#[derive(Debug)]
pub struct TimerError {
    kind: TimerErrorKind,
    desc: &'static str,
}

impl TimerError {
    fn overflow() -> TimerError {
        TimerError {
            kind: TimerOverflow,
            desc: "too many timer entries"
        }
    }
}

#[derive(Debug)]
pub enum TimerErrorKind {
    TimerOverflow,
}

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

    #[test]
    pub fn test_timeout_next_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();

        tick = t.ms_to_tick(50);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(100);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(150);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_clearing_timeout() {
        let mut t = timer();
        let mut tick;

        let to = t.timeout_at_ms("a", 100).unwrap();
        assert!(t.clear(to));

        tick = t.ms_to_tick(100);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_multiple_timeouts_same_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();
        t.timeout_at_ms("b", 100).unwrap();

        let mut rcv = vec![];

        tick = t.ms_to_tick(100);
        rcv.push(t.tick_to(tick).unwrap());
        rcv.push(t.tick_to(tick).unwrap());

        assert_eq!(None, t.tick_to(tick));

        rcv.sort();
        assert!(rcv == ["a", "b"], "actual={:?}", rcv);

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_multiple_timeouts_diff_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 110).unwrap();
        t.timeout_at_ms("b", 220).unwrap();
        t.timeout_at_ms("c", 230).unwrap();
        t.timeout_at_ms("d", 440).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(300);
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(400);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(500);
        assert_eq!(Some("d"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(600);
        assert_eq!(None, t.tick_to(tick));
    }

    #[test]
    pub fn test_catching_up() {
        let mut t = timer();

        t.timeout_at_ms("a", 110).unwrap();
        t.timeout_at_ms("b", 220).unwrap();
        t.timeout_at_ms("c", 230).unwrap();
        t.timeout_at_ms("d", 440).unwrap();

        let tick = t.ms_to_tick(600);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(Some("d"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));
    }

    #[test]
    pub fn test_timeout_hash_collision() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();
        t.timeout_at_ms("b", 100 + TICK * SLOTS as u64).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(1, t.count());

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));
        assert_eq!(1, t.count());

        tick = t.ms_to_tick(100 + TICK * SLOTS as u64);
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(0, t.count());
    }

    #[test]
    pub fn test_clearing_timeout_between_triggers() {
        let mut t = timer();
        let mut tick;

        let a = t.timeout_at_ms("a", 100).unwrap();
        let _ = t.timeout_at_ms("b", 100).unwrap();
        let _ = t.timeout_at_ms("c", 200).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(2, t.count());

        t.clear(a);
        assert_eq!(1, t.count());

        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(0, t.count());
    }

    const TICK: u64 = 100;
    const SLOTS: usize = 16;

    fn timer() -> Timer<&'static str> {
        Timer::new(TICK, SLOTS, 32)
    }
}
<commit_msg>Implement error::Error for TimerError<commit_after>use token::Token;
use util::Slab;
use time::precise_time_ns;
use std::{error, fmt, usize, iter};
use std::cmp::max;

use self::TimerErrorKind::TimerOverflow;

const EMPTY: Token = Token(usize::MAX);
const NS_PER_MS: u64 = 1_000_000;

// Implements coarse-grained timeouts using an algorithm based on hashed timing
// wheels by Varghese & Lauck.
//
// TODO:
// * Handle the case when the timer falls more than an entire wheel behind. There
//   is no point to loop multiple times around the wheel in one go.
// * New type for tick, now() -> Tick
#[derive(Debug)]
pub struct Timer<T> {
    // Size of each tick in milliseconds
    tick_ms: u64,
    // Slab of timeout entries
    entries: Slab<Entry<T>>,
    // Timeout wheel. Each tick, the timer will look at the next slot for
    // timeouts that match the current tick.
    wheel: Vec<Token>,
    // Tick 0's time in milliseconds
    start: u64,
    // The current tick
    tick: u64,
    // The next entry to possibly timeout
    next: Token,
    // Masks the target tick to get the slot
    mask: u64,
}

#[derive(Copy, Clone)]
pub struct Timeout {
    // Reference into the timer entry slab
    token: Token,
    // Tick that it should matchup with
    tick: u64,
}

impl<T> Timer<T> {
    pub fn new(tick_ms: u64, mut slots: usize, mut capacity: usize) -> Timer<T> {
        slots = slots.next_power_of_two();
        capacity = capacity.next_power_of_two();

        Timer {
            tick_ms: tick_ms,
            entries: Slab::new(capacity),
            wheel: iter::repeat(EMPTY).take(slots).collect(),
            start: 0,
            tick: 0,
            next: EMPTY,
            mask: (slots as u64) - 1
        }
    }

    #[cfg(test)]
    pub fn count(&self) -> usize {
        self.entries.count()
    }

    // Number of ms remaining until the next tick
    pub fn next_tick_in_ms(&self) -> Option<u64> {
        if self.entries.count() == 0 {
            return None;
        }

        let now = self.now_ms();
        let nxt = self.start + (self.tick + 1) * self.tick_ms;

        if nxt <= now {
            return Some(0);
        }

        Some(nxt - now)
    }

    /*
     *
     * ===== Initialization =====
     *
     */

    // Sets the starting time of the timer using the current system time
    pub fn setup(&mut self) {
        let now = self.now_ms();
        self.set_start_ms(now);
    }

    fn set_start_ms(&mut self, start: u64) {
        assert!(!self.is_initialized(), "the timer has already started");
        self.start = start;
    }

    /*
     *
     * ===== Timeout create / cancel =====
     *
     */

    pub fn timeout_ms(&mut self, token: T, delay: u64) -> TimerResult<Timeout> {
        let at = self.now_ms() + max(0, delay);
        self.timeout_at_ms(token, at)
    }

    pub fn timeout_at_ms(&mut self, token: T, mut at: u64) -> TimerResult<Timeout> {
        // Make relative to start
        at -= self.start;
        // Calculate tick
        let mut tick = (at + self.tick_ms - 1) / self.tick_ms;

        // Always target at least 1 tick in the future
        if tick <= self.tick {
            tick = self.tick + 1;
        }

        self.insert(token, tick)
    }

    pub fn clear(&mut self, timeout: Timeout) -> bool {
        let links = match self.entries.get(timeout.token) {
            Some(e) => e.links,
            None => return false
        };

        // Sanity check
        if links.tick != timeout.tick {
            return false;
        }

        self.unlink(&links, timeout.token);
        self.entries.remove(timeout.token);
        true
    }

    fn insert(&mut self, token: T, tick: u64) -> TimerResult<Timeout> {
        // Get the slot for the requested tick
        let slot = (tick & self.mask) as usize;
        let curr = self.wheel[slot];

        // Insert the new entry
        let token = try!(
            self.entries.insert(Entry::new(token, tick, curr))
            .map_err(|_| TimerError::overflow()));

        if curr != EMPTY {
            // If there was a previous entry, set its prev pointer to the new
            // entry
            self.entries[curr].links.prev = token;
        }

        // Update the head slot
        self.wheel[slot] = token;

        trace!("inserted timout; slot={}; token={:?}", slot, token);

        // Return the new timeout
        Ok(Timeout {
            token: token,
            tick: tick
        })
    }

    fn unlink(&mut self, links: &EntryLinks, token: Token) {
       trace!("unlinking timeout; slot={}; token={:?}",
               self.slot_for(links.tick), token);

        if links.prev == EMPTY {
            let slot = self.slot_for(links.tick);
            self.wheel[slot] = links.next;
        } else {
            self.entries[links.prev].links.next = links.next;
        }

        if links.next != EMPTY {
            self.entries[links.next].links.prev = links.prev;

            if token == self.next {
                self.next = links.next;
            }
        } else if token == self.next {
            self.next = EMPTY;
        }
    }

    /*
     *
     * ===== Advance time =====
     *
     */

    pub fn now(&self) -> u64 {
        self.ms_to_tick(self.now_ms())
    }

    pub fn tick_to(&mut self, now: u64) -> Option<T> {
        trace!("tick_to; now={}; tick={}", now, self.tick);

        while self.tick <= now {
            let curr = self.next;

            trace!("ticking; curr={:?}", curr);

            if curr == EMPTY {
                self.tick += 1;
                self.next = self.wheel[self.slot_for(self.tick)];
            } else {
                let links = self.entries[curr].links;

                if links.tick <= self.tick {
                    trace!("triggering; token={:?}", curr);

                    // Unlink will also advance self.next
                    self.unlink(&links, curr);

                    // Remove and return the token
                    return self.entries.remove(curr)
                        .map(|e| e.token);
                } else {
                    self.next = links.next;
                }
            }
        }

        None
    }

    /*
     *
     * ===== Misc =====
     *
     */

    // Timers are initialized when either the current time has been advanced or a timeout has been set
    #[inline]
    fn is_initialized(&self) -> bool {
        self.tick > 0 || !self.entries.is_empty()
    }

    #[inline]
    fn slot_for(&self, tick: u64) -> usize {
        (self.mask & tick) as usize
    }

    // Convert a ms duration into a number of ticks, rounds up
    #[inline]
    fn ms_to_tick(&self, ms: u64) -> u64 {
        (ms - self.start) / self.tick_ms
    }

    #[inline]
    fn now_ms(&self) -> u64 {
        precise_time_ns() / NS_PER_MS
    }
}

// Doubly linked list of timer entries. Allows for efficient insertion /
// removal of timeouts.
struct Entry<T> {
    token: T,
    links: EntryLinks,
}

impl<T> Entry<T> {
    fn new(token: T, tick: u64, next: Token) -> Entry<T> {
        Entry {
            token: token,
            links: EntryLinks {
                tick: tick,
                prev: EMPTY,
                next: next,
            },
        }
    }
}

#[derive(Copy, Clone)]
struct EntryLinks {
    tick: u64,
    prev: Token,
    next: Token
}

pub type TimerResult<T> = Result<T, TimerError>;

#[derive(Debug)]
pub struct TimerError {
    kind: TimerErrorKind,
    desc: &'static str,
}

impl fmt::Display for TimerError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}: {}", self.kind, self.desc)
    }
}

impl TimerError {
    fn overflow() -> TimerError {
        TimerError {
            kind: TimerOverflow,
            desc: "too many timer entries"
        }
    }
}

impl error::Error for TimerError {
    fn description(&self) -> &str {
        self.desc
    }
}

#[derive(Debug)]
pub enum TimerErrorKind {
    TimerOverflow,
}

impl fmt::Display for TimerErrorKind {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TimerOverflow => write!(fmt, "TimerOverflow"),
        }
    }
}

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

    #[test]
    pub fn test_timeout_next_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();

        tick = t.ms_to_tick(50);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(100);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(150);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_clearing_timeout() {
        let mut t = timer();
        let mut tick;

        let to = t.timeout_at_ms("a", 100).unwrap();
        assert!(t.clear(to));

        tick = t.ms_to_tick(100);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_multiple_timeouts_same_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();
        t.timeout_at_ms("b", 100).unwrap();

        let mut rcv = vec![];

        tick = t.ms_to_tick(100);
        rcv.push(t.tick_to(tick).unwrap());
        rcv.push(t.tick_to(tick).unwrap());

        assert_eq!(None, t.tick_to(tick));

        rcv.sort();
        assert!(rcv == ["a", "b"], "actual={:?}", rcv);

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));

        assert_eq!(t.count(), 0);
    }

    #[test]
    pub fn test_multiple_timeouts_diff_tick() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 110).unwrap();
        t.timeout_at_ms("b", 220).unwrap();
        t.timeout_at_ms("c", 230).unwrap();
        t.timeout_at_ms("d", 440).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(300);
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(400);
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(500);
        assert_eq!(Some("d"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(600);
        assert_eq!(None, t.tick_to(tick));
    }

    #[test]
    pub fn test_catching_up() {
        let mut t = timer();

        t.timeout_at_ms("a", 110).unwrap();
        t.timeout_at_ms("b", 220).unwrap();
        t.timeout_at_ms("c", 230).unwrap();
        t.timeout_at_ms("d", 440).unwrap();

        let tick = t.ms_to_tick(600);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(Some("d"), t.tick_to(tick));
        assert_eq!(None, t.tick_to(tick));
    }

    #[test]
    pub fn test_timeout_hash_collision() {
        let mut t = timer();
        let mut tick;

        t.timeout_at_ms("a", 100).unwrap();
        t.timeout_at_ms("b", 100 + TICK * SLOTS as u64).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(Some("a"), t.tick_to(tick));
        assert_eq!(1, t.count());

        tick = t.ms_to_tick(200);
        assert_eq!(None, t.tick_to(tick));
        assert_eq!(1, t.count());

        tick = t.ms_to_tick(100 + TICK * SLOTS as u64);
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(0, t.count());
    }

    #[test]
    pub fn test_clearing_timeout_between_triggers() {
        let mut t = timer();
        let mut tick;

        let a = t.timeout_at_ms("a", 100).unwrap();
        let _ = t.timeout_at_ms("b", 100).unwrap();
        let _ = t.timeout_at_ms("c", 200).unwrap();

        tick = t.ms_to_tick(100);
        assert_eq!(Some("b"), t.tick_to(tick));
        assert_eq!(2, t.count());

        t.clear(a);
        assert_eq!(1, t.count());

        assert_eq!(None, t.tick_to(tick));

        tick = t.ms_to_tick(200);
        assert_eq!(Some("c"), t.tick_to(tick));
        assert_eq!(0, t.count());
    }

    const TICK: u64 = 100;
    const SLOTS: usize = 16;

    fn timer() -> Timer<&'static str> {
        Timer::new(TICK, SLOTS, 32)
    }
}
<|endoftext|>


<commit_before>use std::collections::BinaryHeap;
use std::cmp::{PartialOrd, Ord, PartialEq, Eq, Ordering};
use std::sync::{Arc, Mutex, Condvar};
use std::thread::Thread;
use std::thunk::Thunk;
use std::time::Duration;

use time;

enum JobType {
    Once(Thunk),
    FixedRate {
        f: Box<FnMut() + Send>,
        rate: Duration,
    },
}

struct Job {
    type_: JobType,
    time: u64,
}

impl PartialOrd for Job {
    fn partial_cmp(&self, other: &Job) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Job {
    fn cmp(&self, other: &Job) -> Ordering {
        // reverse because BinaryHeap's a max heap
        self.time.cmp(&other.time).reverse()
    }
}

impl PartialEq for Job {
    fn eq(&self, other: &Job) -> bool {
        self.time == other.time
    }
}

impl Eq for Job {}

struct InnerPool {
    queue: BinaryHeap<Job>,
    shutdown: bool,
}

struct SharedPool {
    inner: Mutex<InnerPool>,
    cvar: Condvar,
}

impl SharedPool {
    fn run(&self, job: Job) {
        let mut inner = self.inner.lock().unwrap();

        // Calls from the pool itself will never hit this, but calls from workers might
        if inner.shutdown {
            return;
        }

        match inner.queue.peek() {
            None => self.cvar.notify_all(),
            Some(e) if e.time > job.time => self.cvar.notify_all(),
            _ => {}
        };
        inner.queue.push(job);
    }
}

/// A fixed-size thread pool which allows jobs to be scheduled in the future.
///
/// When the pool falls out of scope, all pending tasks will be executed, after
/// which the worker threads will shut down.
pub struct ScheduledTaskPool {
    shared: Arc<SharedPool>,
}

impl Drop for ScheduledTaskPool {
    fn drop(&mut self) {
        self.shared.inner.lock().unwrap().shutdown = true;
        self.shared.cvar.notify_all();
    }
}

impl ScheduledTaskPool {
    /// Creates a new `ScheduledTaskPool` with the specified number of threads.
    ///
    /// # Panics
    ///
    /// Panics if `size` is 0.
    pub fn new(size: usize) -> ScheduledTaskPool {
        assert!(size > 0, "size must be positive");

        let inner = InnerPool {
            queue: BinaryHeap::new(),
            shutdown: false,
        };

        let shared = SharedPool {
            inner: Mutex::new(inner),
            cvar: Condvar::new(),
        };

        let pool = ScheduledTaskPool {
            shared: Arc::new(shared),
        };

        for _ in (0..size) {
            let mut worker = Worker {
                shared: pool.shared.clone(),
            };

            Thread::spawn(move || worker.run());
        }

        pool
    }

    /// Asynchronously executes `job` with no added delay.
    pub fn run<F>(&self, job: F) where F: FnOnce() + Send {
        self.run_after(Duration::zero(), f)
    }

    /// Asynchronously executes `job` after the specified delay.
    pub fn run_after<F>(&self, dur: Duration, job: F) where F: FnOnce() + Send {
        let job = Job {
            type_: JobType::Once(Thunk::new(job)),
            time: (time::precise_time_ns() as i64 + dur.num_nanoseconds().unwrap()) as u64,
        };
        self.shared.run(job)
    }

    /// Asynchronously executes `job` repeatedly at the specified rate.
    ///
    /// If the job panics, it will no longer be executed. When the pool is
    /// destroyed, the job will no longer be rescheduled for execution, but any
    /// pending execution of the job will be handled as a normal job would.
    pub fn run_at_fixed_rate<F>(&self, rate: Duration, f: F) where F: FnMut() + Send {
        let job = Job {
            type_: JobType::FixedRate { f: Box::new(f), rate: rate },
            time: (time::precise_time_ns() as i64 + rate.num_nanoseconds().unwrap()) as u64,
        };
        self.shared.run(job)
    }

    /// Consumes the `ScheduledTaskPool`, canceling any pending jobs.
    ///
    /// Currently running jobs will continue to run to completion.
    pub fn shutdown_now(self) {
        self.shared.inner.lock().unwrap().queue.clear();
    }
}

struct Worker {
    shared: Arc<SharedPool>,
}

impl Drop for Worker {
    fn drop(&mut self) {
        // Start up a new worker if this one's going away due to a panic from a job
        if Thread::panicking() {
            let mut worker = Worker {
                shared: self.shared.clone(),
            };
            Thread::spawn(move || worker.run());
        }
    }
}

impl Worker {
    fn run(&mut self) {
        loop {
            match self.get_job() {
                Some(job) => self.run_job(job),
                None => break,
            }
        }
    }

    fn get_job(&self) -> Option<Job> {
        enum Need {
            Wait,
            WaitTimeout(Duration),
        }

        let mut inner = self.shared.inner.lock().unwrap();
        loop {
            let now = time::precise_time_ns();

            let need = match inner.queue.peek() {
                None if inner.shutdown => return None,
                None => Need::Wait,
                Some(e) if e.time <= now => break,
                Some(e) => Need::WaitTimeout(Duration::nanoseconds(e.time as i64 - now as i64)),
            };

            inner = match need {
                Need::Wait => self.shared.cvar.wait(inner).unwrap(),
                Need::WaitTimeout(t) => self.shared.cvar.wait_timeout(inner, t).unwrap().0,
            };
        }

        Some(inner.queue.pop().unwrap())
    }

    fn run_job(&self, job: Job) {
        match job.type_ {
            JobType::Once(f) => f.invoke(()),
            JobType::FixedRate { mut f, rate } => {
                f();
                let new_job = Job {
                    type_: JobType::FixedRate { f: f, rate: rate },
                    time: (job.time as i64 + rate.num_nanoseconds().unwrap()) as u64,
                };
                self.shared.run(new_job)
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::iter::AdditiveIterator;
    use std::sync::mpsc::channel;
    use std::sync::{Arc, Barrier};
    use std::time::Duration;

    use super::ScheduledTaskPool;

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            pool.run(move|| {
                tx.send(1us).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    #[should_fail(expected = "size must be positive")]
    fn test_zero_tasks_panic() {
        ScheduledTaskPool::new(0);
    }

    #[test]
    fn test_recovery_from_subtask_panic() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);

        // Panic all the existing threads.
        let waiter = Arc::new(Barrier::new(TEST_TASKS as usize));
        for _ in range(0, TEST_TASKS) {
            let waiter = waiter.clone();
            pool.run(move || -> () {
                waiter.wait();
                panic!();
            });
        }

        // Ensure new threads were spawned to compensate.
        let (tx, rx) = channel();
        let waiter = Arc::new(Barrier::new(TEST_TASKS as usize));
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            let waiter = waiter.clone();
            pool.run(move || {
                waiter.wait();
                tx.send(1us).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    fn test_run_after() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);
        let (tx, rx) = channel();

        let tx1 = tx.clone();
        pool.run_after(Duration::seconds(1), move || tx1.send(1us).unwrap());
        pool.run_after(Duration::milliseconds(500), move || tx.send(2us).unwrap());

        assert_eq!(2, rx.recv().unwrap());
        assert_eq!(1, rx.recv().unwrap());
    }

    #[test]
    fn test_jobs_complete_after_drop() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);
        let (tx, rx) = channel();

        let tx1 = tx.clone();
        pool.run_after(Duration::seconds(1), move || tx1.send(1us).unwrap());
        pool.run_after(Duration::milliseconds(500), move || tx.send(2us).unwrap());

        drop(pool);

        assert_eq!(2, rx.recv().unwrap());
        assert_eq!(1, rx.recv().unwrap());
    }

    #[test]
    fn test_fixed_delay_jobs_stop_after_drop() {
        let pool = Arc::new(ScheduledTaskPool::new(TEST_TASKS));
        let (tx, rx) = channel();
        let (tx2, rx2) = channel();

        let mut pool2 = Some(pool.clone());
        let mut i = 0i32;
        pool.run_at_fixed_rate(Duration::milliseconds(500), move || {
            i += 1;
            tx.send(i).unwrap();
            rx2.recv().unwrap();
            if i == 2 {
                drop(pool2.take().unwrap());
            }
        });
        drop(pool);

        assert_eq!(Ok(1), rx.recv());
        tx2.send(()).unwrap();
        assert_eq!(Ok(2), rx.recv());
        tx2.send(()).unwrap();
        assert!(rx.recv().is_err());
    }
}
<commit_msg>Fix typo<commit_after>use std::collections::BinaryHeap;
use std::cmp::{PartialOrd, Ord, PartialEq, Eq, Ordering};
use std::sync::{Arc, Mutex, Condvar};
use std::thread::Thread;
use std::thunk::Thunk;
use std::time::Duration;

use time;

enum JobType {
    Once(Thunk),
    FixedRate {
        f: Box<FnMut() + Send>,
        rate: Duration,
    },
}

struct Job {
    type_: JobType,
    time: u64,
}

impl PartialOrd for Job {
    fn partial_cmp(&self, other: &Job) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Job {
    fn cmp(&self, other: &Job) -> Ordering {
        // reverse because BinaryHeap's a max heap
        self.time.cmp(&other.time).reverse()
    }
}

impl PartialEq for Job {
    fn eq(&self, other: &Job) -> bool {
        self.time == other.time
    }
}

impl Eq for Job {}

struct InnerPool {
    queue: BinaryHeap<Job>,
    shutdown: bool,
}

struct SharedPool {
    inner: Mutex<InnerPool>,
    cvar: Condvar,
}

impl SharedPool {
    fn run(&self, job: Job) {
        let mut inner = self.inner.lock().unwrap();

        // Calls from the pool itself will never hit this, but calls from workers might
        if inner.shutdown {
            return;
        }

        match inner.queue.peek() {
            None => self.cvar.notify_all(),
            Some(e) if e.time > job.time => self.cvar.notify_all(),
            _ => {}
        };
        inner.queue.push(job);
    }
}

/// A fixed-size thread pool which allows jobs to be scheduled in the future.
///
/// When the pool falls out of scope, all pending tasks will be executed, after
/// which the worker threads will shut down.
pub struct ScheduledTaskPool {
    shared: Arc<SharedPool>,
}

impl Drop for ScheduledTaskPool {
    fn drop(&mut self) {
        self.shared.inner.lock().unwrap().shutdown = true;
        self.shared.cvar.notify_all();
    }
}

impl ScheduledTaskPool {
    /// Creates a new `ScheduledTaskPool` with the specified number of threads.
    ///
    /// # Panics
    ///
    /// Panics if `size` is 0.
    pub fn new(size: usize) -> ScheduledTaskPool {
        assert!(size > 0, "size must be positive");

        let inner = InnerPool {
            queue: BinaryHeap::new(),
            shutdown: false,
        };

        let shared = SharedPool {
            inner: Mutex::new(inner),
            cvar: Condvar::new(),
        };

        let pool = ScheduledTaskPool {
            shared: Arc::new(shared),
        };

        for _ in (0..size) {
            let mut worker = Worker {
                shared: pool.shared.clone(),
            };

            Thread::spawn(move || worker.run());
        }

        pool
    }

    /// Asynchronously executes `job` with no added delay.
    pub fn run<F>(&self, job: F) where F: FnOnce() + Send {
        self.run_after(Duration::zero(), job)
    }

    /// Asynchronously executes `job` after the specified delay.
    pub fn run_after<F>(&self, dur: Duration, job: F) where F: FnOnce() + Send {
        let job = Job {
            type_: JobType::Once(Thunk::new(job)),
            time: (time::precise_time_ns() as i64 + dur.num_nanoseconds().unwrap()) as u64,
        };
        self.shared.run(job)
    }

    /// Asynchronously executes `job` repeatedly at the specified rate.
    ///
    /// If the job panics, it will no longer be executed. When the pool is
    /// destroyed, the job will no longer be rescheduled for execution, but any
    /// pending execution of the job will be handled as a normal job would.
    pub fn run_at_fixed_rate<F>(&self, rate: Duration, f: F) where F: FnMut() + Send {
        let job = Job {
            type_: JobType::FixedRate { f: Box::new(f), rate: rate },
            time: (time::precise_time_ns() as i64 + rate.num_nanoseconds().unwrap()) as u64,
        };
        self.shared.run(job)
    }

    /// Consumes the `ScheduledTaskPool`, canceling any pending jobs.
    ///
    /// Currently running jobs will continue to run to completion.
    pub fn shutdown_now(self) {
        self.shared.inner.lock().unwrap().queue.clear();
    }
}

struct Worker {
    shared: Arc<SharedPool>,
}

impl Drop for Worker {
    fn drop(&mut self) {
        // Start up a new worker if this one's going away due to a panic from a job
        if Thread::panicking() {
            let mut worker = Worker {
                shared: self.shared.clone(),
            };
            Thread::spawn(move || worker.run());
        }
    }
}

impl Worker {
    fn run(&mut self) {
        loop {
            match self.get_job() {
                Some(job) => self.run_job(job),
                None => break,
            }
        }
    }

    fn get_job(&self) -> Option<Job> {
        enum Need {
            Wait,
            WaitTimeout(Duration),
        }

        let mut inner = self.shared.inner.lock().unwrap();
        loop {
            let now = time::precise_time_ns();

            let need = match inner.queue.peek() {
                None if inner.shutdown => return None,
                None => Need::Wait,
                Some(e) if e.time <= now => break,
                Some(e) => Need::WaitTimeout(Duration::nanoseconds(e.time as i64 - now as i64)),
            };

            inner = match need {
                Need::Wait => self.shared.cvar.wait(inner).unwrap(),
                Need::WaitTimeout(t) => self.shared.cvar.wait_timeout(inner, t).unwrap().0,
            };
        }

        Some(inner.queue.pop().unwrap())
    }

    fn run_job(&self, job: Job) {
        match job.type_ {
            JobType::Once(f) => f.invoke(()),
            JobType::FixedRate { mut f, rate } => {
                f();
                let new_job = Job {
                    type_: JobType::FixedRate { f: f, rate: rate },
                    time: (job.time as i64 + rate.num_nanoseconds().unwrap()) as u64,
                };
                self.shared.run(new_job)
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::iter::AdditiveIterator;
    use std::sync::mpsc::channel;
    use std::sync::{Arc, Barrier};
    use std::time::Duration;

    use super::ScheduledTaskPool;

    const TEST_TASKS: usize = 4;

    #[test]
    fn test_works() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            pool.run(move|| {
                tx.send(1us).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    #[should_fail(expected = "size must be positive")]
    fn test_zero_tasks_panic() {
        ScheduledTaskPool::new(0);
    }

    #[test]
    fn test_recovery_from_subtask_panic() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);

        // Panic all the existing threads.
        let waiter = Arc::new(Barrier::new(TEST_TASKS as usize));
        for _ in range(0, TEST_TASKS) {
            let waiter = waiter.clone();
            pool.run(move || -> () {
                waiter.wait();
                panic!();
            });
        }

        // Ensure new threads were spawned to compensate.
        let (tx, rx) = channel();
        let waiter = Arc::new(Barrier::new(TEST_TASKS as usize));
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            let waiter = waiter.clone();
            pool.run(move || {
                waiter.wait();
                tx.send(1us).unwrap();
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    fn test_run_after() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);
        let (tx, rx) = channel();

        let tx1 = tx.clone();
        pool.run_after(Duration::seconds(1), move || tx1.send(1us).unwrap());
        pool.run_after(Duration::milliseconds(500), move || tx.send(2us).unwrap());

        assert_eq!(2, rx.recv().unwrap());
        assert_eq!(1, rx.recv().unwrap());
    }

    #[test]
    fn test_jobs_complete_after_drop() {
        let pool = ScheduledTaskPool::new(TEST_TASKS);
        let (tx, rx) = channel();

        let tx1 = tx.clone();
        pool.run_after(Duration::seconds(1), move || tx1.send(1us).unwrap());
        pool.run_after(Duration::milliseconds(500), move || tx.send(2us).unwrap());

        drop(pool);

        assert_eq!(2, rx.recv().unwrap());
        assert_eq!(1, rx.recv().unwrap());
    }

    #[test]
    fn test_fixed_delay_jobs_stop_after_drop() {
        let pool = Arc::new(ScheduledTaskPool::new(TEST_TASKS));
        let (tx, rx) = channel();
        let (tx2, rx2) = channel();

        let mut pool2 = Some(pool.clone());
        let mut i = 0i32;
        pool.run_at_fixed_rate(Duration::milliseconds(500), move || {
            i += 1;
            tx.send(i).unwrap();
            rx2.recv().unwrap();
            if i == 2 {
                drop(pool2.take().unwrap());
            }
        });
        drop(pool);

        assert_eq!(Ok(1), rx.recv());
        tx2.send(()).unwrap();
        assert_eq!(Ok(2), rx.recv());
        tx2.send(()).unwrap();
        assert!(rx.recv().is_err());
    }
}
<|endoftext|>


<commit_before>use track::*;
use ast::{Binop, BinopTag, Logop, LogopTag, Assop, AssopTag};

// Unconditionally reserved words.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum Reserved {
    // 11.6.2 Reserved Words
    Null,
    True,
    False,

    // 11.6.2.1 Keywords
    Break,
    Case,
    Catch,
    Class,
    Const,
    Continue,
    Debugger,
    Default,
    Delete,
    Do,
    Else,
    Export,
    Extends,
    Finally,
    For,
    Function,
    If,
    Import,
    In,
    Instanceof,
    New,
    Return,
    Super,
    Switch,
    This,
    Throw,
    Try,
    Typeof,
    Var,
    Void,
    While,
    With,

    // 11.6.2.2 Future Reserved Words
    Enum
}

impl Reserved {
    pub fn name(&self) -> &'static str {
        match *self {
            // 11.6.2 Reserved Words
            Reserved::Null       => "null",
            Reserved::True       => "true",
            Reserved::False      => "false",

            // 11.6.2.1 Keywords
            Reserved::Break      => "break",
            Reserved::Case       => "case",
            Reserved::Catch      => "catch",
            Reserved::Class      => "class",
            Reserved::Const      => "const",
            Reserved::Continue   => "continue",
            Reserved::Debugger   => "debugger",
            Reserved::Default    => "default",
            Reserved::Delete     => "delete",
            Reserved::Do         => "do",
            Reserved::Else       => "else",
            Reserved::Export     => "export",
            Reserved::Extends    => "extends",
            Reserved::Finally    => "finally",
            Reserved::For        => "for",
            Reserved::Function   => "function",
            Reserved::If         => "if",
            Reserved::Import     => "import",
            Reserved::In         => "in",
            Reserved::Instanceof => "instanceof",
            Reserved::New        => "new",
            Reserved::Return     => "return",
            Reserved::Super      => "super",
            Reserved::Switch     => "switch",
            Reserved::This       => "this",
            Reserved::Throw      => "throw",
            Reserved::Try        => "try",
            Reserved::Typeof     => "typeof",
            Reserved::Var        => "var",
            Reserved::Void       => "void",
            Reserved::While      => "while",
            Reserved::With       => "with",

            // 11.6.2.2 Future Reserved Words
            Reserved::Enum       => "enum"
        }
    }

    pub fn into_string(self) -> String {
        format!("{}", self.name())
    }
}

// Contextually reserved words and special identifier names.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub enum Atom {
    Arguments,
    Async,
    Await,
    Eval,
    From,
    Get,
    Implements,
    Interface,
    Let,
    Of,
    Package,
    Private,
    Protected,
    Public,
    Set,
    Static,
    Target,
    Yield
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum Name {
    Atom(Atom),
    String(String)
}

impl Name {
    pub fn new(from: String) -> Name {
        match &from[..] {
            "arguments"  => Name::Atom(Atom::Arguments),
            "await"      => Name::Atom(Atom::Await),
            "eval"       => Name::Atom(Atom::Eval),
            "async"      => Name::Atom(Atom::Async),
            "from"       => Name::Atom(Atom::From),
            "get"        => Name::Atom(Atom::Get),
            "implements" => Name::Atom(Atom::Implements),
            "interface"  => Name::Atom(Atom::Interface),
            "let"        => Name::Atom(Atom::Let),
            "of"         => Name::Atom(Atom::Of),
            "package"    => Name::Atom(Atom::Package),
            "private"    => Name::Atom(Atom::Private),
            "protected"  => Name::Atom(Atom::Protected),
            "public"     => Name::Atom(Atom::Public),
            "set"        => Name::Atom(Atom::Set),
            "static"     => Name::Atom(Atom::Static),
            "target"     => Name::Atom(Atom::Target),
            "yield"      => Name::Atom(Atom::Yield),
            _            => Name::String(from)
        }
    }

    pub fn into_string(self) -> String {
        match self {
            Name::Atom(atom) => format!("{}", atom.name()),
            Name::String(s) => s
        }
    }
}

impl Atom {
    pub fn name(self) -> &'static str {
        match self {
            Atom::Arguments  => "arguments",
            Atom::Await      => "await",
            Atom::Eval       => "eval",
            Atom::Async      => "async",
            Atom::From       => "from",
            Atom::Get        => "get",
            Atom::Implements => "implements",
            Atom::Interface  => "interface",
            Atom::Let        => "let",
            Atom::Of         => "of",
            Atom::Package    => "package",
            Atom::Private    => "private",
            Atom::Protected  => "protected",
            Atom::Public     => "public",
            Atom::Set        => "set",
            Atom::Static     => "static",
            Atom::Target     => "target",
            Atom::Yield      => "yield"
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Token {
    pub location: Span,
    pub newline: bool,    // was there a newline between the preceding token and this one?
    pub value: TokenData
}

impl Token {
    pub fn to_binop(&self, allow_in: bool) -> Option<Binop> {
        match self.value {
            TokenData::Star => Some(BinopTag::Times.tracked(self.location())),
            TokenData::Slash => Some(BinopTag::Div.tracked(self.location())),
            TokenData::Mod => Some(BinopTag::Mod.tracked(self.location())),
            TokenData::Plus => Some(BinopTag::Plus.tracked(self.location())),
            TokenData::Minus => Some(BinopTag::Minus.tracked(self.location())),
            TokenData::LShift => Some(BinopTag::LShift.tracked(self.location())),
            TokenData::RShift => Some(BinopTag::RShift.tracked(self.location())),
            TokenData::URShift => Some(BinopTag::URShift.tracked(self.location())),
            TokenData::LAngle => Some(BinopTag::Lt.tracked(self.location())),
            TokenData::RAngle => Some(BinopTag::Gt.tracked(self.location())),
            TokenData::LEq => Some(BinopTag::LEq.tracked(self.location())),
            TokenData::GEq => Some(BinopTag::GEq.tracked(self.location())),
            TokenData::Reserved(Reserved::Instanceof) => Some(BinopTag::Instanceof.tracked(self.location())),
            TokenData::Reserved(Reserved::In) => {
                if allow_in {
                    Some(BinopTag::In.tracked(self.location()))
                } else {
                    None
                }
            }
            TokenData::Eq => Some(BinopTag::Eq.tracked(self.location())),
            TokenData::NEq => Some(BinopTag::NEq.tracked(self.location())),
            TokenData::StrictEq => Some(BinopTag::StrictEq.tracked(self.location())),
            TokenData::StrictNEq => Some(BinopTag::StrictNEq.tracked(self.location())),
            TokenData::BitAnd => Some(BinopTag::BitAnd.tracked(self.location())),
            TokenData::BitXor => Some(BinopTag::BitXor.tracked(self.location())),
            TokenData::BitOr => Some(BinopTag::BitOr.tracked(self.location())),
            _ => None
        }
    }

    pub fn to_logop(&self) -> Option<Logop> {
        match self.value {
            TokenData::LogicalAnd => Some(LogopTag::And.tracked(self.location())),
            TokenData::LogicalOr => Some(LogopTag::Or.tracked(self.location())),
            _ => None
        }
    }

    pub fn to_assop(&self) -> Option<Assop> {
        match self.value {
            TokenData::Assign => Some(AssopTag::Eq.tracked(self.location())),
            TokenData::PlusAssign => Some(AssopTag::PlusEq.tracked(self.location())),
            TokenData::MinusAssign => Some(AssopTag::MinusEq.tracked(self.location())),
            TokenData::StarAssign => Some(AssopTag::TimesEq.tracked(self.location())),
            TokenData::SlashAssign => Some(AssopTag::DivEq.tracked(self.location())),
            TokenData::ModAssign => Some(AssopTag::ModEq.tracked(self.location())),
            TokenData::LShiftAssign => Some(AssopTag::LShiftEq.tracked(self.location())),
            TokenData::RShiftAssign => Some(AssopTag::RShiftEq.tracked(self.location())),
            TokenData::URShiftAssign => Some(AssopTag::URShiftEq.tracked(self.location())),
            TokenData::BitAndAssign => Some(AssopTag::BitAndEq.tracked(self.location())),
            TokenData::BitOrAssign => Some(AssopTag::BitOrEq.tracked(self.location())),
            TokenData::BitXorAssign => Some(AssopTag::BitXorEq.tracked(self.location())),
            _ => None
        }
    }
}

impl Track for Token {
    fn location(&self) -> Option<Span> {
        Some(self.location)
    }
}

impl Token {
    pub fn new(start: Posn, end: Posn, value: TokenData) -> Token {
        Token {
            location: Span { start: start, end: end },
            newline: false,
            value: value
        }
    }
}

#[derive(Debug, PartialEq)]
#[allow(dead_code)]
pub enum TokenData {
    Reserved(Reserved),

    // 11.7 Punctuators
    LBrace,
    RBrace,
    LParen,
    RParen,
    LBrack,
    RBrack,
    Dot,
    //Ellipsis,
    Semi,
    Comma,
    LAngle,
    RAngle,
    LEq,
    GEq,
    Eq,
    NEq,
    StrictEq,
    StrictNEq,
    Plus,
    Minus,
    Star,
    Mod,
    Slash,
    Inc,
    Dec,
    LShift,
    RShift,
    URShift,
    BitAnd,
    BitOr,
    BitXor,
    Bang,
    Tilde,
    LogicalAnd,
    LogicalOr,
    Question,
    Colon,
    Assign,
    PlusAssign,
    MinusAssign,
    StarAssign,
    SlashAssign,
    ModAssign,
    LShiftAssign,
    RShiftAssign,
    URShiftAssign,
    BitAndAssign,
    BitOrAssign,
    BitXorAssign,
    Arrow,

    Number(NumberLiteral),

    String(StringLiteral),
    RegExp(String, Vec<char>),

    Identifier(Name),

    EOF
}

#[derive(Debug, PartialEq)]
pub enum StringDelimiter {
    Single,
    Double
}

#[derive(Debug)]
pub struct StringLiteral {
    pub source: String,
    pub delimiter: StringDelimiter
}

impl PartialEq for StringLiteral {
    fn eq(&self, other: &Self) -> bool {
        let mut l = self.source.chars();
        let mut r = other.source.chars();
        let mut ocl = l.next();
        let mut ocr = r.next();
        loop {
            match (ocl, ocr) {
                (None,       None)                 => { return true; }
                (Some(cl),   Some(cr)) if cl == cr => { ocl = l.next(); ocr = r.next(); }
                (Some('\\'), Some(_))              => { ocl = l.next(); }
                (Some(_),    Some('\\'))           => { ocr = r.next(); }
                _                                  => { return false; }
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum NumberLiteral {
    DecimalInt(String, Option<Exp>),
    RadixInt(Radix, String),
    Float(Option<String>, Option<String>, Option<Exp>)
}

fn format_sign(sign: &Option<Sign>) -> String {
    match *sign {
        Some(Sign::Minus) => format!("-"),
        _ => format!("")
    }
}

fn format_int(src: &Option<String>) -> String {
    match *src {
        None => format!(""),
        Some(ref s) => format!("{}", s)
    }
}

impl NumberLiteral {
    pub fn value(&self) -> Option<f64> {
        match *self {
            NumberLiteral::DecimalInt(ref mantissa, None) => {
                let i: Option<i64> = mantissa.parse().ok();
                i.map(|i| i as f64)
            }
            NumberLiteral::DecimalInt(ref mantissa, Some(Exp { ref sign, ref value, .. })) => {
                let s = format!("{}e{}{}", mantissa, format_sign(sign), value);
                let i: Option<i64> = s.parse().ok();
                i.map(|i| i as f64)
            }
            NumberLiteral::RadixInt(ref radix, ref src) => {
                i64::from_str_radix(&src[..], radix.value()).ok().map(|i| i as f64)
            }
            NumberLiteral::Float(ref ip, ref fp, None) => {
                format!("{}.{}", format_int(ip), format_int(fp)).parse().ok()
            }
            NumberLiteral::Float(ref ip, ref fp, Some(Exp { ref sign, ref value, .. })) => {
                format!("{}.{}e{}{}", format_int(ip), format_int(fp), format_sign(sign), value).parse().ok()
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Exp {
    pub e: CharCase,
    pub sign: Option<Sign>,
    pub value: String
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Radix {
    Bin(CharCase),
    Oct(Option<CharCase>),
    Hex(CharCase)
}

impl Radix {
    pub fn value(&self) -> u32 {
        match *self {
            Radix::Bin(_) => 2,
            Radix::Oct(_) => 8,
            Radix::Hex(_) => 16
        }
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum CharCase {
    LowerCase,
    UpperCase
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Sign {
    Plus,
    Minus
}
<commit_msg>bugfix: only allow elided escape in string literal for the opposite delimiter character<commit_after>use track::*;
use ast::{Binop, BinopTag, Logop, LogopTag, Assop, AssopTag};

// Unconditionally reserved words.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum Reserved {
    // 11.6.2 Reserved Words
    Null,
    True,
    False,

    // 11.6.2.1 Keywords
    Break,
    Case,
    Catch,
    Class,
    Const,
    Continue,
    Debugger,
    Default,
    Delete,
    Do,
    Else,
    Export,
    Extends,
    Finally,
    For,
    Function,
    If,
    Import,
    In,
    Instanceof,
    New,
    Return,
    Super,
    Switch,
    This,
    Throw,
    Try,
    Typeof,
    Var,
    Void,
    While,
    With,

    // 11.6.2.2 Future Reserved Words
    Enum
}

impl Reserved {
    pub fn name(&self) -> &'static str {
        match *self {
            // 11.6.2 Reserved Words
            Reserved::Null       => "null",
            Reserved::True       => "true",
            Reserved::False      => "false",

            // 11.6.2.1 Keywords
            Reserved::Break      => "break",
            Reserved::Case       => "case",
            Reserved::Catch      => "catch",
            Reserved::Class      => "class",
            Reserved::Const      => "const",
            Reserved::Continue   => "continue",
            Reserved::Debugger   => "debugger",
            Reserved::Default    => "default",
            Reserved::Delete     => "delete",
            Reserved::Do         => "do",
            Reserved::Else       => "else",
            Reserved::Export     => "export",
            Reserved::Extends    => "extends",
            Reserved::Finally    => "finally",
            Reserved::For        => "for",
            Reserved::Function   => "function",
            Reserved::If         => "if",
            Reserved::Import     => "import",
            Reserved::In         => "in",
            Reserved::Instanceof => "instanceof",
            Reserved::New        => "new",
            Reserved::Return     => "return",
            Reserved::Super      => "super",
            Reserved::Switch     => "switch",
            Reserved::This       => "this",
            Reserved::Throw      => "throw",
            Reserved::Try        => "try",
            Reserved::Typeof     => "typeof",
            Reserved::Var        => "var",
            Reserved::Void       => "void",
            Reserved::While      => "while",
            Reserved::With       => "with",

            // 11.6.2.2 Future Reserved Words
            Reserved::Enum       => "enum"
        }
    }

    pub fn into_string(self) -> String {
        format!("{}", self.name())
    }
}

// Contextually reserved words and special identifier names.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub enum Atom {
    Arguments,
    Async,
    Await,
    Eval,
    From,
    Get,
    Implements,
    Interface,
    Let,
    Of,
    Package,
    Private,
    Protected,
    Public,
    Set,
    Static,
    Target,
    Yield
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum Name {
    Atom(Atom),
    String(String)
}

impl Name {
    pub fn new(from: String) -> Name {
        match &from[..] {
            "arguments"  => Name::Atom(Atom::Arguments),
            "await"      => Name::Atom(Atom::Await),
            "eval"       => Name::Atom(Atom::Eval),
            "async"      => Name::Atom(Atom::Async),
            "from"       => Name::Atom(Atom::From),
            "get"        => Name::Atom(Atom::Get),
            "implements" => Name::Atom(Atom::Implements),
            "interface"  => Name::Atom(Atom::Interface),
            "let"        => Name::Atom(Atom::Let),
            "of"         => Name::Atom(Atom::Of),
            "package"    => Name::Atom(Atom::Package),
            "private"    => Name::Atom(Atom::Private),
            "protected"  => Name::Atom(Atom::Protected),
            "public"     => Name::Atom(Atom::Public),
            "set"        => Name::Atom(Atom::Set),
            "static"     => Name::Atom(Atom::Static),
            "target"     => Name::Atom(Atom::Target),
            "yield"      => Name::Atom(Atom::Yield),
            _            => Name::String(from)
        }
    }

    pub fn into_string(self) -> String {
        match self {
            Name::Atom(atom) => format!("{}", atom.name()),
            Name::String(s) => s
        }
    }
}

impl Atom {
    pub fn name(self) -> &'static str {
        match self {
            Atom::Arguments  => "arguments",
            Atom::Await      => "await",
            Atom::Eval       => "eval",
            Atom::Async      => "async",
            Atom::From       => "from",
            Atom::Get        => "get",
            Atom::Implements => "implements",
            Atom::Interface  => "interface",
            Atom::Let        => "let",
            Atom::Of         => "of",
            Atom::Package    => "package",
            Atom::Private    => "private",
            Atom::Protected  => "protected",
            Atom::Public     => "public",
            Atom::Set        => "set",
            Atom::Static     => "static",
            Atom::Target     => "target",
            Atom::Yield      => "yield"
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Token {
    pub location: Span,
    pub newline: bool,    // was there a newline between the preceding token and this one?
    pub value: TokenData
}

impl Token {
    pub fn to_binop(&self, allow_in: bool) -> Option<Binop> {
        match self.value {
            TokenData::Star => Some(BinopTag::Times.tracked(self.location())),
            TokenData::Slash => Some(BinopTag::Div.tracked(self.location())),
            TokenData::Mod => Some(BinopTag::Mod.tracked(self.location())),
            TokenData::Plus => Some(BinopTag::Plus.tracked(self.location())),
            TokenData::Minus => Some(BinopTag::Minus.tracked(self.location())),
            TokenData::LShift => Some(BinopTag::LShift.tracked(self.location())),
            TokenData::RShift => Some(BinopTag::RShift.tracked(self.location())),
            TokenData::URShift => Some(BinopTag::URShift.tracked(self.location())),
            TokenData::LAngle => Some(BinopTag::Lt.tracked(self.location())),
            TokenData::RAngle => Some(BinopTag::Gt.tracked(self.location())),
            TokenData::LEq => Some(BinopTag::LEq.tracked(self.location())),
            TokenData::GEq => Some(BinopTag::GEq.tracked(self.location())),
            TokenData::Reserved(Reserved::Instanceof) => Some(BinopTag::Instanceof.tracked(self.location())),
            TokenData::Reserved(Reserved::In) => {
                if allow_in {
                    Some(BinopTag::In.tracked(self.location()))
                } else {
                    None
                }
            }
            TokenData::Eq => Some(BinopTag::Eq.tracked(self.location())),
            TokenData::NEq => Some(BinopTag::NEq.tracked(self.location())),
            TokenData::StrictEq => Some(BinopTag::StrictEq.tracked(self.location())),
            TokenData::StrictNEq => Some(BinopTag::StrictNEq.tracked(self.location())),
            TokenData::BitAnd => Some(BinopTag::BitAnd.tracked(self.location())),
            TokenData::BitXor => Some(BinopTag::BitXor.tracked(self.location())),
            TokenData::BitOr => Some(BinopTag::BitOr.tracked(self.location())),
            _ => None
        }
    }

    pub fn to_logop(&self) -> Option<Logop> {
        match self.value {
            TokenData::LogicalAnd => Some(LogopTag::And.tracked(self.location())),
            TokenData::LogicalOr => Some(LogopTag::Or.tracked(self.location())),
            _ => None
        }
    }

    pub fn to_assop(&self) -> Option<Assop> {
        match self.value {
            TokenData::Assign => Some(AssopTag::Eq.tracked(self.location())),
            TokenData::PlusAssign => Some(AssopTag::PlusEq.tracked(self.location())),
            TokenData::MinusAssign => Some(AssopTag::MinusEq.tracked(self.location())),
            TokenData::StarAssign => Some(AssopTag::TimesEq.tracked(self.location())),
            TokenData::SlashAssign => Some(AssopTag::DivEq.tracked(self.location())),
            TokenData::ModAssign => Some(AssopTag::ModEq.tracked(self.location())),
            TokenData::LShiftAssign => Some(AssopTag::LShiftEq.tracked(self.location())),
            TokenData::RShiftAssign => Some(AssopTag::RShiftEq.tracked(self.location())),
            TokenData::URShiftAssign => Some(AssopTag::URShiftEq.tracked(self.location())),
            TokenData::BitAndAssign => Some(AssopTag::BitAndEq.tracked(self.location())),
            TokenData::BitOrAssign => Some(AssopTag::BitOrEq.tracked(self.location())),
            TokenData::BitXorAssign => Some(AssopTag::BitXorEq.tracked(self.location())),
            _ => None
        }
    }
}

impl Track for Token {
    fn location(&self) -> Option<Span> {
        Some(self.location)
    }
}

impl Token {
    pub fn new(start: Posn, end: Posn, value: TokenData) -> Token {
        Token {
            location: Span { start: start, end: end },
            newline: false,
            value: value
        }
    }
}

#[derive(Debug, PartialEq)]
#[allow(dead_code)]
pub enum TokenData {
    Reserved(Reserved),

    // 11.7 Punctuators
    LBrace,
    RBrace,
    LParen,
    RParen,
    LBrack,
    RBrack,
    Dot,
    //Ellipsis,
    Semi,
    Comma,
    LAngle,
    RAngle,
    LEq,
    GEq,
    Eq,
    NEq,
    StrictEq,
    StrictNEq,
    Plus,
    Minus,
    Star,
    Mod,
    Slash,
    Inc,
    Dec,
    LShift,
    RShift,
    URShift,
    BitAnd,
    BitOr,
    BitXor,
    Bang,
    Tilde,
    LogicalAnd,
    LogicalOr,
    Question,
    Colon,
    Assign,
    PlusAssign,
    MinusAssign,
    StarAssign,
    SlashAssign,
    ModAssign,
    LShiftAssign,
    RShiftAssign,
    URShiftAssign,
    BitAndAssign,
    BitOrAssign,
    BitXorAssign,
    Arrow,

    Number(NumberLiteral),

    String(StringLiteral),
    RegExp(String, Vec<char>),

    Identifier(Name),

    EOF
}

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum StringDelimiter {
    Single,
    Double
}

impl StringDelimiter {
    pub fn opposite(self) -> StringDelimiter {
        match self {
            StringDelimiter::Single => StringDelimiter::Double,
            StringDelimiter::Double => StringDelimiter::Single
        }
    }

    pub fn to_char(self) -> char {
        match self {
            StringDelimiter::Single => '\'',
            StringDelimiter::Double => '"'
        }
    }
}

#[derive(Debug)]
pub struct StringLiteral {
    pub source: String,
    pub delimiter: StringDelimiter
}

impl PartialEq for StringLiteral {
    fn eq(&self, other: &Self) -> bool {
        let ql = self.delimiter.opposite().to_char();
        let qr = other.delimiter.opposite().to_char();
        let mut l = self.source.chars();
        let mut r = other.source.chars();
        let mut ocl = l.next();
        let mut ocr = r.next();
        loop {
            match (ocl, ocr) {
                (None,       None)                   => { return true; }
                (Some(cl),   Some(cr)) if cl == cr   => { ocl = l.next(); ocr = r.next(); }
                (Some('\\'), Some(cr)) if cr == qr   => { ocl = l.next(); }
                (Some(cl),   Some('\\')) if cl == ql => { ocr = r.next(); }
                _                                    => { return false; }
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum NumberLiteral {
    DecimalInt(String, Option<Exp>),
    RadixInt(Radix, String),
    Float(Option<String>, Option<String>, Option<Exp>)
}

fn format_sign(sign: &Option<Sign>) -> String {
    match *sign {
        Some(Sign::Minus) => format!("-"),
        _ => format!("")
    }
}

fn format_int(src: &Option<String>) -> String {
    match *src {
        None => format!(""),
        Some(ref s) => format!("{}", s)
    }
}

impl NumberLiteral {
    pub fn value(&self) -> Option<f64> {
        match *self {
            NumberLiteral::DecimalInt(ref mantissa, None) => {
                let i: Option<i64> = mantissa.parse().ok();
                i.map(|i| i as f64)
            }
            NumberLiteral::DecimalInt(ref mantissa, Some(Exp { ref sign, ref value, .. })) => {
                let s = format!("{}e{}{}", mantissa, format_sign(sign), value);
                let i: Option<i64> = s.parse().ok();
                i.map(|i| i as f64)
            }
            NumberLiteral::RadixInt(ref radix, ref src) => {
                i64::from_str_radix(&src[..], radix.value()).ok().map(|i| i as f64)
            }
            NumberLiteral::Float(ref ip, ref fp, None) => {
                format!("{}.{}", format_int(ip), format_int(fp)).parse().ok()
            }
            NumberLiteral::Float(ref ip, ref fp, Some(Exp { ref sign, ref value, .. })) => {
                format!("{}.{}e{}{}", format_int(ip), format_int(fp), format_sign(sign), value).parse().ok()
            }
        }
    }
}

#[derive(Debug, PartialEq)]
pub struct Exp {
    pub e: CharCase,
    pub sign: Option<Sign>,
    pub value: String
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Radix {
    Bin(CharCase),
    Oct(Option<CharCase>),
    Hex(CharCase)
}

impl Radix {
    pub fn value(&self) -> u32 {
        match *self {
            Radix::Bin(_) => 2,
            Radix::Oct(_) => 8,
            Radix::Hex(_) => 16
        }
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum CharCase {
    LowerCase,
    UpperCase
}

#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Sign {
    Plus,
    Minus
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::VecDeque;
use std::io::{self, ErrorKind, Write};
use std::net::SocketAddr;
use std::rc::Rc;
use std::time::{Duration, Instant};

use bytes::Bytes;
use bytes::buf::FromBuf;
use futures::{Async, Future, Poll};
use futures::task::{self, Task};
use kcp::{Error as KcpError, Kcp, KcpResult, get_conv};
use tokio_core::net::UdpSocket;
use tokio_core::reactor::Handle;

use config::KcpConfig;

struct KcpOutputInner {
    udp: Rc<UdpSocket>,
    task: Option<Task>,
    pkt_queue: VecDeque<(SocketAddr, Bytes)>,
    is_finished: bool,
}

impl KcpOutputInner {
    fn new(udp: Rc<UdpSocket>) -> KcpOutputInner {
        KcpOutputInner {
            udp: udp,
            task: None,
            pkt_queue: VecDeque::new(),
            is_finished: false,
        }
    }

    fn notify(&mut self) {
        if let Some(task) = self.task.take() {
            task.notify();
        }
    }

    fn push_packet(&mut self, pkt: Bytes, peer: SocketAddr) {
        self.pkt_queue.push_back((peer, pkt));
        self.notify();
    }

    fn close(&mut self) {
        self.is_finished = true;
        self.notify();
    }

    fn is_empty(&self) -> bool {
        self.pkt_queue.is_empty()
    }

    fn send_or_push(&mut self, buf: &[u8], peer: &SocketAddr) -> io::Result<usize> {
        if self.is_empty() {
            match self.udp.send_to(buf, peer) {
                Ok(n) => {
                    trace!("[SEND] UDP peer={} conv={} size={} {:?}",
                           peer,
                           get_conv(buf),
                           buf.len(),
                           ::debug::BsDebug(buf));
                    if n != buf.len() {
                        error!("[SEND] UDP Sent size={}, but packet is size={}", n, buf.len());
                    }
                    return Ok(n);
                }
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {}
                Err(err) => return Err(err),
            }
        }

        trace!("[SEND] Delay UDP peer={} conv={} size={} qsize={} {:?}",
               peer,
               get_conv(buf),
               buf.len(),
               self.pkt_queue.len(),
               ::debug::BsDebug(buf));

        self.push_packet(Bytes::from_buf(buf), *peer);
        Ok(buf.len())
    }
}

impl Drop for KcpOutputInner {
    fn drop(&mut self) {
        self.close();
    }
}

struct KcpOutputQueue {
    inner: Rc<RefCell<KcpOutputInner>>,
}

impl Future for KcpOutputQueue {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        let mut inner = self.inner.borrow_mut();

        while !inner.pkt_queue.is_empty() {
            {
                let &(ref peer, ref pkt) = &inner.pkt_queue[0];
                let n = try_nb!(inner.udp.send_to(&*pkt, peer));
                trace!("[SEND] Delayed UDP peer={} conv={} size={} {:?}", peer, get_conv(pkt), pkt.len(), pkt);
                if n != pkt.len() {
                    error!("[SEND] Delayed Sent size={}, but packet is size={}", n, pkt.len());
                }
            }

            let _ = inner.pkt_queue.pop_front();
        }

        if inner.is_finished {
            Ok(Async::Ready(()))
        } else {
            inner.task = Some(task::current());
            Ok(Async::NotReady)
        }
    }
}

#[derive(Clone)]
pub struct KcpOutputHandle {
    inner: Rc<RefCell<KcpOutputInner>>,
}

impl KcpOutputHandle {
    pub fn new(udp: Rc<UdpSocket>, handle: &Handle) -> KcpOutputHandle {
        let inner = KcpOutputInner::new(udp);
        let inner = Rc::new(RefCell::new(inner));
        let queue = KcpOutputQueue { inner: inner.clone() };
        handle.spawn(queue.map_err(move |err| {
                                       error!("[SEND] UDP output queue failed, err: {:?}", err);
                                   }));
        KcpOutputHandle { inner: inner }
    }

    fn send_to(&mut self, buf: &[u8], peer: &SocketAddr) -> io::Result<usize> {
        let mut inner = self.inner.borrow_mut();
        inner.send_or_push(buf, peer)
    }

    fn udp(&self) -> Rc<UdpSocket> {
        let inner = self.inner.borrow();
        inner.udp.clone()
    }
}

pub struct KcpOutput {
    inner: KcpOutputHandle,
    peer: SocketAddr,
}

impl KcpOutput {
    pub fn new(udp: Rc<UdpSocket>, peer: SocketAddr, handle: &Handle) -> KcpOutput {
        KcpOutput::new_with_handle(KcpOutputHandle::new(udp, handle), peer)
    }

    pub fn new_with_handle(h: KcpOutputHandle, peer: SocketAddr) -> KcpOutput {
        KcpOutput {
            inner: h,
            peer: peer,
        }
    }

    fn udp(&self) -> Rc<UdpSocket> {
        self.inner.udp()
    }
}

impl Write for KcpOutput {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.send_to(buf, &self.peer)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

struct KcpCell {
    kcp: Kcp<KcpOutput>,
    last_update: Instant,
    is_closed: bool,
    send_task: Option<Task>,
    udp: Rc<UdpSocket>,
    recv_buf: Vec<u8>,
    expired: bool,
    sent_first: bool,
    flush_write: bool,
    flush_ack_input: bool,
}

impl Drop for KcpCell {
    fn drop(&mut self) {
        let _ = self.kcp.flush();
    }
}

impl KcpCell {
    fn input(&mut self, buf: &[u8]) -> KcpResult<()> {
        match self.kcp.input(buf) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(());
            }
            Err(err) => return Err(err),
        }
        // self.kcp.input(buf)?;
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(())
    }

    fn input_self(&mut self, n: usize) -> KcpResult<()> {
        match self.kcp.input(&self.recv_buf[..n]) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(());
            }
            Err(err) => return Err(err),
        }
        // self.kcp.input(&self.recv_buf[..n])?;
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(())
    }

    fn fetch(&mut self) -> KcpResult<()> {
        if let Async::NotReady = self.udp.poll_read() {
            return Ok(());
        }

        // Initialize with MTU
        if self.recv_buf.len() < self.kcp.mtu() {
            let mtu = self.kcp.mtu();
            self.recv_buf.resize(mtu, 0);
        }

        let n = match self.udp.recv_from(&mut self.recv_buf) {
            Ok((n, _)) => n,
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                return Ok(());
            }
            Err(err) => return Err(From::from(err)),
        };

        // Ah, we got something
        trace!("[RECV] Fetch. SharedKcp recv size={} {:?}", n, ::debug::BsDebug(&self.recv_buf[..n]));
        self.input_self(n)
    }
}

#[derive(Clone)]
pub struct SharedKcp {
    inner: Rc<RefCell<KcpCell>>,
}

impl SharedKcp {
    pub fn new(c: &KcpConfig,
               conv: u32,
               udp: Rc<UdpSocket>,
               peer: SocketAddr,
               handle: &Handle,
               stream: bool)
               -> SharedKcp {
        let output = KcpOutput::new(udp, peer, handle);
        SharedKcp::new_with_output(c, conv, output, stream)
    }

    pub fn new_with_output(c: &KcpConfig, conv: u32, output: KcpOutput, stream: bool) -> SharedKcp {
        let udp = output.udp();
        let mut kcp = if stream {
            Kcp::new_stream(conv, output)
        } else {
            Kcp::new(conv, output)
        };
        c.apply_config(&mut kcp);

        // Ask server to allocate one
        if conv == 0 {
            kcp.input_conv();
        }

        SharedKcp {
            inner: Rc::new(RefCell::new(KcpCell {
                                            kcp: kcp,
                                            last_update: Instant::now(),
                                            is_closed: false,
                                            send_task: None,
                                            udp: udp,
                                            recv_buf: Vec::new(), // Do not initialize it yet.
                                            expired: false,
                                            sent_first: false,
                                            flush_write: c.flush_write,
                                            flush_ack_input: c.flush_acks_input,
                                        })),
        }
    }

    /// Call every time you got data from transmission
    pub fn input(&mut self, buf: &[u8]) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.input(buf)
    }

    /// Recv then input, ignore WouldBlock
    pub fn fetch(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.fetch()
    }

    /// Call if you want to send some data
    pub fn send(&mut self, buf: &[u8]) -> KcpResult<usize> {
        let mut inner = self.inner.borrow_mut();

        // If:
        //     1. Have sent the first packet (asking for conv)
        //     2. Too many pending packets
        if inner.sent_first && (inner.kcp.wait_snd() >= inner.kcp.snd_wnd() as usize || inner.kcp.waiting_conv()) {
            trace!("[SEND] waitsnd={} sndwnd={} excceeded or waiting conv={}",
                   inner.kcp.wait_snd(),
                   inner.kcp.snd_wnd(),
                   inner.kcp.waiting_conv());
            inner.send_task = Some(task::current());
            return Err(From::from(io::Error::new(ErrorKind::WouldBlock, "too many pending packets")));
        }

        if !inner.sent_first {
            assert!(buf.len() <= inner.kcp.mss() as usize, "First packet must be smaller than mss");
        }

        let n = inner.kcp.send(buf)?;
        inner.sent_first = true;
        inner.last_update = Instant::now();

        if inner.flush_write {
            inner.kcp.flush()?;
        }

        Ok(n)
    }

    /// Try to notify writable
    pub fn try_notify_writable(&mut self) {
        let mut inner = self.inner.borrow_mut();
        if inner.kcp.wait_snd() < inner.kcp.snd_wnd() as usize && !inner.kcp.waiting_conv() {
            if let Some(task) = inner.send_task.take() {
                task.notify();
            }
        }
    }

    /// Call if you want to get some data
    /// Always call right after input
    pub fn recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        let mut inner = self.inner.borrow_mut();
        if inner.expired {
            // If it is already expired, return EOF
            return Ok(0);
        }

        let n = inner.kcp.recv(buf)?;
        inner.last_update = Instant::now();
        Ok(n)
    }

    /// Call if you want to flush all pending data in queue
    pub fn flush(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.kcp.flush()?;
        inner.last_update = Instant::now();
        Ok(())
    }

    /// Tell me how long this session have no interaction
    pub fn elapsed(&self) -> Duration {
        let inner = self.inner.borrow();
        inner.last_update.elapsed()
    }

    /// Make this session expire, all read apis will return 0 (EOF)
    /// It will flush the buffer when it is called
    pub fn set_expired(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.expired = true;
        inner.kcp.flush()
    }

    /// Call in every tick
    /// Returns when to call this function again
    pub fn update(&mut self) -> KcpResult<Instant> {
        let mut inner = self.inner.borrow_mut();
        inner.kcp.update(::current())?;
        let next = inner.kcp.check(::current());
        Ok(Instant::now() + Duration::from_millis(next as u64))
        // Ok(Instant::now() + Duration::from_millis(5))
    }

    /// Check if send queue is empty
    pub fn has_waitsnd(&self) -> bool {
        let inner = self.inner.borrow();
        inner.kcp.wait_snd() != 0
    }

    /// Get mtu
    pub fn mtu(&self) -> usize {
        let inner = self.inner.borrow();
        inner.kcp.mtu()
    }

    /// Set is close
    pub fn close(&mut self) {
        let mut inner = self.inner.borrow_mut();
        trace!("[KCP] Set close, conv={}", inner.kcp.conv());
        inner.is_closed = true;
    }

    /// Check if it is closed
    pub fn is_closed(&self) -> bool {
        let inner = self.inner.borrow();
        inner.is_closed
    }

    /// Check if it can read
    pub fn can_read(&self) -> bool {
        let inner = self.inner.borrow();
        inner.kcp.peeksize().unwrap_or(0) != 0
    }

    /// Peek
    pub fn peeksize(&self) -> usize {
        let inner = self.inner.borrow();
        inner.kcp.peeksize().unwrap_or(0)
    }

    /// Get conv
    pub fn conv(&self) -> u32 {
        let inner = self.inner.borrow();
        inner.kcp.conv()
    }
}
<commit_msg>Check send buffer size only if waiting for allocating conv<commit_after>use std::cell::RefCell;
use std::collections::VecDeque;
use std::io::{self, ErrorKind, Write};
use std::net::SocketAddr;
use std::rc::Rc;
use std::time::{Duration, Instant};

use bytes::Bytes;
use bytes::buf::FromBuf;
use futures::{Async, Future, Poll};
use futures::task::{self, Task};
use kcp::{Error as KcpError, Kcp, KcpResult, get_conv};
use tokio_core::net::UdpSocket;
use tokio_core::reactor::Handle;

use config::KcpConfig;

struct KcpOutputInner {
    udp: Rc<UdpSocket>,
    task: Option<Task>,
    pkt_queue: VecDeque<(SocketAddr, Bytes)>,
    is_finished: bool,
}

impl KcpOutputInner {
    fn new(udp: Rc<UdpSocket>) -> KcpOutputInner {
        KcpOutputInner {
            udp: udp,
            task: None,
            pkt_queue: VecDeque::new(),
            is_finished: false,
        }
    }

    fn notify(&mut self) {
        if let Some(task) = self.task.take() {
            task.notify();
        }
    }

    fn push_packet(&mut self, pkt: Bytes, peer: SocketAddr) {
        self.pkt_queue.push_back((peer, pkt));
        self.notify();
    }

    fn close(&mut self) {
        self.is_finished = true;
        self.notify();
    }

    fn is_empty(&self) -> bool {
        self.pkt_queue.is_empty()
    }

    fn send_or_push(&mut self, buf: &[u8], peer: &SocketAddr) -> io::Result<usize> {
        if self.is_empty() {
            match self.udp.send_to(buf, peer) {
                Ok(n) => {
                    trace!("[SEND] UDP peer={} conv={} size={} {:?}",
                           peer,
                           get_conv(buf),
                           buf.len(),
                           ::debug::BsDebug(buf));
                    if n != buf.len() {
                        error!("[SEND] UDP Sent size={}, but packet is size={}", n, buf.len());
                    }
                    return Ok(n);
                }
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {}
                Err(err) => return Err(err),
            }
        }

        trace!("[SEND] Delay UDP peer={} conv={} size={} qsize={} {:?}",
               peer,
               get_conv(buf),
               buf.len(),
               self.pkt_queue.len(),
               ::debug::BsDebug(buf));

        self.push_packet(Bytes::from_buf(buf), *peer);
        Ok(buf.len())
    }
}

impl Drop for KcpOutputInner {
    fn drop(&mut self) {
        self.close();
    }
}

struct KcpOutputQueue {
    inner: Rc<RefCell<KcpOutputInner>>,
}

impl Future for KcpOutputQueue {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        let mut inner = self.inner.borrow_mut();

        while !inner.pkt_queue.is_empty() {
            {
                let &(ref peer, ref pkt) = &inner.pkt_queue[0];
                let n = try_nb!(inner.udp.send_to(&*pkt, peer));
                trace!("[SEND] Delayed UDP peer={} conv={} size={} {:?}", peer, get_conv(pkt), pkt.len(), pkt);
                if n != pkt.len() {
                    error!("[SEND] Delayed Sent size={}, but packet is size={}", n, pkt.len());
                }
            }

            let _ = inner.pkt_queue.pop_front();
        }

        if inner.is_finished {
            Ok(Async::Ready(()))
        } else {
            inner.task = Some(task::current());
            Ok(Async::NotReady)
        }
    }
}

#[derive(Clone)]
pub struct KcpOutputHandle {
    inner: Rc<RefCell<KcpOutputInner>>,
}

impl KcpOutputHandle {
    pub fn new(udp: Rc<UdpSocket>, handle: &Handle) -> KcpOutputHandle {
        let inner = KcpOutputInner::new(udp);
        let inner = Rc::new(RefCell::new(inner));
        let queue = KcpOutputQueue { inner: inner.clone() };
        handle.spawn(queue.map_err(move |err| {
                                       error!("[SEND] UDP output queue failed, err: {:?}", err);
                                   }));
        KcpOutputHandle { inner: inner }
    }

    fn send_to(&mut self, buf: &[u8], peer: &SocketAddr) -> io::Result<usize> {
        let mut inner = self.inner.borrow_mut();
        inner.send_or_push(buf, peer)
    }

    fn udp(&self) -> Rc<UdpSocket> {
        let inner = self.inner.borrow();
        inner.udp.clone()
    }
}

pub struct KcpOutput {
    inner: KcpOutputHandle,
    peer: SocketAddr,
}

impl KcpOutput {
    pub fn new(udp: Rc<UdpSocket>, peer: SocketAddr, handle: &Handle) -> KcpOutput {
        KcpOutput::new_with_handle(KcpOutputHandle::new(udp, handle), peer)
    }

    pub fn new_with_handle(h: KcpOutputHandle, peer: SocketAddr) -> KcpOutput {
        KcpOutput {
            inner: h,
            peer: peer,
        }
    }

    fn udp(&self) -> Rc<UdpSocket> {
        self.inner.udp()
    }
}

impl Write for KcpOutput {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.send_to(buf, &self.peer)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

struct KcpCell {
    kcp: Kcp<KcpOutput>,
    last_update: Instant,
    is_closed: bool,
    send_task: Option<Task>,
    udp: Rc<UdpSocket>,
    recv_buf: Vec<u8>,
    expired: bool,
    sent_first: bool,
    flush_write: bool,
    flush_ack_input: bool,
}

impl Drop for KcpCell {
    fn drop(&mut self) {
        let _ = self.kcp.flush();
    }
}

impl KcpCell {
    fn input(&mut self, buf: &[u8]) -> KcpResult<()> {
        match self.kcp.input(buf) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(());
            }
            Err(err) => return Err(err),
        }
        // self.kcp.input(buf)?;
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(())
    }

    fn input_self(&mut self, n: usize) -> KcpResult<()> {
        match self.kcp.input(&self.recv_buf[..n]) {
            Ok(..) => {}
            Err(KcpError::ConvInconsistent(expected, actual)) => {
                trace!("[INPUT] Conv expected={} actual={} ignored", expected, actual);
                return Ok(());
            }
            Err(err) => return Err(err),
        }
        // self.kcp.input(&self.recv_buf[..n])?;
        self.last_update = Instant::now();

        if self.flush_ack_input {
            self.kcp.flush_ack()?;
        }

        Ok(())
    }

    fn fetch(&mut self) -> KcpResult<()> {
        if let Async::NotReady = self.udp.poll_read() {
            return Ok(());
        }

        // Initialize with MTU
        if self.recv_buf.len() < self.kcp.mtu() {
            let mtu = self.kcp.mtu();
            self.recv_buf.resize(mtu, 0);
        }

        let n = match self.udp.recv_from(&mut self.recv_buf) {
            Ok((n, _)) => n,
            Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                return Ok(());
            }
            Err(err) => return Err(From::from(err)),
        };

        // Ah, we got something
        trace!("[RECV] Fetch. SharedKcp recv size={} {:?}", n, ::debug::BsDebug(&self.recv_buf[..n]));
        self.input_self(n)
    }
}

#[derive(Clone)]
pub struct SharedKcp {
    inner: Rc<RefCell<KcpCell>>,
}

impl SharedKcp {
    pub fn new(c: &KcpConfig,
               conv: u32,
               udp: Rc<UdpSocket>,
               peer: SocketAddr,
               handle: &Handle,
               stream: bool)
               -> SharedKcp {
        let output = KcpOutput::new(udp, peer, handle);
        SharedKcp::new_with_output(c, conv, output, stream)
    }

    pub fn new_with_output(c: &KcpConfig, conv: u32, output: KcpOutput, stream: bool) -> SharedKcp {
        let udp = output.udp();
        let mut kcp = if stream {
            Kcp::new_stream(conv, output)
        } else {
            Kcp::new(conv, output)
        };
        c.apply_config(&mut kcp);

        // Ask server to allocate one
        if conv == 0 {
            kcp.input_conv();
        }

        SharedKcp {
            inner: Rc::new(RefCell::new(KcpCell {
                                            kcp: kcp,
                                            last_update: Instant::now(),
                                            is_closed: false,
                                            send_task: None,
                                            udp: udp,
                                            recv_buf: Vec::new(), // Do not initialize it yet.
                                            expired: false,
                                            sent_first: false,
                                            flush_write: c.flush_write,
                                            flush_ack_input: c.flush_acks_input,
                                        })),
        }
    }

    /// Call every time you got data from transmission
    pub fn input(&mut self, buf: &[u8]) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.input(buf)
    }

    /// Recv then input, ignore WouldBlock
    pub fn fetch(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.fetch()
    }

    /// Call if you want to send some data
    pub fn send(&mut self, buf: &[u8]) -> KcpResult<usize> {
        let mut inner = self.inner.borrow_mut();

        // If:
        //     1. Have sent the first packet (asking for conv)
        //     2. Too many pending packets
        if inner.sent_first && (inner.kcp.wait_snd() >= inner.kcp.snd_wnd() as usize || inner.kcp.waiting_conv()) {
            trace!("[SEND] waitsnd={} sndwnd={} excceeded or waiting conv={}",
                   inner.kcp.wait_snd(),
                   inner.kcp.snd_wnd(),
                   inner.kcp.waiting_conv());
            inner.send_task = Some(task::current());
            return Err(From::from(io::Error::new(ErrorKind::WouldBlock, "too many pending packets")));
        }

        if !inner.sent_first && inner.kcp.waiting_conv() {
            assert!(buf.len() <= inner.kcp.mss() as usize, "First packet must be smaller than mss");
        }

        let n = inner.kcp.send(buf)?;
        inner.sent_first = true;
        inner.last_update = Instant::now();

        if inner.flush_write {
            inner.kcp.flush()?;
        }

        Ok(n)
    }

    /// Try to notify writable
    pub fn try_notify_writable(&mut self) {
        let mut inner = self.inner.borrow_mut();
        if inner.kcp.wait_snd() < inner.kcp.snd_wnd() as usize && !inner.kcp.waiting_conv() {
            if let Some(task) = inner.send_task.take() {
                task.notify();
            }
        }
    }

    /// Call if you want to get some data
    /// Always call right after input
    pub fn recv(&mut self, buf: &mut [u8]) -> KcpResult<usize> {
        let mut inner = self.inner.borrow_mut();
        if inner.expired {
            // If it is already expired, return EOF
            return Ok(0);
        }

        let n = inner.kcp.recv(buf)?;
        inner.last_update = Instant::now();
        Ok(n)
    }

    /// Call if you want to flush all pending data in queue
    pub fn flush(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.kcp.flush()?;
        inner.last_update = Instant::now();
        Ok(())
    }

    /// Tell me how long this session have no interaction
    pub fn elapsed(&self) -> Duration {
        let inner = self.inner.borrow();
        inner.last_update.elapsed()
    }

    /// Make this session expire, all read apis will return 0 (EOF)
    /// It will flush the buffer when it is called
    pub fn set_expired(&mut self) -> KcpResult<()> {
        let mut inner = self.inner.borrow_mut();
        inner.expired = true;
        inner.kcp.flush()
    }

    /// Call in every tick
    /// Returns when to call this function again
    pub fn update(&mut self) -> KcpResult<Instant> {
        let mut inner = self.inner.borrow_mut();
        inner.kcp.update(::current())?;
        let next = inner.kcp.check(::current());
        Ok(Instant::now() + Duration::from_millis(next as u64))
        // Ok(Instant::now() + Duration::from_millis(5))
    }

    /// Check if send queue is empty
    pub fn has_waitsnd(&self) -> bool {
        let inner = self.inner.borrow();
        inner.kcp.wait_snd() != 0
    }

    /// Get mtu
    pub fn mtu(&self) -> usize {
        let inner = self.inner.borrow();
        inner.kcp.mtu()
    }

    /// Set is close
    pub fn close(&mut self) {
        let mut inner = self.inner.borrow_mut();
        trace!("[KCP] Set close, conv={}", inner.kcp.conv());
        inner.is_closed = true;
    }

    /// Check if it is closed
    pub fn is_closed(&self) -> bool {
        let inner = self.inner.borrow();
        inner.is_closed
    }

    /// Check if it can read
    pub fn can_read(&self) -> bool {
        let inner = self.inner.borrow();
        inner.kcp.peeksize().unwrap_or(0) != 0
    }

    /// Peek
    pub fn peeksize(&self) -> usize {
        let inner = self.inner.borrow();
        inner.kcp.peeksize().unwrap_or(0)
    }

    /// Get conv
    pub fn conv(&self) -> u32 {
        let inner = self.inner.borrow();
        inner.kcp.conv()
    }
}
<|endoftext|>