<commit_before>//! Trait solving using Chalk.
use std::sync::{Arc, Mutex};

use rustc_hash::FxHashSet;
use log::debug;
use chalk_ir::cast::Cast;

use crate::{Crate, Trait, db::HirDatabase, ImplBlock};
use super::{TraitRef, Ty, Canonical};

use self::chalk::{ToChalk, from_chalk};

mod chalk;

pub(crate) type Solver = chalk_solve::Solver;

/// This controls the maximum size of types Chalk considers. If we set this too
/// high, we can run into slow edge cases; if we set it too low, Chalk won't
/// find some solutions.
const CHALK_SOLVER_MAX_SIZE: usize = 4;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a, DB> {
    db: &'a DB,
    krate: Crate,
}

pub(crate) fn solver_query(_db: &impl HirDatabase, _krate: Crate) -> Arc<Mutex<Solver>> {
    // krate parameter is just so we cache a unique solver per crate
    let solver_choice = chalk_solve::SolverChoice::SLG { max_size: CHALK_SOLVER_MAX_SIZE };
    debug!("Creating new solver for crate {:?}", _krate);
    Arc::new(Mutex::new(solver_choice.into_solver()))
}

/// Collects impls for the given trait in the whole dependency tree of `krate`.
pub(crate) fn impls_for_trait_query(
    db: &impl HirDatabase,
    krate: Crate,
    trait_: Trait,
) -> Arc<[ImplBlock]> {
    let mut impls = FxHashSet::default();
    // We call the query recursively here. On the one hand, this means we can
    // reuse results from queries for different crates; on the other hand, this
    // will only ever get called for a few crates near the root of the tree (the
    // ones the user is editing), so this may actually be a waste of memory. I'm
    // doing it like this mainly for simplicity for now.
    for dep in krate.dependencies(db) {
        impls.extend(db.impls_for_trait(dep.krate, trait_).iter());
    }
    let crate_impl_blocks = db.impls_in_crate(krate);
    impls.extend(crate_impl_blocks.lookup_impl_blocks_for_trait(&trait_));
    impls.into_iter().collect::<Vec<_>>().into()
}

fn solve(
    db: &impl HirDatabase,
    krate: Crate,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal>>,
) -> Option<chalk_solve::Solution> {
    let context = ChalkContext { db, krate };
    let solver = db.solver(krate);
    debug!("solve goal: {:?}", goal);
    let solution = solver.lock().unwrap().solve(&context, goal);
    debug!("solve({:?}) => {:?}", goal, solution);
    solution
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
}

/// Check using Chalk whether trait is implemented for given parameters including `Self` type.
pub(crate) fn implements_query(
    db: &impl HirDatabase,
    krate: Crate,
    trait_ref: Canonical<TraitRef>,
) -> Option<Solution> {
    let goal: chalk_ir::Goal = trait_ref.value.to_chalk(db).cast();
    debug!("goal: {:?}", goal);
    let env = chalk_ir::Environment::new();
    let in_env = chalk_ir::InEnvironment::new(&env, goal);
    let parameter = chalk_ir::ParameterKind::Ty(chalk_ir::UniverseIndex::ROOT);
    let canonical =
        chalk_ir::Canonical { value: in_env, binders: vec![parameter; trait_ref.num_vars] };
    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solution_from_chalk(db: &impl HirDatabase, solution: chalk_solve::Solution) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution>| {
        let value = subst
            .value
            .parameters
            .into_iter()
            .map(|p| {
                let ty = match p {
                    chalk_ir::Parameter(chalk_ir::ParameterKind::Ty(ty)) => from_chalk(db, ty),
                    chalk_ir::Parameter(chalk_ir::ParameterKind::Lifetime(_)) => unimplemented!(),
                };
                ty
            })
            .collect();
        let result = Canonical { value, num_vars: subst.binders.len() };
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Vec<Ty>>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}
<commit_msg>profile implements query<commit_after>//! Trait solving using Chalk.
use std::sync::{Arc, Mutex};

use rustc_hash::FxHashSet;
use log::debug;
use chalk_ir::cast::Cast;
use ra_prof::profile;

use crate::{Crate, Trait, db::HirDatabase, ImplBlock};
use super::{TraitRef, Ty, Canonical};

use self::chalk::{ToChalk, from_chalk};

mod chalk;

pub(crate) type Solver = chalk_solve::Solver;

/// This controls the maximum size of types Chalk considers. If we set this too
/// high, we can run into slow edge cases; if we set it too low, Chalk won't
/// find some solutions.
const CHALK_SOLVER_MAX_SIZE: usize = 4;

#[derive(Debug, Copy, Clone)]
struct ChalkContext<'a, DB> {
    db: &'a DB,
    krate: Crate,
}

pub(crate) fn solver_query(_db: &impl HirDatabase, _krate: Crate) -> Arc<Mutex<Solver>> {
    // krate parameter is just so we cache a unique solver per crate
    let solver_choice = chalk_solve::SolverChoice::SLG { max_size: CHALK_SOLVER_MAX_SIZE };
    debug!("Creating new solver for crate {:?}", _krate);
    Arc::new(Mutex::new(solver_choice.into_solver()))
}

/// Collects impls for the given trait in the whole dependency tree of `krate`.
pub(crate) fn impls_for_trait_query(
    db: &impl HirDatabase,
    krate: Crate,
    trait_: Trait,
) -> Arc<[ImplBlock]> {
    let mut impls = FxHashSet::default();
    // We call the query recursively here. On the one hand, this means we can
    // reuse results from queries for different crates; on the other hand, this
    // will only ever get called for a few crates near the root of the tree (the
    // ones the user is editing), so this may actually be a waste of memory. I'm
    // doing it like this mainly for simplicity for now.
    for dep in krate.dependencies(db) {
        impls.extend(db.impls_for_trait(dep.krate, trait_).iter());
    }
    let crate_impl_blocks = db.impls_in_crate(krate);
    impls.extend(crate_impl_blocks.lookup_impl_blocks_for_trait(&trait_));
    impls.into_iter().collect::<Vec<_>>().into()
}

fn solve(
    db: &impl HirDatabase,
    krate: Crate,
    goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal>>,
) -> Option<chalk_solve::Solution> {
    let context = ChalkContext { db, krate };
    let solver = db.solver(krate);
    debug!("solve goal: {:?}", goal);
    let solution = solver.lock().unwrap().solve(&context, goal);
    debug!("solve({:?}) => {:?}", goal, solution);
    solution
}

/// Something that needs to be proven (by Chalk) during type checking, e.g. that
/// a certain type implements a certain trait. Proving the Obligation might
/// result in additional information about inference variables.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Obligation {
    /// Prove that a certain type implements a trait (the type is the `Self` type
    /// parameter to the `TraitRef`).
    Trait(TraitRef),
}

/// Check using Chalk whether trait is implemented for given parameters including `Self` type.
pub(crate) fn implements_query(
    db: &impl HirDatabase,
    krate: Crate,
    trait_ref: Canonical<TraitRef>,
) -> Option<Solution> {
    let _p = profile("implements_query");
    let goal: chalk_ir::Goal = trait_ref.value.to_chalk(db).cast();
    debug!("goal: {:?}", goal);
    let env = chalk_ir::Environment::new();
    let in_env = chalk_ir::InEnvironment::new(&env, goal);
    let parameter = chalk_ir::ParameterKind::Ty(chalk_ir::UniverseIndex::ROOT);
    let canonical =
        chalk_ir::Canonical { value: in_env, binders: vec![parameter; trait_ref.num_vars] };
    // We currently don't deal with universes (I think / hope they're not yet
    // relevant for our use cases?)
    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
    let solution = solve(db, krate, &u_canonical);
    solution.map(|solution| solution_from_chalk(db, solution))
}

fn solution_from_chalk(db: &impl HirDatabase, solution: chalk_solve::Solution) -> Solution {
    let convert_subst = |subst: chalk_ir::Canonical<chalk_ir::Substitution>| {
        let value = subst
            .value
            .parameters
            .into_iter()
            .map(|p| {
                let ty = match p {
                    chalk_ir::Parameter(chalk_ir::ParameterKind::Ty(ty)) => from_chalk(db, ty),
                    chalk_ir::Parameter(chalk_ir::ParameterKind::Lifetime(_)) => unimplemented!(),
                };
                ty
            })
            .collect();
        let result = Canonical { value, num_vars: subst.binders.len() };
        SolutionVariables(result)
    };
    match solution {
        chalk_solve::Solution::Unique(constr_subst) => {
            let subst = chalk_ir::Canonical {
                value: constr_subst.value.subst,
                binders: constr_subst.binders,
            };
            Solution::Unique(convert_subst(subst))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Definite(subst)) => {
            Solution::Ambig(Guidance::Definite(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Suggested(subst)) => {
            Solution::Ambig(Guidance::Suggested(convert_subst(subst)))
        }
        chalk_solve::Solution::Ambig(chalk_solve::Guidance::Unknown) => {
            Solution::Ambig(Guidance::Unknown)
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SolutionVariables(pub Canonical<Vec<Ty>>);

#[derive(Clone, Debug, PartialEq, Eq)]
/// A (possible) solution for a proposed goal.
pub enum Solution {
    /// The goal indeed holds, and there is a unique value for all existential
    /// variables.
    Unique(SolutionVariables),

    /// The goal may be provable in multiple ways, but regardless we may have some guidance
    /// for type inference. In this case, we don't return any lifetime
    /// constraints, since we have not "committed" to any particular solution
    /// yet.
    Ambig(Guidance),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// When a goal holds ambiguously (e.g., because there are multiple possible
/// solutions), we issue a set of *guidance* back to type inference.
pub enum Guidance {
    /// The existential variables *must* have the given values if the goal is
    /// ever to hold, but that alone isn't enough to guarantee the goal will
    /// actually hold.
    Definite(SolutionVariables),

    /// There are multiple plausible values for the existentials, but the ones
    /// here are suggested as the preferred choice heuristically. These should
    /// be used for inference fallback only.
    Suggested(SolutionVariables),

    /// There's no useful information to feed back to type inference
    Unknown,
}
<|endoftext|>


<commit_before>//! Waking mechanism for threads blocked on channel operations.

use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread::{self, ThreadId};

use crate::context::Context;
use crate::select::{Operation, Selected};
use crate::utils::Spinlock;

/// Represents a thread blocked on a specific channel operation.
pub(crate) struct Entry {
    /// The operation.
    pub(crate) oper: Operation,

    /// Optional packet.
    pub(crate) packet: *mut (),

    /// Context associated with the thread owning this operation.
    pub(crate) cx: Context,
}

/// A queue of threads blocked on channel operations.
///
/// This data structure is used by threads to register blocking operations and get woken up once
/// an operation becomes ready.
pub(crate) struct Waker {
    /// A list of select operations.
    selectors: Vec<Entry>,

    /// A list of operations waiting to be ready.
    observers: Vec<Entry>,
}

impl Waker {
    /// Creates a new `Waker`.
    #[inline]
    pub(crate) fn new() -> Self {
        Waker {
            selectors: Vec::new(),
            observers: Vec::new(),
        }
    }

    /// Registers a select operation.
    #[inline]
    pub(crate) fn register(&mut self, oper: Operation, cx: &Context) {
        self.register_with_packet(oper, ptr::null_mut(), cx);
    }

    /// Registers a select operation and a packet.
    #[inline]
    pub(crate) fn register_with_packet(&mut self, oper: Operation, packet: *mut (), cx: &Context) {
        self.selectors.push(Entry {
            oper,
            packet,
            cx: cx.clone(),
        });
    }

    /// Unregisters a select operation.
    #[inline]
    pub(crate) fn unregister(&mut self, oper: Operation) -> Option<Entry> {
        if let Some((i, _)) = self
            .selectors
            .iter()
            .enumerate()
            .find(|&(_, entry)| entry.oper == oper)
        {
            let entry = self.selectors.remove(i);
            Some(entry)
        } else {
            None
        }
    }

    /// Attempts to find another thread's entry, select the operation, and wake it up.
    #[inline]
    pub(crate) fn try_select(&mut self) -> Option<Entry> {
        self.selectors
            .iter()
            .position(|selector| {
                // Does the entry belong to a different thread?
                selector.cx.thread_id() != current_thread_id()
                    && selector // Try selecting this operation.
                        .cx
                        .try_select(Selected::Operation(selector.oper))
                        .is_ok()
                    && {
                        // Provide the packet.
                        selector.cx.store_packet(selector.packet);
                        // Wake the thread up.
                        selector.cx.unpark();
                        true
                    }
            })
            // Remove the entry from the queue to keep it clean and improve
            // performance.
            .map(|pos| self.selectors.remove(pos))
    }

    /// Returns `true` if there is an entry which can be selected by the current thread.
    #[inline]
    pub(crate) fn can_select(&self) -> bool {
        if self.selectors.is_empty() {
            false
        } else {
            let thread_id = current_thread_id();

            self.selectors.iter().any(|entry| {
                entry.cx.thread_id() != thread_id && entry.cx.selected() == Selected::Waiting
            })
        }
    }

    /// Registers an operation waiting to be ready.
    #[inline]
    pub(crate) fn watch(&mut self, oper: Operation, cx: &Context) {
        self.observers.push(Entry {
            oper,
            packet: ptr::null_mut(),
            cx: cx.clone(),
        });
    }

    /// Unregisters an operation waiting to be ready.
    #[inline]
    pub(crate) fn unwatch(&mut self, oper: Operation) {
        self.observers.retain(|e| e.oper != oper);
    }

    /// Notifies all operations waiting to be ready.
    #[inline]
    pub(crate) fn notify(&mut self) {
        for entry in self.observers.drain(..) {
            if entry.cx.try_select(Selected::Operation(entry.oper)).is_ok() {
                entry.cx.unpark();
            }
        }
    }

    /// Notifies all registered operations that the channel is disconnected.
    #[inline]
    pub(crate) fn disconnect(&mut self) {
        for entry in self.selectors.iter() {
            if entry.cx.try_select(Selected::Disconnected).is_ok() {
                // Wake the thread up.
                //
                // Here we don't remove the entry from the queue. Registered threads must
                // unregister from the waker by themselves. They might also want to recover the
                // packet value and destroy it, if necessary.
                entry.cx.unpark();
            }
        }

        self.notify();
    }
}

impl Drop for Waker {
    #[inline]
    fn drop(&mut self) {
        debug_assert_eq!(self.selectors.len(), 0);
        debug_assert_eq!(self.observers.len(), 0);
    }
}

/// A waker that can be shared among threads without locking.
///
/// This is a simple wrapper around `Waker` that internally uses a mutex for synchronization.
pub(crate) struct SyncWaker {
    /// The inner `Waker`.
    inner: Spinlock<Waker>,

    /// `true` if the waker is empty.
    is_empty: AtomicBool,
}

impl SyncWaker {
    /// Creates a new `SyncWaker`.
    #[inline]
    pub(crate) fn new() -> Self {
        SyncWaker {
            inner: Spinlock::new(Waker::new()),
            is_empty: AtomicBool::new(true),
        }
    }

    /// Registers the current thread with an operation.
    #[inline]
    pub(crate) fn register(&self, oper: Operation, cx: &Context) {
        let mut inner = self.inner.lock();
        inner.register(oper, cx);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Unregisters an operation previously registered by the current thread.
    #[inline]
    pub(crate) fn unregister(&self, oper: Operation) -> Option<Entry> {
        let mut inner = self.inner.lock();
        let entry = inner.unregister(oper);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
        entry
    }

    /// Attempts to find one thread (not the current one), select its operation, and wake it up.
    #[inline]
    pub(crate) fn notify(&self) {
        if !self.is_empty.load(Ordering::SeqCst) {
            let mut inner = self.inner.lock();
            if !self.is_empty.load(Ordering::SeqCst) {
                inner.try_select();
                inner.notify();
                self.is_empty.store(
                    inner.selectors.is_empty() && inner.observers.is_empty(),
                    Ordering::SeqCst,
                );
            }
        }
    }

    /// Registers an operation waiting to be ready.
    #[inline]
    pub(crate) fn watch(&self, oper: Operation, cx: &Context) {
        let mut inner = self.inner.lock();
        inner.watch(oper, cx);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Unregisters an operation waiting to be ready.
    #[inline]
    pub(crate) fn unwatch(&self, oper: Operation) {
        let mut inner = self.inner.lock();
        inner.unwatch(oper);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Notifies all threads that the channel is disconnected.
    #[inline]
    pub(crate) fn disconnect(&self) {
        let mut inner = self.inner.lock();
        inner.disconnect();
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }
}

impl Drop for SyncWaker {
    #[inline]
    fn drop(&mut self) {
        debug_assert!(self.is_empty.load(Ordering::SeqCst));
    }
}

/// Returns the id of the current thread.
#[inline]
fn current_thread_id() -> ThreadId {
    thread_local! {
        /// Cached thread-local id.
        static THREAD_ID: ThreadId = thread::current().id();
    }

    THREAD_ID
        .try_with(|id| *id)
        .unwrap_or_else(|_| thread::current().id())
}
<commit_msg>channel: Do not call current_thread_id inside iteration<commit_after>//! Waking mechanism for threads blocked on channel operations.

use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread::{self, ThreadId};

use crate::context::Context;
use crate::select::{Operation, Selected};
use crate::utils::Spinlock;

/// Represents a thread blocked on a specific channel operation.
pub(crate) struct Entry {
    /// The operation.
    pub(crate) oper: Operation,

    /// Optional packet.
    pub(crate) packet: *mut (),

    /// Context associated with the thread owning this operation.
    pub(crate) cx: Context,
}

/// A queue of threads blocked on channel operations.
///
/// This data structure is used by threads to register blocking operations and get woken up once
/// an operation becomes ready.
pub(crate) struct Waker {
    /// A list of select operations.
    selectors: Vec<Entry>,

    /// A list of operations waiting to be ready.
    observers: Vec<Entry>,
}

impl Waker {
    /// Creates a new `Waker`.
    #[inline]
    pub(crate) fn new() -> Self {
        Waker {
            selectors: Vec::new(),
            observers: Vec::new(),
        }
    }

    /// Registers a select operation.
    #[inline]
    pub(crate) fn register(&mut self, oper: Operation, cx: &Context) {
        self.register_with_packet(oper, ptr::null_mut(), cx);
    }

    /// Registers a select operation and a packet.
    #[inline]
    pub(crate) fn register_with_packet(&mut self, oper: Operation, packet: *mut (), cx: &Context) {
        self.selectors.push(Entry {
            oper,
            packet,
            cx: cx.clone(),
        });
    }

    /// Unregisters a select operation.
    #[inline]
    pub(crate) fn unregister(&mut self, oper: Operation) -> Option<Entry> {
        if let Some((i, _)) = self
            .selectors
            .iter()
            .enumerate()
            .find(|&(_, entry)| entry.oper == oper)
        {
            let entry = self.selectors.remove(i);
            Some(entry)
        } else {
            None
        }
    }

    /// Attempts to find another thread's entry, select the operation, and wake it up.
    #[inline]
    pub(crate) fn try_select(&mut self) -> Option<Entry> {
        if self.selectors.is_empty() {
            None
        } else {
            let thread_id = current_thread_id();

            self.selectors
                .iter()
                .position(|selector| {
                    // Does the entry belong to a different thread?
                    selector.cx.thread_id() != thread_id
                        && selector // Try selecting this operation.
                            .cx
                            .try_select(Selected::Operation(selector.oper))
                            .is_ok()
                        && {
                            // Provide the packet.
                            selector.cx.store_packet(selector.packet);
                            // Wake the thread up.
                            selector.cx.unpark();
                            true
                        }
                })
                // Remove the entry from the queue to keep it clean and improve
                // performance.
                .map(|pos| self.selectors.remove(pos))
        }
    }

    /// Returns `true` if there is an entry which can be selected by the current thread.
    #[inline]
    pub(crate) fn can_select(&self) -> bool {
        if self.selectors.is_empty() {
            false
        } else {
            let thread_id = current_thread_id();

            self.selectors.iter().any(|entry| {
                entry.cx.thread_id() != thread_id && entry.cx.selected() == Selected::Waiting
            })
        }
    }

    /// Registers an operation waiting to be ready.
    #[inline]
    pub(crate) fn watch(&mut self, oper: Operation, cx: &Context) {
        self.observers.push(Entry {
            oper,
            packet: ptr::null_mut(),
            cx: cx.clone(),
        });
    }

    /// Unregisters an operation waiting to be ready.
    #[inline]
    pub(crate) fn unwatch(&mut self, oper: Operation) {
        self.observers.retain(|e| e.oper != oper);
    }

    /// Notifies all operations waiting to be ready.
    #[inline]
    pub(crate) fn notify(&mut self) {
        for entry in self.observers.drain(..) {
            if entry.cx.try_select(Selected::Operation(entry.oper)).is_ok() {
                entry.cx.unpark();
            }
        }
    }

    /// Notifies all registered operations that the channel is disconnected.
    #[inline]
    pub(crate) fn disconnect(&mut self) {
        for entry in self.selectors.iter() {
            if entry.cx.try_select(Selected::Disconnected).is_ok() {
                // Wake the thread up.
                //
                // Here we don't remove the entry from the queue. Registered threads must
                // unregister from the waker by themselves. They might also want to recover the
                // packet value and destroy it, if necessary.
                entry.cx.unpark();
            }
        }

        self.notify();
    }
}

impl Drop for Waker {
    #[inline]
    fn drop(&mut self) {
        debug_assert_eq!(self.selectors.len(), 0);
        debug_assert_eq!(self.observers.len(), 0);
    }
}

/// A waker that can be shared among threads without locking.
///
/// This is a simple wrapper around `Waker` that internally uses a mutex for synchronization.
pub(crate) struct SyncWaker {
    /// The inner `Waker`.
    inner: Spinlock<Waker>,

    /// `true` if the waker is empty.
    is_empty: AtomicBool,
}

impl SyncWaker {
    /// Creates a new `SyncWaker`.
    #[inline]
    pub(crate) fn new() -> Self {
        SyncWaker {
            inner: Spinlock::new(Waker::new()),
            is_empty: AtomicBool::new(true),
        }
    }

    /// Registers the current thread with an operation.
    #[inline]
    pub(crate) fn register(&self, oper: Operation, cx: &Context) {
        let mut inner = self.inner.lock();
        inner.register(oper, cx);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Unregisters an operation previously registered by the current thread.
    #[inline]
    pub(crate) fn unregister(&self, oper: Operation) -> Option<Entry> {
        let mut inner = self.inner.lock();
        let entry = inner.unregister(oper);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
        entry
    }

    /// Attempts to find one thread (not the current one), select its operation, and wake it up.
    #[inline]
    pub(crate) fn notify(&self) {
        if !self.is_empty.load(Ordering::SeqCst) {
            let mut inner = self.inner.lock();
            if !self.is_empty.load(Ordering::SeqCst) {
                inner.try_select();
                inner.notify();
                self.is_empty.store(
                    inner.selectors.is_empty() && inner.observers.is_empty(),
                    Ordering::SeqCst,
                );
            }
        }
    }

    /// Registers an operation waiting to be ready.
    #[inline]
    pub(crate) fn watch(&self, oper: Operation, cx: &Context) {
        let mut inner = self.inner.lock();
        inner.watch(oper, cx);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Unregisters an operation waiting to be ready.
    #[inline]
    pub(crate) fn unwatch(&self, oper: Operation) {
        let mut inner = self.inner.lock();
        inner.unwatch(oper);
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }

    /// Notifies all threads that the channel is disconnected.
    #[inline]
    pub(crate) fn disconnect(&self) {
        let mut inner = self.inner.lock();
        inner.disconnect();
        self.is_empty.store(
            inner.selectors.is_empty() && inner.observers.is_empty(),
            Ordering::SeqCst,
        );
    }
}

impl Drop for SyncWaker {
    #[inline]
    fn drop(&mut self) {
        debug_assert!(self.is_empty.load(Ordering::SeqCst));
    }
}

/// Returns the id of the current thread.
#[inline]
fn current_thread_id() -> ThreadId {
    thread_local! {
        /// Cached thread-local id.
        static THREAD_ID: ThreadId = thread::current().id();
    }

    THREAD_ID
        .try_with(|id| *id)
        .unwrap_or_else(|_| thread::current().id())
}
<|endoftext|>


<commit_before>//! Helpers for boolean operations

use crate::value::JsonTruthy;

handlebars_helper!(eq: |x: i64, y: i64| x == y);
handlebars_helper!(ne: |x: i64, y: i64| x != y);
handlebars_helper!(gt: |x: i64, y: i64| x > y);
handlebars_helper!(gte: |x: i64, y: i64| x >= y);
handlebars_helper!(lt: |x: i64, y: i64| x < y);
handlebars_helper!(lte: |x: i64, y: i64| x <= y);
handlebars_helper!(and: |x: Json, y: Json| x.is_truthy(false) && y.is_truthy(false));
handlebars_helper!(or: |x: Json, y: Json| x.is_truthy(false) || y.is_truthy(false));
handlebars_helper!(not: |x: Json| !x.is_truthy(false));

#[cfg(test)]
mod test_conditions {
    fn test_condition(condition: &str, expected: bool) {
        let handlebars = crate::Handlebars::new();

        let result = handlebars
            .render_template(
                &format!(
                    "{{{{#if {condition}}}}}lorem{{{{else}}}}ipsum{{{{/if}}}}",
                    condition = condition
                ),
                &json!({}),
            )
            .unwrap();
        assert_eq!(&result, if expected { "lorem" } else { "ipsum" });
    }

    #[test]
    fn foo() {
        test_condition("(gt 5 3)", true);
        test_condition("(gt 3 5)", false);
        test_condition("(or (gt 3 5) (gt 5 3))", true);
        test_condition("(not [])", true);
        test_condition("(and null 4)", false);
        test_condition("(eq 5 5)", true);
        test_condition("(eq 5 6)", false);
        test_condition("(ne 5 6)", true);
        test_condition("(ne 5 5)", false);
    }

    #[test]
    fn nested_conditions() {
        let handlebars = crate::Handlebars::new();

        let result = handlebars
            .render_template("{{#if (gt 5 3)}}lorem{{else}}ipsum{{/if}}", &json!({}))
            .unwrap();
        assert_eq!(&result, "lorem");

        let result = handlebars
            .render_template(
                "{{#if (not (gt 5 3))}}lorem{{else}}ipsum{{/if}}",
                &json!({}),
            )
            .unwrap();
        assert_eq!(&result, "ipsum");
    }

}
<commit_msg>(feat) extend helper "eq" and "ne" to any Json type (number, string, array,...)<commit_after>//! Helpers for boolean operations

use crate::value::JsonTruthy;

handlebars_helper!(eq: |x: Json, y: Json| x == y);
handlebars_helper!(ne: |x: Json, y: Json| x != y);
handlebars_helper!(gt: |x: i64, y: i64| x > y);
handlebars_helper!(gte: |x: i64, y: i64| x >= y);
handlebars_helper!(lt: |x: i64, y: i64| x < y);
handlebars_helper!(lte: |x: i64, y: i64| x <= y);
handlebars_helper!(and: |x: Json, y: Json| x.is_truthy(false) && y.is_truthy(false));
handlebars_helper!(or: |x: Json, y: Json| x.is_truthy(false) || y.is_truthy(false));
handlebars_helper!(not: |x: Json| !x.is_truthy(false));

#[cfg(test)]
mod test_conditions {
    fn test_condition(condition: &str, expected: bool) {
        let handlebars = crate::Handlebars::new();

        let result = handlebars
            .render_template(
                &format!(
                    "{{{{#if {condition}}}}}lorem{{{{else}}}}ipsum{{{{/if}}}}",
                    condition = condition
                ),
                &json!({}),
            )
            .unwrap();
        assert_eq!(&result, if expected { "lorem" } else { "ipsum" });
    }

    #[test]
    fn foo() {
        test_condition("(gt 5 3)", true);
        test_condition("(gt 3 5)", false);
        test_condition("(or (gt 3 5) (gt 5 3))", true);
        test_condition("(not [])", true);
        test_condition("(and null 4)", false);
    }

    #[test]
    fn test_eq() {
        test_condition("(eq 5 5)", true);
        test_condition("(eq 5 6)", false);
        test_condition(r#"(eq "foo" "foo")"#, true);
        test_condition(r#"(eq "foo" "Foo")"#, false);
        test_condition(r#"(eq [5] [5])"#, true);
        test_condition(r#"(eq [5] [4])"#, false);
        test_condition(r#"(eq 5 "5")"#, false);
        test_condition(r#"(eq 5 [5])"#, false);
    }

    #[test]
    fn test_ne() {
        test_condition("(ne 5 6)", true);
        test_condition("(ne 5 5)", false);
        test_condition(r#"(ne "foo" "foo")"#, false);
        test_condition(r#"(ne "foo" "Foo")"#, true);
    }

    #[test]
    fn nested_conditions() {
        let handlebars = crate::Handlebars::new();

        let result = handlebars
            .render_template("{{#if (gt 5 3)}}lorem{{else}}ipsum{{/if}}", &json!({}))
            .unwrap();
        assert_eq!(&result, "lorem");

        let result = handlebars
            .render_template(
                "{{#if (not (gt 5 3))}}lorem{{else}}ipsum{{/if}}",
                &json!({}),
            )
            .unwrap();
        assert_eq!(&result, "ipsum");
    }

}
<|endoftext|>


<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, not(target_os = "emscripten")))]
mod fork {
    use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    #[allow(deprecated)] // Required for compatibility with Rust < 1.24.
    use core::sync::atomic::{ATOMIC_BOOL_INIT, ATOMIC_USIZE_INIT};

    // 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)
    }

    #[allow(deprecated)]
    static FORK_HANDLER_REGISTERED: AtomicBool = ATOMIC_BOOL_INIT;

    extern 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() {
        if !FORK_HANDLER_REGISTERED.load(Ordering::Relaxed) {
            unsafe { libc::pthread_atfork(None, None, Some(fork_handler)) };
            FORK_HANDLER_REGISTERED.store(true, Ordering::Relaxed);
        }
    }
}

#[cfg(not(all(unix, 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>rustfmt: extern "C"<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, not(target_os = "emscripten")))]
mod fork {
    use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    #[allow(deprecated)] // Required for compatibility with Rust < 1.24.
    use core::sync::atomic::{ATOMIC_BOOL_INIT, ATOMIC_USIZE_INIT};

    // 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)
    }

    #[allow(deprecated)]
    static FORK_HANDLER_REGISTERED: AtomicBool = ATOMIC_BOOL_INIT;

    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() {
        if !FORK_HANDLER_REGISTERED.load(Ordering::Relaxed) {
            unsafe { libc::pthread_atfork(None, None, Some(fork_handler)) };
            FORK_HANDLER_REGISTERED.store(true, Ordering::Relaxed);
        }
    }
}

#[cfg(not(all(unix, 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><commit_msg>:zap: Do not stop import on failed to get branch name.<commit_after><|endoftext|>


<commit_before><commit_msg>Unified timestamp handling<commit_after><|endoftext|>


<commit_before>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters};

use gr::{GI, width_box, split_box, midpoint_box, eps_tol};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

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

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_high: f64) {
    let mut max = f_best_high;
    let lq = q.clone();
    for qi in lq.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    let _ = writeln!(&mut std::io::stderr(),
                     "Possible bound: {}, guaranteed bound: {}", f_best_high,
                     max);
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool)
        -> (Flt, Vec<GI>, bool) {
    let mut f_best_high = NINF;
    let mut f_best_low  = NINF;
    let mut best_x = x_0.clone();

    let mut i: u32 = 0;//binade_presplit(&f, &q, &x_0, e_x);
    q.write().unwrap().push(Quple{p: INF, pf: 0, data: x_0.clone(),
                                  fdata: f.call(&x_0)});

    while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire) {
        if sync.load(Ordering::Acquire) {
            // Ugly: Update the update thread's view of the best branch bound.
            *f_best_shared.write().unwrap() = f_best_low;
            b1.wait();
            b2.wait();
        }
        // Take q as writable during an iteration
        let mut q = q.write().unwrap();
        let fbl_orig = f_best_low;
        f_best_low = max!(f_best_low, *f_bestag.read().unwrap());
        if fbl_orig != f_best_low {
            log_max(&q, f_best_high);
        }
        let (ref x, fx) =
            match q.pop() {
                Some(y) => (y.data, y.fdata),
                None    => unreachable!()
            };

        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, e_f) {
                if f_best_high < fx.upper() {
                    f_best_high = fx.upper();
                    best_x = x.clone();
                    if logging {
                        log_max(&q, f_best_high);
                    }
                }
                continue;
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let mid = midpoint_box(&sx);
                let est_m = f.call(&mid);
                let fsx = f.call(&sx);
                let est_max = if est_m.lower() > fsx.lower() {est_m} else {fsx};
                if f_best_low < est_max.lower()  {
                    f_best_low = est_max.lower();
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                i += 1;
                q.push(Quple{p: est_max.upper(),
                             pf: i,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    if !stop.load(Ordering::Acquire) {
        // Exiting normally
        // Tell GA thread to stop
        stop.store(true, Ordering::Release);
        (f_best_high, best_x, true)
    }
    else {
        // Exiting by some other condition
        (f_best_high, best_x, false)
    }
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();

    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
            thread::sleep(Duration::new(1, 0));
            if timeout > 0 &&
                (time::get_time() - start).num_seconds() >= timeout as i64 {
                    let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                    stop.store(true, Ordering::Release);
                    break 'out;
                }
            if stop.load(Ordering::Acquire) { // Check if we've already stopped
                break 'out;
            }
        }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();                                                               // Dead lock here?
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();                                                               // Dead lock here?
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[{}, {{}}]", result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging)
        })};

    let _ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 2000, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.3_f64 // 0.5_f64
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    /*let ea_result = ea_thread.unwrap().join();
    if !ea_result.is_ok() {
        unreachable!();
    }*/

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (mut max, mut interval, completed) = result.unwrap();
        if !completed {
            // Go through all remaining intervals from IBBA to find the true
            // max
            let mut lq = q.write().unwrap();
            while lq.len() != 0 {
                let v = lq.pop().unwrap();
                let ref l_data = &v.data;
                let ub = fo.call(l_data).upper();
                if ub > max {
                    max = ub;
                    interval =v.data.clone();
                }
            }
        }
        println!("[{}, {{", max);
        for i in 0..args.names.len() {
            println!("'{}' : {},", args.names[i], interval[i].to_string());
        }
        println!("}}]");

    }
    else {println!("error")}
}
<commit_msg>More info in logging<commit_after>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters};

use gr::{GI, width_box, split_box, midpoint_box, eps_tol};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

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

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_low: f64,
           f_best_high: f64) {
    let mut max = f_best_high;
    let lq = q.clone();
    for qi in lq.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    let _ = writeln!(&mut std::io::stderr(),
                     "lb: {}, possible ub: {}, guaranteed ub: {}",
                     f_best_low,
                     f_best_high,
                     max);
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool)
        -> (Flt, Vec<GI>, bool) {
    let mut f_best_high = NINF;
    let mut f_best_low  = NINF;
    let mut best_x = x_0.clone();

    let mut i: u32 = 0;//binade_presplit(&f, &q, &x_0, e_x);
    q.write().unwrap().push(Quple{p: INF, pf: 0, data: x_0.clone(),
                                  fdata: f.call(&x_0)});

    while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire) {
        if sync.load(Ordering::Acquire) {
            // Ugly: Update the update thread's view of the best branch bound.
            *f_best_shared.write().unwrap() = f_best_low;
            b1.wait();
            b2.wait();
        }
        // Take q as writable during an iteration
        let mut q = q.write().unwrap();
        let fbl_orig = f_best_low;
        f_best_low = max!(f_best_low, *f_bestag.read().unwrap());
        if fbl_orig != f_best_low {
            log_max(&q, f_best_low, f_best_high);
        }
        let (ref x, fx) =
            match q.pop() {
                Some(y) => (y.data, y.fdata),
                None    => unreachable!()
            };

        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, e_f) {
                if f_best_high < fx.upper() {
                    f_best_high = fx.upper();
                    best_x = x.clone();
                    if logging {
                        log_max(&q, f_best_low, f_best_high);
                    }
                }
                continue;
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let mid = midpoint_box(&sx);
                let est_m = f.call(&mid);
                let fsx = f.call(&sx);
                let est_max = if est_m.lower() > fsx.lower() {est_m} else {fsx};
                if f_best_low < est_max.lower()  {
                    f_best_low = est_max.lower();
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                i += 1;
                q.push(Quple{p: est_max.upper(),
                             pf: i,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    if !stop.load(Ordering::Acquire) {
        // Exiting normally
        // Tell GA thread to stop
        stop.store(true, Ordering::Release);
        (f_best_high, best_x, true)
    }
    else {
        // Exiting by some other condition
        (f_best_high, best_x, false)
    }
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();

    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
            thread::sleep(Duration::new(1, 0));
            if timeout > 0 &&
                (time::get_time() - start).num_seconds() >= timeout as i64 {
                    let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                    stop.store(true, Ordering::Release);
                    break 'out;
                }
            if stop.load(Ordering::Acquire) { // Check if we've already stopped
                break 'out;
            }
        }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();                                                               // Dead lock here?
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();                                                               // Dead lock here?
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[{}, {{}}]", result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging)
        })};

    let _ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 2000, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.3_f64 // 0.5_f64
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    /*let ea_result = ea_thread.unwrap().join();
    if !ea_result.is_ok() {
        unreachable!();
    }*/

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (mut max, mut interval, completed) = result.unwrap();
        if !completed {
            // Go through all remaining intervals from IBBA to find the true
            // max
            let mut lq = q.write().unwrap();
            while lq.len() != 0 {
                let v = lq.pop().unwrap();
                let ref l_data = &v.data;
                let ub = fo.call(l_data).upper();
                if ub > max {
                    max = ub;
                    interval =v.data.clone();
                }
            }
        }
        println!("[{}, {{", max);
        for i in 0..args.names.len() {
            println!("'{}' : {},", args.names[i], interval[i].to_string());
        }
        println!("}}]");

    }
    else {println!("error")}
}
<|endoftext|>