<commit_before>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id`, `guild_id` and `webhook_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivalent in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivalent for the two
//! requests (`10`).
//!
//! Major parameters are why some variants (i.e. all of the channel/guild
//! variants) have an associated u64 as data. This is the Id of the parameter,
//! differentiating between different ratelimits.
//!
//! [Taken from]: https://discordapp.com/developers/docs/topics/rate-limits#rate-limits
pub use super::routing::Route;

use chrono::{DateTime, Utc};
use reqwest::{
    Response,
    header::HeaderMap as Headers,
    StatusCode,
};
use crate::internal::prelude::*;
use parking_lot::Mutex;
use std::{
    sync::Arc,
    time::Duration,
    str,
    thread,
    i64,
};
use super::{Http, HttpError, Request};
use log::debug;

/// Refer to [`offset`].
///
/// [`offset`]: fn.offset.html
static mut OFFSET: Option<i64> = None;

pub(super) fn perform(http: &Http, req: Request<'_>) -> Result<Response> {
    loop {
        // This will block if another thread is trying to send
        // an HTTP-request already (due to receiving an x-ratelimit-global).
        let _ = http.limiter.lock();

        // Destructure the tuple instead of retrieving the third value to
        // take advantage of the type system. If `RouteInfo::deconstruct`
        // returns a different number of tuple elements in the future, directly
        // accessing a certain index (e.g. `req.route.deconstruct().1`) would
        // mean this code would not indicate it might need to be updated for the
        // new tuple element amount.
        //
        // This isn't normally important, but might be for ratelimiting.
        let (_, route, _) = req.route.deconstruct();

        // Perform pre-checking here:
        //
        // - get the route's relevant rate
        // - sleep if that route's already rate-limited until the end of the
        //   'reset' time;
        // - get the global rate;
        // - sleep if there is 0 remaining
        // - then, perform the request
        let bucket = Arc::clone(http.routes
            .lock()
            .entry(route)
            .or_insert_with(|| {
                Arc::new(Mutex::new(RateLimit {
                    limit: i64::MAX,
                    remaining: i64::MAX,
                    reset: i64::MAX,
                }))
            }));

        let mut lock = bucket.lock();
        lock.pre_hook(&route);

        let response = http.retry(&req)?;

        // Check if an offset has been calculated yet to determine the time
        // difference from Discord can the client.
        //
        // Refer to the documentation for `OFFSET` for more information.
        //
        // This should probably only be a one-time check, although we may want
        // to choose to check this often in the future.
        if unsafe { OFFSET }.is_none() {
            calculate_offset(&response.headers().get("date").and_then(|d| Some(d.as_bytes())));
        }

        // Check if the request got ratelimited by checking for status 429,
        // and if so, sleep for the value of the header 'retry-after' -
        // which is in milliseconds - and then `continue` to try again
        //
        // If it didn't ratelimit, subtract one from the RateLimit's
        // 'remaining'
        //
        // Update the 'reset' with the value of the 'x-ratelimit-reset'
        // header
        //
        // It _may_ be possible for the limit to be raised at any time,
        // so check if it did from the value of the 'x-ratelimit-limit'
        // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
        if route == Route::None {
            return Ok(response);
        } else {
            let redo = if response.headers().get("x-ratelimit-global").is_some() {
                let _ = http.limiter.lock();

                Ok(
                    if let Some(retry_after) = parse_header(&response.headers(), "retry-after")? {
                        debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
                        thread::sleep(Duration::from_millis(retry_after as u64));

                        true
                    } else {
                        false
                    },
                )
            } else {
                lock.post_hook(&response, &route)
            };

            if !redo.unwrap_or(true) {
                return Ok(response);
            }
        }
    }
}

/// A set of data containing information about the ratelimits for a particular
/// [`Route`], which is stored in [`Http`].
///
/// See the [Discord docs] on ratelimits for more information.
///
/// **Note**: You should _not_ mutate any of the fields, as this can help cause
/// 429s.
///
/// [`Http`]: ../raw/struct.Http.html#structfield.routes
/// [`Route`]: ../routing/enum.Route.html
/// [Discord docs]: https://discordapp.com/developers/docs/topics/rate-limits
#[derive(Clone, Debug, Default)]
pub struct RateLimit {
    /// The total number of requests that can be made in a period of time.
    pub limit: i64,
    /// The number of requests remaining in the period of time.
    pub remaining: i64,
    /// When the interval resets and the the [`limit`] resets to the value of
    /// [`remaining`].
    ///
    /// [`limit`]: #structfield.limit
    /// [`remaining`]: #structfield.remaining
    pub reset: i64,
}

impl RateLimit {
    pub(crate) fn pre_hook(&mut self, route: &Route) {
        if self.limit == 0 {
            return;
        }

        let offset = unsafe { OFFSET }.unwrap_or(0);
        let now = Utc::now().timestamp();
        let current_time = now - offset;

        // The reset was in the past, so we're probably good.
        if current_time > self.reset {
            self.remaining = self.limit;

            return;
        }

        let diff = (self.reset - current_time) as u64;

        if self.remaining == 0 {
            let delay = diff * 1000;

            debug!(
                "Pre-emptive ratelimit on route {:?} for {:?}ms",
                route,
                delay
            );

            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub(crate) fn post_hook(&mut self, response: &Response, route: &Route) -> Result<bool> {
        if let Some(limit) = parse_header(&response.headers(), "x-ratelimit-limit")? {
            self.limit = limit;
        }

        if let Some(remaining) = parse_header(&response.headers(), "x-ratelimit-remaining")? {
            self.remaining = remaining;
        }

        if let Some(reset) = parse_header(&response.headers(), "x-ratelimit-reset")? {
            self.reset = reset;
        }

        Ok(if response.status() != StatusCode::TOO_MANY_REQUESTS {
            false
        } else if let Some(retry_after) = parse_header(&response.headers(), "retry-after")? {
            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
            thread::sleep(Duration::from_millis(retry_after as u64));

            true
        } else {
            false
        })
    }
}

/// The calculated offset of the time difference between Discord and the client
/// in seconds.
///
/// This does not have millisecond precision as calculating that isn't
/// realistic.
///
/// This is used in ratelimiting to help determine how long to wait for
/// pre-emptive ratelimits. For example, if the client is 2 seconds ahead, then
/// the client would think the ratelimit is over 2 seconds before it actually is
/// and would then send off queued requests. Using an offset, we can know that
/// there's actually still 2 seconds left (+/- some milliseconds).
///
/// This isn't a definitive solution to fix all problems, but it can help with
/// some precision gains.
///
/// This will return `None` if an HTTP request hasn't been made, meaning that
/// no offset could have been calculated.
pub fn offset() -> Option<i64> {
    unsafe { OFFSET }
}

fn calculate_offset(header: &Option<&[u8]>) {
    // Get the current time as soon as possible.
    let now = Utc::now().timestamp();

    let header = header.and_then(|x| str::from_utf8(x).ok());

    if let Some(header) = header {
        // Replace the `GMT` timezone with an offset, and then parse it
        // into a chrono DateTime. If it parses correctly, calculate the
        // diff and then set it as the offset.
        let s = header.replace("GMT", "+0000");

        let parsed = DateTime::parse_from_str(&s, "%a, %d %b %Y %T %z");

        if let Ok(parsed) = parsed {
            let offset = parsed.timestamp();

            let diff = offset - now;

            unsafe {
                OFFSET = Some(diff);

                debug!("[ratelimiting] Set the ratelimit offset to {}", diff);
            }
        }
    }

}

fn parse_header(headers: &Headers, header: &str) -> Result<Option<i64>> {
    headers.get(header).map_or(Ok(None), |header| {
        str::from_utf8(&header.as_bytes())
            .map_err(|_| Error::from(HttpError::RateLimitUtf8))
            .and_then(|v| {
                v.parse::<i64>()
                    .map(Some)
                    .map_err(|_| Error::from(HttpError::RateLimitI64))
            })
    })
}
<commit_msg>ratelimiting: simplify parse_headers, add tests<commit_after>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id`, `guild_id` and `webhook_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivalent in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivalent for the two
//! requests (`10`).
//!
//! Major parameters are why some variants (i.e. all of the channel/guild
//! variants) have an associated u64 as data. This is the Id of the parameter,
//! differentiating between different ratelimits.
//!
//! [Taken from]: https://discordapp.com/developers/docs/topics/rate-limits#rate-limits
pub use super::routing::Route;

use chrono::{DateTime, Utc};
use reqwest::{
    Response,
    header::HeaderMap,
    StatusCode,
};
use crate::internal::prelude::*;
use parking_lot::Mutex;
use std::{
    sync::Arc,
    time::Duration,
    str,
    thread,
    i64,
};
use super::{Http, HttpError, Request};
use log::debug;

/// Refer to [`offset`].
///
/// [`offset`]: fn.offset.html
static mut OFFSET: Option<i64> = None;

pub(super) fn perform(http: &Http, req: Request<'_>) -> Result<Response> {
    loop {
        // This will block if another thread is trying to send
        // an HTTP-request already (due to receiving an x-ratelimit-global).
        let _ = http.limiter.lock();

        // Destructure the tuple instead of retrieving the third value to
        // take advantage of the type system. If `RouteInfo::deconstruct`
        // returns a different number of tuple elements in the future, directly
        // accessing a certain index (e.g. `req.route.deconstruct().1`) would
        // mean this code would not indicate it might need to be updated for the
        // new tuple element amount.
        //
        // This isn't normally important, but might be for ratelimiting.
        let (_, route, _) = req.route.deconstruct();

        // Perform pre-checking here:
        //
        // - get the route's relevant rate
        // - sleep if that route's already rate-limited until the end of the
        //   'reset' time;
        // - get the global rate;
        // - sleep if there is 0 remaining
        // - then, perform the request
        let bucket = Arc::clone(http.routes
            .lock()
            .entry(route)
            .or_insert_with(|| {
                Arc::new(Mutex::new(RateLimit {
                    limit: i64::MAX,
                    remaining: i64::MAX,
                    reset: i64::MAX,
                }))
            }));

        let mut lock = bucket.lock();
        lock.pre_hook(&route);

        let response = http.retry(&req)?;

        // Check if an offset has been calculated yet to determine the time
        // difference from Discord can the client.
        //
        // Refer to the documentation for `OFFSET` for more information.
        //
        // This should probably only be a one-time check, although we may want
        // to choose to check this often in the future.
        if unsafe { OFFSET }.is_none() {
            calculate_offset(&response.headers().get("date").and_then(|d| Some(d.as_bytes())));
        }

        // Check if the request got ratelimited by checking for status 429,
        // and if so, sleep for the value of the header 'retry-after' -
        // which is in milliseconds - and then `continue` to try again
        //
        // If it didn't ratelimit, subtract one from the RateLimit's
        // 'remaining'
        //
        // Update the 'reset' with the value of the 'x-ratelimit-reset'
        // header
        //
        // It _may_ be possible for the limit to be raised at any time,
        // so check if it did from the value of the 'x-ratelimit-limit'
        // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
        if route == Route::None {
            return Ok(response);
        } else {
            let redo = if response.headers().get("x-ratelimit-global").is_some() {
                let _ = http.limiter.lock();

                Ok(
                    if let Some(retry_after) = parse_header(&response.headers(), "retry-after")? {
                        debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
                        thread::sleep(Duration::from_millis(retry_after as u64));

                        true
                    } else {
                        false
                    },
                )
            } else {
                lock.post_hook(&response, &route)
            };

            if !redo.unwrap_or(true) {
                return Ok(response);
            }
        }
    }
}

/// A set of data containing information about the ratelimits for a particular
/// [`Route`], which is stored in [`Http`].
///
/// See the [Discord docs] on ratelimits for more information.
///
/// **Note**: You should _not_ mutate any of the fields, as this can help cause
/// 429s.
///
/// [`Http`]: ../raw/struct.Http.html#structfield.routes
/// [`Route`]: ../routing/enum.Route.html
/// [Discord docs]: https://discordapp.com/developers/docs/topics/rate-limits
#[derive(Clone, Debug, Default)]
pub struct RateLimit {
    /// The total number of requests that can be made in a period of time.
    pub limit: i64,
    /// The number of requests remaining in the period of time.
    pub remaining: i64,
    /// When the interval resets and the the [`limit`] resets to the value of
    /// [`remaining`].
    ///
    /// [`limit`]: #structfield.limit
    /// [`remaining`]: #structfield.remaining
    pub reset: i64,
}

impl RateLimit {
    pub(crate) fn pre_hook(&mut self, route: &Route) {
        if self.limit == 0 {
            return;
        }

        let offset = unsafe { OFFSET }.unwrap_or(0);
        let now = Utc::now().timestamp();
        let current_time = now - offset;

        // The reset was in the past, so we're probably good.
        if current_time > self.reset {
            self.remaining = self.limit;

            return;
        }

        let diff = (self.reset - current_time) as u64;

        if self.remaining == 0 {
            let delay = diff * 1000;

            debug!(
                "Pre-emptive ratelimit on route {:?} for {:?}ms",
                route,
                delay
            );

            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub(crate) fn post_hook(&mut self, response: &Response, route: &Route) -> Result<bool> {
        if let Some(limit) = parse_header(&response.headers(), "x-ratelimit-limit")? {
            self.limit = limit;
        }

        if let Some(remaining) = parse_header(&response.headers(), "x-ratelimit-remaining")? {
            self.remaining = remaining;
        }

        if let Some(reset) = parse_header(&response.headers(), "x-ratelimit-reset")? {
            self.reset = reset;
        }

        Ok(if response.status() != StatusCode::TOO_MANY_REQUESTS {
            false
        } else if let Some(retry_after) = parse_header(&response.headers(), "retry-after")? {
            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
            thread::sleep(Duration::from_millis(retry_after as u64));

            true
        } else {
            false
        })
    }
}

/// The calculated offset of the time difference between Discord and the client
/// in seconds.
///
/// This does not have millisecond precision as calculating that isn't
/// realistic.
///
/// This is used in ratelimiting to help determine how long to wait for
/// pre-emptive ratelimits. For example, if the client is 2 seconds ahead, then
/// the client would think the ratelimit is over 2 seconds before it actually is
/// and would then send off queued requests. Using an offset, we can know that
/// there's actually still 2 seconds left (+/- some milliseconds).
///
/// This isn't a definitive solution to fix all problems, but it can help with
/// some precision gains.
///
/// This will return `None` if an HTTP request hasn't been made, meaning that
/// no offset could have been calculated.
pub fn offset() -> Option<i64> {
    unsafe { OFFSET }
}

fn calculate_offset(header: &Option<&[u8]>) {
    // Get the current time as soon as possible.
    let now = Utc::now().timestamp();

    let header = header.and_then(|x| str::from_utf8(x).ok());

    if let Some(header) = header {
        // Replace the `GMT` timezone with an offset, and then parse it
        // into a chrono DateTime. If it parses correctly, calculate the
        // diff and then set it as the offset.
        let s = header.replace("GMT", "+0000");

        let parsed = DateTime::parse_from_str(&s, "%a, %d %b %Y %T %z");

        if let Ok(parsed) = parsed {
            let offset = parsed.timestamp();

            let diff = offset - now;

            unsafe {
                OFFSET = Some(diff);

                debug!("[ratelimiting] Set the ratelimit offset to {}", diff);
            }
        }
    }

}

fn parse_header(headers: &HeaderMap, header: &str) -> Result<Option<i64>> {
    let header = match headers.get(header) {
        Some(v) => v,
        None => return Ok(None),
    };

    let unicode = str::from_utf8(&header.as_bytes()).map_err(|_| {
        Error::from(HttpError::RateLimitUtf8)
    })?;

    let num = unicode.parse().map_err(|_| {
        Error::from(HttpError::RateLimitI64)
    })?;

    Ok(Some(num))
}

#[cfg(test)]
mod tests {
    use crate::{
        error::Error,
        http::HttpError,
    };
    use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
    use std::{
        error::Error as StdError,
        result::Result as StdResult,
    };
    use super::parse_header;

    type Result<T> = StdResult<T, Box<dyn StdError>>;

    fn headers() -> HeaderMap {
        let pairs = &[
            (
                HeaderName::from_static("x-ratelimit-limit"),
                HeaderValue::from_static("5"),
            ),
            (
                HeaderName::from_static("x-ratelimit-remaining"),
                HeaderValue::from_static("4"),
            ),
            (
                HeaderName::from_static("x-ratelimit-reset"),
                HeaderValue::from_static("1560704880"),
            ),
            (
                HeaderName::from_static("x-bad-num"),
                HeaderValue::from_static("abc"),
            ),
            (
                HeaderName::from_static("x-bad-unicode"),
                HeaderValue::from_bytes(&[255, 255, 255, 255]).unwrap(),
            ),
        ];

        let mut map = HeaderMap::with_capacity(pairs.len());

        for (name, val) in pairs.into_iter() {
            map.insert(name, val.to_owned());
        }

        map
    }

    #[test]
    fn test_parse_header_good() -> Result<()> {
        let headers = headers();

        assert_eq!(parse_header(&headers, "x-ratelimit-limit")?.unwrap(), 5);
        assert_eq!(
            parse_header(&headers, "x-ratelimit-remaining")?.unwrap(),
            4,
        );
        assert_eq!(
            parse_header(&headers, "x-ratelimit-reset")?.unwrap(),
            1_560_704_880,
        );

        Ok(())
    }

    #[test]
    fn test_parse_header_errors() -> Result<()> {
        let headers = headers();

        match parse_header(&headers, "x-bad-num").unwrap_err() {
            Error::Http(x) => match *x {
                HttpError::RateLimitI64 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        match parse_header(&headers, "x-bad-unicode").unwrap_err() {
            Error::Http(http_err) => match *http_err {
                HttpError::RateLimitUtf8 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>/*
 * Implemenation of a bitfield as a vec. Only
 * supports bytes (Vec<u8>) as the underlying
 * storage.
 *
 * A future implementation should be more efficient,
 * this is just to get the job done for now.
 */
extern crate ssz;

use std::cmp::max;

#[derive(Eq, Clone, Default, Debug)]
pub struct BooleanBitfield{
    len: usize,
    vec: Vec<u8>
}

impl BooleanBitfield {
    /// Create a new bitfield with a length of zero.
    pub fn new() -> Self {
        Self {
            len: 0,
            vec: vec![0]
        }
    }

    /// Create a new bitfield of a certain capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let mut vec = Vec::with_capacity(capacity / 8 + 1);
        vec.push(0);
        Self {
            len: 0,
            vec
        }
    }

    /// Read the value of a bit.
    ///
    /// Will return `true` if the bit has been set to `true`
    /// without then being set to `False`.
    pub fn get_bit(&self, i: usize) -> bool {
        let bit = |i: usize| i % 8;
        let byte = |i: usize| i / 8;

        if byte(i) >= self.vec.len() {
            false
        } else {
             self.vec[byte(i)] & (1 << (bit(i) as u8)) != 0
        }
    }

    /// Set the value of a bit.
    ///
    /// If this bit is larger than the length of the underlying byte
    /// array it will be extended.
    pub fn set_bit(&mut self, i: usize, to: bool) {
        let bit = |i: usize| i % 8;
        let byte = |i: usize| i / 8;

        self.len = max(self.len, i + 1);

        if byte(i) >= self.vec.len() {
            self.vec.resize(byte(i) + 1, 0);
        }
        if to {
            self.vec[byte(i)] =
                self.vec[byte(i)] |  (1 << (bit(i) as u8))
        } else {
            self.vec[byte(i)] =
                self.vec[byte(i)] & !(1 << (bit(i) as u8))
        }
    }

    /// Return the "length" of this bitfield. Length is defined as
    /// the highest bit that has been set.
    ///
    /// Note: this is distinct from the length of the underlying
    /// vector.
    pub fn len(&self) -> usize { self.len }

    /// True if no bits have ever been set. A bit that is set and then
    /// unset will still count to the length of the bitfield.
    ///
    /// Note: this is distinct from the length of the underlying
    /// vector.
    pub fn is_empty(&self) -> bool { self.len == 0 }

    /// The number of bytes required to represent the bitfield.
    pub fn num_bytes(&self) -> usize { self.vec.len() }

    /// Iterate through the underlying vector and count the number of
    /// true bits.
    pub fn num_true_bits(&self) -> u64 {
        let mut count: u64 = 0;
        for byte in &self.vec {
            for bit in 0..8 {
                if byte & (1 << (bit as u8)) != 0 {
                    count += 1;
                }
            }
        }
        count
    }

    /// Iterate through the underlying vector and find the highest
    /// set bit. Useful for instantiating a new instance from
    /// some set of bytes.
    pub fn compute_length(bytes: &[u8]) -> usize {
        for byte in (0..bytes.len()).rev() {
            for bit in (0..8).rev() {
                if byte & (1 << (bit as u8)) != 0 {
                    return (byte * 8) + bit
                }
            }
        }
        0
    }

    /// Clone and return the underlying byte array (`Vec<u8>`).
    pub fn to_be_vec(&self) -> Vec<u8> {
        let mut o = self.vec.clone();
        o.reverse();
        o
    }
}

impl<'a> From<&'a [u8]> for BooleanBitfield {
    fn from(input: &[u8]) -> Self {
        let mut vec = input.to_vec();
        vec.reverse();
        BooleanBitfield {
            vec,
            len: BooleanBitfield::compute_length(input)
        }
    }
}

impl PartialEq for BooleanBitfield {
    fn eq(&self, other: &BooleanBitfield) -> bool {
        (self.vec == other.vec) &
            (self.len == other.len)
    }
}

impl ssz::Encodable for BooleanBitfield {
    fn ssz_append(&self, s: &mut ssz::SszStream) {
        s.append_vec(&self.to_be_vec());
    }
}

impl ssz::Decodable for BooleanBitfield {
    fn ssz_decode(bytes: &[u8], index: usize)
        -> Result<(Self, usize), ssz::DecodeError>
    {
        let len = ssz::decode::decode_length(
            bytes,
            index,
            ssz::LENGTH_BYTES)?;
        if (ssz::LENGTH_BYTES + len) > bytes.len() {
            return Err(ssz::DecodeError::TooShort);
        }
        if len == 0 {
            Ok((BooleanBitfield::new(),
                index + ssz::LENGTH_BYTES))
        } else {
            let b = BooleanBitfield::
                from(&bytes[(index + 4)..(index + len + 4)]);
            let index = index + ssz::LENGTH_BYTES + len;
            Ok((b, index))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::Decodable;

    #[test]
    fn test_ssz_encoding() {
        let mut b = BooleanBitfield::new();
        b.set_bit(8, true);

        let mut stream = ssz::SszStream::new();
        stream.append(&b);

        assert_eq!(stream.drain(), vec![0, 0, 0, 2, 1, 0]);
    }

    #[test]
    fn test_ssz_decoding() {
        /*
         * Correct input
         */
        let input = vec![0, 0, 0, 2, 1, 0];
        let (b, i) = BooleanBitfield::ssz_decode(&input, 0).unwrap();
        assert_eq!(i, 6);
        assert_eq!(b.num_true_bits(), 1);
        assert_eq!(b.get_bit(8), true);

        /*
         * Input too long
         */
        let mut input = vec![0, 0, 0, 2, 1, 0];
        input.push(42);
        let (b, i) = BooleanBitfield::ssz_decode(&input, 0).unwrap();
        assert_eq!(i, 6);
        assert_eq!(b.num_true_bits(), 1);
        assert_eq!(b.get_bit(8), true);

        /*
         * Input too short
         */
        let input = vec![0, 0, 0, 2, 1];
        let res = BooleanBitfield::ssz_decode(&input, 0);
        assert_eq!(res, Err(ssz::DecodeError::TooShort));
    }

    #[test]
    fn test_new_bitfield_len() {
        let b = BooleanBitfield::new();
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);

        let b = BooleanBitfield::with_capacity(100);
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);
    }

    #[test]
    fn test_new_bitfield_len() {
        let b = BooleanBitfield::new();
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);

        let b = BooleanBitfield::with_capacity(100);
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);
    }

    #[test]
    fn test_bitfield_set() {
        let mut b = BooleanBitfield::new();
        b.set_bit(0, false);
        assert_eq!(b.to_be_vec(), [0]);

        b = BooleanBitfield::new();
        b.set_bit(7, true);
        assert_eq!(b.to_be_vec(), [128]);
        b.set_bit(7, false);
        assert_eq!(b.to_be_vec(), [0]);
        assert_eq!(b.len(), 8);

        b = BooleanBitfield::new();
        b.set_bit(7, true);
        b.set_bit(0, true);
        assert_eq!(b.to_be_vec(), [129]);
        b.set_bit(7, false);
        assert_eq!(b.to_be_vec(), [1]);
        assert_eq!(b.len(), 8);

        b = BooleanBitfield::new();
        b.set_bit(8, true);
        assert_eq!(b.to_be_vec(), [1, 0]);
        assert_eq!(b.len(), 9);
        b.set_bit(8, false);
        assert_eq!(b.to_be_vec(), [0, 0]);
        assert_eq!(b.len(), 9);

        b = BooleanBitfield::new();
        b.set_bit(15, true);
        assert_eq!(b.to_be_vec(), [128, 0]);
        b.set_bit(15, false);
        assert_eq!(b.to_be_vec(), [0, 0]);
        assert_eq!(b.len(), 16);

        b = BooleanBitfield::new();
        b.set_bit(8, true);
        b.set_bit(15, true);
        assert_eq!(b.to_be_vec(), [129, 0]);
        b.set_bit(15, false);
        assert_eq!(b.to_be_vec(), [1, 0]);
        assert_eq!(b.len(), 16);
    }

    #[test]
    fn test_bitfield_get() {
        let test_nums = vec![0, 8, 15, 42, 1337];
        for i in test_nums {
            let mut b = BooleanBitfield::new();
            assert_eq!(b.get_bit(i), false);
            b.set_bit(i, true);
            assert_eq!(b.get_bit(i), true);
            b.set_bit(i, true);
        }
    }

    #[test]
    fn test_bitfield_num_true_bits() {
        let mut b = BooleanBitfield::new();
        assert_eq!(b.num_true_bits(), 0);
        b.set_bit(15, true);
        assert_eq!(b.num_true_bits(), 1);
        b.set_bit(15, false);
        assert_eq!(b.num_true_bits(), 0);
        b.set_bit(0, true);
        b.set_bit(7, true);
        b.set_bit(8, true);
        b.set_bit(1337, true);
        assert_eq!(b.num_true_bits(), 4);
    }
}
<commit_msg>Add method to bitfield to read specific byte<commit_after>/*
 * Implemenation of a bitfield as a vec. Only
 * supports bytes (Vec<u8>) as the underlying
 * storage.
 *
 * A future implementation should be more efficient,
 * this is just to get the job done for now.
 */
extern crate ssz;

use std::cmp::max;

#[derive(Eq, Clone, Default, Debug)]
pub struct BooleanBitfield{
    len: usize,
    vec: Vec<u8>
}

impl BooleanBitfield {
    /// Create a new bitfield with a length of zero.
    pub fn new() -> Self {
        Self {
            len: 0,
            vec: vec![0]
        }
    }

    /// Create a new bitfield of a certain capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let mut vec = Vec::with_capacity(capacity / 8 + 1);
        vec.push(0);
        Self {
            len: 0,
            vec
        }
    }

    /// Read the value of a bit.
    ///
    /// Will return `true` if the bit has been set to `true`
    /// without then being set to `False`.
    pub fn get_bit(&self, i: usize) -> bool {
        let bit = |i: usize| i % 8;
        let byte = |i: usize| i / 8;

        if byte(i) >= self.vec.len() {
            false
        } else {
             self.vec[byte(i)] & (1 << (bit(i) as u8)) != 0
        }
    }

    /// Set the value of a bit.
    ///
    /// If this bit is larger than the length of the underlying byte
    /// array it will be extended.
    pub fn set_bit(&mut self, i: usize, to: bool) {
        let bit = |i: usize| i % 8;
        let byte = |i: usize| i / 8;

        self.len = max(self.len, i + 1);

        if byte(i) >= self.vec.len() {
            self.vec.resize(byte(i) + 1, 0);
        }
        if to {
            self.vec[byte(i)] =
                self.vec[byte(i)] |  (1 << (bit(i) as u8))
        } else {
            self.vec[byte(i)] =
                self.vec[byte(i)] & !(1 << (bit(i) as u8))
        }
    }

    /// Return the "length" of this bitfield. Length is defined as
    /// the highest bit that has been set.
    ///
    /// Note: this is distinct from the length of the underlying
    /// vector.
    pub fn len(&self) -> usize { self.len }

    /// True if no bits have ever been set. A bit that is set and then
    /// unset will still count to the length of the bitfield.
    ///
    /// Note: this is distinct from the length of the underlying
    /// vector.
    pub fn is_empty(&self) -> bool { self.len == 0 }

    /// The number of bytes required to represent the bitfield.
    pub fn num_bytes(&self) -> usize { self.vec.len() }

    /// Iterate through the underlying vector and count the number of
    /// true bits.
    pub fn num_true_bits(&self) -> u64 {
        let mut count: u64 = 0;
        for byte in &self.vec {
            for bit in 0..8 {
                if byte & (1 << (bit as u8)) != 0 {
                    count += 1;
                }
            }
        }
        count
    }

    /// Iterate through the underlying vector and find the highest
    /// set bit. Useful for instantiating a new instance from
    /// some set of bytes.
    pub fn compute_length(bytes: &[u8]) -> usize {
        for byte in (0..bytes.len()).rev() {
            for bit in (0..8).rev() {
                if byte & (1 << (bit as u8)) != 0 {
                    return (byte * 8) + bit
                }
            }
        }
        0
    }

    /// Get the byte at a position, assuming big-endian encoding.
    pub fn get_byte(&self, n: usize) -> Option<&u8> {
        self.vec.get(n)
    }

    /// Clone and return the underlying byte array (`Vec<u8>`).
    pub fn to_be_vec(&self) -> Vec<u8> {
        let mut o = self.vec.clone();
        o.reverse();
        o
    }
}

impl<'a> From<&'a [u8]> for BooleanBitfield {
    fn from(input: &[u8]) -> Self {
        let mut vec = input.to_vec();
        vec.reverse();
        BooleanBitfield {
            vec,
            len: BooleanBitfield::compute_length(input)
        }
    }
}

impl PartialEq for BooleanBitfield {
    fn eq(&self, other: &BooleanBitfield) -> bool {
        (self.vec == other.vec) &
            (self.len == other.len)
    }
}

impl ssz::Encodable for BooleanBitfield {
    fn ssz_append(&self, s: &mut ssz::SszStream) {
        s.append_vec(&self.to_be_vec());
    }
}

impl ssz::Decodable for BooleanBitfield {
    fn ssz_decode(bytes: &[u8], index: usize)
        -> Result<(Self, usize), ssz::DecodeError>
    {
        let len = ssz::decode::decode_length(
            bytes,
            index,
            ssz::LENGTH_BYTES)?;
        if (ssz::LENGTH_BYTES + len) > bytes.len() {
            return Err(ssz::DecodeError::TooShort);
        }
        if len == 0 {
            Ok((BooleanBitfield::new(),
                index + ssz::LENGTH_BYTES))
        } else {
            let b = BooleanBitfield::
                from(&bytes[(index + 4)..(index + len + 4)]);
            let index = index + ssz::LENGTH_BYTES + len;
            Ok((b, index))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ssz::Decodable;

    #[test]
    fn test_ssz_encoding() {
        let mut b = BooleanBitfield::new();
        b.set_bit(8, true);

        let mut stream = ssz::SszStream::new();
        stream.append(&b);

        assert_eq!(stream.drain(), vec![0, 0, 0, 2, 1, 0]);
    }

    #[test]
    fn test_ssz_decoding() {
        /*
         * Correct input
         */
        let input = vec![0, 0, 0, 2, 1, 0];
        let (b, i) = BooleanBitfield::ssz_decode(&input, 0).unwrap();
        assert_eq!(i, 6);
        assert_eq!(b.num_true_bits(), 1);
        assert_eq!(b.get_bit(8), true);

        /*
         * Input too long
         */
        let mut input = vec![0, 0, 0, 2, 1, 0];
        input.push(42);
        let (b, i) = BooleanBitfield::ssz_decode(&input, 0).unwrap();
        assert_eq!(i, 6);
        assert_eq!(b.num_true_bits(), 1);
        assert_eq!(b.get_bit(8), true);

        /*
         * Input too short
         */
        let input = vec![0, 0, 0, 2, 1];
        let res = BooleanBitfield::ssz_decode(&input, 0);
        assert_eq!(res, Err(ssz::DecodeError::TooShort));
    }

    #[test]
    fn test_new_bitfield_len() {
        let b = BooleanBitfield::new();
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);

        let b = BooleanBitfield::with_capacity(100);
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);
    }

    #[test]
    fn test_new_bitfield_len() {
        let b = BooleanBitfield::new();
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);

        let b = BooleanBitfield::with_capacity(100);
        assert_eq!(b.len(), 0);
        assert_eq!(b.to_be_vec(), vec![0]);
    }

    #[test]
    fn test_bitfield_set() {
        let mut b = BooleanBitfield::new();
        b.set_bit(0, false);
        assert_eq!(b.to_be_vec(), [0]);

        b = BooleanBitfield::new();
        b.set_bit(7, true);
        assert_eq!(b.to_be_vec(), [128]);
        b.set_bit(7, false);
        assert_eq!(b.to_be_vec(), [0]);
        assert_eq!(b.len(), 8);

        b = BooleanBitfield::new();
        b.set_bit(7, true);
        b.set_bit(0, true);
        assert_eq!(b.to_be_vec(), [129]);
        b.set_bit(7, false);
        assert_eq!(b.to_be_vec(), [1]);
        assert_eq!(b.len(), 8);

        b = BooleanBitfield::new();
        b.set_bit(8, true);
        assert_eq!(b.to_be_vec(), [1, 0]);
        assert_eq!(b.len(), 9);
        b.set_bit(8, false);
        assert_eq!(b.to_be_vec(), [0, 0]);
        assert_eq!(b.len(), 9);

        b = BooleanBitfield::new();
        b.set_bit(15, true);
        assert_eq!(b.to_be_vec(), [128, 0]);
        b.set_bit(15, false);
        assert_eq!(b.to_be_vec(), [0, 0]);
        assert_eq!(b.len(), 16);

        b = BooleanBitfield::new();
        b.set_bit(8, true);
        b.set_bit(15, true);
        assert_eq!(b.to_be_vec(), [129, 0]);
        b.set_bit(15, false);
        assert_eq!(b.to_be_vec(), [1, 0]);
        assert_eq!(b.len(), 16);
    }

    #[test]
    fn test_bitfield_get() {
        let test_nums = vec![0, 8, 15, 42, 1337];
        for i in test_nums {
            let mut b = BooleanBitfield::new();
            assert_eq!(b.get_bit(i), false);
            b.set_bit(i, true);
            assert_eq!(b.get_bit(i), true);
            b.set_bit(i, true);
        }
    }

    #[test]
    fn test_bitfield_num_true_bits() {
        let mut b = BooleanBitfield::new();
        assert_eq!(b.num_true_bits(), 0);
        b.set_bit(15, true);
        assert_eq!(b.num_true_bits(), 1);
        b.set_bit(15, false);
        assert_eq!(b.num_true_bits(), 0);
        b.set_bit(0, true);
        b.set_bit(7, true);
        b.set_bit(8, true);
        b.set_bit(1337, true);
        assert_eq!(b.num_true_bits(), 4);
    }
}
<|endoftext|>


<commit_before>#![feature(async_closure)]

use futures::stream::StreamExt;
use std::io;
use std::io::Write;

use term_input::{Event, Input, Key};

fn main() {
    // put the terminal in non-buffering, no-enchoing mode
    let mut old_term: libc::termios = unsafe { std::mem::zeroed() };
    unsafe {
        libc::tcgetattr(libc::STDIN_FILENO, &mut old_term);
    }

    let mut new_term: libc::termios = old_term;
    new_term.c_iflag &= !(libc::IGNBRK
        | libc::BRKINT
        | libc::PARMRK
        | libc::ISTRIP
        | libc::INLCR
        | libc::IGNCR
        | libc::ICRNL
        | libc::IXON);
    new_term.c_lflag &= !(libc::ICANON | libc::ECHO | libc::ISIG | libc::IEXTEN);
    unsafe { libc::tcsetattr(libc::STDIN_FILENO, libc::TCSAFLUSH, &new_term) };

    // enable focus events
    {
        let stdout = io::stdout();
        stdout.lock().write_all(b"\x1b[?1004h").unwrap();
        stdout.lock().flush().unwrap();
    }

    /* DO THE BUSINESS HERE */
    let mut input = Input::new();
    let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();
    executor.spawn(async move {
        while let Some(mb_ev) = input.next().await {
            match mb_ev {
                Ok(ev) => {
                    println!("{:?}", ev);
                    if ev == Event::Key(Key::Esc) {
                        break;
                    }
                }
                Err(io_err) => {
                    println!("Error: {:?}", io_err);
                    break;
                }
            }
        }
    });
    executor.run().unwrap();

    // restore the old settings
    unsafe { libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &old_term) };

    // disable focus events
    {
        let stdout = io::stdout();
        stdout.lock().write_all(b"\x1b[?1004l").unwrap();
    }
}
<commit_msg>term_input: Fix example<commit_after>#![feature(async_closure)]

use futures::stream::StreamExt;
use std::io;
use std::io::Write;

use term_input::{Event, Input, Key};

fn main() {
    // put the terminal in non-buffering, no-enchoing mode
    let mut old_term: libc::termios = unsafe { std::mem::zeroed() };
    unsafe {
        libc::tcgetattr(libc::STDIN_FILENO, &mut old_term);
    }

    let mut new_term: libc::termios = old_term;
    new_term.c_iflag &= !(libc::IGNBRK
        | libc::BRKINT
        | libc::PARMRK
        | libc::ISTRIP
        | libc::INLCR
        | libc::IGNCR
        | libc::ICRNL
        | libc::IXON);
    new_term.c_lflag &= !(libc::ICANON | libc::ECHO | libc::ISIG | libc::IEXTEN);
    unsafe { libc::tcsetattr(libc::STDIN_FILENO, libc::TCSAFLUSH, &new_term) };

    // enable focus events
    {
        let stdout = io::stdout();
        stdout.lock().write_all(b"\x1b[?1004h").unwrap();
        stdout.lock().flush().unwrap();
    }

    /* DO THE BUSINESS HERE */
    let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();
    executor.spawn(async move {
        let mut input = Input::new();
        while let Some(mb_ev) = input.next().await {
            match mb_ev {
                Ok(ev) => {
                    println!("{:?}", ev);
                    if ev == Event::Key(Key::Esc) {
                        break;
                    }
                }
                Err(io_err) => {
                    println!("Error: {:?}", io_err);
                    break;
                }
            }
        }
    });
    executor.run().unwrap();

    // restore the old settings
    unsafe { libc::tcsetattr(libc::STDIN_FILENO, libc::TCSANOW, &old_term) };

    // disable focus events
    {
        let stdout = io::stdout();
        stdout.lock().write_all(b"\x1b[?1004l").unwrap();
    }
}
<|endoftext|>


<commit_before>///! Inline of `https://github.com/bltavares/atomic-shim`

#[cfg(not(any(
    target_arch = "mips",
    target_arch = "powerpc",
    feature = "mutex"
)))]
pub use std::sync::atomic::{AtomicI64, AtomicU64};
#[cfg(any(target_arch = "mips", target_arch = "powerpc", feature = "mutex"))]
mod shim {
    use parking_lot::{const_rwlock, RwLock};
    use std::sync::atomic::Ordering;

    #[derive(Debug, Default)]
    pub struct AtomicU64 {
        value: RwLock<u64>,
    }

    impl AtomicU64 {
        pub const fn new(v: u64) -> Self {
            Self { value: const_rwlock(v) }
        }

        #[allow(dead_code)]
        pub fn load(&self, _: Ordering) -> u64 {
            *self.value.read()
        }

        #[allow(dead_code)]
        pub fn store(&self, value: u64, _: Ordering) {
            let mut lock = self.value.write();
            *lock = value;
        }

        #[allow(dead_code)]
        pub fn swap(&self, value: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = value;
            prev
        }

        #[allow(dead_code)]
        pub fn compare_exchange(
            &self,
            current: u64,
            new: u64,
            _: Ordering,
            _: Ordering,
        ) -> Result<u64, u64> {
            let mut lock = self.value.write();
            let prev = *lock;
            if prev == current {
                *lock = new;
                Ok(current)
            } else {
                Err(prev)
            }
        }

        #[allow(dead_code)]
        pub fn compare_exchange_weak(
            &self,
            current: u64,
            new: u64,
            success: Ordering,
            failure: Ordering,
        ) -> Result<u64, u64> {
            self.compare_exchange(current, new, success, failure)
        }

        #[allow(dead_code)]
        pub fn fetch_add(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_add(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_sub(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_sub(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_and(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev & val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_nand(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = !(prev & val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_or(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev | val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_xor(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev ^ val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_max(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.max(val);
            prev
        }
    }

    impl From<u64> for AtomicU64 {
        fn from(value: u64) -> Self {
            AtomicU64::new(value)
        }
    }

    #[derive(Debug, Default)]
    pub struct AtomicI64 {
        value: RwLock<i64>,
    }

    impl AtomicI64 {
        pub const fn new(v: i64) -> Self {
            Self { value: const_rwlock(v) }
        }

        #[allow(dead_code)]
        pub fn load(&self, _: Ordering) -> i64 {
            *self.value.read()
        }

        #[allow(dead_code)]
        pub fn store(&self, value: i64, _: Ordering) {
            let mut lock = self.value.write();
            *lock = value;
        }

        #[allow(dead_code)]
        pub fn swap(&self, value: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = value;
            prev
        }

        #[allow(dead_code)]
        pub fn compare_exchange(
            &self,
            current: i64,
            new: i64,
            _: Ordering,
            _: Ordering,
        ) -> Result<i64, i64> {
            let mut lock = self.value.write();
            let prev = *lock;
            if prev == current {
                *lock = new;
                Ok(current)
            } else {
                Err(prev)
            }
        }

        #[allow(dead_code)]
        pub fn compare_exchange_weak(
            &self,
            current: i64,
            new: i64,
            success: Ordering,
            failure: Ordering,
        ) -> Result<i64, i64> {
            self.compare_exchange(current, new, success, failure)
        }

        #[allow(dead_code)]
        pub fn fetch_add(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_add(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_sub(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_sub(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_and(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev & val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_nand(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = !(prev & val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_or(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev | val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_xor(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev ^ val;
            prev
        }
    }

    impl From<i64> for AtomicI64 {
        fn from(value: i64) -> Self {
            AtomicI64::new(value)
        }
    }
}

#[cfg(any(
    target_arch = "mips",
    target_arch = "powerpc",
    feature = "mutex"
))]
pub use shim::{AtomicI64, AtomicU64};
<commit_msg>Add atomic_shim impl for fetch_max on AtomicI64 as well<commit_after>///! Inline of `https://github.com/bltavares/atomic-shim`

#[cfg(not(any(
    target_arch = "mips",
    target_arch = "powerpc",
    feature = "mutex"
)))]
pub use std::sync::atomic::{AtomicI64, AtomicU64};
#[cfg(any(target_arch = "mips", target_arch = "powerpc", feature = "mutex"))]
mod shim {
    use parking_lot::{const_rwlock, RwLock};
    use std::sync::atomic::Ordering;

    #[derive(Debug, Default)]
    pub struct AtomicU64 {
        value: RwLock<u64>,
    }

    impl AtomicU64 {
        pub const fn new(v: u64) -> Self {
            Self { value: const_rwlock(v) }
        }

        #[allow(dead_code)]
        pub fn load(&self, _: Ordering) -> u64 {
            *self.value.read()
        }

        #[allow(dead_code)]
        pub fn store(&self, value: u64, _: Ordering) {
            let mut lock = self.value.write();
            *lock = value;
        }

        #[allow(dead_code)]
        pub fn swap(&self, value: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = value;
            prev
        }

        #[allow(dead_code)]
        pub fn compare_exchange(
            &self,
            current: u64,
            new: u64,
            _: Ordering,
            _: Ordering,
        ) -> Result<u64, u64> {
            let mut lock = self.value.write();
            let prev = *lock;
            if prev == current {
                *lock = new;
                Ok(current)
            } else {
                Err(prev)
            }
        }

        #[allow(dead_code)]
        pub fn compare_exchange_weak(
            &self,
            current: u64,
            new: u64,
            success: Ordering,
            failure: Ordering,
        ) -> Result<u64, u64> {
            self.compare_exchange(current, new, success, failure)
        }

        #[allow(dead_code)]
        pub fn fetch_add(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_add(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_sub(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_sub(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_and(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev & val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_nand(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = !(prev & val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_or(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev | val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_xor(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev ^ val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_max(&self, val: u64, _: Ordering) -> u64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.max(val);
            prev
        }
    }

    impl From<u64> for AtomicU64 {
        fn from(value: u64) -> Self {
            AtomicU64::new(value)
        }
    }

    #[derive(Debug, Default)]
    pub struct AtomicI64 {
        value: RwLock<i64>,
    }

    impl AtomicI64 {
        pub const fn new(v: i64) -> Self {
            Self { value: const_rwlock(v) }
        }

        #[allow(dead_code)]
        pub fn load(&self, _: Ordering) -> i64 {
            *self.value.read()
        }

        #[allow(dead_code)]
        pub fn store(&self, value: i64, _: Ordering) {
            let mut lock = self.value.write();
            *lock = value;
        }

        #[allow(dead_code)]
        pub fn swap(&self, value: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = value;
            prev
        }

        #[allow(dead_code)]
        pub fn compare_exchange(
            &self,
            current: i64,
            new: i64,
            _: Ordering,
            _: Ordering,
        ) -> Result<i64, i64> {
            let mut lock = self.value.write();
            let prev = *lock;
            if prev == current {
                *lock = new;
                Ok(current)
            } else {
                Err(prev)
            }
        }

        #[allow(dead_code)]
        pub fn compare_exchange_weak(
            &self,
            current: i64,
            new: i64,
            success: Ordering,
            failure: Ordering,
        ) -> Result<i64, i64> {
            self.compare_exchange(current, new, success, failure)
        }

        #[allow(dead_code)]
        pub fn fetch_add(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_add(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_sub(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.wrapping_sub(val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_and(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev & val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_nand(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = !(prev & val);
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_or(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev | val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_xor(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev ^ val;
            prev
        }

        #[allow(dead_code)]
        pub fn fetch_max(&self, val: i64, _: Ordering) -> i64 {
            let mut lock = self.value.write();
            let prev = *lock;
            *lock = prev.max(val);
            prev
        }
    }

    impl From<i64> for AtomicI64 {
        fn from(value: i64) -> Self {
            AtomicI64::new(value)
        }
    }
}

#[cfg(any(
    target_arch = "mips",
    target_arch = "powerpc",
    feature = "mutex"
))]
pub use shim::{AtomicI64, AtomicU64};
<|endoftext|>


<commit_before><commit_msg>cargo<commit_after>fn main() {

}
<|endoftext|>


<commit_before>use proc_macro2::{Span, TokenStream};
use shared::{map_lifetimes, map_type_params, split_for_impl};
use syn::{
    self, Data, DataEnum, DataStruct, DeriveInput, Field, Fields, FieldsNamed, FieldsUnnamed,
    Generics, Ident, Variant,
};

use attr;

pub fn derive(input: TokenStream) -> TokenStream {
    let derive_input = syn::parse2(input).expect("Input is checked by rustc");

    let container = attr::Container::from_ast(&derive_input);

    let DeriveInput {
        ident,
        data,
        generics,
        ..
    } = derive_input;

    let tokens = match data {
        Data::Struct(ast) => derive_struct(&container, ast, ident, generics),
        Data::Enum(ast) => derive_enum(&container, ast, ident, generics),
        Data::Union(_) => panic!("Unions are not supported"),
    };

    tokens.into()
}

fn derive_struct(
    container: &attr::Container,
    ast: DataStruct,
    ident: Ident,
    generics: Generics,
) -> TokenStream {
    let cons = match ast.fields {
        Fields::Named(FieldsNamed { named, .. }) => gen_struct_cons(&ident, named),
        Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => gen_tuple_struct_cons(&ident, unnamed),
        Fields::Unit => quote! { #ident },
    };

    gen_impl(container, ident, generics, cons)
}

fn gen_struct_cons<I>(ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = Field>,
{
    // lookup each field by its name and then convert to its type using the Getable
    // impl of the fields type
    let field_initializers = fields.into_iter().map(|field| {
        let field_ty = &field.ty;
        let ident = field
            .ident
            .as_ref()
            .expect("Struct fields always have names");
        let quoted_ident = format!("{}", quote! { #ident });

        quote! {
            #ident: if let Some(val) = data.lookup_field(vm, #quoted_ident) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Cannot find the field '{}'. Do the type definitions match?", #quoted_ident);
            }
        }
    });

    let unpack_data = unpack_data();
    quote! {
        #unpack_data

        #ident {
            #(#field_initializers,)*
        }
    }
}

fn gen_tuple_struct_cons<I>(ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = Field>,
{
    let mut fields = fields.into_iter().fuse();

    // Treat newtype structs as just their inner type
    let (first, second) = (fields.next(), fields.next());
    match (&first, &second) {
        (Some(field), None) => {
            let field_ty = &field.ty;
            return quote! {
                #ident (
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, variants)
                )
            };
        }
        _ => (),
    }

    // do the lookup using the tag, because tuple structs don't have field names
    let field_initializers = first
        .into_iter()
        .chain(second)
        .chain(fields)
        .enumerate()
        .map(|(tag, field)| {
            let field_ty = &field.ty;

            quote! {
                if let Some(val) = data.get_variant(#tag) {
                    <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
                } else {
                    panic!("Cannot find the field with tag '{}'. Do the type definitions match?", #tag);
                }
            }
        });

    let unpack_data = unpack_data();
    quote! {
        #unpack_data

        #ident (
            #(#field_initializers,)*
        )
    }
}

fn derive_enum(
    container: &attr::Container,
    ast: DataEnum,
    ident: Ident,
    generics: Generics,
) -> TokenStream {
    let cons;
    {
        let variants = ast
            .variants
            .iter()
            .enumerate()
            .map(|(tag, variant)| gen_variant_match(&ident, tag, variant));

        let unpack_data = unpack_data();

        // data contains the the data for each field of a variant; the variant of the passed value
        // is defined by the tag(), which is defined by order of the variants (the first variant is 0)
        cons = quote! {
            #unpack_data

            match data.tag() as usize {
                #(#variants,)*
                tag => panic!("Unexpected tag: '{}'. Do the type definitions match?", tag)
            }
        };
    }

    gen_impl(container, ident, generics, cons)
}

fn gen_impl(
    container: &attr::Container,
    ident: Ident,
    generics: Generics,
    push_impl: TokenStream,
) -> TokenStream {
    // lifetime bounds like '__vm: 'a, 'a: '__vm (which implies => 'a == '__vm)
    // writing bounds like this is a lot easier than actually replacing all lifetimes
    // with '__vm
    let lifetime_bounds = create_lifetime_bounds(&generics);

    // generate bounds like T: Getable for every type parameter
    let getable_bounds = create_getable_bounds(&generics);

    let (impl_generics, ty_generics, where_clause) =
        split_for_impl(&generics, &["'__vm", "'__value"]);

    let dummy_const = Ident::new(&format!("_IMPL_GETABLE_FOR_{}", ident), Span::call_site());

    let gluon = match container.crate_name {
        attr::CrateName::Some(ref ident) => quote! {
            use #ident::api as _gluon_api;
            use #ident::thread as _gluon_thread;
            use #ident::Variants as _GluonVariants;
            use #ident::Result as _GluonResult;
        },
        attr::CrateName::GluonVm => quote! {
            use crate::api as _gluon_api;
            use crate::thread as _gluon_thread;
            use crate::Variants as _GluonVariants;
            use crate::Result as _GluonResult;
        },
        attr::CrateName::None => quote! {
            use gluon::vm::api as _gluon_api;
            use gluon::vm::thread as _gluon_thread;
            use gluon::vm::Variants as _GluonVariants;
            use gluon::vm::Result as _GluonResult;
        },
    };

    quote! {
        #[allow(non_upper_case_globals)]
        const #dummy_const: () = {
            #gluon

            #[automatically_derived]
            #[allow(unused_attributes, unused_variables)]
            impl #impl_generics _gluon_api::Getable<'__vm, '__value> for #ident #ty_generics
            #where_clause #(#getable_bounds,)* #(#lifetime_bounds),*
            {
                type Proxy = _GluonVariants<'__value>;

                fn to_proxy(vm: &'__vm _gluon_thread::Thread, value: _GluonVariants<'__value>) -> _GluonResult<Self::Proxy> {
                    Ok(value)
                }

                fn from_proxy(vm: &'__vm _gluon_thread::Thread, proxy: &'__value mut Self::Proxy) -> Self {
                    Self::from_value(vm, proxy.clone())
                }

                fn from_value(vm: &'__vm _gluon_thread::Thread, variants: _GluonVariants<'__value>) -> Self {
                    #push_impl
                }
            }
        };
    }
}

fn gen_variant_match(ident: &Ident, tag: usize, variant: &Variant) -> TokenStream {
    let variant_ident = &variant.ident;

    // depending on the type of the variant we need to generate different constructors
    // for the enum
    match &variant.fields {
        Fields::Unit => quote! {
            #tag => #ident::#variant_ident
        },
        // both constructors that need to marshall values extract them by using the index
        // of the field to get the content from Data::get_variant;
        // the data variable was assigned in the function body above
        Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
            let cons = gen_tuple_variant_cons(unnamed);

            quote! {
                #tag => #ident::#variant_ident#cons
            }
        }
        Fields::Named(FieldsNamed { named, .. }) => {
            let cons = gen_struct_variant_cons(ident, variant_ident, named);

            quote! {
                #tag => #cons
            }
        }
    }
}

fn gen_tuple_variant_cons<'a, I>(fields: I) -> TokenStream
where
    I: IntoIterator<Item = &'a Field>,
{
    let fields = fields.into_iter().enumerate().map(|(idx, field)| {
        let field_ty = &field.ty;

        quote! {
            if let Some(val) = data.get_variant(#idx) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Enum does not contain data at index '{}'. Do the type definitions match?", #idx)
            }
        }
    });

    quote! {
        (#(#fields),*)
    }
}

fn gen_struct_variant_cons<'a, I>(ident: &Ident, variant_ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = &'a Field>,
{
    let fields = fields.into_iter().enumerate().map(|(idx, field)| {
        let field_ty = &field.ty;
        let field_ident = field
            .ident
            .as_ref()
            .expect("Struct fields always have names");
        let quoted_field_ident = format!("{}", quote! { #field_ident });
        quote! {
            #field_ident: if let Some(val) = inner_data.lookup_field(vm, #quoted_field_ident) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Enum does not contain data at index '{}'. Do the type definitions match?", #idx)
            }
        }
    });

    quote! {{
        if let Some(val) = data.get_variant(0) {
            let inner_data = match val.as_ref() {
                _gluon_api::ValueRef::Data(data) => data,
                val => panic!("Unexpected value: '{:?}'. Do the type definitions match?", val),
            };
            #ident::#variant_ident{#(#fields),*}
        } else {
            panic!("Enum does not contain data at index '0'. Do the type definitions match?")
        }
    }}
}

fn create_getable_bounds(generics: &Generics) -> Vec<TokenStream> {
    map_type_params(generics, |ty| {
        quote! {
            #ty: _gluon_api::Getable<'__vm, '__value>
        }
    })
}

fn create_lifetime_bounds(generics: &Generics) -> Vec<TokenStream> {
    map_lifetimes(generics, |lifetime| {
        quote! { #lifetime: '__value, '__value: #lifetime }
    })
}

fn unpack_data() -> TokenStream {
    quote! {
        let data = match variants.as_ref() {
            _gluon_api::ValueRef::Data(data) => data,
            val => panic!("Unexpected value: '{:?}'. Do the type definitions match?", val),
        };
    }
}
<commit_msg>Fix error message on struct extraction<commit_after>use proc_macro2::{Span, TokenStream};
use shared::{map_lifetimes, map_type_params, split_for_impl};
use syn::{
    self, Data, DataEnum, DataStruct, DeriveInput, Field, Fields, FieldsNamed, FieldsUnnamed,
    Generics, Ident, Variant,
};

use attr;

pub fn derive(input: TokenStream) -> TokenStream {
    let derive_input = syn::parse2(input).expect("Input is checked by rustc");

    let container = attr::Container::from_ast(&derive_input);

    let DeriveInput {
        ident,
        data,
        generics,
        ..
    } = derive_input;

    let tokens = match data {
        Data::Struct(ast) => derive_struct(&container, ast, ident, generics),
        Data::Enum(ast) => derive_enum(&container, ast, ident, generics),
        Data::Union(_) => panic!("Unions are not supported"),
    };

    tokens.into()
}

fn derive_struct(
    container: &attr::Container,
    ast: DataStruct,
    ident: Ident,
    generics: Generics,
) -> TokenStream {
    let cons = match ast.fields {
        Fields::Named(FieldsNamed { named, .. }) => gen_struct_cons(&ident, named),
        Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => gen_tuple_struct_cons(&ident, unnamed),
        Fields::Unit => quote! { #ident },
    };

    gen_impl(container, ident, generics, cons)
}

fn gen_struct_cons<I>(ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = Field>,
{
    // lookup each field by its name and then convert to its type using the Getable
    // impl of the fields type
    let field_initializers = fields.into_iter().map(|field| {
        let field_ty = &field.ty;
        let ident = field
            .ident
            .as_ref()
            .expect("Struct fields always have names");
        let quoted_ident = format!("{}", quote! { #ident });

        quote! {
            #ident: if let Some(val) = data.lookup_field(vm, #quoted_ident) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Cannot find the field '{}'. Do the type definitions match?", #quoted_ident);
            }
        }
    });

    let unpack_data = unpack_data();
    quote! {
        #unpack_data

        #ident {
            #(#field_initializers,)*
        }
    }
}

fn gen_tuple_struct_cons<I>(ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = Field>,
{
    let mut fields = fields.into_iter().fuse();

    // Treat newtype structs as just their inner type
    let (first, second) = (fields.next(), fields.next());
    match (&first, &second) {
        (Some(field), None) => {
            let field_ty = &field.ty;
            return quote! {
                #ident (
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, variants)
                )
            };
        }
        _ => (),
    }

    // do the lookup using the tag, because tuple structs don't have field names
    let field_initializers = first
        .into_iter()
        .chain(second)
        .chain(fields)
        .enumerate()
        .map(|(tag, field)| {
            let field_ty = &field.ty;

            quote! {
                if let Some(val) = data.get_variant(#tag) {
                    <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
                } else {
                    panic!("Cannot find the field with tag '{}'. Do the type definitions match?", #tag);
                }
            }
        });

    let unpack_data = unpack_data();
    quote! {
        #unpack_data

        #ident (
            #(#field_initializers,)*
        )
    }
}

fn derive_enum(
    container: &attr::Container,
    ast: DataEnum,
    ident: Ident,
    generics: Generics,
) -> TokenStream {
    let cons;
    {
        let variants = ast
            .variants
            .iter()
            .enumerate()
            .map(|(tag, variant)| gen_variant_match(&ident, tag, variant));

        let unpack_data = unpack_data();

        // data contains the the data for each field of a variant; the variant of the passed value
        // is defined by the tag(), which is defined by order of the variants (the first variant is 0)
        cons = quote! {
            #unpack_data

            match data.tag() as usize {
                #(#variants,)*
                tag => panic!("Unexpected tag: '{}'. Do the type definitions match?", tag)
            }
        };
    }

    gen_impl(container, ident, generics, cons)
}

fn gen_impl(
    container: &attr::Container,
    ident: Ident,
    generics: Generics,
    push_impl: TokenStream,
) -> TokenStream {
    // lifetime bounds like '__vm: 'a, 'a: '__vm (which implies => 'a == '__vm)
    // writing bounds like this is a lot easier than actually replacing all lifetimes
    // with '__vm
    let lifetime_bounds = create_lifetime_bounds(&generics);

    // generate bounds like T: Getable for every type parameter
    let getable_bounds = create_getable_bounds(&generics);

    let (impl_generics, ty_generics, where_clause) =
        split_for_impl(&generics, &["'__vm", "'__value"]);

    let dummy_const = Ident::new(&format!("_IMPL_GETABLE_FOR_{}", ident), Span::call_site());

    let gluon = match container.crate_name {
        attr::CrateName::Some(ref ident) => quote! {
            use #ident::api as _gluon_api;
            use #ident::thread as _gluon_thread;
            use #ident::Variants as _GluonVariants;
            use #ident::Result as _GluonResult;
        },
        attr::CrateName::GluonVm => quote! {
            use crate::api as _gluon_api;
            use crate::thread as _gluon_thread;
            use crate::Variants as _GluonVariants;
            use crate::Result as _GluonResult;
        },
        attr::CrateName::None => quote! {
            use gluon::vm::api as _gluon_api;
            use gluon::vm::thread as _gluon_thread;
            use gluon::vm::Variants as _GluonVariants;
            use gluon::vm::Result as _GluonResult;
        },
    };

    quote! {
        #[allow(non_upper_case_globals)]
        const #dummy_const: () = {
            #gluon

            #[automatically_derived]
            #[allow(unused_attributes, unused_variables)]
            impl #impl_generics _gluon_api::Getable<'__vm, '__value> for #ident #ty_generics
            #where_clause #(#getable_bounds,)* #(#lifetime_bounds),*
            {
                type Proxy = _GluonVariants<'__value>;

                fn to_proxy(vm: &'__vm _gluon_thread::Thread, value: _GluonVariants<'__value>) -> _GluonResult<Self::Proxy> {
                    Ok(value)
                }

                fn from_proxy(vm: &'__vm _gluon_thread::Thread, proxy: &'__value mut Self::Proxy) -> Self {
                    Self::from_value(vm, proxy.clone())
                }

                fn from_value(vm: &'__vm _gluon_thread::Thread, variants: _GluonVariants<'__value>) -> Self {
                    #push_impl
                }
            }
        };
    }
}

fn gen_variant_match(ident: &Ident, tag: usize, variant: &Variant) -> TokenStream {
    let variant_ident = &variant.ident;

    // depending on the type of the variant we need to generate different constructors
    // for the enum
    match &variant.fields {
        Fields::Unit => quote! {
            #tag => #ident::#variant_ident
        },
        // both constructors that need to marshall values extract them by using the index
        // of the field to get the content from Data::get_variant;
        // the data variable was assigned in the function body above
        Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => {
            let cons = gen_tuple_variant_cons(unnamed);

            quote! {
                #tag => #ident::#variant_ident#cons
            }
        }
        Fields::Named(FieldsNamed { named, .. }) => {
            let cons = gen_struct_variant_cons(ident, variant_ident, named);

            quote! {
                #tag => #cons
            }
        }
    }
}

fn gen_tuple_variant_cons<'a, I>(fields: I) -> TokenStream
where
    I: IntoIterator<Item = &'a Field>,
{
    let fields = fields.into_iter().enumerate().map(|(idx, field)| {
        let field_ty = &field.ty;

        quote! {
            if let Some(val) = data.get_variant(#idx) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Enum does not contain data at index '{}'. Do the type definitions match?", #idx)
            }
        }
    });

    quote! {
        (#(#fields),*)
    }
}

fn gen_struct_variant_cons<'a, I>(ident: &Ident, variant_ident: &Ident, fields: I) -> TokenStream
where
    I: IntoIterator<Item = &'a Field>,
{
    let fields = fields.into_iter().map(|field| {
        let field_ty = &field.ty;
        let field_ident = field
            .ident
            .as_ref()
            .expect("Struct fields always have names");
        let quoted_field_ident = format!("{}", quote! { #field_ident });
        quote! {
            #field_ident: if let Some(val) = inner_data.lookup_field(vm, #quoted_field_ident) {
                <#field_ty as _gluon_api::Getable<'__vm, '__value>>::from_value(vm, val)
            } else {
                panic!("Enum does not contain the field `{}`. Do the type definitions match?", #quoted_field_ident)
            }
        }
    });

    quote! {{
        if let Some(val) = data.get_variant(0) {
            let inner_data = match val.as_ref() {
                _gluon_api::ValueRef::Data(data) => data,
                val => panic!("Unexpected value: '{:?}'. Do the type definitions match?", val),
            };
            #ident::#variant_ident{#(#fields),*}
        } else {
            panic!("Enum does not contain data at index '0'. Do the type definitions match?")
        }
    }}
}

fn create_getable_bounds(generics: &Generics) -> Vec<TokenStream> {
    map_type_params(generics, |ty| {
        quote! {
            #ty: _gluon_api::Getable<'__vm, '__value>
        }
    })
}

fn create_lifetime_bounds(generics: &Generics) -> Vec<TokenStream> {
    map_lifetimes(generics, |lifetime| {
        quote! { #lifetime: '__value, '__value: #lifetime }
    })
}

fn unpack_data() -> TokenStream {
    quote! {
        let data = match variants.as_ref() {
            _gluon_api::ValueRef::Data(data) => data,
            val => panic!("Unexpected value: '{:?}'. Do the type definitions match?", val),
        };
    }
}
<|endoftext|>


<commit_before>//! A small library meant to be used as a build dependency with Cargo for easily
//! integrating [ISPC](https://ispc.github.io/) code into Rust projects.
//!
//! DOCS TODO for runtime lib

#![allow(dead_code)]

extern crate libc;
extern crate aligned_alloc;
extern crate num_cpus;

pub mod task;
pub mod exec;
pub mod instrument;

use std::mem;
use std::sync::{Once, ONCE_INIT, Arc};
use std::ffi::CStr;

pub use task::ISPCTaskFn;
pub use exec::{TaskSystem, Parallel};
pub use instrument::{Instrument, SimpleInstrument};

/// Convenience macro for generating the module to hold the raw/unsafe ISPC bindings.
///
/// In addition to building the library with ISPC we use rust-bindgen to generate
/// a rust module containing bindings to the functions exported from ISPC. These
/// can be imported by passing the name of your library to the `ispc_module` macro.
///
/// # Example
///
/// ```ignore
/// #[macro_use]
/// extern crate ispc;
///
/// // Functions exported from foo will be callable under foo::*
/// ispc_module!(foo);
/// ```
#[macro_export]
macro_rules! ispc_module {
    ($lib:ident) => (
        include!(concat!(env!("OUT_DIR"), "/", stringify!($lib), ".rs"));
    )
}

static mut TASK_SYSTEM: Option<&'static TaskSystem> = None;
static TASK_INIT: Once = ONCE_INIT;

static mut INSTRUMENT: Option<&'static Instrument> = None;
static INSTRUMENT_INIT: Once = ONCE_INIT;

/// If you have implemented your own task system you can provide it for use instead
/// of the default threaded one. This must be done prior to calling ISPC code which
/// spawns tasks otherwise the task system will have already been initialized to
/// `Parallel`, which you can also see as an example for implementing a task system.
///
/// Use the function to do any extra initialization for your task system. Note that
/// the task system will be leaked and not destroyed until the program exits and the
/// memory space is cleaned up.
pub fn set_task_system<F: FnOnce() -> Arc<TaskSystem>>(f: F) {
    TASK_INIT.call_once(|| {
        let task_sys = f();
        unsafe {
            let s: *const TaskSystem = mem::transmute(&*task_sys);
            mem::forget(task_sys);
            TASK_SYSTEM = Some(&*s);
        }
    });
}

fn get_task_system() -> &'static TaskSystem {
    // TODO: This is a bit nasty, but I'm not sure on a nicer solution. Maybe something that
    // would let the user register the desired (or default) task system? But if
    // mutable statics can't have destructors we still couldn't have an Arc or Box to something?
    TASK_INIT.call_once(|| {
        unsafe {
            let task_sys = Parallel::new() as Arc<TaskSystem>;
            let s: *const TaskSystem = mem::transmute(&*task_sys);
            mem::forget(task_sys);
            TASK_SYSTEM = Some(&*s);
        }
    });
    unsafe { TASK_SYSTEM.unwrap() }
}

/// If you have implemented your own instrument for logging ISPC performance
/// data you can use this function to provide it for use instead of the
/// default one. This function **must** be called before calling into ISPC code,
/// otherwise the instrumenter will already be set to the default.
pub fn set_instrument<F: FnOnce() -> Arc<Instrument>>(f: F) {
    INSTRUMENT_INIT.call_once(|| {
        let instrument = f();
        unsafe {
            let s: *const Instrument = mem::transmute(&*instrument);
            mem::forget(instrument);
            INSTRUMENT = Some(&*s);
        }
    });
}

/// Print out a summary of performace data gathered from instrumenting ISPC.
/// Must enable instrumenting to have this record and print data, see
/// `Config::instrument`.
pub fn print_instrumenting_summary() {
    get_instrument().print_summary();
}

fn get_instrument() -> &'static Instrument {
    // TODO: This is a bit nasty, like above
    INSTRUMENT_INIT.call_once(|| {
        unsafe {
            let instrument = Arc::new(SimpleInstrument) as Arc<Instrument>;
            let s: *const Instrument = mem::transmute(&*instrument);
            mem::forget(instrument);
            INSTRUMENT = Some(&*s);
        }
    });
    unsafe { INSTRUMENT.unwrap() }
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCAlloc(handle_ptr: *mut *mut libc::c_void, size: libc::int64_t,
                                   align: libc::int32_t) -> *mut libc::c_void {
    get_task_system().alloc(handle_ptr, size, align)
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCLaunch(handle_ptr: *mut *mut libc::c_void, f: *mut libc::c_void,
                                    data: *mut libc::c_void, count0: libc::c_int,
                                    count1: libc::c_int, count2: libc::c_int) {
    let task_fn: ISPCTaskFn = mem::transmute(f);
    get_task_system().launch(handle_ptr, task_fn, data, count0 as i32, count1 as i32, count2 as i32);
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCSync(handle: *mut libc::c_void){
    get_task_system().sync(handle);
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCInstrument(cfile: *const libc::c_char, cnote: *const libc::c_char,
                                        line: libc::c_int, mask: libc::uint64_t) {

    let file_name = CStr::from_ptr(cfile);
    let note = CStr::from_ptr(cnote);
    let active_count = (mask as u64).count_ones();
    get_instrument().instrument(file_name, note, line as i32, mask as u64, active_count);
}

<commit_msg>Add utilities to link a pre-compiled packaged ISPC library<commit_after>//! A small library meant to be used as a build dependency with Cargo for easily
//! integrating [ISPC](https://ispc.github.io/) code into Rust projects.
//!
//! DOCS TODO for runtime lib

#![allow(dead_code)]

extern crate libc;
extern crate aligned_alloc;
extern crate num_cpus;

pub mod task;
pub mod exec;
pub mod instrument;

use std::mem;
use std::env;
use std::sync::{Once, ONCE_INIT, Arc};
use std::ffi::CStr;
use std::path::{Path, PathBuf};

pub use task::ISPCTaskFn;
pub use exec::{TaskSystem, Parallel};
pub use instrument::{Instrument, SimpleInstrument};

/// Convenience macro for generating the module to hold the raw/unsafe ISPC bindings.
///
/// In addition to building the library with ISPC we use rust-bindgen to generate
/// a rust module containing bindings to the functions exported from ISPC. These
/// can be imported by passing the name of your library to the `ispc_module` macro.
///
/// # Example
///
/// ```ignore
/// #[macro_use]
/// extern crate ispc;
///
/// // Functions exported from foo will be callable under foo::*
/// ispc_module!(foo);
/// ```
#[macro_export]
macro_rules! ispc_module {
    ($lib:ident) => (
        include!(concat!(env!("OUT_DIR"), "/", stringify!($lib), ".rs"));
    )
}

#[macro_export]
macro_rules! packaged_ispc_module {
    ($lib:ident) => (
        include!(concat!(stringify!($lib), ".rs"));
    )
}

pub struct PackagedModule {
    path: Option<PathBuf>,
    lib: String,
}

impl PackagedModule {
    pub fn new(lib: &str) -> PackagedModule {
        PackagedModule {path: None, lib: lib.to_owned()}
    }
    /// Specify the lib path to search for the packaged ISPC modules
    pub fn lib_path<P: AsRef<Path>>(&mut self, path: P) -> &mut PackagedModule {
        self.path = Some(path.as_ref().to_path_buf());
        self
    }
    /// Link with a previously built ISPC library packaged with the crate
    pub fn link(&self) {
        let libfile = self.lib.clone() + &env::var("HOST").unwrap();
        println!("cargo:rustc-link-lib=static={}", libfile);
        if let Some(ref p) = self.path {
            println!("cargo:rustc-link-search=native={}", p.display());
        } else {
            println!("cargo:rustc-link-search=native=./");
        }
    }
}

static mut TASK_SYSTEM: Option<&'static TaskSystem> = None;
static TASK_INIT: Once = ONCE_INIT;

static mut INSTRUMENT: Option<&'static Instrument> = None;
static INSTRUMENT_INIT: Once = ONCE_INIT;

/// If you have implemented your own task system you can provide it for use instead
/// of the default threaded one. This must be done prior to calling ISPC code which
/// spawns tasks otherwise the task system will have already been initialized to
/// `Parallel`, which you can also see as an example for implementing a task system.
///
/// Use the function to do any extra initialization for your task system. Note that
/// the task system will be leaked and not destroyed until the program exits and the
/// memory space is cleaned up.
pub fn set_task_system<F: FnOnce() -> Arc<TaskSystem>>(f: F) {
    TASK_INIT.call_once(|| {
        let task_sys = f();
        unsafe {
            let s: *const TaskSystem = mem::transmute(&*task_sys);
            mem::forget(task_sys);
            TASK_SYSTEM = Some(&*s);
        }
    });
}

fn get_task_system() -> &'static TaskSystem {
    // TODO: This is a bit nasty, but I'm not sure on a nicer solution. Maybe something that
    // would let the user register the desired (or default) task system? But if
    // mutable statics can't have destructors we still couldn't have an Arc or Box to something?
    TASK_INIT.call_once(|| {
        unsafe {
            let task_sys = Parallel::new() as Arc<TaskSystem>;
            let s: *const TaskSystem = mem::transmute(&*task_sys);
            mem::forget(task_sys);
            TASK_SYSTEM = Some(&*s);
        }
    });
    unsafe { TASK_SYSTEM.unwrap() }
}

/// If you have implemented your own instrument for logging ISPC performance
/// data you can use this function to provide it for use instead of the
/// default one. This function **must** be called before calling into ISPC code,
/// otherwise the instrumenter will already be set to the default.
pub fn set_instrument<F: FnOnce() -> Arc<Instrument>>(f: F) {
    INSTRUMENT_INIT.call_once(|| {
        let instrument = f();
        unsafe {
            let s: *const Instrument = mem::transmute(&*instrument);
            mem::forget(instrument);
            INSTRUMENT = Some(&*s);
        }
    });
}

/// Print out a summary of performace data gathered from instrumenting ISPC.
/// Must enable instrumenting to have this record and print data, see
/// `Config::instrument`.
pub fn print_instrumenting_summary() {
    get_instrument().print_summary();
}

fn get_instrument() -> &'static Instrument {
    // TODO: This is a bit nasty, like above
    INSTRUMENT_INIT.call_once(|| {
        unsafe {
            let instrument = Arc::new(SimpleInstrument) as Arc<Instrument>;
            let s: *const Instrument = mem::transmute(&*instrument);
            mem::forget(instrument);
            INSTRUMENT = Some(&*s);
        }
    });
    unsafe { INSTRUMENT.unwrap() }
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCAlloc(handle_ptr: *mut *mut libc::c_void, size: libc::int64_t,
                                   align: libc::int32_t) -> *mut libc::c_void {
    get_task_system().alloc(handle_ptr, size, align)
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCLaunch(handle_ptr: *mut *mut libc::c_void, f: *mut libc::c_void,
                                    data: *mut libc::c_void, count0: libc::c_int,
                                    count1: libc::c_int, count2: libc::c_int) {
    let task_fn: ISPCTaskFn = mem::transmute(f);
    get_task_system().launch(handle_ptr, task_fn, data, count0 as i32, count1 as i32, count2 as i32);
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCSync(handle: *mut libc::c_void){
    get_task_system().sync(handle);
}

#[allow(non_snake_case)]
#[doc(hidden)]
#[no_mangle]
pub unsafe extern "C" fn ISPCInstrument(cfile: *const libc::c_char, cnote: *const libc::c_char,
                                        line: libc::c_int, mask: libc::uint64_t) {

    let file_name = CStr::from_ptr(cfile);
    let note = CStr::from_ptr(cnote);
    let active_count = (mask as u64).count_ones();
    get_instrument().instrument(file_name, note, line as i32, mask as u64, active_count);
}

<|endoftext|>