<commit_before>use std::sync::atomic::Ordering::{Acquire, Release, Relaxed};
use std::sync::atomic::AtomicBool;
use std::{ptr, mem};
use std::thread::{self, Thread};

use epoch::{self, Atomic, Owned, Shared};
use CachePadded;

/// A Michael-Scott lock-free queue, with support for blocking `pop`s.
///
/// Usable with any number of producers and consumers.
// The representation here is a singly-linked list, with a sentinel
// node at the front. In general the `tail` pointer may lag behind the
// actual tail. Non-sentinel nodes are either all `Data` or all
// `Blocked` (requests for data from blocked threads).
#[derive(Debug)]
pub struct MsQueue<T> {
    head: CachePadded<Atomic<Node<T>>>,
    tail: CachePadded<Atomic<Node<T>>>,
}

#[derive(Debug)]
struct Node<T> {
    payload: Payload<T>,
    next: Atomic<Node<T>>,
}

#[derive(Debug)]
enum Payload<T> {
    /// A node with actual data that can be popped.
    Data(T),
    /// A node representing a blocked request for data.
    Blocked(*mut Signal<T>),
}

/// A blocked request for data, which includes a slot to write the data.
#[derive(Debug)]
struct Signal<T> {
    /// Thread to unpark when data is ready.
    thread: Thread,
    /// The actual data, when available.
    data: Option<T>,
    /// Is the data ready? Needed to cope with spurious wakeups.
    ready: AtomicBool,
}

impl<T> Node<T> {
    fn is_data(&self) -> bool {
        if let Payload::Data(_) = self.payload { true } else { false }
    }
}

// Any particular `T` should never accessed concurrently, so no need
// for Sync.
unsafe impl<T: Send> Sync for MsQueue<T> {}
unsafe impl<T: Send> Send for MsQueue<T> {}

impl<T> MsQueue<T> {
    /// Create a new, empty queue.
    pub fn new() -> MsQueue<T> {
        let q = MsQueue {
            head: CachePadded::new(Atomic::null()),
            tail: CachePadded::new(Atomic::null()),
        };
        let sentinel = Owned::new(Node {
            payload: Payload::Data(unsafe { mem::uninitialized() }),
            next: Atomic::null(),
        });
        let guard = epoch::pin();
        let sentinel = q.head.store_and_ref(sentinel, Relaxed, &guard);
        q.tail.store_shared(Some(sentinel), Relaxed);
        q
    }

    #[inline(always)]
    /// Attempt to atomically place `n` into the `next` pointer of `onto`.
    ///
    /// If unsuccessful, returns ownership of `n`, possibly updating
    /// the queue's `tail` pointer.
    fn push_internal(&self,
                     guard: &epoch::Guard,
                     onto: Shared<Node<T>>,
                     n: Owned<Node<T>>)
                     -> Result<(), Owned<Node<T>>>
    {
        // is `onto` the actual tail?
        if let Some(next) = onto.next.load(Acquire, guard) {
            // if not, try to "help" by moving the tail pointer forward
            self.tail.cas_shared(Some(onto), Some(next), Release);
            Err(n)
        } else {
            // looks like the actual tail; attempt to link in `n`
            onto.next.cas_and_ref(None, n, Release, guard).map(|shared| {
                // try to move the tail pointer forward
                self.tail.cas_shared(Some(onto), Some(shared), Release);
            })
        }
    }

    /// Add `t` to the back of the queue, possibly waking up threads
    /// blocked on `pop`.
    pub fn push(&self, t: T) {
        /// We may or may not need to allocate a node; once we do,
        /// we cache that allocation.
        enum Cache<T> {
            Data(T),
            Node(Owned<Node<T>>),
        }

        impl<T> Cache<T> {
            /// Extract the node if cached, or allocate if not.
            fn into_node(self) -> Owned<Node<T>> {
                match self {
                    Cache::Data(t) => {
                        Owned::new(Node {
                            payload: Payload::Data(t),
                            next: Atomic::null()
                        })
                    }
                    Cache::Node(n) => n
                }
            }

            /// Extract the data from the cache, deallocating any cached node.
            fn into_data(self) -> T {
                match self {
                    Cache::Data(t) => t,
                    Cache::Node(node) => {
                        match node.into_inner().payload {
                            Payload::Data(t) => t,
                            _ => unreachable!(),
                        }
                    }
                }
            }
        }

        let mut cache = Cache::Data(t); // don't allocate up front
        let guard = epoch::pin();

        loop {
            // We push onto the tail, so we'll start optimistically by looking
            // there first.
            let tail = self.tail.load(Acquire, &guard).unwrap();

            // Is the queue in Data mode (empty queues can be viewed as either mode)?
            if tail.is_data() ||
                self.head.load(Relaxed, &guard).unwrap().as_raw() == tail.as_raw()
            {
                // Attempt to push onto the `tail` snapshot; fails if
                // `tail.next` has changed, which will always be the case if the
                // queue has transitioned to blocking mode.
                match self.push_internal(&guard, tail, cache.into_node()) {
                    Ok(_) => return,
                    Err(n) => {
                        // replace the cache, retry whole thing
                        cache = Cache::Node(n)
                    }
                }
            } else {
                // Queue is in blocking mode. Attempt to unblock a thread.
                let head = self.head.load(Acquire, &guard).unwrap();
                // Get a handle on the first blocked node. Racy, so queue might
                // be empty or in data mode by the time we see it.
                let request = head.next.load(Acquire, &guard).and_then(|next| {
                    match next.payload {
                        Payload::Blocked(signal) => Some((next, signal)),
                        Payload::Data(_) => None,
                    }
                });
                if let Some((blocked_node, signal)) = request {
                    // race to dequeue the node
                    if self.head.cas_shared(Some(head), Some(blocked_node), Release) {
                        unsafe {
                            // signal the thread
                            (*signal).data = Some(cache.into_data());
                            (*signal).ready.store(true, Release);
                            (*signal).thread.unpark();
                            guard.unlinked(head);
                            return;
                        }
                    }
                }
            }
        }
    }

    #[inline(always)]
    // Attempt to pop a data node. `Ok(None)` if queue is empty or in blocking
    // mode; `Err(())` if lost race to pop.
    fn pop_internal(&self, guard: &epoch::Guard) -> Result<Option<T>, ()> {
        let head = self.head.load(Acquire, guard).unwrap();
        if let Some(next) = head.next.load(Acquire, guard) {
            if let Payload::Data(ref t) = next.payload {
                unsafe {
                    if self.head.cas_shared(Some(head), Some(next), Release) {
                        guard.unlinked(head);
                        Ok(Some(ptr::read(t)))
                    } else {
                        Err(())
                    }
                }
            } else {
                Ok(None)
            }
        } else {
            Ok(None)
        }
    }

    /// Check if this queue is empty.
    pub fn is_empty(&self) -> bool {
        let guard = epoch::pin();
        let head = self.head.load(Acquire, &guard).unwrap();

        if let Some(next) = head.next.load(Acquire, &guard) {
            if let Payload::Data(_) = next.payload {
                false
            } else {
                true
            }
        } else {
            true
        }
    }

    /// Attempt to dequeue from the front.
    ///
    /// Returns `None` if the queue is observed to be empty.
    pub fn try_pop(&self) -> Option<T> {
        let guard = epoch::pin();
        loop {
            if let Ok(r) = self.pop_internal(&guard) {
                return r;
            }
        }
    }

    /// Dequeue an element from the front of the queue, blocking if the queue is
    /// empty.
    pub fn pop(&self) -> T {
        let guard = epoch::pin();

        // Fast path: keep retrying until we observe that the queue has no data,
        // avoiding the allocation of a blocked node.
        loop {
            match self.pop_internal(&guard) {
                Ok(Some(r)) => {
                    return r;
                }
                Ok(None) => {
                    break;
                }
                Err(()) => {}
            }
        }

        // The signal gets to live on the stack, since this stack frame will be
        // blocked until receiving the signal.
        let mut signal = Signal {
            thread: thread::current(),
            data: None,
            ready: AtomicBool::new(false),
        };

        // Go ahead and allocate the blocked node; chances are, we'll need it.
        let mut node = Owned::new(Node {
            payload: Payload::Blocked(&mut signal),
            next: Atomic::null(),
        });

        loop {
            // try a normal pop
            if let Ok(Some(r)) = self.pop_internal(&guard) {
                return r;
            }

            // At this point, we believe the queue is empty/blocked.
            // Snapshot the tail, onto which we want to push a blocked node.
            let tail = self.tail.load(Relaxed, &guard).unwrap();

            // Double-check that we're in blocking mode
            if tail.is_data() {
                // The current tail is in data mode, so we probably need to abort.
                // BUT, it might be the sentinel, so check for that first.
                let head = self.head.load(Relaxed, &guard).unwrap();
                if tail.is_data() && tail.as_raw() != head.as_raw() { continue; }
            }

            // At this point, the tail snapshot is either a blocked node deep in
            // the queue, the sentinel, or no longer accessible from the queue.
            // In *ALL* of these cases, if we succeed in pushing onto the
            // snapshot, we know we are maintaining the core invariant: all
            // reachable, non-sentinel nodes have the same payload mode, in this
            // case, blocked.
            match self.push_internal(&guard, tail, node) {
                Ok(()) => {
                    while !signal.ready.load(Acquire) {
                        thread::park();
                    }
                    return signal.data.unwrap();
                }
                Err(n) => {
                    node = n;
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    const CONC_COUNT: i64 = 1000000;

    use scope;
    use super::*;

    #[test]
    fn push_try_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(37));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        q.push(48);
        assert_eq!(q.try_pop(), Some(37));
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(48));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.try_pop(), Some(i));
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.pop(), 37);
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), 37);
        assert_eq!(q.pop(), 48);
    }

    #[test]
    fn push_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.pop(), i);
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());

        scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.try_pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
    }

    #[test]
    fn push_try_pop_many_spmc() {
        fn recv(_t: i32, q: &MsQueue<i64>) {
            let mut cur = -1;
            for _i in 0..CONC_COUNT {
                if let Some(elem) = q.try_pop() {
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 { break }
                }
            }
        }

        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        let qr = &q;
        scope(|scope| {
            for i in 0..3 {
                scope.spawn(move || recv(i, qr));
            }

            scope.spawn(|| {
                for i in 0..CONC_COUNT {
                    q.push(i);
                }
            })
        });
    }

    #[test]
    fn push_try_pop_many_mpmc() {
        enum LR { Left(i64), Right(i64) }

        let q: MsQueue<LR> = MsQueue::new();
        assert!(q.is_empty());

        scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.try_pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        });
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();

        scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;
                while next < CONC_COUNT {
                    assert_eq!(q.pop(), next);
                    next += 1;
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
        assert!(q.is_empty());
    }

    #[test]
    fn is_empty_dont_pop() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(20);
        q.push(20);
        assert!(!q.is_empty());
        assert!(!q.is_empty());
        assert!(q.try_pop().is_some());
    }
}
<commit_msg>Fix ordering bug in MS queue (#100)<commit_after>use std::sync::atomic::Ordering::{Acquire, Release, Relaxed};
use std::sync::atomic::AtomicBool;
use std::{ptr, mem};
use std::thread::{self, Thread};

use epoch::{self, Atomic, Owned, Shared};
use CachePadded;

/// A Michael-Scott lock-free queue, with support for blocking `pop`s.
///
/// Usable with any number of producers and consumers.
// The representation here is a singly-linked list, with a sentinel
// node at the front. In general the `tail` pointer may lag behind the
// actual tail. Non-sentinel nodes are either all `Data` or all
// `Blocked` (requests for data from blocked threads).
#[derive(Debug)]
pub struct MsQueue<T> {
    head: CachePadded<Atomic<Node<T>>>,
    tail: CachePadded<Atomic<Node<T>>>,
}

#[derive(Debug)]
struct Node<T> {
    payload: Payload<T>,
    next: Atomic<Node<T>>,
}

#[derive(Debug)]
enum Payload<T> {
    /// A node with actual data that can be popped.
    Data(T),
    /// A node representing a blocked request for data.
    Blocked(*mut Signal<T>),
}

/// A blocked request for data, which includes a slot to write the data.
#[derive(Debug)]
struct Signal<T> {
    /// Thread to unpark when data is ready.
    thread: Thread,
    /// The actual data, when available.
    data: Option<T>,
    /// Is the data ready? Needed to cope with spurious wakeups.
    ready: AtomicBool,
}

impl<T> Node<T> {
    fn is_data(&self) -> bool {
        if let Payload::Data(_) = self.payload { true } else { false }
    }
}

// Any particular `T` should never accessed concurrently, so no need
// for Sync.
unsafe impl<T: Send> Sync for MsQueue<T> {}
unsafe impl<T: Send> Send for MsQueue<T> {}

impl<T> MsQueue<T> {
    /// Create a new, empty queue.
    pub fn new() -> MsQueue<T> {
        let q = MsQueue {
            head: CachePadded::new(Atomic::null()),
            tail: CachePadded::new(Atomic::null()),
        };
        let sentinel = Owned::new(Node {
            payload: Payload::Data(unsafe { mem::uninitialized() }),
            next: Atomic::null(),
        });
        let guard = epoch::pin();
        let sentinel = q.head.store_and_ref(sentinel, Relaxed, &guard);
        q.tail.store_shared(Some(sentinel), Relaxed);
        q
    }

    #[inline(always)]
    /// Attempt to atomically place `n` into the `next` pointer of `onto`.
    ///
    /// If unsuccessful, returns ownership of `n`, possibly updating
    /// the queue's `tail` pointer.
    fn push_internal(&self,
                     guard: &epoch::Guard,
                     onto: Shared<Node<T>>,
                     n: Owned<Node<T>>)
                     -> Result<(), Owned<Node<T>>>
    {
        // is `onto` the actual tail?
        if let Some(next) = onto.next.load(Acquire, guard) {
            // if not, try to "help" by moving the tail pointer forward
            self.tail.cas_shared(Some(onto), Some(next), Release);
            Err(n)
        } else {
            // looks like the actual tail; attempt to link in `n`
            onto.next.cas_and_ref(None, n, Release, guard).map(|shared| {
                // try to move the tail pointer forward
                self.tail.cas_shared(Some(onto), Some(shared), Release);
            })
        }
    }

    /// Add `t` to the back of the queue, possibly waking up threads
    /// blocked on `pop`.
    pub fn push(&self, t: T) {
        /// We may or may not need to allocate a node; once we do,
        /// we cache that allocation.
        enum Cache<T> {
            Data(T),
            Node(Owned<Node<T>>),
        }

        impl<T> Cache<T> {
            /// Extract the node if cached, or allocate if not.
            fn into_node(self) -> Owned<Node<T>> {
                match self {
                    Cache::Data(t) => {
                        Owned::new(Node {
                            payload: Payload::Data(t),
                            next: Atomic::null()
                        })
                    }
                    Cache::Node(n) => n
                }
            }

            /// Extract the data from the cache, deallocating any cached node.
            fn into_data(self) -> T {
                match self {
                    Cache::Data(t) => t,
                    Cache::Node(node) => {
                        match node.into_inner().payload {
                            Payload::Data(t) => t,
                            _ => unreachable!(),
                        }
                    }
                }
            }
        }

        let mut cache = Cache::Data(t); // don't allocate up front
        let guard = epoch::pin();

        loop {
            // We push onto the tail, so we'll start optimistically by looking
            // there first.
            let tail = self.tail.load(Acquire, &guard).unwrap();

            // Is the queue in Data mode (empty queues can be viewed as either mode)?
            if tail.is_data() ||
                self.head.load(Relaxed, &guard).unwrap().as_raw() == tail.as_raw()
            {
                // Attempt to push onto the `tail` snapshot; fails if
                // `tail.next` has changed, which will always be the case if the
                // queue has transitioned to blocking mode.
                match self.push_internal(&guard, tail, cache.into_node()) {
                    Ok(_) => return,
                    Err(n) => {
                        // replace the cache, retry whole thing
                        cache = Cache::Node(n)
                    }
                }
            } else {
                // Queue is in blocking mode. Attempt to unblock a thread.
                let head = self.head.load(Acquire, &guard).unwrap();
                // Get a handle on the first blocked node. Racy, so queue might
                // be empty or in data mode by the time we see it.
                let request = head.next.load(Acquire, &guard).and_then(|next| {
                    match next.payload {
                        Payload::Blocked(signal) => Some((next, signal)),
                        Payload::Data(_) => None,
                    }
                });
                if let Some((blocked_node, signal)) = request {
                    // race to dequeue the node
                    if self.head.cas_shared(Some(head), Some(blocked_node), Release) {
                        unsafe {
                            // signal the thread
                            (*signal).data = Some(cache.into_data());
                            (*signal).ready.store(true, Release);
                            (*signal).thread.unpark();
                            guard.unlinked(head);
                            return;
                        }
                    }
                }
            }
        }
    }

    #[inline(always)]
    // Attempt to pop a data node. `Ok(None)` if queue is empty or in blocking
    // mode; `Err(())` if lost race to pop.
    fn pop_internal(&self, guard: &epoch::Guard) -> Result<Option<T>, ()> {
        let head = self.head.load(Acquire, guard).unwrap();
        if let Some(next) = head.next.load(Acquire, guard) {
            if let Payload::Data(ref t) = next.payload {
                unsafe {
                    if self.head.cas_shared(Some(head), Some(next), Release) {
                        guard.unlinked(head);
                        Ok(Some(ptr::read(t)))
                    } else {
                        Err(())
                    }
                }
            } else {
                Ok(None)
            }
        } else {
            Ok(None)
        }
    }

    /// Check if this queue is empty.
    pub fn is_empty(&self) -> bool {
        let guard = epoch::pin();
        let head = self.head.load(Acquire, &guard).unwrap();

        if let Some(next) = head.next.load(Acquire, &guard) {
            if let Payload::Data(_) = next.payload {
                false
            } else {
                true
            }
        } else {
            true
        }
    }

    /// Attempt to dequeue from the front.
    ///
    /// Returns `None` if the queue is observed to be empty.
    pub fn try_pop(&self) -> Option<T> {
        let guard = epoch::pin();
        loop {
            if let Ok(r) = self.pop_internal(&guard) {
                return r;
            }
        }
    }

    /// Dequeue an element from the front of the queue, blocking if the queue is
    /// empty.
    pub fn pop(&self) -> T {
        let guard = epoch::pin();

        // Fast path: keep retrying until we observe that the queue has no data,
        // avoiding the allocation of a blocked node.
        loop {
            match self.pop_internal(&guard) {
                Ok(Some(r)) => {
                    return r;
                }
                Ok(None) => {
                    break;
                }
                Err(()) => {}
            }
        }

        // The signal gets to live on the stack, since this stack frame will be
        // blocked until receiving the signal.
        let mut signal = Signal {
            thread: thread::current(),
            data: None,
            ready: AtomicBool::new(false),
        };

        // Go ahead and allocate the blocked node; chances are, we'll need it.
        let mut node = Owned::new(Node {
            payload: Payload::Blocked(&mut signal),
            next: Atomic::null(),
        });

        loop {
            // try a normal pop
            if let Ok(Some(r)) = self.pop_internal(&guard) {
                return r;
            }

            // At this point, we believe the queue is empty/blocked.
            // Snapshot the tail, onto which we want to push a blocked node.
            let tail = self.tail.load(Acquire, &guard).unwrap();

            // Double-check that we're in blocking mode
            if tail.is_data() {
                // The current tail is in data mode, so we probably need to abort.
                // BUT, it might be the sentinel, so check for that first.
                let head = self.head.load(Relaxed, &guard).unwrap();
                if tail.is_data() && tail.as_raw() != head.as_raw() { continue; }
            }

            // At this point, the tail snapshot is either a blocked node deep in
            // the queue, the sentinel, or no longer accessible from the queue.
            // In *ALL* of these cases, if we succeed in pushing onto the
            // snapshot, we know we are maintaining the core invariant: all
            // reachable, non-sentinel nodes have the same payload mode, in this
            // case, blocked.
            match self.push_internal(&guard, tail, node) {
                Ok(()) => {
                    while !signal.ready.load(Acquire) {
                        thread::park();
                    }
                    return signal.data.unwrap();
                }
                Err(n) => {
                    node = n;
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    const CONC_COUNT: i64 = 1000000;

    use scope;
    use super::*;

    #[test]
    fn push_try_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(37));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        q.push(48);
        assert_eq!(q.try_pop(), Some(37));
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(48));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.try_pop(), Some(i));
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_1() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.pop(), 37);
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_2() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), 37);
        assert_eq!(q.pop(), 48);
    }

    #[test]
    fn push_pop_many_seq() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.pop(), i);
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());

        scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.try_pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
    }

    #[test]
    fn push_try_pop_many_spmc() {
        fn recv(_t: i32, q: &MsQueue<i64>) {
            let mut cur = -1;
            for _i in 0..CONC_COUNT {
                if let Some(elem) = q.try_pop() {
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 { break }
                }
            }
        }

        let q: MsQueue<i64> = MsQueue::new();
        assert!(q.is_empty());
        let qr = &q;
        scope(|scope| {
            for i in 0..3 {
                scope.spawn(move || recv(i, qr));
            }

            scope.spawn(|| {
                for i in 0..CONC_COUNT {
                    q.push(i);
                }
            })
        });
    }

    #[test]
    fn push_try_pop_many_mpmc() {
        enum LR { Left(i64), Right(i64) }

        let q: MsQueue<LR> = MsQueue::new();
        assert!(q.is_empty());

        scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|| {
                    for i in CONC_COUNT-1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.try_pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        });
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: MsQueue<i64> = MsQueue::new();

        scope(|scope| {
            scope.spawn(|| {
                let mut next = 0;
                while next < CONC_COUNT {
                    assert_eq!(q.pop(), next);
                    next += 1;
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        });
        assert!(q.is_empty());
    }

    #[test]
    fn is_empty_dont_pop() {
        let q: MsQueue<i64> = MsQueue::new();
        q.push(20);
        q.push(20);
        assert!(!q.is_empty());
        assert!(!q.is_empty());
        assert!(q.try_pop().is_some());
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::BTreeMap,
    path::{Path, PathBuf},
};

use crate::{background_jobs::Environment, uploaders::Uploader, util::errors::std_error_no_send};

use scopeguard::defer;
use swirl::PerformError;

/// Create CSV dumps of the public information in the database, wrap them in a
/// tarball and upload to S3.
#[swirl::background_job]
pub fn dump_db(
    env: &Environment,
    database_url: String,
    target_name: String,
) -> Result<(), PerformError> {
    let timestamp = chrono::Utc::now().format("%Y-%m-%d-%H%M%S").to_string();
    let export_dir = std::env::temp_dir().join("dump-db").join(timestamp);
    std::fs::create_dir_all(&export_dir)?;
    defer! {{
        std::fs::remove_dir_all(&export_dir).unwrap();
    }}
    let visibility_config = toml::from_str(include_str!("dump-db.toml")).unwrap();
    run_psql(&visibility_config, &database_url, &export_dir)?;
    let tarball = create_tarball(&export_dir)?;
    defer! {{
        std::fs::remove_file(&tarball).unwrap();
    }}
    upload_tarball(&tarball, &target_name, &env.uploader)?;
    println!("Database dump uploaded to {}.", &target_name);
    Ok(())
}

/// An enum indicating whether a column is included in the database dumps.
/// Public columns are included, private are not.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
enum ColumnVisibility {
    Private,
    Public,
}

/// Filtering information for a single table. The `filter` field is a valid SQL
/// expression used in a `WHERE` clause to filter the rows of the table. The
/// `columns` field maps column names to their respective visibilities.
#[derive(Clone, Debug, Deserialize)]
struct TableConfig {
    filter: Option<String>,
    columns: BTreeMap<String, ColumnVisibility>,
}

impl TableConfig {
    fn columns_str(&self) -> String {
        self.columns
            .iter()
            .filter(|&(_, &vis)| vis == ColumnVisibility::Public)
            .map(|(col, _)| format!("\"{}\"", col))
            .collect::<Vec<String>>()
            .join(", ")
    }

    fn view_sql(&self, table: &str) -> String {
        self.filter
            .as_ref()
            .map(|filter| {
                format!(
                    r#"
                    CREATE TEMPORARY VIEW "dump_db_{table}" AS (
                        SELECT {columns}
                        FROM "{table}"
                        WHERE {filter}
                    );
                    "#,
                    table = table,
                    columns = self.columns_str(),
                    filter = filter,
                )
            })
            .unwrap_or_default()
    }

    fn copy_sql(&self, table: &str) -> String {
        if self.filter.is_some() {
            format!(
                r#"
                \copy (SELECT * FROM "dump_db_{table}") TO '{table}.csv' WITH CSV HEADER
                "#,
                table = table,
            )
        } else {
            let cols_str = self.columns_str();
            if cols_str.is_empty() {
                String::new()
            } else {
                format!(
                    r#"
                    \copy "{table}" ({columns}) TO '{table}.csv' WITH CSV HEADER
                    "#,
                    table = table,
                    columns = cols_str,
                )
            }
        }
    }
}

/// Maps table names to the respective configurations. Used to load `dump_db.toml`.
#[derive(Clone, Debug, Deserialize)]
#[serde(transparent)]
struct VisibilityConfig(BTreeMap<String, TableConfig>);

impl VisibilityConfig {
    fn gen_psql_script(&self) -> String {
        let view_sql = self
            .0
            .iter()
            .map(|(table, config)| config.view_sql(table))
            .collect::<Vec<String>>()
            .concat();
        let copy_sql = self
            .0
            .iter()
            .map(|(table, config)| config.copy_sql(table))
            .collect::<Vec<String>>()
            .concat();
        format!(
            r#"
            BEGIN;
            {view_sql}
            COMMIT;
            BEGIN ISOLATION LEVEL SERIALIZABLE READ ONLY DEFERRABLE;
            {copy_sql}
            COMMIT;
            "#,
            view_sql = view_sql,
            copy_sql = copy_sql,
        )
    }
}

fn run_psql(
    config: &VisibilityConfig,
    database_url: &str,
    export_dir: &Path,
) -> Result<(), PerformError> {
    use std::io::prelude::*;
    use std::process::{Command, Stdio};

    let psql_script = config.gen_psql_script();
    // TODO Redirect stdout and stderr to avoid polluting the worker logs.
    let mut psql = Command::new("psql")
        .arg(database_url)
        .current_dir(export_dir)
        .stdin(Stdio::piped())
        .spawn()?;
    let mut stdin = psql.stdin.take().unwrap();
    stdin.write_all(psql_script.as_bytes())?;
    drop(stdin);
    psql.wait()?;
    Ok(())
}

fn create_tarball(export_dir: &Path) -> Result<PathBuf, PerformError> {
    let tarball_name = export_dir.with_extension("tar.gz");
    let tarball = std::fs::File::create(&tarball_name)?;
    let encoder = flate2::write::GzEncoder::new(tarball, flate2::Compression::default());
    let mut archive = tar::Builder::new(encoder);
    archive.append_dir_all(export_dir.file_name().unwrap(), &export_dir)?;
    Ok(tarball_name)
}

fn upload_tarball(
    tarball: &Path,
    target_name: &str,
    uploader: &Uploader,
) -> Result<(), PerformError> {
    let client = reqwest::Client::new();
    let tarfile = std::fs::File::open(tarball)?;
    let content_length = tarfile.metadata()?.len();
    // TODO Figure out the correct content type.
    uploader
        .upload(
            &client,
            target_name,
            tarfile,
            content_length,
            "application/gzip",
        )
        .map_err(std_error_no_send)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_util::pg_connection;
    use diesel::prelude::*;
    use std::collections::HashSet;

    /// Test whether the schema in the visibility configuration matches the test database.
    #[test]
    fn check_visibility_config() {
        let conn = pg_connection();
        let db_columns: HashSet<_> = get_db_columns(&conn)
            .into_iter()
            .map(|c| (c.table_name, c.column_name))
            .collect();
        let visibility_config: VisibilityConfig =
            toml::from_str(include_str!("dump-db.toml")).unwrap();
        let vis_columns: HashSet<_> = visibility_config
            .0
            .iter()
            .flat_map(|(table, config)| {
                config
                    .columns
                    .iter()
                    .map(move |(column, _)| (table.clone(), column.clone()))
            })
            .collect();
        let mut errors = vec![];
        for (table, col) in db_columns.difference(&vis_columns) {
            errors.push(format!(
                "No visibility information for columns {}.{}.",
                table, col
            ));
        }
        for (table, col) in vis_columns.difference(&db_columns) {
            errors.push(format!(
                "Column {}.{} does not exist in the database.",
                table, col
            ));
        }
        assert!(
            errors.is_empty(),
            "The visibility configuration does not match the database schema:\n{}",
            errors.join("\n  - "),
        );
    }

    mod information_schema {
        table! {
            information_schema.columns (table_schema, table_name, column_name) {
                table_schema -> Text,
                table_name -> Text,
                column_name -> Text,
                ordinal_position -> Integer,
            }
        }
    }

    #[derive(Debug, PartialEq, Queryable)]
    struct Column {
        table_name: String,
        column_name: String,
    }

    fn get_db_columns(conn: &PgConnection) -> Vec<Column> {
        use information_schema::columns::dsl::*;
        columns
            .select((table_name, column_name))
            .filter(table_schema.eq("public"))
            .order_by((table_name, ordinal_position))
            .load(conn)
            .unwrap()
    }
}
<commit_msg>Correct error handling for the psql call.<commit_after>use std::{
    collections::BTreeMap,
    path::{Path, PathBuf},
};

use crate::{background_jobs::Environment, uploaders::Uploader, util::errors::std_error_no_send};

use scopeguard::defer;
use swirl::PerformError;

/// Create CSV dumps of the public information in the database, wrap them in a
/// tarball and upload to S3.
#[swirl::background_job]
pub fn dump_db(
    env: &Environment,
    database_url: String,
    target_name: String,
) -> Result<(), PerformError> {
    let timestamp = chrono::Utc::now().format("%Y-%m-%d-%H%M%S").to_string();
    let export_dir = std::env::temp_dir().join("dump-db").join(timestamp);
    std::fs::create_dir_all(&export_dir)?;
    defer! {{
        std::fs::remove_dir_all(&export_dir).unwrap();
    }}
    let visibility_config = toml::from_str(include_str!("dump-db.toml")).unwrap();
    run_psql(&visibility_config, &database_url, &export_dir)?;
    let tarball = create_tarball(&export_dir)?;
    defer! {{
        std::fs::remove_file(&tarball).unwrap();
    }}
    upload_tarball(&tarball, &target_name, &env.uploader)?;
    println!("Database dump uploaded to {}.", &target_name);
    Ok(())
}

/// An enum indicating whether a column is included in the database dumps.
/// Public columns are included, private are not.
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
enum ColumnVisibility {
    Private,
    Public,
}

/// Filtering information for a single table. The `filter` field is a valid SQL
/// expression used in a `WHERE` clause to filter the rows of the table. The
/// `columns` field maps column names to their respective visibilities.
#[derive(Clone, Debug, Deserialize)]
struct TableConfig {
    filter: Option<String>,
    columns: BTreeMap<String, ColumnVisibility>,
}

impl TableConfig {
    fn columns_str(&self) -> String {
        self.columns
            .iter()
            .filter(|&(_, &vis)| vis == ColumnVisibility::Public)
            .map(|(col, _)| format!("\"{}\"", col))
            .collect::<Vec<String>>()
            .join(", ")
    }

    fn view_sql(&self, table: &str) -> String {
        self.filter
            .as_ref()
            .map(|filter| {
                format!(
                    r#"
                    CREATE TEMPORARY VIEW "dump_db_{table}" AS (
                        SELECT {columns}
                        FROM "{table}"
                        WHERE {filter}
                    );
                    "#,
                    table = table,
                    columns = self.columns_str(),
                    filter = filter,
                )
            })
            .unwrap_or_default()
    }

    fn copy_sql(&self, table: &str) -> String {
        if self.filter.is_some() {
            format!(
                r#"
                \copy (SELECT * FROM "dump_db_{table}") TO '{table}.csv' WITH CSV HEADER
                "#,
                table = table,
            )
        } else {
            let cols_str = self.columns_str();
            if cols_str.is_empty() {
                String::new()
            } else {
                format!(
                    r#"
                    \copy "{table}" ({columns}) TO '{table}.csv' WITH CSV HEADER
                    "#,
                    table = table,
                    columns = cols_str,
                )
            }
        }
    }
}

/// Maps table names to the respective configurations. Used to load `dump_db.toml`.
#[derive(Clone, Debug, Deserialize)]
#[serde(transparent)]
struct VisibilityConfig(BTreeMap<String, TableConfig>);

impl VisibilityConfig {
    fn gen_psql_script(&self) -> String {
        let view_sql = self
            .0
            .iter()
            .map(|(table, config)| config.view_sql(table))
            .collect::<Vec<String>>()
            .concat();
        let copy_sql = self
            .0
            .iter()
            .map(|(table, config)| config.copy_sql(table))
            .collect::<Vec<String>>()
            .concat();
        format!(
            r#"
            BEGIN;
            {view_sql}
            COMMIT;
            BEGIN ISOLATION LEVEL SERIALIZABLE READ ONLY DEFERRABLE;
            {copy_sql}
            COMMIT;
            "#,
            view_sql = view_sql,
            copy_sql = copy_sql,
        )
    }
}

fn run_psql(
    config: &VisibilityConfig,
    database_url: &str,
    export_dir: &Path,
) -> Result<(), PerformError> {
    use std::io::prelude::*;
    use std::process::{Command, Stdio};

    let psql_script = config.gen_psql_script();
    let mut psql = Command::new("psql")
        .arg(database_url)
        .current_dir(export_dir)
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::piped())
        .spawn()?;
    let mut stdin = psql.stdin.take().unwrap();
    let input_thread = std::thread::spawn(move || -> std::io::Result<()> {
        stdin.write_all(psql_script.as_bytes())?;
        Ok(())
    });
    let output = psql.wait_with_output()?;
    input_thread.join().unwrap()?;
    if !output.stderr.is_empty() {
        Err(format!(
            "Error while executing psql: {}",
            String::from_utf8_lossy(&output.stderr)
        ))?;
    }
    if !output.status.success() {
        Err("psql did not finish successfully.")?;
    }
    Ok(())
}

fn create_tarball(export_dir: &Path) -> Result<PathBuf, PerformError> {
    let tarball_name = export_dir.with_extension("tar.gz");
    let tarball = std::fs::File::create(&tarball_name)?;
    let encoder = flate2::write::GzEncoder::new(tarball, flate2::Compression::default());
    let mut archive = tar::Builder::new(encoder);
    archive.append_dir_all(export_dir.file_name().unwrap(), &export_dir)?;
    Ok(tarball_name)
}

fn upload_tarball(
    tarball: &Path,
    target_name: &str,
    uploader: &Uploader,
) -> Result<(), PerformError> {
    let client = reqwest::Client::new();
    let tarfile = std::fs::File::open(tarball)?;
    let content_length = tarfile.metadata()?.len();
    // TODO Figure out the correct content type.
    uploader
        .upload(
            &client,
            target_name,
            tarfile,
            content_length,
            "application/gzip",
        )
        .map_err(std_error_no_send)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_util::pg_connection;
    use diesel::prelude::*;
    use std::collections::HashSet;

    /// Test whether the schema in the visibility configuration matches the test database.
    #[test]
    fn check_visibility_config() {
        let conn = pg_connection();
        let db_columns: HashSet<_> = get_db_columns(&conn)
            .into_iter()
            .map(|c| (c.table_name, c.column_name))
            .collect();
        let visibility_config: VisibilityConfig =
            toml::from_str(include_str!("dump-db.toml")).unwrap();
        let vis_columns: HashSet<_> = visibility_config
            .0
            .iter()
            .flat_map(|(table, config)| {
                config
                    .columns
                    .iter()
                    .map(move |(column, _)| (table.clone(), column.clone()))
            })
            .collect();
        let mut errors = vec![];
        for (table, col) in db_columns.difference(&vis_columns) {
            errors.push(format!(
                "No visibility information for columns {}.{}.",
                table, col
            ));
        }
        for (table, col) in vis_columns.difference(&db_columns) {
            errors.push(format!(
                "Column {}.{} does not exist in the database.",
                table, col
            ));
        }
        assert!(
            errors.is_empty(),
            "The visibility configuration does not match the database schema:\n{}",
            errors.join("\n  - "),
        );
    }

    mod information_schema {
        table! {
            information_schema.columns (table_schema, table_name, column_name) {
                table_schema -> Text,
                table_name -> Text,
                column_name -> Text,
                ordinal_position -> Integer,
            }
        }
    }

    #[derive(Debug, PartialEq, Queryable)]
    struct Column {
        table_name: String,
        column_name: String,
    }

    fn get_db_columns(conn: &PgConnection) -> Vec<Column> {
        use information_schema::columns::dsl::*;
        columns
            .select((table_name, column_name))
            .filter(table_schema.eq("public"))
            .order_by((table_name, ordinal_position))
            .load(conn)
            .unwrap()
    }
}
<|endoftext|>


<commit_before>//! Term simplification.

pub use std::cmp::Ordering ;
use std::ops::Deref ;

use common::* ;

lazy_static! {
  static ref simpl_solver: RwLock<
    Option<::rsmt2::Solver<()>>
  > = RwLock::new(
    if conf.check_simpl {
      Some(
        conf.solver.spawn("check_simpl", (), & Instance::new()).unwrap()
      )
    } else {
      None
    }
  ) ;
}

/// Result of a simplification check between two terms.
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum SimplRes {
  /// Nothing.
  None,
  /// The two terms are directly related.
  Cmp(Ordering),
  /// The two terms are equivalent to another term.
  Yields(Term),
}
impl SimplRes {
  /// Inverts the result.
  pub fn invert(self) -> Self {
    use self::SimplRes::* ;
    match self {
      None => None,
      Cmp(ord) => Cmp( ord.reverse() ),
      Yields(term) => Yields(term),
    }
  }
  /// True if `self` is not `None`.
  #[inline]
  pub fn is_some(& self) -> bool {
    self != & SimplRes::None
  }
  /// True if `self` is `None`.
  #[inline]
  pub fn is_none(& self) -> bool {
    ! self.is_some()
  }
}

/// Simplifies a conjunction.
pub fn conj_vec_simpl(mut terms: Vec<Term>) -> Vec<Term> {
  let mut res = Vec::with_capacity( terms.len() ) ;

  'add_terms: while let Some(term) = terms.pop() {

    let mut cnt = 0 ;

    'check_current: while cnt < res.len() {
      use self::SimplRes::* ;
      use std::cmp::Ordering::* ;

      match conj_simpl(& term, & res[cnt]) {
        None => cnt += 1,

        Cmp(Less) |
        Cmp(Equal) => continue 'add_terms,

        Cmp(Greater) => {
          res.swap_remove(cnt) ;
        },

        Yields(term) => {
          res.swap_remove(cnt) ;
          terms.push(term) ;
          continue 'add_terms
        },
      }
    }

    res.push(term)

  }

  res.shrink_to_fit() ;

  res
}

/// Checks the conjunction of two terms for simplification.
///
/// Returns
///
/// - `None` if no conclusion was reached,
/// - `Cmp(Greater)` if `lhs => rhs`,
/// - `Cmp(Less)` if `lhs <= rhs`,
/// - `Cmp(Equal)` if `lhs` and `rhs` are equivalent.
/// - `Yields(term)` if `lhs and rhs` is equivalent to `term`.
///
/// So *greater* really means *stronger*.
///
/// # Examples
///
/// ```
/// use std::cmp::Ordering::* ;
/// use hoice::term::simplify::SimplRes::* ;
/// use hoice::term::simplify::conj_simpl ;
/// use hoice::term ;
///
/// let lhs = term::not(
///   term::lt(
///     term::int(0),
///     term::sub( vec![ term::int(3), term::int_var(0) ] )
///   )
/// ) ;
/// # println!("   {}", lhs) ;
/// let rhs = term::ge( term::int_var(0), term::int(3) ) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Equal) }
///
/// # println!("   {}", lhs) ;
/// let rhs = term::ge( term::int_var(0), term::int(7) ) ;
/// # println!("<= {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Less) }
///
/// # println!("   {}", rhs) ;
/// # println!("=> {}\n\n", lhs) ;
/// debug_assert_eq! { conj_simpl(& rhs, & lhs), Cmp(Greater) }
///
/// let lhs = term::gt( term::int_var(0), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Greater) }
///
/// let lhs = term::le( term::int_var(0), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// let expected = term::eq( term::int_var(0), term::int(7) ) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Yields(expected) }
///
/// let lhs = term::le( term::int_var(1), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), None }
///
/// let lhs = term::le( term::int_var(1), term::int(7) ) ;
/// let rhs = term::or(
///   vec![ lhs.clone(), term::eq( term::int_var(3), term::int(7) ) ]
/// ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Greater) }
///
/// # println!("   {}", rhs) ;
/// # println!("=> {}\n\n", lhs) ;
/// debug_assert_eq! { conj_simpl(& rhs, & lhs), Cmp(Less) }
/// ```
pub fn conj_simpl<T1, T2>(lhs: & T1, rhs: & T2) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {

  if conf.term_simpl == 0 {
    return SimplRes::None
  }

  let res = conj_simpl_2(lhs, rhs) ;

  if res.is_some() {
    if let Some(solver) = simpl_solver.write().unwrap().as_mut() {
      solver.push(1).unwrap() ;
      let mut vars = VarSet::new() ;
      lhs.iter(
        |term| if let Some(idx) = term.var_idx() {
          let is_new = vars.insert(idx) ;
          if is_new {
            solver.declare_const(& format!("{}", term), & term.typ()).unwrap()
          }
        }
      ) ;
      rhs.iter(
        |term| if let Some(idx) = term.var_idx() {
          let is_new = vars.insert(idx) ;
          if is_new {
            solver.declare_const(& format!("{}", term), & term.typ()).unwrap()
          }
        }
      ) ;
      use std::cmp::Ordering::* ;
      let check = match res {
        SimplRes::Cmp(Equal) => format!(
          "(= {} {})", lhs.deref(), rhs.deref()
        ),
        SimplRes::Cmp(Less) => format!(
          "(=> {} {})", rhs.deref(), lhs.deref()
        ),
        SimplRes::Cmp(Greater) => format!(
          "(=> {} {})", lhs.deref(), rhs.deref()
        ),
        SimplRes::Yields(ref term) => format!(
          "(= (and {} {}) {})", lhs.deref(), rhs.deref(), term
        ),
        SimplRes::None => unreachable!(),
      } ;

      solver.assert(& format!("(not {})", check)).unwrap() ;

      if solver.check_sat().unwrap() {
        log! { @0
          " " ;
          "{}", lhs.deref() ;
          "{}", rhs.deref() ;
          "{:?}", res ;
          " "
        }
        panic! { "simplification failure" }
      }

      solver.pop(1).unwrap()
    }
  }

  match res {
    SimplRes::Yields(_) if conf.term_simpl > 1 => SimplRes::None,
    res => res,
  }
}
pub fn conj_simpl_2<T1, T2>(lhs: & T1, rhs: & T2) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {
  use std::cmp::Ordering::* ;

  if let Some(args) = lhs.disj_inspect() {
    for lhs in args {
      let mut greater_count = 0 ;
      let mut yields = vec![] ;
      match int_conj_simpl(lhs, rhs, true) {
        SimplRes::Cmp(Equal) |
        SimplRes::Cmp(Less) => return SimplRes::Cmp(Less),
        SimplRes::Cmp(Greater) => greater_count += 1,

        SimplRes::Yields(term) => yields.push(term),
        SimplRes::None => break,
      }
      if yields.len() == args.len() {
        return SimplRes::Yields( term::or(yields) )
      } else if greater_count == args.len() {
        return SimplRes::Cmp(Greater)
      }
    }
  } else if let Some(args) = rhs.disj_inspect() {
    for rhs in args {
      let mut less_count = 0 ;
      let mut yields = vec![] ;
      match int_conj_simpl(lhs, rhs, true) {
        SimplRes::Cmp(Equal) |
        SimplRes::Cmp(Greater) => return SimplRes::Cmp(Greater),
        SimplRes::Cmp(Less) => less_count += 1,

        SimplRes::Yields(term) => yields.push(term),
        SimplRes::None => break,
      }
      if yields.len() == args.len() {
        return SimplRes::Yields( term::or(yields) )
      } else if less_count == args.len() {
        return SimplRes::Cmp(Greater)
      }
    }
  }

  let res = int_conj_simpl(lhs, rhs, true) ;
  // if res.is_some() {
  //   log! { @0
  //     "\n\n{}", lhs.deref() ;
  //     "{}", rhs.deref() ;
  //     "{:?}", res
  //   }
  // }

  res
}

fn int_conj_simpl<T1, T2>(lhs: & T1, rhs: & T2, conj: bool) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {
  use std::cmp::Ordering::* ;
  // Input boolean is true (false) for `lhs` => `rhs` (reversed).
  macro_rules! ord_of_bool {
    ($b:expr) => (
      if $b {
        SimplRes::Cmp(Greater)
      } else {
        SimplRes::Cmp(Less)
      }
    ) ;
  }

  let (lhs, rhs) = ( lhs.deref(), rhs.deref() ) ;

  // A term implies itself.
  if lhs == rhs { return SimplRes::Cmp(Equal) }

  match ( lhs.bool(), rhs.bool() ) {
    // True can only imply true.
    (Some(true), rhs) => return ord_of_bool!(
      rhs.unwrap_or(false)
    ),
    // False implies anything.
    (Some(false), _) => return ord_of_bool!(true),
    // False can only be implied by false.
    (lhs, Some(false)) => return ord_of_bool!(
      ! lhs.unwrap_or(true)
    ),
    // True is implied by anything.
    (_, Some(true)) => return ord_of_bool!(true),
    // Otherwise we don't know (yet).
    (None, None) => (),
  }

  // Only legal atoms are `vars >= cst` and `vars > cst`.
  let (
    lhs_op, lhs_vars, lhs_cst
  ) = if let Some((op, args)) = lhs.app_inspect() {
    if ! args[0].typ().is_arith() {
      return SimplRes::None
    }
    (op, & args[0], args[1].val().unwrap())
  } else {
    return SimplRes::None
  } ;
  let (
    rhs_op, rhs_vars, rhs_cst
  ) = if let Some((op, args)) = rhs.app_inspect() {
    if ! args[0].typ().is_arith() {
      return SimplRes::None
    }
    (op, & args[0], args[1].val().unwrap())
  } else {
    return SimplRes::None
  } ;

  if lhs_vars == rhs_vars {

    match (lhs_op, rhs_op) {
      (Op::Gt, Op::Gt) |
      (Op::Ge, Op::Ge) => return SimplRes::Cmp(
        lhs_cst.cmp(& rhs_cst).unwrap()
      ),

      (Op::Gt, Op::Ge) |
      (Op::Ge, Op::Gt) => if lhs_cst == rhs_cst {
        if lhs_op == Op::Gt {
          return SimplRes::Cmp(Less)
        } else {
          return SimplRes::Cmp(Greater)
        }
      } else {
        return SimplRes::Cmp(
          lhs_cst.cmp(& rhs_cst).unwrap().reverse()
        )
      },

      (Op::Eql, Op::Ge) |
      (Op::Eql, Op::Gt) => match lhs_cst.cmp(& rhs_cst) {
        Some(Less) => return SimplRes::Yields( term::fls() ),
        Some(Equal) if rhs_op == Op::Gt => if conj {
          return SimplRes::Yields( term::fls() )
        } else {
        },
        Some(Equal) |
        Some(Greater) => return SimplRes::Cmp(Greater),
        None => unreachable!(),
      },

      (Op::Ge, Op::Eql) |
      (Op::Gt, Op::Eql) => match rhs_cst.cmp(& lhs_cst) {
        Some(Less) => return SimplRes::Yields( term::fls() ),
        Some(Equal) if rhs_op == Op::Gt => return SimplRes::Yields(
          term::fls()
        ),
        Some(Equal) |
        Some(Greater) => return SimplRes::Cmp(Less),
        None => unreachable!(),
      },

      (Op::Eql, Op::Eql) => if rhs_cst == lhs_cst {
        return SimplRes::Cmp(Greater)
      } else {
        return SimplRes::Yields( term::fls() )
      },

      _ => unreachable!(),
    }

  } else

  if lhs_op == Op::Ge && rhs_op == Op::Ge
  && lhs_cst == rhs_cst.minus().unwrap()
  && rhs_vars == & term::u_minus( lhs_vars.clone() ) {

    if conj {
      return SimplRes::Yields(
        term::eq(lhs_vars.clone(), lhs_cst.to_term().unwrap())
      )
    } else {
      return SimplRes::Yields(
        term::tru()
      )
    }
  }

  SimplRes::None
}

<commit_msg>bug fix<commit_after>//! Term simplification.

pub use std::cmp::Ordering ;
use std::ops::Deref ;

use common::* ;

lazy_static! {
  static ref simpl_solver: RwLock<
    Option<::rsmt2::Solver<()>>
  > = RwLock::new(
    if conf.check_simpl {
      Some(
        conf.solver.spawn("check_simpl", (), & Instance::new()).unwrap()
      )
    } else {
      None
    }
  ) ;
}

/// Result of a simplification check between two terms.
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum SimplRes {
  /// Nothing.
  None,
  /// The two terms are directly related.
  Cmp(Ordering),
  /// The two terms are equivalent to another term.
  Yields(Term),
}
impl SimplRes {
  /// Inverts the result.
  pub fn invert(self) -> Self {
    use self::SimplRes::* ;
    match self {
      None => None,
      Cmp(ord) => Cmp( ord.reverse() ),
      Yields(term) => Yields(term),
    }
  }
  /// True if `self` is not `None`.
  #[inline]
  pub fn is_some(& self) -> bool {
    self != & SimplRes::None
  }
  /// True if `self` is `None`.
  #[inline]
  pub fn is_none(& self) -> bool {
    ! self.is_some()
  }
}

/// Simplifies a conjunction.
pub fn conj_vec_simpl(mut terms: Vec<Term>) -> Vec<Term> {
  let mut res = Vec::with_capacity( terms.len() ) ;

  'add_terms: while let Some(term) = terms.pop() {

    let mut cnt = 0 ;

    'check_current: while cnt < res.len() {
      use self::SimplRes::* ;
      use std::cmp::Ordering::* ;

      match conj_simpl(& term, & res[cnt]) {
        None => cnt += 1,

        Cmp(Less) |
        Cmp(Equal) => continue 'add_terms,

        Cmp(Greater) => {
          res.swap_remove(cnt) ;
        },

        Yields(term) => {
          res.swap_remove(cnt) ;
          terms.push(term) ;
          continue 'add_terms
        },
      }
    }

    res.push(term)

  }

  res.shrink_to_fit() ;

  res
}

/// Checks the conjunction of two terms for simplification.
///
/// Returns
///
/// - `None` if no conclusion was reached,
/// - `Cmp(Greater)` if `lhs => rhs`,
/// - `Cmp(Less)` if `lhs <= rhs`,
/// - `Cmp(Equal)` if `lhs` and `rhs` are equivalent.
/// - `Yields(term)` if `lhs and rhs` is equivalent to `term`.
///
/// So *greater* really means *stronger*.
///
/// # Examples
///
/// ```
/// use std::cmp::Ordering::* ;
/// use hoice::term::simplify::SimplRes::* ;
/// use hoice::term::simplify::conj_simpl ;
/// use hoice::term ;
///
/// let lhs = term::not(
///   term::lt(
///     term::int(0),
///     term::sub( vec![ term::int(3), term::int_var(0) ] )
///   )
/// ) ;
/// # println!("   {}", lhs) ;
/// let rhs = term::ge( term::int_var(0), term::int(3) ) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Equal) }
///
/// # println!("   {}", lhs) ;
/// let rhs = term::ge( term::int_var(0), term::int(7) ) ;
/// # println!("<= {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Less) }
///
/// # println!("   {}", rhs) ;
/// # println!("=> {}\n\n", lhs) ;
/// debug_assert_eq! { conj_simpl(& rhs, & lhs), Cmp(Greater) }
///
/// let lhs = term::gt( term::int_var(0), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Greater) }
///
/// let lhs = term::le( term::int_var(0), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// let expected = term::eq( term::int_var(0), term::int(7) ) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Yields(expected) }
///
/// let lhs = term::le( term::int_var(1), term::int(7) ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), None }
///
/// let lhs = term::le( term::int_var(1), term::int(7) ) ;
/// let rhs = term::or(
///   vec![ lhs.clone(), term::eq( term::int_var(3), term::int(7) ) ]
/// ) ;
/// # println!("   {}", lhs) ;
/// # println!("=> {}\n\n", rhs) ;
/// debug_assert_eq! { conj_simpl(& lhs, & rhs), Cmp(Greater) }
///
/// # println!("   {}", rhs) ;
/// # println!("=> {}\n\n", lhs) ;
/// debug_assert_eq! { conj_simpl(& rhs, & lhs), Cmp(Less) }
/// ```
pub fn conj_simpl<T1, T2>(lhs: & T1, rhs: & T2) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {

  if conf.term_simpl == 0 {
    return SimplRes::None
  }

  let res = conj_simpl_2(lhs, rhs) ;

  if res.is_some() {
    if let Some(solver) = simpl_solver.write().unwrap().as_mut() {
      solver.push(1).unwrap() ;
      let mut vars = VarSet::new() ;
      lhs.iter(
        |term| if let Some(idx) = term.var_idx() {
          let is_new = vars.insert(idx) ;
          if is_new {
            solver.declare_const(& format!("{}", term), & term.typ()).unwrap()
          }
        }
      ) ;
      rhs.iter(
        |term| if let Some(idx) = term.var_idx() {
          let is_new = vars.insert(idx) ;
          if is_new {
            solver.declare_const(& format!("{}", term), & term.typ()).unwrap()
          }
        }
      ) ;
      use std::cmp::Ordering::* ;
      let check = match res {
        SimplRes::Cmp(Equal) => format!(
          "(= {} {})", lhs.deref(), rhs.deref()
        ),
        SimplRes::Cmp(Less) => format!(
          "(=> {} {})", rhs.deref(), lhs.deref()
        ),
        SimplRes::Cmp(Greater) => format!(
          "(=> {} {})", lhs.deref(), rhs.deref()
        ),
        SimplRes::Yields(ref term) => format!(
          "(= (and {} {}) {})", lhs.deref(), rhs.deref(), term
        ),
        SimplRes::None => unreachable!(),
      } ;

      solver.assert(& format!("(not {})", check)).unwrap() ;

      if solver.check_sat().unwrap() {
        log! { @0
          " " ;
          "{}", lhs.deref() ;
          "{}", rhs.deref() ;
          "{:?}", res ;
          " "
        }
        panic! { "simplification failure" }
      }

      solver.pop(1).unwrap()
    }
  }

  match res {
    SimplRes::Yields(_) if conf.term_simpl <= 1 => SimplRes::None,
    res => res,
  }
}
pub fn conj_simpl_2<T1, T2>(lhs: & T1, rhs: & T2) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {
  use std::cmp::Ordering::* ;

  if let Some(args) = lhs.disj_inspect() {
    for lhs in args {
      let mut greater_count = 0 ;
      let mut yields = vec![] ;
      match int_conj_simpl(lhs, rhs, true) {
        SimplRes::Cmp(Equal) |
        SimplRes::Cmp(Less) => return SimplRes::Cmp(Less),
        SimplRes::Cmp(Greater) => greater_count += 1,

        SimplRes::Yields(term) => yields.push(term),
        SimplRes::None => break,
      }
      if yields.len() == args.len() {
        return SimplRes::Yields( term::or(yields) )
      } else if greater_count == args.len() {
        return SimplRes::Cmp(Greater)
      }
    }
  } else if let Some(args) = rhs.disj_inspect() {
    for rhs in args {
      let mut less_count = 0 ;
      let mut yields = vec![] ;
      match int_conj_simpl(lhs, rhs, true) {
        SimplRes::Cmp(Equal) |
        SimplRes::Cmp(Greater) => return SimplRes::Cmp(Greater),
        SimplRes::Cmp(Less) => less_count += 1,

        SimplRes::Yields(term) => yields.push(term),
        SimplRes::None => break,
      }
      if yields.len() == args.len() {
        return SimplRes::Yields( term::or(yields) )
      } else if less_count == args.len() {
        return SimplRes::Cmp(Greater)
      }
    }
  }

  let res = int_conj_simpl(lhs, rhs, true) ;
  // if res.is_some() {
  //   log! { @0
  //     "\n\n{}", lhs.deref() ;
  //     "{}", rhs.deref() ;
  //     "{:?}", res
  //   }
  // }

  res
}

fn int_conj_simpl<T1, T2>(lhs: & T1, rhs: & T2, conj: bool) -> SimplRes
where T1: Deref<Target=RTerm>, T2: Deref<Target=RTerm> {
  use std::cmp::Ordering::* ;
  // Input boolean is true (false) for `lhs` => `rhs` (reversed).
  macro_rules! ord_of_bool {
    ($b:expr) => (
      if $b {
        SimplRes::Cmp(Greater)
      } else {
        SimplRes::Cmp(Less)
      }
    ) ;
  }

  let (lhs, rhs) = ( lhs.deref(), rhs.deref() ) ;

  // A term implies itself.
  if lhs == rhs { return SimplRes::Cmp(Equal) }

  match ( lhs.bool(), rhs.bool() ) {
    // True can only imply true.
    (Some(true), rhs) => return ord_of_bool!(
      rhs.unwrap_or(false)
    ),
    // False implies anything.
    (Some(false), _) => return ord_of_bool!(true),
    // False can only be implied by false.
    (lhs, Some(false)) => return ord_of_bool!(
      ! lhs.unwrap_or(true)
    ),
    // True is implied by anything.
    (_, Some(true)) => return ord_of_bool!(true),
    // Otherwise we don't know (yet).
    (None, None) => (),
  }

  // Only legal atoms are `vars >= cst` and `vars > cst`.
  let (
    lhs_op, lhs_vars, lhs_cst
  ) = if let Some((op, args)) = lhs.app_inspect() {
    if ! args[0].typ().is_arith() {
      return SimplRes::None
    }
    (op, & args[0], args[1].val().unwrap())
  } else {
    return SimplRes::None
  } ;
  let (
    rhs_op, rhs_vars, rhs_cst
  ) = if let Some((op, args)) = rhs.app_inspect() {
    if ! args[0].typ().is_arith() {
      return SimplRes::None
    }
    (op, & args[0], args[1].val().unwrap())
  } else {
    return SimplRes::None
  } ;

  if lhs_vars == rhs_vars {

    match (lhs_op, rhs_op) {
      (Op::Gt, Op::Gt) |
      (Op::Ge, Op::Ge) => return SimplRes::Cmp(
        lhs_cst.cmp(& rhs_cst).unwrap()
      ),

      (Op::Gt, Op::Ge) |
      (Op::Ge, Op::Gt) => if lhs_cst == rhs_cst {
        if lhs_op == Op::Gt {
          return SimplRes::Cmp(Less)
        } else {
          return SimplRes::Cmp(Greater)
        }
      } else {
        return SimplRes::Cmp(
          lhs_cst.cmp(& rhs_cst).unwrap().reverse()
        )
      },

      (Op::Eql, Op::Ge) |
      (Op::Eql, Op::Gt) => match lhs_cst.cmp(& rhs_cst) {
        Some(Less) => return SimplRes::Yields( term::fls() ),
        Some(Equal) if rhs_op == Op::Gt => if conj {
          return SimplRes::Yields( term::fls() )
        } else {
        },
        Some(Equal) |
        Some(Greater) => return SimplRes::Cmp(Greater),
        None => unreachable!(),
      },

      (Op::Ge, Op::Eql) |
      (Op::Gt, Op::Eql) => match rhs_cst.cmp(& lhs_cst) {
        Some(Less) => return SimplRes::Yields( term::fls() ),
        Some(Equal) if rhs_op == Op::Gt => return SimplRes::Yields(
          term::fls()
        ),
        Some(Equal) |
        Some(Greater) => return SimplRes::Cmp(Less),
        None => unreachable!(),
      },

      (Op::Eql, Op::Eql) => if rhs_cst == lhs_cst {
        return SimplRes::Cmp(Greater)
      } else {
        return SimplRes::Yields( term::fls() )
      },

      _ => unreachable!(),
    }

  } else

  if lhs_op == Op::Ge && rhs_op == Op::Ge
  && lhs_cst == rhs_cst.minus().unwrap()
  && rhs_vars == & term::u_minus( lhs_vars.clone() ) {

    if conj {
      return SimplRes::Yields(
        term::eq(lhs_vars.clone(), lhs_cst.to_term().unwrap())
      )
    } else {
      return SimplRes::Yields(
        term::tru()
      )
    }
  }

  SimplRes::None
}

<|endoftext|>


<commit_before>
extern crate chrono;
extern crate rustc_serialize;

pub mod daterecord;
pub mod direction;
pub mod error;
pub mod timeentry;
mod timeentrypair;

use chrono::*;
use csv;
use self::daterecord::DateRecord;
use self::direction::Direction;
use self::error::TimeClockError;
use self::timeentry::TimeEntry;
pub use self::timeentrypair::{TimeEntryPair, TimeEntryPairsIter,
                              timeentry_pairs};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::SeekFrom;
use std::io::prelude::*;

pub fn read_timesheet<R: Read>(file: R)
                               -> Result<Vec<TimeEntry>, TimeClockError> {
    let mut rdr = csv::Reader::from_reader(file).has_headers(false);
    let mut in_v = rdr.decode().collect::<csv::Result<Vec<TimeEntry>>>()?;
    in_v.sort_by_key(|k| k.time);

    return Ok(in_v);
}

/// Match TimeEntrys into pairs of In, Out records, inserting new ones where
/// complete pairs can't be made.
pub fn pair_time_entries(records: Vec<TimeEntry>) -> Vec<TimeEntry> {
    let mut v: Vec<TimeEntry> = Vec::new();

    let mut dir = Direction::In;
    let mut prev_time: Option<DateTime<FixedOffset>> = None;

    for rec in records {
        match (dir, rec.dir) {
            (Direction::In, Direction::Out) => {
                // expected an In, got an Out
                v.push(TimeEntry {
                    dir: Direction::In,
                    time: rec.time,
                    memo: String::from("Missing clock in."),
                });
                v.push(rec.clone());
                dir = Direction::Out;
            }
            (Direction::Out, Direction::In) => {
                // expected an Out, got an In
                v.push(TimeEntry {
                    dir: Direction::Out,
                    time: prev_time.or(Some(rec.time)).unwrap(),
                    memo: String::from("Missing clock out."),
                });
                v.push(rec.clone());
                dir = Direction::Out;
            }
            (_, _) => v.push(rec.clone()),
        }
        dir = dir.reverse();
        prev_time = Some(rec.time);
    }

    // if the for loop ends on a Direction::In record, then we need to add
    // an out record to balance it.
    if v.len() % 2 != 0 {
        v.push(TimeEntry {
            dir: Direction::Out,
            time: now(),
            memo: String::from("Still clocked in."),
        });
    }

    return v;
}

/// Reduce pairs of TimeEntrys into DateRecords
pub fn collect_date_records(records: Vec<TimeEntry>) -> Vec<DateRecord> {
    let mut date_duration_map = BTreeMap::new();

    for tep in timeentry_pairs(records.into_iter()) {
        let r = DateRecord::from(tep);

        if !date_duration_map.contains_key(&r.date()) {
            date_duration_map.insert(r.date(), r);
        } else {
            let mut rec = date_duration_map.remove(&r.date()).unwrap();
            rec.add_seconds(r.seconds());
            rec.append_memo(&r.memo());

            date_duration_map.insert(rec.date(), rec);
        }
    }

    let v: Vec<DateRecord> = date_duration_map.iter()
        .map(|(_, r)| r.clone())
        .collect();

    return v;
}

/// Marks the time.
pub fn mark_time(d: Direction,
                 time: DateTime<FixedOffset>,
                 memo: String,
                 mut file: File) {
    let record = TimeEntry {
        dir: d,
        time: time,
        memo: memo,
    };
    // seek in case we write without reading first
    let _ = file.seek(SeekFrom::End(0));
    let mut wtr = csv::Writer::from_writer(file);
    let _ = wtr.encode(record);
}

pub fn now() -> DateTime<FixedOffset> {
    let lt: DateTime<Local> = Local::now();
    DateTime::parse_from_rfc3339(lt.to_rfc3339().as_ref()).unwrap()
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use super::*;
    use timeclock::direction::Direction;

    #[test]
    fn read_timesheet_test() {
        // test for sorting and general function
        let s = "Out,2016-12-18T16:53:33-0600,\n\
                 In,2016-12-19T20:54:53-0600,\n\
                 In,2016-12-18T13:01:50-0600,\n\
                 Out,2016-12-19T20:54:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let records = read_timesheet(buff).unwrap();
        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        // test ordering
        assert!(records[0].time < records[1].time);
        assert!(records[1].time < records[2].time);
        assert!(records[2].time < records[3].time);
    }

    #[test]
    #[should_panic(expected = "called `Result::unwrap()` on an `Err` value")]
    fn read_timesheet_haggis_test() {
        // test for sorting and general function
        let s = "Fair fa' your honest, sonsie face,\n\
                 Great chieftain o' the pudding-race!\n\
                 Aboon them a' ye tak your place,\n\
                 Painch, tripe, or thairm:\n\
                 Weel are ye wordy o' a grace\n\
                 As lang's my arm.";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let records = read_timesheet(buff).unwrap();
        assert!(records[0].dir == Direction::In);
    }

    #[test]
    fn pair_time_entries_test() {
        let s = "Out,2016-12-18T13:01:50-0600,\n\
                 In,2016-12-19T20:54:57-0600,\n\
                 Out,2016-12-20T20:50:57-0600,\n\
                 Out,2016-12-20T20:50:59-0600,\n\
                 In,2016-12-20T21:04:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let raw_records = read_timesheet(buff).unwrap();
        let records = pair_time_entries(raw_records);

        println!("\nlen={}", records.len());
        for record in records.clone() {
            println!("{}", record);
        }

        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        assert!(records.len() == 8);
    }

    #[test]
    fn pair_time_entries_double_in_test() {
        let s = "In,2016-12-20T21:01:57-0600,\n\
                 In,2016-12-20T21:04:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let raw_records = read_timesheet(buff).unwrap();
        let records = pair_time_entries(raw_records);

        println!("\nlen={}", records.len());
        for record in records.clone() {
            println!("{}", record);
        }

        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        assert!(records.len() == 4);
    }

    #[test]
    fn collect_date_records_tz_test() {
        // test mismatched timezones
        let s = "In,2016-12-18T13:01:50-0500,\n\
                 Out,2016-12-18T13:01:50-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let entries = read_timesheet(buff).unwrap();
        let paired_records = pair_time_entries(entries);
        let records = collect_date_records(paired_records);

        assert!(records.len() == 1);
        println!("\n{}", records[0].seconds());
        assert!(records[0].seconds() == 3600.0);
    }
}
<commit_msg>Added test for collect_date_records()<commit_after>
extern crate chrono;
extern crate rustc_serialize;

pub mod daterecord;
pub mod direction;
pub mod error;
pub mod timeentry;
mod timeentrypair;

use chrono::*;
use csv;
use self::daterecord::DateRecord;
use self::direction::Direction;
use self::error::TimeClockError;
pub use self::timeentry::TimeEntry;
pub use self::timeentrypair::{TimeEntryPair, TimeEntryPairsIter,
                              timeentry_pairs};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::SeekFrom;
use std::io::prelude::*;

pub fn read_timesheet<R: Read>(file: R)
                               -> Result<Vec<TimeEntry>, TimeClockError> {
    let mut rdr = csv::Reader::from_reader(file).has_headers(false);
    let mut in_v = rdr.decode().collect::<csv::Result<Vec<TimeEntry>>>()?;
    in_v.sort_by_key(|k| k.time);

    return Ok(in_v);
}

/// Match TimeEntrys into pairs of In, Out records, inserting new ones where
/// complete pairs can't be made.
pub fn pair_time_entries(records: Vec<TimeEntry>) -> Vec<TimeEntry> {
    let mut v: Vec<TimeEntry> = Vec::new();

    let mut dir = Direction::In;
    let mut prev_time: Option<DateTime<FixedOffset>> = None;

    for rec in records {
        match (dir, rec.dir) {
            (Direction::In, Direction::Out) => {
                // expected an In, got an Out
                v.push(TimeEntry {
                    dir: Direction::In,
                    time: rec.time,
                    memo: String::from("Missing clock in."),
                });
                v.push(rec.clone());
                dir = Direction::Out;
            }
            (Direction::Out, Direction::In) => {
                // expected an Out, got an In
                v.push(TimeEntry {
                    dir: Direction::Out,
                    time: prev_time.or(Some(rec.time)).unwrap(),
                    memo: String::from("Missing clock out."),
                });
                v.push(rec.clone());
                dir = Direction::Out;
            }
            (_, _) => v.push(rec.clone()),
        }
        dir = dir.reverse();
        prev_time = Some(rec.time);
    }

    // if the for loop ends on a Direction::In record, then we need to add
    // an out record to balance it.
    if v.len() % 2 != 0 {
        v.push(TimeEntry {
            dir: Direction::Out,
            time: now(),
            memo: String::from("Still clocked in."),
        });
    }

    return v;
}

/// Reduce pairs of TimeEntrys into DateRecords
pub fn collect_date_records(records: Vec<TimeEntry>) -> Vec<DateRecord> {
    let mut date_duration_map = BTreeMap::new();

    for tep in timeentry_pairs(records.into_iter()) {
        let r = DateRecord::from(tep);

        if !date_duration_map.contains_key(&r.date()) {
            date_duration_map.insert(r.date(), r);
        } else {
            let mut rec = date_duration_map.remove(&r.date()).unwrap();
            rec.add_seconds(r.seconds());
            rec.append_memo(&r.memo());

            date_duration_map.insert(rec.date(), rec);
        }
    }

    let v: Vec<DateRecord> = date_duration_map.iter()
        .map(|(_, r)| r.clone())
        .collect();

    return v;
}

/// Marks the time.
pub fn mark_time(d: Direction,
                 time: DateTime<FixedOffset>,
                 memo: String,
                 mut file: File) {
    let record = TimeEntry {
        dir: d,
        time: time,
        memo: memo,
    };
    // seek in case we write without reading first
    let _ = file.seek(SeekFrom::End(0));
    let mut wtr = csv::Writer::from_writer(file);
    let _ = wtr.encode(record);
}

pub fn now() -> DateTime<FixedOffset> {
    let lt: DateTime<Local> = Local::now();
    DateTime::parse_from_rfc3339(lt.to_rfc3339().as_ref()).unwrap()
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use super::*;
    use timeclock::direction::Direction;

    #[test]
    fn read_timesheet_test() {
        // test for sorting and general function
        let s = "Out,2016-12-18T16:53:33-0600,\n\
                 In,2016-12-19T20:54:53-0600,\n\
                 In,2016-12-18T13:01:50-0600,\n\
                 Out,2016-12-19T20:54:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let records = read_timesheet(buff).unwrap();
        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        // test ordering
        assert!(records[0].time < records[1].time);
        assert!(records[1].time < records[2].time);
        assert!(records[2].time < records[3].time);
    }

    #[test]
    #[should_panic(expected = "called `Result::unwrap()` on an `Err` value")]
    fn read_timesheet_haggis_test() {
        // test for sorting and general function
        let s = "Fair fa' your honest, sonsie face,\n\
                 Great chieftain o' the pudding-race!\n\
                 Aboon them a' ye tak your place,\n\
                 Painch, tripe, or thairm:\n\
                 Weel are ye wordy o' a grace\n\
                 As lang's my arm.";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let records = read_timesheet(buff).unwrap();
        assert!(records[0].dir == Direction::In);
    }

    #[test]
    fn pair_time_entries_test() {
        let s = "Out,2016-12-18T13:01:50-0600,\n\
                 In,2016-12-19T20:54:57-0600,\n\
                 Out,2016-12-20T20:50:57-0600,\n\
                 Out,2016-12-20T20:50:59-0600,\n\
                 In,2016-12-20T21:04:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let raw_records = read_timesheet(buff).unwrap();
        let records = pair_time_entries(raw_records);

        println!("\nlen={}", records.len());
        for record in records.clone() {
            println!("{}", record);
        }

        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        assert!(records.len() == 8);
    }

    #[test]
    fn pair_time_entries_double_in_test() {
        let s = "In,2016-12-20T21:01:57-0600,\n\
                 In,2016-12-20T21:04:57-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let raw_records = read_timesheet(buff).unwrap();
        let records = pair_time_entries(raw_records);

        println!("\nlen={}", records.len());
        for record in records.clone() {
            println!("{}", record);
        }

        for slc in records.chunks(2) {
            assert!(slc[0].dir == Direction::In);
            assert!(slc[1].dir == Direction::Out);
        }
        assert!(records.len() == 4);
    }

    #[test]
    fn collect_date_records_test() {
        let time = now();
        let records = vec![TimeEntry::new(Direction::In, time, "In"),
                           TimeEntry::new(Direction::Out, time, "Out")];

        let res = collect_date_records(records);
        assert_eq!(res.len(), 1);
        let dr = res.first().unwrap();
        assert_eq!(dr.date(), time.date());
        assert_eq!(dr.hours(), 0.0);
    }

    #[test]
    fn collect_date_records_tz_test() {
        // test mismatched timezones
        let s = "In,2016-12-18T13:01:50-0500,\n\
                 Out,2016-12-18T13:01:50-0600,";
        let vs = s.as_bytes();
        let buff = Cursor::new(vs);
        let entries = read_timesheet(buff).unwrap();
        let paired_records = pair_time_entries(entries);
        let records = collect_date_records(paired_records);

        assert!(records.len() == 1);
        println!("\n{}", records[0].seconds());
        assert!(records[0].seconds() == 3600.0);
    }
}
<|endoftext|>


<commit_before>use crate::sys;

use tokio_io::{AsyncRead, AsyncWrite};

use futures_core::ready;
use std::cmp;
use std::future::Future;
use std::io;
use std::io::prelude::*;
use std::pin::Pin;
use std::task::Poll::*;
use std::task::{Context, Poll};

use self::State::*;

/// `T` should not implement _both_ Read and Write.
#[derive(Debug)]
pub(crate) struct Blocking<T> {
    inner: Option<T>,
    state: State<T>,
}

#[derive(Debug)]
pub(crate) struct Buf {
    buf: Vec<u8>,
    pos: usize,
}

pub(crate) const MAX_BUF: usize = 16 * 1024;

#[derive(Debug)]
enum State<T> {
    Idle(Option<Buf>),
    Busy(sys::Blocking<(io::Result<usize>, Buf, T)>),
}

impl<T> Blocking<T> {
    pub(crate) fn new(inner: T) -> Blocking<T> {
        Blocking {
            inner: Some(inner),
            state: State::Idle(Some(Buf::with_capacity(0))),
        }
    }
}

impl<T> AsyncRead for Blocking<T>
where
    T: Read + Unpin + Send + 'static,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        dst: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        loop {
            match self.state {
                Idle(ref mut buf_cell) => {
                    let mut buf = buf_cell.take().unwrap();

                    if !buf.is_empty() {
                        let n = buf.copy_to(dst);
                        *buf_cell = Some(buf);
                        return Ready(Ok(n));
                    }

                    buf.ensure_capacity_for(dst);
                    let mut inner = self.inner.take().unwrap();

                    self.state = Busy(sys::run(move || {
                        let res = buf.read_from(&mut inner);
                        (res, buf, inner)
                    }));
                }
                Busy(ref mut rx) => {
                    let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx));
                    self.inner = Some(inner);

                    match res {
                        Ok(_) => {
                            let n = buf.copy_to(dst);
                            self.state = Idle(Some(buf));
                            return Ready(Ok(n));
                        }
                        Err(e) => {
                            assert!(buf.is_empty());

                            self.state = Idle(Some(buf));
                            return Ready(Err(e));
                        }
                    }
                }
            }
        }
    }
}

impl<T> AsyncWrite for Blocking<T>
where
    T: Write + Unpin + Send + 'static,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        src: &[u8],
    ) -> Poll<io::Result<usize>> {
        loop {
            match self.state {
                Idle(ref mut buf_cell) => {
                    let mut buf = buf_cell.take().unwrap();

                    assert!(buf.is_empty());

                    let n = buf.copy_from(src);
                    let mut inner = self.inner.take().unwrap();

                    self.state = Busy(sys::run(move || {
                        let n = buf.len();
                        let res = buf.write_to(&mut inner).map(|_| n);

                        (res, buf, inner)
                    }));

                    return Ready(Ok(n));
                }
                Busy(ref mut rx) => {
                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx));
                    self.state = Idle(Some(buf));
                    self.inner = Some(inner);

                    // If error, return
                    res?;
                }
            }
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        let (res, buf, inner) = match self.state {
            Idle(_) => return Ready(Ok(())),
            Busy(ref mut rx) => ready!(Pin::new(rx).poll(cx)),
        };

        // The buffer is not used here
        self.state = Idle(Some(buf));
        self.inner = Some(inner);

        Ready(res.map(|_| ()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Poll::Ready(Ok(()))
    }
}

/// Repeates operations that are interrupted
macro_rules! uninterruptibly {
    ($e:expr) => {{
        loop {
            match $e {
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                res => break res,
            }
        }
    }};
}

impl Buf {
    pub(crate) fn with_capacity(n: usize) -> Buf {
        Buf {
            buf: Vec::with_capacity(n),
            pos: 0,
        }
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub(crate) fn len(&self) -> usize {
        self.buf.len() - self.pos
    }

    pub(crate) fn copy_to(&mut self, dst: &mut [u8]) -> usize {
        let n = cmp::min(self.len(), dst.len());
        dst[..n].copy_from_slice(&self.bytes()[..n]);
        self.pos += n;

        if self.pos == self.buf.len() {
            self.buf.truncate(0);
            self.pos = 0;
        }

        n
    }

    pub(crate) fn copy_from(&mut self, src: &[u8]) -> usize {
        assert!(self.is_empty());

        let n = cmp::min(src.len(), MAX_BUF);

        self.buf.extend_from_slice(&src[..n]);
        n
    }

    pub(crate) fn bytes(&self) -> &[u8] {
        &self.buf[self.pos..]
    }

    pub(crate) fn ensure_capacity_for(&mut self, bytes: &[u8]) {
        assert!(self.is_empty());

        let len = cmp::min(bytes.len(), MAX_BUF);

        if self.buf.len() < len {
            self.buf.reserve(len - self.buf.len());
        }

        unsafe {
            self.buf.set_len(len);
        }
    }

    pub(crate) fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> {
        let res = uninterruptibly!(rd.read(&mut self.buf));

        if let Ok(n) = res {
            self.buf.truncate(n);
        } else {
            self.buf.clear();
        }

        assert_eq!(self.pos, 0);

        res
    }

    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
        assert_eq!(self.pos, 0);

        // `write_all` already ignores interrupts
        let res = wr.write_all(&self.buf);
        self.buf.clear();
        res
    }

    pub(crate) fn discard_read(&mut self) -> i64 {
        let ret = -(self.bytes().len() as i64);
        self.pos = 0;
        self.buf.truncate(0);
        ret
    }
}
<commit_msg>fs: propagate flush for stdout / stderr. (#1528)<commit_after>use crate::sys;

use tokio_io::{AsyncRead, AsyncWrite};

use futures_core::ready;
use std::cmp;
use std::future::Future;
use std::io;
use std::io::prelude::*;
use std::pin::Pin;
use std::task::Poll::*;
use std::task::{Context, Poll};

use self::State::*;

/// `T` should not implement _both_ Read and Write.
#[derive(Debug)]
pub(crate) struct Blocking<T> {
    inner: Option<T>,
    state: State<T>,
    /// true if the lower IO layer needs flushing
    need_flush: bool,
}

#[derive(Debug)]
pub(crate) struct Buf {
    buf: Vec<u8>,
    pos: usize,
}

pub(crate) const MAX_BUF: usize = 16 * 1024;

#[derive(Debug)]
enum State<T> {
    Idle(Option<Buf>),
    Busy(sys::Blocking<(io::Result<usize>, Buf, T)>),
}

impl<T> Blocking<T> {
    pub(crate) fn new(inner: T) -> Blocking<T> {
        Blocking {
            inner: Some(inner),
            state: State::Idle(Some(Buf::with_capacity(0))),
            need_flush: false,
        }
    }
}

impl<T> AsyncRead for Blocking<T>
where
    T: Read + Unpin + Send + 'static,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        dst: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        loop {
            match self.state {
                Idle(ref mut buf_cell) => {
                    let mut buf = buf_cell.take().unwrap();

                    if !buf.is_empty() {
                        let n = buf.copy_to(dst);
                        *buf_cell = Some(buf);
                        return Ready(Ok(n));
                    }

                    buf.ensure_capacity_for(dst);
                    let mut inner = self.inner.take().unwrap();

                    self.state = Busy(sys::run(move || {
                        let res = buf.read_from(&mut inner);
                        (res, buf, inner)
                    }));
                }
                Busy(ref mut rx) => {
                    let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx));
                    self.inner = Some(inner);

                    match res {
                        Ok(_) => {
                            let n = buf.copy_to(dst);
                            self.state = Idle(Some(buf));
                            return Ready(Ok(n));
                        }
                        Err(e) => {
                            assert!(buf.is_empty());

                            self.state = Idle(Some(buf));
                            return Ready(Err(e));
                        }
                    }
                }
            }
        }
    }
}

impl<T> AsyncWrite for Blocking<T>
where
    T: Write + Unpin + Send + 'static,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        src: &[u8],
    ) -> Poll<io::Result<usize>> {
        loop {
            match self.state {
                Idle(ref mut buf_cell) => {
                    let mut buf = buf_cell.take().unwrap();

                    assert!(buf.is_empty());

                    let n = buf.copy_from(src);
                    let mut inner = self.inner.take().unwrap();

                    self.state = Busy(sys::run(move || {
                        let n = buf.len();
                        let res = buf.write_to(&mut inner).map(|_| n);

                        (res, buf, inner)
                    }));
                    self.need_flush = true;

                    return Ready(Ok(n));
                }
                Busy(ref mut rx) => {
                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx));
                    self.state = Idle(Some(buf));
                    self.inner = Some(inner);

                    // If error, return
                    res?;
                }
            }
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        loop {
            let need_flush = self.need_flush;
            match self.state {
                // The buffer is not used here
                Idle(ref mut buf_cell) => {
                    if need_flush {
                        let buf = buf_cell.take().unwrap();
                        let mut inner = self.inner.take().unwrap();

                        self.state = Busy(sys::run(move || {
                            let res = inner.flush().map(|_| 0);
                            (res, buf, inner)
                        }));

                        self.need_flush = false;
                    } else {
                        return Ready(Ok(()));
                    }
                }
                Busy(ref mut rx) => {
                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx));
                    self.state = Idle(Some(buf));
                    self.inner = Some(inner);

                    // If error, return
                    res?;
                }
            }
        }
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        Poll::Ready(Ok(()))
    }
}

/// Repeates operations that are interrupted
macro_rules! uninterruptibly {
    ($e:expr) => {{
        loop {
            match $e {
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                res => break res,
            }
        }
    }};
}

impl Buf {
    pub(crate) fn with_capacity(n: usize) -> Buf {
        Buf {
            buf: Vec::with_capacity(n),
            pos: 0,
        }
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub(crate) fn len(&self) -> usize {
        self.buf.len() - self.pos
    }

    pub(crate) fn copy_to(&mut self, dst: &mut [u8]) -> usize {
        let n = cmp::min(self.len(), dst.len());
        dst[..n].copy_from_slice(&self.bytes()[..n]);
        self.pos += n;

        if self.pos == self.buf.len() {
            self.buf.truncate(0);
            self.pos = 0;
        }

        n
    }

    pub(crate) fn copy_from(&mut self, src: &[u8]) -> usize {
        assert!(self.is_empty());

        let n = cmp::min(src.len(), MAX_BUF);

        self.buf.extend_from_slice(&src[..n]);
        n
    }

    pub(crate) fn bytes(&self) -> &[u8] {
        &self.buf[self.pos..]
    }

    pub(crate) fn ensure_capacity_for(&mut self, bytes: &[u8]) {
        assert!(self.is_empty());

        let len = cmp::min(bytes.len(), MAX_BUF);

        if self.buf.len() < len {
            self.buf.reserve(len - self.buf.len());
        }

        unsafe {
            self.buf.set_len(len);
        }
    }

    pub(crate) fn read_from<T: Read>(&mut self, rd: &mut T) -> io::Result<usize> {
        let res = uninterruptibly!(rd.read(&mut self.buf));

        if let Ok(n) = res {
            self.buf.truncate(n);
        } else {
            self.buf.clear();
        }

        assert_eq!(self.pos, 0);

        res
    }

    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
        assert_eq!(self.pos, 0);

        // `write_all` already ignores interrupts
        let res = wr.write_all(&self.buf);
        self.buf.clear();
        res
    }

    pub(crate) fn discard_read(&mut self) -> i64 {
        let ret = -(self.bytes().len() as i64);
        self.pos = 0;
        self.buf.truncate(0);
        ret
    }
}
<|endoftext|>