<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use core::mem::size_of;

use rand_core::block::{BlockRng, BlockRngCore};
use rand_core::{CryptoRng, Error, RngCore, SeedableRng};

/// A wrapper around any PRNG that implements [`BlockRngCore`], that adds the
/// ability to reseed it.
///
/// `ReseedingRng` reseeds the underlying PRNG in the following cases:
///
/// - On a manual call to [`reseed()`].
/// - After `clone()`, the clone will be reseeded on first use.
/// - After a process is forked, the RNG in the child process is reseeded within
///   the next few generated values, depending on the block size of the
///   underlying PRNG. For ChaCha and Hc128 this is a maximum of
///   15 `u32` values before reseeding.
/// - After the PRNG has generated a configurable number of random bytes.
///
/// # When should reseeding after a fixed number of generated bytes be used?
///
/// Reseeding after a fixed number of generated bytes is never strictly
/// *necessary*. Cryptographic PRNGs don't have a limited number of bytes they
/// can output, or at least not a limit reachable in any practical way. There is
/// no such thing as 'running out of entropy'.
///
/// Occasionally reseeding can be seen as some form of 'security in depth'. Even
/// if in the future a cryptographic weakness is found in the CSPRNG being used,
/// or a flaw in the implementation, occasionally reseeding should make
/// exploiting it much more difficult or even impossible.
///
/// Use [`ReseedingRng::new`] with a `threshold` of `0` to disable reseeding
/// after a fixed number of generated bytes.
///
/// # Error handling
///
/// Although unlikely, reseeding the wrapped PRNG can fail. `ReseedingRng` will
/// never panic but try to handle the error intelligently through some
/// combination of retrying and delaying reseeding until later.
/// If handling the source error fails `ReseedingRng` will continue generating
/// data from the wrapped PRNG without reseeding.
///
/// Manually calling [`reseed()`] will not have this retry or delay logic, but
/// reports the error.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand_chacha::ChaCha20Core; // Internal part of ChaChaRng that
///                              // implements BlockRngCore
/// use rand::rngs::OsRng;
/// use rand::rngs::adapter::ReseedingRng;
///
/// let prng = ChaCha20Core::from_entropy();
/// let mut reseeding_rng = ReseedingRng::new(prng, 0, OsRng);
///
/// println!("{}", reseeding_rng.gen::<u64>());
///
/// let mut cloned_rng = reseeding_rng.clone();
/// assert!(reseeding_rng.gen::<u64>() != cloned_rng.gen::<u64>());
/// ```
///
/// [`BlockRngCore`]: rand_core::block::BlockRngCore
/// [`ReseedingRng::new`]: ReseedingRng::new
/// [`reseed()`]: ReseedingRng::reseed
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr>(BlockRng<ReseedingCore<R, Rsdr>>)
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore;

impl<R, Rsdr> ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    /// Create a new `ReseedingRng` from an existing PRNG, combined with a RNG
    /// to use as reseeder.
    ///
    /// `threshold` sets the number of generated bytes after which to reseed the
    /// PRNG. Set it to zero to never reseed based on the number of generated
    /// values.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
        ReseedingRng(BlockRng::new(ReseedingCore::new(rng, threshold, reseeder)))
    }

    /// Reseed the internal PRNG.
    pub fn reseed(&mut self) -> Result<(), Error> {
        self.0.core.reseed()
    }
}

// TODO: this should be implemented for any type where the inner type
// implements RngCore, but we can't specify that because ReseedingCore is private
impl<R, Rsdr: RngCore> RngCore for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore<Item = u32> + SeedableRng,
    <R as BlockRngCore>::Results: AsRef<[u32]> + AsMut<[u32]>,
{
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        self.0.next_u32()
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        self.0.next_u64()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.0.fill_bytes(dest)
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

impl<R, Rsdr> Clone for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + Clone,
    Rsdr: RngCore + Clone,
{
    fn clone(&self) -> ReseedingRng<R, Rsdr> {
        // Recreating `BlockRng` seems easier than cloning it and resetting
        // the index.
        ReseedingRng(BlockRng::new(self.0.core.clone()))
    }
}

impl<R, Rsdr> CryptoRng for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + CryptoRng,
    Rsdr: RngCore + CryptoRng,
{
}

#[derive(Debug)]
struct ReseedingCore<R, Rsdr> {
    inner: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
    fork_counter: usize,
}

impl<R, Rsdr> BlockRngCore for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    type Item = <R as BlockRngCore>::Item;
    type Results = <R as BlockRngCore>::Results;

    fn generate(&mut self, results: &mut Self::Results) {
        let global_fork_counter = fork::get_fork_counter();
        if self.bytes_until_reseed <= 0 || self.is_forked(global_fork_counter) {
            // We get better performance by not calling only `reseed` here
            // and continuing with the rest of the function, but by directly
            // returning from a non-inlined function.
            return self.reseed_and_generate(results, global_fork_counter);
        }
        let num_bytes = results.as_ref().len() * size_of::<Self::Item>();
        self.bytes_until_reseed -= num_bytes as i64;
        self.inner.generate(results);
    }
}

impl<R, Rsdr> ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    /// Create a new `ReseedingCore`.
    fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
        use ::core::i64::MAX;
        fork::register_fork_handler();

        // Because generating more values than `i64::MAX` takes centuries on
        // current hardware, we just clamp to that value.
        // Also we set a threshold of 0, which indicates no limit, to that
        // value.
        let threshold = if threshold == 0 {
            MAX
        } else if threshold <= MAX as u64 {
            threshold as i64
        } else {
            MAX
        };

        ReseedingCore {
            inner: rng,
            reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
            fork_counter: 0,
        }
    }

    /// Reseed the internal PRNG.
    fn reseed(&mut self) -> Result<(), Error> {
        R::from_rng(&mut self.reseeder).map(|result| {
            self.bytes_until_reseed = self.threshold;
            self.inner = result
        })
    }

    fn is_forked(&self, global_fork_counter: usize) -> bool {
        // In theory, on 32-bit platforms, it is possible for
        // `global_fork_counter` to wrap around after ~4e9 forks.
        //
        // This check will detect a fork in the normal case where
        // `fork_counter < global_fork_counter`, and also when the difference
        // between both is greater than `isize::MAX` (wrapped around).
        //
        // It will still fail to detect a fork if there have been more than
        // `isize::MAX` forks, without any reseed in between. Seems unlikely
        // enough.
        (self.fork_counter.wrapping_sub(global_fork_counter) as isize) < 0
    }

    #[inline(never)]
    fn reseed_and_generate(
        &mut self, results: &mut <Self as BlockRngCore>::Results, global_fork_counter: usize,
    ) {
        #![allow(clippy::if_same_then_else)] // false positive
        if self.is_forked(global_fork_counter) {
            info!("Fork detected, reseeding RNG");
        } else {
            trace!("Reseeding RNG (periodic reseed)");
        }

        let num_bytes = results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();

        if let Err(e) = self.reseed() {
            warn!("Reseeding RNG failed: {}", e);
            let _ = e;
        }
        self.fork_counter = global_fork_counter;

        self.bytes_until_reseed = self.threshold - num_bytes as i64;
        self.inner.generate(results);
    }
}

impl<R, Rsdr> Clone for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + Clone,
    Rsdr: RngCore + Clone,
{
    fn clone(&self) -> ReseedingCore<R, Rsdr> {
        ReseedingCore {
            inner: self.inner.clone(),
            reseeder: self.reseeder.clone(),
            threshold: self.threshold,
            bytes_until_reseed: 0, // reseed clone on first use
            fork_counter: self.fork_counter,
        }
    }
}

impl<R, Rsdr> CryptoRng for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + CryptoRng,
    Rsdr: RngCore + CryptoRng,
{
}

#[cfg(all(unix, feature = "std", not(target_os = "emscripten")))]
mod fork {
    use core::sync::atomic::{AtomicUsize, Ordering};
    #[allow(deprecated)] // Required for compatibility with Rust < 1.24.
    use core::sync::atomic::{ATOMIC_USIZE_INIT};
    use std::sync::Once;

    // Fork protection
    //
    // We implement fork protection on Unix using `pthread_atfork`.
    // When the process is forked, we increment `RESEEDING_RNG_FORK_COUNTER`.
    // Every `ReseedingRng` stores the last known value of the static in
    // `fork_counter`. If the cached `fork_counter` is less than
    // `RESEEDING_RNG_FORK_COUNTER`, it is time to reseed this RNG.
    //
    // If reseeding fails, we don't deal with this by setting a delay, but just
    // don't update `fork_counter`, so a reseed is attempted as soon as
    // possible.

    #[allow(deprecated)]
    static RESEEDING_RNG_FORK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

    pub fn get_fork_counter() -> usize {
        RESEEDING_RNG_FORK_COUNTER.load(Ordering::Relaxed)
    }

    extern "C" fn fork_handler() {
        // Note: fetch_add is defined to wrap on overflow
        // (which is what we want).
        RESEEDING_RNG_FORK_COUNTER.fetch_add(1, Ordering::Relaxed);
    }

    pub fn register_fork_handler() {
        static REGISTER: Once = Once::new();
        REGISTER.call_once(|| unsafe {
            libc::pthread_atfork(None, None, Some(fork_handler));
        });
    }
}

#[cfg(not(all(unix, feature = "std", not(target_os = "emscripten"))))]
mod fork {
    pub fn get_fork_counter() -> usize {
        0
    }
    pub fn register_fork_handler() {}
}

#[cfg(test)]
mod test {
    use super::ReseedingRng;
    use crate::rngs::mock::StepRng;
    use crate::rngs::std::Core;
    use crate::{Rng, SeedableRng};

    #[test]
    fn test_reseeding() {
        let mut zero = StepRng::new(0, 0);
        let rng = Core::from_rng(&mut zero).unwrap();
        let thresh = 1; // reseed every time the buffer is exhausted
        let mut reseeding = ReseedingRng::new(rng, thresh, zero);

        // RNG buffer size is [u32; 64]
        // Debug is only implemented up to length 32 so use two arrays
        let mut buf = ([0u32; 32], [0u32; 32]);
        reseeding.fill(&mut buf.0);
        reseeding.fill(&mut buf.1);
        let seq = buf;
        for _ in 0..10 {
            reseeding.fill(&mut buf.0);
            reseeding.fill(&mut buf.1);
            assert_eq!(buf, seq);
        }
    }

    #[test]
    fn test_clone_reseeding() {
        let mut zero = StepRng::new(0, 0);
        let rng = Core::from_rng(&mut zero).unwrap();
        let mut rng1 = ReseedingRng::new(rng, 32 * 4, zero);

        let first: u32 = rng1.gen();
        for _ in 0..10 {
            let _ = rng1.gen::<u32>();
        }

        let mut rng2 = rng1.clone();
        assert_eq!(first, rng2.gen::<u32>());
    }
}
<commit_msg>Remove use of deprecated ATOMIC_USIZE_INIT<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2013 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use core::mem::size_of;

use rand_core::block::{BlockRng, BlockRngCore};
use rand_core::{CryptoRng, Error, RngCore, SeedableRng};

/// A wrapper around any PRNG that implements [`BlockRngCore`], that adds the
/// ability to reseed it.
///
/// `ReseedingRng` reseeds the underlying PRNG in the following cases:
///
/// - On a manual call to [`reseed()`].
/// - After `clone()`, the clone will be reseeded on first use.
/// - After a process is forked, the RNG in the child process is reseeded within
///   the next few generated values, depending on the block size of the
///   underlying PRNG. For ChaCha and Hc128 this is a maximum of
///   15 `u32` values before reseeding.
/// - After the PRNG has generated a configurable number of random bytes.
///
/// # When should reseeding after a fixed number of generated bytes be used?
///
/// Reseeding after a fixed number of generated bytes is never strictly
/// *necessary*. Cryptographic PRNGs don't have a limited number of bytes they
/// can output, or at least not a limit reachable in any practical way. There is
/// no such thing as 'running out of entropy'.
///
/// Occasionally reseeding can be seen as some form of 'security in depth'. Even
/// if in the future a cryptographic weakness is found in the CSPRNG being used,
/// or a flaw in the implementation, occasionally reseeding should make
/// exploiting it much more difficult or even impossible.
///
/// Use [`ReseedingRng::new`] with a `threshold` of `0` to disable reseeding
/// after a fixed number of generated bytes.
///
/// # Error handling
///
/// Although unlikely, reseeding the wrapped PRNG can fail. `ReseedingRng` will
/// never panic but try to handle the error intelligently through some
/// combination of retrying and delaying reseeding until later.
/// If handling the source error fails `ReseedingRng` will continue generating
/// data from the wrapped PRNG without reseeding.
///
/// Manually calling [`reseed()`] will not have this retry or delay logic, but
/// reports the error.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand_chacha::ChaCha20Core; // Internal part of ChaChaRng that
///                              // implements BlockRngCore
/// use rand::rngs::OsRng;
/// use rand::rngs::adapter::ReseedingRng;
///
/// let prng = ChaCha20Core::from_entropy();
/// let mut reseeding_rng = ReseedingRng::new(prng, 0, OsRng);
///
/// println!("{}", reseeding_rng.gen::<u64>());
///
/// let mut cloned_rng = reseeding_rng.clone();
/// assert!(reseeding_rng.gen::<u64>() != cloned_rng.gen::<u64>());
/// ```
///
/// [`BlockRngCore`]: rand_core::block::BlockRngCore
/// [`ReseedingRng::new`]: ReseedingRng::new
/// [`reseed()`]: ReseedingRng::reseed
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr>(BlockRng<ReseedingCore<R, Rsdr>>)
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore;

impl<R, Rsdr> ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    /// Create a new `ReseedingRng` from an existing PRNG, combined with a RNG
    /// to use as reseeder.
    ///
    /// `threshold` sets the number of generated bytes after which to reseed the
    /// PRNG. Set it to zero to never reseed based on the number of generated
    /// values.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
        ReseedingRng(BlockRng::new(ReseedingCore::new(rng, threshold, reseeder)))
    }

    /// Reseed the internal PRNG.
    pub fn reseed(&mut self) -> Result<(), Error> {
        self.0.core.reseed()
    }
}

// TODO: this should be implemented for any type where the inner type
// implements RngCore, but we can't specify that because ReseedingCore is private
impl<R, Rsdr: RngCore> RngCore for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore<Item = u32> + SeedableRng,
    <R as BlockRngCore>::Results: AsRef<[u32]> + AsMut<[u32]>,
{
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        self.0.next_u32()
    }

    #[inline(always)]
    fn next_u64(&mut self) -> u64 {
        self.0.next_u64()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.0.fill_bytes(dest)
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.0.try_fill_bytes(dest)
    }
}

impl<R, Rsdr> Clone for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + Clone,
    Rsdr: RngCore + Clone,
{
    fn clone(&self) -> ReseedingRng<R, Rsdr> {
        // Recreating `BlockRng` seems easier than cloning it and resetting
        // the index.
        ReseedingRng(BlockRng::new(self.0.core.clone()))
    }
}

impl<R, Rsdr> CryptoRng for ReseedingRng<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + CryptoRng,
    Rsdr: RngCore + CryptoRng,
{
}

#[derive(Debug)]
struct ReseedingCore<R, Rsdr> {
    inner: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
    fork_counter: usize,
}

impl<R, Rsdr> BlockRngCore for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    type Item = <R as BlockRngCore>::Item;
    type Results = <R as BlockRngCore>::Results;

    fn generate(&mut self, results: &mut Self::Results) {
        let global_fork_counter = fork::get_fork_counter();
        if self.bytes_until_reseed <= 0 || self.is_forked(global_fork_counter) {
            // We get better performance by not calling only `reseed` here
            // and continuing with the rest of the function, but by directly
            // returning from a non-inlined function.
            return self.reseed_and_generate(results, global_fork_counter);
        }
        let num_bytes = results.as_ref().len() * size_of::<Self::Item>();
        self.bytes_until_reseed -= num_bytes as i64;
        self.inner.generate(results);
    }
}

impl<R, Rsdr> ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng,
    Rsdr: RngCore,
{
    /// Create a new `ReseedingCore`.
    fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
        use ::core::i64::MAX;
        fork::register_fork_handler();

        // Because generating more values than `i64::MAX` takes centuries on
        // current hardware, we just clamp to that value.
        // Also we set a threshold of 0, which indicates no limit, to that
        // value.
        let threshold = if threshold == 0 {
            MAX
        } else if threshold <= MAX as u64 {
            threshold as i64
        } else {
            MAX
        };

        ReseedingCore {
            inner: rng,
            reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
            fork_counter: 0,
        }
    }

    /// Reseed the internal PRNG.
    fn reseed(&mut self) -> Result<(), Error> {
        R::from_rng(&mut self.reseeder).map(|result| {
            self.bytes_until_reseed = self.threshold;
            self.inner = result
        })
    }

    fn is_forked(&self, global_fork_counter: usize) -> bool {
        // In theory, on 32-bit platforms, it is possible for
        // `global_fork_counter` to wrap around after ~4e9 forks.
        //
        // This check will detect a fork in the normal case where
        // `fork_counter < global_fork_counter`, and also when the difference
        // between both is greater than `isize::MAX` (wrapped around).
        //
        // It will still fail to detect a fork if there have been more than
        // `isize::MAX` forks, without any reseed in between. Seems unlikely
        // enough.
        (self.fork_counter.wrapping_sub(global_fork_counter) as isize) < 0
    }

    #[inline(never)]
    fn reseed_and_generate(
        &mut self, results: &mut <Self as BlockRngCore>::Results, global_fork_counter: usize,
    ) {
        #![allow(clippy::if_same_then_else)] // false positive
        if self.is_forked(global_fork_counter) {
            info!("Fork detected, reseeding RNG");
        } else {
            trace!("Reseeding RNG (periodic reseed)");
        }

        let num_bytes = results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();

        if let Err(e) = self.reseed() {
            warn!("Reseeding RNG failed: {}", e);
            let _ = e;
        }
        self.fork_counter = global_fork_counter;

        self.bytes_until_reseed = self.threshold - num_bytes as i64;
        self.inner.generate(results);
    }
}

impl<R, Rsdr> Clone for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + Clone,
    Rsdr: RngCore + Clone,
{
    fn clone(&self) -> ReseedingCore<R, Rsdr> {
        ReseedingCore {
            inner: self.inner.clone(),
            reseeder: self.reseeder.clone(),
            threshold: self.threshold,
            bytes_until_reseed: 0, // reseed clone on first use
            fork_counter: self.fork_counter,
        }
    }
}

impl<R, Rsdr> CryptoRng for ReseedingCore<R, Rsdr>
where
    R: BlockRngCore + SeedableRng + CryptoRng,
    Rsdr: RngCore + CryptoRng,
{
}

#[cfg(all(unix, feature = "std", not(target_os = "emscripten")))]
mod fork {
    use core::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Once;

    // Fork protection
    //
    // We implement fork protection on Unix using `pthread_atfork`.
    // When the process is forked, we increment `RESEEDING_RNG_FORK_COUNTER`.
    // Every `ReseedingRng` stores the last known value of the static in
    // `fork_counter`. If the cached `fork_counter` is less than
    // `RESEEDING_RNG_FORK_COUNTER`, it is time to reseed this RNG.
    //
    // If reseeding fails, we don't deal with this by setting a delay, but just
    // don't update `fork_counter`, so a reseed is attempted as soon as
    // possible.

    static RESEEDING_RNG_FORK_COUNTER: AtomicUsize = AtomicUsize::new(0);

    pub fn get_fork_counter() -> usize {
        RESEEDING_RNG_FORK_COUNTER.load(Ordering::Relaxed)
    }

    extern "C" fn fork_handler() {
        // Note: fetch_add is defined to wrap on overflow
        // (which is what we want).
        RESEEDING_RNG_FORK_COUNTER.fetch_add(1, Ordering::Relaxed);
    }

    pub fn register_fork_handler() {
        static REGISTER: Once = Once::new();
        REGISTER.call_once(|| unsafe {
            libc::pthread_atfork(None, None, Some(fork_handler));
        });
    }
}

#[cfg(not(all(unix, feature = "std", not(target_os = "emscripten"))))]
mod fork {
    pub fn get_fork_counter() -> usize {
        0
    }
    pub fn register_fork_handler() {}
}

#[cfg(test)]
mod test {
    use super::ReseedingRng;
    use crate::rngs::mock::StepRng;
    use crate::rngs::std::Core;
    use crate::{Rng, SeedableRng};

    #[test]
    fn test_reseeding() {
        let mut zero = StepRng::new(0, 0);
        let rng = Core::from_rng(&mut zero).unwrap();
        let thresh = 1; // reseed every time the buffer is exhausted
        let mut reseeding = ReseedingRng::new(rng, thresh, zero);

        // RNG buffer size is [u32; 64]
        // Debug is only implemented up to length 32 so use two arrays
        let mut buf = ([0u32; 32], [0u32; 32]);
        reseeding.fill(&mut buf.0);
        reseeding.fill(&mut buf.1);
        let seq = buf;
        for _ in 0..10 {
            reseeding.fill(&mut buf.0);
            reseeding.fill(&mut buf.1);
            assert_eq!(buf, seq);
        }
    }

    #[test]
    fn test_clone_reseeding() {
        let mut zero = StepRng::new(0, 0);
        let rng = Core::from_rng(&mut zero).unwrap();
        let mut rng1 = ReseedingRng::new(rng, 32 * 4, zero);

        let first: u32 = rng1.gen();
        for _ in 0..10 {
            let _ = rng1.gen::<u32>();
        }

        let mut rng2 = rng1.clone();
        assert_eq!(first, rng2.gen::<u32>());
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::cdef::*;
use crate::cpu_features::CpuFeatureLevel;
use crate::frame::*;
use crate::tiling::PlaneRegionMut;
use crate::util::*;

type CdefFilterFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  tmp: *const u16,
  tmp_stride: isize,
  pri_strength: i32,
  sec_strength: i32,
  dir: i32,
  damping: i32,
);

type CdefFilterHBDFn = unsafe extern fn(
  dst: *mut u16,
  dst_stride: isize,
  tmp: *const u16,
  tmp_stride: isize,
  pri_strength: i32,
  sec_strength: i32,
  dir: i32,
  damping: i32,
  bitdepth_max: i32,
);

#[inline(always)]
const fn decimate_index(xdec: usize, ydec: usize) -> usize {
  ((ydec << 1) | xdec) & 3
}

pub(crate) unsafe fn cdef_filter_block<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, src: *const T, src_stride: isize,
  pri_strength: i32, sec_strength: i32, dir: usize, damping: i32,
  bit_depth: usize, xdec: usize, ydec: usize, edges: u8, cpu: CpuFeatureLevel,
) {
  let call_rust = |dst: &mut PlaneRegionMut<T>| {
    rust::cdef_filter_block(
      dst,
      src,
      src_stride,
      pri_strength,
      sec_strength,
      dir,
      damping,
      bit_depth,
      xdec,
      ydec,
      edges,
      cpu,
    );
  };

  // TODO: handle padding in the fast path
  if edges != CDEF_HAVE_ALL {
    call_rust(dst);
  } else {
    #[cfg(feature = "check_asm")]
    let ref_dst = {
      let mut copy = dst.scratch_copy();
      call_rust(&mut copy.as_region_mut());
      copy
    };
    match T::type_enum() {
      PixelType::U8 => {
        match CDEF_FILTER_FNS[cpu.as_index()][decimate_index(xdec, ydec)] {
          Some(func) => {
            // current cdef_filter_block asm does 16->8 for historical
            // reasons.  Copy into tmp space for now (also handling
            // padding) until asm is updated
            const TMPSTRIDE: isize =
              std::mem::align_of::<Aligned<u16>>() as isize;
            /* 256 or 512-bit alignment, greater than 2 * (8>>xdec) + 2 */
            const TMPSIZE: usize =
              ((2 + 8 + 2) * TMPSTRIDE + TMPSTRIDE) as usize;
            let mut tmp: Aligned<[u16; TMPSIZE]> =
              Aligned::new([CDEF_VERY_LARGE; TMPSIZE]);
            rust::pad_into_tmp16(
              tmp.data.as_mut_ptr().offset(TMPSTRIDE - 2), // points to
              // *padding* upper left; the -2 is to make sure the
              // block area is SIMD-aligned, not the padding
              TMPSTRIDE,
              src, // points to *block* upper left
              src_stride,
              8 >> xdec,
              8 >> ydec,
              edges,
            );
            (func)(
              dst.data_ptr_mut() as *mut _,
              T::to_asm_stride(dst.plane_cfg.stride),
              tmp.data.as_ptr().offset(3 * TMPSTRIDE) as *const u16,
              TMPSTRIDE,
              pri_strength,
              sec_strength,
              dir as i32,
              damping,
            );
          }
          None => call_rust(dst),
        }
      }
      PixelType::U16 => {
        match CDEF_FILTER_HBD_FNS[cpu.as_index()][decimate_index(xdec, ydec)] {
          Some(func) => {
            (func)(
              dst.data_ptr_mut() as *mut _,
              T::to_asm_stride(dst.plane_cfg.stride),
              src as *const _,
              src_stride,
              pri_strength,
              sec_strength,
              dir as i32,
              damping,
              (1 << bit_depth) - 1,
            );
          }
          None => call_rust(dst),
        }
      }
    }
    #[cfg(feature = "check_asm")]
    {
      for (dst_row, ref_row) in
        dst.rows_iter().zip(ref_dst.as_region().rows_iter())
      {
        for (dst, reference) in dst_row.iter().zip(ref_row) {
          assert_eq!(*dst, *reference);
        }
      }
    }
  }
}

extern {
  fn rav1e_cdef_filter_4x4_neon(
    dst: *mut u8, dst_stride: isize, tmp: *const u16, tmp_stride: isize,
    pri_strength: i32, sec_strength: i32, dir: i32, damping: i32,
  );

  fn rav1e_cdef_filter_4x8_neon(
    dst: *mut u8, dst_stride: isize, tmp: *const u16, tmp_stride: isize,
    pri_strength: i32, sec_strength: i32, dir: i32, damping: i32,
  );

  fn rav1e_cdef_filter_8x8_neon(
    dst: *mut u8, dst_stride: isize, tmp: *const u16, tmp_stride: isize,
    pri_strength: i32, sec_strength: i32, dir: i32, damping: i32,
  );
}

static CDEF_FILTER_FNS_NEON: [Option<CdefFilterFn>; 4] = {
  let mut out: [Option<CdefFilterFn>; 4] = [None; 4];
  out[decimate_index(1, 1)] = Some(rav1e_cdef_filter_4x4_neon);
  out[decimate_index(1, 0)] = Some(rav1e_cdef_filter_4x8_neon);
  out[decimate_index(0, 0)] = Some(rav1e_cdef_filter_8x8_neon);
  out
};

cpu_function_lookup_table!(
  CDEF_FILTER_FNS: [[Option<CdefFilterFn>; 4]],
  default: [None; 4],
  // [NEON]
  []
);

cpu_function_lookup_table!(
  CDEF_FILTER_HBD_FNS: [[Option<CdefFilterHBDFn>; 4]],
  default: [None; 4],
  []
);

type CdefDirLBDFn =
  unsafe extern fn(tmp: *const u8, tmp_stride: isize, var: *mut u32) -> i32;
type CdefDirHBDFn =
  unsafe extern fn(tmp: *const u16, tmp_stride: isize, var: *mut u32) -> i32;

#[inline(always)]
#[allow(clippy::let_and_return)]
pub(crate) fn cdef_find_dir<T: Pixel>(
  img: &PlaneSlice<'_, T>, var: &mut u32, coeff_shift: usize,
  cpu: CpuFeatureLevel,
) -> i32 {
  let call_rust =
    |var: &mut u32| rust::cdef_find_dir::<T>(img, var, coeff_shift, cpu);

  #[cfg(feature = "check_asm")]
  let (ref_dir, ref_var) = {
    let mut var: u32 = 0;
    let dir = call_rust(&mut var);
    (dir, var)
  };

  let dir = match T::type_enum() {
    PixelType::U8 => {
      if let Some(func) = CDEF_DIR_LBD_FNS[cpu.as_index()] {
        unsafe {
          (func)(
            img.as_ptr() as *const _,
            T::to_asm_stride(img.plane.cfg.stride),
            var as *mut u32,
          )
        }
      } else {
        call_rust(var)
      }
    }
    PixelType::U16 => {
      if let Some(func) = CDEF_DIR_HBD_FNS[cpu.as_index()] {
        unsafe {
          (func)(
            img.as_ptr() as *const _,
            T::to_asm_stride(img.plane.cfg.stride),
            var as *mut u32,
          )
        }
      } else {
        call_rust(var)
      }
    }
  };

  #[cfg(feature = "check_asm")]
  {
    assert_eq!(dir, ref_dir);
    assert_eq!(*var, ref_var);
  }

  dir
}

extern {
  fn rav1e_cdef_find_dir_8bpc_neon(
    tmp: *const u8, tmp_stride: isize, var: *mut u32,
  ) -> i32;
}

extern {
  fn rav1e_cdef_find_dir_16bpc_neon(
    tmp: *const u16, tmp_stride: isize, var: *mut u32,
  ) -> i32;
}

cpu_function_lookup_table!(
  CDEF_DIR_LBD_FNS: [Option<CdefDirLBDFn>],
  default: None,
  [(NEON, Some(rav1e_cdef_find_dir_8bpc_neon))]
);

cpu_function_lookup_table!(
  CDEF_DIR_HBD_FNS: [Option<CdefDirHBDFn>],
  default: None,
  [(NEON, Some(rav1e_cdef_find_dir_16bpc_neon))]
);

<commit_msg>Enable full aarch64 assembly for 8-bit CDEF<commit_after>// Copyright (c) 2020, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

use crate::cdef::*;
use crate::cpu_features::CpuFeatureLevel;
use crate::frame::*;
use crate::tiling::PlaneRegionMut;
use crate::util::*;

type CdefPaddingFn = unsafe extern fn(
  tmp: *mut u16,
  src: *const u8,
  src_stride: isize,
  left: *const [u8; 2],
  top: *const u8,
  h: i32,
  edges: isize,
);

type CdefFilterFn = unsafe extern fn(
  dst: *mut u8,
  dst_stride: isize,
  tmp: *const u16,
  pri_strength: i32,
  sec_strength: i32,
  dir: i32,
  damping: i32,
  h: i32,
  edges: isize,
);

#[inline(always)]
const fn decimate_index(xdec: usize, ydec: usize) -> usize {
  ((ydec << 1) | xdec) & 3
}

pub(crate) unsafe fn cdef_filter_block<T: Pixel>(
  dst: &mut PlaneRegionMut<'_, T>, src: *const T, src_stride: isize,
  pri_strength: i32, sec_strength: i32, dir: usize, damping: i32,
  bit_depth: usize, xdec: usize, ydec: usize, edges: u8, cpu: CpuFeatureLevel,
) {
  let call_rust = |dst: &mut PlaneRegionMut<T>| {
    rust::cdef_filter_block(
      dst,
      src,
      src_stride,
      pri_strength,
      sec_strength,
      dir,
      damping,
      bit_depth,
      xdec,
      ydec,
      edges,
      cpu,
    );
  };

  #[cfg(feature = "check_asm")]
  let ref_dst = {
    let mut copy = dst.scratch_copy();
    call_rust(&mut copy.as_region_mut());
    copy
  };
  match T::type_enum() {
    PixelType::U8 => {
      match (CDEF_PAD_FNS[cpu.as_index()][decimate_index(xdec, ydec)],
             CDEF_FILTER_FNS[cpu.as_index()][decimate_index(xdec, ydec)]) {
        (Some(pad), Some(func)) => {
	  let h = if ydec == 1 {4} else {8} as i32;
	  let tmpstride = if xdec == 1 {8} else {16} as isize;
	  const MAXTMPSTRIDE: isize = 16;
          const TMPSIZE: usize =
            (12 * MAXTMPSTRIDE + 8) as usize;
          let mut tmp: Aligned<[u16; TMPSIZE]> =
            Aligned::new([CDEF_VERY_LARGE; TMPSIZE]);
          let top = src.offset( -2*src_stride );
          let mut left: Aligned<[[u8; 2]; 8]> =
            Aligned::new([[0; 2]; 8]);

          // Rather than modify the dav1d assembly, just swallow our
	  // pride and copy relevant portions of src into a left
	  // array.  The array is a monolithic, packed [x=2][y=8],
	  // though it is aligned to start on a 16-byte boundary.
          if edges & CDEF_HAVE_LEFT != 0 {
            let mut wptr = src.offset(-2) as *const u8;
            for i in 0..h {
	      left.data[i as usize][0] = *wptr;
	      left.data[i as usize][1] = *wptr.add(1);
	      wptr = wptr.offset(src_stride);
            }
	  }

	  // dav1d's implicit indexing as of this version: tmp array
	  // working pointer points to the upper-left of the current
	  // coded block, not the upper-left of the tmp array storage,
	  // with an adjustment to ensure UL of the block is 16-byte
	  // aligned.  src, similarly, points to upper left of src
	  // block being coded.  top points to coded block minus two
	  // rows (that is, src.x, src.y-2).  It does _not_ point to
	  // src.x-2, src.y-2.
          (pad) (
            tmp.data.as_mut_ptr().offset(2 * tmpstride + 8) as *mut u16,
            src as *const u8,
            src_stride,
            left.data.as_ptr() as *const [u8; 2],
            top as *const u8,
            8 >> ydec,
            edges.into(),
          );

          (func) (
            dst.data_ptr_mut() as *mut u8,
            T::to_asm_stride(dst.plane_cfg.stride),
            tmp.data.as_ptr().offset(2 * tmpstride + 8) as *const u16,
            pri_strength,
            sec_strength,
            dir as i32,
            damping,
            8 >> ydec,
            edges.into(),
	  );
	}
        _ => call_rust(dst),
      }
    }
    PixelType::U16 => {
      call_rust(dst);
    }
  }
  #[cfg(feature = "check_asm")]
  {
    for (dst_row, ref_row) in
      dst.rows_iter().zip(ref_dst.as_region().rows_iter())
    {
      for (dst, reference) in dst_row.iter().zip(ref_row) {
        assert_eq!(*dst, *reference);
        }
    }
  }
}

extern {
  fn rav1e_cdef_filter4_8bpc_neon(
    dst: *mut u8, dst_stride: isize, tmp: *const u16,
    pri_strength: i32, sec_strength: i32, dir: i32, damping: i32,
    h: i32, edges: isize,
  );

  fn rav1e_cdef_padding4_8bpc_neon(
    tmp: *mut u16, src: *const u8, src_stride: isize, left: *const [u8; 2],
    top: *const u8, h: i32, edges: isize,
  );

  fn rav1e_cdef_filter8_8bpc_neon(
    dst: *mut u8, dst_stride: isize, tmp: *const u16,
    pri_strength: i32, sec_strength: i32, dir: i32, damping: i32,
    h: i32, edges: isize,
  );

  fn rav1e_cdef_padding8_8bpc_neon(
    tmp: *mut u16, src: *const u8, src_stride: isize, left: *const [u8; 2],
    top: *const u8, h: i32, edges: isize,
  );
}

static CDEF_PAD_FNS_NEON: [Option<CdefPaddingFn>; 4] = {
  let mut out: [Option<CdefPaddingFn>; 4] = [None; 4];
  out[decimate_index(1, 1)] = Some(rav1e_cdef_padding4_8bpc_neon);
  out[decimate_index(1, 0)] = Some(rav1e_cdef_padding4_8bpc_neon);
  out[decimate_index(0, 0)] = Some(rav1e_cdef_padding8_8bpc_neon);
  out
};

static CDEF_FILTER_FNS_NEON: [Option<CdefFilterFn>; 4] = {
  let mut out: [Option<CdefFilterFn>; 4] = [None; 4];
  out[decimate_index(1, 1)] = Some(rav1e_cdef_filter4_8bpc_neon);
  out[decimate_index(1, 0)] = Some(rav1e_cdef_filter4_8bpc_neon);
  out[decimate_index(0, 0)] = Some(rav1e_cdef_filter8_8bpc_neon);
  out
};

cpu_function_lookup_table!(
  CDEF_PAD_FNS: [[Option<CdefPaddingFn>; 4]],
  default: [None; 4],
  [NEON]
);

cpu_function_lookup_table!(
  CDEF_FILTER_FNS: [[Option<CdefFilterFn>; 4]],
  default: [None; 4],
  [NEON]
);

type CdefDirLBDFn =
  unsafe extern fn(tmp: *const u8, tmp_stride: isize, var: *mut u32) -> i32;
type CdefDirHBDFn =
  unsafe extern fn(tmp: *const u16, tmp_stride: isize, var: *mut u32) -> i32;

#[inline(always)]
#[allow(clippy::let_and_return)]
pub(crate) fn cdef_find_dir<T: Pixel>(
  img: &PlaneSlice<'_, T>, var: &mut u32, coeff_shift: usize,
  cpu: CpuFeatureLevel,
) -> i32 {
  let call_rust =
    |var: &mut u32| rust::cdef_find_dir::<T>(img, var, coeff_shift, cpu);

  #[cfg(feature = "check_asm")]
  let (ref_dir, ref_var) = {
    let mut var: u32 = 0;
    let dir = call_rust(&mut var);
    (dir, var)
  };

  let dir = match T::type_enum() {
    PixelType::U8 => {
      if let Some(func) = CDEF_DIR_LBD_FNS[cpu.as_index()] {
        unsafe {
          (func)(
            img.as_ptr() as *const _,
            T::to_asm_stride(img.plane.cfg.stride),
            var as *mut u32,
          )
        }
      } else {
        call_rust(var)
      }
    }
    PixelType::U16 => {
      if let Some(func) = CDEF_DIR_HBD_FNS[cpu.as_index()] {
        unsafe {
          (func)(
            img.as_ptr() as *const _,
            T::to_asm_stride(img.plane.cfg.stride),
            var as *mut u32,
          )
        }
      } else {
        call_rust(var)
      }
    }
  };

  #[cfg(feature = "check_asm")]
  {
    assert_eq!(dir, ref_dir);
    assert_eq!(*var, ref_var);
  }

  dir
}

extern {
  fn rav1e_cdef_find_dir_8bpc_neon(
    tmp: *const u8, tmp_stride: isize, var: *mut u32,
  ) -> i32;
}

extern {
  fn rav1e_cdef_find_dir_16bpc_neon(
    tmp: *const u16, tmp_stride: isize, var: *mut u32,
  ) -> i32;
}

cpu_function_lookup_table!(
  CDEF_DIR_LBD_FNS: [Option<CdefDirLBDFn>],
  default: None,
  [(NEON, Some(rav1e_cdef_find_dir_8bpc_neon))]
);

cpu_function_lookup_table!(
  CDEF_DIR_HBD_FNS: [Option<CdefDirHBDFn>],
  default: None,
  [(NEON, Some(rav1e_cdef_find_dir_16bpc_neon))]
);

<|endoftext|>


<commit_before>extern crate chrono;
extern crate clap;
extern crate fern;
#[macro_use] extern crate log;
extern crate rand;
extern crate thud;

use clap::App;
use thud::game;
use thud::mcts;
use thud::console_ui;
use thud::util;

fn main() {
    // Set up arg handling.
    let matches = {
        let app = util::set_common_args(
            App::new("console_play")
                .version("0.1.0")
                .author("Stu Black <trurl@freeshell.org>")
                .about("Play against Thud AI"),
            &[util::ITERATION_COUNT_FLAG,
              util::SIMULATION_COUNT_FLAG,
              util::EXPLORATION_BIAS_FLAG,
              util::INITIAL_BOARD_FLAG,
              util::AI_PLAYER_FLAG,
              util::LOG_LEVEL_FLAG]);
        app.get_matches()
    };
    let iteration_count =
        match matches.value_of(util::ITERATION_COUNT_FLAG).unwrap().parse::<usize>() {
            Ok(x) => x,
            Err(e) => panic!("Bad iteration count: {}", e),
        };
    let simulation_count =
        match matches.value_of(util::SIMULATION_COUNT_FLAG).unwrap().parse::<usize>() {
            Ok(x) => x,
            Err(e) => panic!("Bad simulation count: {}", e),
        };
    let exploration_bias =
        match matches.value_of(util::EXPLORATION_BIAS_FLAG).unwrap().parse::<f64>() {
            Ok(x) => x,
            Err(e) => panic!("Bad exploration bias: {}", e),
        };
    let initial_cells =
        match matches.value_of(util::INITIAL_BOARD_FLAG).map(|x| x.parse::<util::InitialBoard>()) {
            None => game::board::Cells::default(),
            Some(Ok(x)) => x.cells(),
            Some(Err(e)) => panic!("Bad initial board configuration: {}", e),
        };
    let ai_role = match matches.value_of(util::INITIAL_PLAYER_FLAG).map(|x| x.parse::<game::Role>()) {
        None | Some(Ok(game::Role::Dwarf)) => game::Role::Dwarf,
        Some(Ok(game::Role::Troll)) => game::Role::Troll,
        Some(Err(x)) => panic!("{}", x),
    };
    let human_role = ai_role.toggle();
    let logging_level =
        match matches.value_of(::thud::util::LOG_LEVEL_FLAG).map(|x| x.parse::<log::LogLevelFilter>()) {
            Some(Ok(x)) => x,
            Some(Err(_)) => panic!("Bad logging level '{}'", matches.value_of(::thud::util::LOG_LEVEL_FLAG).unwrap()),
            None => log::LogLevelFilter::Info,
        };

    // Set up logging.
    let logger_config = fern::DispatchConfig {
        format: Box::new(|msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
            format!("[{}][{}] {}",
                    chrono::Local::now().format("%Y-%m-%d %T%.3f%z").to_string(), level, msg)
        }),
        output: vec![fern::OutputConfig::stdout()],
        level: logging_level,
    };
    if let Err(e) = fern::init_global_logger(logger_config, log::LogLevelFilter::Trace) {
        panic!("Failed to initialize global logger: {}", e);
    }

    let mut state = mcts::State::new(initial_cells);
    let mut graph = mcts::Graph::new();
    let mut search_state = mcts::SearchState::new(rand::thread_rng(), exploration_bias);
    loop {
        console_ui::write_board(state.board());
        if state.active_role() == ai_role {
            println!("{:?} player's turn. Thinking...", ai_role);
            if graph.get_node(&state).is_none() {
                util::initialize_search(state.clone(), &mut graph);
                info!("Manually added current game state to graph");
            }
            let mut best_action = None;
            for iteration in 0..iteration_count {
                if iteration % 100 == 0 {
                    info!("iteration: {} / {} = {:.02}%", iteration, iteration_count,
                          ((100 * iteration) as f64) / (iteration_count as f64));
                }
                match search_state.search(
                    &mut graph, state.clone(),
                    |_: usize| mcts::SearchSettings {
                        simulation_count: simulation_count,
                    }) {
                    Ok(stats) => {
                        if iteration % 1000 == 0 || iteration + 1 == iteration_count {
                            info!("root stats:");
                            let mut best_visits = ::std::usize::MIN;
                            for (action, stats, ucb) in stats.into_iter() {
                                info!("{:?}: [{}, {}] = {:?} / {}; UCB = {:?}", action,
                                      (stats.payoff.values[0] as f64) / (stats.visits as f64),
                                      (stats.payoff.values[1] as f64) / (stats.visits as f64),
                                      stats.payoff, stats.visits, ucb);
                                best_action = match best_action {
                                    None => Some(action),
                                    Some(_) if best_visits < stats.visits => {
                                        best_visits = stats.visits;
                                        Some(action)
                                    },
                                    _ => best_action,
                                }
                            }
                        }
                    },
                    Err(e) => {
                        error!("Error in search: {}", e);
                        break
                    }
                }
            }
            match best_action {
                Some(a) => {
                    info!("best action: {:?}", a);
                    state.do_action(&a);
                    graph.retain_reachable_from(&[&state]);
                },
                None => {
                    info!("No best action. Exiting.");
                    return
                }
            }
        } else {
            // Prompt for play.
            loop {
                println!("{:?} player's turn. Enter coordinate of piece to move.", human_role);
                let c = console_ui::prompt_for_piece(state.board(), human_role);
                let piece_actions: Vec<game::Action> = state.position_actions(c).collect();
                if piece_actions.is_empty() {
                    println!("Piece at {:?} has no actions.", c);
                } else {
                    if let Some(action) = console_ui::select_one(&piece_actions) {
                        let mut moved_state = state.clone();
                        moved_state.do_action(&action);
                        println!("After action, board: {}",
                                 game::board::format_board(moved_state.board()));
                        println!("Is this okay?");
                        match console_ui::select_one(&["y", "n"]) {
                            Some(&"y") => {
                                state = moved_state;
                                graph.retain_reachable_from(&[&state]);
                                break
                            },
                            _ => continue,
                        }
                    }
                }
            }
        }
    }
}
<commit_msg>Fix command line flags for console_play.<commit_after>extern crate chrono;
extern crate clap;
extern crate fern;
#[macro_use] extern crate log;
extern crate rand;
extern crate thud;

use clap::App;
use thud::game;
use thud::mcts;
use thud::console_ui;
use thud::util;

fn main() {
    // Set up arg handling.
    let matches = {
        let app = util::set_common_args(
            App::new("console_play")
                .version("0.1.0")
                .author("Stu Black <trurl@freeshell.org>")
                .about("Play against Thud AI"),
            &[util::ITERATION_COUNT_FLAG,
              util::SIMULATION_COUNT_FLAG,
              util::EXPLORATION_BIAS_FLAG,
              util::INITIAL_BOARD_FLAG,
              util::AI_PLAYER_FLAG,
              util::LOG_LEVEL_FLAG]);
        app.get_matches()
    };
    let iteration_count =
        match matches.value_of(util::ITERATION_COUNT_FLAG).unwrap().parse::<usize>() {
            Ok(x) => x,
            Err(e) => panic!("Bad iteration count: {}", e),
        };
    let simulation_count =
        match matches.value_of(util::SIMULATION_COUNT_FLAG).unwrap().parse::<usize>() {
            Ok(x) => x,
            Err(e) => panic!("Bad simulation count: {}", e),
        };
    let exploration_bias =
        match matches.value_of(util::EXPLORATION_BIAS_FLAG).unwrap().parse::<f64>() {
            Ok(x) => x,
            Err(e) => panic!("Bad exploration bias: {}", e),
        };
    let initial_cells =
        match matches.value_of(util::INITIAL_BOARD_FLAG).map(|x| x.parse::<util::InitialBoard>()) {
            None => game::board::Cells::default(),
            Some(Ok(x)) => x.cells(),
            Some(Err(e)) => panic!("Bad initial board configuration: {}", e),
        };
    let ai_role = match matches.value_of(util::AI_PLAYER_FLAG).map(|x| x.parse::<game::Role>()) {
        None | Some(Ok(game::Role::Dwarf)) => game::Role::Dwarf,
        Some(Ok(game::Role::Troll)) => game::Role::Troll,
        Some(Err(x)) => panic!("{}", x),
    };
    let human_role = ai_role.toggle();
    let logging_level =
        match matches.value_of(::thud::util::LOG_LEVEL_FLAG).map(|x| x.parse::<log::LogLevelFilter>()) {
            Some(Ok(x)) => x,
            Some(Err(_)) => panic!("Bad logging level '{}'", matches.value_of(::thud::util::LOG_LEVEL_FLAG).unwrap()),
            None => log::LogLevelFilter::Info,
        };

    // Set up logging.
    let logger_config = fern::DispatchConfig {
        format: Box::new(|msg: &str, level: &log::LogLevel, _location: &log::LogLocation| {
            format!("[{}][{}] {}",
                    chrono::Local::now().format("%Y-%m-%d %T%.3f%z").to_string(), level, msg)
        }),
        output: vec![fern::OutputConfig::stdout()],
        level: logging_level,
    };
    if let Err(e) = fern::init_global_logger(logger_config, log::LogLevelFilter::Trace) {
        panic!("Failed to initialize global logger: {}", e);
    }

    let mut state = mcts::State::new(initial_cells);
    let mut graph = mcts::Graph::new();
    let mut search_state = mcts::SearchState::new(rand::thread_rng(), exploration_bias);
    loop {
        console_ui::write_board(state.board());
        if state.active_role() == ai_role {
            println!("{:?} player's turn. Thinking...", ai_role);
            if graph.get_node(&state).is_none() {
                util::initialize_search(state.clone(), &mut graph);
                info!("Manually added current game state to graph");
            }
            let mut best_action = None;
            for iteration in 0..iteration_count {
                if iteration % 100 == 0 {
                    info!("iteration: {} / {} = {:.02}%", iteration, iteration_count,
                          ((100 * iteration) as f64) / (iteration_count as f64));
                }
                match search_state.search(
                    &mut graph, state.clone(),
                    |_: usize| mcts::SearchSettings {
                        simulation_count: simulation_count,
                    }) {
                    Ok(stats) => {
                        if iteration % 1000 == 0 || iteration + 1 == iteration_count {
                            info!("root stats:");
                            let mut best_visits = ::std::usize::MIN;
                            for (action, stats, ucb) in stats.into_iter() {
                                info!("{:?}: [{}, {}] = {:?} / {}; UCB = {:?}", action,
                                      (stats.payoff.values[0] as f64) / (stats.visits as f64),
                                      (stats.payoff.values[1] as f64) / (stats.visits as f64),
                                      stats.payoff, stats.visits, ucb);
                                best_action = match best_action {
                                    None => Some(action),
                                    Some(_) if best_visits < stats.visits => {
                                        best_visits = stats.visits;
                                        Some(action)
                                    },
                                    _ => best_action,
                                }
                            }
                        }
                    },
                    Err(e) => {
                        error!("Error in search: {}", e);
                        break
                    }
                }
            }
            match best_action {
                Some(a) => {
                    info!("best action: {:?}", a);
                    state.do_action(&a);
                    graph.retain_reachable_from(&[&state]);
                },
                None => {
                    info!("No best action. Exiting.");
                    return
                }
            }
        } else {
            // Prompt for play.
            loop {
                println!("{:?} player's turn. Enter coordinate of piece to move.", human_role);
                let c = console_ui::prompt_for_piece(state.board(), human_role);
                let piece_actions: Vec<game::Action> = state.position_actions(c).collect();
                if piece_actions.is_empty() {
                    println!("Piece at {:?} has no actions.", c);
                } else {
                    if let Some(action) = console_ui::select_one(&piece_actions) {
                        let mut moved_state = state.clone();
                        moved_state.do_action(&action);
                        println!("After action, board: {}",
                                 game::board::format_board(moved_state.board()));
                        println!("Is this okay?");
                        match console_ui::select_one(&["y", "n"]) {
                            Some(&"y") => {
                                state = moved_state;
                                graph.retain_reachable_from(&[&state]);
                                break
                            },
                            _ => continue,
                        }
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::ffi;

use sass_sys::{self, Sass_Output_Style};

use bindings::ptr::Unique;

use options::OutputStyle;

/// The internal options we will pass to libsass
#[derive(Debug)]
pub struct SassOptions {
    pub raw: Unique<sass_sys::Sass_Options>
}

impl SassOptions {
    pub fn set_output_style(&mut self, style: OutputStyle) {
        let style = match style {
            OutputStyle::Nested => Sass_Output_Style::SASS_STYLE_NESTED,
            OutputStyle::Expanded => Sass_Output_Style::SASS_STYLE_EXPANDED,
            OutputStyle::Compact => Sass_Output_Style::SASS_STYLE_COMPACT,
            OutputStyle::Compressed => Sass_Output_Style::SASS_STYLE_COMPRESSED,
        };

        unsafe {
            sass_sys::sass_option_set_output_style(self.raw.get_mut(), style);
        }
    }

    pub fn set_precision(&mut self, precision: usize) {
        unsafe {
            sass_sys::sass_option_set_precision(self.raw.get_mut(), precision as i32);
        }
    }

    pub fn set_is_indented_syntax(&mut self) {
        unsafe {
            sass_sys::sass_option_set_is_indented_syntax_src(self.raw.get_mut(), true);
        }
    }

    pub fn set_include_path(&mut self, paths: Vec<String>) {
        let include_path = if cfg!(windows) {
            paths.join(";")
        } else {
            paths.join(",")
        };
        let c_str = ffi::CString::new(include_path).unwrap();
        let ptr = c_str.into_raw();
        unsafe {
            sass_sys::sass_option_set_include_path(self.raw.get_mut(), ptr);
            let _ = ffi::CString::from_raw(ptr);
        }
    }
}
<commit_msg>Join include_paths with colon on Unix<commit_after>use std::ffi;

use sass_sys::{self, Sass_Output_Style};

use bindings::ptr::Unique;

use options::OutputStyle;

/// The internal options we will pass to libsass
#[derive(Debug)]
pub struct SassOptions {
    pub raw: Unique<sass_sys::Sass_Options>
}

impl SassOptions {
    pub fn set_output_style(&mut self, style: OutputStyle) {
        let style = match style {
            OutputStyle::Nested => Sass_Output_Style::SASS_STYLE_NESTED,
            OutputStyle::Expanded => Sass_Output_Style::SASS_STYLE_EXPANDED,
            OutputStyle::Compact => Sass_Output_Style::SASS_STYLE_COMPACT,
            OutputStyle::Compressed => Sass_Output_Style::SASS_STYLE_COMPRESSED,
        };

        unsafe {
            sass_sys::sass_option_set_output_style(self.raw.get_mut(), style);
        }
    }

    pub fn set_precision(&mut self, precision: usize) {
        unsafe {
            sass_sys::sass_option_set_precision(self.raw.get_mut(), precision as i32);
        }
    }

    pub fn set_is_indented_syntax(&mut self) {
        unsafe {
            sass_sys::sass_option_set_is_indented_syntax_src(self.raw.get_mut(), true);
        }
    }

    pub fn set_include_path(&mut self, paths: Vec<String>) {
        let include_path = if cfg!(windows) {
            paths.join(";")
        } else {
            paths.join(":")
        };
        let c_str = ffi::CString::new(include_path).unwrap();
        let ptr = c_str.into_raw();
        unsafe {
            sass_sys::sass_option_set_include_path(self.raw.get_mut(), ptr);
            let _ = ffi::CString::from_raw(ptr);
        }
    }
}
<|endoftext|>


<commit_before>#![allow(deprecated)] // for SipHasher

use std::collections::hash_map::{Entry, HashMap};
use std::env;
use std::ffi::{OsStr, OsString};
use std::hash::{Hash, Hasher, SipHasher};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Mutex;

use log::{debug, info, warn};
use serde::{Deserialize, Serialize};

use crate::util::paths;
use crate::util::{self, internal, profile, CargoResult, ProcessBuilder};

/// Information on the `rustc` executable
#[derive(Debug)]
pub struct Rustc {
    /// The location of the exe
    pub path: PathBuf,
    /// An optional program that will be passed the path of the rust exe as its first argument, and
    /// rustc args following this.
    pub wrapper: Option<RustcWrapper>,
    /// Verbose version information (the output of `rustc -vV`)
    pub verbose_version: String,
    /// The host triple (arch-platform-OS), this comes from verbose_version.
    pub host: String,
    cache: Mutex<Cache>,
}

/// Information on the `rustc` wrapper
#[derive(Debug)]
pub struct RustcWrapper {
    path: PathBuf,
    env: HashMap<String, Option<OsString>>,
}

impl RustcWrapper {
    pub fn new<T: Into<PathBuf>>(path: T) -> RustcWrapper {
        RustcWrapper {
            path: path.into(),
            env: HashMap::new(),
        }
    }

    /// (chainable) Sets an environment variable for the process.
    pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut RustcWrapper {
        self.env
            .insert(key.to_string(), Some(val.as_ref().to_os_string()));
        self
    }

    /// (chainable) Unsets an environment variable for the process.
    pub fn env_remove(&mut self, key: &str) -> &mut RustcWrapper {
        self.env.insert(key.to_string(), None);
        self
    }

    pub fn process(&self) -> ProcessBuilder {
        let mut cmd = util::process(&self.path);

        for (k, v) in &self.env {
            match *v {
                Some(ref v) => {
                    cmd.env(k, v);
                }
                None => {
                    cmd.env_remove(k);
                }
            }
        }

        cmd
    }

    pub fn is_empty(&self) -> bool {
        self.path.as_os_str().is_empty()
    }
}

impl Rustc {
    /// Runs the compiler at `path` to learn various pieces of information about
    /// it, with an optional wrapper.
    ///
    /// If successful this function returns a description of the compiler along
    /// with a list of its capabilities.
    pub fn new(
        path: PathBuf,
        wrapper: Option<PathBuf>,
        rustup_rustc: &Path,
        cache_location: Option<PathBuf>,
    ) -> CargoResult<Rustc> {
        let _p = profile::start("Rustc::new");

        let mut cache = Cache::load(&path, rustup_rustc, cache_location);

        let mut cmd = util::process(&path);
        cmd.arg("-vV");
        let verbose_version = cache.cached_output(&cmd)?.0;

        let host = {
            let triple = verbose_version
                .lines()
                .find(|l| l.starts_with("host: "))
                .map(|l| &l[6..])
                .ok_or_else(|| internal("rustc -v didn't have a line for `host:`"))?;
            triple.to_string()
        };

        Ok(Rustc {
            path,
            wrapper: wrapper.map(|w| RustcWrapper::new(w)),
            verbose_version,
            host,
            cache: Mutex::new(cache),
        })
    }

    /// Gets a process builder set up to use the found rustc version, with a wrapper if `Some`.
    pub fn process(&self) -> ProcessBuilder {
        match self.wrapper {
            Some(ref wrapper) if !wrapper.is_empty() => {
                let mut cmd = wrapper.process();
                cmd.arg(&self.path);
                cmd
            }
            _ => self.process_no_wrapper(),
        }
    }

    pub fn process_no_wrapper(&self) -> ProcessBuilder {
        util::process(&self.path)
    }

    pub fn cached_output(&self, cmd: &ProcessBuilder) -> CargoResult<(String, String)> {
        self.cache.lock().unwrap().cached_output(cmd)
    }

    pub fn cached_success(&self, cmd: &ProcessBuilder) -> CargoResult<bool> {
        self.cache.lock().unwrap().cached_success(cmd)
    }

    pub fn push_wrapper(&mut self, wrapper: RustcWrapper) {
        self.wrapper = Some(wrapper);
    }
}

/// It is a well known that `rustc` is not the fastest compiler in the world.
/// What is less known is that even `rustc --version --verbose` takes about a
/// hundred milliseconds! Because we need compiler version info even for no-op
/// builds, we cache it here, based on compiler's mtime and rustup's current
/// toolchain.
///
/// https://github.com/rust-lang/cargo/issues/5315
/// https://github.com/rust-lang/rust/issues/49761
#[derive(Debug)]
struct Cache {
    cache_location: Option<PathBuf>,
    dirty: bool,
    data: CacheData,
}

#[derive(Serialize, Deserialize, Debug, Default)]
struct CacheData {
    rustc_fingerprint: u64,
    outputs: HashMap<u64, (String, String)>,
    successes: HashMap<u64, bool>,
}

impl Cache {
    fn load(rustc: &Path, rustup_rustc: &Path, cache_location: Option<PathBuf>) -> Cache {
        match (cache_location, rustc_fingerprint(rustc, rustup_rustc)) {
            (Some(cache_location), Ok(rustc_fingerprint)) => {
                let empty = CacheData {
                    rustc_fingerprint,
                    outputs: HashMap::new(),
                    successes: HashMap::new(),
                };
                let mut dirty = true;
                let data = match read(&cache_location) {
                    Ok(data) => {
                        if data.rustc_fingerprint == rustc_fingerprint {
                            info!("reusing existing rustc info cache");
                            dirty = false;
                            data
                        } else {
                            info!("different compiler, creating new rustc info cache");
                            empty
                        }
                    }
                    Err(e) => {
                        info!("failed to read rustc info cache: {}", e);
                        empty
                    }
                };
                return Cache {
                    cache_location: Some(cache_location),
                    dirty,
                    data,
                };

                fn read(path: &Path) -> CargoResult<CacheData> {
                    let json = paths::read(path)?;
                    Ok(serde_json::from_str(&json)?)
                }
            }
            (_, fingerprint) => {
                if let Err(e) = fingerprint {
                    warn!("failed to calculate rustc fingerprint: {}", e);
                }
                info!("rustc info cache disabled");
                Cache {
                    cache_location: None,
                    dirty: false,
                    data: CacheData::default(),
                }
            }
        }
    }

    fn cached_output(&mut self, cmd: &ProcessBuilder) -> CargoResult<(String, String)> {
        let key = process_fingerprint(cmd);
        match self.data.outputs.entry(key) {
            Entry::Occupied(entry) => {
                info!("rustc info cache hit");
                Ok(entry.get().clone())
            }
            Entry::Vacant(entry) => {
                info!("rustc info cache miss");
                let output = cmd.exec_with_output()?;
                let stdout = String::from_utf8(output.stdout)
                    .map_err(|_| internal("rustc didn't return utf8 output"))?;
                let stderr = String::from_utf8(output.stderr)
                    .map_err(|_| internal("rustc didn't return utf8 output"))?;
                let output = (stdout, stderr);
                entry.insert(output.clone());
                self.dirty = true;
                Ok(output)
            }
        }
    }

    fn cached_success(&mut self, cmd: &ProcessBuilder) -> CargoResult<bool> {
        let key = process_fingerprint(cmd);
        match self.data.successes.entry(key) {
            Entry::Occupied(entry) => {
                info!("rustc info cache hit");
                Ok(*entry.get())
            }
            Entry::Vacant(entry) => {
                info!("rustc info cache miss");
                let success = cmd
                    .build_command()
                    .stdout(Stdio::null())
                    .stderr(Stdio::null())
                    .status()?
                    .success();
                entry.insert(success);
                self.dirty = true;
                Ok(success)
            }
        }
    }
}

impl Drop for Cache {
    fn drop(&mut self) {
        if !self.dirty {
            return;
        }
        if let Some(ref path) = self.cache_location {
            let json = serde_json::to_string(&self.data).unwrap();
            match paths::write(path, json.as_bytes()) {
                Ok(()) => info!("updated rustc info cache"),
                Err(e) => warn!("failed to update rustc info cache: {}", e),
            }
        }
    }
}

fn rustc_fingerprint(path: &Path, rustup_rustc: &Path) -> CargoResult<u64> {
    let mut hasher = SipHasher::new_with_keys(0, 0);

    let path = paths::resolve_executable(path)?;
    path.hash(&mut hasher);

    paths::mtime(&path)?.hash(&mut hasher);

    // Rustup can change the effective compiler without touching
    // the `rustc` binary, so we try to account for this here.
    // If we see rustup's env vars, we mix them into the fingerprint,
    // but we also mix in the mtime of the actual compiler (and not
    // the rustup shim at `~/.cargo/bin/rustup`), because `RUSTUP_TOOLCHAIN`
    // could be just `stable-x86_64-unknown-linux-gnu`, i.e, it could
    // not mention the version of Rust at all, which changes after
    // `rustup update`.
    //
    // If we don't see rustup env vars, but it looks like the compiler
    // is managed by rustup, we conservatively bail out.
    let maybe_rustup = rustup_rustc == path;
    match (
        maybe_rustup,
        env::var("RUSTUP_HOME"),
        env::var("RUSTUP_TOOLCHAIN"),
    ) {
        (_, Ok(rustup_home), Ok(rustup_toolchain)) => {
            debug!("adding rustup info to rustc fingerprint");
            rustup_toolchain.hash(&mut hasher);
            rustup_home.hash(&mut hasher);
            let real_rustc = Path::new(&rustup_home)
                .join("toolchains")
                .join(rustup_toolchain)
                .join("bin")
                .join("rustc")
                .with_extension(env::consts::EXE_EXTENSION);
            paths::mtime(&real_rustc)?.hash(&mut hasher);
        }
        (true, _, _) => failure::bail!("probably rustup rustc, but without rustup's env vars"),
        _ => (),
    }

    Ok(hasher.finish())
}

fn process_fingerprint(cmd: &ProcessBuilder) -> u64 {
    let mut hasher = SipHasher::new_with_keys(0, 0);
    cmd.get_args().hash(&mut hasher);
    let mut env = cmd.get_envs().iter().collect::<Vec<_>>();
    env.sort_unstable();
    env.hash(&mut hasher);
    hasher.finish()
}
<commit_msg>totally unnecessary template type because its fun<commit_after>#![allow(deprecated)] // for SipHasher

use std::collections::hash_map::{Entry, HashMap};
use std::env;
use std::ffi::{OsStr, OsString};
use std::hash::{Hash, Hasher, SipHasher};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Mutex;

use log::{debug, info, warn};
use serde::{Deserialize, Serialize};

use crate::util::paths;
use crate::util::{self, internal, profile, CargoResult, ProcessBuilder};

/// Information on the `rustc` executable
#[derive(Debug)]
pub struct Rustc {
    /// The location of the exe
    pub path: PathBuf,
    /// An optional program that will be passed the path of the rust exe as its first argument, and
    /// rustc args following this.
    pub wrapper: Option<RustcWrapper>,
    /// Verbose version information (the output of `rustc -vV`)
    pub verbose_version: String,
    /// The host triple (arch-platform-OS), this comes from verbose_version.
    pub host: String,
    cache: Mutex<Cache>,
}

/// Information on the `rustc` wrapper
#[derive(Debug)]
pub struct RustcWrapper {
    path: PathBuf,
    env: HashMap<String, Option<OsString>>,
}

impl RustcWrapper {
    pub fn new<T: Into<PathBuf>>(path: T) -> RustcWrapper {
        RustcWrapper {
            path: path.into(),
            env: HashMap::new(),
        }
    }

    /// (chainable) Sets an environment variable for the process.
    pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut RustcWrapper {
        self.env
            .insert(key.to_string(), Some(val.as_ref().to_os_string()));
        self
    }

    /// (chainable) Unsets an environment variable for the process.
    pub fn env_remove(&mut self, key: &str) -> &mut RustcWrapper {
        self.env.insert(key.to_string(), None);
        self
    }

    pub fn process(&self) -> ProcessBuilder {
        let mut cmd = util::process(&self.path);

        for (k, v) in &self.env {
            match *v {
                Some(ref v) => {
                    cmd.env(k, v);
                }
                None => {
                    cmd.env_remove(k);
                }
            }
        }

        cmd
    }

    pub fn is_empty(&self) -> bool {
        self.path.as_os_str().is_empty()
    }
}

impl Rustc {
    /// Runs the compiler at `path` to learn various pieces of information about
    /// it, with an optional wrapper.
    ///
    /// If successful this function returns a description of the compiler along
    /// with a list of its capabilities.
    pub fn new(
        path: PathBuf,
        wrapper: Option<PathBuf>,
        rustup_rustc: &Path,
        cache_location: Option<PathBuf>,
    ) -> CargoResult<Rustc> {
        let _p = profile::start("Rustc::new");

        let mut cache = Cache::load(&path, rustup_rustc, cache_location);

        let mut cmd = util::process(&path);
        cmd.arg("-vV");
        let verbose_version = cache.cached_output(&cmd)?.0;

        let host = {
            let triple = verbose_version
                .lines()
                .find(|l| l.starts_with("host: "))
                .map(|l| &l[6..])
                .ok_or_else(|| internal("rustc -v didn't have a line for `host:`"))?;
            triple.to_string()
        };

        Ok(Rustc {
            path,
            wrapper: wrapper.map(RustcWrapper::new),
            verbose_version,
            host,
            cache: Mutex::new(cache),
        })
    }

    /// Gets a process builder set up to use the found rustc version, with a wrapper if `Some`.
    pub fn process(&self) -> ProcessBuilder {
        match self.wrapper {
            Some(ref wrapper) if !wrapper.is_empty() => {
                let mut cmd = wrapper.process();
                cmd.arg(&self.path);
                cmd
            }
            _ => self.process_no_wrapper(),
        }
    }

    pub fn process_no_wrapper(&self) -> ProcessBuilder {
        util::process(&self.path)
    }

    pub fn cached_output(&self, cmd: &ProcessBuilder) -> CargoResult<(String, String)> {
        self.cache.lock().unwrap().cached_output(cmd)
    }

    pub fn cached_success(&self, cmd: &ProcessBuilder) -> CargoResult<bool> {
        self.cache.lock().unwrap().cached_success(cmd)
    }

    pub fn push_wrapper<T: Into<Option<RustcWrapper>>>(&mut self, wrapper: T) {
        self.wrapper = wrapper.into();
    }
}

/// It is a well known that `rustc` is not the fastest compiler in the world.
/// What is less known is that even `rustc --version --verbose` takes about a
/// hundred milliseconds! Because we need compiler version info even for no-op
/// builds, we cache it here, based on compiler's mtime and rustup's current
/// toolchain.
///
/// https://github.com/rust-lang/cargo/issues/5315
/// https://github.com/rust-lang/rust/issues/49761
#[derive(Debug)]
struct Cache {
    cache_location: Option<PathBuf>,
    dirty: bool,
    data: CacheData,
}

#[derive(Serialize, Deserialize, Debug, Default)]
struct CacheData {
    rustc_fingerprint: u64,
    outputs: HashMap<u64, (String, String)>,
    successes: HashMap<u64, bool>,
}

impl Cache {
    fn load(rustc: &Path, rustup_rustc: &Path, cache_location: Option<PathBuf>) -> Cache {
        match (cache_location, rustc_fingerprint(rustc, rustup_rustc)) {
            (Some(cache_location), Ok(rustc_fingerprint)) => {
                let empty = CacheData {
                    rustc_fingerprint,
                    outputs: HashMap::new(),
                    successes: HashMap::new(),
                };
                let mut dirty = true;
                let data = match read(&cache_location) {
                    Ok(data) => {
                        if data.rustc_fingerprint == rustc_fingerprint {
                            info!("reusing existing rustc info cache");
                            dirty = false;
                            data
                        } else {
                            info!("different compiler, creating new rustc info cache");
                            empty
                        }
                    }
                    Err(e) => {
                        info!("failed to read rustc info cache: {}", e);
                        empty
                    }
                };
                return Cache {
                    cache_location: Some(cache_location),
                    dirty,
                    data,
                };

                fn read(path: &Path) -> CargoResult<CacheData> {
                    let json = paths::read(path)?;
                    Ok(serde_json::from_str(&json)?)
                }
            }
            (_, fingerprint) => {
                if let Err(e) = fingerprint {
                    warn!("failed to calculate rustc fingerprint: {}", e);
                }
                info!("rustc info cache disabled");
                Cache {
                    cache_location: None,
                    dirty: false,
                    data: CacheData::default(),
                }
            }
        }
    }

    fn cached_output(&mut self, cmd: &ProcessBuilder) -> CargoResult<(String, String)> {
        let key = process_fingerprint(cmd);
        match self.data.outputs.entry(key) {
            Entry::Occupied(entry) => {
                info!("rustc info cache hit");
                Ok(entry.get().clone())
            }
            Entry::Vacant(entry) => {
                info!("rustc info cache miss");
                let output = cmd.exec_with_output()?;
                let stdout = String::from_utf8(output.stdout)
                    .map_err(|_| internal("rustc didn't return utf8 output"))?;
                let stderr = String::from_utf8(output.stderr)
                    .map_err(|_| internal("rustc didn't return utf8 output"))?;
                let output = (stdout, stderr);
                entry.insert(output.clone());
                self.dirty = true;
                Ok(output)
            }
        }
    }

    fn cached_success(&mut self, cmd: &ProcessBuilder) -> CargoResult<bool> {
        let key = process_fingerprint(cmd);
        match self.data.successes.entry(key) {
            Entry::Occupied(entry) => {
                info!("rustc info cache hit");
                Ok(*entry.get())
            }
            Entry::Vacant(entry) => {
                info!("rustc info cache miss");
                let success = cmd
                    .build_command()
                    .stdout(Stdio::null())
                    .stderr(Stdio::null())
                    .status()?
                    .success();
                entry.insert(success);
                self.dirty = true;
                Ok(success)
            }
        }
    }
}

impl Drop for Cache {
    fn drop(&mut self) {
        if !self.dirty {
            return;
        }
        if let Some(ref path) = self.cache_location {
            let json = serde_json::to_string(&self.data).unwrap();
            match paths::write(path, json.as_bytes()) {
                Ok(()) => info!("updated rustc info cache"),
                Err(e) => warn!("failed to update rustc info cache: {}", e),
            }
        }
    }
}

fn rustc_fingerprint(path: &Path, rustup_rustc: &Path) -> CargoResult<u64> {
    let mut hasher = SipHasher::new_with_keys(0, 0);

    let path = paths::resolve_executable(path)?;
    path.hash(&mut hasher);

    paths::mtime(&path)?.hash(&mut hasher);

    // Rustup can change the effective compiler without touching
    // the `rustc` binary, so we try to account for this here.
    // If we see rustup's env vars, we mix them into the fingerprint,
    // but we also mix in the mtime of the actual compiler (and not
    // the rustup shim at `~/.cargo/bin/rustup`), because `RUSTUP_TOOLCHAIN`
    // could be just `stable-x86_64-unknown-linux-gnu`, i.e, it could
    // not mention the version of Rust at all, which changes after
    // `rustup update`.
    //
    // If we don't see rustup env vars, but it looks like the compiler
    // is managed by rustup, we conservatively bail out.
    let maybe_rustup = rustup_rustc == path;
    match (
        maybe_rustup,
        env::var("RUSTUP_HOME"),
        env::var("RUSTUP_TOOLCHAIN"),
    ) {
        (_, Ok(rustup_home), Ok(rustup_toolchain)) => {
            debug!("adding rustup info to rustc fingerprint");
            rustup_toolchain.hash(&mut hasher);
            rustup_home.hash(&mut hasher);
            let real_rustc = Path::new(&rustup_home)
                .join("toolchains")
                .join(rustup_toolchain)
                .join("bin")
                .join("rustc")
                .with_extension(env::consts::EXE_EXTENSION);
            paths::mtime(&real_rustc)?.hash(&mut hasher);
        }
        (true, _, _) => failure::bail!("probably rustup rustc, but without rustup's env vars"),
        _ => (),
    }

    Ok(hasher.finish())
}

fn process_fingerprint(cmd: &ProcessBuilder) -> u64 {
    let mut hasher = SipHasher::new_with_keys(0, 0);
    cmd.get_args().hash(&mut hasher);
    let mut env = cmd.get_envs().iter().collect::<Vec<_>>();
    env.sort_unstable();
    env.hash(&mut hasher);
    hasher.finish()
}
<|endoftext|>


<commit_before>extern crate sorting;

/// Verify order of input vector as sorted by indices.
///
/// # Details
///
/// This function takes as arguments a vector of data and a vector representing the indices of the
/// first vector making it sorted.
///
/// The elements of the sorted vector are compared to the next ones to make sure they are less
/// or equal.
///
/// # Panics
///
/// The sorting of the data vector is enforced using `assert!()`.
///
/// # Examples
///
/// ```
/// let data: Vec<i32> = vec![4, 2, 3, 1];
/// let sorted_indices: Vec<usize> = vec![3, 1, 2, 0];
/// verify_sorting(&data, &sorted_indices);
/// ```
///
fn verify_sorting<T: PartialOrd>(array: &Vec<T>, sorted_indices: &Vec<usize>) {
    let n = array.len();
    assert_eq!(n , sorted_indices.len());

    if !array.is_empty() {
        for i in 0..n-1 {
            assert!(array[sorted_indices[i]] <= array[sorted_indices[i+1]]);
        }
    }
}

/// Call the specified sorting function on an empty vector of specific type.
///
/// # Details
///
/// This function takes as arguments a pointer to a sorting function acting and call
/// it on an empty vector of type T. The call should succeed and the index vector should
/// be empty as well.
///
/// # Panics
///
/// The emptiness of the sorted list is enforced using `assert!()`.
///
/// # Examples
///
/// ```
/// test_empty_vec::<i8>(sorting::simplesorts::insertion::sort);
/// ```
///
fn test_empty_vec<T: PartialOrd>(sorting_fct: fn(&Vec<T>) -> Vec<usize>) {
    let to_sort: Vec<T> = vec![];
    let sorted_indices = sorting_fct(&to_sort);
    assert_eq!(sorted_indices, vec![]);
    assert!(sorted_indices.is_empty());
}

fn test_sort_vec_float<T: PartialOrd>(to_sort: &Vec<T>, expected_indices: &Vec<usize>,
                                      sorting_fct: fn(&Vec<T>) -> Vec<usize>) {
    let sorted_indices = sorting_fct(&to_sort);

    verify_sorting(&to_sort, &sorted_indices);

    assert_eq!(sorted_indices.len(), expected_indices.len());

    for i in 0..to_sort.len() {
        assert_eq!(sorted_indices[i], expected_indices[i]);
    }
}

/// Validate against empty vector (i8).
#[test]
fn simple_insertion_empty_vec_i8() {
    test_empty_vec::<i8>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i32).
#[test]
fn simple_insertion_empty_vec_i16() {
    test_empty_vec::<i16>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i32).
#[test]
fn simple_insertion_empty_vec_i32() {
    test_empty_vec::<i32>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i64).
#[test]
fn simple_insertion_empty_vec_i64() {
    test_empty_vec::<i64>(sorting::simplesorts::insertion::sort);
}

/// Validate against empty vector (u8).
#[test]
fn simple_insertion_empty_vec_u8() {
    test_empty_vec::<u8>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u32).
#[test]
fn simple_insertion_empty_vec_u16() {
    test_empty_vec::<u16>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u32).
#[test]
fn simple_insertion_empty_vec_u32() {
    test_empty_vec::<u32>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u64).
#[test]
fn simple_insertion_empty_vec_u64() {
    test_empty_vec::<u64>(sorting::simplesorts::insertion::sort);
}

/// Validate sorting a vector of single precision values (f32).
#[test]
fn simple_insertion_vec_f32() {
    let to_sort: Vec<f32> = vec![6.0,   5.0,  3.0,  1.0,  2.4, 4.0, 10.0, 7.0,
                                 3.42, 32.2, 44.2, 56.3, 67.9, 3.2, 44.2, 2.0];

    let known_sorted_indices = vec![3, 15, 4, 2, 13, 8, 5, 1, 0, 7, 6, 9, 10, 14, 11, 12];

    test_sort_vec_float::<f32>(&to_sort, &known_sorted_indices,
                               sorting::simplesorts::insertion::sort);
}

/// Validate sorting a vector of double precision values (f64).
#[test]
fn simple_insertion_vec_f64() {
    let to_sort: Vec<f64> = vec![6.0,   5.0,  3.0,  1.0,  2.4, 4.0, 10.0, 7.0,
                                 3.42, 32.2, 44.2, 56.3, 67.9, 3.2, 44.2, 2.0];

    let known_sorted_indices = vec![3, 15, 4, 2, 13, 8, 5, 1, 0, 7, 6, 9, 10, 14, 11, 12];

    test_sort_vec_float::<f64>(&to_sort, &known_sorted_indices,
                               sorting::simplesorts::insertion::sort);
}
<commit_msg>Docs comments<commit_after>extern crate sorting;

/// Verify order of input vector as sorted by indices.
///
/// # Details
///
/// This function takes as arguments a vector of data and a vector representing the indices of the
/// first vector making it sorted.
///
/// The elements of the sorted vector are compared to the next ones to make sure they are less
/// or equal.
///
/// # Panics
///
/// The sorting of the data vector is enforced using `assert!()`.
///
/// # Examples
///
/// ```
/// let data: Vec<i32> = vec![4, 2, 3, 1];
/// let sorted_indices: Vec<usize> = vec![3, 1, 2, 0];
/// verify_sorting(&data, &sorted_indices);
/// ```
///
fn verify_sorting<T: PartialOrd>(array: &Vec<T>, sorted_indices: &Vec<usize>) {
    let n = array.len();
    assert_eq!(n , sorted_indices.len());

    if !array.is_empty() {
        for i in 0..n-1 {
            assert!(array[sorted_indices[i]] <= array[sorted_indices[i+1]]);
        }
    }
}

/// Call the specified sorting function on an empty vector of specific type.
///
/// # Details
///
/// This function takes as arguments a pointer to a sorting function acting and call
/// it on an empty vector of type T. The call should succeed and the index vector should
/// be empty as well.
///
/// # Panics
///
/// The emptiness of the sorted list is enforced using `assert!()`.
///
/// # Examples
///
/// ```
/// test_empty_vec::<i8>(sorting::simplesorts::insertion::sort);
/// ```
///
fn test_empty_vec<T: PartialOrd>(sorting_fct: fn(&Vec<T>) -> Vec<usize>) {
    let to_sort: Vec<T> = vec![];
    let sorted_indices = sorting_fct(&to_sort);
    assert_eq!(sorted_indices, vec![]);
    assert!(sorted_indices.is_empty());
}

/// Call the specified sorting function on an specific vector of floats and validate the output.
///
/// # Details
///
/// This function takes as arguments:
///    1. An immutable vector of floats, to be sorted;
///    2. An immutable vector of indices as the reference to compare with;
///    3. A function pointer to be called on the first vector.
///
/// # Panics
///
/// This calls the function `verify_sorting()` which will assert if the result is not in order.
/// It also compares the resulting indices with expected values using asserts.
///
/// # Examples
///
/// ```
/// let to_sort: Vec<f32> = vec![6.0,   5.0,  3.0,  1.0,  2.4, 4.0, 10.0, 7.0,
///                              3.42, 32.2, 44.2, 56.3, 67.9, 3.2, 44.2, 2.0];
///
/// let known_sorted_indices = vec![3, 15, 4, 2, 13, 8, 5, 1, 0, 7, 6, 9, 10, 14, 11, 12];
///
/// test_sort_vec_float::<f32>(&to_sort, &known_sorted_indices,
///                            sorting::simplesorts::insertion::sort);
/// ```
///
fn test_sort_vec_float<T: PartialOrd>(to_sort: &Vec<T>, expected_indices: &Vec<usize>,
                                      sorting_fct: fn(&Vec<T>) -> Vec<usize>) {
    let sorted_indices = sorting_fct(&to_sort);

    verify_sorting(&to_sort, &sorted_indices);

    assert_eq!(sorted_indices.len(), expected_indices.len());

    for i in 0..to_sort.len() {
        assert_eq!(sorted_indices[i], expected_indices[i]);
    }
}

/// Validate against empty vector (i8).
#[test]
fn simple_insertion_empty_vec_i8() {
    test_empty_vec::<i8>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i32).
#[test]
fn simple_insertion_empty_vec_i16() {
    test_empty_vec::<i16>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i32).
#[test]
fn simple_insertion_empty_vec_i32() {
    test_empty_vec::<i32>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (i64).
#[test]
fn simple_insertion_empty_vec_i64() {
    test_empty_vec::<i64>(sorting::simplesorts::insertion::sort);
}

/// Validate against empty vector (u8).
#[test]
fn simple_insertion_empty_vec_u8() {
    test_empty_vec::<u8>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u32).
#[test]
fn simple_insertion_empty_vec_u16() {
    test_empty_vec::<u16>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u32).
#[test]
fn simple_insertion_empty_vec_u32() {
    test_empty_vec::<u32>(sorting::simplesorts::insertion::sort);
}
/// Validate against empty vector (u64).
#[test]
fn simple_insertion_empty_vec_u64() {
    test_empty_vec::<u64>(sorting::simplesorts::insertion::sort);
}

/// Validate sorting a vector of single precision values (f32).
#[test]
fn simple_insertion_vec_f32() {
    let to_sort: Vec<f32> = vec![6.0,   5.0,  3.0,  1.0,  2.4, 4.0, 10.0, 7.0,
                                 3.42, 32.2, 44.2, 56.3, 67.9, 3.2, 44.2, 2.0];

    let known_sorted_indices = vec![3, 15, 4, 2, 13, 8, 5, 1, 0, 7, 6, 9, 10, 14, 11, 12];

    test_sort_vec_float::<f32>(&to_sort, &known_sorted_indices,
                               sorting::simplesorts::insertion::sort);
}

/// Validate sorting a vector of double precision values (f64).
#[test]
fn simple_insertion_vec_f64() {
    let to_sort: Vec<f64> = vec![6.0,   5.0,  3.0,  1.0,  2.4, 4.0, 10.0, 7.0,
                                 3.42, 32.2, 44.2, 56.3, 67.9, 3.2, 44.2, 2.0];

    let known_sorted_indices = vec![3, 15, 4, 2, 13, 8, 5, 1, 0, 7, 6, 9, 10, 14, 11, 12];

    test_sort_vec_float::<f64>(&to_sort, &known_sorted_indices,
                               sorting::simplesorts::insertion::sort);
}
<|endoftext|>