<commit_before>//! Implements `StdHashTable` and `StdHashTableEntry`.

use libc;
use libc::c_void;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::slice;
use std::isize;
use std::cell::Cell;
use std::cmp::{min, max};
use std::mem;
use value::*;
use depth::*;
use hash_table::*;
use moves::MoveDigest;

/// Implements the `HashTableEntry` trait.
#[derive(Copy, Clone, Debug)]
pub struct StdHashTableEntry {
    value: Value,

    // The transposition table maintains a generation number for each
    // entry, which is used to implement an efficient replacement
    // strategy. This field stores the entry's generation (the highest
    // 6 bits) and the bound type (the lowest 2 bits).
    gen_bound: u8,

    depth: Depth,
    move_digest: MoveDigest,
    static_eval: Value,
}

impl HashTableEntry for StdHashTableEntry {
    #[inline]
    fn new(value: Value,
           bound: BoundType,
           depth: Depth,
           move_digest: MoveDigest)
           -> StdHashTableEntry {
        Self::with_static_eval(value, bound, depth, move_digest, VALUE_UNKNOWN)
    }

    #[inline]
    fn with_static_eval(value: Value,
                        bound: BoundType,
                        depth: Depth,
                        move_digest: MoveDigest,
                        static_eval: Value)
                        -> StdHashTableEntry {
        debug_assert!(value != VALUE_UNKNOWN);
        debug_assert!(bound <= 0b11);
        debug_assert!(DEPTH_MIN <= depth && depth <= DEPTH_MAX);
        StdHashTableEntry {
            value: value,
            gen_bound: bound,
            depth: depth,
            move_digest: move_digest,
            static_eval: static_eval,
        }
    }

    #[inline]
    fn value(&self) -> Value {
        self.value
    }

    #[inline]
    fn bound(&self) -> BoundType {
        self.gen_bound & 0b11
    }

    #[inline]
    fn depth(&self) -> Depth {
        self.depth
    }

    #[inline]
    fn move_digest(&self) -> MoveDigest {
        self.move_digest
    }

    /// Returns the `static_eval` passed to the constructor.
    #[inline]
    fn static_eval(&self) -> Value {
        self.static_eval
    }
}

impl StdHashTableEntry {
    /// Returns the contained data as one `u64` value.
    #[inline]
    fn as_u64(&self) -> u64 {
        unsafe { mem::transmute(*self) }
    }
}

/// Implements the `HashTable` trait.
pub struct StdHashTable {
    /// The current generation number. The lowest 2 bits will always
    /// be zeros.
    generation: Cell<u8>,

    /// The number of buckets in the table.
    bucket_count: usize,

    /// The raw pointer obtained from `libc::calloc`. It will be
    /// passed to `libc::free` before the transposition table is
    /// dropped.
    alloc_ptr: AtomicPtr<c_void>,

    /// The transposition table consists of a vector of buckets. Each
    /// buckets stores 4 records.
    table_ptr: AtomicPtr<Bucket>,
}

impl HashTable for StdHashTable {
    type Entry = StdHashTableEntry;

    fn new(size_mb: Option<usize>) -> StdHashTable {
        let size_mb = size_mb.unwrap_or(16);
        let bucket_size = mem::size_of::<Bucket>();
        let bucket_count = {
            let n = max(1, ((size_mb * 1024 * 1024) / bucket_size) as u64);
            1 << (63 - n.leading_zeros())
        };
        let alloc_ptr;
        let table_ptr = unsafe {
            // Make sure that the first bucket is optimally
            // aligned. This may improve performance when cache line's
            // size is divisible to bucket's size or vice versa.
            alloc_ptr = libc::calloc(bucket_count + 1, bucket_size);
            let mut addr = mem::transmute::<*mut c_void, usize>(alloc_ptr);
            addr += bucket_size;
            addr &= !(bucket_size - 1);
            mem::transmute::<usize, *mut Bucket>(addr)
        };

        StdHashTable {
            generation: Cell::new(0),
            bucket_count: bucket_count,
            alloc_ptr: AtomicPtr::new(alloc_ptr),
            table_ptr: AtomicPtr::new(table_ptr),
        }
    }

    fn new_search(&self) {
        const N: usize = 128;

        loop {
            // Increment `self.generation` (with wrapping).
            self.generation.set(self.generation.get().wrapping_add(0b100));
            debug_assert_eq!(self.generation.get() & 0b11, 0);

            // Count how many staled records from this generation
            // there are among the first `N` buckets.
            let mut staled = 0;
            let mut bucket_iter = self.table().iter();
            for _ in 0..min(N, self.bucket_count) {
                for record in bucket_iter.next().unwrap() {
                    if record.key != 0 && record.generation() == self.generation.get() {
                        staled += 1;
                    }
                }
            }

            if staled < N {
                // Note that we will continue to increment
                // `self.generation` if the staled records from this
                // generation are too many. (This may happen if a very
                // long search was executed long time ago.)
                break;
            }
        }
    }

    fn store(&self, key: u64, mut data: Self::Entry) {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        // Set entry's generation.
        data.gen_bound = self.generation.get() | data.bound();

        // Choose a slot to which to write the data. (Each bucket can
        // have several slots.)
        let mut bucket = self.bucket_mut(key);
        let mut replace_index = 0;
        let mut replace_score = isize::MAX;
        for (i, record) in bucket.iter_mut().enumerate() {
            // Check if this is an empty slot, or an old record for
            // the same key. If this this is the case we will use this
            // slot for the new record.
            if record.key == 0 || record.key ^ record.data.as_u64() == key {
                if data.move_digest == MoveDigest::invalid() {
                    data.move_digest = record.data.move_digest; // Preserve any existing move.
                }
                replace_index = i;
                break;
            }

            // Calculate the score for this record. If we can not find
            // an empty slot or an old record, the replaced record
            // will be the record with the lowest score.
            let record_score = self.calc_score(record);
            if record_score < replace_score {
                replace_index = i;
                replace_score = record_score;
            }
        }

        // Write the data to the chosen slot.
        bucket[replace_index] = Record {
            key: key ^ data.as_u64(),
            data: data,
        };
    }

    #[inline]
    fn probe(&self, key: u64) -> Option<Self::Entry> {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        let bucket = self.bucket_mut(key);
        for record in bucket.iter_mut() {
            if record.key ^ record.data.as_u64() == key {
                // If `key` and `data` were written simultaneously by
                // different search instances with different keys,
                // this will yield in a mismatch of the above
                // comparison (except for the rare and inherent key
                // collisions).
                record.set_generation(self.generation.get());
                return Some(record.data);
            }
        }
        None
    }

    fn clear(&self) {
        for bucket in self.table_mut() {
            for record in bucket.iter_mut() {
                *record = Default::default();
            }
        }
        self.generation.set(0);
    }
}

impl StdHashTable {
    /// A helper method for `store`. It implements the record
    /// replacement strategy.
    #[inline]
    fn calc_score(&self, record: &Record) -> isize {
        // Here we try to return higher values for the records that
        // are move likely to save CPU work in the future:

        // Positions from the current generation are always scored
        // higher than positions from older generations.
        (if record.generation() == self.generation.get() {
            DEPTH_MAX as isize + 2
        } else {
            0
        })

        // Positions with higher search depths are scored higher.
        + record.data.depth() as isize

        // Positions with exact evaluations are given slight advantage.
        + (if record.data.bound() == BOUND_EXACT {
            1
        } else {
            0
        })
    }

    /// A helper method for `probe` and `store`. It returns the bucket
    /// for a given key.
    #[inline]
    fn bucket_mut(&self, key: u64) -> &mut Bucket {
        unsafe {
            let index = (key & (self.bucket_count - 1) as u64) as usize;
            self.table_mut().get_unchecked_mut(index)
        }
    }

    #[inline]
    fn table(&self) -> &[Bucket] {
        unsafe { slice::from_raw_parts(self.table_ptr.load(Ordering::Relaxed), self.bucket_count) }
    }

    #[inline]
    fn table_mut(&self) -> &mut [Bucket] {
        unsafe {
            slice::from_raw_parts_mut(self.table_ptr.load(Ordering::Relaxed), self.bucket_count)
        }
    }
}

impl Drop for StdHashTable {
    fn drop(&mut self) {
        unsafe {
            libc::free(self.alloc_ptr.load(Ordering::Relaxed));
        }
    }
}

unsafe impl Sync for StdHashTable {}

/// Represents a record in the transposition table.
///
/// It consists of 16 bytes, and is laid out the following way:
///
/// * key         64 bit
/// * move_digest 16 bit
/// * value       16 bit
/// * eval value  16 bit
/// * depth        8 bit
/// * generation   6 bit
/// * bound type   2 bit
#[derive(Copy, Clone)]
struct Record {
    key: u64,
    data: StdHashTableEntry,
}

impl Default for Record {
    fn default() -> Record {
        Record {
            key: 0,
            data: unsafe { mem::transmute(0u64) },
        }
    }
}

impl Record {
    #[inline]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }

    #[inline]
    fn set_generation(&mut self, generation: u8) {
        debug_assert_eq!(generation & 0b11, 0);

        // Since the `key` is saved XOR-ed with the data, when we
        // change the data, we have to change the stored `key` as
        // well.
        let old_data_as_u64 = self.data.as_u64();
        self.data.gen_bound = generation | self.data.bound();
        self.key ^= old_data_as_u64 ^ self.data.as_u64();
    }
}

type Bucket = [Record; 4];

#[cfg(test)]
mod tests {
    use super::*;
    use super::Bucket;
    use std;
    use depth::*;
    use hash_table::*;
    use moves::*;

    #[test]
    fn bucket_size() {
        let n = std::mem::size_of::<Bucket>();
        assert!(n % 64 == 0 || 64 % n == 0);
    }

    #[test]
    fn store_and_probe() {
        let tt = StdHashTable::new(None);
        assert!(tt.probe(1).is_none());
        let data = StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid());
        assert_eq!(data.depth(), 50);
        assert_eq!(data.move_digest(), MoveDigest::invalid());
        tt.store(1, data);
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        tt.store(1, StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid()));
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(1).unwrap().move_digest(), MoveDigest::invalid());
        let digest = MoveDigest::invalid();
        for i in 2..50 {
            tt.store(i, StdHashTableEntry::new(i as i16, 0, i as Depth, digest));
        }
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(49).unwrap().depth(), 49);
        assert_eq!(tt.probe(48).unwrap().depth(), 48);
        assert_eq!(tt.probe(47).unwrap().depth(), 47);
        tt.clear();
        assert!(tt.probe(1).is_none());
        tt.store(1, data);
        tt.new_search();
        tt.probe(1);
        assert!(tt.probe(1).is_some());
    }

    #[test]
    fn new_search() {
        let tt = StdHashTable::new(None);
        assert_eq!(tt.generation.get(), 0 << 2);
        tt.new_search();
        assert_eq!(tt.generation.get(), 1 << 2);
        for _ in 0..64 {
            tt.new_search();
        }
        assert_eq!(tt.generation.get(), 1 << 2);
    }
}
<commit_msg>implement Cluster struct<commit_after>//! Implements `StdHashTable` and `StdHashTableEntry`.

use libc;
use libc::c_void;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::slice;
use std::ops::{Deref, DerefMut};
use std::isize;
use std::cell::Cell;
use std::cmp::{min, max};
use std::mem;
use value::*;
use depth::*;
use hash_table::*;
use moves::MoveDigest;

/// Implements the `HashTableEntry` trait.
#[derive(Copy, Clone, Debug)]
pub struct StdHashTableEntry {
    value: Value,

    // The transposition table maintains a generation number for each
    // entry, which is used to implement an efficient replacement
    // strategy. This field stores the entry's generation (the highest
    // 6 bits) and the bound type (the lowest 2 bits).
    gen_bound: u8,

    depth: Depth,
    move_digest: MoveDigest,
    static_eval: Value,
}

impl HashTableEntry for StdHashTableEntry {
    #[inline]
    fn new(value: Value,
           bound: BoundType,
           depth: Depth,
           move_digest: MoveDigest)
           -> StdHashTableEntry {
        Self::with_static_eval(value, bound, depth, move_digest, VALUE_UNKNOWN)
    }

    #[inline]
    fn with_static_eval(value: Value,
                        bound: BoundType,
                        depth: Depth,
                        move_digest: MoveDigest,
                        static_eval: Value)
                        -> StdHashTableEntry {
        debug_assert!(value != VALUE_UNKNOWN);
        debug_assert!(bound <= 0b11);
        debug_assert!(DEPTH_MIN <= depth && depth <= DEPTH_MAX);
        StdHashTableEntry {
            value: value,
            gen_bound: bound,
            depth: depth,
            move_digest: move_digest,
            static_eval: static_eval,
        }
    }

    #[inline]
    fn value(&self) -> Value {
        self.value
    }

    #[inline]
    fn bound(&self) -> BoundType {
        self.gen_bound & 0b11
    }

    #[inline]
    fn depth(&self) -> Depth {
        self.depth
    }

    #[inline]
    fn move_digest(&self) -> MoveDigest {
        self.move_digest
    }

    /// Returns the `static_eval` passed to the constructor.
    #[inline]
    fn static_eval(&self) -> Value {
        self.static_eval
    }
}

impl StdHashTableEntry {
    /// Returns the contained data as one `u64` value.
    #[inline]
    fn as_u64(&self) -> u64 {
        unsafe { mem::transmute(*self) }
    }
}

/// Implements the `HashTable` trait.
pub struct StdHashTable {
    /// The current generation number. The lowest 2 bits will always
    /// be zeros.
    generation: Cell<u8>,

    /// The number of buckets in the table.
    bucket_count: usize,

    /// The raw pointer obtained from `libc::calloc`. It will be
    /// passed to `libc::free` before the transposition table is
    /// dropped.
    alloc_ptr: AtomicPtr<c_void>,

    /// The transposition table consists of a vector of buckets. Each
    /// buckets stores 4 records.
    table_ptr: AtomicPtr<Bucket>,
}

impl HashTable for StdHashTable {
    type Entry = StdHashTableEntry;

    fn new(size_mb: Option<usize>) -> StdHashTable {
        let size_mb = size_mb.unwrap_or(16);
        let bucket_size = mem::size_of::<Bucket>();
        let bucket_count = {
            let n = max(1, ((size_mb * 1024 * 1024) / bucket_size) as u64);
            1 << (63 - n.leading_zeros())
        };
        let alloc_ptr;
        let table_ptr = unsafe {
            // Make sure that the first bucket is optimally
            // aligned. This may improve performance when cache line's
            // size is divisible to bucket's size or vice versa.
            alloc_ptr = libc::calloc(bucket_count + 1, bucket_size);
            let mut addr = mem::transmute::<*mut c_void, usize>(alloc_ptr);
            addr += bucket_size;
            addr &= !(bucket_size - 1);
            mem::transmute::<usize, *mut Bucket>(addr)
        };

        StdHashTable {
            generation: Cell::new(0),
            bucket_count: bucket_count,
            alloc_ptr: AtomicPtr::new(alloc_ptr),
            table_ptr: AtomicPtr::new(table_ptr),
        }
    }

    fn new_search(&self) {
        const N: usize = 128;

        loop {
            // Increment `self.generation` (with wrapping).
            self.generation.set(self.generation.get().wrapping_add(0b100));
            debug_assert_eq!(self.generation.get() & 0b11, 0);

            // Count how many staled records from this generation
            // there are among the first `N` buckets.
            let mut staled = 0;
            let mut bucket_iter = self.table().iter();
            for _ in 0..min(N, self.bucket_count) {
                for record in bucket_iter.next().unwrap() {
                    if record.key != 0 && record.generation() == self.generation.get() {
                        staled += 1;
                    }
                }
            }

            if staled < N {
                // Note that we will continue to increment
                // `self.generation` if the staled records from this
                // generation are too many. (This may happen if a very
                // long search was executed long time ago.)
                break;
            }
        }
    }

    fn store(&self, key: u64, mut data: Self::Entry) {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        // Set entry's generation.
        data.gen_bound = self.generation.get() | data.bound();

        // Choose a slot to which to write the data. (Each bucket can
        // have several slots.)
        let mut bucket = self.bucket_mut(key);
        let mut replace_index = 0;
        let mut replace_score = isize::MAX;
        for (i, record) in bucket.iter_mut().enumerate() {
            // Check if this is an empty slot, or an old record for
            // the same key. If this this is the case we will use this
            // slot for the new record.
            if record.key == 0 || record.key ^ record.data.as_u64() == key {
                if data.move_digest == MoveDigest::invalid() {
                    data.move_digest = record.data.move_digest; // Preserve any existing move.
                }
                replace_index = i;
                break;
            }

            // Calculate the score for this record. If we can not find
            // an empty slot or an old record, the replaced record
            // will be the record with the lowest score.
            let record_score = self.calc_score(record);
            if record_score < replace_score {
                replace_index = i;
                replace_score = record_score;
            }
        }

        // Write the data to the chosen slot.
        bucket[replace_index] = Record {
            key: key ^ data.as_u64(),
            data: data,
        };
    }

    #[inline]
    fn probe(&self, key: u64) -> Option<Self::Entry> {
        // `store` and `probe` jointly implement a clever lock-less
        // hashing strategy. Rather than storing two disjoint items,
        // the key is stored XOR-ed with data, while data is stored
        // additionally as usual.

        let bucket = self.bucket_mut(key);
        for record in bucket.iter_mut() {
            if record.key ^ record.data.as_u64() == key {
                // If `key` and `data` were written simultaneously by
                // different search instances with different keys,
                // this will yield in a mismatch of the above
                // comparison (except for the rare and inherent key
                // collisions).
                record.set_generation(self.generation.get());
                return Some(record.data);
            }
        }
        None
    }

    fn clear(&self) {
        for bucket in self.table_mut() {
            for record in bucket.iter_mut() {
                *record = Default::default();
            }
        }
        self.generation.set(0);
    }
}

impl StdHashTable {
    /// A helper method for `store`. It implements the record
    /// replacement strategy.
    #[inline]
    fn calc_score(&self, record: &Record) -> isize {
        // Here we try to return higher values for the records that
        // are move likely to save CPU work in the future:

        // Positions from the current generation are always scored
        // higher than positions from older generations.
        (if record.generation() == self.generation.get() {
            DEPTH_MAX as isize + 2
        } else {
            0
        })

        // Positions with higher search depths are scored higher.
        + record.data.depth() as isize

        // Positions with exact evaluations are given slight advantage.
        + (if record.data.bound() == BOUND_EXACT {
            1
        } else {
            0
        })
    }

    /// A helper method for `probe` and `store`. It returns the bucket
    /// for a given key.
    #[inline]
    fn bucket_mut(&self, key: u64) -> &mut Bucket {
        unsafe {
            let index = (key & (self.bucket_count - 1) as u64) as usize;
            self.table_mut().get_unchecked_mut(index)
        }
    }

    #[inline]
    fn table(&self) -> &[Bucket] {
        unsafe { slice::from_raw_parts(self.table_ptr.load(Ordering::Relaxed), self.bucket_count) }
    }

    #[inline]
    fn table_mut(&self) -> &mut [Bucket] {
        unsafe {
            slice::from_raw_parts_mut(self.table_ptr.load(Ordering::Relaxed), self.bucket_count)
        }
    }
}

impl Drop for StdHashTable {
    fn drop(&mut self) {
        unsafe {
            libc::free(self.alloc_ptr.load(Ordering::Relaxed));
        }
    }
}

unsafe impl Sync for StdHashTable {}

/// Represents a record in the transposition table.
///
/// It consists of 16 bytes, and is laid out the following way:
///
/// * key         64 bit
/// * move_digest 16 bit
/// * value       16 bit
/// * eval value  16 bit
/// * depth        8 bit
/// * generation   6 bit
/// * bound type   2 bit
#[derive(Copy, Clone)]
struct Record {
    key: u64,
    data: StdHashTableEntry,
}

impl Default for Record {
    fn default() -> Record {
        Record {
            key: 0,
            data: unsafe { mem::transmute(0u64) },
        }
    }
}

impl Record {
    #[inline]
    fn generation(&self) -> u8 {
        self.data.gen_bound & 0b11111100
    }

    #[inline]
    fn set_generation(&mut self, generation: u8) {
        debug_assert_eq!(generation & 0b11, 0);

        // Since the `key` is saved XOR-ed with the data, when we
        // change the data, we have to change the stored `key` as
        // well.
        let old_data_as_u64 = self.data.as_u64();
        self.data.gen_bound = generation | self.data.bound();
        self.key ^= old_data_as_u64 ^ self.data.as_u64();
    }
}

type Bucket = [Record; 4];

struct Cluster<'a, T: Sized + 'a> {
    info: *mut u32,
    entries: &'a mut [T],
}

const BUCKET_INFO_SHIFTS: [u8; 6] = [0, 5, 10, 15, 20, 25];

impl<'a, T: Sized + 'a> Cluster<'a, T> {
    #[inline]
    pub unsafe fn new(p: *mut c_void) -> Cluster<'a, T> {
        Cluster {
            info: p.offset(60) as *mut u32,
            entries: slice::from_raw_parts_mut(p as *mut T, 60 / mem::size_of::<T>()),
        }
    }

    #[inline]
    pub fn get_generation(&self, index: usize) -> u32 {
        let info = unsafe { self.info.as_ref().unwrap() };
        *info >> BUCKET_INFO_SHIFTS[index] & 31
    }

    #[inline]
    pub fn set_generation(&self, index: usize, generation: u32) {
        debug_assert!(generation <= 31);
        let info = unsafe { self.info.as_mut().unwrap() };
        *info &= [!(31 << 0),
                  !(31 << 5),
                  !(31 << 10),
                  !(31 << 15),
                  !(31 << 20),
                  !(31 << 25)][index];
        *info |= generation << BUCKET_INFO_SHIFTS[index];
    }
}

impl<'a, T: Sized + 'a> Deref for Cluster<'a, T> {
    type Target = &'a mut [T];

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.entries
    }
}

impl<'a, T: Sized + 'a> DerefMut for Cluster<'a, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.entries
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::Bucket;
    use std;
    use depth::*;
    use hash_table::*;
    use moves::*;

    #[test]
    fn bucket() {
        use libc;
        use moves::MoveDigest;
        use super::Cluster;
        unsafe {
            let p = libc::malloc(64);
            let mut b = Cluster::<(u32, StdHashTableEntry)>::new(p);
            b[0] = (0, StdHashTableEntry::new(0, BOUND_NONE, 10, MoveDigest::invalid()));
            b.set_generation(0, 12);
            b.set_generation(1, 13);
            assert_eq!(b[0].1.depth, 10);
            assert_eq!(b.get_generation(0), 12);
            assert_eq!(b.get_generation(1), 13);
            assert_eq!(b.iter().count(), 5);
            libc::free(p);
        }
    }

    #[test]
    fn bucket_size() {
        let n = std::mem::size_of::<Bucket>();
        assert!(n % 64 == 0 || 64 % n == 0);
    }

    #[test]
    fn store_and_probe() {
        let tt = StdHashTable::new(None);
        assert!(tt.probe(1).is_none());
        let data = StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid());
        assert_eq!(data.depth(), 50);
        assert_eq!(data.move_digest(), MoveDigest::invalid());
        tt.store(1, data);
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        tt.store(1, StdHashTableEntry::new(0, 0, 50, MoveDigest::invalid()));
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(1).unwrap().move_digest(), MoveDigest::invalid());
        let digest = MoveDigest::invalid();
        for i in 2..50 {
            tt.store(i, StdHashTableEntry::new(i as i16, 0, i as Depth, digest));
        }
        assert_eq!(tt.probe(1).unwrap().depth(), 50);
        assert_eq!(tt.probe(49).unwrap().depth(), 49);
        assert_eq!(tt.probe(48).unwrap().depth(), 48);
        assert_eq!(tt.probe(47).unwrap().depth(), 47);
        tt.clear();
        assert!(tt.probe(1).is_none());
        tt.store(1, data);
        tt.new_search();
        tt.probe(1);
        assert!(tt.probe(1).is_some());
    }

    #[test]
    fn new_search() {
        let tt = StdHashTable::new(None);
        assert_eq!(tt.generation.get(), 0 << 2);
        tt.new_search();
        assert_eq!(tt.generation.get(), 1 << 2);
        for _ in 0..64 {
            tt.new_search();
        }
        assert_eq!(tt.generation.get(), 1 << 2);
    }
}
<|endoftext|>


<commit_before>//! The Rust Prelude.
//!
//! Rust comes with a variety of things in its standard library. However, if
//! you had to manually import every single thing that you used, it would be
//! very verbose. But importing a lot of things that a program never uses isn't
//! good either. A balance needs to be struck.
//!
//! The *prelude* is the list of things that Rust automatically imports into
//! every Rust program. It's kept as small as possible, and is focused on
//! things, particularly traits, which are used in almost every single Rust
//! program.
//!
//! # Other preludes
//!
//! Preludes can be seen as a pattern to make using multiple types more
//! convenient. As such, you'll find other preludes in the standard library,
//! such as [`std::io::prelude`]. Various libraries in the Rust ecosystem may
//! also define their own preludes.
//!
//! [`std::io::prelude`]: crate::io::prelude
//!
//! The difference between 'the prelude' and these other preludes is that they
//! are not automatically `use`'d, and must be imported manually. This is still
//! easier than importing all of their constituent components.
//!
//! # Prelude contents
//!
//! The current version of the prelude (version 1) lives in
//! [`std::prelude::v1`], and re-exports the following.
//!
//! * [`std::marker`]::{[`Copy`], [`Send`], [`Sized`], [`Sync`], [`Unpin`]}. The
//!   marker traits indicate fundamental properties of types.
//! * [`std::ops`]::{[`Drop`], [`Fn`], [`FnMut`], [`FnOnce`]}. Various
//!   operations for both destructors and overloading `()`.
//! * [`std::mem`]::[`drop`][`mem::drop`], a convenience function for explicitly
//!   dropping a value.
//! * [`std::boxed`]::[`Box`], a way to allocate values on the heap.
//! * [`std::borrow`]::[`ToOwned`], The conversion trait that defines
//!   [`to_owned`], the generic method for creating an owned type from a
//!   borrowed type.
//! * [`std::clone`]::[`Clone`], the ubiquitous trait that defines
//!   [`clone`][`Clone::clone`], the method for producing a copy of a value.
//! * [`std::cmp`]::{[`PartialEq`], [`PartialOrd`], [`Eq`], [`Ord`] }. The
//!   comparison traits, which implement the comparison operators and are often
//!   seen in trait bounds.
//! * [`std::convert`]::{[`AsRef`], [`AsMut`], [`Into`], [`From`]}. Generic
//!   conversions, used by savvy API authors to create overloaded methods.
//! * [`std::default`]::[`Default`], types that have default values.
//! * [`std::iter`]::{[`Iterator`], [`Extend`], [`IntoIterator`],
//!   [`DoubleEndedIterator`], [`ExactSizeIterator`]}. Iterators of various
//!   kinds.
//! * [`std::option`]::[`Option`]::{[`self`][`Option`], [`Some`], [`None`]}. A
//!   type which expresses the presence or absence of a value. This type is so
//!   commonly used, its variants are also exported.
//! * [`std::result`]::[`Result`]::{[`self`][`Result`], [`Ok`], [`Err`]}. A type
//!   for functions that may succeed or fail. Like [`Option`], its variants are
//!   exported as well.
//! * [`std::string`]::{[`String`], [`ToString`]}, heap allocated strings.
//! * [`std::vec`]::[`Vec`], a growable, heap-allocated
//!   vector.
//!
//! [`mem::drop`]: crate::mem::drop
//! [`std::borrow`]: crate::borrow
//! [`std::boxed`]: crate::boxed
//! [`std::clone`]: crate::clone
//! [`std::cmp`]: crate::cmp
//! [`std::convert`]: crate::convert
//! [`std::default`]: crate::default
//! [`std::iter`]: crate::iter
//! [`std::marker`]: crate::marker
//! [`std::mem`]: crate::mem
//! [`std::ops`]: crate::ops
//! [`std::option`]: crate::option
//! [`std::prelude::v1`]: v1
//! [`std::result`]: crate::result
//! [`std::slice`]: crate::slice
//! [`std::string`]: crate::string
//! [`std::vec`]: mod@crate::vec
//! [`to_owned`]: crate::borrow::ToOwned::to_owned
//! [book-closures]: ../../book/ch13-01-closures.html
//! [book-dtor]: ../../book/ch15-03-drop.html
//! [book-enums]: ../../book/ch06-01-defining-an-enum.html
//! [book-iter]: ../../book/ch13-02-iterators.html

#![stable(feature = "rust1", since = "1.0.0")]

pub mod v1;
<commit_msg>Rollup merge of #80146 - pierwill:pierwill-prelude-mod-docs, r=lcnr<commit_after>//! The Rust Prelude.
//!
//! Rust comes with a variety of things in its standard library. However, if
//! you had to manually import every single thing that you used, it would be
//! very verbose. But importing a lot of things that a program never uses isn't
//! good either. A balance needs to be struck.
//!
//! The *prelude* is the list of things that Rust automatically imports into
//! every Rust program. It's kept as small as possible, and is focused on
//! things, particularly traits, which are used in almost every single Rust
//! program.
//!
//! # Other preludes
//!
//! Preludes can be seen as a pattern to make using multiple types more
//! convenient. As such, you'll find other preludes in the standard library,
//! such as [`std::io::prelude`]. Various libraries in the Rust ecosystem may
//! also define their own preludes.
//!
//! [`std::io::prelude`]: crate::io::prelude
//!
//! The difference between 'the prelude' and these other preludes is that they
//! are not automatically `use`'d, and must be imported manually. This is still
//! easier than importing all of their constituent components.
//!
//! # Prelude contents
//!
//! The current version of the prelude (version 1) lives in
//! [`std::prelude::v1`], and re-exports the following:
//!
//! * [`std::marker`]::{[`Copy`], [`Send`], [`Sized`], [`Sync`], [`Unpin`]},
//!   marker traits that indicate fundamental properties of types.
//! * [`std::ops`]::{[`Drop`], [`Fn`], [`FnMut`], [`FnOnce`]}, various
//!   operations for both destructors and overloading `()`.
//! * [`std::mem`]::[`drop`][`mem::drop`], a convenience function for explicitly
//!   dropping a value.
//! * [`std::boxed`]::[`Box`], a way to allocate values on the heap.
//! * [`std::borrow`]::[`ToOwned`], the conversion trait that defines
//!   [`to_owned`], the generic method for creating an owned type from a
//!   borrowed type.
//! * [`std::clone`]::[`Clone`], the ubiquitous trait that defines
//!   [`clone`][`Clone::clone`], the method for producing a copy of a value.
//! * [`std::cmp`]::{[`PartialEq`], [`PartialOrd`], [`Eq`], [`Ord`] }, the
//!   comparison traits, which implement the comparison operators and are often
//!   seen in trait bounds.
//! * [`std::convert`]::{[`AsRef`], [`AsMut`], [`Into`], [`From`]}, generic
//!   conversions, used by savvy API authors to create overloaded methods.
//! * [`std::default`]::[`Default`], types that have default values.
//! * [`std::iter`]::{[`Iterator`], [`Extend`], [`IntoIterator`]
//!   [`DoubleEndedIterator`], [`ExactSizeIterator`]}, iterators of various
//!   kinds.
//! * [`std::option`]::[`Option`]::{[`self`][`Option`], [`Some`], [`None`]}, a
//!   type which expresses the presence or absence of a value. This type is so
//!   commonly used, its variants are also exported.
//! * [`std::result`]::[`Result`]::{[`self`][`Result`], [`Ok`], [`Err`]}, a type
//!   for functions that may succeed or fail. Like [`Option`], its variants are
//!   exported as well.
//! * [`std::string`]::{[`String`], [`ToString`]}, heap allocated strings.
//! * [`std::vec`]::[`Vec`], a growable, heap-allocated vector.
//!
//! [`mem::drop`]: crate::mem::drop
//! [`std::borrow`]: crate::borrow
//! [`std::boxed`]: crate::boxed
//! [`std::clone`]: crate::clone
//! [`std::cmp`]: crate::cmp
//! [`std::convert`]: crate::convert
//! [`std::default`]: crate::default
//! [`std::iter`]: crate::iter
//! [`std::marker`]: crate::marker
//! [`std::mem`]: crate::mem
//! [`std::ops`]: crate::ops
//! [`std::option`]: crate::option
//! [`std::prelude::v1`]: v1
//! [`std::result`]: crate::result
//! [`std::slice`]: crate::slice
//! [`std::string`]: crate::string
//! [`std::vec`]: mod@crate::vec
//! [`to_owned`]: crate::borrow::ToOwned::to_owned
//! [book-closures]: ../../book/ch13-01-closures.html
//! [book-dtor]: ../../book/ch15-03-drop.html
//! [book-enums]: ../../book/ch06-01-defining-an-enum.html
//! [book-iter]: ../../book/ch13-02-iterators.html

#![stable(feature = "rust1", since = "1.0.0")]

pub mod v1;
<|endoftext|>


<commit_before>extern crate indy_crypto;
extern crate zeroize;

use self::zeroize::Zeroize;

use named_type::NamedType;

#[derive(Derivative)]
#[derivative(Debug)]
#[derive(Serialize, Deserialize, Clone, NamedType)]
pub struct Key {
    pub verkey: String,
    #[cfg(not(test))]
    #[derivative(Debug="ignore")]
    pub signkey: String,
    #[cfg(test)]
    pub signkey: String
}

impl Key {
    pub fn new(verkey: String, signkey: String) -> Key {
        Key {
            verkey,
            signkey
        }
    }
}

//Can't figure out why this macro can't be found in `cargo test` because
//if works for src/utils/crypto/randombytes/sodium
memzeroize!(Key, signkey);

//impl Zeroize for Key {
//    fn zeroize(&mut self) { self.signkey.zeroize(); }
//}
//
//impl Drop for Key {
//    fn drop(&mut self) {
//        self.signkey.zeroize();
//    }
//}

#[derive(Serialize, Deserialize, Debug)]
pub struct KeyInfo {
    pub seed: Option<String>,
    pub crypto_type: Option<String>
}

#[derive(Serialize, Deserialize, Debug, NamedType)]
pub struct KeyMetadata {
    pub value: String
}
<commit_msg>Remove comment<commit_after>extern crate indy_crypto;
extern crate zeroize;

use self::zeroize::Zeroize;

use named_type::NamedType;

#[derive(Derivative)]
#[derivative(Debug)]
#[derive(Serialize, Deserialize, Clone, NamedType)]
pub struct Key {
    pub verkey: String,
    #[cfg(not(test))]
    #[derivative(Debug="ignore")]
    pub signkey: String,
    #[cfg(test)]
    pub signkey: String
}

impl Key {
    pub fn new(verkey: String, signkey: String) -> Key {
        Key {
            verkey,
            signkey
        }
    }
}

memzeroize!(Key, signkey);

#[derive(Serialize, Deserialize, Debug)]
pub struct KeyInfo {
    pub seed: Option<String>,
    pub crypto_type: Option<String>
}

#[derive(Serialize, Deserialize, Debug, NamedType)]
pub struct KeyMetadata {
    pub value: String
}
<|endoftext|>


<commit_before>use std::rc::Rc;
use std::io::{self, Read, Write};
use std::net::Shutdown;

use futures::prelude::*;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_io::io::shutdown;
use bytes::{BytesMut, Bytes, IntoBuf};
use h2::{self, server as h2s, client as h2c};
use tokio_core::net::TcpStream;

#[async]
pub fn copy_from_h2<
    W: AsyncWrite + 'static,
    B: Stream<Item = Bytes, Error = h2::Error> + 'static,
>(
    src: B,
    mut dst: W,
) -> Result<usize, h2::Error> {
    let mut counter = 0;
    #[async]
    for buf in src {
        let mut buf = buf.into_buf();
        let n = poll!(dst.write_buf(&mut buf))?;
        counter += n;
    }
    await!(shutdown(dst))?;
    Ok(counter)
}

pub trait SendData<B: IntoBuf> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error>;
}

impl<B: IntoBuf> SendData<B> for h2c::Stream<B> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error> {
        h2c::Stream::send_data(self, data, end_of_stream)
    }
}
impl<B: IntoBuf> SendData<B> for h2s::Stream<B> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error> {
        h2s::Stream::send_data(self, data, end_of_stream)
    }
}

#[async]
pub fn copy_to_h2<R: AsyncRead + 'static, H: SendData<Bytes> + 'static>(
    mut src: R,
    mut dst: H,
) -> Result<usize, h2::Error> {
    let mut counter = 0;
    let mut buf = BytesMut::with_capacity(1024);
    loop {
        let n = poll!(src.read_buf(&mut buf))?;
        if n == 0 {
            dst.send_data(buf.take().freeze(), true)?;
            break;
        } else {
            dst.send_data(buf.take().freeze(), false)?;
        }
        counter += n;
    }
    Ok(counter)
}

// This is a custom type used to have a custom implementation of the
// `AsyncWrite::shutdown` method which actually calls `TcpStream::shutdown` to
// notify the remote end that we're done writing.
#[derive(Clone)]
pub struct Socket(Rc<TcpStream>);

impl Socket {
    pub fn new(s: TcpStream) -> Socket {
        Socket(Rc::new(s))
    }
}

impl Read for Socket {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (&*self.0).read(buf)
    }
}

impl Write for Socket {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        (&*self.0).write(buf)
    }

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

impl AsyncRead for Socket {}

impl AsyncWrite for Socket {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        try!(self.0.shutdown(Shutdown::Write));
        Ok(().into())
    }
}
<commit_msg>Fix `io::copy_to_h2` only receive 1024 bytes<commit_after>use std::rc::Rc;
use std::io::{self, Read, Write};
use std::net::Shutdown;

use futures::prelude::*;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_io::io::shutdown;
use bytes::{Bytes, IntoBuf};
use h2::{self, server as h2s, client as h2c};
use tokio_core::net::TcpStream;

#[async]
pub fn copy_from_h2<
    W: AsyncWrite + 'static,
    B: Stream<Item = Bytes, Error = h2::Error> + 'static,
>(
    src: B,
    mut dst: W,
) -> Result<usize, h2::Error> {
    let mut counter = 0;
    #[async]
    for buf in src {
        let mut buf = buf.into_buf();
        let n = poll!(dst.write_buf(&mut buf))?;
        counter += n;
    }
    await!(shutdown(dst))?;
    Ok(counter)
}

pub trait SendData<B: IntoBuf> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error>;
}

impl<B: IntoBuf> SendData<B> for h2c::Stream<B> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error> {
        h2c::Stream::send_data(self, data, end_of_stream)
    }
}
impl<B: IntoBuf> SendData<B> for h2s::Stream<B> {
    fn send_data(&mut self, data: B, end_of_stream: bool) -> Result<(), h2::Error> {
        h2s::Stream::send_data(self, data, end_of_stream)
    }
}

#[async]
pub fn copy_to_h2<R: AsyncRead + 'static, H: SendData<Bytes> + 'static>(
    mut src: R,
    mut dst: H,
) -> Result<usize, h2::Error> {
    let mut counter = 0;
    let mut buf = Vec::with_capacity(1024);
    loop {
        let n = poll!(src.read_buf(&mut buf))?;
        if n == 0 {
            dst.send_data(Bytes::from(&buf[..n]), true)?;
            break;
        } else {
            dst.send_data(Bytes::from(&buf[..n]), false)?;
        }
        counter += n;
    }
    Ok(counter)
}

// This is a custom type used to have a custom implementation of the
// `AsyncWrite::shutdown` method which actually calls `TcpStream::shutdown` to
// notify the remote end that we're done writing.
#[derive(Clone)]
pub struct Socket(Rc<TcpStream>);

impl Socket {
    pub fn new(s: TcpStream) -> Socket {
        Socket(Rc::new(s))
    }
}

impl Read for Socket {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (&*self.0).read(buf)
    }
}

impl Write for Socket {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        (&*self.0).write(buf)
    }

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

impl AsyncRead for Socket {}

impl AsyncWrite for Socket {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        try!(self.0.shutdown(Shutdown::Write));
        Ok(().into())
    }
}
<|endoftext|>


<commit_before>use crate::cmp::Ordering;
use crate::time::Duration;

use core::hash::{Hash, Hasher};

pub use self::inner::{Instant, SystemTime, UNIX_EPOCH};
use crate::convert::TryInto;

const NSEC_PER_SEC: u64 = 1_000_000_000;

#[derive(Copy, Clone)]
struct Timespec {
    t: libc::timespec,
}

impl Timespec {
    const fn zero() -> Timespec {
        Timespec { t: libc::timespec { tv_sec: 0, tv_nsec: 0 } }
    }

    fn sub_timespec(&self, other: &Timespec) -> Result<Duration, Duration> {
        if self >= other {
            Ok(if self.t.tv_nsec >= other.t.tv_nsec {
                Duration::new(
                    (self.t.tv_sec - other.t.tv_sec) as u64,
                    (self.t.tv_nsec - other.t.tv_nsec) as u32,
                )
            } else {
                Duration::new(
                    (self.t.tv_sec - 1 - other.t.tv_sec) as u64,
                    self.t.tv_nsec as u32 + (NSEC_PER_SEC as u32) - other.t.tv_nsec as u32,
                )
            })
        } else {
            match other.sub_timespec(self) {
                Ok(d) => Err(d),
                Err(d) => Ok(d),
            }
        }
    }

    fn checked_add_duration(&self, other: &Duration) -> Option<Timespec> {
        let mut secs = other
            .as_secs()
            .try_into() // <- target type would be `libc::time_t`
            .ok()
            .and_then(|secs| self.t.tv_sec.checked_add(secs))?;

        // Nano calculations can't overflow because nanos are <1B which fit
        // in a u32.
        let mut nsec = other.subsec_nanos() + self.t.tv_nsec as u32;
        if nsec >= NSEC_PER_SEC as u32 {
            nsec -= NSEC_PER_SEC as u32;
            secs = secs.checked_add(1)?;
        }
        Some(Timespec { t: libc::timespec { tv_sec: secs, tv_nsec: nsec as _ } })
    }

    fn checked_sub_duration(&self, other: &Duration) -> Option<Timespec> {
        let mut secs = other
            .as_secs()
            .try_into() // <- target type would be `libc::time_t`
            .ok()
            .and_then(|secs| self.t.tv_sec.checked_sub(secs))?;

        // Similar to above, nanos can't overflow.
        let mut nsec = self.t.tv_nsec as i32 - other.subsec_nanos() as i32;
        if nsec < 0 {
            nsec += NSEC_PER_SEC as i32;
            secs = secs.checked_sub(1)?;
        }
        Some(Timespec { t: libc::timespec { tv_sec: secs, tv_nsec: nsec as _ } })
    }
}

impl PartialEq for Timespec {
    fn eq(&self, other: &Timespec) -> bool {
        self.t.tv_sec == other.t.tv_sec && self.t.tv_nsec == other.t.tv_nsec
    }
}

impl Eq for Timespec {}

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

impl Ord for Timespec {
    fn cmp(&self, other: &Timespec) -> Ordering {
        let me = (self.t.tv_sec, self.t.tv_nsec);
        let other = (other.t.tv_sec, other.t.tv_nsec);
        me.cmp(&other)
    }
}

impl Hash for Timespec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.t.tv_sec.hash(state);
        self.t.tv_nsec.hash(state);
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
mod inner {
    use crate::fmt;
    use crate::mem;
    use crate::sync::atomic::{AtomicUsize, Ordering::SeqCst};
    use crate::sys::cvt;
    use crate::sys_common::mul_div_u64;
    use crate::time::Duration;

    use super::Timespec;
    use super::NSEC_PER_SEC;

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
    pub struct Instant {
        t: u64,
    }

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct SystemTime {
        t: Timespec,
    }

    pub const UNIX_EPOCH: SystemTime = SystemTime { t: Timespec::zero() };

    #[repr(C)]
    #[derive(Copy, Clone)]
    struct mach_timebase_info {
        numer: u32,
        denom: u32,
    }
    type mach_timebase_info_t = *mut mach_timebase_info;
    type kern_return_t = libc::c_int;

    impl Instant {
        pub fn now() -> Instant {
            extern "C" {
                fn mach_absolute_time() -> u64;
            }
            Instant { t: unsafe { mach_absolute_time() } }
        }

        pub const fn zero() -> Instant {
            Instant { t: 0 }
        }

        pub fn actually_monotonic() -> bool {
            true
        }

        pub fn checked_sub_instant(&self, other: &Instant) -> Option<Duration> {
            let diff = self.t.checked_sub(other.t)?;
            let info = info();
            let nanos = mul_div_u64(diff, info.numer as u64, info.denom as u64);
            Some(Duration::new(nanos / NSEC_PER_SEC, (nanos % NSEC_PER_SEC) as u32))
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_add(checked_dur2intervals(other)?)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_sub(checked_dur2intervals(other)?)? })
        }
    }

    impl SystemTime {
        pub fn now() -> SystemTime {
            use crate::ptr;

            let mut s = libc::timeval { tv_sec: 0, tv_usec: 0 };
            cvt(unsafe { libc::gettimeofday(&mut s, ptr::null_mut()) }).unwrap();
            return SystemTime::from(s);
        }

        pub fn sub_time(&self, other: &SystemTime) -> Result<Duration, Duration> {
            self.t.sub_timespec(&other.t)
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl From<libc::timeval> for SystemTime {
        fn from(t: libc::timeval) -> SystemTime {
            SystemTime::from(libc::timespec {
                tv_sec: t.tv_sec,
                tv_nsec: (t.tv_usec * 1000) as libc::c_long,
            })
        }
    }

    impl From<libc::timespec> for SystemTime {
        fn from(t: libc::timespec) -> SystemTime {
            SystemTime { t: Timespec { t } }
        }
    }

    impl fmt::Debug for SystemTime {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("SystemTime")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    fn checked_dur2intervals(dur: &Duration) -> Option<u64> {
        let nanos =
            dur.as_secs().checked_mul(NSEC_PER_SEC)?.checked_add(dur.subsec_nanos() as u64)?;
        let info = info();
        Some(mul_div_u64(nanos, info.denom as u64, info.numer as u64))
    }

    fn info() -> mach_timebase_info {
        static mut INFO: mach_timebase_info = mach_timebase_info { numer: 0, denom: 0 };
        static STATE: AtomicUsize = AtomicUsize::new(0);

        unsafe {
            // If a previous thread has filled in this global state, use that.
            if STATE.load(SeqCst) == 2 {
                return INFO;
            }

            // ... otherwise learn for ourselves ...
            let mut info = mem::zeroed();
            extern "C" {
                fn mach_timebase_info(info: mach_timebase_info_t) -> kern_return_t;
            }

            mach_timebase_info(&mut info);

            // ... and attempt to be the one thread that stores it globally for
            // all other threads
            if STATE.compare_exchange(0, 1, SeqCst, SeqCst).is_ok() {
                INFO = info;
                STATE.store(2, SeqCst);
            }
            return info;
        }
    }
}

#[cfg(not(any(target_os = "macos", target_os = "ios")))]
mod inner {
    use crate::fmt;
    use crate::sys::cvt;
    use crate::time::Duration;

    use super::Timespec;

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct Instant {
        t: Timespec,
    }

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct SystemTime {
        t: Timespec,
    }

    pub const UNIX_EPOCH: SystemTime = SystemTime { t: Timespec::zero() };

    impl Instant {
        pub fn now() -> Instant {
            Instant { t: now(libc::CLOCK_MONOTONIC) }
        }

        pub const fn zero() -> Instant {
            Instant { t: Timespec::zero() }
        }

        pub fn actually_monotonic() -> bool {
            (cfg!(target_os = "linux") && cfg!(target_arch = "x86_64"))
                || (cfg!(target_os = "linux") && cfg!(target_arch = "x86"))
                || cfg!(target_os = "fuchsia")
        }

        pub fn checked_sub_instant(&self, other: &Instant) -> Option<Duration> {
            self.t.sub_timespec(&other.t).ok()
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl fmt::Debug for Instant {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Instant")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    impl SystemTime {
        pub fn now() -> SystemTime {
            SystemTime { t: now(libc::CLOCK_REALTIME) }
        }

        pub fn sub_time(&self, other: &SystemTime) -> Result<Duration, Duration> {
            self.t.sub_timespec(&other.t)
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl From<libc::timespec> for SystemTime {
        fn from(t: libc::timespec) -> SystemTime {
            SystemTime { t: Timespec { t } }
        }
    }

    impl fmt::Debug for SystemTime {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("SystemTime")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    #[cfg(not(target_os = "dragonfly"))]
    pub type clock_t = libc::c_int;
    #[cfg(target_os = "dragonfly")]
    pub type clock_t = libc::c_ulong;

    fn now(clock: clock_t) -> Timespec {
        let mut t = Timespec { t: libc::timespec { tv_sec: 0, tv_nsec: 0 } };
        cvt(unsafe { libc::clock_gettime(clock, &mut t.t) }).unwrap();
        t
    }
}
<commit_msg>std/sys/unix/time: make it easier for LLVM to optimize `Instant` subtraction.<commit_after>use crate::cmp::Ordering;
use crate::time::Duration;

use core::hash::{Hash, Hasher};

pub use self::inner::{Instant, SystemTime, UNIX_EPOCH};
use crate::convert::TryInto;

const NSEC_PER_SEC: u64 = 1_000_000_000;

#[derive(Copy, Clone)]
struct Timespec {
    t: libc::timespec,
}

impl Timespec {
    const fn zero() -> Timespec {
        Timespec { t: libc::timespec { tv_sec: 0, tv_nsec: 0 } }
    }

    fn sub_timespec(&self, other: &Timespec) -> Result<Duration, Duration> {
        if self >= other {
            // NOTE(eddyb) two aspects of this `if`-`else` are required for LLVM
            // to optimize it into a branchless form (see also #75545):
            //
            // 1. `self.t.tv_sec - other.t.tv_sec` shows up as a common expression
            //    in both branches, i.e. the `else` must have its `- 1`
            //    subtraction after the common one, not interleaved with it
            //    (it used to be `self.t.tv_sec - 1 - other.t.tv_sec`)
            //
            // 2. the `Duration::new` call (or any other additional complexity)
            //    is outside of the `if`-`else`, not duplicated in both branches
            //
            // Ideally this code could be rearranged such that it more
            // directly expresses the lower-cost behavior we want from it.
            let (secs, nsec) = if self.t.tv_nsec >= other.t.tv_nsec {
                ((self.t.tv_sec - other.t.tv_sec) as u64, (self.t.tv_nsec - other.t.tv_nsec) as u32)
            } else {
                (
                    (self.t.tv_sec - other.t.tv_sec - 1) as u64,
                    self.t.tv_nsec as u32 + (NSEC_PER_SEC as u32) - other.t.tv_nsec as u32,
                )
            };

            Ok(Duration::new(secs, nsec))
        } else {
            match other.sub_timespec(self) {
                Ok(d) => Err(d),
                Err(d) => Ok(d),
            }
        }
    }

    fn checked_add_duration(&self, other: &Duration) -> Option<Timespec> {
        let mut secs = other
            .as_secs()
            .try_into() // <- target type would be `libc::time_t`
            .ok()
            .and_then(|secs| self.t.tv_sec.checked_add(secs))?;

        // Nano calculations can't overflow because nanos are <1B which fit
        // in a u32.
        let mut nsec = other.subsec_nanos() + self.t.tv_nsec as u32;
        if nsec >= NSEC_PER_SEC as u32 {
            nsec -= NSEC_PER_SEC as u32;
            secs = secs.checked_add(1)?;
        }
        Some(Timespec { t: libc::timespec { tv_sec: secs, tv_nsec: nsec as _ } })
    }

    fn checked_sub_duration(&self, other: &Duration) -> Option<Timespec> {
        let mut secs = other
            .as_secs()
            .try_into() // <- target type would be `libc::time_t`
            .ok()
            .and_then(|secs| self.t.tv_sec.checked_sub(secs))?;

        // Similar to above, nanos can't overflow.
        let mut nsec = self.t.tv_nsec as i32 - other.subsec_nanos() as i32;
        if nsec < 0 {
            nsec += NSEC_PER_SEC as i32;
            secs = secs.checked_sub(1)?;
        }
        Some(Timespec { t: libc::timespec { tv_sec: secs, tv_nsec: nsec as _ } })
    }
}

impl PartialEq for Timespec {
    fn eq(&self, other: &Timespec) -> bool {
        self.t.tv_sec == other.t.tv_sec && self.t.tv_nsec == other.t.tv_nsec
    }
}

impl Eq for Timespec {}

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

impl Ord for Timespec {
    fn cmp(&self, other: &Timespec) -> Ordering {
        let me = (self.t.tv_sec, self.t.tv_nsec);
        let other = (other.t.tv_sec, other.t.tv_nsec);
        me.cmp(&other)
    }
}

impl Hash for Timespec {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.t.tv_sec.hash(state);
        self.t.tv_nsec.hash(state);
    }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
mod inner {
    use crate::fmt;
    use crate::mem;
    use crate::sync::atomic::{AtomicUsize, Ordering::SeqCst};
    use crate::sys::cvt;
    use crate::sys_common::mul_div_u64;
    use crate::time::Duration;

    use super::Timespec;
    use super::NSEC_PER_SEC;

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
    pub struct Instant {
        t: u64,
    }

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct SystemTime {
        t: Timespec,
    }

    pub const UNIX_EPOCH: SystemTime = SystemTime { t: Timespec::zero() };

    #[repr(C)]
    #[derive(Copy, Clone)]
    struct mach_timebase_info {
        numer: u32,
        denom: u32,
    }
    type mach_timebase_info_t = *mut mach_timebase_info;
    type kern_return_t = libc::c_int;

    impl Instant {
        pub fn now() -> Instant {
            extern "C" {
                fn mach_absolute_time() -> u64;
            }
            Instant { t: unsafe { mach_absolute_time() } }
        }

        pub const fn zero() -> Instant {
            Instant { t: 0 }
        }

        pub fn actually_monotonic() -> bool {
            true
        }

        pub fn checked_sub_instant(&self, other: &Instant) -> Option<Duration> {
            let diff = self.t.checked_sub(other.t)?;
            let info = info();
            let nanos = mul_div_u64(diff, info.numer as u64, info.denom as u64);
            Some(Duration::new(nanos / NSEC_PER_SEC, (nanos % NSEC_PER_SEC) as u32))
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_add(checked_dur2intervals(other)?)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_sub(checked_dur2intervals(other)?)? })
        }
    }

    impl SystemTime {
        pub fn now() -> SystemTime {
            use crate::ptr;

            let mut s = libc::timeval { tv_sec: 0, tv_usec: 0 };
            cvt(unsafe { libc::gettimeofday(&mut s, ptr::null_mut()) }).unwrap();
            return SystemTime::from(s);
        }

        pub fn sub_time(&self, other: &SystemTime) -> Result<Duration, Duration> {
            self.t.sub_timespec(&other.t)
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl From<libc::timeval> for SystemTime {
        fn from(t: libc::timeval) -> SystemTime {
            SystemTime::from(libc::timespec {
                tv_sec: t.tv_sec,
                tv_nsec: (t.tv_usec * 1000) as libc::c_long,
            })
        }
    }

    impl From<libc::timespec> for SystemTime {
        fn from(t: libc::timespec) -> SystemTime {
            SystemTime { t: Timespec { t } }
        }
    }

    impl fmt::Debug for SystemTime {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("SystemTime")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    fn checked_dur2intervals(dur: &Duration) -> Option<u64> {
        let nanos =
            dur.as_secs().checked_mul(NSEC_PER_SEC)?.checked_add(dur.subsec_nanos() as u64)?;
        let info = info();
        Some(mul_div_u64(nanos, info.denom as u64, info.numer as u64))
    }

    fn info() -> mach_timebase_info {
        static mut INFO: mach_timebase_info = mach_timebase_info { numer: 0, denom: 0 };
        static STATE: AtomicUsize = AtomicUsize::new(0);

        unsafe {
            // If a previous thread has filled in this global state, use that.
            if STATE.load(SeqCst) == 2 {
                return INFO;
            }

            // ... otherwise learn for ourselves ...
            let mut info = mem::zeroed();
            extern "C" {
                fn mach_timebase_info(info: mach_timebase_info_t) -> kern_return_t;
            }

            mach_timebase_info(&mut info);

            // ... and attempt to be the one thread that stores it globally for
            // all other threads
            if STATE.compare_exchange(0, 1, SeqCst, SeqCst).is_ok() {
                INFO = info;
                STATE.store(2, SeqCst);
            }
            return info;
        }
    }
}

#[cfg(not(any(target_os = "macos", target_os = "ios")))]
mod inner {
    use crate::fmt;
    use crate::sys::cvt;
    use crate::time::Duration;

    use super::Timespec;

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct Instant {
        t: Timespec,
    }

    #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
    pub struct SystemTime {
        t: Timespec,
    }

    pub const UNIX_EPOCH: SystemTime = SystemTime { t: Timespec::zero() };

    impl Instant {
        pub fn now() -> Instant {
            Instant { t: now(libc::CLOCK_MONOTONIC) }
        }

        pub const fn zero() -> Instant {
            Instant { t: Timespec::zero() }
        }

        pub fn actually_monotonic() -> bool {
            (cfg!(target_os = "linux") && cfg!(target_arch = "x86_64"))
                || (cfg!(target_os = "linux") && cfg!(target_arch = "x86"))
                || cfg!(target_os = "fuchsia")
        }

        pub fn checked_sub_instant(&self, other: &Instant) -> Option<Duration> {
            self.t.sub_timespec(&other.t).ok()
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<Instant> {
            Some(Instant { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl fmt::Debug for Instant {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("Instant")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    impl SystemTime {
        pub fn now() -> SystemTime {
            SystemTime { t: now(libc::CLOCK_REALTIME) }
        }

        pub fn sub_time(&self, other: &SystemTime) -> Result<Duration, Duration> {
            self.t.sub_timespec(&other.t)
        }

        pub fn checked_add_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_add_duration(other)? })
        }

        pub fn checked_sub_duration(&self, other: &Duration) -> Option<SystemTime> {
            Some(SystemTime { t: self.t.checked_sub_duration(other)? })
        }
    }

    impl From<libc::timespec> for SystemTime {
        fn from(t: libc::timespec) -> SystemTime {
            SystemTime { t: Timespec { t } }
        }
    }

    impl fmt::Debug for SystemTime {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            f.debug_struct("SystemTime")
                .field("tv_sec", &self.t.t.tv_sec)
                .field("tv_nsec", &self.t.t.tv_nsec)
                .finish()
        }
    }

    #[cfg(not(target_os = "dragonfly"))]
    pub type clock_t = libc::c_int;
    #[cfg(target_os = "dragonfly")]
    pub type clock_t = libc::c_ulong;

    fn now(clock: clock_t) -> Timespec {
        let mut t = Timespec { t: libc::timespec { tv_sec: 0, tv_nsec: 0 } };
        cvt(unsafe { libc::clock_gettime(clock, &mut t.t) }).unwrap();
        t
    }
}
<|endoftext|>


<commit_before>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    is_finish: bool,
    show_bar: bool,
    show_speed: bool,
    show_percent: bool,
    show_counter: bool,
    show_time_left: bool,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pb::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: false,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pb::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    /// Add to current value
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        return self.add(1);
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if !self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{}/s ", speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<commit_msg>docs(pb): add<commit_after>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    is_finish: bool,
    show_bar: bool,
    show_speed: bool,
    show_percent: bool,
    show_counter: bool,
    show_time_left: bool,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pb::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: false,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pb::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        return self.add(1);
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if !self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{}/s ", speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use std::ops::{Range, RangeFrom, RangeTo, RangeFull};
use super::Ixs;

// [a:b:s] syntax for example [:3], [::-1]
// [0,:] -- first row of matrix
// [:,0] -- first column of matrix

#[derive(Copy, PartialEq, Eq, Hash, Debug)]
/// A slice, a description of a range of an array axis.
///
/// Fields are `begin`, `end` and `stride`, where
/// negative `begin` or `end` indexes are counted from the back
/// of the axis.
///
/// If `end` is `None`, the slice extends to the end of the axis.
///
/// See also the [`s![] macro`](macro.s!.html), a convenient way to specify
/// an array of `Si`.
///
/// ## Examples
///
/// `Si(0, None, 1)` is the full range of an axis.
/// Python equivalent is `[:]`. Macro equivalent is `s![..]`.
///
/// `Si(a, Some(b), 2)` is every second element from `a` until `b`.
/// Python equivalent is `[a:b:2]`. Macro equivalent is `s![a..b;2]`.
///
/// `Si(a, None, -1)` is every element, from `a`
/// until the end, in reverse order. Python equivalent is `[a::-1]`.
/// Macro equivalent is `s![a..;-1]`.
///
/// The constant [`S`] is a shorthand for the full range of an axis.
/// [`S`]: constant.S.html
pub struct Si(pub Ixs, pub Option<Ixs>, pub Ixs);

impl From<Range<Ixs>> for Si {
    #[inline]
    fn from(r: Range<Ixs>) -> Si {
        Si(r.start, Some(r.end), 1)
    }
}

impl From<RangeFrom<Ixs>> for Si {
    #[inline]
    fn from(r: RangeFrom<Ixs>) -> Si {
        Si(r.start, None, 1)
    }
}

impl From<RangeTo<Ixs>> for Si {
    #[inline]
    fn from(r: RangeTo<Ixs>) -> Si {
        Si(0, Some(r.end), 1)
    }
}

impl From<RangeFull> for Si {
    #[inline]
    fn from(_: RangeFull) -> Si {
        S
    }
}

impl Si {
    #[inline]
    pub fn step(self, step: Ixs) -> Self {
        Si(self.0, self.1, self.2 * step)
    }
}

impl Clone for Si {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}

/// Slice value for the full range of an axis.
pub const S: Si = Si(0, None, 1);

/// Slice argument constructor.
///
/// `s![]` takes a list of ranges, separated by comma, with optional strides
/// that are separated from the range by a semicolon.
/// It is converted into a slice argument with type `&[Si; N]`.
///
/// Each range uses signed indices, where a negative value is counted from
/// the end of the axis. Strides are also signed and may be negative, but
/// must not be zero.
///
/// For example, if an array has two axes, the slice argument is passed as
/// type `&[Si; 2]`.
///
/// For example `s![a..b;c, d..e]`
/// is equivalent to `&[Si(a, Some(b), c), Si(d, Some(e), 1)]`.
///
/// ```
/// #[macro_use]
/// extern crate ndarray;
///
/// use ndarray::{
///     ArrayView,
///     Ix,
///     OwnedArray,
/// };
///
/// fn laplacian(v: &ArrayView<f32, (Ix, Ix)>) -> OwnedArray<f32, (Ix, Ix)> {
///     -4. * &v.slice(s![1..-1, 1..-1])
///     + v.slice(s![ ..-2, 1..-1])
///     + v.slice(s![1..-1,  ..-2])
///     + v.slice(s![1..-1, 2..  ])
///     + v.slice(s![2..  , 1..-1])
/// }
/// # fn main() { }
/// ```
#[macro_export]
macro_rules! s(
    (@as_expr $e:expr) => ($e);
    // convert a..b;c into @step(a..b, c), final item
    (@parse [$($stack:tt)*] $r:expr;$s:expr) => {
        s![@as_expr &[$($stack)* s!(@step $r, $s)]]
    };
    // convert a..b into @step(a..b, 1), final item
    (@parse [$($stack:tt)*] $r:expr) => {
        s![@as_expr &[$($stack)* s!(@step $r, 1)]]
    };
    // convert a..b;c into @step(a..b, c)
    (@parse [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
        s![@parse [$($stack)* s!(@step $r, $s),] $($t)*]
    };
    // convert a..b into @step(a..b, 1)
    (@parse [$($stack:tt)*] $r:expr, $($t:tt)*) => {
        s![@parse [$($stack)* s!(@step $r, 1),] $($t)*]
    };
    // convert range, step into Si
    (@step $r:expr, $s:expr) => {
        <$crate::Si as ::std::convert::From<_>>::from($r).step($s)
    };
    ($($t:tt)*) => {
        s![@parse [] $($t)*]
    };
);
<commit_msg>Add another example to s![] docs<commit_after>// Copyright 2014-2016 bluss and ndarray developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use std::ops::{Range, RangeFrom, RangeTo, RangeFull};
use super::Ixs;

// [a:b:s] syntax for example [:3], [::-1]
// [0,:] -- first row of matrix
// [:,0] -- first column of matrix

#[derive(Copy, PartialEq, Eq, Hash, Debug)]
/// A slice, a description of a range of an array axis.
///
/// Fields are `begin`, `end` and `stride`, where
/// negative `begin` or `end` indexes are counted from the back
/// of the axis.
///
/// If `end` is `None`, the slice extends to the end of the axis.
///
/// See also the [`s![] macro`](macro.s!.html), a convenient way to specify
/// an array of `Si`.
///
/// ## Examples
///
/// `Si(0, None, 1)` is the full range of an axis.
/// Python equivalent is `[:]`. Macro equivalent is `s![..]`.
///
/// `Si(a, Some(b), 2)` is every second element from `a` until `b`.
/// Python equivalent is `[a:b:2]`. Macro equivalent is `s![a..b;2]`.
///
/// `Si(a, None, -1)` is every element, from `a`
/// until the end, in reverse order. Python equivalent is `[a::-1]`.
/// Macro equivalent is `s![a..;-1]`.
///
/// The constant [`S`] is a shorthand for the full range of an axis.
/// [`S`]: constant.S.html
pub struct Si(pub Ixs, pub Option<Ixs>, pub Ixs);

impl From<Range<Ixs>> for Si {
    #[inline]
    fn from(r: Range<Ixs>) -> Si {
        Si(r.start, Some(r.end), 1)
    }
}

impl From<RangeFrom<Ixs>> for Si {
    #[inline]
    fn from(r: RangeFrom<Ixs>) -> Si {
        Si(r.start, None, 1)
    }
}

impl From<RangeTo<Ixs>> for Si {
    #[inline]
    fn from(r: RangeTo<Ixs>) -> Si {
        Si(0, Some(r.end), 1)
    }
}

impl From<RangeFull> for Si {
    #[inline]
    fn from(_: RangeFull) -> Si {
        S
    }
}

impl Si {
    #[inline]
    pub fn step(self, step: Ixs) -> Self {
        Si(self.0, self.1, self.2 * step)
    }
}

impl Clone for Si {
    #[inline]
    fn clone(&self) -> Self {
        *self
    }
}

/// Slice value for the full range of an axis.
pub const S: Si = Si(0, None, 1);

/// Slice argument constructor.
///
/// `s![]` takes a list of ranges, separated by comma, with optional strides
/// that are separated from the range by a semicolon.
/// It is converted into a slice argument with type `&[Si; N]`.
///
/// Each range uses signed indices, where a negative value is counted from
/// the end of the axis. Strides are also signed and may be negative, but
/// must not be zero.
///
/// For example `s![0..4;2, 1..5]` is a slice of rows 0..4 with step size 2,
/// and columns 1..5 with default step size 1. The slice would have
/// shape `[2, 4]`.
///
/// If an array has two axes, the slice argument is passed as
/// type `&[Si; 2]`.  The macro expansion of `s![a..b;c, d..e]`
/// is equivalent to `&[Si(a, Some(b), c), Si(d, Some(e), 1)]`.
///
/// ```
/// #[macro_use]
/// extern crate ndarray;
///
/// use ndarray::{
///     ArrayView,
///     Ix,
///     OwnedArray,
/// };
///
/// fn laplacian(v: &ArrayView<f32, (Ix, Ix)>) -> OwnedArray<f32, (Ix, Ix)> {
///     -4. * &v.slice(s![1..-1, 1..-1])
///     + v.slice(s![ ..-2, 1..-1])
///     + v.slice(s![1..-1,  ..-2])
///     + v.slice(s![1..-1, 2..  ])
///     + v.slice(s![2..  , 1..-1])
/// }
/// # fn main() { }
/// ```
#[macro_export]
macro_rules! s(
    (@as_expr $e:expr) => ($e);
    // convert a..b;c into @step(a..b, c), final item
    (@parse [$($stack:tt)*] $r:expr;$s:expr) => {
        s![@as_expr &[$($stack)* s!(@step $r, $s)]]
    };
    // convert a..b into @step(a..b, 1), final item
    (@parse [$($stack:tt)*] $r:expr) => {
        s![@as_expr &[$($stack)* s!(@step $r, 1)]]
    };
    // convert a..b;c into @step(a..b, c)
    (@parse [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
        s![@parse [$($stack)* s!(@step $r, $s),] $($t)*]
    };
    // convert a..b into @step(a..b, 1)
    (@parse [$($stack:tt)*] $r:expr, $($t:tt)*) => {
        s![@parse [$($stack)* s!(@step $r, 1),] $($t)*]
    };
    // convert range, step into Si
    (@step $r:expr, $s:expr) => {
        <$crate::Si as ::std::convert::From<_>>::from($r).step($s)
    };
    ($($t:tt)*) => {
        s![@parse [] $($t)*]
    };
);
<|endoftext|>