<commit_before>//! The thread-local state.

use std::{mem, thread};
use std::cell::RefCell;
use {global, hazard, guard};
use garbage::Garbage;

thread_local! {
    /// The state of this thread.
    static STATE: RefCell<State> = RefCell::new(State::default());
}

/// Add new garbage to be deleted.
///
/// This garbage is pushed to a thread-local queue. When enough garbage is accumulated in the
/// thread, it is exported to the global state.
pub fn add_garbage(garbage: Garbage) {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    if STATE.state() == thread::LocalKeyState::Destroyed {
        // The state was deinitialized, so we must rely on the global state for queueing garbage.
        global::export_garbage(vec![garbage]);
    } else {
        // Add the garbage.
        if STATE.with(|s| s.borrow_mut().add_garbage(garbage)) {
            // The local state exported garbage to the global state, hence we must tick in order to
            // ensure that the garbage is periodically collected.
            global::tick();
        }
    }
}

/// Get a blocked hazard.
///
/// If possible, this will simply pop one of the thread-local cache of hazards. Otherwise, one must
/// be registered in the global state.
///
/// # Fence
///
/// This does not fence, and you must thus be careful with updating the value afterwards, as
/// reordering can happen, meaning that the hazard has not been blocked yet.
pub fn get_hazard() -> hazard::Writer {
    if STATE.state() == thread::LocalKeyState::Destroyed {
        // The state was deinitialized, so we must rely on the global state for creating new
        // hazards.
        global::create_hazard()
    } else {
        STATE.with(|s| s.borrow_mut().get_hazard())
    }
}

/// Free a hazard.
///
/// This frees a hazard to the thread-local cache of hazards.
///
/// It is important that the hazard is **not** in blocked state, as such thing can cause infinite
/// looping.
///
/// # Panics
///
/// This might panic in debug mode if the hazard given is in blocked state.
pub fn free_hazard(hazard: hazard::Writer) {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    debug_assert!(!hazard.is_blocked(), "Freeing a blocked hazards. See docs.");

    if STATE.state() == thread::LocalKeyState::Destroyed {
        // Since the state was deinitialized, we cannot store it for later reuse, so we are forced
        // to simply kill the hazard.
        hazard.kill();
    } else {
        STATE.with(|s| s.borrow_mut().free_hazard(hazard));
    }
}

/// Export the garbage of this thread to the global state.
///
/// This is useful for propagating accumulated garbage such that it can be destroyed by the next
/// garbage collection.
pub fn export_garbage() {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    // We can only export when the TLS variable isn't destroyed. Otherwise, there would be nothing
    // to export!
    if STATE.state() != thread::LocalKeyState::Destroyed {
        STATE.with(|s| s.borrow_mut().export_garbage());
        // We tick after the state is no longer reserved, as the tick could potentially call
        // destructor that access the TLS variable.
        global::tick();
    }
}

/// A thread-local state.
#[derive(Default)]
struct State {
    /// The cached garbage waiting to be exported to the global state.
    garbage: Vec<Garbage>,
    /// The cache of currently available hazards.
    ///
    /// We maintain this cache to avoid the performance hit of creating new hazards.
    ///
    /// The hazards in this vector are not necessarily in state "free". Only when a sufficient
    /// amount of available hazards has accumulated, they will be set to free. This means that we
    /// don't have to reset the state of a hazard after usage, giving a quite significant speed-up.
    available_hazards: Vec<hazard::Writer>,
    /// The hazards in the cache before this index are free.
    ///
    /// This number keeps track what hazards in `self.available_hazard` are set to state "free".
    /// Before this index, every hazard must be set to "free".
    ///
    /// It is useful for knowing when to free the hazards to allow garbage collection.
    available_hazards_free_before: usize,
}

impl State {
    /// Get the number of hazards in the cache which are not in state "free".
    fn non_free_hazards(&self) -> usize {
        self.available_hazards.len() - self.available_hazards_free_before
    }

    /// See `get_hazard()`.
    fn get_hazard(&mut self) -> hazard::Writer {
        // Check if there is hazards in the cache.
        if let Some(hazard) = self.available_hazards.pop() {
            // There is; we don't need to create a new hazard.

            // Since the hazard popped from the cache is not blocked, we must block the hazard to
            // satisfy the requirements of this function.
            hazard.block();
            hazard
        } else {
            // There is not; we must create a new hazard.
            global::create_hazard()
        }
    }

    /// See `free_hazard()`.
    fn free_hazard(&mut self, hazard: hazard::Writer) {
        /// The maximal amount of hazards before cleaning up.
        ///
        /// With "cleaning up" we mean setting the state of the hazards to "free" in order to allow
        /// garbage collection of the object it is currently protecting.
        const MAX_NON_FREE_HAZARDS: usize = 16;

        // Push the given hazard to the cache.
        self.available_hazards.push(hazard);

        // Check if we exceeded the limit.
        if self.non_free_hazards() > MAX_NON_FREE_HAZARDS {
            // We did; we must now set the non-free hazards to "free".
            for i in &self.available_hazards[self.available_hazards_free_before..] {
                i.set(hazard::State::Free);
            }

            // Update the counter such that we mark the new hazards set to "free".
            self.available_hazards_free_before = self.available_hazards.len();
        }
    }

    /// Queues garbage to destroy.
    ///
    /// Eventually the added garbage will be exported to the global state through
    /// `global::add_garbage()`.
    ///
    /// See `add_garbage` for more information.
    ///
    /// When this happens (i.e. the global state gets the garbage), it returns `true`. Otherwise,
    /// it returns `false`.
    fn add_garbage(&mut self, garbage: Garbage) -> bool {
        /// The maximal amount of garbage before exportation to the global state.
        const MAX_GARBAGE: usize = 128;

        // Push the garbage to the cache of garbage.
        self.garbage.push(garbage);

        // Export the garbage if it exceeds the limit.
        // TODO: use memory instead of items as a metric.
        if self.garbage.len() > MAX_GARBAGE {
            self.export_garbage();
            true
        } else { false }
    }

    /// See `export_garbage()` for more information.
    fn export_garbage(&mut self) {
        // Clear the vector and export the garbage.
        global::export_garbage(mem::replace(&mut self.garbage, Vec::new()));
    }
}

impl Drop for State {
    fn drop(&mut self) {
        // Clear every hazard to "dead" state.
        for hazard in self.available_hazards.drain(..) {
            hazard.kill();
        }

        // The thread is exiting, thus we must export the garbage to the global state to avoid
        // memory leaks. It is very important that this does indeed not tick, as causing garbage
        // collection means accessing RNG state, a TLS variable, which cannot be done when, we are
        // here, after it has deinitialized.
        // TODO: Figure out a way we can tick anyway.
        self.export_garbage();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use garbage::Garbage;
    use hazard;
    use std::thread;

    #[test]
    fn dtor_runs() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = Box::new(0);
            let h = get_hazard();
            h.set(hazard::State::Protect(&*b));
            add_garbage(Garbage::new(&*b, dtor));
            ::gc();
            assert_eq!(*b, 0);
            ::gc();
            h.set(hazard::State::Free);
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[test]
    fn dtor_runs_cross_thread() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = Box::new(0);
            let bptr = &*b as *const _ as usize;
            let h = thread::spawn(move || {
                let h = get_hazard();
                h.set(hazard::State::Protect(bptr as *const u8));
                h
            }).join().unwrap();
            add_garbage(Garbage::new(&*b, dtor));
            ::gc();
            assert_eq!(*b, 0);
            ::gc();
            h.set(hazard::State::Free);
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[test]
    fn kill_hazards() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = thread::spawn(move || {
                let b = Box::new(0);
                let h = get_hazard();
                h.set(hazard::State::Protect(&*b));
                add_garbage(Garbage::new(&*b, dtor));
                ::gc();
                assert_eq!(*b, 0);
                b
            }).join().unwrap();
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic]
    fn debug_free_blocked() {
        use std::mem;

        let (writer, reader) = hazard::create();
        mem::forget(reader);

        free_hazard(writer);
    }
}
<commit_msg>Add tests for `conc::local`'s hazard clearing behavior.<commit_after>//! The thread-local state.

use std::{mem, thread};
use std::cell::RefCell;
use {global, hazard, guard};
use garbage::Garbage;

thread_local! {
    /// The state of this thread.
    static STATE: RefCell<State> = RefCell::new(State::default());
}

/// Add new garbage to be deleted.
///
/// This garbage is pushed to a thread-local queue. When enough garbage is accumulated in the
/// thread, it is exported to the global state.
pub fn add_garbage(garbage: Garbage) {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    if STATE.state() == thread::LocalKeyState::Destroyed {
        // The state was deinitialized, so we must rely on the global state for queueing garbage.
        global::export_garbage(vec![garbage]);
    } else {
        // Add the garbage.
        if STATE.with(|s| s.borrow_mut().add_garbage(garbage)) {
            // The local state exported garbage to the global state, hence we must tick in order to
            // ensure that the garbage is periodically collected.
            global::tick();
        }
    }
}

/// Get a blocked hazard.
///
/// If possible, this will simply pop one of the thread-local cache of hazards. Otherwise, one must
/// be registered in the global state.
///
/// # Fence
///
/// This does not fence, and you must thus be careful with updating the value afterwards, as
/// reordering can happen, meaning that the hazard has not been blocked yet.
pub fn get_hazard() -> hazard::Writer {
    if STATE.state() == thread::LocalKeyState::Destroyed {
        // The state was deinitialized, so we must rely on the global state for creating new
        // hazards.
        global::create_hazard()
    } else {
        STATE.with(|s| s.borrow_mut().get_hazard())
    }
}

/// Free a hazard.
///
/// This frees a hazard to the thread-local cache of hazards.
///
/// It is important that the hazard is **not** in blocked state, as such thing can cause infinite
/// looping.
///
/// # Panics
///
/// This might panic in debug mode if the hazard given is in blocked state.
pub fn free_hazard(hazard: hazard::Writer) {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    debug_assert!(!hazard.is_blocked(), "Freeing a blocked hazards. See docs.");

    if STATE.state() == thread::LocalKeyState::Destroyed {
        // Since the state was deinitialized, we cannot store it for later reuse, so we are forced
        // to simply kill the hazard.
        hazard.kill();
    } else {
        STATE.with(|s| s.borrow_mut().free_hazard(hazard));
    }
}

/// Export the garbage of this thread to the global state.
///
/// This is useful for propagating accumulated garbage such that it can be destroyed by the next
/// garbage collection.
pub fn export_garbage() {
    // Since this function can trigger a GC, it must not be called inside a guard constructor.
    guard::debug_assert_no_create();

    // We can only export when the TLS variable isn't destroyed. Otherwise, there would be nothing
    // to export!
    if STATE.state() != thread::LocalKeyState::Destroyed {
        STATE.with(|s| s.borrow_mut().export_garbage());
        // We tick after the state is no longer reserved, as the tick could potentially call
        // destructor that access the TLS variable.
        global::tick();
    }
}

/// A thread-local state.
#[derive(Default)]
struct State {
    /// The cached garbage waiting to be exported to the global state.
    garbage: Vec<Garbage>,
    /// The cache of currently available hazards.
    ///
    /// We maintain this cache to avoid the performance hit of creating new hazards.
    ///
    /// The hazards in this vector are not necessarily in state "free". Only when a sufficient
    /// amount of available hazards has accumulated, they will be set to free. This means that we
    /// don't have to reset the state of a hazard after usage, giving a quite significant speed-up.
    available_hazards: Vec<hazard::Writer>,
    /// The hazards in the cache before this index are free.
    ///
    /// This number keeps track what hazards in `self.available_hazard` are set to state "free".
    /// Before this index, every hazard must be set to "free".
    ///
    /// It is useful for knowing when to free the hazards to allow garbage collection.
    available_hazards_free_before: usize,
}

impl State {
    /// Get the number of hazards in the cache which are not in state "free".
    fn non_free_hazards(&self) -> usize {
        self.available_hazards.len() - self.available_hazards_free_before
    }

    /// See `get_hazard()`.
    fn get_hazard(&mut self) -> hazard::Writer {
        // Check if there is hazards in the cache.
        if let Some(hazard) = self.available_hazards.pop() {
            // There is; we don't need to create a new hazard.

            // Since the hazard popped from the cache is not blocked, we must block the hazard to
            // satisfy the requirements of this function.
            hazard.block();
            hazard
        } else {
            // There is not; we must create a new hazard.
            global::create_hazard()
        }
    }

    /// See `free_hazard()`.
    fn free_hazard(&mut self, hazard: hazard::Writer) {
        /// The maximal amount of hazards before cleaning up.
        ///
        /// With "cleaning up" we mean setting the state of the hazards to "free" in order to allow
        /// garbage collection of the object it is currently protecting.
        const MAX_NON_FREE_HAZARDS: usize = 16;

        // FIXME: This can lead to some subtle bugs, since the dtor is unpredictable as there is no
        //        way of predicting when the hazard is cleared.

        // Push the given hazard to the cache.
        self.available_hazards.push(hazard);

        // Check if we exceeded the limit.
        if self.non_free_hazards() > MAX_NON_FREE_HAZARDS {
            // We did; we must now set the non-free hazards to "free".
            for i in &self.available_hazards[self.available_hazards_free_before..] {
                i.set(hazard::State::Free);
            }

            // Update the counter such that we mark the new hazards set to "free".
            self.available_hazards_free_before = self.available_hazards.len();
        }
    }

    /// Queues garbage to destroy.
    ///
    /// Eventually the added garbage will be exported to the global state through
    /// `global::add_garbage()`.
    ///
    /// See `add_garbage` for more information.
    ///
    /// When this happens (i.e. the global state gets the garbage), it returns `true`. Otherwise,
    /// it returns `false`.
    fn add_garbage(&mut self, garbage: Garbage) -> bool {
        /// The maximal amount of garbage before exportation to the global state.
        const MAX_GARBAGE: usize = 128;

        // Push the garbage to the cache of garbage.
        self.garbage.push(garbage);

        // Export the garbage if it exceeds the limit.
        // TODO: use memory instead of items as a metric.
        if self.garbage.len() > MAX_GARBAGE {
            self.export_garbage();
            true
        } else { false }
    }

    /// See `export_garbage()` for more information.
    fn export_garbage(&mut self) {
        // Clear the vector and export the garbage.
        global::export_garbage(mem::replace(&mut self.garbage, Vec::new()));
    }
}

impl Drop for State {
    fn drop(&mut self) {
        // Clear every hazard to "dead" state.
        for hazard in self.available_hazards.drain(..) {
            hazard.kill();
        }

        // The thread is exiting, thus we must export the garbage to the global state to avoid
        // memory leaks. It is very important that this does indeed not tick, as causing garbage
        // collection means accessing RNG state, a TLS variable, which cannot be done when, we are
        // here, after it has deinitialized.
        // TODO: Figure out a way we can tick anyway.
        self.export_garbage();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use garbage::Garbage;
    use hazard;
    use std::thread;

    #[test]
    fn dtor_runs() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = Box::new(0);
            let h = get_hazard();
            h.set(hazard::State::Protect(&*b));
            add_garbage(Garbage::new(&*b, dtor));
            ::gc();
            assert_eq!(*b, 0);
            ::gc();
            h.set(hazard::State::Free);
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[test]
    fn dtor_runs_cross_thread() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = Box::new(0);
            let bptr = &*b as *const _ as usize;
            let h = thread::spawn(move || {
                let h = get_hazard();
                h.set(hazard::State::Protect(bptr as *const u8));
                h
            }).join().unwrap();
            add_garbage(Garbage::new(&*b, dtor));
            ::gc();
            assert_eq!(*b, 0);
            ::gc();
            h.set(hazard::State::Free);
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[test]
    fn clear_hazards() {
        let mut s = State::default();
        let mut v = Vec::new();
        for _ in 0..100 {
            let (w, r) = hazard::create();
            w.set(hazard::State::Protect(0x1 as *const u8));
            v.push(r);
            s.free_hazard(w);
        }

        for i in &v[0..16] {
            assert_eq!(i.get(), hazard::State::Free);
        }

        mem::forget(v);
    }

    #[test]
    fn kill_hazards() {
        fn dtor(x: *const u8) {
            unsafe {
                *(x as *mut u8) = 1;
            }
        }

        for _ in 0..1000 {
            let b = thread::spawn(move || {
                let b = Box::new(0);
                let h = get_hazard();
                h.set(hazard::State::Protect(&*b));
                add_garbage(Garbage::new(&*b, dtor));
                ::gc();
                assert_eq!(*b, 0);
                b
            }).join().unwrap();
            ::gc();
            assert_eq!(*b, 1);
        }
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic]
    fn debug_free_blocked() {
        use std::mem;

        let (writer, reader) = hazard::create();
        mem::forget(reader);

        free_hazard(writer);
    }
}
<|endoftext|>


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

#![feature(test)]

extern crate test;
extern crate rand;

const RAND_BENCH_N: u64 = 1000;

use std::mem::size_of;
use test::Bencher;
use std::time::Duration;

use rand::{Rng, FromEntropy};
use rand::rngs::SmallRng;
use rand::distributions::*;

macro_rules! distr_int {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0 as $ty;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum = accum.wrapping_add(x);
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr_float {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0.0;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum += x;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr_duration {
    ($fnn:ident, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = Duration::new(0, 0);
                for _ in 0..::RAND_BENCH_N {
                    let x: Duration = distr.sample(&mut rng);
                    accum = accum.checked_add(x).unwrap_or(Duration::new(u64::max_value(), 999_999_999));
                }
                accum
            });
            b.bytes = size_of::<Duration>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0u32;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum = accum.wrapping_add(x as u32);
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

// uniform
distr_int!(distr_uniform_i8, i8, Uniform::new(20i8, 100));
distr_int!(distr_uniform_i16, i16, Uniform::new(-500i16, 2000));
distr_int!(distr_uniform_i32, i32, Uniform::new(-200_000_000i32, 800_000_000));
distr_int!(distr_uniform_i64, i64, Uniform::new(3i64, 123_456_789_123));
distr_int!(distr_uniform_i128, i128, Uniform::new(-123_456_789_123i128, 123_456_789_123_456_789));

distr_float!(distr_uniform_f32, f32, Uniform::new(2.26f32, 2.319));
distr_float!(distr_uniform_f64, f64, Uniform::new(2.26f64, 2.319));

const LARGE_SEC: u64 = u64::max_value() / 1000;

distr_duration!(distr_uniform_duration_largest,
    Uniform::new_inclusive(Duration::new(0, 0), Duration::new(u64::max_value(), 999_999_999))
);
distr_duration!(distr_uniform_duration_large,
    Uniform::new(Duration::new(0, 0), Duration::new(LARGE_SEC, 1_000_000_000 / 2))
);
distr_duration!(distr_uniform_duration_one,
    Uniform::new(Duration::new(0, 0), Duration::new(1, 0))
);
distr_duration!(distr_uniform_duration_variety,
    Uniform::new(Duration::new(10000, 423423), Duration::new(200000, 6969954))
);
distr_duration!(distr_uniform_duration_edge,
    Uniform::new_inclusive(Duration::new(LARGE_SEC, 999_999_999), Duration::new(LARGE_SEC + 1, 1))
);

// standard
distr_int!(distr_standard_i8, i8, Standard);
distr_int!(distr_standard_i16, i16, Standard);
distr_int!(distr_standard_i32, i32, Standard);
distr_int!(distr_standard_i64, i64, Standard);
distr_int!(distr_standard_i128, i128, Standard);

distr!(distr_standard_bool, bool, Standard);
distr!(distr_standard_alphanumeric, char, Alphanumeric);
distr!(distr_standard_codepoint, char, Standard);

distr_float!(distr_standard_f32, f32, Standard);
distr_float!(distr_standard_f64, f64, Standard);
distr_float!(distr_open01_f32, f32, Open01);
distr_float!(distr_open01_f64, f64, Open01);
distr_float!(distr_openclosed01_f32, f32, OpenClosed01);
distr_float!(distr_openclosed01_f64, f64, OpenClosed01);

// distributions
distr_float!(distr_exp, f64, Exp::new(1.23 * 4.56));
distr_float!(distr_normal, f64, Normal::new(-1.23, 4.56));
distr_float!(distr_log_normal, f64, LogNormal::new(-1.23, 4.56));
distr_float!(distr_gamma_large_shape, f64, Gamma::new(10., 1.0));
distr_float!(distr_gamma_small_shape, f64, Gamma::new(0.1, 1.0));
distr_float!(distr_cauchy, f64, Cauchy::new(4.2, 6.9));
distr_int!(distr_binomial, u64, Binomial::new(20, 0.7));
distr_int!(distr_poisson, u64, Poisson::new(4.0));
distr!(distr_bernoulli, bool, Bernoulli::new(0.18));

// Weighted
distr_int!(distr_weighted_i8, usize, WeightedIndex::new(&[1i8, 2, 3, 4, 12, 0, 2, 1]).unwrap());
distr_int!(distr_weighted_u32, usize, WeightedIndex::new(&[1u32, 2, 3, 4, 12, 0, 2, 1]).unwrap());
distr_int!(distr_weighted_f64, usize, WeightedIndex::new(&[1.0f64, 0.001, 1.0/3.0, 4.01, 0.0, 3.3, 22.0, 0.001]).unwrap());
distr_int!(distr_weighted_large_set, usize, WeightedIndex::new((0..10000).rev().chain(1..10001)).unwrap());

// construct and sample from a range
macro_rules! gen_range_int {
    ($fnn:ident, $ty:ident, $low:expr, $high:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();

            b.iter(|| {
                let mut high = $high;
                let mut accum: $ty = 0;
                for _ in 0..::RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen_range($low, high));
                    // force recalculation of range each time
                    high = high.wrapping_add(1) & std::$ty::MAX;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

gen_range_int!(gen_range_i8, i8, -20i8, 100);
gen_range_int!(gen_range_i16, i16, -500i16, 2000);
gen_range_int!(gen_range_i32, i32, -200_000_000i32, 800_000_000);
gen_range_int!(gen_range_i64, i64, 3i64, 123_456_789_123);
gen_range_int!(gen_range_i128, i128, -12345678901234i128, 123_456_789_123_456_789);

// construct and sample from a floating-point range
macro_rules! gen_range_float {
    ($fnn:ident, $ty:ident, $low:expr, $high:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();

            b.iter(|| {
                let mut high = $high;
                let mut low = $low;
                let mut accum: $ty = 0.0;
                for _ in 0..::RAND_BENCH_N {
                    accum += rng.gen_range(low, high);
                    // force recalculation of range each time
                    low += 0.9;
                    high += 1.1;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

gen_range_float!(gen_range_f32, f32, -20000.0f32, 100000.0);
gen_range_float!(gen_range_f64, f64, 123.456f64, 7890.12);

#[bench]
fn dist_iter(b: &mut Bencher) {
    let mut rng = SmallRng::from_entropy();
    let distr = Normal::new(-2.71828, 3.14159);
    let mut iter = distr.sample_iter(&mut rng);

    b.iter(|| {
        let mut accum = 0.0;
        for _ in 0..::RAND_BENCH_N {
            accum += iter.next().unwrap();
        }
        accum
    });
    b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
}
<commit_msg>Add benchmarks of unit circle and sphere surface sampling<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(test)]

extern crate test;
extern crate rand;

const RAND_BENCH_N: u64 = 1000;

use std::mem::size_of;
use test::Bencher;
use std::time::Duration;

use rand::{Rng, FromEntropy};
use rand::rngs::SmallRng;
use rand::distributions::*;

macro_rules! distr_int {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0 as $ty;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum = accum.wrapping_add(x);
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr_float {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0.0;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum += x;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr_duration {
    ($fnn:ident, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = Duration::new(0, 0);
                for _ in 0..::RAND_BENCH_N {
                    let x: Duration = distr.sample(&mut rng);
                    accum = accum.checked_add(x).unwrap_or(Duration::new(u64::max_value(), 999_999_999));
                }
                accum
            });
            b.bytes = size_of::<Duration>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0u32;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum = accum.wrapping_add(x as u32);
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

macro_rules! distr_arr {
    ($fnn:ident, $ty:ty, $distr:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();
            let distr = $distr;

            b.iter(|| {
                let mut accum = 0u32;
                for _ in 0..::RAND_BENCH_N {
                    let x: $ty = distr.sample(&mut rng);
                    accum = accum.wrapping_add(x[0] as u32);
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

// uniform
distr_int!(distr_uniform_i8, i8, Uniform::new(20i8, 100));
distr_int!(distr_uniform_i16, i16, Uniform::new(-500i16, 2000));
distr_int!(distr_uniform_i32, i32, Uniform::new(-200_000_000i32, 800_000_000));
distr_int!(distr_uniform_i64, i64, Uniform::new(3i64, 123_456_789_123));
distr_int!(distr_uniform_i128, i128, Uniform::new(-123_456_789_123i128, 123_456_789_123_456_789));

distr_float!(distr_uniform_f32, f32, Uniform::new(2.26f32, 2.319));
distr_float!(distr_uniform_f64, f64, Uniform::new(2.26f64, 2.319));

const LARGE_SEC: u64 = u64::max_value() / 1000;

distr_duration!(distr_uniform_duration_largest,
    Uniform::new_inclusive(Duration::new(0, 0), Duration::new(u64::max_value(), 999_999_999))
);
distr_duration!(distr_uniform_duration_large,
    Uniform::new(Duration::new(0, 0), Duration::new(LARGE_SEC, 1_000_000_000 / 2))
);
distr_duration!(distr_uniform_duration_one,
    Uniform::new(Duration::new(0, 0), Duration::new(1, 0))
);
distr_duration!(distr_uniform_duration_variety,
    Uniform::new(Duration::new(10000, 423423), Duration::new(200000, 6969954))
);
distr_duration!(distr_uniform_duration_edge,
    Uniform::new_inclusive(Duration::new(LARGE_SEC, 999_999_999), Duration::new(LARGE_SEC + 1, 1))
);

// standard
distr_int!(distr_standard_i8, i8, Standard);
distr_int!(distr_standard_i16, i16, Standard);
distr_int!(distr_standard_i32, i32, Standard);
distr_int!(distr_standard_i64, i64, Standard);
distr_int!(distr_standard_i128, i128, Standard);

distr!(distr_standard_bool, bool, Standard);
distr!(distr_standard_alphanumeric, char, Alphanumeric);
distr!(distr_standard_codepoint, char, Standard);

distr_float!(distr_standard_f32, f32, Standard);
distr_float!(distr_standard_f64, f64, Standard);
distr_float!(distr_open01_f32, f32, Open01);
distr_float!(distr_open01_f64, f64, Open01);
distr_float!(distr_openclosed01_f32, f32, OpenClosed01);
distr_float!(distr_openclosed01_f64, f64, OpenClosed01);

// distributions
distr_float!(distr_exp, f64, Exp::new(1.23 * 4.56));
distr_float!(distr_normal, f64, Normal::new(-1.23, 4.56));
distr_float!(distr_log_normal, f64, LogNormal::new(-1.23, 4.56));
distr_float!(distr_gamma_large_shape, f64, Gamma::new(10., 1.0));
distr_float!(distr_gamma_small_shape, f64, Gamma::new(0.1, 1.0));
distr_float!(distr_cauchy, f64, Cauchy::new(4.2, 6.9));
distr_int!(distr_binomial, u64, Binomial::new(20, 0.7));
distr_int!(distr_poisson, u64, Poisson::new(4.0));
distr!(distr_bernoulli, bool, Bernoulli::new(0.18));
distr_arr!(distr_circle, [f64; 2], UnitCircle::new());
distr_arr!(distr_sphere_surface, [f64; 3], UnitSphereSurface::new());

// Weighted
distr_int!(distr_weighted_i8, usize, WeightedIndex::new(&[1i8, 2, 3, 4, 12, 0, 2, 1]).unwrap());
distr_int!(distr_weighted_u32, usize, WeightedIndex::new(&[1u32, 2, 3, 4, 12, 0, 2, 1]).unwrap());
distr_int!(distr_weighted_f64, usize, WeightedIndex::new(&[1.0f64, 0.001, 1.0/3.0, 4.01, 0.0, 3.3, 22.0, 0.001]).unwrap());
distr_int!(distr_weighted_large_set, usize, WeightedIndex::new((0..10000).rev().chain(1..10001)).unwrap());

// construct and sample from a range
macro_rules! gen_range_int {
    ($fnn:ident, $ty:ident, $low:expr, $high:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();

            b.iter(|| {
                let mut high = $high;
                let mut accum: $ty = 0;
                for _ in 0..::RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen_range($low, high));
                    // force recalculation of range each time
                    high = high.wrapping_add(1) & std::$ty::MAX;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

gen_range_int!(gen_range_i8, i8, -20i8, 100);
gen_range_int!(gen_range_i16, i16, -500i16, 2000);
gen_range_int!(gen_range_i32, i32, -200_000_000i32, 800_000_000);
gen_range_int!(gen_range_i64, i64, 3i64, 123_456_789_123);
gen_range_int!(gen_range_i128, i128, -12345678901234i128, 123_456_789_123_456_789);

// construct and sample from a floating-point range
macro_rules! gen_range_float {
    ($fnn:ident, $ty:ident, $low:expr, $high:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = SmallRng::from_entropy();

            b.iter(|| {
                let mut high = $high;
                let mut low = $low;
                let mut accum: $ty = 0.0;
                for _ in 0..::RAND_BENCH_N {
                    accum += rng.gen_range(low, high);
                    // force recalculation of range each time
                    low += 0.9;
                    high += 1.1;
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * ::RAND_BENCH_N;
        }
    }
}

gen_range_float!(gen_range_f32, f32, -20000.0f32, 100000.0);
gen_range_float!(gen_range_f64, f64, 123.456f64, 7890.12);

#[bench]
fn dist_iter(b: &mut Bencher) {
    let mut rng = SmallRng::from_entropy();
    let distr = Normal::new(-2.71828, 3.14159);
    let mut iter = distr.sample_iter(&mut rng);

    b.iter(|| {
        let mut accum = 0.0;
        for _ in 0..::RAND_BENCH_N {
            accum += iter.next().unwrap();
        }
        accum
    });
    b.bytes = size_of::<f64>() as u64 * ::RAND_BENCH_N;
}
<|endoftext|>


<commit_before>mod common;

use self::common::*;
use std::io;

#[test]
fn test_downcast() {
    assert_eq!(
        "oh no!",
        bail_literal().unwrap_err().downcast::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast::<io::Error>()
            .unwrap()
            .to_string(),
    );
}

#[test]
fn test_downcast_ref() {
    assert_eq!(
        "oh no!",
        *bail_literal().unwrap_err().downcast_ref::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast_ref::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast_ref::<io::Error>()
            .unwrap()
            .to_string(),
    );
}

#[test]
fn test_downcast_mut() {
    assert_eq!(
        "oh no!",
        *bail_literal().unwrap_err().downcast_mut::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast_mut::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast_mut::<io::Error>()
            .unwrap()
            .to_string(),
    );
}
<commit_msg>Add test of drop after downcast<commit_after>mod common;
mod drop;

use self::common::*;
use self::drop::DetectDrop;
use anyhow::Error;
use std::io;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::Arc;

#[test]
fn test_downcast() {
    assert_eq!(
        "oh no!",
        bail_literal().unwrap_err().downcast::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast::<io::Error>()
            .unwrap()
            .to_string(),
    );
}

#[test]
fn test_downcast_ref() {
    assert_eq!(
        "oh no!",
        *bail_literal().unwrap_err().downcast_ref::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast_ref::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast_ref::<io::Error>()
            .unwrap()
            .to_string(),
    );
}

#[test]
fn test_downcast_mut() {
    assert_eq!(
        "oh no!",
        *bail_literal().unwrap_err().downcast_mut::<&str>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_fmt().unwrap_err().downcast_mut::<String>().unwrap(),
    );
    assert_eq!(
        "oh no!",
        bail_error()
            .unwrap_err()
            .downcast_mut::<io::Error>()
            .unwrap()
            .to_string(),
    );
}

#[test]
fn test_drop() {
    let has_dropped = Arc::new(AtomicBool::new(false));
    let error = Error::new(DetectDrop::new(&has_dropped));
    drop(error.downcast::<DetectDrop>().unwrap());
    assert!(has_dropped.load(SeqCst));
}
<|endoftext|>


<commit_before>use crate::utils::{iter_input_pats, snippet, span_lint, type_is_unsafe_function};
use matches::matches;
use rustc::hir;
use rustc::hir::def::Def;
use rustc::hir::intravisit;
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::ty;
use rustc::{declare_tool_lint, lint_array};
use rustc_data_structures::fx::FxHashSet;
use rustc_target::spec::abi::Abi;
use syntax::ast;
use syntax::source_map::Span;

/// **What it does:** Checks for functions with too many parameters.
///
/// **Why is this bad?** Functions with lots of parameters are considered bad
/// style and reduce readability (“what does the 5th parameter mean?”). Consider
/// grouping some parameters into a new type.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// fn foo(x: u32, y: u32, name: &str, c: Color, w: f32, h: f32, a: f32, b: f32) {
///     ..
/// }
/// ```
declare_clippy_lint! {
    pub TOO_MANY_ARGUMENTS,
    complexity,
    "functions with too many arguments"
}

/// **What it does:** Checks for functions with a large amount of lines.
///
/// **Why is this bad?** Functions with a lot of lines are harder to understand
/// due to having to look at a larger amount of code to understand what the
/// function is doing. Consider splitting the body of the function into
/// multiple functions.
///
/// **Known problems:** None.
///
/// **Example:**
/// ``` rust
/// fn im_too_long() {
/// println!("");
/// // ... 100 more LoC
/// println!("");
/// }
/// ```
declare_clippy_lint! {
    pub TOO_MANY_LINES,
    pedantic,
    "functions with too many lines"
}

/// **What it does:** Checks for public functions that dereferences raw pointer
/// arguments but are not marked unsafe.
///
/// **Why is this bad?** The function should probably be marked `unsafe`, since
/// for an arbitrary raw pointer, there is no way of telling for sure if it is
/// valid.
///
/// **Known problems:**
///
/// * It does not check functions recursively so if the pointer is passed to a
/// private non-`unsafe` function which does the dereferencing, the lint won't
/// trigger.
/// * It only checks for arguments whose type are raw pointers, not raw pointers
/// got from an argument in some other way (`fn foo(bar: &[*const u8])` or
/// `some_argument.get_raw_ptr()`).
///
/// **Example:**
/// ```rust
/// pub fn foo(x: *const u8) {
///     println!("{}", unsafe { *x });
/// }
/// ```
declare_clippy_lint! {
    pub NOT_UNSAFE_PTR_ARG_DEREF,
    correctness,
    "public functions dereferencing raw pointer arguments but not marked `unsafe`"
}

#[derive(Copy, Clone)]
pub struct Functions {
    threshold: u64,
    max_lines: u64
}

impl Functions {
    pub fn new(threshold: u64, max_lines: u64) -> Self {
        Self {
            threshold,
            max_lines
        }
    }
}

impl LintPass for Functions {
    fn get_lints(&self) -> LintArray {
        lint_array!(TOO_MANY_ARGUMENTS, TOO_MANY_LINES, NOT_UNSAFE_PTR_ARG_DEREF)
    }

    fn name(&self) -> &'static str {
        "Functions"
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for Functions {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        kind: intravisit::FnKind<'tcx>,
        decl: &'tcx hir::FnDecl,
        body: &'tcx hir::Body,
        span: Span,
        nodeid: ast::NodeId,
    ) {
        let is_impl = if let Some(hir::Node::Item(item)) = cx.tcx.hir().find(cx.tcx.hir().get_parent_node(nodeid)) {
            matches!(item.node, hir::ItemKind::Impl(_, _, _, _, Some(_), _, _))
        } else {
            false
        };

        let unsafety = match kind {
            hir::intravisit::FnKind::ItemFn(_, _, hir::FnHeader { unsafety, .. }, _, _) => unsafety,
            hir::intravisit::FnKind::Method(_, sig, _, _) => sig.header.unsafety,
            hir::intravisit::FnKind::Closure(_) => return,
        };

        // don't warn for implementations, it's not their fault
        if !is_impl {
            // don't lint extern functions decls, it's not their fault either
            match kind {
                hir::intravisit::FnKind::Method(
                    _,
                    &hir::MethodSig {
                        header: hir::FnHeader { abi: Abi::Rust, .. },
                        ..
                    },
                    _,
                    _,
                )
                | hir::intravisit::FnKind::ItemFn(_, _, hir::FnHeader { abi: Abi::Rust, .. }, _, _) => {
                    self.check_arg_number(cx, decl, span)
                },
                _ => {},
            }
        }

        self.check_raw_ptr(cx, unsafety, decl, body, nodeid);
        self.check_line_number(cx, span);
    }

    fn check_trait_item(&mut self, cx: &LateContext<'a, 'tcx>, item: &'tcx hir::TraitItem) {
        if let hir::TraitItemKind::Method(ref sig, ref eid) = item.node {
            // don't lint extern functions decls, it's not their fault
            if sig.header.abi == Abi::Rust {
                self.check_arg_number(cx, &sig.decl, item.span);
            }

            if let hir::TraitMethod::Provided(eid) = *eid {
                let body = cx.tcx.hir().body(eid);
                self.check_raw_ptr(cx, sig.header.unsafety, &sig.decl, body, item.id);
            }
        }
    }
}

impl<'a, 'tcx> Functions {
    fn check_arg_number(self, cx: &LateContext<'_, '_>, decl: &hir::FnDecl, span: Span) {
        let args = decl.inputs.len() as u64;
        if args > self.threshold {
            span_lint(
                cx,
                TOO_MANY_ARGUMENTS,
                span,
                &format!("this function has too many arguments ({}/{})", args, self.threshold),
            );
        }
    }

    fn check_line_number(self, cx: &LateContext<'_, '_>, span: Span) {
        let code_snippet = snippet(cx, span, "..");
        let mut line_count: u64 = 0;
        let mut in_comment = false;
        let mut code_in_line;

        // Skip the surrounding function decl.
        let start_brace_idx = match code_snippet.find("{") {
            Some(i) => i + 1,
            None => 0
        };
        let end_brace_idx = match code_snippet.find("}") {
            Some(i) => i,
            None => code_snippet.len()
        };
        let function_lines = code_snippet[start_brace_idx..end_brace_idx].lines();

        for mut line in function_lines {
            code_in_line = false;
            loop {
                line = line.trim_start();
                if line.is_empty() { break; }
                if in_comment {
                    match line.find("*/") {
                        Some(i) => {
                            line = &line[i + 2..];
                            in_comment = false;
                            continue;
                        },
                        None => break
                    }
                } else {
                    let multi_idx = match line.find("/*") {
                        Some(i) => i,
                        None => line.len()
                    };
                    let single_idx = match line.find("//") {
                        Some(i) => i,
                        None => line.len()
                    };
                    code_in_line |= multi_idx > 0 && single_idx > 0;
                    // Implies multi_idx is below line.len()
                    if multi_idx < single_idx {
                        line = &line[multi_idx + 2..];
                        in_comment = true;
                        continue;
                    }
                    break;
                }
            }
            if code_in_line { line_count += 1; }
        }

        if line_count > self.max_lines {
            span_lint(cx, TOO_MANY_LINES, span,
                      "This function has a large number of lines.")
        }
    }

    fn check_raw_ptr(
        self,
        cx: &LateContext<'a, 'tcx>,
        unsafety: hir::Unsafety,
        decl: &'tcx hir::FnDecl,
        body: &'tcx hir::Body,
        nodeid: ast::NodeId,
    ) {
        let expr = &body.value;
        if unsafety == hir::Unsafety::Normal && cx.access_levels.is_exported(nodeid) {
            let raw_ptrs = iter_input_pats(decl, body)
                .zip(decl.inputs.iter())
                .filter_map(|(arg, ty)| raw_ptr_arg(arg, ty))
                .collect::<FxHashSet<_>>();

            if !raw_ptrs.is_empty() {
                let tables = cx.tcx.body_tables(body.id());
                let mut v = DerefVisitor {
                    cx,
                    ptrs: raw_ptrs,
                    tables,
                };

                hir::intravisit::walk_expr(&mut v, expr);
            }
        }
    }
}

fn raw_ptr_arg(arg: &hir::Arg, ty: &hir::Ty) -> Option<ast::NodeId> {
    if let (&hir::PatKind::Binding(_, id, _, _), &hir::TyKind::Ptr(_)) = (&arg.pat.node, &ty.node) {
        Some(id)
    } else {
        None
    }
}

struct DerefVisitor<'a, 'tcx: 'a> {
    cx: &'a LateContext<'a, 'tcx>,
    ptrs: FxHashSet<ast::NodeId>,
    tables: &'a ty::TypeckTables<'tcx>,
}

impl<'a, 'tcx> hir::intravisit::Visitor<'tcx> for DerefVisitor<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'tcx hir::Expr) {
        match expr.node {
            hir::ExprKind::Call(ref f, ref args) => {
                let ty = self.tables.expr_ty(f);

                if type_is_unsafe_function(self.cx, ty) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            },
            hir::ExprKind::MethodCall(_, _, ref args) => {
                let def_id = self.tables.type_dependent_defs()[expr.hir_id].def_id();
                let base_type = self.cx.tcx.type_of(def_id);

                if type_is_unsafe_function(self.cx, base_type) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            },
            hir::ExprKind::Unary(hir::UnDeref, ref ptr) => self.check_arg(ptr),
            _ => (),
        }

        hir::intravisit::walk_expr(self, expr);
    }
    fn nested_visit_map<'this>(&'this mut self) -> intravisit::NestedVisitorMap<'this, 'tcx> {
        intravisit::NestedVisitorMap::None
    }
}

impl<'a, 'tcx: 'a> DerefVisitor<'a, 'tcx> {
    fn check_arg(&self, ptr: &hir::Expr) {
        if let hir::ExprKind::Path(ref qpath) = ptr.node {
            if let Def::Local(id) = self.cx.tables.qpath_def(qpath, ptr.hir_id) {
                if self.ptrs.contains(&id) {
                    span_lint(
                        self.cx,
                        NOT_UNSAFE_PTR_ARG_DEREF,
                        ptr.span,
                        "this public function dereferences a raw pointer but is not marked `unsafe`",
                    );
                }
            }
        }
    }
}
<commit_msg>rustfmt<commit_after>use crate::utils::{iter_input_pats, snippet, span_lint, type_is_unsafe_function};
use matches::matches;
use rustc::hir;
use rustc::hir::def::Def;
use rustc::hir::intravisit;
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::ty;
use rustc::{declare_tool_lint, lint_array};
use rustc_data_structures::fx::FxHashSet;
use rustc_target::spec::abi::Abi;
use syntax::ast;
use syntax::source_map::Span;

/// **What it does:** Checks for functions with too many parameters.
///
/// **Why is this bad?** Functions with lots of parameters are considered bad
/// style and reduce readability (“what does the 5th parameter mean?”). Consider
/// grouping some parameters into a new type.
///
/// **Known problems:** None.
///
/// **Example:**
/// ```rust
/// fn foo(x: u32, y: u32, name: &str, c: Color, w: f32, h: f32, a: f32, b: f32) {
///     ..
/// }
/// ```
declare_clippy_lint! {
    pub TOO_MANY_ARGUMENTS,
    complexity,
    "functions with too many arguments"
}

/// **What it does:** Checks for functions with a large amount of lines.
///
/// **Why is this bad?** Functions with a lot of lines are harder to understand
/// due to having to look at a larger amount of code to understand what the
/// function is doing. Consider splitting the body of the function into
/// multiple functions.
///
/// **Known problems:** None.
///
/// **Example:**
/// ``` rust
/// fn im_too_long() {
/// println!("");
/// // ... 100 more LoC
/// println!("");
/// }
/// ```
declare_clippy_lint! {
    pub TOO_MANY_LINES,
    pedantic,
    "functions with too many lines"
}

/// **What it does:** Checks for public functions that dereferences raw pointer
/// arguments but are not marked unsafe.
///
/// **Why is this bad?** The function should probably be marked `unsafe`, since
/// for an arbitrary raw pointer, there is no way of telling for sure if it is
/// valid.
///
/// **Known problems:**
///
/// * It does not check functions recursively so if the pointer is passed to a
/// private non-`unsafe` function which does the dereferencing, the lint won't
/// trigger.
/// * It only checks for arguments whose type are raw pointers, not raw pointers
/// got from an argument in some other way (`fn foo(bar: &[*const u8])` or
/// `some_argument.get_raw_ptr()`).
///
/// **Example:**
/// ```rust
/// pub fn foo(x: *const u8) {
///     println!("{}", unsafe { *x });
/// }
/// ```
declare_clippy_lint! {
    pub NOT_UNSAFE_PTR_ARG_DEREF,
    correctness,
    "public functions dereferencing raw pointer arguments but not marked `unsafe`"
}

#[derive(Copy, Clone)]
pub struct Functions {
    threshold: u64,
    max_lines: u64,
}

impl Functions {
    pub fn new(threshold: u64, max_lines: u64) -> Self {
        Self { threshold, max_lines }
    }
}

impl LintPass for Functions {
    fn get_lints(&self) -> LintArray {
        lint_array!(TOO_MANY_ARGUMENTS, TOO_MANY_LINES, NOT_UNSAFE_PTR_ARG_DEREF)
    }

    fn name(&self) -> &'static str {
        "Functions"
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for Functions {
    fn check_fn(
        &mut self,
        cx: &LateContext<'a, 'tcx>,
        kind: intravisit::FnKind<'tcx>,
        decl: &'tcx hir::FnDecl,
        body: &'tcx hir::Body,
        span: Span,
        nodeid: ast::NodeId,
    ) {
        let is_impl = if let Some(hir::Node::Item(item)) = cx.tcx.hir().find(cx.tcx.hir().get_parent_node(nodeid)) {
            matches!(item.node, hir::ItemKind::Impl(_, _, _, _, Some(_), _, _))
        } else {
            false
        };

        let unsafety = match kind {
            hir::intravisit::FnKind::ItemFn(_, _, hir::FnHeader { unsafety, .. }, _, _) => unsafety,
            hir::intravisit::FnKind::Method(_, sig, _, _) => sig.header.unsafety,
            hir::intravisit::FnKind::Closure(_) => return,
        };

        // don't warn for implementations, it's not their fault
        if !is_impl {
            // don't lint extern functions decls, it's not their fault either
            match kind {
                hir::intravisit::FnKind::Method(
                    _,
                    &hir::MethodSig {
                        header: hir::FnHeader { abi: Abi::Rust, .. },
                        ..
                    },
                    _,
                    _,
                )
                | hir::intravisit::FnKind::ItemFn(_, _, hir::FnHeader { abi: Abi::Rust, .. }, _, _) => {
                    self.check_arg_number(cx, decl, span)
                },
                _ => {},
            }
        }

        self.check_raw_ptr(cx, unsafety, decl, body, nodeid);
        self.check_line_number(cx, span);
    }

    fn check_trait_item(&mut self, cx: &LateContext<'a, 'tcx>, item: &'tcx hir::TraitItem) {
        if let hir::TraitItemKind::Method(ref sig, ref eid) = item.node {
            // don't lint extern functions decls, it's not their fault
            if sig.header.abi == Abi::Rust {
                self.check_arg_number(cx, &sig.decl, item.span);
            }

            if let hir::TraitMethod::Provided(eid) = *eid {
                let body = cx.tcx.hir().body(eid);
                self.check_raw_ptr(cx, sig.header.unsafety, &sig.decl, body, item.id);
            }
        }
    }
}

impl<'a, 'tcx> Functions {
    fn check_arg_number(self, cx: &LateContext<'_, '_>, decl: &hir::FnDecl, span: Span) {
        let args = decl.inputs.len() as u64;
        if args > self.threshold {
            span_lint(
                cx,
                TOO_MANY_ARGUMENTS,
                span,
                &format!("this function has too many arguments ({}/{})", args, self.threshold),
            );
        }
    }

    fn check_line_number(self, cx: &LateContext<'_, '_>, span: Span) {
        let code_snippet = snippet(cx, span, "..");
        let mut line_count: u64 = 0;
        let mut in_comment = false;
        let mut code_in_line;

        // Skip the surrounding function decl.
        let start_brace_idx = match code_snippet.find("{") {
            Some(i) => i + 1,
            None => 0,
        };
        let end_brace_idx = match code_snippet.find("}") {
            Some(i) => i,
            None => code_snippet.len(),
        };
        let function_lines = code_snippet[start_brace_idx..end_brace_idx].lines();

        for mut line in function_lines {
            code_in_line = false;
            loop {
                line = line.trim_start();
                if line.is_empty() {
                    break;
                }
                if in_comment {
                    match line.find("*/") {
                        Some(i) => {
                            line = &line[i + 2..];
                            in_comment = false;
                            continue;
                        },
                        None => break,
                    }
                } else {
                    let multi_idx = match line.find("/*") {
                        Some(i) => i,
                        None => line.len(),
                    };
                    let single_idx = match line.find("//") {
                        Some(i) => i,
                        None => line.len(),
                    };
                    code_in_line |= multi_idx > 0 && single_idx > 0;
                    // Implies multi_idx is below line.len()
                    if multi_idx < single_idx {
                        line = &line[multi_idx + 2..];
                        in_comment = true;
                        continue;
                    }
                    break;
                }
            }
            if code_in_line {
                line_count += 1;
            }
        }

        if line_count > self.max_lines {
            span_lint(cx, TOO_MANY_LINES, span, "This function has a large number of lines.")
        }
    }

    fn check_raw_ptr(
        self,
        cx: &LateContext<'a, 'tcx>,
        unsafety: hir::Unsafety,
        decl: &'tcx hir::FnDecl,
        body: &'tcx hir::Body,
        nodeid: ast::NodeId,
    ) {
        let expr = &body.value;
        if unsafety == hir::Unsafety::Normal && cx.access_levels.is_exported(nodeid) {
            let raw_ptrs = iter_input_pats(decl, body)
                .zip(decl.inputs.iter())
                .filter_map(|(arg, ty)| raw_ptr_arg(arg, ty))
                .collect::<FxHashSet<_>>();

            if !raw_ptrs.is_empty() {
                let tables = cx.tcx.body_tables(body.id());
                let mut v = DerefVisitor {
                    cx,
                    ptrs: raw_ptrs,
                    tables,
                };

                hir::intravisit::walk_expr(&mut v, expr);
            }
        }
    }
}

fn raw_ptr_arg(arg: &hir::Arg, ty: &hir::Ty) -> Option<ast::NodeId> {
    if let (&hir::PatKind::Binding(_, id, _, _), &hir::TyKind::Ptr(_)) = (&arg.pat.node, &ty.node) {
        Some(id)
    } else {
        None
    }
}

struct DerefVisitor<'a, 'tcx: 'a> {
    cx: &'a LateContext<'a, 'tcx>,
    ptrs: FxHashSet<ast::NodeId>,
    tables: &'a ty::TypeckTables<'tcx>,
}

impl<'a, 'tcx> hir::intravisit::Visitor<'tcx> for DerefVisitor<'a, 'tcx> {
    fn visit_expr(&mut self, expr: &'tcx hir::Expr) {
        match expr.node {
            hir::ExprKind::Call(ref f, ref args) => {
                let ty = self.tables.expr_ty(f);

                if type_is_unsafe_function(self.cx, ty) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            },
            hir::ExprKind::MethodCall(_, _, ref args) => {
                let def_id = self.tables.type_dependent_defs()[expr.hir_id].def_id();
                let base_type = self.cx.tcx.type_of(def_id);

                if type_is_unsafe_function(self.cx, base_type) {
                    for arg in args {
                        self.check_arg(arg);
                    }
                }
            },
            hir::ExprKind::Unary(hir::UnDeref, ref ptr) => self.check_arg(ptr),
            _ => (),
        }

        hir::intravisit::walk_expr(self, expr);
    }
    fn nested_visit_map<'this>(&'this mut self) -> intravisit::NestedVisitorMap<'this, 'tcx> {
        intravisit::NestedVisitorMap::None
    }
}

impl<'a, 'tcx: 'a> DerefVisitor<'a, 'tcx> {
    fn check_arg(&self, ptr: &hir::Expr) {
        if let hir::ExprKind::Path(ref qpath) = ptr.node {
            if let Def::Local(id) = self.cx.tables.qpath_def(qpath, ptr.hir_id) {
                if self.ptrs.contains(&id) {
                    span_lint(
                        self.cx,
                        NOT_UNSAFE_PTR_ARG_DEREF,
                        ptr.span,
                        "this public function dereferences a raw pointer but is not marked `unsafe`",
                    );
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>pub mod bldr;
pub mod cli;
pub mod config;
pub mod file;
pub mod license;
pub mod origin;
pub mod pkg;
pub mod plan;
pub mod ring;
pub mod studio;
pub mod sup;
pub mod svc;
#[cfg(test)]
mod tests;
pub mod user;
pub mod util;

use self::{bldr::{Bldr,
                  ConfigOptBldr},
           cli::{Cli,
                 ConfigOptCli},
           config::{ConfigOptServiceConfig,
                    ConfigOptServiceConfigApply,
                    ServiceConfig,
                    ServiceConfigApply},
           file::{ConfigOptFile,
                  File},
           license::{ConfigOptLicense,
                     License},
           origin::{ConfigOptOrigin,
                    Origin},
           pkg::{ConfigOptPkg,
                 ConfigOptPkgInstall,
                 Pkg,
                 PkgInstall},
           plan::{ConfigOptPlan,
                  Plan},
           ring::{ConfigOptRing,
                  Ring},
           studio::{ConfigOptStudio,
                    Studio},
           sup::{ConfigOptHabSup,
                 ConfigOptSupRun,
                 HabSup,
                 SupRun},
           svc::{ConfigOptSvc,
                 ConfigOptSvcStart,
                 ConfigOptSvcStop,
                 Svc,
                 SvcStart,
                 SvcStop},
           user::{ConfigOptUserKeyGenerate,
                  UserKeyGenerate},
           util::{CacheKeyPath,
                  ConfigOptCacheKeyPath}};
use crate::{cli::AFTER_HELP,
            VERSION};
use configopt::ConfigOpt;
use structopt::{clap::AppSettings,
                StructOpt};

#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "hab",
            version = VERSION,
            about = "Patents: https://chef.io/patents\n\"A Habitat is the natural environment for your services\" - Alan Turing",
            author = "\nThe Habitat Maintainers <humans@habitat.sh>\n",
            settings = &[AppSettings::GlobalVersion],
            after_help = AFTER_HELP
        )]
#[allow(clippy::large_enum_variant)]
pub enum Hab {
    #[structopt(no_version)]
    Bldr(Bldr),
    #[structopt(no_version)]
    Cli(Cli),
    #[structopt(no_version)]
    Config(ServiceConfig),
    #[structopt(no_version)]
    File(File),
    #[structopt(no_version, settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
    License(License),
    #[structopt(no_version)]
    Origin(Origin),
    #[structopt(no_version)]
    Pkg(Pkg),
    #[structopt(no_version)]
    Plan(Plan),
    #[structopt(no_version)]
    Ring(Ring),
    #[structopt(no_version, aliases = &["stu", "stud", "studi"])]
    Studio(Studio),
    #[structopt(no_version)]
    Sup(HabSup),
    /// Create a tarball of Habitat Supervisor data to send to support
    #[structopt(no_version)]
    Supportbundle,
    #[structopt(no_version)]
    Svc(Svc),
    #[structopt(no_version)]
    User(User),

    /// Alias for 'config apply'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Apply(ServiceConfigApply),
    /// Alias for 'pkg install'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["i", "in", "ins", "inst", "insta", "instal"])]
    Install(PkgInstall),
    #[cfg(not(target_os = "macos"))]
    /// Alias for 'sup run'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Run(SupRun),
    /// Alias for 'cli setup'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["set", "setu"])]
    Setup(CacheKeyPath),
    /// Alias for 'svc start'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["sta", "star"])]
    Start(SvcStart),
    /// Alias for 'svc stop'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["sto"])]
    Stop(SvcStop),
    #[cfg(not(target_os = "macos"))]
    /// Alias for 'sup term'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Term,
}

/// Commands relating to Habitat users
#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "user", no_version, aliases = &["u", "us", "use"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
pub enum User {
    Key(UserKey),
}

#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "key", no_version, aliases = &["k", "ke"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
/// Commands relating to Habitat user keys
pub enum UserKey {
    #[structopt(no_version, aliases = &["g", "ge", "gen", "gene", "gener", "genera", "generat"])]
    Generate(UserKeyGenerate),
}<commit_msg>[Refactor] added bldr wrapper to hab<commit_after>pub mod bldr;
pub mod cli;
pub mod config;
pub mod file;
pub mod license;
pub mod origin;
pub mod pkg;
pub mod plan;
pub mod ring;
pub mod studio;
pub mod sup;
pub mod svc;
#[cfg(test)]
mod tests;
pub mod user;
pub mod util;

use self::{bldr::{ConfigOptChannelCreate,
                  ConfigOptChannelDemote,
                  ConfigOptChannelDestroy,
                  ConfigOptChannelList,
                  ConfigOptChannelPromote,
                  ConfigOptJobCancel,
                  ConfigOptJobDemote,
                  ConfigOptJobPromote,
                  ConfigOptJobStart,
                  ConfigOptJobStatus,
                  ChannelCreate,
                  ChannelDemote,
                  ChannelDestroy,
                  ChannelList,
                  ChannelPromote,
                  JobCancel,
                  JobDemote,
                  JobPromote,
                  JobStart,
                  JobStatus},
           cli::{Cli,
                 ConfigOptCli},
           config::{ConfigOptServiceConfig,
                    ConfigOptServiceConfigApply,
                    ServiceConfig,
                    ServiceConfigApply},
           file::{ConfigOptFile,
                  File},
           license::{ConfigOptLicense,
                     License},
           origin::{ConfigOptOrigin,
                    Origin},
           pkg::{ConfigOptPkg,
                 ConfigOptPkgInstall,
                 Pkg,
                 PkgInstall},
           plan::{ConfigOptPlan,
                  Plan},
           ring::{ConfigOptRing,
                  Ring},
           studio::{ConfigOptStudio,
                    Studio},
           sup::{ConfigOptHabSup,
                 ConfigOptSupRun,
                 HabSup,
                 SupRun},
           svc::{ConfigOptSvc,
                 ConfigOptSvcStart,
                 ConfigOptSvcStop,
                 Svc,
                 SvcStart,
                 SvcStop},
           user::{ConfigOptUserKeyGenerate,
                  UserKeyGenerate},
           util::{CacheKeyPath,
                  ConfigOptCacheKeyPath}};
use crate::{cli::AFTER_HELP,
            VERSION};
use configopt::ConfigOpt;
use structopt::{clap::AppSettings,
                StructOpt};

#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "hab",
            version = VERSION,
            about = "Patents: https://chef.io/patents\n\"A Habitat is the natural environment for your services\" - Alan Turing",
            author = "\nThe Habitat Maintainers <humans@habitat.sh>\n",
            settings = &[AppSettings::GlobalVersion],
            after_help = AFTER_HELP
        )]
#[allow(clippy::large_enum_variant)]
pub enum Hab {
    #[structopt(no_version)]
    Bldr(Bldr),
    #[structopt(no_version)]
    Cli(Cli),
    #[structopt(no_version)]
    Config(ServiceConfig),
    #[structopt(no_version)]
    File(File),
    #[structopt(no_version, settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
    License(License),
    #[structopt(no_version)]
    Origin(Origin),
    #[structopt(no_version)]
    Pkg(Pkg),
    #[structopt(no_version)]
    Plan(Plan),
    #[structopt(no_version)]
    Ring(Ring),
    #[structopt(no_version, aliases = &["stu", "stud", "studi"])]
    Studio(Studio),
    #[structopt(no_version)]
    Sup(HabSup),
    /// Create a tarball of Habitat Supervisor data to send to support
    #[structopt(no_version)]
    Supportbundle,
    #[structopt(no_version)]
    Svc(Svc),
    #[structopt(no_version)]
    User(User),

    /// Alias for 'config apply'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Apply(ServiceConfigApply),
    /// Alias for 'pkg install'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["i", "in", "ins", "inst", "insta", "instal"])]
    Install(PkgInstall),
    #[cfg(not(target_os = "macos"))]
    /// Alias for 'sup run'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Run(SupRun),
    /// Alias for 'cli setup'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["set", "setu"])]
    Setup(CacheKeyPath),
    /// Alias for 'svc start'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["sta", "star"])]
    Start(SvcStart),
    /// Alias for 'svc stop'
    #[structopt(no_version, settings = &[AppSettings::Hidden], aliases = &["sto"])]
    Stop(SvcStop),
    #[cfg(not(target_os = "macos"))]
    /// Alias for 'sup term'
    #[structopt(no_version, settings = &[AppSettings::Hidden])]
    Term,
}

#[derive(ConfigOpt, StructOpt)]
#[structopt(no_version, aliases = &["b", "bl", "bld"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
/// Commands relating to Habitat Builder
pub enum Bldr {
    #[structopt(no_version)]
    Channel(Channel),
    #[structopt(no_version)]
    Job(Job),
}

#[derive(ConfigOpt, StructOpt)]
#[structopt(no_version, aliases = &["c", "ch", "cha", "chan", "chann", "channe"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
/// Commands relating to Habitat Builder channels
pub enum Channel {
    #[structopt(no_version, aliases = &["c", "cr", "cre", "crea", "creat"])]
    Create(ChannelCreate),
    Demote(ChannelDemote),
    #[structopt(no_version, aliases = &["d", "de", "des", "dest", "destr", "destro"])]
    Destroy(ChannelDestroy),
    #[structopt(no_version, aliases = &["l", "li", "lis"])]
    List(ChannelList),
    Promote(ChannelPromote),
}

#[derive(ConfigOpt, StructOpt)]
#[structopt(no_version, aliases = &["j", "jo"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
/// Commands relating to Habitat Builder jobs
pub enum Job {
    #[structopt(no_version, aliases = &["c", "ca", "can", "cance", "cancel"])]
    Cancel(JobCancel),
    #[structopt(no_version, aliases = &["d", "de", "dem", "demo", "demot"])]
    Demote(JobDemote),
    #[structopt(no_version, aliases = &["p", "pr", "pro", "prom", "promo", "promot"])]
    Promote(JobPromote),
    #[structopt(no_version, aliases = &["s", "st", "sta", "star"])]
    Start(JobStart),
    #[structopt(no_version, aliases = &["stat", "statu"])]
    Status(JobStatus),
}

/// Commands relating to Habitat users
#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "user", no_version, aliases = &["u", "us", "use"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
pub enum User {
    Key(UserKey),
}

#[derive(ConfigOpt, StructOpt)]
#[structopt(name = "key", no_version, aliases = &["k", "ke"], settings = &[AppSettings::ArgRequiredElseHelp, AppSettings::SubcommandRequiredElseHelp])]
/// Commands relating to Habitat user keys
pub enum UserKey {
    #[structopt(no_version, aliases = &["g", "ge", "gen", "gene", "gener", "genera", "generat"])]
    Generate(UserKeyGenerate),
}<|endoftext|>


<commit_before>use core::sync::atomic::{AtomicUsize, Ordering};

use crate::context::timeout;
use crate::device::{local_apic, ioapic, pic};
use crate::device::serial::{COM1, COM2};
use crate::ipi::{ipi, IpiKind, IpiTarget};
use crate::scheme::debug::debug_input;
use crate::{context, ptrace, time};

//resets to 0 in context::switch()
#[thread_local]
pub static PIT_TICKS: AtomicUsize = AtomicUsize::new(0);

#[repr(u8)]
pub enum IrqMethod {
    Pic = 0,
    Apic = 1,
}

static IRQ_METHOD: AtomicUsize = AtomicUsize::new(IrqMethod::Pic as usize);

pub fn set_irq_method(method: IrqMethod) {
    IRQ_METHOD.store(method as usize, core::sync::atomic::Ordering::Release);
}

fn irq_method() -> IrqMethod {
    let raw = IRQ_METHOD.load(core::sync::atomic::Ordering::Acquire);

    match raw {
        0 => IrqMethod::Pic,
        1 => IrqMethod::Apic,
        _ => unreachable!(),
    }
}

extern {
    // triggers irq scheme
    fn irq_trigger(irq: u8);
}

/// Notify the IRQ scheme that an IRQ has been registered. This should mask the IRQ until the
/// scheme user unmasks it ("acknowledges" it).
unsafe fn trigger(irq: u8) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_mask(irq) },
        IrqMethod::Apic => ioapic_mask(irq),
    }
    irq_trigger(irq);
}

/// Unmask the IRQ. This is called from the IRQ scheme, which does this when a user process has
/// processed the IRQ.
pub unsafe fn acknowledge(irq: usize) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_unmask(irq) },
        IrqMethod::Apic => ioapic_unmask(irq),
    }
}

/// Sends an end-of-interrupt, so that the interrupt controller can go on to the next one.
pub unsafe fn eoi(irq: u8) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_eoi(irq) },
        IrqMethod::Apic => lapic_eoi(),
    }
}

unsafe fn pic_mask(irq: u8) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::SLAVE.mask_set(irq - 8);
    } else {
        pic::MASTER.mask_set(irq);
    }
}

unsafe fn ioapic_mask(irq: u8) {
    ioapic::mask(irq);
}

unsafe fn pic_eoi(irq: u8) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::MASTER.ack();
        if irq == 15 {
            //TODO: check spurious
            return;
        }
        pic::SLAVE.ack();
    } else {
        if irq == 7 {
            //TODO: check spurious
            return;
        }
        pic::MASTER.ack();
    }
}

unsafe fn lapic_eoi() {
    local_apic::LOCAL_APIC.eoi()
}

unsafe fn pic_unmask(irq: usize) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::SLAVE.mask_clear(irq as u8 - 8);
    } else {
        pic::MASTER.mask_clear(irq as u8);
    }
}

unsafe fn ioapic_unmask(irq: usize) {
    ioapic::unmask(irq as u8);
}

interrupt_stack!(pit, stack, {
    // Saves CPU time by not sending IRQ event irq_trigger(0);

    const PIT_RATE: u64 = 2_250_286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1_000_000_000;
        offset.0 += sum / 1_000_000_000;
    }

    eoi(0);

    // Wake up other CPUs
    ipi(IpiKind::Pit, IpiTarget::Other);

    // Any better way of doing this?
    timeout::trigger();

    if PIT_TICKS.fetch_add(1, Ordering::SeqCst) >= 10 {
        let _guard = ptrace::set_process_regs(stack);
        let _ = context::switch();
    }
});

interrupt!(keyboard, {
    trigger(1);
    eoi(1);
});

interrupt!(cascade, {
    // No need to do any operations on cascade
    eoi(2);
});

interrupt!(com2, {
    while let Some(c) = COM2.lock().receive() {
        debug_input(c);
    }
    eoi(3);
});

interrupt!(com1, {
    while let Some(c) = COM1.lock().receive() {
        debug_input(c);
    }
    eoi(4);
});

interrupt!(lpt2, {
    trigger(5);
    eoi(5);
});

interrupt!(floppy, {
    trigger(6);
    eoi(6);
});

interrupt!(lpt1, {
    trigger(7);
    eoi(7);
});

interrupt!(rtc, {
    trigger(8);
    eoi(8);
});

interrupt!(pci1, {
    trigger(9);
    eoi(9);
});

interrupt!(pci2, {
    trigger(10);
    eoi(10);
});

interrupt!(pci3, {
    trigger(11);
    eoi(11);
});

interrupt!(mouse, {
    trigger(12);
    eoi(12);
});

interrupt!(fpu, {
    trigger(13);
    eoi(13);
});

interrupt!(ata1, {
    trigger(14);
    eoi(14);
});

interrupt!(ata2, {
    trigger(15);
    eoi(15);
});

interrupt!(lapic_timer, {
    println!("Local apic timer interrupt");
    lapic_eoi();
});

interrupt!(lapic_error, {
    println!("Local apic internal error: ESR={:#0x}", local_apic::LOCAL_APIC.esr());
    lapic_eoi();
});

interrupt!(calib_pit, {
    const PIT_RATE: u64 = 2_250_286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1_000_000_000;
        offset.0 += sum / 1_000_000_000;
    }

    eoi(0);
});
// XXX: This would look way prettier using const generics.

macro_rules! allocatable_irq(
    ( $number:literal, $name:ident ) => {
        interrupt!($name, {
            allocatable_irq_generic($number);
        });
    }
);

pub unsafe fn allocatable_irq_generic(number: u8) {
    irq_trigger(number - 32);
    lapic_eoi();
}

define_default_irqs!();
<commit_msg>Kernel reading of ps2 bytes<commit_after>use core::sync::atomic::{AtomicUsize, Ordering};

use crate::context::timeout;
use crate::device::{local_apic, ioapic, pic};
use crate::device::serial::{COM1, COM2};
use crate::ipi::{ipi, IpiKind, IpiTarget};
use crate::scheme::debug::debug_input;
use crate::{context, ptrace, time};

//resets to 0 in context::switch()
#[thread_local]
pub static PIT_TICKS: AtomicUsize = AtomicUsize::new(0);

// The only way to read PS2 data without race conditions is to allow a keyboard interrupt to happen
// and then read data while reading mouse data, since keyboard data overrides mouse data and
// reading the status register is not done atomically with reading the data. This is not possible
// from userspace, so we do this minimal part of the PS2 driver in the kernel.
#[inline(always)]
unsafe fn ps2_interrupt(index: usize) {
    use crate::scheme::serio::serio_input;

    let data: u8;
    let status: u8;
    asm!("
        sti
        nop
        cli
        in al, 0x64
        mov ah, al
        in al, 0x60
        "
        : "={al}"(data), "={ah}"(status)
        :
        : "memory"
        : "intel", "volatile"
    );

    if status & 1 != 0 {
        let status_index = if status & (1 << 5) == 0 {
            // Keyboard, according to status
            0
        } else {
            // Mouse, according to status
            1
        };
        serio_input(status_index, data);
    }
}

#[repr(u8)]
pub enum IrqMethod {
    Pic = 0,
    Apic = 1,
}

static IRQ_METHOD: AtomicUsize = AtomicUsize::new(IrqMethod::Pic as usize);

pub fn set_irq_method(method: IrqMethod) {
    IRQ_METHOD.store(method as usize, core::sync::atomic::Ordering::Release);
}

fn irq_method() -> IrqMethod {
    let raw = IRQ_METHOD.load(core::sync::atomic::Ordering::Acquire);

    match raw {
        0 => IrqMethod::Pic,
        1 => IrqMethod::Apic,
        _ => unreachable!(),
    }
}

extern {
    // triggers irq scheme
    fn irq_trigger(irq: u8);
}

/// Notify the IRQ scheme that an IRQ has been registered. This should mask the IRQ until the
/// scheme user unmasks it ("acknowledges" it).
unsafe fn trigger(irq: u8) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_mask(irq) },
        IrqMethod::Apic => ioapic_mask(irq),
    }
    irq_trigger(irq);
}

/// Unmask the IRQ. This is called from the IRQ scheme, which does this when a user process has
/// processed the IRQ.
pub unsafe fn acknowledge(irq: usize) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_unmask(irq) },
        IrqMethod::Apic => ioapic_unmask(irq),
    }
}

/// Sends an end-of-interrupt, so that the interrupt controller can go on to the next one.
pub unsafe fn eoi(irq: u8) {
    match irq_method() {
        IrqMethod::Pic => if irq < 16 { pic_eoi(irq) },
        IrqMethod::Apic => lapic_eoi(),
    }
}

unsafe fn pic_mask(irq: u8) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::SLAVE.mask_set(irq - 8);
    } else {
        pic::MASTER.mask_set(irq);
    }
}

unsafe fn ioapic_mask(irq: u8) {
    ioapic::mask(irq);
}

unsafe fn pic_eoi(irq: u8) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::MASTER.ack();
        if irq == 15 {
            //TODO: check spurious
            return;
        }
        pic::SLAVE.ack();
    } else {
        if irq == 7 {
            //TODO: check spurious
            return;
        }
        pic::MASTER.ack();
    }
}

unsafe fn lapic_eoi() {
    local_apic::LOCAL_APIC.eoi()
}

unsafe fn pic_unmask(irq: usize) {
    debug_assert!(irq < 16);

    if irq >= 8 {
        pic::SLAVE.mask_clear(irq as u8 - 8);
    } else {
        pic::MASTER.mask_clear(irq as u8);
    }
}

unsafe fn ioapic_unmask(irq: usize) {
    ioapic::unmask(irq as u8);
}

interrupt_stack!(pit, stack, {
    // Saves CPU time by not sending IRQ event irq_trigger(0);

    const PIT_RATE: u64 = 2_250_286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1_000_000_000;
        offset.0 += sum / 1_000_000_000;
    }

    eoi(0);

    // Wake up other CPUs
    ipi(IpiKind::Pit, IpiTarget::Other);

    // Any better way of doing this?
    timeout::trigger();

    if PIT_TICKS.fetch_add(1, Ordering::SeqCst) >= 10 {
        let _guard = ptrace::set_process_regs(stack);
        let _ = context::switch();
    }
});

interrupt!(keyboard, {
    ps2_interrupt(0);
    eoi(1);
});

interrupt!(cascade, {
    // No need to do any operations on cascade
    eoi(2);
});

interrupt!(com2, {
    while let Some(c) = COM2.lock().receive() {
        debug_input(c);
    }
    eoi(3);
});

interrupt!(com1, {
    while let Some(c) = COM1.lock().receive() {
        debug_input(c);
    }
    eoi(4);
});

interrupt!(lpt2, {
    trigger(5);
    eoi(5);
});

interrupt!(floppy, {
    trigger(6);
    eoi(6);
});

interrupt!(lpt1, {
    trigger(7);
    eoi(7);
});

interrupt!(rtc, {
    trigger(8);
    eoi(8);
});

interrupt!(pci1, {
    trigger(9);
    eoi(9);
});

interrupt!(pci2, {
    trigger(10);
    eoi(10);
});

interrupt!(pci3, {
    trigger(11);
    eoi(11);
});

interrupt!(mouse, {
    ps2_interrupt(1);
    eoi(12);
});

interrupt!(fpu, {
    trigger(13);
    eoi(13);
});

interrupt!(ata1, {
    trigger(14);
    eoi(14);
});

interrupt!(ata2, {
    trigger(15);
    eoi(15);
});

interrupt!(lapic_timer, {
    println!("Local apic timer interrupt");
    lapic_eoi();
});

interrupt!(lapic_error, {
    println!("Local apic internal error: ESR={:#0x}", local_apic::LOCAL_APIC.esr());
    lapic_eoi();
});

interrupt!(calib_pit, {
    const PIT_RATE: u64 = 2_250_286;

    {
        let mut offset = time::OFFSET.lock();
        let sum = offset.1 + PIT_RATE;
        offset.1 = sum % 1_000_000_000;
        offset.0 += sum / 1_000_000_000;
    }

    eoi(0);
});
// XXX: This would look way prettier using const generics.

macro_rules! allocatable_irq(
    ( $number:literal, $name:ident ) => {
        interrupt!($name, {
            allocatable_irq_generic($number);
        });
    }
);

pub unsafe fn allocatable_irq_generic(number: u8) {
    irq_trigger(number - 32);
    lapic_eoi();
}

define_default_irqs!();
<|endoftext|>