<commit_before>//! Provides two lazy sort variants that can be faster than a full sort in certain situations.
//! Complexity of taking the first k elements:
//!
//! | HeapSort      | QuickSort     |
//! | ------------- | ------------- |
//! | O(n + klog(n) | O(n + klog(k) |
//!
//! However, quicksort allocates, whereas heapsort does not,
//! so for values of k that are a significant fraction of n,
//! heapsort may perform better than both quicksort and
//! regular sorting.

#![deny(missing_docs)]
#![feature(slice_splits, core)]
#![cfg_attr(test, feature(test))]
extern crate core;
extern crate itertools;
extern crate rand;

use core::ptr;
use itertools::partition;
use std::cmp::Ordering::{self, Less};
use std::mem;

/// An iterator extension trait that provides two methods for lazily sorting.
pub trait LazySortIterator: Iterator
    where Self: Sized,
          Self::Item: Ord
{
    /// Lazily sort using quicksort.
    fn quick_sort(self) -> QuickSort<Self::Item> {
        QuickSort { inner: QuickSortInternal::new(self.collect()) }
    }

    /// Lazily sort using heapsort.
    fn heap_sort(self) -> HeapSort<Self::Item> {
        HeapSort(self.map(|el| ReverseOrder(el)).collect())
    }
}

impl<T> LazySortIterator for T
    where T: Iterator,
          T::Item: Ord { }

/// An iterator that lazily sorts its input using quicksort.
#[derive(Debug, Clone)]
pub struct QuickSort<T> {
    inner: QuickSortInternal<T>,
}

impl<T: Ord> Iterator for QuickSort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

#[derive(Debug, Clone)]
enum QuickSortInternal<T> {
    Base(Vec<T>),
    Recursive(Recursive<T>),
}

impl<T: Ord> QuickSortInternal<T> {
    fn new(mut v: Vec<T>) -> QuickSortInternal<T> {
        if v.len() <= 32 {
            insertion_sort(&mut v, |a, b| b.cmp(a));
            QuickSortInternal::Base(v)
        } else {
            QuickSortInternal::Recursive(Recursive::new(v))
        }
    }
}

impl<T: Ord> Iterator for QuickSortInternal<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match *self {
            QuickSortInternal::Base(ref mut v) => v.pop(),
            QuickSortInternal::Recursive(ref mut r) => r.next(),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match *self {
            QuickSortInternal::Base(ref v) => (v.len(), Some(v.len())),
            QuickSortInternal::Recursive(ref r) => r.size_hint(),
        }
    }
}

#[derive(Clone, Debug)]
struct Recursive<T> {
    greater: Vec<T>,
    less: Option<Box<QuickSortInternal<T>>>,
}

impl<T: Ord> Recursive<T> {
    fn new(v: Vec<T>) -> Recursive<T> {
        Recursive {
            greater: v,
            less: None,
        }
    }

    fn split_greater(&mut self) -> Option<T> {
        match self.greater.len() {
            0 => None,
            1 => self.greater.pop(),
            _ => {
                let pivot_idx = self.greater.len() - 1;
                let split_idx = {
                    let mid_idx = self.greater.len() / 2;
                    // I've chosen the element in the middle of the vec as the pivot.
                    // However, we first swap the pivot with the last element so that there is
                    // a contiguous space in memory to be partitioned.
                    self.greater.swap(pivot_idx, mid_idx);
                    let (pivot, rest) = self.greater.split_last_mut().unwrap();
                    // partition all but the last element, which is the pivot. This makes the vec
                    // look like [greater, greater, ..., greater, less, less, ..., less, pivot]
                    partition(rest, |el| el > pivot)
                };
                // Swapping the pivot with the first less element allows us to split off
                // vec[split_idx + 1..] to create a new vec with all the elements less than pivot.
                self.greater.swap(pivot_idx, split_idx);
                let split_off_idx = split_idx + 1;
                if split_off_idx < self.greater.len() {
                    let mut less = Box::new(QuickSortInternal::new(self.greater
                                                                      .split_off(split_off_idx)));
                    // Recursively compute the next element from the QuickSortInternal struct
                    // containing the elements less than the pivot.
                    let next = less.next();
                    self.less = Some(less);
                    next
                } else {
                    // If there were no elements less than the pivot, then return the pivot.
                    self.greater.pop()
                }
            }
        }
    }
}

impl<T: Ord> Iterator for Recursive<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let next = if let Some(ref mut less) = self.less {
            less.next()
        } else {
            return self.split_greater();
        };
        if next.is_some() {
            next
        } else {
            self.less = None;
            // The pivot is always the last element in the vec, and it's the first element
            // to be returned once all of the elements less than it have been returned.
            self.greater.pop()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.less {
            None => (self.greater.len(), Some(self.greater.len())),
            Some(ref less) => {
                let (lower, upper) = less.size_hint();
                (lower + self.greater.len(),
                 upper.map(|upper| upper + self.greater.len()))
            }
        }
    }
}

#[test]
fn test_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn test_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn test_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().quick_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

// This is copied from libcollections/slice.rs
fn insertion_sort<T, F>(v: &mut [T], mut compare: F)
    where F: FnMut(&T, &T) -> Ordering
{
    let len = v.len() as isize;
    let buf_v = v.as_mut_ptr();

    // 1 <= i < len;
    for i in 1..len {
        // j satisfies: 0 <= j <= i;
        let mut j = i;
        unsafe {
            // `i` is in bounds.
            let read_ptr = buf_v.offset(i) as *const T;

            // find where to insert, we need to do strict <,
            // rather than <=, to maintain stability.

            // 0 <= j - 1 < len, so .offset(j - 1) is in bounds.
            while j > 0 && compare(&*read_ptr, &*buf_v.offset(j - 1)) == Less {
                j -= 1;
            }

            // shift everything to the right, to make space to
            // insert this value.

            // j + 1 could be `len` (for the last `i`), but in
            // that case, `i == j` so we don't copy. The
            // `.offset(j)` is always in bounds.

            if i != j {
                let tmp = ptr::read(read_ptr);
                ptr::copy(&*buf_v.offset(j), buf_v.offset(j + 1), (i - j) as usize);
                ptr::copy_nonoverlapping(&tmp, buf_v.offset(j), 1);
                mem::forget(tmp);
            }
        }
    }
}

use std::collections::BinaryHeap;
/// An iterator that lazily sorts its input using quicksort.
pub struct HeapSort<T>(BinaryHeap<ReverseOrder<T>>);

#[derive(Eq, PartialEq)]
struct ReverseOrder<T>(T);

impl <T: PartialOrd> PartialOrd for ReverseOrder<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.0.partial_cmp(&self.0)
    }
}

impl <T: Ord> Ord for ReverseOrder<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.0.cmp(&self.0)
    }
}

impl<T: Ord> Iterator for HeapSort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.0.pop().map(|ReverseOrder(el)| el)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.0.len();
        (len, Some(len))
    }
}

#[test]
fn quick_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn heap_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().heap_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn quick_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn heap_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().heap_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn quick_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().quick_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

#[test]
fn heap_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().heap_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

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

    use self::test::Bencher;
    use rand::{thread_rng, Rng};
    use super::LazySortIterator;

    fn take_quick(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().quick_sort().take(k).collect::<Vec<_>>());
    }

    fn take_heap(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().heap_sort().take(k).collect::<Vec<_>>());
    }

    fn take_eager(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(k).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_10_quick(b: &mut Bencher) {
        take_quick(b, 10);
    }

    #[bench]
    fn take_10_heap(b: &mut Bencher) {
        take_heap(b, 10);
    }

    #[bench]
    fn take_10_eager(b: &mut Bencher) {
        take_eager(b, 10);
    }

    #[bench]
    fn take_1000_quick(b: &mut Bencher) {
        take_quick(b, 1000)
    }

    #[bench]
    fn take_1000_heap(b: &mut Bencher) {
        take_heap(b, 1000);
    }

    #[bench]
    fn take_1000_eager(b: &mut Bencher) {
        take_eager(b, 1000);
    }

    #[bench]
    fn take_10_000_quick(b: &mut Bencher) {
        take_quick(b, 10_000);
    }

    #[bench]
    fn take_10_000_heap(b: &mut Bencher) {
        take_heap(b, 10_000);
    }

    #[bench]
    fn take_10_000_eager(b: &mut Bencher) {
        take_eager(b, 10_000);
    }

    #[bench]
    fn take_50_000_quick(b: &mut Bencher) {
        take_quick(b, 50_000);
    }

    #[bench]
    fn take_50_000_heap(b: &mut Bencher) {
        take_heap(b, 50_000);
    }

    #[bench]
    fn take_50_000_eager(b: &mut Bencher) {
        take_eager(b, 50_000);
    }
}
<commit_msg>Change wording a bit<commit_after>//! Provides two lazy sort variants that can be faster than a full sort in certain situations.
//! Complexity of taking the first k elements:
//!
//! | HeapSort      | QuickSort     |
//! | ------------- | ------------- |
//! | O(n + klog(n) | O(n + klog(k) |
//!
//! However, quicksort allocates to track its recursive state, whereas heapsort allocates entirely
//! up front, so for values of k that are a significant fraction of n, heapsort will perform
//! better.

#![deny(missing_docs)]
#![feature(slice_splits, core)]
#![cfg_attr(test, feature(test))]
extern crate core;
extern crate itertools;
extern crate rand;

use core::ptr;
use itertools::partition;
use std::cmp::Ordering::{self, Less};
use std::mem;

/// An iterator extension trait that provides two methods for lazily sorting.
pub trait LazySortIterator: Iterator
    where Self: Sized,
          Self::Item: Ord
{
    /// Lazily sort using quicksort.
    fn quick_sort(self) -> QuickSort<Self::Item> {
        QuickSort { inner: QuickSortInternal::new(self.collect()) }
    }

    /// Lazily sort using heapsort.
    fn heap_sort(self) -> HeapSort<Self::Item> {
        HeapSort(self.map(|el| ReverseOrder(el)).collect())
    }
}

impl<T> LazySortIterator for T
    where T: Iterator,
          T::Item: Ord { }

/// An iterator that lazily sorts its input using quicksort.
#[derive(Debug, Clone)]
pub struct QuickSort<T> {
    inner: QuickSortInternal<T>,
}

impl<T: Ord> Iterator for QuickSort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.inner.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint()
    }
}

#[derive(Debug, Clone)]
enum QuickSortInternal<T> {
    Base(Vec<T>),
    Recursive(Recursive<T>),
}

impl<T: Ord> QuickSortInternal<T> {
    fn new(mut v: Vec<T>) -> QuickSortInternal<T> {
        if v.len() <= 32 {
            insertion_sort(&mut v, |a, b| b.cmp(a));
            QuickSortInternal::Base(v)
        } else {
            QuickSortInternal::Recursive(Recursive::new(v))
        }
    }
}

impl<T: Ord> Iterator for QuickSortInternal<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match *self {
            QuickSortInternal::Base(ref mut v) => v.pop(),
            QuickSortInternal::Recursive(ref mut r) => r.next(),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match *self {
            QuickSortInternal::Base(ref v) => (v.len(), Some(v.len())),
            QuickSortInternal::Recursive(ref r) => r.size_hint(),
        }
    }
}

#[derive(Clone, Debug)]
struct Recursive<T> {
    greater: Vec<T>,
    less: Option<Box<QuickSortInternal<T>>>,
}

impl<T: Ord> Recursive<T> {
    fn new(v: Vec<T>) -> Recursive<T> {
        Recursive {
            greater: v,
            less: None,
        }
    }

    fn split_greater(&mut self) -> Option<T> {
        match self.greater.len() {
            0 => None,
            1 => self.greater.pop(),
            _ => {
                let pivot_idx = self.greater.len() - 1;
                let split_idx = {
                    let mid_idx = self.greater.len() / 2;
                    // I've chosen the element in the middle of the vec as the pivot.
                    // However, we first swap the pivot with the last element so that there is
                    // a contiguous space in memory to be partitioned.
                    self.greater.swap(pivot_idx, mid_idx);
                    let (pivot, rest) = self.greater.split_last_mut().unwrap();
                    // partition all but the last element, which is the pivot. This makes the vec
                    // look like [greater, greater, ..., greater, less, less, ..., less, pivot]
                    partition(rest, |el| el > pivot)
                };
                // Swapping the pivot with the first less element allows us to split off
                // vec[split_idx + 1..] to create a new vec with all the elements less than pivot.
                self.greater.swap(pivot_idx, split_idx);
                let split_off_idx = split_idx + 1;
                if split_off_idx < self.greater.len() {
                    let mut less = Box::new(QuickSortInternal::new(self.greater
                                                                      .split_off(split_off_idx)));
                    // Recursively compute the next element from the QuickSortInternal struct
                    // containing the elements less than the pivot.
                    let next = less.next();
                    self.less = Some(less);
                    next
                } else {
                    // If there were no elements less than the pivot, then return the pivot.
                    self.greater.pop()
                }
            }
        }
    }
}

impl<T: Ord> Iterator for Recursive<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        let next = if let Some(ref mut less) = self.less {
            less.next()
        } else {
            return self.split_greater();
        };
        if next.is_some() {
            next
        } else {
            self.less = None;
            // The pivot is always the last element in the vec, and it's the first element
            // to be returned once all of the elements less than it have been returned.
            self.greater.pop()
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self.less {
            None => (self.greater.len(), Some(self.greater.len())),
            Some(ref less) => {
                let (lower, upper) = less.size_hint();
                (lower + self.greater.len(),
                 upper.map(|upper| upper + self.greater.len()))
            }
        }
    }
}

#[test]
fn test_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn test_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn test_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().quick_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

// This is copied from libcollections/slice.rs
fn insertion_sort<T, F>(v: &mut [T], mut compare: F)
    where F: FnMut(&T, &T) -> Ordering
{
    let len = v.len() as isize;
    let buf_v = v.as_mut_ptr();

    // 1 <= i < len;
    for i in 1..len {
        // j satisfies: 0 <= j <= i;
        let mut j = i;
        unsafe {
            // `i` is in bounds.
            let read_ptr = buf_v.offset(i) as *const T;

            // find where to insert, we need to do strict <,
            // rather than <=, to maintain stability.

            // 0 <= j - 1 < len, so .offset(j - 1) is in bounds.
            while j > 0 && compare(&*read_ptr, &*buf_v.offset(j - 1)) == Less {
                j -= 1;
            }

            // shift everything to the right, to make space to
            // insert this value.

            // j + 1 could be `len` (for the last `i`), but in
            // that case, `i == j` so we don't copy. The
            // `.offset(j)` is always in bounds.

            if i != j {
                let tmp = ptr::read(read_ptr);
                ptr::copy(&*buf_v.offset(j), buf_v.offset(j + 1), (i - j) as usize);
                ptr::copy_nonoverlapping(&tmp, buf_v.offset(j), 1);
                mem::forget(tmp);
            }
        }
    }
}

use std::collections::BinaryHeap;
/// An iterator that lazily sorts its input using quicksort.
pub struct HeapSort<T>(BinaryHeap<ReverseOrder<T>>);

#[derive(Eq, PartialEq)]
struct ReverseOrder<T>(T);

impl <T: PartialOrd> PartialOrd for ReverseOrder<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.0.partial_cmp(&self.0)
    }
}

impl <T: Ord> Ord for ReverseOrder<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.0.cmp(&self.0)
    }
}

impl<T: Ord> Iterator for HeapSort<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.0.pop().map(|ReverseOrder(el)| el)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.0.len();
        (len, Some(len))
    }
}

#[test]
fn quick_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn heap_sort() {
    let mut v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let v2: Vec<_> = v.iter().cloned().heap_sort().collect();
    v.sort();
    assert_eq!(v, v2);
}

#[test]
fn quick_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().quick_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn heap_empty() {
    let v: Vec<u64> = vec![];
    let v2: Vec<_> = v.iter().cloned().heap_sort().collect();
    assert_eq!(v, v2);
}

#[test]
fn quick_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().quick_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

#[test]
fn heap_size_hint() {
    let v = vec![2, 4, 2, 5, 8, 4, 3, 4, 6];
    let mut sort_iter = v.iter().cloned().heap_sort();
    for i in 0..v.len() {
        let (lower, upper) = sort_iter.size_hint();
        assert_eq!(v.len() - i, lower);
        assert_eq!(Some(v.len() - i), upper);
        sort_iter.next();
    }
}

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

    use self::test::Bencher;
    use rand::{thread_rng, Rng};
    use super::LazySortIterator;

    fn take_quick(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().quick_sort().take(k).collect::<Vec<_>>());
    }

    fn take_heap(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| v.iter().cloned().heap_sort().take(k).collect::<Vec<_>>());
    }

    fn take_eager(b: &mut Bencher, k: usize) {
        let mut rng = thread_rng();
        let v: Vec<u32> = rng.gen_iter().take(50_000).collect();
        b.iter(|| {
            let mut v = v.clone();
            v.sort();
            v.iter().cloned().take(k).collect::<Vec<_>>();
        });
    }

    #[bench]
    fn take_10_quick(b: &mut Bencher) {
        take_quick(b, 10);
    }

    #[bench]
    fn take_10_heap(b: &mut Bencher) {
        take_heap(b, 10);
    }

    #[bench]
    fn take_10_eager(b: &mut Bencher) {
        take_eager(b, 10);
    }

    #[bench]
    fn take_1000_quick(b: &mut Bencher) {
        take_quick(b, 1000)
    }

    #[bench]
    fn take_1000_heap(b: &mut Bencher) {
        take_heap(b, 1000);
    }

    #[bench]
    fn take_1000_eager(b: &mut Bencher) {
        take_eager(b, 1000);
    }

    #[bench]
    fn take_10_000_quick(b: &mut Bencher) {
        take_quick(b, 10_000);
    }

    #[bench]
    fn take_10_000_heap(b: &mut Bencher) {
        take_heap(b, 10_000);
    }

    #[bench]
    fn take_10_000_eager(b: &mut Bencher) {
        take_eager(b, 10_000);
    }

    #[bench]
    fn take_50_000_quick(b: &mut Bencher) {
        take_quick(b, 50_000);
    }

    #[bench]
    fn take_50_000_heap(b: &mut Bencher) {
        take_heap(b, 50_000);
    }

    #[bench]
    fn take_50_000_eager(b: &mut Bencher) {
        take_eager(b, 50_000);
    }
}
<|endoftext|>


<commit_before>#![no_std]
#![feature(globs, phase, asm, macro_rules, lang_items, default_type_params, unboxed_closures)]

#[cfg(test)] #[phase(plugin, link)] extern crate log;

#[cfg(not(test))] extern crate "os_alloc" as alloc;
#[cfg(test)] extern crate alloc;
extern crate unicode;
#[phase(plugin, link)] extern crate core;
#[cfg(not(test))] extern crate "os_collections" as core_collections;
#[cfg(test)] extern crate "collections" as core_collections;
extern crate "rand" as core_rand;
extern crate rlibc;
extern crate spinlock;

// Make std testable by not duplicating lang items. See #2912
#[cfg(test)] extern crate "std" as realstd;
#[cfg(test)] pub use realstd::kinds;
#[cfg(test)] pub use realstd::ops;
#[cfg(test)] pub use realstd::cmp;
#[cfg(test)] pub use realstd::boxed;

// NB: These reexports are in the order they should be listed in rustdoc

pub use core::any;
pub use core::borrow;
pub use core::cell;
pub use core::clone;
#[cfg(not(test))] pub use core::cmp;
pub use core::default;
pub use core::finally;
pub use core::intrinsics;
pub use core::iter;
#[cfg(not(test))] pub use core::kinds;
pub use core::mem;
#[cfg(not(test))] pub use core::ops;
pub use core::ptr;
pub use core::raw;
pub use core::simd;
pub use core::result;
pub use core::option;

#[cfg(not(test))] pub use alloc::boxed;
pub use alloc::rc;

pub use core_collections::slice;
pub use core_collections::str;
pub use core_collections::string;
pub use core_collections::vec;

pub use unicode::char;

use core::prelude::*;

/* Exported macros */

pub mod macros;

/* new for os */
mod multiboot;
mod init;
mod vga_buffer;
mod fn_box;
mod scheduler;
mod global;

pub mod fmt;

/* Common data structures */

pub mod collections;

// Documentation for primitive types

mod bool;
mod unit;
mod tuple;

#[no_mangle]
pub fn main(multiboot: *const multiboot::Information) {

    unsafe{init::frame_stack(multiboot)};
    global::init();
    unsafe{enable_interrupts()};

    print!("test\n\niuaeiae");
    let x = box 5i;

    let y = 0xb8000 as *mut u64;
    unsafe{*y = 0xffffffffffffffff};

    print!("test");
    print!("test\n");
    println!("newline {}", x);

    scheduler::spawn(|| print!("I'm #1!\n"));

    fn test(name: &str) {
        loop {
            let mut x = 0u;
            for i in range(0,100000) {
                x = i;
            }
            print!("{}", name);
        }
    }

    scheduler::spawn(|| test("1"));
    scheduler::spawn(|| test("2"));
    scheduler::spawn(|| test("3"));
    scheduler::spawn(|| test("4"));
    scheduler::spawn(|| test("5"));
    scheduler::spawn(|| test("6"));
    loop{
        let mut x = 0u;
            for i in range(0,100000) {
                x = i;
            }
        print!("m");
    }

    loop{}
    panic!("end of os!");
}

// A curious inner-module that's not exported that contains the binding
// 'std' so that macro-expanded references to std::error and such
// can be resolved within libstd.
#[doc(hidden)]
mod std {
    // mods used for deriving
    pub use clone;
    pub use cmp;
    //pub use hash;

    //pub use comm; // used for select!()
    //pub use error; // used for try!()
    pub use fmt; // used for any formatting strings
    //pub use io; // used for println!()
    pub use option; // used for bitflags!{}
    //pub use rt; // used for panic!()
    pub use vec; // used for vec![]
    pub use cell; // used for tls!
    //pub use thread_local; // used for thread_local!
    pub use kinds; // used for tls!

    // The test runner calls ::std::os::args() but really wants realstd
    #[cfg(test)] pub use realstd::os as os;
    // The test runner requires std::slice::Vector, so re-export std::slice just for it.
    //
    // It is also used in vec![]
    pub use slice;

    pub use boxed; // used for vec![]

}

/* Interrupt Handlers */

extern {
//    fn send_eoi(interrupt_number: u64) -> int; //return value to mark rax as used
}

unsafe fn out_byte(port: u16, data: u8) {
    asm!("outb %al, %dx" :: "{dx}"(port), "{al}"(data) :: "volatile");
}
unsafe fn in_byte(port: u16) -> u8 {
    let ret: u8;
    asm!("inb %dx, %al" : "={al}"(ret) : "{dx}"(port) :: "volatile");
    ret
}
unsafe fn enable_interrupts() {
    asm!("sti" :::: "volatile");
}
unsafe fn disable_interrupts() {
    asm!("cli" :::: "volatile");
}

unsafe fn send_eoi(interrupt_number: u64) {
    unsafe fn send_master_eoi() {out_byte(0x20, 0x20)}
    unsafe fn send_slave_eoi() {out_byte(0xA0, 0x20)}

    match interrupt_number {
        i if i >= 40 => {
            send_slave_eoi();
            send_master_eoi();
        },
        32...40 => send_master_eoi(),
        _ => {},
    }
}

#[no_mangle]
pub extern "C" fn interrupt_handler(interrupt_number: u64, error_code: u64, rsp:uint) {
    match interrupt_number {
        13 if error_code != 0 => panic!(
            "General Protection Fault: Segment error at segment 0x{:x}", error_code),
        32 => {},
        33 => print!("k"),
        50 => panic!("out of memory"),
        66 => println!("ending thread..."),
        _ => panic!("unknown interrupt! number: {}, error_code: {:x}",interrupt_number, error_code),
    };
    unsafe{send_eoi(interrupt_number)};

    match interrupt_number {
        32 => unsafe{scheduler::reschedule(rsp)},
        66 => unsafe{scheduler::schedule()},
        _ => {},
    }
}

#[no_mangle]
pub extern "C" fn pagefault_handler(address: u64, error_code: u64, rsp:uint) {
    panic!("page fault: address: {:x}, error_code: {:b}, rsp: {:x}", address, error_code, rsp);
}

#[cfg(not(test))]
#[lang = "stack_exhausted"]
extern fn stack_exhausted() {panic!("stack exhausted");}

#[cfg(not(test))]
#[lang = "eh_personality"]
extern fn eh_personality() {unimplemented!();}

#[cfg(not(test))]
#[lang = "panic_fmt"]
fn panic_fmt() -> ! { unimplemented!();}

#[no_mangle]
#[allow(non_snake_case)]
pub extern "C" fn _Unwind_Resume() -> ! {
    unimplemented!();
}
<commit_msg>remove test code redundancy<commit_after>#![no_std]
#![feature(globs, phase, asm, macro_rules, lang_items, default_type_params, unboxed_closures)]

#[cfg(test)] #[phase(plugin, link)] extern crate log;

#[cfg(not(test))] extern crate "os_alloc" as alloc;
#[cfg(test)] extern crate alloc;
extern crate unicode;
#[phase(plugin, link)] extern crate core;
#[cfg(not(test))] extern crate "os_collections" as core_collections;
#[cfg(test)] extern crate "collections" as core_collections;
extern crate "rand" as core_rand;
extern crate rlibc;
extern crate spinlock;

// Make std testable by not duplicating lang items. See #2912
#[cfg(test)] extern crate "std" as realstd;
#[cfg(test)] pub use realstd::kinds;
#[cfg(test)] pub use realstd::ops;
#[cfg(test)] pub use realstd::cmp;
#[cfg(test)] pub use realstd::boxed;

// NB: These reexports are in the order they should be listed in rustdoc

pub use core::any;
pub use core::borrow;
pub use core::cell;
pub use core::clone;
#[cfg(not(test))] pub use core::cmp;
pub use core::default;
pub use core::finally;
pub use core::intrinsics;
pub use core::iter;
#[cfg(not(test))] pub use core::kinds;
pub use core::mem;
#[cfg(not(test))] pub use core::ops;
pub use core::ptr;
pub use core::raw;
pub use core::simd;
pub use core::result;
pub use core::option;

#[cfg(not(test))] pub use alloc::boxed;
pub use alloc::rc;

pub use core_collections::slice;
pub use core_collections::str;
pub use core_collections::string;
pub use core_collections::vec;

pub use unicode::char;

use core::prelude::*;

/* Exported macros */

pub mod macros;

/* new for os */
mod multiboot;
mod init;
mod vga_buffer;
mod fn_box;
mod scheduler;
mod global;

pub mod fmt;

/* Common data structures */

pub mod collections;

// Documentation for primitive types

mod bool;
mod unit;
mod tuple;

#[no_mangle]
pub fn main(multiboot: *const multiboot::Information) {

    unsafe{init::frame_stack(multiboot)};
    global::init();
    unsafe{enable_interrupts()};

    print!("test\n\niuaeiae");
    let x = box 5i;

    let y = 0xb8000 as *mut u64;
    unsafe{*y = 0xffffffffffffffff};

    print!("test");
    print!("test\n");
    println!("newline {}", x);

    scheduler::spawn(|| print!("I'm #1!\n"));

    fn test(name: &str) {
        loop {
            let mut x = 0u;
            for i in range(0,100000) {
                x = i;
            }
            print!("{}", name);
        }
    }

    scheduler::spawn(|| test("1"));
    scheduler::spawn(|| test("2"));
    scheduler::spawn(|| test("3"));
    scheduler::spawn(|| test("4"));
    scheduler::spawn(|| test("5"));
    scheduler::spawn(|| test("6"));

    test("m");

    loop{}
    panic!("end of os!");
}

// A curious inner-module that's not exported that contains the binding
// 'std' so that macro-expanded references to std::error and such
// can be resolved within libstd.
#[doc(hidden)]
mod std {
    // mods used for deriving
    pub use clone;
    pub use cmp;
    //pub use hash;

    //pub use comm; // used for select!()
    //pub use error; // used for try!()
    pub use fmt; // used for any formatting strings
    //pub use io; // used for println!()
    pub use option; // used for bitflags!{}
    //pub use rt; // used for panic!()
    pub use vec; // used for vec![]
    pub use cell; // used for tls!
    //pub use thread_local; // used for thread_local!
    pub use kinds; // used for tls!

    // The test runner calls ::std::os::args() but really wants realstd
    #[cfg(test)] pub use realstd::os as os;
    // The test runner requires std::slice::Vector, so re-export std::slice just for it.
    //
    // It is also used in vec![]
    pub use slice;

    pub use boxed; // used for vec![]

}

/* Interrupt Handlers */

extern {
//    fn send_eoi(interrupt_number: u64) -> int; //return value to mark rax as used
}

unsafe fn out_byte(port: u16, data: u8) {
    asm!("outb %al, %dx" :: "{dx}"(port), "{al}"(data) :: "volatile");
}
unsafe fn in_byte(port: u16) -> u8 {
    let ret: u8;
    asm!("inb %dx, %al" : "={al}"(ret) : "{dx}"(port) :: "volatile");
    ret
}
unsafe fn enable_interrupts() {
    asm!("sti" :::: "volatile");
}
unsafe fn disable_interrupts() {
    asm!("cli" :::: "volatile");
}

unsafe fn send_eoi(interrupt_number: u64) {
    unsafe fn send_master_eoi() {out_byte(0x20, 0x20)}
    unsafe fn send_slave_eoi() {out_byte(0xA0, 0x20)}

    match interrupt_number {
        i if i >= 40 => {
            send_slave_eoi();
            send_master_eoi();
        },
        32...40 => send_master_eoi(),
        _ => {},
    }
}

#[no_mangle]
pub extern "C" fn interrupt_handler(interrupt_number: u64, error_code: u64, rsp:uint) {
    match interrupt_number {
        13 if error_code != 0 => panic!(
            "General Protection Fault: Segment error at segment 0x{:x}", error_code),
        32 => {},
        33 => print!("k"),
        50 => panic!("out of memory"),
        66 => println!("ending thread..."),
        _ => panic!("unknown interrupt! number: {}, error_code: {:x}",interrupt_number, error_code),
    };
    unsafe{send_eoi(interrupt_number)};

    match interrupt_number {
        32 => unsafe{scheduler::reschedule(rsp)},
        66 => unsafe{scheduler::schedule()},
        _ => {},
    }
}

#[no_mangle]
pub extern "C" fn pagefault_handler(address: u64, error_code: u64, rsp:uint) {
    panic!("page fault: address: {:x}, error_code: {:b}, rsp: {:x}", address, error_code, rsp);
}

#[cfg(not(test))]
#[lang = "stack_exhausted"]
extern fn stack_exhausted() {panic!("stack exhausted");}

#[cfg(not(test))]
#[lang = "eh_personality"]
extern fn eh_personality() {unimplemented!();}

#[cfg(not(test))]
#[lang = "panic_fmt"]
fn panic_fmt() -> ! { unimplemented!();}

#[no_mangle]
#[allow(non_snake_case)]
pub extern "C" fn _Unwind_Resume() -> ! {
    unimplemented!();
}
<|endoftext|>


<commit_before>//! Compute timezones of points on the Earth.
//!
//! This is a direct port of
//! [bradfitz/latlong](https://github.com/bradfitz/latlong), and hence
//! the same bonuses/caveats apply:
//!
//! > It tries to have a small binary size (~360 KB), low memory footprint
//! > (~1 MB), and incredibly fast lookups (~0.5 microseconds). It does not
//! > try to be perfectly accurate when very close to borders.
//!
//! # Installation
//!
//! Add the following to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! tz-search = "0.1"
//! ```
//!
//! # Examples
//!
//! ```rust
//! assert_eq!(tz_search::lookup(-33.79, 151.17).unwrap(),
//!            "Australia/Sydney");
//!
//! // in the ocean
//! assert_eq!(tz_search::lookup(0.0, 0.0), None);
//! ```

extern crate flate2;
extern crate rustc_serialize;
extern crate byteorder;

use std::{cmp, mem, sync};
use std::sync::atomic;
use std::io::prelude::*;
use std::io::BufReader;
use byteorder::{BigEndian, ReadBytesExt};

#[allow(warnings)]
mod tables;

/// Attempt to compute the timezone that the point `lat`, `long`
/// lies in.
///
/// The latitude `lat` should lie in the range `[-90, 90]` with
/// negative representing south, and the longitude should lie in
/// the range `[-180, 180]` with negative representing west. This
/// will fail (return `None`) if the point lies in the ocean.
///
/// # Panics
///
/// `lookup` will panic if either of the two ranges above are
/// violated.
///
/// # Examples
///
/// ```rust
/// assert_eq!(tz_search::lookup(-33.79, 151.17).unwrap(),
///            "Australia/Sydney");
///
/// // in the ocean
/// assert_eq!(tz_search::lookup(0.0, 0.0), None);
/// ```
pub fn lookup(lat: f64, lon: f64) -> Option<String> {
    static SHARED: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;
    static ONCE: sync::Once = sync::ONCE_INIT;

    ONCE.call_once(|| {
        let s = Box::new(TzSearch::new());
        SHARED.store(unsafe {mem::transmute(s)}, atomic::Ordering::Relaxed);
    });

    let ptr = SHARED.load(atomic::Ordering::Relaxed);
    assert!(ptr != 0);
    let s = unsafe {&*(ptr as *const TzSearch)};
    s.lookup(lat, lon)
}

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
struct TileKey(u32);
impl TileKey {
    fn new(size: u8, x: u16, y: u16) -> TileKey {
        let size = size as u32;
        let x = x as u32;
        let y = y as u32;
        TileKey((size % 8) << 28 | (y % (1<<14)) << 14 | (x % (1<<14)))
    }
}

#[derive(Debug)]
struct TileLooker {
    tile: TileKey,
    idx: u16,
}
#[derive(Debug)]
struct ZoomLevel {
    tiles: Vec<TileLooker>
}
#[derive(Debug)]

/// All the information required for efficient time-zone lookups.
///
/// Unless you need absolutely strict control over memory use, you
/// probably want to call the top-level `lookup`, rather than going
/// via this.
pub struct TzSearch {
    leaves: Vec<Zone>,
    zoom_levels: Vec<ZoomLevel>,
}

enum Zone {
    StaticZone(String),
    OneBitTile([u16; 2], [u8; 8]),
    Pixmap([u8; 128])
}
impl std::fmt::Debug for Zone {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match *self {
            Zone::StaticZone(ref s) => write!(f, "StaticZone({:?})", s),
            Zone::OneBitTile(a, b) => write!(f, "OneBitTile({:?}, {:?})", &a[..], &b[..]),
            Zone::Pixmap(c) => write!(f, "Pixmap({:?})", &c[..]),
        }
    }
}

impl TzSearch {
    /// Create a new `TzSearch`.
    ///
    /// This is very expensive: the initialisation routine takes a
    /// long time, and the resulting structure uses a lot of
    /// memory. Hence, this should be called as rarely as
    /// possible.
    ///
    /// The free-standing `lookup` function internally manages
    /// creating exactly one of these, and should be preferred.
    pub fn new() -> TzSearch {
        let mut zoom_levels = vec![];
        for data in tables::REV_ZOOM_LEVELS.iter().rev() {
            let unb64d = rustc_serialize::base64::FromBase64::from_base64(*data).unwrap();
            let mut slurp = vec![];
            flate2::read::GzDecoder::new(&*unb64d).unwrap().read_to_end(&mut slurp).unwrap();
            assert_eq!(slurp.len() % 6, 0);
            let count = slurp.len() / 6;

            let mut slurp = &slurp[..];
            let mut tiles = Vec::with_capacity(count);
            for _ in 0..count {
                tiles.push(TileLooker {
                    tile: TileKey(slurp.read_u32::<BigEndian>().unwrap()),
                    idx: slurp.read_u16::<BigEndian>().unwrap(),
                })
            }

            zoom_levels.push(ZoomLevel { tiles: tiles })
        }
        let mut leaves = Vec::with_capacity(tables::NUM_LEAVES);
        let unb64d = rustc_serialize::base64::FromBase64::from_base64(tables::UNIQUE_LEAVES_PACKED)
            .unwrap();
        let mut ungz = BufReader::new(flate2::read::GzDecoder::new(&*unb64d).unwrap());
        let mut buf = [0; 128];
        for _ in 0..tables::NUM_LEAVES {
            let zone = match ungz.read_u8().unwrap() {
                b'S' => {
                    let mut zone_name = vec![];
                    ungz.read_until(0, &mut zone_name).unwrap();
                    if zone_name.last() == Some(&0) { zone_name.pop(); }
                    Zone::StaticZone(String::from_utf8(zone_name).unwrap())
                }
                b'2' => {
                    let idx = [ungz.read_u16::<BigEndian>().unwrap(),
                               ungz.read_u16::<BigEndian>().unwrap()];
                    let bits = ungz.read_u64::<BigEndian>().unwrap();
                    let mut rows = [0; 8];
                    for (y, place) in rows.iter_mut().enumerate() {
                        for x in 0..8 {
                            if bits & (1 << (y * 8 + x)) != 0 {
                                *place |= 1 << x
                            }
                        }
                    }
                    Zone::OneBitTile(idx, rows)
                }
                b'P' => {
                    let mut i = 0;
                    while i < 128 {
                        i += ungz.read(&mut buf[i..]).unwrap()
                    }
                    Zone::Pixmap(buf)
                }
                _ => panic!("unknown leaf type")
            };
            leaves.push(zone)
        }

        TzSearch {
            zoom_levels: zoom_levels,
            leaves: leaves
        }
    }

    /// Attempt to compute the timezone that the point `lat`, `long`
    /// lies in.
    ///
    /// The latitude `lat` should lie in the range `[-90, 90]` with
    /// negative representing south, and the longitude should lie in
    /// the range `[-180, 180]` with negative representing west. This
    /// will fail (return `None`) if the point lies in the ocean.
    ///
    /// See also: the `lookup` function at the top-level.
    ///
    /// # Panics
    ///
    /// `lookup` will panic if either of the two ranges above are
    /// violated.
    ///
    /// # Examples
    ///
    /// ```rust
    /// let s = tz_search::TzSearch::new();
    /// let (lat, long) = (-33.79, 151.17);
    /// assert_eq!(s.lookup(-33.79, 151.17).unwrap(),
    ///            "Australia/Sydney");
    /// ```
    pub fn lookup(&self, lat: f64, long: f64) -> Option<String> {
        fn clamp(x: isize, lim: isize) -> usize {
            cmp::max(0, cmp::min(lim * tables::DEG_PIXELS as isize, x)) as usize
        }
        assert!(-90.0 <= lat && lat <= 90.0);
        assert!(-180.0 <= long && long <= 180.0);
        let x = ((long + 180.0) * (tables::DEG_PIXELS as f64)) as isize;
        let y = ((90.0 - lat) * (tables::DEG_PIXELS as f64)) as isize;
        let x = clamp(x, 360);
        let y = clamp(y, 180);

        self.lookup_pixel(x, y)
    }
    fn lookup_pixel(&self, x: usize, y: usize) -> Option<String> {
        for level in (0..6).rev() {
            let shift = 3 + level;
            let xt = x >> shift;
            let yt = y >> shift;
            let tk = TileKey::new(level, xt as u16, yt as u16);
            if let Some(ret) = self.zoom_level_lookup(&self.zoom_levels[level as usize], x, y, tk) {
                return ret
            }
        }
        None
    }

    fn zone_lookup(&self, zone: &Zone, x: usize, y: usize, tk: TileKey) -> Option<Option<String>> {
        match *zone {
            Zone::StaticZone(ref s) => Some(Some(s.clone())),
            Zone::OneBitTile(idxs, rows) => {
                let idx = if rows[y & 7] & (1 << (x & 7)) != 0 {
                    idxs[1]
                } else {
                    idxs[0]
                };
                self.zone_lookup(&self.leaves[idx as usize], x, y, tk)
            }
            Zone::Pixmap(ref p) => {
                let xx = x & 7;
                let yy = y & 7;
                let i = 2 * (yy * 8 + xx);
                let idx = ((p[i] as usize) << 8) + p[i+1] as usize;
                const OCEAN_INDEX: usize = 0xFFFF;
                if idx == OCEAN_INDEX {
                    Some(None)
                } else {
                    self.zone_lookup(&self.leaves[idx], x, y, tk)
                }
            }
        }
    }

    fn zoom_level_lookup(&self, zl: &ZoomLevel, x: usize, y: usize, tk: TileKey)
                         -> Option<Option<String>>
    {
        let pos = zl.tiles.binary_search_by(|t| t.tile.cmp(&tk)).unwrap_or_else(|x| x);

        match zl.tiles.get(pos) {
            Some(tl) if tl.tile == tk => self.zone_lookup(&self.leaves[tl.idx as usize], x, y, tk),
            _ => None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{lookup, TzSearch};

    #[test]
    fn loads_ok() {
        let _searcher = TzSearch::new();
    }

    #[test]
    fn test_lookup_lat_long() {
        let searcher = TzSearch::new();
        let tests = [(37.7833, -122.4167, Some("America/Los_Angeles")),
                     (-33.79, 151.17, Some("Australia/Sydney"))];
        for &(lat, lon, want) in &tests {
            let want = want.map(|s| s.to_string());
            assert_eq!(searcher.lookup(lat, lon), want);
            assert_eq!(lookup(lat, lon), want)
        }
    }
    #[test]
    fn test_lookup_pixel() {
        let searcher = TzSearch::new();
        let tests = [
            (9200, 2410, Some("Asia/Phnom_Penh")),
            (9047, 2488, Some("Asia/Phnom_Penh")),

            // one-bit leaf tile:
            (9290, 530, Some("Asia/Krasnoyarsk")),
            (9290, 531, Some("Asia/Yakutsk")),

            // four-bit tile:
            (2985, 1654, Some("America/Indiana/Vincennes")),
            (2986, 1654, Some("America/Indiana/Marengo")),
            (2986, 1655, Some("America/Indiana/Tell_City")),

            // Empty tile:
            (4000, 2000, None),

            // Big 1-color tile in ocean with island:
            (3687, 1845, Some("Atlantic/Bermuda")),
            // Same, but off Oregon coast:
            (1747, 1486, Some("America/Los_Angeles")),

            // Little solid tile:
            (2924, 2316, Some("America/Belize")),
            ];

        for &(lat, lon, ref want) in &tests {
            assert_eq!(searcher.lookup_pixel(lat, lon), want.map(|s| s.to_string()));
        }
    }
}
<commit_msg>Link to github.<commit_after>//! Compute timezones of points on the Earth.
//!
//! This is a direct port of
//! [bradfitz/latlong](https://github.com/bradfitz/latlong), and hence
//! the same bonuses/caveats apply:
//!
//! > It tries to have a small binary size (~360 KB), low memory footprint
//! > (~1 MB), and incredibly fast lookups (~0.5 microseconds). It does not
//! > try to be perfectly accurate when very close to borders.
//!
//! [Source](https://github.com/huonw/tz-search).
//!
//! # Installation
//!
//! Add the following to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! tz-search = "0.1"
//! ```
//!
//! # Examples
//!
//! ```rust
//! assert_eq!(tz_search::lookup(-33.79, 151.17).unwrap(),
//!            "Australia/Sydney");
//!
//! // in the ocean
//! assert_eq!(tz_search::lookup(0.0, 0.0), None);
//! ```

extern crate flate2;
extern crate rustc_serialize;
extern crate byteorder;

use std::{cmp, mem, sync};
use std::sync::atomic;
use std::io::prelude::*;
use std::io::BufReader;
use byteorder::{BigEndian, ReadBytesExt};

#[allow(warnings)]
mod tables;

/// Attempt to compute the timezone that the point `lat`, `long`
/// lies in.
///
/// The latitude `lat` should lie in the range `[-90, 90]` with
/// negative representing south, and the longitude should lie in
/// the range `[-180, 180]` with negative representing west. This
/// will fail (return `None`) if the point lies in the ocean.
///
/// # Panics
///
/// `lookup` will panic if either of the two ranges above are
/// violated.
///
/// # Examples
///
/// ```rust
/// assert_eq!(tz_search::lookup(-33.79, 151.17).unwrap(),
///            "Australia/Sydney");
///
/// // in the ocean
/// assert_eq!(tz_search::lookup(0.0, 0.0), None);
/// ```
pub fn lookup(lat: f64, lon: f64) -> Option<String> {
    static SHARED: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT;
    static ONCE: sync::Once = sync::ONCE_INIT;

    ONCE.call_once(|| {
        let s = Box::new(TzSearch::new());
        SHARED.store(unsafe {mem::transmute(s)}, atomic::Ordering::Relaxed);
    });

    let ptr = SHARED.load(atomic::Ordering::Relaxed);
    assert!(ptr != 0);
    let s = unsafe {&*(ptr as *const TzSearch)};
    s.lookup(lat, lon)
}

#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
struct TileKey(u32);
impl TileKey {
    fn new(size: u8, x: u16, y: u16) -> TileKey {
        let size = size as u32;
        let x = x as u32;
        let y = y as u32;
        TileKey((size % 8) << 28 | (y % (1<<14)) << 14 | (x % (1<<14)))
    }
}

#[derive(Debug)]
struct TileLooker {
    tile: TileKey,
    idx: u16,
}
#[derive(Debug)]
struct ZoomLevel {
    tiles: Vec<TileLooker>
}
#[derive(Debug)]

/// All the information required for efficient time-zone lookups.
///
/// Unless you need absolutely strict control over memory use, you
/// probably want to call the top-level `lookup`, rather than going
/// via this.
pub struct TzSearch {
    leaves: Vec<Zone>,
    zoom_levels: Vec<ZoomLevel>,
}

enum Zone {
    StaticZone(String),
    OneBitTile([u16; 2], [u8; 8]),
    Pixmap([u8; 128])
}
impl std::fmt::Debug for Zone {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match *self {
            Zone::StaticZone(ref s) => write!(f, "StaticZone({:?})", s),
            Zone::OneBitTile(a, b) => write!(f, "OneBitTile({:?}, {:?})", &a[..], &b[..]),
            Zone::Pixmap(c) => write!(f, "Pixmap({:?})", &c[..]),
        }
    }
}

impl TzSearch {
    /// Create a new `TzSearch`.
    ///
    /// This is very expensive: the initialisation routine takes a
    /// long time, and the resulting structure uses a lot of
    /// memory. Hence, this should be called as rarely as
    /// possible.
    ///
    /// The free-standing `lookup` function internally manages
    /// creating exactly one of these, and should be preferred.
    pub fn new() -> TzSearch {
        let mut zoom_levels = vec![];
        for data in tables::REV_ZOOM_LEVELS.iter().rev() {
            let unb64d = rustc_serialize::base64::FromBase64::from_base64(*data).unwrap();
            let mut slurp = vec![];
            flate2::read::GzDecoder::new(&*unb64d).unwrap().read_to_end(&mut slurp).unwrap();
            assert_eq!(slurp.len() % 6, 0);
            let count = slurp.len() / 6;

            let mut slurp = &slurp[..];
            let mut tiles = Vec::with_capacity(count);
            for _ in 0..count {
                tiles.push(TileLooker {
                    tile: TileKey(slurp.read_u32::<BigEndian>().unwrap()),
                    idx: slurp.read_u16::<BigEndian>().unwrap(),
                })
            }

            zoom_levels.push(ZoomLevel { tiles: tiles })
        }
        let mut leaves = Vec::with_capacity(tables::NUM_LEAVES);
        let unb64d = rustc_serialize::base64::FromBase64::from_base64(tables::UNIQUE_LEAVES_PACKED)
            .unwrap();
        let mut ungz = BufReader::new(flate2::read::GzDecoder::new(&*unb64d).unwrap());
        let mut buf = [0; 128];
        for _ in 0..tables::NUM_LEAVES {
            let zone = match ungz.read_u8().unwrap() {
                b'S' => {
                    let mut zone_name = vec![];
                    ungz.read_until(0, &mut zone_name).unwrap();
                    if zone_name.last() == Some(&0) { zone_name.pop(); }
                    Zone::StaticZone(String::from_utf8(zone_name).unwrap())
                }
                b'2' => {
                    let idx = [ungz.read_u16::<BigEndian>().unwrap(),
                               ungz.read_u16::<BigEndian>().unwrap()];
                    let bits = ungz.read_u64::<BigEndian>().unwrap();
                    let mut rows = [0; 8];
                    for (y, place) in rows.iter_mut().enumerate() {
                        for x in 0..8 {
                            if bits & (1 << (y * 8 + x)) != 0 {
                                *place |= 1 << x
                            }
                        }
                    }
                    Zone::OneBitTile(idx, rows)
                }
                b'P' => {
                    let mut i = 0;
                    while i < 128 {
                        i += ungz.read(&mut buf[i..]).unwrap()
                    }
                    Zone::Pixmap(buf)
                }
                _ => panic!("unknown leaf type")
            };
            leaves.push(zone)
        }

        TzSearch {
            zoom_levels: zoom_levels,
            leaves: leaves
        }
    }

    /// Attempt to compute the timezone that the point `lat`, `long`
    /// lies in.
    ///
    /// The latitude `lat` should lie in the range `[-90, 90]` with
    /// negative representing south, and the longitude should lie in
    /// the range `[-180, 180]` with negative representing west. This
    /// will fail (return `None`) if the point lies in the ocean.
    ///
    /// See also: the `lookup` function at the top-level.
    ///
    /// # Panics
    ///
    /// `lookup` will panic if either of the two ranges above are
    /// violated.
    ///
    /// # Examples
    ///
    /// ```rust
    /// let s = tz_search::TzSearch::new();
    /// let (lat, long) = (-33.79, 151.17);
    /// assert_eq!(s.lookup(-33.79, 151.17).unwrap(),
    ///            "Australia/Sydney");
    /// ```
    pub fn lookup(&self, lat: f64, long: f64) -> Option<String> {
        fn clamp(x: isize, lim: isize) -> usize {
            cmp::max(0, cmp::min(lim * tables::DEG_PIXELS as isize, x)) as usize
        }
        assert!(-90.0 <= lat && lat <= 90.0);
        assert!(-180.0 <= long && long <= 180.0);
        let x = ((long + 180.0) * (tables::DEG_PIXELS as f64)) as isize;
        let y = ((90.0 - lat) * (tables::DEG_PIXELS as f64)) as isize;
        let x = clamp(x, 360);
        let y = clamp(y, 180);

        self.lookup_pixel(x, y)
    }
    fn lookup_pixel(&self, x: usize, y: usize) -> Option<String> {
        for level in (0..6).rev() {
            let shift = 3 + level;
            let xt = x >> shift;
            let yt = y >> shift;
            let tk = TileKey::new(level, xt as u16, yt as u16);
            if let Some(ret) = self.zoom_level_lookup(&self.zoom_levels[level as usize], x, y, tk) {
                return ret
            }
        }
        None
    }

    fn zone_lookup(&self, zone: &Zone, x: usize, y: usize, tk: TileKey) -> Option<Option<String>> {
        match *zone {
            Zone::StaticZone(ref s) => Some(Some(s.clone())),
            Zone::OneBitTile(idxs, rows) => {
                let idx = if rows[y & 7] & (1 << (x & 7)) != 0 {
                    idxs[1]
                } else {
                    idxs[0]
                };
                self.zone_lookup(&self.leaves[idx as usize], x, y, tk)
            }
            Zone::Pixmap(ref p) => {
                let xx = x & 7;
                let yy = y & 7;
                let i = 2 * (yy * 8 + xx);
                let idx = ((p[i] as usize) << 8) + p[i+1] as usize;
                const OCEAN_INDEX: usize = 0xFFFF;
                if idx == OCEAN_INDEX {
                    Some(None)
                } else {
                    self.zone_lookup(&self.leaves[idx], x, y, tk)
                }
            }
        }
    }

    fn zoom_level_lookup(&self, zl: &ZoomLevel, x: usize, y: usize, tk: TileKey)
                         -> Option<Option<String>>
    {
        let pos = zl.tiles.binary_search_by(|t| t.tile.cmp(&tk)).unwrap_or_else(|x| x);

        match zl.tiles.get(pos) {
            Some(tl) if tl.tile == tk => self.zone_lookup(&self.leaves[tl.idx as usize], x, y, tk),
            _ => None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{lookup, TzSearch};

    #[test]
    fn loads_ok() {
        let _searcher = TzSearch::new();
    }

    #[test]
    fn test_lookup_lat_long() {
        let searcher = TzSearch::new();
        let tests = [(37.7833, -122.4167, Some("America/Los_Angeles")),
                     (-33.79, 151.17, Some("Australia/Sydney"))];
        for &(lat, lon, want) in &tests {
            let want = want.map(|s| s.to_string());
            assert_eq!(searcher.lookup(lat, lon), want);
            assert_eq!(lookup(lat, lon), want)
        }
    }
    #[test]
    fn test_lookup_pixel() {
        let searcher = TzSearch::new();
        let tests = [
            (9200, 2410, Some("Asia/Phnom_Penh")),
            (9047, 2488, Some("Asia/Phnom_Penh")),

            // one-bit leaf tile:
            (9290, 530, Some("Asia/Krasnoyarsk")),
            (9290, 531, Some("Asia/Yakutsk")),

            // four-bit tile:
            (2985, 1654, Some("America/Indiana/Vincennes")),
            (2986, 1654, Some("America/Indiana/Marengo")),
            (2986, 1655, Some("America/Indiana/Tell_City")),

            // Empty tile:
            (4000, 2000, None),

            // Big 1-color tile in ocean with island:
            (3687, 1845, Some("Atlantic/Bermuda")),
            // Same, but off Oregon coast:
            (1747, 1486, Some("America/Los_Angeles")),

            // Little solid tile:
            (2924, 2316, Some("America/Belize")),
            ];

        for &(lat, lon, ref want) in &tests {
            assert_eq!(searcher.lookup_pixel(lat, lon), want.map(|s| s.to_string()));
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate log;
extern crate serde;
extern crate serde_json;

use serde::{Serialize, Serializer};
use serde_json::to_string;
use std::fmt::Debug;
use std::io::{Write, stderr, stdout};

#[macro_export]
macro_rules! json_object {
    ( $( $key:expr => $value:expr ),* ) => {
        {
            use serde_json::{Value, to_value, Map};
            let mut m: Map<String, Value> = Map::new();
            $(
                m.insert($key.to_owned(), to_value(&$value));
            )*
            m
        }
    }
}

#[macro_export]
macro_rules! json_format {
    ( $( $key:expr => $value:expr ),* ) => {
        {
            let mut s = String::with_capacity(128);
            s.push('{');
            $(
                s.push_str(&format!(r#""{}":{},"#, $key, $value));
            )*
            let _ = s.pop();
            s.push('}');
            s
        }
    };
    ( $( $value:expr ),* ) => {
        {
            let mut s = String::with_capacity(128);
            s.push('[');
            $(
                s.push_str(&format!("{},", $value));
            )*
            let _ = s.pop();
            s.push(']');
            s
        }
    }
}

#[macro_export]
macro_rules! q {
    ( $value:expr ) => {
        format!("\"{}\"", $value)
    }
}

pub trait StructuredLog {
    fn slog(&self) -> String;
}

#[inline]
pub fn serialize<T>(value: &T) -> String
    where T: StructuredLog
{
    value.slog()
}

#[derive(Debug)]
pub struct SLogJson<'a, T: 'a>(&'a T);

impl<'a, T: Serialize> SLogJson<'a, T> {
    pub fn new<'b>(value: &'b T) -> SLogJson<'b, T> {
        SLogJson(value)
    }
}

impl<'a, T: Serialize> Serialize for SLogJson<'a, T> {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer
    {
        self.0.serialize(serializer)
    }
}

impl<'a, T: Serialize + Debug> StructuredLog for SLogJson<'a, T> {
    fn slog(&self) -> String {
        to_string(self).unwrap_or_else(|err| {
            json_format! {
                "format_error" => q!(escape_str(&format!("{:?}", err))),
                "value" => q!(escape_str(&format!("{:?}", self)))
            }
        })
    }
}

#[macro_export]
macro_rules! s_error {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_error!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_warn {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_warn!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_info {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_info!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_debug {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_debug!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_trace {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_trace!(target: module_path!(), $value);
        }
    };
}

pub fn escape_str(x: &str) -> String {
    let mut v = Vec::new();
    let mut l = 0;
    let bytes = x.as_bytes();
    for (i, b) in bytes.iter().enumerate() {
        let escaped = match *b {
            b'\x00' => r"\u0000",
            b'\x01' => r"\u0001",
            b'\x02' => r"\u0002",
            b'\x03' => r"\u0003",
            b'\x04' => r"\u0004",
            b'\x05' => r"\u0005",
            b'\x06' => r"\u0006",
            b'\x07' => r"\u0007",
            b'\x08' => r"\b",
            b'\x09' => r"\t",
            b'\x10' => r"\n",
            b'\x11' => r"\u0011",
            b'\x12' => r"\f",
            b'\x13' => r"\r",
            b'\x14' => r"\u0014",
            b'\x15' => r"\u0015",
            b'\x16' => r"\u0016",
            b'\x17' => r"\u0017",
            b'\x18' => r"\u0018",
            b'\x19' => r"\u0019",
            b'\\' => r"\\",
            b'"' => r#"\""#,
            _ => {
                continue;
            }
        };

        if l < i {
            v.extend_from_slice(&bytes[l..i]);
        }

        v.extend_from_slice(escaped.as_bytes());

        l = i + 1;
    }

    String::from_utf8(v).unwrap()
}

pub enum LoggerOutput {
    Stdout,
    Stderr,
}

struct JsonLogger {
    filter: log::LogLevelFilter,
    output: LoggerOutput,
}

impl log::Log for JsonLogger {
    fn enabled(&self, metadata: &log::LogMetadata) -> bool {
        metadata.level() <= self.filter
    }

    fn log(&self, record: &log::LogRecord) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let json = json_format! {
            "level" => q!(record.level()),
            "meta" => json_format! {
                "target" => q!(record.target()),
                "location" => json_format! {
                    "module_path" => q!(record.location().module_path()),
                    "file" => q!(record.location().file()),
                    "line" => record.location().line()
                }
            },
            "value" => if record.target().starts_with("json:") {
                format!("{}", record.args())
            } else {
                q!(escape_str(&record.args().to_string()))
            }
        };

        let _ = match self.output {
            LoggerOutput::Stderr => writeln!(stderr(), "{}", json),
            LoggerOutput::Stdout => writeln!(stdout(), "{}", json),
        };
    }
}

pub fn init(output: LoggerOutput, log_level: log::LogLevelFilter) {
    let logger = JsonLogger {
        filter: log_level,
        output: output,
    };

    log::set_logger(|max_log_level| {
            max_log_level.set(logger.filter);
            Box::new(logger)
        })
        .unwrap();
}

#[cfg(feature = "documentation")]
pub mod doc;

#[cfg(test)]
mod tests {
    #[test]
    fn simple_logger() {}

    #[test]
    fn escape_str() {
        let x = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\\\"";
        let expected = r#"\u0000\u0001\u0002\u0003\u0004\u0005\u0006\u0007\b\t\n\u0011\f\r\u0014\u0015\u0016\u0017\u0018\u0019\\\""#.to_owned();
        assert_eq!(::escape_str(x), expected);
    }

    #[test]
    fn json_format() {
        let obj = json_format! {
            "key1" => q!("value1"),
            "key2" => 1
        };

        let array = json_format![q!("value1"), 1];

        assert_eq!(obj, r#"{"key1":"value1","key2":1}"#);
        assert_eq!(array, r#"["value1",1]"#);
    }
}
<commit_msg>Fix the all s_*! calling trace! bug<commit_after>#[macro_use]
extern crate log;
extern crate serde;
extern crate serde_json;

use serde::{Serialize, Serializer};
use serde_json::to_string;
use std::fmt::Debug;
use std::io::{Write, stderr, stdout};

#[macro_export]
macro_rules! json_object {
    ( $( $key:expr => $value:expr ),* ) => {
        {
            use serde_json::{Value, to_value, Map};
            let mut m: Map<String, Value> = Map::new();
            $(
                m.insert($key.to_owned(), to_value(&$value));
            )*
            m
        }
    }
}

#[macro_export]
macro_rules! json_format {
    ( $( $key:expr => $value:expr ),* ) => {
        {
            let mut s = String::with_capacity(128);
            s.push('{');
            $(
                s.push_str(&format!(r#""{}":{},"#, $key, $value));
            )*
            let _ = s.pop();
            s.push('}');
            s
        }
    };
    ( $( $value:expr ),* ) => {
        {
            let mut s = String::with_capacity(128);
            s.push('[');
            $(
                s.push_str(&format!("{},", $value));
            )*
            let _ = s.pop();
            s.push(']');
            s
        }
    }
}

#[macro_export]
macro_rules! q {
    ( $value:expr ) => {
        format!("\"{}\"", $value)
    }
}

pub trait StructuredLog {
    fn slog(&self) -> String;
}

#[inline]
pub fn serialize<T>(value: &T) -> String
    where T: StructuredLog
{
    value.slog()
}

#[derive(Debug)]
pub struct SLogJson<'a, T: 'a>(&'a T);

impl<'a, T: Serialize> SLogJson<'a, T> {
    pub fn new<'b>(value: &'b T) -> SLogJson<'b, T> {
        SLogJson(value)
    }
}

impl<'a, T: Serialize> Serialize for SLogJson<'a, T> {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer
    {
        self.0.serialize(serializer)
    }
}

impl<'a, T: Serialize + Debug> StructuredLog for SLogJson<'a, T> {
    fn slog(&self) -> String {
        to_string(self).unwrap_or_else(|err| {
            json_format! {
                "format_error" => q!(escape_str(&format!("{:?}", err))),
                "value" => q!(escape_str(&format!("{:?}", self)))
            }
        })
    }
}

#[macro_export]
macro_rules! s_error {
    (target: $target:expr, $value:expr) => {
        {
            error!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_error!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_warn {
    (target: $target:expr, $value:expr) => {
        {
            warn!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_warn!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_info {
    (target: $target:expr, $value:expr) => {
        {
            info!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_info!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_debug {
    (target: $target:expr, $value:expr) => {
        {
            debug!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_debug!(target: module_path!(), $value);
        }
    };
}

#[macro_export]
macro_rules! s_trace {
    (target: $target:expr, $value:expr) => {
        {
            trace!(target: &format!("json:{}", $target), "{}", $crate::serialize(&$crate::SLogJson::new(&$value)));
        }
    };
    ($value:expr) => {
        {
            s_trace!(target: module_path!(), $value);
        }
    };
}

pub fn escape_str(x: &str) -> String {
    let mut v = Vec::new();
    let mut l = 0;
    let bytes = x.as_bytes();
    for (i, b) in bytes.iter().enumerate() {
        let escaped = match *b {
            b'\x00' => r"\u0000",
            b'\x01' => r"\u0001",
            b'\x02' => r"\u0002",
            b'\x03' => r"\u0003",
            b'\x04' => r"\u0004",
            b'\x05' => r"\u0005",
            b'\x06' => r"\u0006",
            b'\x07' => r"\u0007",
            b'\x08' => r"\b",
            b'\x09' => r"\t",
            b'\x10' => r"\n",
            b'\x11' => r"\u0011",
            b'\x12' => r"\f",
            b'\x13' => r"\r",
            b'\x14' => r"\u0014",
            b'\x15' => r"\u0015",
            b'\x16' => r"\u0016",
            b'\x17' => r"\u0017",
            b'\x18' => r"\u0018",
            b'\x19' => r"\u0019",
            b'\\' => r"\\",
            b'"' => r#"\""#,
            _ => {
                continue;
            }
        };

        if l < i {
            v.extend_from_slice(&bytes[l..i]);
        }

        v.extend_from_slice(escaped.as_bytes());

        l = i + 1;
    }

    String::from_utf8(v).unwrap()
}

pub enum LoggerOutput {
    Stdout,
    Stderr,
}

struct JsonLogger {
    filter: log::LogLevelFilter,
    output: LoggerOutput,
}

impl log::Log for JsonLogger {
    fn enabled(&self, metadata: &log::LogMetadata) -> bool {
        metadata.level() <= self.filter
    }

    fn log(&self, record: &log::LogRecord) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let json = json_format! {
            "level" => q!(record.level()),
            "meta" => json_format! {
                "target" => q!(record.target()),
                "location" => json_format! {
                    "module_path" => q!(record.location().module_path()),
                    "file" => q!(record.location().file()),
                    "line" => record.location().line()
                }
            },
            "value" => if record.target().starts_with("json:") {
                format!("{}", record.args())
            } else {
                q!(escape_str(&record.args().to_string()))
            }
        };

        let _ = match self.output {
            LoggerOutput::Stderr => writeln!(stderr(), "{}", json),
            LoggerOutput::Stdout => writeln!(stdout(), "{}", json),
        };
    }
}

pub fn init(output: LoggerOutput, log_level: log::LogLevelFilter) {
    let logger = JsonLogger {
        filter: log_level,
        output: output,
    };

    log::set_logger(|max_log_level| {
            max_log_level.set(logger.filter);
            Box::new(logger)
        })
        .unwrap();
}

#[cfg(feature = "documentation")]
pub mod doc;

#[cfg(test)]
mod tests {
    #[test]
    fn simple_logger() {}

    #[test]
    fn escape_str() {
        let x = "\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\\\"";
        let expected = r#"\u0000\u0001\u0002\u0003\u0004\u0005\u0006\u0007\b\t\n\u0011\f\r\u0014\u0015\u0016\u0017\u0018\u0019\\\""#.to_owned();
        assert_eq!(::escape_str(x), expected);
    }

    #[test]
    fn json_format() {
        let obj = json_format! {
            "key1" => q!("value1"),
            "key2" => 1
        };

        let array = json_format![q!("value1"), 1];

        assert_eq!(obj, r#"{"key1":"value1","key2":1}"#);
        assert_eq!(array, r#"["value1",1]"#);
    }
}
<|endoftext|>


<commit_before>#![allow(unused_features)]
#![feature(asm, libc, rustc_private, page_size, core, alloc, rt, fnbox, box_raw)]

#[macro_use]
extern crate log;
extern crate libc;
<<<<<<< HEAD
extern crate memmap;
=======
extern crate mmap;
extern crate simd;
>>>>>>> feature-fnbox

pub use context::Context;
pub use stack::Stack;

pub mod context;
pub mod stack;
mod sys;
<commit_msg>forget to merge lib.rs<commit_after>#![allow(unused_features)]
#![feature(asm, libc, rustc_private, page_size, core, alloc, rt, fnbox, box_raw)]

#[macro_use]
extern crate log;
extern crate libc;
extern crate memmap;
extern crate simd;

pub use context::Context;
pub use stack::Stack;

pub mod context;
pub mod stack;
mod sys;
<|endoftext|>


<commit_before>#![feature(optin_builtin_traits)]
#![feature(unsafe_destructor)]
#![feature(core)]
#![feature(std_misc)]
#![feature(test)]
extern crate time;

use std::cell::UnsafeCell;
pub use std::sync::{
    TryLockError,
    LockResult,
    TryLockResult,
    PoisonError,
};
use std::thread;
use std::ops::{Deref, DerefMut};

use std::sync::atomic::{Ordering, AtomicUsize};

// Panic guards, taken from std::sync

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

struct Flag { failed: UnsafeCell<bool> }
unsafe impl Send for Flag {}
unsafe impl Sync for Flag {}

impl Flag {
    #[inline]
    fn new() -> Flag {
        Flag { failed : UnsafeCell { value : false }}
    }

    #[inline]
    fn borrow(&self) -> LockResult<Guard> {
        let ret = Guard { panicking: thread::panicking() };
        if unsafe { *self.failed.get() } {
            Err(PoisonError::new(ret))
        } else {
            Ok(ret)
        }
    }

    #[inline]
    fn done(&self, guard: &Guard) {
        if !guard.panicking && thread::panicking() {
            unsafe { *self.failed.get() = true; }
        }
    }

    #[inline]
    fn get(&self) -> bool {
        unsafe { *self.failed.get() }
    }
}

struct Guard {
    panicking: bool,
}

fn map_result<T, U, F>(result: LockResult<T>, f: F)
                           -> LockResult<U>
                           where F: FnOnce(T) -> U {
    match result {
        Ok(t) => Ok(f(t)),
        Err(err) => Err(PoisonError::new(f(err.into_inner())))
    }
}

// SpinLock implementation
//
const SHARED : usize = 0x80000000;
const EXCLWAIT : usize = 0x00100000;
const MAXWAIT : u64 = 2_000_000_000;
const SPIN_COUNT : u32 = 0xFF;

pub struct SpinLock<T> {
    count : AtomicUsize,
    data : UnsafeCell<T>,
    poison : Flag,
}

unsafe impl<T: Send + Sync> Send for SpinLock<T> {}
unsafe impl<T: Send + Sync> Sync for SpinLock<T> {}

impl<T : Send + Sync> SpinLock<T> {
    pub fn new(data : T) -> SpinLock<T> {
        SpinLock {
            count : AtomicUsize::new(0),
            data : UnsafeCell::new(data),
            poison : Flag::new(),
        }
    }

    #[inline]
    pub fn write(&self) -> LockResult<SpinLockWriteGuard<T>> {
        if self.count.fetch_add(1, Ordering::Acquire) != 0 {
            return self.write_contested()
        }
        SpinLockWriteGuard::new(self, &self.data)
    }

    pub fn write_contested(&self) -> LockResult<SpinLockWriteGuard<T>> {
        let mut base_time : u64 = 0;
        let mut time : u64;
        let mut i = 0u32;

        self.count.fetch_add(EXCLWAIT - 1, Ordering::Relaxed);
        self.count.fetch_and(! SHARED, Ordering::Relaxed);
        loop {
            // Use relaxed ordering to avoid trashing the cache coherency handling for free
            let count = self.count.load(Ordering::Relaxed);
            if count & (EXCLWAIT - 1) == 0 &&
                self.count.compare_and_swap(count, count - EXCLWAIT + 1, Ordering::Acquire) == count {
                    break;
                }

            if i >= SPIN_COUNT  {
                time = time::precise_time_ns();
                if base_time == 0 {
                    base_time = time;
                } else if time - base_time > MAXWAIT {
                    // XXX we converted the SHARED locks to exclusive ones, i'm not sure how bad
                    // actually is
                    self.count.fetch_sub(EXCLWAIT, Ordering::Relaxed);
                    panic!("Spinning on a spin lock for too long")
                }
                thread::yield_now();
            } else {
                i += 1;
            }
        }
        SpinLockWriteGuard::new(self, &self.data)
    }

    #[inline]
    pub fn read(&self) -> LockResult<SpinLockReadGuard<T>> {
        if self.count.compare_and_swap(0, SHARED | 1, Ordering::Acquire) != 0 {
            return self.read_contested()
        }
        SpinLockReadGuard::new(self, &self.data)
    }

    pub fn read_contested(&self) -> LockResult<SpinLockReadGuard<T>> {
        let mut i = 0u32;
        let mut base_time : u64 = 0;
        let mut time : u64;
        loop {
            let count = self.count.load(Ordering::Relaxed);
            if count == 0 {
                if self.count.compare_and_swap(0, SHARED | 1, Ordering::Acquire) == 0 {
                    break
                }
            } else if count & SHARED == SHARED {
                if self.count.compare_and_swap(count, count + 1, Ordering::Acquire) == count {
                    break
                }
            }
            if i >= SPIN_COUNT  {
                time = time::precise_time_ns();
                if base_time == 0 {
                    base_time = time;
                } else if time - base_time > MAXWAIT {
                    panic!("Spinning on a spin lock for too long")
                }
                thread::yield_now();
            } else {
                i+=1;
            }
        }
        SpinLockReadGuard::new(self, &self.data)
    }

    #[inline]
    pub fn try_read(&self) -> TryLockResult<SpinLockReadGuard<T>> {
        if self.count.compare_and_swap(0, 1 | SHARED, Ordering::Acquire) == 0 {
            return Ok(try!(SpinLockReadGuard::new(self, &self.data)))
        } else {
            let count = self.count.load(Ordering::Relaxed);
            if count & SHARED == SHARED &&
                self.count.compare_and_swap(count, count + 1, Ordering::Acquire) == count {
                return Ok(try!(SpinLockReadGuard::new(self, &self.data)))
            }
            return Err(TryLockError::WouldBlock)
        }
    }

    pub fn try_write(&self) -> TryLockResult<SpinLockWriteGuard<T>> {
        if self.count.compare_and_swap(0, 1, Ordering::Acquire) == 0 {
            Ok(try!(SpinLockWriteGuard::new(self, &self.data)))
        } else {
            Err(TryLockError::WouldBlock)
        }
    }

    #[inline]
    pub fn is_poisonned(&self) -> bool {
        self.poison.get()
    }

}

impl<T> SpinLock<T> {

    #[inline]
    fn read_unlock(&self) {
        self.count.fetch_sub(1, Ordering::Release);

        while self.count.load(Ordering::Relaxed) == SHARED {
            if self.count.compare_and_swap(SHARED, 0, Ordering::Relaxed) == SHARED {
                break
            }
        }
    }

    fn write_unlock(&self) {
        self.count.fetch_sub(1, Ordering::Release);
    }
}

#[must_use]
pub struct SpinLockReadGuard<'a, T : 'a> {
    lock : &'a SpinLock<T>,
    data : &'a UnsafeCell<T>,
}

impl<'a, T> !Send for SpinLockReadGuard<'a, T> {}

impl<'a, T> SpinLockReadGuard<'a, T> {
    fn new(lock : &'a SpinLock<T>, data : &'a UnsafeCell<T>) -> LockResult<SpinLockReadGuard<'a, T>> {
        map_result(lock.poison.borrow(), |_| {
            SpinLockReadGuard {
                lock : lock,
                data : data,
            }
        })
    }
}

impl<'a, T> Deref for SpinLockReadGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &T { unsafe { &*self.data.get() } }
}

#[unsafe_destructor]
impl<'a, T> Drop for SpinLockReadGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.read_unlock();
    }
}

#[must_use]
pub struct SpinLockWriteGuard<'a, T : 'a> {
    lock : &'a SpinLock<T>,
    data : &'a UnsafeCell<T>,
    poison :  Guard,
}

impl<'a, T> !Send for SpinLockWriteGuard<'a, T> {}

impl<'a, T> SpinLockWriteGuard<'a, T> {
    fn new(lock : &'a SpinLock<T>, data : &'a UnsafeCell<T>) -> LockResult<SpinLockWriteGuard<'a, T>> {
        map_result(lock.poison.borrow(), |guard| {
            SpinLockWriteGuard {
                lock : lock,
                data : data,
                poison : guard,
            }
        })
    }
}

impl<'a, T> Deref for SpinLockWriteGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &T { unsafe { &*self.data.get() } }
}

impl<'a, T> DerefMut for SpinLockWriteGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.data.get() } }
}

#[unsafe_destructor]
impl<'a, T> Drop for SpinLockWriteGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.poison.done(&self.poison);
        self.lock.write_unlock();
    }
}

#[cfg(test)]
mod tests;
<commit_msg>Use Relaxed ordering to release a read lock<commit_after>#![feature(optin_builtin_traits)]
#![feature(unsafe_destructor)]
#![feature(core)]
#![feature(std_misc)]
#![feature(test)]
extern crate time;

use std::cell::UnsafeCell;
pub use std::sync::{
    TryLockError,
    LockResult,
    TryLockResult,
    PoisonError,
};
use std::thread;
use std::ops::{Deref, DerefMut};

use std::sync::atomic::{Ordering, AtomicUsize};

// Panic guards, taken from std::sync

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

struct Flag { failed: UnsafeCell<bool> }
unsafe impl Send for Flag {}
unsafe impl Sync for Flag {}

impl Flag {
    #[inline]
    fn new() -> Flag {
        Flag { failed : UnsafeCell { value : false }}
    }

    #[inline]
    fn borrow(&self) -> LockResult<Guard> {
        let ret = Guard { panicking: thread::panicking() };
        if unsafe { *self.failed.get() } {
            Err(PoisonError::new(ret))
        } else {
            Ok(ret)
        }
    }

    #[inline]
    fn done(&self, guard: &Guard) {
        if !guard.panicking && thread::panicking() {
            unsafe { *self.failed.get() = true; }
        }
    }

    #[inline]
    fn get(&self) -> bool {
        unsafe { *self.failed.get() }
    }
}

struct Guard {
    panicking: bool,
}

fn map_result<T, U, F>(result: LockResult<T>, f: F)
                           -> LockResult<U>
                           where F: FnOnce(T) -> U {
    match result {
        Ok(t) => Ok(f(t)),
        Err(err) => Err(PoisonError::new(f(err.into_inner())))
    }
}

// SpinLock implementation
//
const SHARED : usize = 0x80000000;
const EXCLWAIT : usize = 0x00100000;
const MAXWAIT : u64 = 2_000_000_000;
const SPIN_COUNT : u32 = 0xFF;

pub struct SpinLock<T> {
    count : AtomicUsize,
    data : UnsafeCell<T>,
    poison : Flag,
}

unsafe impl<T: Send + Sync> Send for SpinLock<T> {}
unsafe impl<T: Send + Sync> Sync for SpinLock<T> {}

impl<T : Send + Sync> SpinLock<T> {
    pub fn new(data : T) -> SpinLock<T> {
        SpinLock {
            count : AtomicUsize::new(0),
            data : UnsafeCell::new(data),
            poison : Flag::new(),
        }
    }

    #[inline]
    pub fn write(&self) -> LockResult<SpinLockWriteGuard<T>> {
        if self.count.fetch_add(1, Ordering::Acquire) != 0 {
            return self.write_contested()
        }
        SpinLockWriteGuard::new(self, &self.data)
    }

    pub fn write_contested(&self) -> LockResult<SpinLockWriteGuard<T>> {
        let mut base_time : u64 = 0;
        let mut time : u64;
        let mut i = 0u32;

        self.count.fetch_add(EXCLWAIT - 1, Ordering::Relaxed);
        self.count.fetch_and(! SHARED, Ordering::Relaxed);
        loop {
            // Use relaxed ordering to avoid trashing the cache coherency handling for free
            let count = self.count.load(Ordering::Relaxed);
            if count & (EXCLWAIT - 1) == 0 &&
                self.count.compare_and_swap(count, count - EXCLWAIT + 1, Ordering::Acquire) == count {
                    break;
                }

            if i >= SPIN_COUNT  {
                time = time::precise_time_ns();
                if base_time == 0 {
                    base_time = time;
                } else if time - base_time > MAXWAIT {
                    // XXX we converted the SHARED locks to exclusive ones, i'm not sure how bad
                    // actually is
                    self.count.fetch_sub(EXCLWAIT, Ordering::Relaxed);
                    panic!("Spinning on a spin lock for too long")
                }
                thread::yield_now();
            } else {
                i += 1;
            }
        }
        SpinLockWriteGuard::new(self, &self.data)
    }

    #[inline]
    pub fn read(&self) -> LockResult<SpinLockReadGuard<T>> {
        if self.count.compare_and_swap(0, SHARED | 1, Ordering::Acquire) != 0 {
            return self.read_contested()
        }
        SpinLockReadGuard::new(self, &self.data)
    }

    pub fn read_contested(&self) -> LockResult<SpinLockReadGuard<T>> {
        let mut i = 0u32;
        let mut base_time : u64 = 0;
        let mut time : u64;
        loop {
            let count = self.count.load(Ordering::Relaxed);
            if count == 0 {
                if self.count.compare_and_swap(0, SHARED | 1, Ordering::Acquire) == 0 {
                    break
                }
            } else if count & SHARED == SHARED {
                if self.count.compare_and_swap(count, count + 1, Ordering::Acquire) == count {
                    break
                }
            }
            if i >= SPIN_COUNT  {
                time = time::precise_time_ns();
                if base_time == 0 {
                    base_time = time;
                } else if time - base_time > MAXWAIT {
                    panic!("Spinning on a spin lock for too long")
                }
                thread::yield_now();
            } else {
                i+=1;
            }
        }
        SpinLockReadGuard::new(self, &self.data)
    }

    #[inline]
    pub fn try_read(&self) -> TryLockResult<SpinLockReadGuard<T>> {
        if self.count.compare_and_swap(0, 1 | SHARED, Ordering::Acquire) == 0 {
            return Ok(try!(SpinLockReadGuard::new(self, &self.data)))
        } else {
            let count = self.count.load(Ordering::Relaxed);
            if count & SHARED == SHARED &&
                self.count.compare_and_swap(count, count + 1, Ordering::Acquire) == count {
                return Ok(try!(SpinLockReadGuard::new(self, &self.data)))
            }
            return Err(TryLockError::WouldBlock)
        }
    }

    pub fn try_write(&self) -> TryLockResult<SpinLockWriteGuard<T>> {
        if self.count.compare_and_swap(0, 1, Ordering::Acquire) == 0 {
            Ok(try!(SpinLockWriteGuard::new(self, &self.data)))
        } else {
            Err(TryLockError::WouldBlock)
        }
    }

}

impl<T> SpinLock<T> {

    #[inline]
    fn read_unlock(&self) {
        self.count.fetch_sub(1, Ordering::Relaxed);

        while self.count.load(Ordering::Relaxed) == SHARED {
            if self.count.compare_and_swap(SHARED, 0, Ordering::Relaxed) == SHARED {
                break
            }
        }
    }

    fn write_unlock(&self) {
        self.count.fetch_sub(1, Ordering::Release);
    }

    #[inline]
    pub fn is_poisonned(&self) -> bool {
        self.poison.get()
    }
}

#[must_use]
pub struct SpinLockReadGuard<'a, T : 'a> {
    lock : &'a SpinLock<T>,
    data : &'a UnsafeCell<T>,
}

impl<'a, T> !Send for SpinLockReadGuard<'a, T> {}

impl<'a, T> SpinLockReadGuard<'a, T> {
    fn new(lock : &'a SpinLock<T>, data : &'a UnsafeCell<T>) -> LockResult<SpinLockReadGuard<'a, T>> {
        map_result(lock.poison.borrow(), |_| {
            SpinLockReadGuard {
                lock : lock,
                data : data,
            }
        })
    }
}

impl<'a, T> Deref for SpinLockReadGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &T { unsafe { &*self.data.get() } }
}

#[unsafe_destructor]
impl<'a, T> Drop for SpinLockReadGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.read_unlock();
    }
}

#[must_use]
pub struct SpinLockWriteGuard<'a, T : 'a> {
    lock : &'a SpinLock<T>,
    data : &'a UnsafeCell<T>,
    poison :  Guard,
}

impl<'a, T> !Send for SpinLockWriteGuard<'a, T> {}

impl<'a, T> SpinLockWriteGuard<'a, T> {
    fn new(lock : &'a SpinLock<T>, data : &'a UnsafeCell<T>) -> LockResult<SpinLockWriteGuard<'a, T>> {
        map_result(lock.poison.borrow(), |guard| {
            SpinLockWriteGuard {
                lock : lock,
                data : data,
                poison : guard,
            }
        })
    }
}

impl<'a, T> Deref for SpinLockWriteGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &T { unsafe { &*self.data.get() } }
}

impl<'a, T> DerefMut for SpinLockWriteGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.data.get() } }
}

#[unsafe_destructor]
impl<'a, T> Drop for SpinLockWriteGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.poison.done(&self.poison);
        self.lock.write_unlock();
    }
}

#[cfg(test)]
mod tests;
<|endoftext|>


<commit_before>#[test]
fn it_works() {
}
<commit_msg>Add some inline doc<commit_after>//! Bit utilities for Rust
//!
//! This will eventually expand to include Nibbles and packing them into bytes and integers.

<|endoftext|>