<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.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return unsafe { &*tls };
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

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

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            unsafe { thread_data.parker.prepare_park(); }

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            unsafe { thread_data.parker.park(); }

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = unsafe { (*current).queue_tail.get() };
                if !queue_tail.is_null() {
                    break;
                }
                unsafe {
                    let next = (*current).next.get();
                    (*next).prev.set(current);
                    current = next;
                }
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            unsafe {
                (*queue_head).queue_tail.set(queue_tail);
            }

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = unsafe { (*queue_tail).prev.get() };
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                unsafe {
                    (*queue_head).queue_tail.set(new_tail);
                }
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            unsafe {
                (*queue_tail).parker.unpark_lock().unpark();
            }
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<commit_msg>Add warning about why unlock is unsafe<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.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return unsafe { &*tls };
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    /// Must not be called on an already unlocked `WordLock`!
    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

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

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            unsafe { thread_data.parker.prepare_park(); }

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            unsafe { thread_data.parker.park(); }

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = unsafe { (*current).queue_tail.get() };
                if !queue_tail.is_null() {
                    break;
                }
                unsafe {
                    let next = (*current).next.get();
                    (*next).prev.set(current);
                    current = next;
                }
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            unsafe {
                (*queue_head).queue_tail.set(queue_tail);
            }

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = unsafe { (*queue_tail).prev.get() };
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                unsafe {
                    (*queue_head).queue_tail.set(new_tail);
                }
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            unsafe {
                (*queue_tail).parker.unpark_lock().unpark();
            }
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<|endoftext|>


<commit_before>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        // let mut dst = vec![0.0; nodes];                         // local rank accumulation

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    let mut session = output.session(&iter);

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    let mut counter = 0u64;

                    for node in 0..src.len() {
                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            session.give((b, value));
                            counter += 1;
                        }
                    }
                    // \------ end familiar part ------/

                    println!("iteration {:?}: {}s; sent {} pairs", iter, time::precise_time_s() - start, counter);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        src[node as usize / peers] += rank;
                    }
                }
            }
        );

        // let local_index = index as usize % workers;
        // let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }
}
<commit_msg>update exotic pagerank<commit_after>// #![feature(scoped)]
// #![feature(collections)]

extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate dataflow_join;

extern crate docopt;
use docopt::Docopt;

use std::thread;
use std::mem;

use dataflow_join::graph::{GraphTrait, GraphMMap};

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_static::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      list of host:port for workers
";

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    // let workers = if let Ok(threads) = args.get_str("<workers>").parse() { threads }
    //               else { panic!("invalid setting for workers: {}", args.get_str("<workers>")) };
    // println!("starting pagerank dataflow with {:?} worker{}", workers, if workers == 1 { "" } else { "s" });
    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_multi(communicators, source);
    }
    else if workers > 1 {
        println!("Initializing ProcessCommunicator");
        pagerank_multi(ProcessCommunicator::new_vector(workers), source);
    }
    else {
        println!("Initializing ThreadCommunicator");
        pagerank_multi(vec![ThreadCommunicator], source);
    };
}

fn pagerank_multi<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

fn pagerank<C>(communicator: C, filename: String, workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    {   // new scope avoids long borrow on root
        let mut builder = root.new_subgraph();

        // establish the beginnings of a loop,
        // 20 iterations, each time around += 1.
        let (helper, stream) = builder.loop_variable::<(u32, f32)>(RootTimestamp::new(20), Local(1));

        let graph = GraphMMap::<u32>::new(&filename);

        let nodes = graph.nodes();

        let mut src = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        let mut tmp = vec![1.0; 1 + (nodes / peers as usize)];  // local rank accumulation
        // let mut dst = vec![0.0; nodes];                         // local rank accumulation

        let mut start = time::precise_time_s();

        // from feedback, place an operator that
        // aggregates and broadcasts ranks along edges.
        let ranks = stream.enable(builder).unary_notify(

            Exchange::new(|x: &(u32, f32)| x.0 as u64),     // 1. how data should be exchanged
            format!("PageRank"),                            // 2. a tasteful, descriptive name
            vec![RootTimestamp::new(0)],                    // 3. indicate an initial capability
            move |input, output, iterator| {                // 4. provide the operator logic

                while let Some((iter, _)) = iterator.next() {

                    mem::swap(&mut src, &mut tmp);

                    // /---- should look familiar! ----\
                    for node in 0..src.len() {
                        src[node] = 0.15 + 0.85 * src[node];
                    }

                    for node in 0..src.len() {
                        let mut session = output.session(&iter);

                        let edges = graph.edges(index + peers * node);
                        let value = src[node] / edges.len() as f32;
                        for &b in edges {
                            session.give((b, value));
                        }

                        if let Some((iter, data)) = input.pull() {
                            iterator.notify_at(&iter);
                            for (node, rank) in data.drain_temp() {
                                tmp[node as usize / peers] += rank;
                            }
                        }
                    }
                    // \------ end familiar part ------/

                    println!("iteration {:?}: {}s", iter, time::precise_time_s() - start);
                    start = time::precise_time_s();
                }

                while let Some((iter, data)) = input.pull() {
                    iterator.notify_at(&iter);
                    for (node, rank) in data.drain_temp() {
                        tmp[node as usize / peers] += rank;
                    }
                }
            }
        );

        // let local_index = index as usize % workers;
        // let mut acc = vec![0.0; 1 + (nodes / workers)];

        ranks
        // .unary_notify(
        //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
        //     format!("Aggregation"),
        //     vec![],
        //     move |input, output, iterator| {
        //         while let Some((iter, data)) = input.pull() {
        //             iterator.notify_at(&iter);
        //             for (node, rank) in data.drain_temp() {
        //                 acc[node as usize / workers] += rank;
        //             }
        //         }
        //
        //         while let Some((item, _)) = iterator.next() {
        //
        //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
        //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
        //
        //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
        //             assert!(acc.len() == (1 + (nodes/workers)));
        //         }
        //     }
        // )
        .connect_loop(helper);
    }

    while root.step() { }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

//! Core primitive types for relm.
//!
//! The primary type is `EventStream`.

#![warn(
    missing_docs,
    trivial_casts,
    trivial_numeric_casts,
    unused_extern_crates,
    unused_import_braces,
    unused_qualifications,
)]

extern crate glib;
extern crate glib_sys;
extern crate libc;

mod source;

use std::cell::RefCell;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::rc::Rc;
use std::sync::mpsc::{self, Receiver, Sender};

use source::{SourceFuncs, new_source, source_get};

use glib::{MainContext, Source};

/// A lock is used to temporarily stop emitting messages.
#[must_use]
pub struct Lock<MSG> {
    stream: Rc<RefCell<_EventStream<MSG>>>,
}

impl<MSG> Drop for Lock<MSG> {
    fn drop(&mut self) {
        self.stream.borrow_mut().locked = false;
    }
}

struct ChannelData<MSG> {
    callback: Box<FnMut(MSG)>,
    peeked_value: Option<MSG>,
    receiver: Receiver<MSG>,
}

/// A channel to send a message to a relm widget from another thread.
pub struct Channel<MSG> {
    _source: Source,
    _phantom: PhantomData<MSG>,
}

impl<MSG> Channel<MSG> {
    /// Create a new channel with a callback that will be called when a message is received.
    pub fn new<CALLBACK: FnMut(MSG) + 'static>(callback: CALLBACK) -> (Self, Sender<MSG>) {
        let (sender, receiver) = mpsc::channel();
        let source = new_source(RefCell::new(ChannelData {
            callback: Box::new(callback),
            peeked_value: None,
            receiver,
        }));
        let main_context = MainContext::default().expect("no main context");
        source.attach(&main_context);
        (Self {
            _source: source,
            _phantom: PhantomData,
        }, sender)
    }
}

impl<MSG> SourceFuncs for RefCell<ChannelData<MSG>> {
    fn dispatch(&self) -> bool {
        // TODO: show errors.
        let msg = self.borrow_mut().peeked_value.take().or_else(|| {
            self.borrow().receiver.try_recv().ok()
        });
        if let Some(msg) = msg {
            let callback = &mut self.borrow_mut().callback;
            callback(msg);
        }
        true
    }

    fn prepare(&self) -> (bool, Option<u32>) {
        if self.borrow().peeked_value.is_some() {
            return (true, None);
        }
        let peek_val = self.borrow().receiver.try_recv().ok();
        self.borrow_mut().peeked_value = peek_val;
        (self.borrow().peeked_value.is_some(), None)
    }

}

struct _EventStream<MSG> {
    events: VecDeque<MSG>,
    locked: bool,
    observers: Vec<Rc<Fn(&MSG)>>,
}

impl<MSG> SourceFuncs for SourceData<MSG> {
    fn dispatch(&self) -> bool {
        let event = self.stream.borrow_mut().events.pop_front();
        if let (Some(event), Some(callback)) = (event, self.callback.borrow_mut().as_mut()) {
            callback(event);
        }
        true
    }

    fn prepare(&self) -> (bool, Option<u32>) {
        (!self.stream.borrow().events.is_empty(), None)
    }

}

struct SourceData<MSG> {
    callback: Rc<RefCell<Option<Box<FnMut(MSG)>>>>,
    stream: Rc<RefCell<_EventStream<MSG>>>,
}

/// A stream of messages to be used for widget/signal communication and inter-widget communication.
pub struct EventStream<MSG> {
    source: Source,
    _phantom: PhantomData<MSG>,
}

impl<MSG> Clone for EventStream<MSG> {
    fn clone(&self) -> Self {
        EventStream {
            source: self.source.clone(),
            _phantom: PhantomData,
        }
    }
}

impl<MSG> EventStream<MSG> {
    fn get_callback(&self) -> Rc<RefCell<Option<Box<FnMut(MSG)>>>> {
        source_get::<SourceData<MSG>>(&self.source).callback.clone()
    }

    fn get_stream(&self) -> Rc<RefCell<_EventStream<MSG>>> {
        source_get::<SourceData<MSG>>(&self.source).stream.clone()
    }
}

impl<MSG> EventStream<MSG> {
    /// Create a new event stream.
    pub fn new() -> Self {
        let event_stream: _EventStream<MSG> = _EventStream {
            events: VecDeque::new(),
            locked: false,
            observers: vec![],
        };
        let source = new_source(SourceData {
            callback: Rc::new(RefCell::new(None)),
            stream: Rc::new(RefCell::new(event_stream)),
        });
        let main_context = MainContext::default().expect("no main context");
        source.attach(&main_context);
        EventStream {
            source,
            _phantom: PhantomData,
        }
    }

    /// Close the event stream, i.e. stop processing messages.
    pub fn close(&self) {
        self.source.destroy();
    }

    /// Send the `event` message to the stream and the observers.
    pub fn emit(&self, event: MSG) {
        let stream = self.get_stream();
        if !stream.borrow().locked {
            let len = stream.borrow().observers.len();
            for i in 0..len {
                let observer = stream.borrow().observers[i].clone();
                observer(&event);
            }

            stream.borrow_mut().events.push_back(event);
        }
    }

    /// Lock the stream (don't emit message) until the `Lock` goes out of scope.
    pub fn lock(&self) -> Lock<MSG> {
        let stream = self.get_stream();
        stream.borrow_mut().locked = true;
        Lock {
            stream: self.get_stream().clone(),
        }
    }

    /// Add an observer to the event stream.
    /// This callback will be called every time a message is emmited.
    pub fn observe<CALLBACK: Fn(&MSG) + 'static>(&self, callback: CALLBACK) {
        let stream = self.get_stream();
        stream.borrow_mut().observers.push(Rc::new(callback));
    }

    /// Add a callback to the event stream.
    /// This is the main callback and received a owned version of the message, in contrast to
    /// observe().
    pub fn set_callback<CALLBACK: FnMut(MSG) + 'static>(&self, callback: CALLBACK) {
        let source_callback = self.get_callback();
        *source_callback.borrow_mut() = Some(Box::new(callback));
    }
}
<commit_msg>Fix channel to wakeup the event loop when sending a message<commit_after>/*
 * Copyright (c) 2017 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

//! Core primitive types for relm.
//!
//! The primary type is `EventStream`.

#![warn(
    missing_docs,
    trivial_casts,
    trivial_numeric_casts,
    unused_extern_crates,
    unused_import_braces,
    unused_qualifications,
)]

extern crate glib;
extern crate glib_sys;
extern crate libc;

mod source;

use std::cell::RefCell;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::rc::Rc;
use std::sync::mpsc::{self, Receiver, SendError};

use source::{SourceFuncs, new_source, source_get};

use glib::{MainContext, Source};

/// A lock is used to temporarily stop emitting messages.
#[must_use]
pub struct Lock<MSG> {
    stream: Rc<RefCell<_EventStream<MSG>>>,
}

impl<MSG> Drop for Lock<MSG> {
    fn drop(&mut self) {
        self.stream.borrow_mut().locked = false;
    }
}

struct ChannelData<MSG> {
    callback: Box<FnMut(MSG)>,
    peeked_value: Option<MSG>,
    receiver: Receiver<MSG>,
}

/// A wrapper over a `std::sync::mpsc::Sender` to wakeup the glib event loop when sending a
/// message.
#[derive(Clone)]
pub struct Sender<MSG> {
    sender: mpsc::Sender<MSG>,
}

impl<MSG> Sender<MSG> {
    /// Send a message and wakeup the event loop.
    pub fn send(&self, msg: MSG) -> Result<(), SendError<MSG>> {
        let result = self.sender.send(msg);
        if let Some(context) = MainContext::default() {
            context.wakeup();
        }
        result
    }
}

/// A channel to send a message to a relm widget from another thread.
pub struct Channel<MSG> {
    _source: Source,
    _phantom: PhantomData<MSG>,
}

impl<MSG> Channel<MSG> {
    /// Create a new channel with a callback that will be called when a message is received.
    pub fn new<CALLBACK: FnMut(MSG) + 'static>(callback: CALLBACK) -> (Self, Sender<MSG>) {
        let (sender, receiver) = mpsc::channel();
        let source = new_source(RefCell::new(ChannelData {
            callback: Box::new(callback),
            peeked_value: None,
            receiver,
        }));
        let main_context = MainContext::default().expect("no main context");
        source.attach(&main_context);
        (Self {
            _source: source,
            _phantom: PhantomData,
        }, Sender {
            sender,
        })
    }
}

impl<MSG> SourceFuncs for RefCell<ChannelData<MSG>> {
    fn dispatch(&self) -> bool {
        // TODO: show errors.
        let msg = self.borrow_mut().peeked_value.take().or_else(|| {
            self.borrow().receiver.try_recv().ok()
        });
        if let Some(msg) = msg {
            let callback = &mut self.borrow_mut().callback;
            callback(msg);
        }
        true
    }

    fn prepare(&self) -> (bool, Option<u32>) {
        if self.borrow().peeked_value.is_some() {
            return (true, None);
        }
        let peek_val = self.borrow().receiver.try_recv().ok();
        self.borrow_mut().peeked_value = peek_val;
        (self.borrow().peeked_value.is_some(), None)
    }

}

struct _EventStream<MSG> {
    events: VecDeque<MSG>,
    locked: bool,
    observers: Vec<Rc<Fn(&MSG)>>,
}

impl<MSG> SourceFuncs for SourceData<MSG> {
    fn dispatch(&self) -> bool {
        let event = self.stream.borrow_mut().events.pop_front();
        if let (Some(event), Some(callback)) = (event, self.callback.borrow_mut().as_mut()) {
            callback(event);
        }
        true
    }

    fn prepare(&self) -> (bool, Option<u32>) {
        (!self.stream.borrow().events.is_empty(), None)
    }

}

struct SourceData<MSG> {
    callback: Rc<RefCell<Option<Box<FnMut(MSG)>>>>,
    stream: Rc<RefCell<_EventStream<MSG>>>,
}

/// A stream of messages to be used for widget/signal communication and inter-widget communication.
pub struct EventStream<MSG> {
    source: Source,
    _phantom: PhantomData<MSG>,
}

impl<MSG> Clone for EventStream<MSG> {
    fn clone(&self) -> Self {
        EventStream {
            source: self.source.clone(),
            _phantom: PhantomData,
        }
    }
}

impl<MSG> EventStream<MSG> {
    fn get_callback(&self) -> Rc<RefCell<Option<Box<FnMut(MSG)>>>> {
        source_get::<SourceData<MSG>>(&self.source).callback.clone()
    }

    fn get_stream(&self) -> Rc<RefCell<_EventStream<MSG>>> {
        source_get::<SourceData<MSG>>(&self.source).stream.clone()
    }
}

impl<MSG> EventStream<MSG> {
    /// Create a new event stream.
    pub fn new() -> Self {
        let event_stream: _EventStream<MSG> = _EventStream {
            events: VecDeque::new(),
            locked: false,
            observers: vec![],
        };
        let source = new_source(SourceData {
            callback: Rc::new(RefCell::new(None)),
            stream: Rc::new(RefCell::new(event_stream)),
        });
        let main_context = MainContext::default().expect("no main context");
        source.attach(&main_context);
        EventStream {
            source,
            _phantom: PhantomData,
        }
    }

    /// Close the event stream, i.e. stop processing messages.
    pub fn close(&self) {
        self.source.destroy();
    }

    /// Send the `event` message to the stream and the observers.
    pub fn emit(&self, event: MSG) {
        let stream = self.get_stream();
        if !stream.borrow().locked {
            let len = stream.borrow().observers.len();
            for i in 0..len {
                let observer = stream.borrow().observers[i].clone();
                observer(&event);
            }

            stream.borrow_mut().events.push_back(event);
        }
    }

    /// Lock the stream (don't emit message) until the `Lock` goes out of scope.
    pub fn lock(&self) -> Lock<MSG> {
        let stream = self.get_stream();
        stream.borrow_mut().locked = true;
        Lock {
            stream: self.get_stream().clone(),
        }
    }

    /// Add an observer to the event stream.
    /// This callback will be called every time a message is emmited.
    pub fn observe<CALLBACK: Fn(&MSG) + 'static>(&self, callback: CALLBACK) {
        let stream = self.get_stream();
        stream.borrow_mut().observers.push(Rc::new(callback));
    }

    /// Add a callback to the event stream.
    /// This is the main callback and received a owned version of the message, in contrast to
    /// observe().
    pub fn set_callback<CALLBACK: FnMut(MSG) + 'static>(&self, callback: CALLBACK) {
        let source_callback = self.get_callback();
        *source_callback.borrow_mut() = Some(Box::new(callback));
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::default::Default;
use std::os;
use syntax::ast;
use syntax::codemap;
use syntax::ext::base;
use syntax::parse;
use syntax::parse::token;
use syntax::util::small_vector::SmallVector;

use super::{generate_bindings, BindgenOptions, Logger};

pub fn bindgen_macro(cx: &mut base::ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<base::MacResult> {
    let mut visit = BindgenArgsVisitor {
        options: Default::default(),
        seen_named: false
    };

    visit.options.builtins = true;
    if !parse_macro_opts(cx, tts, &mut visit) {
        return base::DummyResult::any(sp);
    }

    // Reparse clang_args as it is passed in string form
    let clang_args = visit.options.clang_args.connect(" ");
    visit.options.clang_args = parse_process_args(clang_args.as_slice());

    // Set the working dir to the directory containing the invoking rs file so
    // that clang searches for headers relative to it rather than the crate root
    let mod_dir = Vec::from_slice(Path::new(cx.codemap().span_to_filename(sp)).dirname());
    let cwd = os::getcwd();
    os::change_dir(&Path::new(mod_dir));

    // We want the span for errors to just match the bindgen! symbol
    // instead of the whole invocation which can span multiple lines
    let mut short_span = sp;
    short_span.hi = short_span.lo + codemap::BytePos(8);

    let logger = MacroLogger { sp: short_span, cx: cx };

    let ret = match generate_bindings(visit.options, Some(&logger as &Logger), short_span) {
        Ok(items) => {
            box BindgenResult { items: RefCell::new(Some(SmallVector::many(items))) } as Box<base::MacResult>
        }
        Err(_) => base::DummyResult::any(sp)
    };

    os::change_dir(&Path::new(cwd));

    ret
}

trait MacroArgsVisitor {
    fn visit_str(&mut self, name: Option<&str>, val: &str) -> bool;
    fn visit_int(&mut self, name: Option<&str>, val: i64) -> bool;
    fn visit_bool(&mut self, name: Option<&str>, val: bool) -> bool;
    fn visit_ident(&mut self, name: Option<&str>, ident: &str) -> bool;
}

struct BindgenArgsVisitor {
    pub options: BindgenOptions,
    seen_named: bool
}

impl MacroArgsVisitor for BindgenArgsVisitor {
    fn visit_str(&mut self, mut name: Option<&str>, val: &str) -> bool {
        if name.is_some() { self.seen_named = true; }
        else if !self.seen_named { name = Some("clang_args") }
        match name {
            Some("link") => self.options.links.push((val.to_string(), None)),
            Some("link_static") => self.options.links.push((val.to_string(), Some("static".to_string()))),
            Some("link_framework") => self.options.links.push((val.to_string(), Some("framework".to_string()))),
            Some("abi") => self.options.abi = val.to_string(),
            Some("match") => self.options.match_pat.push(val.to_string()),
            Some("clang_args") => self.options.clang_args.push(val.to_string()),
            _ => return false
        }
        true
    }

    #[allow(unused_variable)]
    fn visit_int(&mut self, name: Option<&str>, val: i64) -> bool {
        if name.is_some() { self.seen_named = true; }
        false
    }

    fn visit_bool(&mut self, name: Option<&str>, val: bool) -> bool {
        if name.is_some() { self.seen_named = true; }
        match name {
            Some("allow_bitfields") => self.options.fail_on_bitfield = !val,
            Some("allow_unknown_types") => self.options.fail_on_unknown_type = !val,
            Some("emit_builtins") => self.options.builtins = val,
            _ => return false
        }
        true
    }

    #[allow(unused_variable)]
    fn visit_ident(&mut self, name: Option<&str>, val: &str) -> bool {
        if name.is_some() { self.seen_named = true; }
        false
    }
}

// Parses macro invocations in the form [ident=|:]value where value is an ident or literal
// e.g. bindgen!(module_name, "header.h", emit_builtins=false, clang_args:"-I /usr/local/include")
fn parse_macro_opts(cx: &mut base::ExtCtxt, tts: &[ast::TokenTree], visit: &mut MacroArgsVisitor) -> bool {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(), Vec::from_slice(tts));
    let mut args_good = true;

    loop {
        let mut name: Option<String> = None;
        let mut span = parser.span;

        // Check for [ident=]value and if found save ident to name
        if parser.look_ahead(1, |t| t == &token::EQ) {
            match parser.bump_and_get() {
                token::IDENT(ident, _) => {
                    let ident = parser.id_to_interned_str(ident);
                    name = Some(ident.get().to_string());
                    parser.expect(&token::EQ);
                },
                _ => {
                    cx.span_err(span, "invalid argument format");
                    return false
                }
            }
        }

        match parser.token {
            // Match [ident]
            token::IDENT(val, _) => {
                let val = parser.id_to_interned_str(val);
                span.hi = parser.span.hi;
                parser.bump();

                // Bools are simply encoded as idents
                let ret = match val.get() {
                    "true" => visit.visit_bool(as_str(&name), true),
                    "false" => visit.visit_bool(as_str(&name), false),
                    val => visit.visit_ident(as_str(&name), val)
                };
                if !ret {
                    cx.span_err(span, "invalid argument");
                    args_good = false;
                }
            }
            // Match [literal] and parse as an expression so we can expand macros
            _ => {
                let expr = cx.expand_expr(parser.parse_expr());
                span.hi = expr.span.hi;
                match expr.node {
                    ast::ExprLit(lit) => {
                        let ret = match lit.node {
                            ast::LitStr(ref s, _) => visit.visit_str(as_str(&name), s.get()),
                            ast::LitBool(b) => visit.visit_bool(as_str(&name), b),
                            ast::LitIntUnsuffixed(i) |
                            ast::LitInt(i, _) => visit.visit_int(as_str(&name), i),
                            ast::LitUint(i, _) => visit.visit_int(as_str(&name), i as i64),
                            _ => {
                                cx.span_err(span, "invalid argument format");
                                return false
                            }
                        };
                        if !ret {
                            cx.span_err(span, "invalid argument");
                            args_good = false;
                        }
                    },
                    _ => {
                        cx.span_err(span, "invalid argument format");
                        return false
                    }
                }
            }
        }

        if parser.eat(&token::EOF) {
            return args_good
        }

        if !parser.eat(&token::COMMA) {
            cx.span_err(parser.span, "invalid argument format");
            return false
        }
    }
}

// I'm sure there's a nicer way of doing it
fn as_str<'a>(owned: &'a Option<String>) -> Option<&'a str> {
    match owned {
        &Some(ref s) => Some(s.as_slice()),
        &None => None
    }
}

#[deriving(PartialEq, Eq)]
enum QuoteState {
    InNone,
    InSingleQuotes,
    InDoubleQuotes
}

fn parse_process_args(s: &str) -> Vec<String> {
    let s = s.trim();
    let mut parts = Vec::new();
    let mut quote_state = InNone;
    let mut positions = vec!(0);
    let mut last = ' ';
    for (i, c) in s.chars().chain(" ".chars()).enumerate() {
        match (last, c) {
            // Match \" set has_escaped and skip
            ('\\', '\"') => (),
            // Match \'
            ('\\', '\'') => (),
            // Match \<space>
            ('\\', ' ') => (),
            // Match \\
            ('\\', '\\') => (),
            // Match <any>"
            (_, '\"') => {
                match quote_state {
                    InNone => {
                        quote_state = InDoubleQuotes;
                        positions.push(i);
                        positions.push(i + 1);
                    },
                    InDoubleQuotes => {
                        quote_state = InNone;
                        positions.push(i);
                        positions.push(i + 1);
                    }
                    InSingleQuotes => ()
                }
            },
            // Match <any>'
            (_, '\'') =>  {
                match quote_state {
                    InNone => {
                        quote_state = InSingleQuotes;
                        positions.push(i);
                        positions.push(i + 1);
                    },
                    InSingleQuotes => {
                        quote_state = InNone;
                        positions.push(i);
                        positions.push(i + 1);
                    }
                    InDoubleQuotes => ()
                }
            }
            // Match <any><space>
            (_, ' ') => {
                // If we are at the end of the string close any open quotes
                if i >= s.len() {
                    quote_state = InNone;
                }

                if quote_state == InNone {
                    {
                        positions.push(i);

                        let starts = positions.iter().enumerate().filter(|&(i, _)| i % 2 == 0);
                        let ends = positions.iter().enumerate().filter(|&(i, _)| i % 2 == 1);

                        let part: Vec<String> = starts.zip(ends).map(|((_, start), (_, end))| s.slice(*start, *end).to_string()).collect();
                        let part = part.connect("");

                        if part.len() > 0 {
                            // Remove any extra whitespace outside the quotes
                            let part = part.as_slice().trim();
                            // Replace quoted characters
                            let part = part.replace("\\\"", "\"");
                            let part = part.replace("\\\'", "\'");
                            let part = part.replace("\\ ", " ");
                            let part = part.replace("\\\\", "\\");
                            error!("part: {}", part);
                            parts.push(part);
                        }
                    }

                    positions.clear();
                    positions.push(i + 1);
                }
            },
            (_, _) => ()
        }
        last = c;
    }
    parts
}

struct MacroLogger<'a, 'b> {
    sp: codemap::Span,
    cx: &'a base::ExtCtxt<'b>
}

impl<'a, 'b> Logger for MacroLogger<'a, 'b> {
    fn error(&self, msg: &str) {
        self.cx.span_err(self.sp, msg)
    }

    fn warn(&self, msg: &str) {
        self.cx.span_warn(self.sp, msg)
    }
}

struct BindgenResult {
    items: RefCell<Option<SmallVector<@ast::Item>>>
}

impl base::MacResult for BindgenResult {
    fn make_items(&self) -> Option<SmallVector<@ast::Item>> {
        self.items.borrow_mut().take()
    }
}
<commit_msg>Tidy macro::parse_process_args and add tests<commit_after>use std::cell::RefCell;
use std::default::Default;
use std::os;
use syntax::ast;
use syntax::codemap;
use syntax::ext::base;
use syntax::parse;
use syntax::parse::token;
use syntax::util::small_vector::SmallVector;

use super::{generate_bindings, BindgenOptions, Logger};

pub fn bindgen_macro(cx: &mut base::ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<base::MacResult> {
    let mut visit = BindgenArgsVisitor {
        options: Default::default(),
        seen_named: false
    };

    visit.options.builtins = true;
    if !parse_macro_opts(cx, tts, &mut visit) {
        return base::DummyResult::any(sp);
    }

    // Reparse clang_args as it is passed in string form
    let clang_args = visit.options.clang_args.connect(" ");
    visit.options.clang_args = parse_process_args(clang_args.as_slice());

    // Set the working dir to the directory containing the invoking rs file so
    // that clang searches for headers relative to it rather than the crate root
    let mod_dir = Vec::from_slice(Path::new(cx.codemap().span_to_filename(sp)).dirname());
    let cwd = os::getcwd();
    os::change_dir(&Path::new(mod_dir));

    // We want the span for errors to just match the bindgen! symbol
    // instead of the whole invocation which can span multiple lines
    let mut short_span = sp;
    short_span.hi = short_span.lo + codemap::BytePos(8);

    let logger = MacroLogger { sp: short_span, cx: cx };

    let ret = match generate_bindings(visit.options, Some(&logger as &Logger), short_span) {
        Ok(items) => {
            box BindgenResult { items: RefCell::new(Some(SmallVector::many(items))) } as Box<base::MacResult>
        }
        Err(_) => base::DummyResult::any(sp)
    };

    os::change_dir(&Path::new(cwd));

    ret
}

trait MacroArgsVisitor {
    fn visit_str(&mut self, name: Option<&str>, val: &str) -> bool;
    fn visit_int(&mut self, name: Option<&str>, val: i64) -> bool;
    fn visit_bool(&mut self, name: Option<&str>, val: bool) -> bool;
    fn visit_ident(&mut self, name: Option<&str>, ident: &str) -> bool;
}

struct BindgenArgsVisitor {
    pub options: BindgenOptions,
    seen_named: bool
}

impl MacroArgsVisitor for BindgenArgsVisitor {
    fn visit_str(&mut self, mut name: Option<&str>, val: &str) -> bool {
        if name.is_some() { self.seen_named = true; }
        else if !self.seen_named { name = Some("clang_args") }
        match name {
            Some("link") => self.options.links.push((val.to_string(), None)),
            Some("link_static") => self.options.links.push((val.to_string(), Some("static".to_string()))),
            Some("link_framework") => self.options.links.push((val.to_string(), Some("framework".to_string()))),
            Some("abi") => self.options.abi = val.to_string(),
            Some("match") => self.options.match_pat.push(val.to_string()),
            Some("clang_args") => self.options.clang_args.push(val.to_string()),
            _ => return false
        }
        true
    }

    #[allow(unused_variable)]
    fn visit_int(&mut self, name: Option<&str>, val: i64) -> bool {
        if name.is_some() { self.seen_named = true; }
        false
    }

    fn visit_bool(&mut self, name: Option<&str>, val: bool) -> bool {
        if name.is_some() { self.seen_named = true; }
        match name {
            Some("allow_bitfields") => self.options.fail_on_bitfield = !val,
            Some("allow_unknown_types") => self.options.fail_on_unknown_type = !val,
            Some("emit_builtins") => self.options.builtins = val,
            _ => return false
        }
        true
    }

    #[allow(unused_variable)]
    fn visit_ident(&mut self, name: Option<&str>, val: &str) -> bool {
        if name.is_some() { self.seen_named = true; }
        false
    }
}

// Parses macro invocations in the form [ident=|:]value where value is an ident or literal
// e.g. bindgen!(module_name, "header.h", emit_builtins=false, clang_args:"-I /usr/local/include")
fn parse_macro_opts(cx: &mut base::ExtCtxt, tts: &[ast::TokenTree], visit: &mut MacroArgsVisitor) -> bool {
    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(), Vec::from_slice(tts));
    let mut args_good = true;

    loop {
        let mut name: Option<String> = None;
        let mut span = parser.span;

        // Check for [ident=]value and if found save ident to name
        if parser.look_ahead(1, |t| t == &token::EQ) {
            match parser.bump_and_get() {
                token::IDENT(ident, _) => {
                    let ident = parser.id_to_interned_str(ident);
                    name = Some(ident.get().to_string());
                    parser.expect(&token::EQ);
                },
                _ => {
                    cx.span_err(span, "invalid argument format");
                    return false
                }
            }
        }

        match parser.token {
            // Match [ident]
            token::IDENT(val, _) => {
                let val = parser.id_to_interned_str(val);
                span.hi = parser.span.hi;
                parser.bump();

                // Bools are simply encoded as idents
                let ret = match val.get() {
                    "true" => visit.visit_bool(as_str(&name), true),
                    "false" => visit.visit_bool(as_str(&name), false),
                    val => visit.visit_ident(as_str(&name), val)
                };
                if !ret {
                    cx.span_err(span, "invalid argument");
                    args_good = false;
                }
            }
            // Match [literal] and parse as an expression so we can expand macros
            _ => {
                let expr = cx.expand_expr(parser.parse_expr());
                span.hi = expr.span.hi;
                match expr.node {
                    ast::ExprLit(lit) => {
                        let ret = match lit.node {
                            ast::LitStr(ref s, _) => visit.visit_str(as_str(&name), s.get()),
                            ast::LitBool(b) => visit.visit_bool(as_str(&name), b),
                            ast::LitIntUnsuffixed(i) |
                            ast::LitInt(i, _) => visit.visit_int(as_str(&name), i),
                            ast::LitUint(i, _) => visit.visit_int(as_str(&name), i as i64),
                            _ => {
                                cx.span_err(span, "invalid argument format");
                                return false
                            }
                        };
                        if !ret {
                            cx.span_err(span, "invalid argument");
                            args_good = false;
                        }
                    },
                    _ => {
                        cx.span_err(span, "invalid argument format");
                        return false
                    }
                }
            }
        }

        if parser.eat(&token::EOF) {
            return args_good
        }

        if !parser.eat(&token::COMMA) {
            cx.span_err(parser.span, "invalid argument format");
            return false
        }
    }
}

// I'm sure there's a nicer way of doing it
fn as_str<'a>(owned: &'a Option<String>) -> Option<&'a str> {
    match owned {
        &Some(ref s) => Some(s.as_slice()),
        &None => None
    }
}

#[deriving(PartialEq, Eq)]
enum QuoteState {
    InNone,
    InSingleQuotes,
    InDoubleQuotes
}

fn parse_process_args(s: &str) -> Vec<String> {
    let s = s.trim();
    let mut parts = Vec::new();
    let mut quote_state = InNone;
    let mut positions = vec!(0);
    let mut last = ' ';
    for (i, c) in s.chars().chain(" ".chars()).enumerate() {
        match (last, c) {
            // Match \" set has_escaped and skip
            ('\\', '\"') => (),
            // Match \'
            ('\\', '\'') => (),
            // Match \<space>
            // Check we don't escape the final added space
            ('\\', ' ') if i < s.len() => (),
            // Match \\
            ('\\', '\\') => (),
            // Match <any>"
            (_, '\"') if quote_state == InNone => {
                quote_state = InDoubleQuotes;
                positions.push(i);
                positions.push(i + 1);
            },
            (_, '\"') if quote_state == InDoubleQuotes => {
                quote_state = InNone;
                positions.push(i);
                positions.push(i + 1);
            },
            // Match <any>'
            (_, '\'') if quote_state == InNone => {
                quote_state = InSingleQuotes;
                positions.push(i);
                positions.push(i + 1);
            },
            (_, '\'') if quote_state == InSingleQuotes => {
                quote_state = InNone;
                positions.push(i);
                positions.push(i + 1);
            },
            // Match <any><space>
            // If we are at the end of the string close any open quotes
            (_, ' ') if quote_state == InNone || i >= s.len() => {
                {
                    positions.push(i);

                    let starts = positions.iter().enumerate().filter(|&(i, _)| i % 2 == 0);
                    let ends = positions.iter().enumerate().filter(|&(i, _)| i % 2 == 1);

                    let part: Vec<String> = starts.zip(ends).map(|((_, start), (_, end))| s.slice(*start, *end).to_string()).collect();

                    let part = part.connect("");

                    if part.len() > 0 {
                        // Remove any extra whitespace outside the quotes
                        let part = part.as_slice().trim();
                        // Replace quoted characters
                        let part = part.replace("\\\"", "\"");
                        let part = part.replace("\\\'", "\'");
                        let part = part.replace("\\ ", " ");
                        let part = part.replace("\\\\", "\\");
                        parts.push(part);
                    }
                }

                positions.clear();
                positions.push(i + 1);
            },
            (_, _) => ()
        }
        last = c;
    }
    parts
}

struct MacroLogger<'a, 'b> {
    sp: codemap::Span,
    cx: &'a base::ExtCtxt<'b>
}

impl<'a, 'b> Logger for MacroLogger<'a, 'b> {
    fn error(&self, msg: &str) {
        self.cx.span_err(self.sp, msg)
    }

    fn warn(&self, msg: &str) {
        self.cx.span_warn(self.sp, msg)
    }
}

struct BindgenResult {
    items: RefCell<Option<SmallVector<@ast::Item>>>
}

impl base::MacResult for BindgenResult {
    fn make_items(&self) -> Option<SmallVector<@ast::Item>> {
        self.items.borrow_mut().take()
    }
}

#[cfg(test)]
fn make_string_vec(parts: &[&str]) -> Vec<String> {
    parts.iter().map(|p| p.to_string()).collect()
}

#[test]
fn test_parse_process_args() {
    assert_eq!(parse_process_args("a b c"), make_string_vec(["a", "b", "c"]));
    assert_eq!(parse_process_args("a \"b\" c"), make_string_vec(["a", "b", "c"]));
    assert_eq!(parse_process_args("a \'b\' c"), make_string_vec(["a", "b", "c"]));
    assert_eq!(parse_process_args("a \"b c\""), make_string_vec(["a", "b c"]));
    assert_eq!(parse_process_args("a \'\"b\"\' c"), make_string_vec(["a", "\"b\"", "c"]));
    assert_eq!(parse_process_args("a b\\ c"), make_string_vec(["a", "b c"]));
    assert_eq!(parse_process_args("a b c\\"), make_string_vec(["a", "b", "c\\"]));
}
<|endoftext|>


<commit_before>//! Copied from my historian crate. - Tyler Neely
//!
//! A zero-config simple histogram collector
//!
//! for use in instrumented optimization.
//! Uses logarithmic bucketing rather than sampling,
//! and has bounded (generally <0.5%) error on percentiles.
//! Performs no allocations after initial creation.
//! Uses Relaxed atomics during collection.
//!
//! When you create it, it allocates 65k `AtomicUsize`'s
//! that it uses for incrementing. Generating reports
//! after running workloads on dozens of `Histogram`'s
//! does not result in a perceptible delay, but it
//! might not be acceptable for use in low-latency
//! reporting paths.
//!
//! The trade-offs taken in this are to minimize latency
//! during collection, while initial allocation and
//! postprocessing delays are acceptable.
//!
//! Future work to further reduce collection latency
//! may include using thread-local caches that perform
//! no atomic operations until they are dropped, when
//! they may atomically aggregate their measurements
//! into the shared collector that will be used for
//! reporting.
#![allow(unused)]
#![allow(unused_results)]
#![allow(clippy::print_stdout)]
#![allow(clippy::float_arithmetic)]

use std::convert::TryFrom;
use std::fmt::{self, Debug};
use std::sync::atomic::{AtomicUsize, Ordering};

const PRECISION: f64 = 100.;
const BUCKETS: usize = 1 << 16;

/// A histogram collector that uses zero-configuration logarithmic buckets.
pub struct Histogram {
    vals: Vec<AtomicUsize>,
    sum: AtomicUsize,
    count: AtomicUsize,
}

impl Default for Histogram {
    #[allow(unsafe_code)]
    fn default() -> Histogram {
        #[cfg(not(feature = "miri_optimizations"))]
        {
            let mut vals = Vec::with_capacity(BUCKETS);
            vals.resize_with(BUCKETS, Default::default);

            Histogram { vals, sum: AtomicUsize::new(0), count: AtomicUsize::new(0) }
        }

        #[cfg(feature = "miri_optimizations")]
        {
            // Avoid calling Vec::resize_with with a large length because its
            // internals cause stacked borrows tracking information to add an
            // item for each element of the vector.
            let mut vals = std::mem::ManuallyDrop::new(vec![0usize; BUCKETS]);
            let ptr: *mut usize = vals.as_mut_ptr();
            let len = vals.len();
            let capacity = vals.capacity();

            let vals: Vec<AtomicUsize> = unsafe {
                Vec::from_raw_parts(ptr as *mut AtomicUsize, len, capacity)
            };

            Histogram { vals, sum: AtomicUsize::new(0), count: AtomicUsize::new(0) }
        }
    }
}

#[allow(unsafe_code)]
unsafe impl Send for Histogram {}

impl Debug for Histogram {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        const PS: [f64; 10] =
            [0., 50., 75., 90., 95., 97.5, 99., 99.9, 99.99, 100.];
        f.write_str("Histogramgram[")?;

        for p in &PS {
            let res = self.percentile(*p).round();
            let line = format!("({} -> {}) ", p, res);
            f.write_str(&*line)?;
        }

        f.write_str("]")
    }
}

impl Histogram {
    /// Record a value.
    #[inline]
    pub fn measure(&self, raw_value: u64) {
        #[cfg(not(feature = "no_metrics"))]
        {
            let value_float: f64 = raw_value as f64;
            self.sum.fetch_add(value_float.round() as usize, Ordering::Relaxed);

            self.count.fetch_add(1, Ordering::Relaxed);

            // compress the value to one of 2**16 values
            // using logarithmic bucketing
            let compressed: u16 = compress(value_float);

            // increment the counter for this compressed value
            self.vals[compressed as usize].fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Retrieve a percentile [0-100]. Returns NAN if no metrics have been
    /// collected yet.
    pub fn percentile(&self, p: f64) -> f64 {
        #[cfg(not(feature = "no_metrics"))]
        {
            assert!(p <= 100., "percentiles must not exceed 100.0");

            let count = self.count.load(Ordering::Acquire);

            if count == 0 {
                return std::f64::NAN;
            }

            let mut target = count as f64 * (p / 100.);
            if target == 0. {
                target = 1.;
            }

            let mut sum = 0.;

            for (idx, val) in self.vals.iter().enumerate() {
                let count = val.load(Ordering::Acquire);
                sum += count as f64;

                if sum >= target {
                    return decompress(idx as u16);
                }
            }
        }

        std::f64::NAN
    }

    /// Dump out some common percentiles.
    pub fn print_percentiles(&self) {
        println!("{:?}", self);
    }

    /// Return the sum of all observations in this histogram.
    pub fn sum(&self) -> usize {
        self.sum.load(Ordering::Acquire)
    }

    /// Return the count of observations in this histogram.
    pub fn count(&self) -> usize {
        self.count.load(Ordering::Acquire)
    }
}

// compress takes a value and lossily shrinks it to an u16 to facilitate
// bucketing of histogram values, staying roughly within 1% of the true
// value. This fails for large values of 1e142 and above, and is
// inaccurate for values closer to 0 than +/- 0.51 or +/- math.Inf.
#[allow(clippy::cast_sign_loss)]
#[allow(clippy::cast_possible_truncation)]
#[inline]
fn compress<T: Into<f64>>(input_value: T) -> u16 {
    let value: f64 = input_value.into();
    let abs = value.abs();
    let boosted = 1. + abs;
    let ln = boosted.ln();
    let compressed = PRECISION.mul_add(ln, 0.5);
    assert!(compressed <= f64::from(u16::max_value()));

    compressed as u16
}

// decompress takes a lossily shrunken u16 and returns an f64 within 1% of
// the original passed to compress.
#[inline]
fn decompress(compressed: u16) -> f64 {
    let unboosted = f64::from(compressed) / PRECISION;
    (unboosted.exp() - 1.)
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn it_works() {
    let c = Histogram::default();
    c.measure(2);
    c.measure(2);
    c.measure(3);
    c.measure(3);
    c.measure(4);
    assert_eq!(c.percentile(0.).round() as usize, 2);
    assert_eq!(c.percentile(40.).round() as usize, 2);
    assert_eq!(c.percentile(40.1).round() as usize, 3);
    assert_eq!(c.percentile(80.).round() as usize, 3);
    assert_eq!(c.percentile(80.1).round() as usize, 4);
    assert_eq!(c.percentile(100.).round() as usize, 4);
    c.print_percentiles();
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn high_percentiles() {
    let c = Histogram::default();
    for _ in 0..9000 {
        c.measure(10);
    }
    for _ in 0..900 {
        c.measure(25);
    }
    for _ in 0..90 {
        c.measure(33);
    }
    for _ in 0..9 {
        c.measure(47);
    }
    c.measure(500);
    assert_eq!(c.percentile(0.).round() as usize, 10);
    assert_eq!(c.percentile(99.).round() as usize, 25);
    assert_eq!(c.percentile(99.89).round() as usize, 33);
    assert_eq!(c.percentile(99.91).round() as usize, 47);
    assert_eq!(c.percentile(99.99).round() as usize, 47);
    assert_eq!(c.percentile(100.).round() as usize, 502);
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn multithreaded() {
    use std::sync::Arc;
    use std::thread;

    let h = Arc::new(Histogram::default());
    let mut threads = vec![];

    for _ in 0..10 {
        let h = h.clone();
        threads.push(thread::spawn(move || {
            h.measure(20);
        }));
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    assert_eq!(h.percentile(50.).round() as usize, 20);
}
<commit_msg>Clippy fix<commit_after>//! Copied from my historian crate. - Tyler Neely
//!
//! A zero-config simple histogram collector
//!
//! for use in instrumented optimization.
//! Uses logarithmic bucketing rather than sampling,
//! and has bounded (generally <0.5%) error on percentiles.
//! Performs no allocations after initial creation.
//! Uses Relaxed atomics during collection.
//!
//! When you create it, it allocates 65k `AtomicUsize`'s
//! that it uses for incrementing. Generating reports
//! after running workloads on dozens of `Histogram`'s
//! does not result in a perceptible delay, but it
//! might not be acceptable for use in low-latency
//! reporting paths.
//!
//! The trade-offs taken in this are to minimize latency
//! during collection, while initial allocation and
//! postprocessing delays are acceptable.
//!
//! Future work to further reduce collection latency
//! may include using thread-local caches that perform
//! no atomic operations until they are dropped, when
//! they may atomically aggregate their measurements
//! into the shared collector that will be used for
//! reporting.
#![allow(unused)]
#![allow(unused_results)]
#![allow(clippy::print_stdout)]
#![allow(clippy::float_arithmetic)]

use std::convert::TryFrom;
use std::fmt::{self, Debug};
use std::sync::atomic::{AtomicUsize, Ordering};

const PRECISION: f64 = 100.;
const BUCKETS: usize = 1 << 16;

/// A histogram collector that uses zero-configuration logarithmic buckets.
pub struct Histogram {
    vals: Vec<AtomicUsize>,
    sum: AtomicUsize,
    count: AtomicUsize,
}

impl Default for Histogram {
    #[allow(unsafe_code)]
    fn default() -> Histogram {
        #[cfg(not(feature = "miri_optimizations"))]
        {
            let mut vals = Vec::with_capacity(BUCKETS);
            vals.resize_with(BUCKETS, Default::default);

            Histogram { vals, sum: AtomicUsize::new(0), count: AtomicUsize::new(0) }
        }

        #[cfg(feature = "miri_optimizations")]
        {
            // Avoid calling Vec::resize_with with a large length because its
            // internals cause stacked borrows tracking information to add an
            // item for each element of the vector.
            let mut vals = std::mem::ManuallyDrop::new(vec![0_usize; BUCKETS]);
            let ptr: *mut usize = vals.as_mut_ptr();
            let len = vals.len();
            let capacity = vals.capacity();

            let vals: Vec<AtomicUsize> = unsafe {
                Vec::from_raw_parts(ptr as *mut AtomicUsize, len, capacity)
            };

            Histogram { vals, sum: AtomicUsize::new(0), count: AtomicUsize::new(0) }
        }
    }
}

#[allow(unsafe_code)]
unsafe impl Send for Histogram {}

impl Debug for Histogram {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        const PS: [f64; 10] =
            [0., 50., 75., 90., 95., 97.5, 99., 99.9, 99.99, 100.];
        f.write_str("Histogramgram[")?;

        for p in &PS {
            let res = self.percentile(*p).round();
            let line = format!("({} -> {}) ", p, res);
            f.write_str(&*line)?;
        }

        f.write_str("]")
    }
}

impl Histogram {
    /// Record a value.
    #[inline]
    pub fn measure(&self, raw_value: u64) {
        #[cfg(not(feature = "no_metrics"))]
        {
            let value_float: f64 = raw_value as f64;
            self.sum.fetch_add(value_float.round() as usize, Ordering::Relaxed);

            self.count.fetch_add(1, Ordering::Relaxed);

            // compress the value to one of 2**16 values
            // using logarithmic bucketing
            let compressed: u16 = compress(value_float);

            // increment the counter for this compressed value
            self.vals[compressed as usize].fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Retrieve a percentile [0-100]. Returns NAN if no metrics have been
    /// collected yet.
    pub fn percentile(&self, p: f64) -> f64 {
        #[cfg(not(feature = "no_metrics"))]
        {
            assert!(p <= 100., "percentiles must not exceed 100.0");

            let count = self.count.load(Ordering::Acquire);

            if count == 0 {
                return std::f64::NAN;
            }

            let mut target = count as f64 * (p / 100.);
            if target == 0. {
                target = 1.;
            }

            let mut sum = 0.;

            for (idx, val) in self.vals.iter().enumerate() {
                let count = val.load(Ordering::Acquire);
                sum += count as f64;

                if sum >= target {
                    return decompress(idx as u16);
                }
            }
        }

        std::f64::NAN
    }

    /// Dump out some common percentiles.
    pub fn print_percentiles(&self) {
        println!("{:?}", self);
    }

    /// Return the sum of all observations in this histogram.
    pub fn sum(&self) -> usize {
        self.sum.load(Ordering::Acquire)
    }

    /// Return the count of observations in this histogram.
    pub fn count(&self) -> usize {
        self.count.load(Ordering::Acquire)
    }
}

// compress takes a value and lossily shrinks it to an u16 to facilitate
// bucketing of histogram values, staying roughly within 1% of the true
// value. This fails for large values of 1e142 and above, and is
// inaccurate for values closer to 0 than +/- 0.51 or +/- math.Inf.
#[allow(clippy::cast_sign_loss)]
#[allow(clippy::cast_possible_truncation)]
#[inline]
fn compress<T: Into<f64>>(input_value: T) -> u16 {
    let value: f64 = input_value.into();
    let abs = value.abs();
    let boosted = 1. + abs;
    let ln = boosted.ln();
    let compressed = PRECISION.mul_add(ln, 0.5);
    assert!(compressed <= f64::from(u16::max_value()));

    compressed as u16
}

// decompress takes a lossily shrunken u16 and returns an f64 within 1% of
// the original passed to compress.
#[inline]
fn decompress(compressed: u16) -> f64 {
    let unboosted = f64::from(compressed) / PRECISION;
    (unboosted.exp() - 1.)
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn it_works() {
    let c = Histogram::default();
    c.measure(2);
    c.measure(2);
    c.measure(3);
    c.measure(3);
    c.measure(4);
    assert_eq!(c.percentile(0.).round() as usize, 2);
    assert_eq!(c.percentile(40.).round() as usize, 2);
    assert_eq!(c.percentile(40.1).round() as usize, 3);
    assert_eq!(c.percentile(80.).round() as usize, 3);
    assert_eq!(c.percentile(80.1).round() as usize, 4);
    assert_eq!(c.percentile(100.).round() as usize, 4);
    c.print_percentiles();
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn high_percentiles() {
    let c = Histogram::default();
    for _ in 0..9000 {
        c.measure(10);
    }
    for _ in 0..900 {
        c.measure(25);
    }
    for _ in 0..90 {
        c.measure(33);
    }
    for _ in 0..9 {
        c.measure(47);
    }
    c.measure(500);
    assert_eq!(c.percentile(0.).round() as usize, 10);
    assert_eq!(c.percentile(99.).round() as usize, 25);
    assert_eq!(c.percentile(99.89).round() as usize, 33);
    assert_eq!(c.percentile(99.91).round() as usize, 47);
    assert_eq!(c.percentile(99.99).round() as usize, 47);
    assert_eq!(c.percentile(100.).round() as usize, 502);
}

#[cfg(not(feature = "no_metrics"))]
#[test]
fn multithreaded() {
    use std::sync::Arc;
    use std::thread;

    let h = Arc::new(Histogram::default());
    let mut threads = vec![];

    for _ in 0..10 {
        let h = h.clone();
        threads.push(thread::spawn(move || {
            h.measure(20);
        }));
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    assert_eq!(h.percentile(50.).round() as usize, 20);
}
<|endoftext|>