<commit_before>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#[cfg(feature = "nightly")]
use std::sync::atomic::{AtomicU8, ATOMIC_U8_INIT, Ordering, fence};
#[cfg(feature = "nightly")]
type U8 = u8;
#[cfg(not(feature = "nightly"))]
use stable::{AtomicU8, ATOMIC_U8_INIT, Ordering, fence};
#[cfg(not(feature = "nightly"))]
type U8 = usize;
use std::mem;
use parking_lot_core::{self, SpinWait, DEFAULT_PARK_TOKEN, DEFAULT_UNPARK_TOKEN};
use util::UncheckedOptionExt;

const DONE_BIT: U8 = 1;
const POISON_BIT: U8 = 2;
const LOCKED_BIT: U8 = 4;
const PARKED_BIT: U8 = 8;

/// Current state of a `Once`.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum OnceState {
    /// A closure has not been executed yet
    New,

    /// A closure was executed but panicked.
    Poisoned,

    /// A thread is currently executing a closure.
    InProgress,

    /// A closure has completed sucessfully.
    Done,
}

impl OnceState {
    /// Returns whether the associated `Once` has been poisoned.
    ///
    /// Once an initalization routine for a `Once` has panicked it will forever
    /// indicate to future forced initialization routines that it is poisoned.
    #[inline]
    pub fn poisoned(&self) -> bool {
        match *self {
            OnceState::Poisoned => true,
            _ => false,
        }
    }

    /// Returns whether the associated `Once` has successfullly executed a
    /// closure.
    ///
    /// Note that you generally need an `Acquire` barrier after checking this to
    /// ensure that operations are correctly ordered with regards to a
    /// concurrently executing closure.
    #[inline]
    pub fn done(&self) -> bool {
        match *self {
            OnceState::Done => true,
            _ => false,
        }
    }
}

/// A synchronization primitive which can be used to run a one-time
/// initialization. Useful for one-time initialization for globals, FFI or
/// related functionality.
///
/// # Differences from the standard library `Once`
///
/// - Only requires 1 byte of space, instead of 1 word.
/// - Not required to be `'static`.
/// - Relaxed memory barriers in the fast path, which can significantly improve
///   performance on some architectures.
/// - Efficient handling of micro-contention using adaptive spinning.
///
/// # Examples
///
/// ```
/// use parking_lot::{Once, ONCE_INIT};
///
/// static START: Once = ONCE_INIT;
///
/// START.call_once(|| {
///     // run initialization here
/// });
/// ```
pub struct Once(AtomicU8);

/// Initialization value for static `Once` values.
pub const ONCE_INIT: Once = Once(ATOMIC_U8_INIT);

impl Once {
    /// Creates a new `Once` value.
    #[cfg(feature = "nightly")]
    #[inline]
    pub const fn new() -> Once {
        Once(AtomicU8::new(0))
    }

    /// Creates a new `Once` value.
    #[cfg(not(feature = "nightly"))]
    #[inline]
    pub fn new() -> Once {
        Once(AtomicU8::new(0))
    }

    /// Returns the current state of this `Once`.
    #[inline]
    pub fn state(&self) -> OnceState {
        let state = self.0.load(Ordering::Relaxed);
        if state & DONE_BIT != 0 {
            OnceState::Done
        } else if state & LOCKED_BIT != 0 {
            OnceState::InProgress
        } else if state & POISON_BIT != 0 {
            OnceState::Poisoned
        } else {
            OnceState::New
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). It is also
    /// guaranteed that any memory writes performed by the executed closure can
    /// be reliably observed by other threads at this point (there is a
    /// happens-before relation between the closure and code executing after the
    /// return).
    ///
    /// # Examples
    ///
    /// ```
    /// use parking_lot::{Once, ONCE_INIT};
    ///
    /// static mut VAL: usize = 0;
    /// static INIT: Once = ONCE_INIT;
    ///
    /// // Accessing a `static mut` is unsafe much of the time, but if we do so
    /// // in a synchronized fashion (e.g. write once or read all) then we're
    /// // good to go!
    /// //
    /// // This function will only call `expensive_computation` once, and will
    /// // otherwise always return the value returned from the first invocation.
    /// fn get_cached_val() -> usize {
    ///     unsafe {
    ///         INIT.call_once(|| {
    ///             VAL = expensive_computation();
    ///         });
    ///         VAL
    ///     }
    /// }
    ///
    /// fn expensive_computation() -> usize {
    ///     // ...
    /// # 2
    /// }
    /// ```
    ///
    /// # Panics
    ///
    /// The closure `f` will only be executed once if this is called
    /// concurrently amongst many threads. If that closure panics, however, then
    /// it will *poison* this `Once` instance, causing all future invocations of
    /// `call_once` to also panic.
    #[inline]
    pub fn call_once<F>(&self, f: F)
        where F: FnOnce()
    {
        if self.0.load(Ordering::Acquire) == DONE_BIT {
            return;
        }

        let mut f = Some(f);
        self.call_once_slow(false, &mut |_| unsafe { f.take().unchecked_unwrap()() });
    }

    /// Performs the same function as `call_once` except ignores poisoning.
    ///
    /// If this `Once` has been poisoned (some initialization panicked) then
    /// this function will continue to attempt to call initialization functions
    /// until one of them doesn't panic.
    ///
    /// The closure `f` is yielded a structure which can be used to query the
    /// state of this `Once` (whether initialization has previously panicked or
    /// not).
    #[inline]
    pub fn call_once_force<F>(&self, f: F)
        where F: FnOnce(OnceState)
    {
        if self.0.load(Ordering::Acquire) == DONE_BIT {
            return;
        }

        let mut f = Some(f);
        self.call_once_slow(true,
                            &mut |state| unsafe { f.take().unchecked_unwrap()(state) });
    }

    // This is a non-generic function to reduce the monomorphization cost of
    // using `call_once` (this isn't exactly a trivial or small implementation).
    //
    // Additionally, this is tagged with `#[cold]` as it should indeed be cold
    // and it helps let LLVM know that calls to this function should be off the
    // fast path. Essentially, this should help generate more straight line code
    // in LLVM.
    //
    // Finally, this takes an `FnMut` instead of a `FnOnce` because there's
    // currently no way to take an `FnOnce` and call it via virtual dispatch
    // without some allocation overhead.
    #[cold]
    #[inline(never)]
    fn call_once_slow(&self, ignore_poison: bool, f: &mut FnMut(OnceState)) {
        let mut spinwait = SpinWait::new();
        let mut state = self.0.load(Ordering::Relaxed);
        loop {
            // If another thread called the closure, we're done
            if state & DONE_BIT != 0 {
                // An acquire fence is needed here since we didn't load the
                // state with Ordering::Acquire.
                fence(Ordering::Acquire);
                return;
            }

            // If the state has been poisoned and we aren't forcing, then panic
            if state & POISON_BIT != 0 && !ignore_poison {
                // Need the fence here as well for the same reason
                fence(Ordering::Acquire);
                panic!("Once instance has previously been poisoned");
            }

            // Grab the lock if it isn't locked, even if there is a queue on it.
            // We also clear the poison bit since we are going to try running
            // the closure again.
            if state & LOCKED_BIT == 0 {
                match self.0
                    .compare_exchange_weak(state,
                                           (state | LOCKED_BIT) & !POISON_BIT,
                                           Ordering::Acquire,
                                           Ordering::Relaxed) {
                    Ok(_) => break,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state & PARKED_BIT == 0 && spinwait.spin() {
                state = self.0.load(Ordering::Relaxed);
                continue;
            }

            // Set the parked bit
            if state & PARKED_BIT == 0 {
                if let Err(x) = self.0.compare_exchange_weak(state,
                                                             state | PARKED_BIT,
                                                             Ordering::Relaxed,
                                                             Ordering::Relaxed) {
                    state = x;
                    continue;
                }
            }

            // Park our thread until we are woken up by the thread that owns the
            // lock.
            unsafe {
                let addr = self as *const _ as usize;
                let validate = || self.0.load(Ordering::Relaxed) == LOCKED_BIT | PARKED_BIT;
                let before_sleep = || {};
                let timed_out = |_, _| unreachable!();
                parking_lot_core::park(addr,
                                       validate,
                                       before_sleep,
                                       timed_out,
                                       DEFAULT_PARK_TOKEN,
                                       None);
            }

            // Loop back and check if the done bit was set
            spinwait.reset();
            state = self.0.load(Ordering::Relaxed);
        }

        struct PanicGuard<'a>(&'a Once);
        impl<'a> Drop for PanicGuard<'a> {
            fn drop(&mut self) {
                // Mark the state as poisoned, unlock it and unpark all threads.
                let once = self.0;
                let state = once.0.swap(POISON_BIT, Ordering::Release);
                if state & PARKED_BIT != 0 {
                    unsafe {
                        let addr = once as *const _ as usize;
                        parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
                    }
                }
            }
        }

        // At this point we have the lock, so run the closure. Make sure we
        // properly clean up if the closure panicks.
        let guard = PanicGuard(self);
        let once_state = if state & POISON_BIT != 0 {
            OnceState::Poisoned
        } else {
            OnceState::New
        };
        f(once_state);
        mem::forget(guard);

        // Now unlock the state, set the done bit and unpark all threads
        let state = self.0.swap(DONE_BIT, Ordering::Release);
        if state & PARKED_BIT != 0 {
            unsafe {
                let addr = self as *const _ as usize;
                parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
            }
        }
    }
}

impl Default for Once {
    #[inline]
    fn default() -> Once {
        Once::new()
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "nightly")]
    use std::panic;
    use std::sync::mpsc::channel;
    use std::thread;
    use {Once, ONCE_INIT};

    #[test]
    fn smoke_once() {
        static O: Once = ONCE_INIT;
        let mut a = 0;
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
    }

    #[test]
    fn stampede_once() {
        static O: Once = ONCE_INIT;
        static mut RUN: bool = false;

        let (tx, rx) = channel();
        for _ in 0..10 {
            let tx = tx.clone();
            thread::spawn(move || {
                for _ in 0..4 {
                    thread::yield_now()
                }
                unsafe {
                    O.call_once(|| {
                        assert!(!RUN);
                        RUN = true;
                    });
                    assert!(RUN);
                }
                tx.send(()).unwrap();
            });
        }

        unsafe {
            O.call_once(|| {
                assert!(!RUN);
                RUN = true;
            });
            assert!(RUN);
        }

        for _ in 0..10 {
            rx.recv().unwrap();
        }
    }

    #[cfg(feature = "nightly")]
    #[test]
    fn poison_bad() {
        static O: Once = ONCE_INIT;

        // poison the once
        let t = panic::catch_unwind(|| { O.call_once(|| panic!()); });
        assert!(t.is_err());

        // poisoning propagates
        let t = panic::catch_unwind(|| { O.call_once(|| {}); });
        assert!(t.is_err());

        // we can subvert poisoning, however
        let mut called = false;
        O.call_once_force(|p| {
            called = true;
            assert!(p.poisoned())
        });
        assert!(called);

        // once any success happens, we stop propagating the poison
        O.call_once(|| {});
    }

    #[cfg(feature = "nightly")]
    #[test]
    fn wait_for_force_to_finish() {
        static O: Once = ONCE_INIT;

        // poison the once
        let t = panic::catch_unwind(|| { O.call_once(|| panic!()); });
        assert!(t.is_err());

        // make sure someone's waiting inside the once via a force
        let (tx1, rx1) = channel();
        let (tx2, rx2) = channel();
        let t1 = thread::spawn(move || {
            O.call_once_force(|p| {
                assert!(p.poisoned());
                tx1.send(()).unwrap();
                rx2.recv().unwrap();
            });
        });

        rx1.recv().unwrap();

        // put another waiter on the once
        let t2 = thread::spawn(|| {
            let mut called = false;
            O.call_once(|| { called = true; });
            assert!(!called);
        });

        tx2.send(()).unwrap();

        assert!(t1.join().is_ok());
        assert!(t2.join().is_ok());

    }
}
<commit_msg>Remove the need for a separate Acquire fence after calling Once::state<commit_after>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#[cfg(feature = "nightly")]
use std::sync::atomic::{AtomicU8, ATOMIC_U8_INIT, Ordering, fence};
#[cfg(feature = "nightly")]
type U8 = u8;
#[cfg(not(feature = "nightly"))]
use stable::{AtomicU8, ATOMIC_U8_INIT, Ordering, fence};
#[cfg(not(feature = "nightly"))]
type U8 = usize;
use std::mem;
use parking_lot_core::{self, SpinWait, DEFAULT_PARK_TOKEN, DEFAULT_UNPARK_TOKEN};
use util::UncheckedOptionExt;

const DONE_BIT: U8 = 1;
const POISON_BIT: U8 = 2;
const LOCKED_BIT: U8 = 4;
const PARKED_BIT: U8 = 8;

/// Current state of a `Once`.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum OnceState {
    /// A closure has not been executed yet
    New,

    /// A closure was executed but panicked.
    Poisoned,

    /// A thread is currently executing a closure.
    InProgress,

    /// A closure has completed sucessfully.
    Done,
}

impl OnceState {
    /// Returns whether the associated `Once` has been poisoned.
    ///
    /// Once an initalization routine for a `Once` has panicked it will forever
    /// indicate to future forced initialization routines that it is poisoned.
    #[inline]
    pub fn poisoned(&self) -> bool {
        match *self {
            OnceState::Poisoned => true,
            _ => false,
        }
    }

    /// Returns whether the associated `Once` has successfullly executed a
    /// closure.
    #[inline]
    pub fn done(&self) -> bool {
        match *self {
            OnceState::Done => true,
            _ => false,
        }
    }
}

/// A synchronization primitive which can be used to run a one-time
/// initialization. Useful for one-time initialization for globals, FFI or
/// related functionality.
///
/// # Differences from the standard library `Once`
///
/// - Only requires 1 byte of space, instead of 1 word.
/// - Not required to be `'static`.
/// - Relaxed memory barriers in the fast path, which can significantly improve
///   performance on some architectures.
/// - Efficient handling of micro-contention using adaptive spinning.
///
/// # Examples
///
/// ```
/// use parking_lot::{Once, ONCE_INIT};
///
/// static START: Once = ONCE_INIT;
///
/// START.call_once(|| {
///     // run initialization here
/// });
/// ```
pub struct Once(AtomicU8);

/// Initialization value for static `Once` values.
pub const ONCE_INIT: Once = Once(ATOMIC_U8_INIT);

impl Once {
    /// Creates a new `Once` value.
    #[cfg(feature = "nightly")]
    #[inline]
    pub const fn new() -> Once {
        Once(AtomicU8::new(0))
    }

    /// Creates a new `Once` value.
    #[cfg(not(feature = "nightly"))]
    #[inline]
    pub fn new() -> Once {
        Once(AtomicU8::new(0))
    }

    /// Returns the current state of this `Once`.
    #[inline]
    pub fn state(&self) -> OnceState {
        let state = self.0.load(Ordering::Acquire);
        if state & DONE_BIT != 0 {
            OnceState::Done
        } else if state & LOCKED_BIT != 0 {
            OnceState::InProgress
        } else if state & POISON_BIT != 0 {
            OnceState::Poisoned
        } else {
            OnceState::New
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). It is also
    /// guaranteed that any memory writes performed by the executed closure can
    /// be reliably observed by other threads at this point (there is a
    /// happens-before relation between the closure and code executing after the
    /// return).
    ///
    /// # Examples
    ///
    /// ```
    /// use parking_lot::{Once, ONCE_INIT};
    ///
    /// static mut VAL: usize = 0;
    /// static INIT: Once = ONCE_INIT;
    ///
    /// // Accessing a `static mut` is unsafe much of the time, but if we do so
    /// // in a synchronized fashion (e.g. write once or read all) then we're
    /// // good to go!
    /// //
    /// // This function will only call `expensive_computation` once, and will
    /// // otherwise always return the value returned from the first invocation.
    /// fn get_cached_val() -> usize {
    ///     unsafe {
    ///         INIT.call_once(|| {
    ///             VAL = expensive_computation();
    ///         });
    ///         VAL
    ///     }
    /// }
    ///
    /// fn expensive_computation() -> usize {
    ///     // ...
    /// # 2
    /// }
    /// ```
    ///
    /// # Panics
    ///
    /// The closure `f` will only be executed once if this is called
    /// concurrently amongst many threads. If that closure panics, however, then
    /// it will *poison* this `Once` instance, causing all future invocations of
    /// `call_once` to also panic.
    #[inline]
    pub fn call_once<F>(&self, f: F)
        where F: FnOnce()
    {
        if self.0.load(Ordering::Acquire) == DONE_BIT {
            return;
        }

        let mut f = Some(f);
        self.call_once_slow(false, &mut |_| unsafe { f.take().unchecked_unwrap()() });
    }

    /// Performs the same function as `call_once` except ignores poisoning.
    ///
    /// If this `Once` has been poisoned (some initialization panicked) then
    /// this function will continue to attempt to call initialization functions
    /// until one of them doesn't panic.
    ///
    /// The closure `f` is yielded a structure which can be used to query the
    /// state of this `Once` (whether initialization has previously panicked or
    /// not).
    #[inline]
    pub fn call_once_force<F>(&self, f: F)
        where F: FnOnce(OnceState)
    {
        if self.0.load(Ordering::Acquire) == DONE_BIT {
            return;
        }

        let mut f = Some(f);
        self.call_once_slow(true,
                            &mut |state| unsafe { f.take().unchecked_unwrap()(state) });
    }

    // This is a non-generic function to reduce the monomorphization cost of
    // using `call_once` (this isn't exactly a trivial or small implementation).
    //
    // Additionally, this is tagged with `#[cold]` as it should indeed be cold
    // and it helps let LLVM know that calls to this function should be off the
    // fast path. Essentially, this should help generate more straight line code
    // in LLVM.
    //
    // Finally, this takes an `FnMut` instead of a `FnOnce` because there's
    // currently no way to take an `FnOnce` and call it via virtual dispatch
    // without some allocation overhead.
    #[cold]
    #[inline(never)]
    fn call_once_slow(&self, ignore_poison: bool, f: &mut FnMut(OnceState)) {
        let mut spinwait = SpinWait::new();
        let mut state = self.0.load(Ordering::Relaxed);
        loop {
            // If another thread called the closure, we're done
            if state & DONE_BIT != 0 {
                // An acquire fence is needed here since we didn't load the
                // state with Ordering::Acquire.
                fence(Ordering::Acquire);
                return;
            }

            // If the state has been poisoned and we aren't forcing, then panic
            if state & POISON_BIT != 0 && !ignore_poison {
                // Need the fence here as well for the same reason
                fence(Ordering::Acquire);
                panic!("Once instance has previously been poisoned");
            }

            // Grab the lock if it isn't locked, even if there is a queue on it.
            // We also clear the poison bit since we are going to try running
            // the closure again.
            if state & LOCKED_BIT == 0 {
                match self.0
                    .compare_exchange_weak(state,
                                           (state | LOCKED_BIT) & !POISON_BIT,
                                           Ordering::Acquire,
                                           Ordering::Relaxed) {
                    Ok(_) => break,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state & PARKED_BIT == 0 && spinwait.spin() {
                state = self.0.load(Ordering::Relaxed);
                continue;
            }

            // Set the parked bit
            if state & PARKED_BIT == 0 {
                if let Err(x) = self.0.compare_exchange_weak(state,
                                                             state | PARKED_BIT,
                                                             Ordering::Relaxed,
                                                             Ordering::Relaxed) {
                    state = x;
                    continue;
                }
            }

            // Park our thread until we are woken up by the thread that owns the
            // lock.
            unsafe {
                let addr = self as *const _ as usize;
                let validate = || self.0.load(Ordering::Relaxed) == LOCKED_BIT | PARKED_BIT;
                let before_sleep = || {};
                let timed_out = |_, _| unreachable!();
                parking_lot_core::park(addr,
                                       validate,
                                       before_sleep,
                                       timed_out,
                                       DEFAULT_PARK_TOKEN,
                                       None);
            }

            // Loop back and check if the done bit was set
            spinwait.reset();
            state = self.0.load(Ordering::Relaxed);
        }

        struct PanicGuard<'a>(&'a Once);
        impl<'a> Drop for PanicGuard<'a> {
            fn drop(&mut self) {
                // Mark the state as poisoned, unlock it and unpark all threads.
                let once = self.0;
                let state = once.0.swap(POISON_BIT, Ordering::Release);
                if state & PARKED_BIT != 0 {
                    unsafe {
                        let addr = once as *const _ as usize;
                        parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
                    }
                }
            }
        }

        // At this point we have the lock, so run the closure. Make sure we
        // properly clean up if the closure panicks.
        let guard = PanicGuard(self);
        let once_state = if state & POISON_BIT != 0 {
            OnceState::Poisoned
        } else {
            OnceState::New
        };
        f(once_state);
        mem::forget(guard);

        // Now unlock the state, set the done bit and unpark all threads
        let state = self.0.swap(DONE_BIT, Ordering::Release);
        if state & PARKED_BIT != 0 {
            unsafe {
                let addr = self as *const _ as usize;
                parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
            }
        }
    }
}

impl Default for Once {
    #[inline]
    fn default() -> Once {
        Once::new()
    }
}

#[cfg(test)]
mod tests {
    #[cfg(feature = "nightly")]
    use std::panic;
    use std::sync::mpsc::channel;
    use std::thread;
    use {Once, ONCE_INIT};

    #[test]
    fn smoke_once() {
        static O: Once = ONCE_INIT;
        let mut a = 0;
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
        O.call_once(|| a += 1);
        assert_eq!(a, 1);
    }

    #[test]
    fn stampede_once() {
        static O: Once = ONCE_INIT;
        static mut RUN: bool = false;

        let (tx, rx) = channel();
        for _ in 0..10 {
            let tx = tx.clone();
            thread::spawn(move || {
                for _ in 0..4 {
                    thread::yield_now()
                }
                unsafe {
                    O.call_once(|| {
                        assert!(!RUN);
                        RUN = true;
                    });
                    assert!(RUN);
                }
                tx.send(()).unwrap();
            });
        }

        unsafe {
            O.call_once(|| {
                assert!(!RUN);
                RUN = true;
            });
            assert!(RUN);
        }

        for _ in 0..10 {
            rx.recv().unwrap();
        }
    }

    #[cfg(feature = "nightly")]
    #[test]
    fn poison_bad() {
        static O: Once = ONCE_INIT;

        // poison the once
        let t = panic::catch_unwind(|| { O.call_once(|| panic!()); });
        assert!(t.is_err());

        // poisoning propagates
        let t = panic::catch_unwind(|| { O.call_once(|| {}); });
        assert!(t.is_err());

        // we can subvert poisoning, however
        let mut called = false;
        O.call_once_force(|p| {
            called = true;
            assert!(p.poisoned())
        });
        assert!(called);

        // once any success happens, we stop propagating the poison
        O.call_once(|| {});
    }

    #[cfg(feature = "nightly")]
    #[test]
    fn wait_for_force_to_finish() {
        static O: Once = ONCE_INIT;

        // poison the once
        let t = panic::catch_unwind(|| { O.call_once(|| panic!()); });
        assert!(t.is_err());

        // make sure someone's waiting inside the once via a force
        let (tx1, rx1) = channel();
        let (tx2, rx2) = channel();
        let t1 = thread::spawn(move || {
            O.call_once_force(|p| {
                assert!(p.poisoned());
                tx1.send(()).unwrap();
                rx2.recv().unwrap();
            });
        });

        rx1.recv().unwrap();

        // put another waiter on the once
        let t2 = thread::spawn(|| {
            let mut called = false;
            O.call_once(|| { called = true; });
            assert!(!called);
        });

        tx2.send(()).unwrap();

        assert!(t1.join().is_ok());
        assert!(t2.join().is_ok());

    }
}
<|endoftext|>


<commit_before>use crate::ptr;
use crate::ops::{Deref, DerefMut};

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used, and that this function is called at most once
    /// for a given instance of ManuallyDrop
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<commit_msg>Improve formatting of 'ManuallyDrop'<commit_after>use crate::ptr;
use crate::ops::{Deref, DerefMut};

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used, and that this function is called at most once
    /// for a given instance of `ManuallyDrop<T>`.
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<|endoftext|>


<commit_before>use crate::ich::StableHashingContext;
use crate::mir::{BasicBlock, BasicBlockData, Body, LocalDecls, Location, Successors};
use rustc_data_structures::graph::dominators::{dominators, Dominators};
use rustc_data_structures::graph::{self, GraphPredecessors, GraphSuccessors};
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use rustc_index::vec::IndexVec;
use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
use std::iter;
use std::ops::{Deref, DerefMut, Index, IndexMut};
use std::vec::IntoIter;

#[derive(Clone, Debug)]
pub struct Cache {
    predecessors: Option<IndexVec<BasicBlock, Vec<BasicBlock>>>,
}

impl rustc_serialize::Encodable for Cache {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        Encodable::encode(&(), s)
    }
}

impl rustc_serialize::Decodable for Cache {
    fn decode<D: Decoder>(d: &mut D) -> Result<Self, D::Error> {
        Decodable::decode(d).map(|_v: ()| Self::new())
    }
}

impl<'a> HashStable<StableHashingContext<'a>> for Cache {
    fn hash_stable(&self, _: &mut StableHashingContext<'a>, _: &mut StableHasher) {
        // Do nothing.
    }
}

impl Cache {
    pub fn new() -> Self {
        Self { predecessors: None }
    }

    pub fn invalidate_predecessors(&mut self) {
        // FIXME: consider being more fine-grained
        self.predecessors = None;
    }

    pub fn ensure_predecessors(&mut self, body: &Body<'_>) {
        if self.predecessors.is_none() {
            let mut result = IndexVec::from_elem(vec![], body.basic_blocks());
            for (bb, data) in body.basic_blocks().iter_enumerated() {
                if let Some(ref term) = data.terminator {
                    for &tgt in term.successors() {
                        result[tgt].push(bb);
                    }
                }
            }

            self.predecessors = Some(result)
        }
    }

    /// This will recompute the predecessors cache if it is not available
    fn predecessors(&mut self, body: &Body<'_>) -> &IndexVec<BasicBlock, Vec<BasicBlock>> {
        self.ensure_predecessors(body);
        self.predecessors.as_ref().unwrap()
    }

    fn unwrap_predecessors_for(&self, bb: BasicBlock) -> &[BasicBlock] {
        &self.predecessors.as_ref().unwrap()[bb]
    }

    fn unwrap_predecessor_locations<'a>(
        &'a self,
        loc: Location,
        body: &'a Body<'a>,
    ) -> impl Iterator<Item = Location> + 'a {
        let if_zero_locations = if loc.statement_index == 0 {
            let predecessor_blocks = self.unwrap_predecessors_for(loc.block);
            let num_predecessor_blocks = predecessor_blocks.len();
            Some(
                (0..num_predecessor_blocks)
                    .map(move |i| predecessor_blocks[i])
                    .map(move |bb| body.terminator_loc(bb)),
            )
        } else {
            None
        };

        let if_not_zero_locations = if loc.statement_index == 0 {
            None
        } else {
            Some(Location { block: loc.block, statement_index: loc.statement_index - 1 })
        };

        if_zero_locations.into_iter().flatten().chain(if_not_zero_locations)
    }

    pub fn basic_blocks_mut<'a, 'tcx>(
        &mut self,
        body: &'a mut Body<'tcx>,
    ) -> &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        debug!("bbm: Clearing predecessors cache for body at: {:?}", body.span.data());
        self.invalidate_predecessors();
        &mut body.basic_blocks
    }

    pub fn basic_blocks_and_local_decls_mut<'a, 'tcx>(
        &mut self,
        body: &'a mut Body<'tcx>,
    ) -> (&'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &'a mut LocalDecls<'tcx>) {
        debug!("bbaldm: Clearing predecessors cache for body at: {:?}", body.span.data());
        self.invalidate_predecessors();
        (&mut body.basic_blocks, &mut body.local_decls)
    }
}

#[derive(Clone, Debug, HashStable, RustcEncodable, RustcDecodable, TypeFoldable)]
pub struct BodyAndCache<'tcx> {
    body: Body<'tcx>,
    cache: Cache,
}

impl BodyAndCache<'tcx> {
    pub fn new(body: Body<'tcx>) -> Self {
        Self { body, cache: Cache::new() }
    }
}

#[macro_export]
macro_rules! read_only {
    ($body:expr) => {{
        $body.ensure_predecessors();
        $body.unwrap_read_only()
    }};
}

impl BodyAndCache<'tcx> {
    pub fn ensure_predecessors(&mut self) {
        self.cache.ensure_predecessors(&self.body);
    }

    pub fn predecessors(&mut self) -> &IndexVec<BasicBlock, Vec<BasicBlock>> {
        self.cache.predecessors(&self.body)
    }

    pub fn unwrap_read_only(&self) -> ReadOnlyBodyAndCache<'_, 'tcx> {
        ReadOnlyBodyAndCache::new(&self.body, &self.cache)
    }

    pub fn basic_blocks_mut(&mut self) -> &mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        self.cache.basic_blocks_mut(&mut self.body)
    }

    pub fn basic_blocks_and_local_decls_mut(
        &mut self,
    ) -> (&mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &mut LocalDecls<'tcx>) {
        self.cache.basic_blocks_and_local_decls_mut(&mut self.body)
    }
}

impl<'tcx> Index<BasicBlock> for BodyAndCache<'tcx> {
    type Output = BasicBlockData<'tcx>;

    fn index(&self, index: BasicBlock) -> &BasicBlockData<'tcx> {
        &self.body[index]
    }
}

impl<'tcx> IndexMut<BasicBlock> for BodyAndCache<'tcx> {
    fn index_mut(&mut self, index: BasicBlock) -> &mut Self::Output {
        &mut self.basic_blocks_mut()[index]
    }
}

impl<'tcx> Deref for BodyAndCache<'tcx> {
    type Target = Body<'tcx>;

    fn deref(&self) -> &Self::Target {
        &self.body
    }
}

impl<'tcx> DerefMut for BodyAndCache<'tcx> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.body
    }
}

#[derive(Copy, Clone, Debug)]
pub struct ReadOnlyBodyAndCache<'a, 'tcx> {
    body: &'a Body<'tcx>,
    cache: &'a Cache,
}

impl ReadOnlyBodyAndCache<'a, 'tcx> {
    fn new(body: &'a Body<'tcx>, cache: &'a Cache) -> Self {
        assert!(
            cache.predecessors.is_some(),
            "Cannot construct ReadOnlyBodyAndCache without computed predecessors"
        );
        Self { body, cache }
    }

    pub fn predecessors(&self) -> &IndexVec<BasicBlock, Vec<BasicBlock>> {
        self.cache.predecessors.as_ref().unwrap()
    }

    pub fn predecessors_for(&self, bb: BasicBlock) -> &[BasicBlock] {
        self.cache.unwrap_predecessors_for(bb)
    }

    pub fn predecessor_locations(&self, loc: Location) -> impl Iterator<Item = Location> + '_ {
        self.cache.unwrap_predecessor_locations(loc, self.body)
    }

    pub fn basic_blocks(&self) -> &IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        &self.body.basic_blocks
    }

    pub fn dominators(&self) -> Dominators<BasicBlock> {
        dominators(self)
    }
}

impl graph::DirectedGraph for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Node = BasicBlock;
}

impl graph::GraphPredecessors<'graph> for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Item = BasicBlock;
    type Iter = IntoIter<BasicBlock>;
}

impl graph::WithPredecessors for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn predecessors(&self, node: Self::Node) -> <Self as GraphPredecessors<'_>>::Iter {
        self.cache.unwrap_predecessors_for(node).to_vec().into_iter()
    }
}

impl graph::WithNumNodes for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn num_nodes(&self) -> usize {
        self.body.num_nodes()
    }
}

impl graph::WithStartNode for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn start_node(&self) -> Self::Node {
        self.body.start_node()
    }
}

impl graph::WithSuccessors for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn successors(&self, node: Self::Node) -> <Self as GraphSuccessors<'_>>::Iter {
        self.body.successors(node)
    }
}

impl<'a, 'b, 'tcx> graph::GraphSuccessors<'b> for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Item = BasicBlock;
    type Iter = iter::Cloned<Successors<'b>>;
}

impl Deref for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Target = &'a Body<'tcx>;

    fn deref(&self) -> &Self::Target {
        &self.body
    }
}

CloneTypeFoldableAndLiftImpls! {
    Cache,
}
<commit_msg>Use a `SmallVec` for `Cache::predecessors`.<commit_after>use crate::ich::StableHashingContext;
use crate::mir::{BasicBlock, BasicBlockData, Body, LocalDecls, Location, Successors};
use rustc_data_structures::graph::dominators::{dominators, Dominators};
use rustc_data_structures::graph::{self, GraphPredecessors, GraphSuccessors};
use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
use rustc_index::vec::IndexVec;
use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
use smallvec::SmallVec;
use std::iter;
use std::ops::{Deref, DerefMut, Index, IndexMut};
use std::vec::IntoIter;

#[derive(Clone, Debug)]
pub struct Cache {
    // Typically 95%+ of the inner vectors have 4 or fewer elements.
    predecessors: Option<IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>>>,
}

impl rustc_serialize::Encodable for Cache {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        Encodable::encode(&(), s)
    }
}

impl rustc_serialize::Decodable for Cache {
    fn decode<D: Decoder>(d: &mut D) -> Result<Self, D::Error> {
        Decodable::decode(d).map(|_v: ()| Self::new())
    }
}

impl<'a> HashStable<StableHashingContext<'a>> for Cache {
    fn hash_stable(&self, _: &mut StableHashingContext<'a>, _: &mut StableHasher) {
        // Do nothing.
    }
}

impl Cache {
    pub fn new() -> Self {
        Self { predecessors: None }
    }

    pub fn invalidate_predecessors(&mut self) {
        // FIXME: consider being more fine-grained
        self.predecessors = None;
    }

    pub fn ensure_predecessors(&mut self, body: &Body<'_>) {
        if self.predecessors.is_none() {
            let mut result = IndexVec::from_elem(smallvec![], body.basic_blocks());
            for (bb, data) in body.basic_blocks().iter_enumerated() {
                if let Some(ref term) = data.terminator {
                    for &tgt in term.successors() {
                        result[tgt].push(bb);
                    }
                }
            }

            self.predecessors = Some(result)
        }
    }

    /// This will recompute the predecessors cache if it is not available
    // njn: typedef?
    fn predecessors(
        &mut self,
        body: &Body<'_>,
    ) -> &IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>> {
        self.ensure_predecessors(body);
        self.predecessors.as_ref().unwrap()
    }

    fn unwrap_predecessors_for(&self, bb: BasicBlock) -> &[BasicBlock] {
        &self.predecessors.as_ref().unwrap()[bb]
    }

    fn unwrap_predecessor_locations<'a>(
        &'a self,
        loc: Location,
        body: &'a Body<'a>,
    ) -> impl Iterator<Item = Location> + 'a {
        let if_zero_locations = if loc.statement_index == 0 {
            let predecessor_blocks = self.unwrap_predecessors_for(loc.block);
            let num_predecessor_blocks = predecessor_blocks.len();
            Some(
                (0..num_predecessor_blocks)
                    .map(move |i| predecessor_blocks[i])
                    .map(move |bb| body.terminator_loc(bb)),
            )
        } else {
            None
        };

        let if_not_zero_locations = if loc.statement_index == 0 {
            None
        } else {
            Some(Location { block: loc.block, statement_index: loc.statement_index - 1 })
        };

        if_zero_locations.into_iter().flatten().chain(if_not_zero_locations)
    }

    pub fn basic_blocks_mut<'a, 'tcx>(
        &mut self,
        body: &'a mut Body<'tcx>,
    ) -> &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        debug!("bbm: Clearing predecessors cache for body at: {:?}", body.span.data());
        self.invalidate_predecessors();
        &mut body.basic_blocks
    }

    pub fn basic_blocks_and_local_decls_mut<'a, 'tcx>(
        &mut self,
        body: &'a mut Body<'tcx>,
    ) -> (&'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &'a mut LocalDecls<'tcx>) {
        debug!("bbaldm: Clearing predecessors cache for body at: {:?}", body.span.data());
        self.invalidate_predecessors();
        (&mut body.basic_blocks, &mut body.local_decls)
    }
}

#[derive(Clone, Debug, HashStable, RustcEncodable, RustcDecodable, TypeFoldable)]
pub struct BodyAndCache<'tcx> {
    body: Body<'tcx>,
    cache: Cache,
}

impl BodyAndCache<'tcx> {
    pub fn new(body: Body<'tcx>) -> Self {
        Self { body, cache: Cache::new() }
    }
}

#[macro_export]
macro_rules! read_only {
    ($body:expr) => {{
        $body.ensure_predecessors();
        $body.unwrap_read_only()
    }};
}

impl BodyAndCache<'tcx> {
    pub fn ensure_predecessors(&mut self) {
        self.cache.ensure_predecessors(&self.body);
    }

    pub fn predecessors(&mut self) -> &IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>> {
        self.cache.predecessors(&self.body)
    }

    pub fn unwrap_read_only(&self) -> ReadOnlyBodyAndCache<'_, 'tcx> {
        ReadOnlyBodyAndCache::new(&self.body, &self.cache)
    }

    pub fn basic_blocks_mut(&mut self) -> &mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        self.cache.basic_blocks_mut(&mut self.body)
    }

    pub fn basic_blocks_and_local_decls_mut(
        &mut self,
    ) -> (&mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &mut LocalDecls<'tcx>) {
        self.cache.basic_blocks_and_local_decls_mut(&mut self.body)
    }
}

impl<'tcx> Index<BasicBlock> for BodyAndCache<'tcx> {
    type Output = BasicBlockData<'tcx>;

    fn index(&self, index: BasicBlock) -> &BasicBlockData<'tcx> {
        &self.body[index]
    }
}

impl<'tcx> IndexMut<BasicBlock> for BodyAndCache<'tcx> {
    fn index_mut(&mut self, index: BasicBlock) -> &mut Self::Output {
        &mut self.basic_blocks_mut()[index]
    }
}

impl<'tcx> Deref for BodyAndCache<'tcx> {
    type Target = Body<'tcx>;

    fn deref(&self) -> &Self::Target {
        &self.body
    }
}

impl<'tcx> DerefMut for BodyAndCache<'tcx> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.body
    }
}

#[derive(Copy, Clone, Debug)]
pub struct ReadOnlyBodyAndCache<'a, 'tcx> {
    body: &'a Body<'tcx>,
    cache: &'a Cache,
}

impl ReadOnlyBodyAndCache<'a, 'tcx> {
    fn new(body: &'a Body<'tcx>, cache: &'a Cache) -> Self {
        assert!(
            cache.predecessors.is_some(),
            "Cannot construct ReadOnlyBodyAndCache without computed predecessors"
        );
        Self { body, cache }
    }

    pub fn predecessors(&self) -> &IndexVec<BasicBlock, SmallVec<[BasicBlock; 4]>> {
        self.cache.predecessors.as_ref().unwrap()
    }

    pub fn predecessors_for(&self, bb: BasicBlock) -> &[BasicBlock] {
        self.cache.unwrap_predecessors_for(bb)
    }

    pub fn predecessor_locations(&self, loc: Location) -> impl Iterator<Item = Location> + '_ {
        self.cache.unwrap_predecessor_locations(loc, self.body)
    }

    pub fn basic_blocks(&self) -> &IndexVec<BasicBlock, BasicBlockData<'tcx>> {
        &self.body.basic_blocks
    }

    pub fn dominators(&self) -> Dominators<BasicBlock> {
        dominators(self)
    }
}

impl graph::DirectedGraph for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Node = BasicBlock;
}

impl graph::GraphPredecessors<'graph> for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Item = BasicBlock;
    type Iter = IntoIter<BasicBlock>;
}

impl graph::WithPredecessors for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn predecessors(&self, node: Self::Node) -> <Self as GraphPredecessors<'_>>::Iter {
        self.cache.unwrap_predecessors_for(node).to_vec().into_iter()
    }
}

impl graph::WithNumNodes for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn num_nodes(&self) -> usize {
        self.body.num_nodes()
    }
}

impl graph::WithStartNode for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn start_node(&self) -> Self::Node {
        self.body.start_node()
    }
}

impl graph::WithSuccessors for ReadOnlyBodyAndCache<'a, 'tcx> {
    fn successors(&self, node: Self::Node) -> <Self as GraphSuccessors<'_>>::Iter {
        self.body.successors(node)
    }
}

impl<'a, 'b, 'tcx> graph::GraphSuccessors<'b> for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Item = BasicBlock;
    type Iter = iter::Cloned<Successors<'b>>;
}

impl Deref for ReadOnlyBodyAndCache<'a, 'tcx> {
    type Target = &'a Body<'tcx>;

    fn deref(&self) -> &Self::Target {
        &self.body
    }
}

CloneTypeFoldableAndLiftImpls! {
    Cache,
}
<|endoftext|>


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

//! Sampling from random distributions

// Some implementations use the Ziggurat method
// https://en.wikipedia.org/wiki/Ziggurat_algorithm
//
// The version used here is ZIGNOR [Doornik 2005, "An Improved
// Ziggurat Method to Generate Normal Random Samples"] which is slower
// (about double, it generates an extra random number) than the
// canonical version [Marsaglia & Tsang 2000, "The Ziggurat Method for
// Generating Random Variables"], but more robust. If one wanted, one
// could implement VIZIGNOR the ZIGNOR paper for more speed.

use num;
use rand::{Rng,Rand};

mod ziggurat_tables;

// inlining should mean there is no performance penalty for this
#[inline]
fn ziggurat<R:Rng>(rng: &mut R,
                   center_u: bool,
                   X: ziggurat_tables::ZigTable,
                   F: ziggurat_tables::ZigTable,
                   F_DIFF: ziggurat_tables::ZigTable,
                   pdf: &'static fn(f64) -> f64, // probability density function
                   zero_case: &'static fn(&mut R, f64) -> f64) -> f64 {
    loop {
        let u = if center_u {2.0 * rng.gen() - 1.0} else {rng.gen()};
        let i: uint = rng.gen::<uint>() & 0xff;
        let x = u * X[i];

        let test_x = if center_u {num::abs(x)} else {x};

        // algebraically equivalent to |u| < X[i+1]/X[i] (or u < X[i+1]/X[i])
        if test_x < X[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if F[i+1] + F_DIFF[i+1] * rng.gen() < pdf(x) {
            return x;
        }
    }
}

/// A wrapper around an `f64` to generate N(0, 1) random numbers (a.k.a.  a
/// standard normal, or Gaussian). Multiplying the generated values by the
/// desired standard deviation `sigma` then adding the desired mean `mu` will
/// give N(mu, sigma^2) distributed random numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ~~~
/// use core::rand::distributions::StandardNormal;
///
/// fn main() {
///     let normal = 2.0 + (*rand::random::<StandardNormal>()) * 3.0;
///     printfln!("%f is from a N(2, 9) distribution", normal)
/// }
/// ~~~
pub struct StandardNormal(f64);

impl Rand for StandardNormal {
    fn rand<R:Rng>(rng: &mut R) -> StandardNormal {
        #[inline]
        fn pdf(x: f64) -> f64 {
            ((-x*x/2.0) as f64).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            // XXX infinities?
            while -2.0 * y < x * x {
                x = rng.gen::<f64>().ln() / ziggurat_tables::ZIG_NORM_R;
                y = rng.gen::<f64>().ln();
            }

            if u < 0.0 { x - ziggurat_tables::ZIG_NORM_R } else { ziggurat_tables::ZIG_NORM_R - x }
        }

        StandardNormal(ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F, &ziggurat_tables::ZIG_NORM_F_DIFF,
            pdf, zero_case))
    }
}

/// A wrapper around an `f64` to generate Exp(1) random numbers. Dividing by
/// the desired rate `lambda` will give Exp(lambda) distributed random
/// numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ~~~
/// use core::rand::distributions::Exp1;
///
/// fn main() {
///     let exp2 = (*rand::random::<Exp1>()) * 0.5;
///     printfln!("%f is from a Exp(2) distribution", exp2);
/// }
/// ~~~
pub struct Exp1(f64);

// This could be done via `-rng.gen::<f64>().ln()` but that is slower.
impl Rand for Exp1 {
    #[inline]
    fn rand<R:Rng>(rng: &mut R) -> Exp1 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, _u: f64) -> f64 {
            ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
        }

        Exp1(ziggurat(rng, false,
                      &ziggurat_tables::ZIG_EXP_X,
                      &ziggurat_tables::ZIG_EXP_F, &ziggurat_tables::ZIG_EXP_F_DIFF,
                      pdf, zero_case))
    }
}
<commit_msg>doc: convert remaining uses of core:: to std::.<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Sampling from random distributions

// Some implementations use the Ziggurat method
// https://en.wikipedia.org/wiki/Ziggurat_algorithm
//
// The version used here is ZIGNOR [Doornik 2005, "An Improved
// Ziggurat Method to Generate Normal Random Samples"] which is slower
// (about double, it generates an extra random number) than the
// canonical version [Marsaglia & Tsang 2000, "The Ziggurat Method for
// Generating Random Variables"], but more robust. If one wanted, one
// could implement VIZIGNOR the ZIGNOR paper for more speed.

use num;
use rand::{Rng,Rand};

mod ziggurat_tables;

// inlining should mean there is no performance penalty for this
#[inline]
fn ziggurat<R:Rng>(rng: &mut R,
                   center_u: bool,
                   X: ziggurat_tables::ZigTable,
                   F: ziggurat_tables::ZigTable,
                   F_DIFF: ziggurat_tables::ZigTable,
                   pdf: &'static fn(f64) -> f64, // probability density function
                   zero_case: &'static fn(&mut R, f64) -> f64) -> f64 {
    loop {
        let u = if center_u {2.0 * rng.gen() - 1.0} else {rng.gen()};
        let i: uint = rng.gen::<uint>() & 0xff;
        let x = u * X[i];

        let test_x = if center_u {num::abs(x)} else {x};

        // algebraically equivalent to |u| < X[i+1]/X[i] (or u < X[i+1]/X[i])
        if test_x < X[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if F[i+1] + F_DIFF[i+1] * rng.gen() < pdf(x) {
            return x;
        }
    }
}

/// A wrapper around an `f64` to generate N(0, 1) random numbers (a.k.a.  a
/// standard normal, or Gaussian). Multiplying the generated values by the
/// desired standard deviation `sigma` then adding the desired mean `mu` will
/// give N(mu, sigma^2) distributed random numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ~~~
/// use std::rand::distributions::StandardNormal;
///
/// fn main() {
///     let normal = 2.0 + (*rand::random::<StandardNormal>()) * 3.0;
///     printfln!("%f is from a N(2, 9) distribution", normal)
/// }
/// ~~~
pub struct StandardNormal(f64);

impl Rand for StandardNormal {
    fn rand<R:Rng>(rng: &mut R) -> StandardNormal {
        #[inline]
        fn pdf(x: f64) -> f64 {
            ((-x*x/2.0) as f64).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, u: f64) -> f64 {
            // compute a random number in the tail by hand

            // strange initial conditions, because the loop is not
            // do-while, so the condition should be true on the first
            // run, they get overwritten anyway (0 < 1, so these are
            // good).
            let mut x = 1.0f64;
            let mut y = 0.0f64;

            // XXX infinities?
            while -2.0 * y < x * x {
                x = rng.gen::<f64>().ln() / ziggurat_tables::ZIG_NORM_R;
                y = rng.gen::<f64>().ln();
            }

            if u < 0.0 { x - ziggurat_tables::ZIG_NORM_R } else { ziggurat_tables::ZIG_NORM_R - x }
        }

        StandardNormal(ziggurat(
            rng,
            true, // this is symmetric
            &ziggurat_tables::ZIG_NORM_X,
            &ziggurat_tables::ZIG_NORM_F, &ziggurat_tables::ZIG_NORM_F_DIFF,
            pdf, zero_case))
    }
}

/// A wrapper around an `f64` to generate Exp(1) random numbers. Dividing by
/// the desired rate `lambda` will give Exp(lambda) distributed random
/// numbers.
///
/// Note that this has to be unwrapped before use as an `f64` (using either
/// `*` or `cast::transmute` is safe).
///
/// # Example
///
/// ~~~
/// use std::rand::distributions::Exp1;
///
/// fn main() {
///     let exp2 = (*rand::random::<Exp1>()) * 0.5;
///     printfln!("%f is from a Exp(2) distribution", exp2);
/// }
/// ~~~
pub struct Exp1(f64);

// This could be done via `-rng.gen::<f64>().ln()` but that is slower.
impl Rand for Exp1 {
    #[inline]
    fn rand<R:Rng>(rng: &mut R) -> Exp1 {
        #[inline]
        fn pdf(x: f64) -> f64 {
            (-x).exp()
        }
        #[inline]
        fn zero_case<R:Rng>(rng: &mut R, _u: f64) -> f64 {
            ziggurat_tables::ZIG_EXP_R - rng.gen::<f64>().ln()
        }

        Exp1(ziggurat(rng, false,
                      &ziggurat_tables::ZIG_EXP_X,
                      &ziggurat_tables::ZIG_EXP_F, &ziggurat_tables::ZIG_EXP_F_DIFF,
                      pdf, zero_case))
    }
}
<|endoftext|>


<commit_before>//! HTTP response builder.

use std::{
    cell::{Ref, RefMut},
    fmt,
    future::Future,
    pin::Pin,
    str,
    task::{Context, Poll},
};

use bytes::Bytes;
use futures_core::Stream;

use crate::{
    body::{Body, BodyStream, ResponseBody},
    error::Error,
    extensions::Extensions,
    header::{IntoHeaderPair, IntoHeaderValue},
    http::{header, Error as HttpError, StatusCode},
    message::{BoxedResponseHead, ConnectionType, ResponseHead},
    Response,
};

/// An HTTP response builder.
///
/// This type can be used to construct an instance of `Response` using a builder pattern.
pub struct ResponseBuilder {
    head: Option<BoxedResponseHead>,
    err: Option<HttpError>,
}

impl ResponseBuilder {
    #[inline]
    /// Create response builder
    pub fn new(status: StatusCode) -> Self {
        ResponseBuilder {
            head: Some(BoxedResponseHead::new(status)),
            err: None,
        }
    }

    /// Set HTTP status code of this response.
    #[inline]
    pub fn status(&mut self, status: StatusCode) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.status = status;
        }
        self
    }

    /// Insert a header, replacing any that were set with an equivalent field name.
    ///
    /// ```
    /// # use actix_http::Response;
    /// use actix_http::http::{header, StatusCode};
    ///
    /// Response::build(StatusCode::OK)
    ///     .insert_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .insert_header(("X-TEST", "value"))
    ///     .finish();
    /// ```
    pub fn insert_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => {
                    parts.headers.insert(key, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Append a header, keeping any that were set with an equivalent field name.
    ///
    /// ```
    /// # use actix_http::Response;
    /// use actix_http::http::{header, StatusCode};
    ///
    /// Response::build(StatusCode::OK)
    ///     .append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .append_header(("X-TEST", "value1"))
    ///     .append_header(("X-TEST", "value2"))
    ///     .finish();
    /// ```
    pub fn append_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => parts.headers.append(key, value),
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Set the custom reason for the response.
    #[inline]
    pub fn reason(&mut self, reason: &'static str) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.reason = Some(reason);
        }
        self
    }

    /// Set connection type to KeepAlive
    #[inline]
    pub fn keep_alive(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::KeepAlive);
        }
        self
    }

    /// Set connection type to Upgrade
    #[inline]
    pub fn upgrade<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Upgrade);
        }

        if let Ok(value) = value.try_into_value() {
            self.insert_header((header::UPGRADE, value));
        }

        self
    }

    /// Force close connection, even if it is marked as keep-alive
    #[inline]
    pub fn force_close(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Close);
        }
        self
    }

    /// Disable chunked transfer encoding for HTTP/1.1 streaming responses.
    #[inline]
    pub fn no_chunking(&mut self, len: u64) -> &mut Self {
        let mut buf = itoa::Buffer::new();
        self.insert_header((header::CONTENT_LENGTH, buf.format(len)));

        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.no_chunking(true);
        }
        self
    }

    /// Set response content type.
    #[inline]
    pub fn content_type<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match value.try_into_value() {
                Ok(value) => {
                    parts.headers.insert(header::CONTENT_TYPE, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }
        self
    }

    /// Responses extensions
    #[inline]
    pub fn extensions(&self) -> Ref<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow()
    }

    /// Mutable reference to a the response's extensions
    #[inline]
    pub fn extensions_mut(&mut self) -> RefMut<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow_mut()
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn body<B: Into<Body>>(&mut self, body: B) -> Response<Body> {
        self.message_body(body.into())
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    pub fn message_body<B>(&mut self, body: B) -> Response<B> {
        if let Some(e) = self.err.take() {
            return Response::from(Error::from(e)).into_body();
        }

        let response = self.head.take().expect("cannot reuse response builder");

        Response {
            head: response,
            body: ResponseBody::Body(body),
            error: None,
        }
    }

    /// Set a streaming body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn streaming<S, E>(&mut self, stream: S) -> Response<Body>
    where
        S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
        E: Into<Error> + 'static,
    {
        self.body(Body::from_message(BodyStream::new(stream)))
    }

    /// Set an empty body and generate `Response`
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn finish(&mut self) -> Response<Body> {
        self.body(Body::Empty)
    }

    /// This method construct new `ResponseBuilder`
    pub fn take(&mut self) -> ResponseBuilder {
        ResponseBuilder {
            head: self.head.take(),
            err: self.err.take(),
        }
    }
}

#[inline]
fn parts<'a>(
    parts: &'a mut Option<BoxedResponseHead>,
    err: &Option<HttpError>,
) -> Option<&'a mut ResponseHead> {
    if err.is_some() {
        return None;
    }
    parts.as_mut().map(|r| &mut **r)
}

/// Convert `Response` to a `ResponseBuilder`. Body get dropped.
impl<B> From<Response<B>> for ResponseBuilder {
    fn from(res: Response<B>) -> ResponseBuilder {
        ResponseBuilder {
            head: Some(res.head),
            err: None,
        }
    }
}

/// Convert `ResponseHead` to a `ResponseBuilder`
impl<'a> From<&'a ResponseHead> for ResponseBuilder {
    fn from(head: &'a ResponseHead) -> ResponseBuilder {
        let mut msg = BoxedResponseHead::new(head.status);
        msg.version = head.version;
        msg.reason = head.reason;

        for (k, v) in head.headers.iter() {
            msg.headers.append(k.clone(), v.clone());
        }

        msg.no_chunking(!head.chunked());

        ResponseBuilder {
            head: Some(msg),
            err: None,
        }
    }
}

impl Future for ResponseBuilder {
    type Output = Result<Response<Body>, Error>;

    fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(Ok(self.finish()))
    }
}

impl fmt::Debug for ResponseBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let head = self.head.as_ref().unwrap();

        let res = writeln!(
            f,
            "\nResponseBuilder {:?} {}{}",
            head.version,
            head.status,
            head.reason.unwrap_or(""),
        );
        let _ = writeln!(f, "  headers:");
        for (key, val) in head.headers.iter() {
            let _ = writeln!(f, "    {:?}: {:?}", key, val);
        }
        res
    }
}

#[cfg(test)]
mod tests {
    use bytes::BytesMut;

    use super::*;
    use crate::body::Body;
    use crate::http::header::{HeaderName, HeaderValue, CONTENT_TYPE, COOKIE};

    #[test]
    fn test_debug() {
        let resp = Response::build(StatusCode::OK)
            .append_header((COOKIE, HeaderValue::from_static("cookie1=value1; ")))
            .append_header((COOKIE, HeaderValue::from_static("cookie2=value2; ")))
            .finish();
        let dbg = format!("{:?}", resp);
        assert!(dbg.contains("Response"));
    }

    #[test]
    fn test_basic_builder() {
        let resp = Response::build(StatusCode::OK)
            .insert_header(("X-TEST", "value"))
            .finish();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn test_upgrade() {
        let resp = Response::build(StatusCode::OK)
            .upgrade("websocket")
            .finish();
        assert!(resp.upgrade());
        assert_eq!(
            resp.headers().get(header::UPGRADE).unwrap(),
            HeaderValue::from_static("websocket")
        );
    }

    #[test]
    fn test_force_close() {
        let resp = Response::build(StatusCode::OK).force_close().finish();
        assert!(!resp.keep_alive())
    }

    #[test]
    fn test_content_type() {
        let resp = Response::build(StatusCode::OK)
            .content_type("text/plain")
            .body(Body::Empty);
        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "text/plain")
    }

    #[test]
    fn test_into_response() {
        let resp: Response<Body> = "test".into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain; charset=utf-8")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let resp: Response<Body> = b"test".as_ref().into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("application/octet-stream")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let resp: Response<Body> = "test".to_owned().into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain; charset=utf-8")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let resp: Response<Body> = (&"test".to_owned()).into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("text/plain; charset=utf-8")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let b = Bytes::from_static(b"test");
        let resp: Response<Body> = b.into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("application/octet-stream")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let b = Bytes::from_static(b"test");
        let resp: Response<Body> = b.into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("application/octet-stream")
        );
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");

        let b = BytesMut::from("test");
        let resp: Response<Body> = b.into();
        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(
            resp.headers().get(CONTENT_TYPE).unwrap(),
            HeaderValue::from_static("application/octet-stream")
        );

        assert_eq!(resp.status(), StatusCode::OK);
        assert_eq!(resp.body().get_ref(), b"test");
    }

    #[test]
    fn test_into_builder() {
        let mut resp: Response<Body> = "test".into();
        assert_eq!(resp.status(), StatusCode::OK);

        resp.headers_mut().insert(
            HeaderName::from_static("cookie"),
            HeaderValue::from_static("cookie1=val100"),
        );

        let mut builder: ResponseBuilder = resp.into();
        let resp = builder.status(StatusCode::BAD_REQUEST).finish();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let cookie = resp.headers().get_all("Cookie").next().unwrap();
        assert_eq!(cookie.to_str().unwrap(), "cookie1=val100");
    }

    #[test]
    fn response_builder_header_insert_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header(("Content-Type", "application/octet-stream"));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_insert_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_append_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header(("Content-Type", "application/octet-stream"));
        res.append_header(("Content-Type", "application/json"));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }

    #[test]
    fn response_builder_header_append_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }
}
<commit_msg>add some doc examples to response builder<commit_after>//! HTTP response builder.

use std::{
    cell::{Ref, RefMut},
    fmt,
    future::Future,
    pin::Pin,
    str,
    task::{Context, Poll},
};

use bytes::Bytes;
use futures_core::Stream;

use crate::{
    body::{Body, BodyStream, ResponseBody},
    error::Error,
    extensions::Extensions,
    header::{IntoHeaderPair, IntoHeaderValue},
    http::{header, Error as HttpError, StatusCode},
    message::{BoxedResponseHead, ConnectionType, ResponseHead},
    Response,
};

/// An HTTP response builder.
///
/// Used to construct an instance of `Response` using a builder pattern. Response builders are often
/// created using [`Response::build`].
///
/// # Examples
/// ```
/// use actix_http::{Response, ResponseBuilder, body, http::StatusCode, http::header};
///
/// # actix_rt::System::new().block_on(async {
/// let mut res: Response<_> = Response::build(StatusCode::OK)
///     .content_type(mime::APPLICATION_JSON)
///     .insert_header((header::SERVER, "my-app/1.0"))
///     .append_header((header::SET_COOKIE, "a=1"))
///     .append_header((header::SET_COOKIE, "b=2"))
///     .body("1234");
///
/// assert_eq!(res.status(), StatusCode::OK);
/// assert_eq!(body::to_bytes(res.take_body()).await.unwrap(), &b"1234"[..]);
///
/// assert!(res.headers().contains_key("server"));
/// assert_eq!(res.headers().get_all("set-cookie").count(), 2);
/// # })
/// ```
pub struct ResponseBuilder {
    head: Option<BoxedResponseHead>,
    err: Option<HttpError>,
}

impl ResponseBuilder {
    /// Create response builder
    ///
    /// # Examples
    /// ```
    /// use actix_http::{Response, ResponseBuilder, http::StatusCode};
    ///
    /// let res: Response<_> = ResponseBuilder::default().finish();
    /// assert_eq!(res.status(), StatusCode::OK);
    /// ```
    #[inline]
    pub fn new(status: StatusCode) -> Self {
        ResponseBuilder {
            head: Some(BoxedResponseHead::new(status)),
            err: None,
        }
    }

    /// Set HTTP status code of this response.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::StatusCode};
    ///
    /// let res = ResponseBuilder::default().status(StatusCode::NOT_FOUND).finish();
    /// assert_eq!(res.status(), StatusCode::NOT_FOUND);
    /// ```
    #[inline]
    pub fn status(&mut self, status: StatusCode) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.status = status;
        }
        self
    }

    /// Insert a header, replacing any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .insert_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .insert_header(("X-TEST", "value"))
    ///     .finish();
    ///
    /// assert!(res.headers().contains_key("content-type"));
    /// assert!(res.headers().contains_key("x-test"));
    /// ```
    pub fn insert_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => {
                    parts.headers.insert(key, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Append a header, keeping any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .append_header(("X-TEST", "value1"))
    ///     .append_header(("X-TEST", "value2"))
    ///     .finish();
    ///
    /// assert_eq!(res.headers().get_all("content-type").count(), 1);
    /// assert_eq!(res.headers().get_all("x-test").count(), 2);
    /// ```
    pub fn append_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => parts.headers.append(key, value),
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Set the custom reason for the response.
    #[inline]
    pub fn reason(&mut self, reason: &'static str) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.reason = Some(reason);
        }
        self
    }

    /// Set connection type to KeepAlive
    #[inline]
    pub fn keep_alive(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::KeepAlive);
        }
        self
    }

    /// Set connection type to Upgrade
    #[inline]
    pub fn upgrade<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Upgrade);
        }

        if let Ok(value) = value.try_into_value() {
            self.insert_header((header::UPGRADE, value));
        }

        self
    }

    /// Force close connection, even if it is marked as keep-alive
    #[inline]
    pub fn force_close(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Close);
        }
        self
    }

    /// Disable chunked transfer encoding for HTTP/1.1 streaming responses.
    #[inline]
    pub fn no_chunking(&mut self, len: u64) -> &mut Self {
        let mut buf = itoa::Buffer::new();
        self.insert_header((header::CONTENT_LENGTH, buf.format(len)));

        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.no_chunking(true);
        }
        self
    }

    /// Set response content type.
    #[inline]
    pub fn content_type<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match value.try_into_value() {
                Ok(value) => {
                    parts.headers.insert(header::CONTENT_TYPE, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }
        self
    }

    /// Responses extensions
    #[inline]
    pub fn extensions(&self) -> Ref<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow()
    }

    /// Mutable reference to a the response's extensions
    #[inline]
    pub fn extensions_mut(&mut self) -> RefMut<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow_mut()
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn body<B: Into<Body>>(&mut self, body: B) -> Response<Body> {
        self.message_body(body.into())
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    pub fn message_body<B>(&mut self, body: B) -> Response<B> {
        if let Some(e) = self.err.take() {
            return Response::from(Error::from(e)).into_body();
        }

        let response = self.head.take().expect("cannot reuse response builder");

        Response {
            head: response,
            body: ResponseBody::Body(body),
            error: None,
        }
    }

    /// Set a streaming body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn streaming<S, E>(&mut self, stream: S) -> Response<Body>
    where
        S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
        E: Into<Error> + 'static,
    {
        self.body(Body::from_message(BodyStream::new(stream)))
    }

    /// Set an empty body and generate `Response`
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn finish(&mut self) -> Response<Body> {
        self.body(Body::Empty)
    }

    /// This method construct new `ResponseBuilder`
    pub fn take(&mut self) -> ResponseBuilder {
        ResponseBuilder {
            head: self.head.take(),
            err: self.err.take(),
        }
    }
}

#[inline]
fn parts<'a>(
    parts: &'a mut Option<BoxedResponseHead>,
    err: &Option<HttpError>,
) -> Option<&'a mut ResponseHead> {
    if err.is_some() {
        return None;
    }
    parts.as_mut().map(|r| &mut **r)
}

impl Default for ResponseBuilder {
    fn default() -> Self {
        Self::new(StatusCode::OK)
    }
}

/// Convert `Response` to a `ResponseBuilder`. Body get dropped.
impl<B> From<Response<B>> for ResponseBuilder {
    fn from(res: Response<B>) -> ResponseBuilder {
        ResponseBuilder {
            head: Some(res.head),
            err: None,
        }
    }
}

/// Convert `ResponseHead` to a `ResponseBuilder`
impl<'a> From<&'a ResponseHead> for ResponseBuilder {
    fn from(head: &'a ResponseHead) -> ResponseBuilder {
        let mut msg = BoxedResponseHead::new(head.status);
        msg.version = head.version;
        msg.reason = head.reason;

        for (k, v) in head.headers.iter() {
            msg.headers.append(k.clone(), v.clone());
        }

        msg.no_chunking(!head.chunked());

        ResponseBuilder {
            head: Some(msg),
            err: None,
        }
    }
}

impl Future for ResponseBuilder {
    type Output = Result<Response<Body>, Error>;

    fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(Ok(self.finish()))
    }
}

impl fmt::Debug for ResponseBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let head = self.head.as_ref().unwrap();

        let res = writeln!(
            f,
            "\nResponseBuilder {:?} {}{}",
            head.version,
            head.status,
            head.reason.unwrap_or(""),
        );
        let _ = writeln!(f, "  headers:");
        for (key, val) in head.headers.iter() {
            let _ = writeln!(f, "    {:?}: {:?}", key, val);
        }
        res
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::body::Body;
    use crate::http::header::{HeaderName, HeaderValue, CONTENT_TYPE};

    #[test]
    fn test_basic_builder() {
        let resp = Response::build(StatusCode::OK)
            .insert_header(("X-TEST", "value"))
            .finish();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn test_upgrade() {
        let resp = Response::build(StatusCode::OK)
            .upgrade("websocket")
            .finish();
        assert!(resp.upgrade());
        assert_eq!(
            resp.headers().get(header::UPGRADE).unwrap(),
            HeaderValue::from_static("websocket")
        );
    }

    #[test]
    fn test_force_close() {
        let resp = Response::build(StatusCode::OK).force_close().finish();
        assert!(!resp.keep_alive())
    }

    #[test]
    fn test_content_type() {
        let resp = Response::build(StatusCode::OK)
            .content_type("text/plain")
            .body(Body::Empty);
        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "text/plain")
    }

    #[test]
    fn test_into_builder() {
        let mut resp: Response<Body> = "test".into();
        assert_eq!(resp.status(), StatusCode::OK);

        resp.headers_mut().insert(
            HeaderName::from_static("cookie"),
            HeaderValue::from_static("cookie1=val100"),
        );

        let mut builder: ResponseBuilder = resp.into();
        let resp = builder.status(StatusCode::BAD_REQUEST).finish();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let cookie = resp.headers().get_all("Cookie").next().unwrap();
        assert_eq!(cookie.to_str().unwrap(), "cookie1=val100");
    }

    #[test]
    fn response_builder_header_insert_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header(("Content-Type", "application/octet-stream"));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_insert_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_append_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header(("Content-Type", "application/octet-stream"));
        res.append_header(("Content-Type", "application/json"));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }

    #[test]
    fn response_builder_header_append_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }
}
<|endoftext|>