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

//! Functions for randomly accessing and sampling sequences.

use super::Rng;
use std::collections::hash_map::HashMap;

/// Randomly sample `amount` elements from a finite iterator.
///
/// The following can be returned:
/// - `Ok`: `Vec` of `amount` non-repeating randomly sampled elements. The order is not random.
/// - `Err`: `Vec` of all the elements from `iterable` in sequential order. This happens when the
///   length of `iterable` was less than `amount`. This is considered an error since exactly
///   `amount` elements is typically expected.
///
/// This implementation uses `O(len(iterable))` time and `O(amount)` memory.
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let sample = seq::sample_iter(&mut rng, 1..100, 5).unwrap();
/// println!("{:?}", sample);
/// ```
pub fn sample_iter<T, I, R>(rng: &mut R, iterable: I, amount: usize) -> Result<Vec<T>, Vec<T>>
    where I: IntoIterator<Item=T>,
          R: Rng,
{
    let mut iter = iterable.into_iter();
    let mut reservoir = Vec::with_capacity(amount);
    reservoir.extend(iter.by_ref().take(amount));

    // Continue unless the iterator was exhausted
    //
    // note: this prevents iterators that "restart" from causing problems.
    // If the iterator stops once, then so do we.
    if reservoir.len() == amount {
        for (i, elem) in iter.enumerate() {
            let k = rng.gen_range(0, i + 1 + amount);
            if let Some(spot) = reservoir.get_mut(k) {
                *spot = elem;
            }
        }
        Ok(reservoir)
    } else {
        // Don't hang onto extra memory. There is a corner case where
        // `amount` was much less than `len(iterable)`.
        reservoir.shrink_to_fit();
        Err(reservoir)
    }
}

/// Randomly sample exactly `amount` values from `slice`.
///
/// The values are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// Panics if `amount > self.len()`
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let values = vec![5, 6, 1, 3, 4, 6, 7];
/// println!("{:?}", seq::sample_slice(&mut rng, &values, 3));
/// ```
pub fn sample_slice<R, T>(rng: &mut R, slice: &[T], amount: usize) -> Vec<T>
    where R: Rng,
          T: Clone
{
    let indices = sample_indices(rng, slice.len(), amount);

    let mut out = Vec::with_capacity(amount);
    out.extend(indices.iter().map(|i| slice[*i].clone()));
    out
}

/// Randomly sample exactly `amount` references from `slice`.
///
/// The references are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// Panics if `amount > self.len()`
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let values = vec![5, 6, 1, 3, 4, 6, 7];
/// println!("{:?}", seq::sample_slice_ref(&mut rng, &values, 3));
/// ```
pub fn sample_slice_ref<'a, R, T>(rng: &mut R, slice: &'a [T], amount: usize) -> Vec<&'a T>
    where R: Rng
{
    let indices = sample_indices(rng, slice.len(), amount);

    let mut out = Vec::with_capacity(amount);
    out.extend(indices.iter().map(|i| &slice[*i]));
    out
}

/// Randomly sample exactly `amount` indices from `0..length`.
///
/// The values are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// This method is used internally by the slice sampling methods, but it can sometimes be useful to
/// have the indices themselves so this is provided as an alternative.
///
/// Panics if `amount > self.len()`
pub fn sample_indices<R>(rng: &mut R, length: usize, amount: usize) -> Vec<usize>
    where R: Rng,
{
    if amount > length {
        panic!("`amount` must be less than or equal to `slice.len()`");
    }

    // We are going to have to allocate at least `amount` for the output no matter what. However,
    // if we use the `cached` version we will have to allocate `amount` as a HashMap as well since
    // it inserts an element for every loop.
    //
    // Therefore, if `amount >= length / 2` then inplace will be both faster and use less memory.
    //
    // TODO: there is probably even more fine-tuning that can be done here since
    // `HashMap::with_capacity(amount)` probably allocates more than `amount` in practice,
    // and a trade off could probably be made between memory/cpu, since hashmap operations
    // are slower than array index swapping.
    if amount >= length / 2 {
        sample_indices_inplace(rng, length, amount)
    } else {
        sample_indices_cache(rng, length, amount)
    }
}

/// Sample an amount of indices using an inplace partial fisher yates method.
///
/// This allocates the entire `length` of indices and randomizes only the first `amount`.
/// It then truncates to `amount` and returns.
///
/// This is better than using a HashMap "cache" when `amount >= length / 2` since it does not
/// require allocating an extra cache and is much faster.
fn sample_indices_inplace<R>(rng: &mut R, length: usize, amount: usize) -> Vec<usize>
    where R: Rng,
{
    debug_assert!(amount <= length);
    let mut indices: Vec<usize> = Vec::with_capacity(length);
    indices.extend(0..length);
    for i in 0..amount {
        let j: usize = rng.gen_range(i, length);
        let tmp = indices[i];
        indices[i] = indices[j];
        indices[j] = tmp;
    }
    indices.truncate(amount);
    debug_assert_eq!(indices.len(), amount);
    indices
}

/// This method performs a partial fisher-yates on a range of indices using a HashMap
/// as a cache to record potential collisions.
///
/// The cache avoids allocating the entire `length` of values. This is especially useful when
/// `amount <<< length`, i.e. select 3 non-repeating from 1_000_000
fn sample_indices_cache<R>(
    rng: &mut R,
    length: usize,
    amount: usize,
) -> Vec<usize>
    where R: Rng,
{
    debug_assert!(amount <= length);
    let mut cache = HashMap::with_capacity(amount);
    let mut out = Vec::with_capacity(amount);
    for i in 0..amount {
        let j: usize = rng.gen_range(i, length);

        // equiv: let tmp = slice[i];
        let tmp = match cache.get(&i) {
            Some(e) => *e,
            None => i,
        };

        // equiv: slice[i] = slice[j];
        let x = match cache.get(&j) {
            Some(x) => *x,
            None => j,
        };

        // equiv: slice[j] = tmp;
        cache.insert(j, tmp);

        // note that in the inplace version, slice[i] is automatically "returned" value
        out.push(x);
    }
    debug_assert_eq!(out.len(), amount);
    out
}

#[cfg(test)]
mod test {
    use super::*;
    use {thread_rng, XorShiftRng, SeedableRng};

    #[test]
    fn test_sample_iter() {
        let min_val = 1;
        let max_val = 100;

        let mut r = thread_rng();
        let vals = (min_val..max_val).collect::<Vec<i32>>();
        let small_sample = sample_iter(&mut r, vals.iter(), 5).unwrap();
        let large_sample = sample_iter(&mut r, vals.iter(), vals.len() + 5).unwrap_err();

        assert_eq!(small_sample.len(), 5);
        assert_eq!(large_sample.len(), vals.len());
        // no randomization happens when amount >= len
        assert_eq!(large_sample, vals.iter().collect::<Vec<_>>());

        assert!(small_sample.iter().all(|e| {
            **e >= min_val && **e <= max_val
        }));
    }
    #[test]
    fn test_sample_slice_boundaries() {
        let empty: &[u8] = &[];

        let mut r = thread_rng();

        // sample 0 items
        assert_eq!(sample_slice(&mut r, empty, 0), vec![]);
        assert_eq!(sample_slice(&mut r, &[42, 2, 42], 0), vec![]);

        // sample 1 item
        assert_eq!(sample_slice(&mut r, &[42], 1), vec![42]);
        let v = sample_slice(&mut r, &[1, 42], 1)[0];
        assert!(v == 1 || v == 42);

        // sample "all" the items
        let v = sample_slice(&mut r, &[42, 133], 2);
        assert!(v == vec![42, 133] || v == vec![133, 42]);

        assert_eq!(sample_indices_inplace(&mut r, 0, 0), vec![]);
        assert_eq!(sample_indices_inplace(&mut r, 1, 0), vec![]);
        assert_eq!(sample_indices_inplace(&mut r, 1, 1), vec![0]);

        assert_eq!(sample_indices_cache(&mut r, 0, 0), vec![]);
        assert_eq!(sample_indices_cache(&mut r, 1, 0), vec![]);
        assert_eq!(sample_indices_cache(&mut r, 1, 1), vec![0]);

        // Make sure lucky 777's aren't lucky
        let slice = &[42, 777];
        let mut num_42 = 0;
        let total = 1000;
        for _ in 0..total {
            let v = sample_slice(&mut r, slice, 1);
            assert_eq!(v.len(), 1);
            let v = v[0];
            assert!(v == 42 || v == 777);
            if v == 42 {
                num_42 += 1;
            }
        }
        let ratio_42 = num_42 as f64 / 1000 as f64;
        assert!(0.4 <= ratio_42 || ratio_42 <= 0.6, "{}", ratio_42);
    }

    #[test]
    fn test_sample_slice() {
        let xor_rng = XorShiftRng::from_seed;

        let max_range = 100;
        let mut r = thread_rng();

        for length in 1usize..max_range {
            let amount = r.gen_range(0, length);
            let seed: [u32; 4] = [
                r.next_u32(), r.next_u32(), r.next_u32(), r.next_u32()
            ];

            println!("Selecting indices: len={}, amount={}, seed={:?}", length, amount, seed);

            // assert that the two index methods give exactly the same result
            let inplace = sample_indices_inplace(
                &mut xor_rng(seed), length, amount);
            let cache = sample_indices_cache(
                &mut xor_rng(seed), length, amount);
            assert_eq!(inplace, cache);

            // assert the basics work
            let regular = sample_indices(
                &mut xor_rng(seed), length, amount);
            assert_eq!(regular.len(), amount);
            assert!(regular.iter().all(|e| *e < length));
            assert_eq!(regular, inplace);

            // also test that sampling the slice works
            let vec: Vec<usize> = (0..length).collect();
            {
                let result = sample_slice(&mut xor_rng(seed), &vec, amount);
                assert_eq!(result, regular);
            }

            {
                let result = sample_slice_ref(&mut xor_rng(seed), &vec, amount);
                let expected = regular.iter().map(|v| v).collect::<Vec<_>>();
                assert_eq!(result, expected);
            }
        }
    }
}
<commit_msg>clean up len docs<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Functions for randomly accessing and sampling sequences.

use super::Rng;
use std::collections::hash_map::HashMap;

/// Randomly sample `amount` elements from a finite iterator.
///
/// The following can be returned:
/// - `Ok`: `Vec` of `amount` non-repeating randomly sampled elements. The order is not random.
/// - `Err`: `Vec` of all the elements from `iterable` in sequential order. This happens when the
///   length of `iterable` was less than `amount`. This is considered an error since exactly
///   `amount` elements is typically expected.
///
/// This implementation uses `O(len(iterable))` time and `O(amount)` memory.
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let sample = seq::sample_iter(&mut rng, 1..100, 5).unwrap();
/// println!("{:?}", sample);
/// ```
pub fn sample_iter<T, I, R>(rng: &mut R, iterable: I, amount: usize) -> Result<Vec<T>, Vec<T>>
    where I: IntoIterator<Item=T>,
          R: Rng,
{
    let mut iter = iterable.into_iter();
    let mut reservoir = Vec::with_capacity(amount);
    reservoir.extend(iter.by_ref().take(amount));

    // Continue unless the iterator was exhausted
    //
    // note: this prevents iterators that "restart" from causing problems.
    // If the iterator stops once, then so do we.
    if reservoir.len() == amount {
        for (i, elem) in iter.enumerate() {
            let k = rng.gen_range(0, i + 1 + amount);
            if let Some(spot) = reservoir.get_mut(k) {
                *spot = elem;
            }
        }
        Ok(reservoir)
    } else {
        // Don't hang onto extra memory. There is a corner case where
        // `amount` was much less than `len(iterable)`.
        reservoir.shrink_to_fit();
        Err(reservoir)
    }
}

/// Randomly sample exactly `amount` values from `slice`.
///
/// The values are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// Panics if `amount > slice.len()`
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let values = vec![5, 6, 1, 3, 4, 6, 7];
/// println!("{:?}", seq::sample_slice(&mut rng, &values, 3));
/// ```
pub fn sample_slice<R, T>(rng: &mut R, slice: &[T], amount: usize) -> Vec<T>
    where R: Rng,
          T: Clone
{
    let indices = sample_indices(rng, slice.len(), amount);

    let mut out = Vec::with_capacity(amount);
    out.extend(indices.iter().map(|i| slice[*i].clone()));
    out
}

/// Randomly sample exactly `amount` references from `slice`.
///
/// The references are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// Panics if `amount > slice.len()`
///
/// # Example
///
/// ```rust
/// use rand::{thread_rng, seq};
///
/// let mut rng = thread_rng();
/// let values = vec![5, 6, 1, 3, 4, 6, 7];
/// println!("{:?}", seq::sample_slice_ref(&mut rng, &values, 3));
/// ```
pub fn sample_slice_ref<'a, R, T>(rng: &mut R, slice: &'a [T], amount: usize) -> Vec<&'a T>
    where R: Rng
{
    let indices = sample_indices(rng, slice.len(), amount);

    let mut out = Vec::with_capacity(amount);
    out.extend(indices.iter().map(|i| &slice[*i]));
    out
}

/// Randomly sample exactly `amount` indices from `0..length`.
///
/// The values are non-repeating and in random order.
///
/// This implementation uses `O(amount)` time and memory.
///
/// This method is used internally by the slice sampling methods, but it can sometimes be useful to
/// have the indices themselves so this is provided as an alternative.
///
/// Panics if `amount > length`
pub fn sample_indices<R>(rng: &mut R, length: usize, amount: usize) -> Vec<usize>
    where R: Rng,
{
    if amount > length {
        panic!("`amount` must be less than or equal to `slice.len()`");
    }

    // We are going to have to allocate at least `amount` for the output no matter what. However,
    // if we use the `cached` version we will have to allocate `amount` as a HashMap as well since
    // it inserts an element for every loop.
    //
    // Therefore, if `amount >= length / 2` then inplace will be both faster and use less memory.
    //
    // TODO: there is probably even more fine-tuning that can be done here since
    // `HashMap::with_capacity(amount)` probably allocates more than `amount` in practice,
    // and a trade off could probably be made between memory/cpu, since hashmap operations
    // are slower than array index swapping.
    if amount >= length / 2 {
        sample_indices_inplace(rng, length, amount)
    } else {
        sample_indices_cache(rng, length, amount)
    }
}

/// Sample an amount of indices using an inplace partial fisher yates method.
///
/// This allocates the entire `length` of indices and randomizes only the first `amount`.
/// It then truncates to `amount` and returns.
///
/// This is better than using a HashMap "cache" when `amount >= length / 2` since it does not
/// require allocating an extra cache and is much faster.
fn sample_indices_inplace<R>(rng: &mut R, length: usize, amount: usize) -> Vec<usize>
    where R: Rng,
{
    debug_assert!(amount <= length);
    let mut indices: Vec<usize> = Vec::with_capacity(length);
    indices.extend(0..length);
    for i in 0..amount {
        let j: usize = rng.gen_range(i, length);
        let tmp = indices[i];
        indices[i] = indices[j];
        indices[j] = tmp;
    }
    indices.truncate(amount);
    debug_assert_eq!(indices.len(), amount);
    indices
}

/// This method performs a partial fisher-yates on a range of indices using a HashMap
/// as a cache to record potential collisions.
///
/// The cache avoids allocating the entire `length` of values. This is especially useful when
/// `amount <<< length`, i.e. select 3 non-repeating from 1_000_000
fn sample_indices_cache<R>(
    rng: &mut R,
    length: usize,
    amount: usize,
) -> Vec<usize>
    where R: Rng,
{
    debug_assert!(amount <= length);
    let mut cache = HashMap::with_capacity(amount);
    let mut out = Vec::with_capacity(amount);
    for i in 0..amount {
        let j: usize = rng.gen_range(i, length);

        // equiv: let tmp = slice[i];
        let tmp = match cache.get(&i) {
            Some(e) => *e,
            None => i,
        };

        // equiv: slice[i] = slice[j];
        let x = match cache.get(&j) {
            Some(x) => *x,
            None => j,
        };

        // equiv: slice[j] = tmp;
        cache.insert(j, tmp);

        // note that in the inplace version, slice[i] is automatically "returned" value
        out.push(x);
    }
    debug_assert_eq!(out.len(), amount);
    out
}

#[cfg(test)]
mod test {
    use super::*;
    use {thread_rng, XorShiftRng, SeedableRng};

    #[test]
    fn test_sample_iter() {
        let min_val = 1;
        let max_val = 100;

        let mut r = thread_rng();
        let vals = (min_val..max_val).collect::<Vec<i32>>();
        let small_sample = sample_iter(&mut r, vals.iter(), 5).unwrap();
        let large_sample = sample_iter(&mut r, vals.iter(), vals.len() + 5).unwrap_err();

        assert_eq!(small_sample.len(), 5);
        assert_eq!(large_sample.len(), vals.len());
        // no randomization happens when amount >= len
        assert_eq!(large_sample, vals.iter().collect::<Vec<_>>());

        assert!(small_sample.iter().all(|e| {
            **e >= min_val && **e <= max_val
        }));
    }
    #[test]
    fn test_sample_slice_boundaries() {
        let empty: &[u8] = &[];

        let mut r = thread_rng();

        // sample 0 items
        assert_eq!(sample_slice(&mut r, empty, 0), vec![]);
        assert_eq!(sample_slice(&mut r, &[42, 2, 42], 0), vec![]);

        // sample 1 item
        assert_eq!(sample_slice(&mut r, &[42], 1), vec![42]);
        let v = sample_slice(&mut r, &[1, 42], 1)[0];
        assert!(v == 1 || v == 42);

        // sample "all" the items
        let v = sample_slice(&mut r, &[42, 133], 2);
        assert!(v == vec![42, 133] || v == vec![133, 42]);

        assert_eq!(sample_indices_inplace(&mut r, 0, 0), vec![]);
        assert_eq!(sample_indices_inplace(&mut r, 1, 0), vec![]);
        assert_eq!(sample_indices_inplace(&mut r, 1, 1), vec![0]);

        assert_eq!(sample_indices_cache(&mut r, 0, 0), vec![]);
        assert_eq!(sample_indices_cache(&mut r, 1, 0), vec![]);
        assert_eq!(sample_indices_cache(&mut r, 1, 1), vec![0]);

        // Make sure lucky 777's aren't lucky
        let slice = &[42, 777];
        let mut num_42 = 0;
        let total = 1000;
        for _ in 0..total {
            let v = sample_slice(&mut r, slice, 1);
            assert_eq!(v.len(), 1);
            let v = v[0];
            assert!(v == 42 || v == 777);
            if v == 42 {
                num_42 += 1;
            }
        }
        let ratio_42 = num_42 as f64 / 1000 as f64;
        assert!(0.4 <= ratio_42 || ratio_42 <= 0.6, "{}", ratio_42);
    }

    #[test]
    fn test_sample_slice() {
        let xor_rng = XorShiftRng::from_seed;

        let max_range = 100;
        let mut r = thread_rng();

        for length in 1usize..max_range {
            let amount = r.gen_range(0, length);
            let seed: [u32; 4] = [
                r.next_u32(), r.next_u32(), r.next_u32(), r.next_u32()
            ];

            println!("Selecting indices: len={}, amount={}, seed={:?}", length, amount, seed);

            // assert that the two index methods give exactly the same result
            let inplace = sample_indices_inplace(
                &mut xor_rng(seed), length, amount);
            let cache = sample_indices_cache(
                &mut xor_rng(seed), length, amount);
            assert_eq!(inplace, cache);

            // assert the basics work
            let regular = sample_indices(
                &mut xor_rng(seed), length, amount);
            assert_eq!(regular.len(), amount);
            assert!(regular.iter().all(|e| *e < length));
            assert_eq!(regular, inplace);

            // also test that sampling the slice works
            let vec: Vec<usize> = (0..length).collect();
            {
                let result = sample_slice(&mut xor_rng(seed), &vec, amount);
                assert_eq!(result, regular);
            }

            {
                let result = sample_slice_ref(&mut xor_rng(seed), &vec, amount);
                let expected = regular.iter().map(|v| v).collect::<Vec<_>>();
                assert_eq!(result, expected);
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::num::NonZeroU8;
use std::{fmt, ops};

/// Interests used in registering.
///
/// Interests are used in [registering] [`event::Source`]s with [`Poll`], they
/// indicate what readiness should be monitored for. For example if a socket is
/// registered with [readable] interests and the socket becomes writable, no
/// event will be returned from a call to [`poll`].
///
/// The size of `Option<Interests>` should be identical to itself.
///
/// ```
/// use std::mem::size_of;
/// use mio::Interests;
///
/// assert_eq!(size_of::<Option<Interests>>(), size_of::<Interests>());
/// ```
///
/// [registering]: crate::Registry::register
/// [`event::Source`]: crate::event::Source
/// [`Poll`]: crate::Poll
/// [readable]: Interests::READABLE
/// [`poll`]: crate::Poll::poll
#[derive(Copy, PartialEq, Eq, Clone, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Interests(NonZeroU8);

// These must be unique.
const READABLE: u8 = 0b0_001;
const WRITABLE: u8 = 0b0_010;
// The following are not available on all platforms.
#[cfg_attr(
    not(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos"
    )),
    allow(dead_code)
)]
const AIO: u8 = 0b0_100;
#[cfg_attr(not(target_os = "freebsd"), allow(dead_code))]
const LIO: u8 = 0b1_000;

impl Interests {
    /// Returns a `Interests` set representing readable interests.
    pub const READABLE: Interests = Interests(unsafe { NonZeroU8::new_unchecked(READABLE) });

    /// Returns a `Interests` set representing writable interests.
    pub const WRITABLE: Interests = Interests(unsafe { NonZeroU8::new_unchecked(WRITABLE) });

    /// Returns a `Interests` set representing AIO completion interests.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos"
    ))]
    pub const AIO: Interests = Interests(unsafe { NonZeroU8::new_unchecked(AIO) });

    /// Returns a `Interests` set representing LIO completion interests.
    #[cfg(target_os = "freebsd")]
    pub const LIO: Interests = Interests(unsafe { NonZeroU8::new_unchecked(LIO) });

    /// Returns true if the value includes readable readiness.
    pub fn is_readable(&self) -> bool {
        (self.0.get() & READABLE) != 0
    }

    /// Returns true if the value includes writable readiness.
    pub fn is_writable(&self) -> bool {
        (self.0.get() & WRITABLE) != 0
    }

    /// Returns true if `Interests` contains AIO readiness
    pub fn is_aio(&self) -> bool {
        (self.0.get() & AIO) != 0
    }

    /// Returns true if `Interests` contains LIO readiness
    pub fn is_lio(&self) -> bool {
        (self.0.get() & LIO) != 0
    }

    #[cfg(windows)]
    pub(crate) fn as_u8(self) -> u8 {
        self.0.get()
    }
}

impl ops::BitOr for Interests {
    type Output = Self;

    #[inline]
    fn bitor(self, other: Self) -> Self {
        Interests(unsafe { NonZeroU8::new_unchecked(self.0.get() | other.0.get()) })
    }
}

impl ops::BitOrAssign for Interests {
    #[inline]
    fn bitor_assign(&mut self, other: Self) {
        self.0 = (*self | other).0;
    }
}

impl fmt::Debug for Interests {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut one = false;
        if self.is_readable() {
            if one {
                write!(fmt, " | ")?
            }
            write!(fmt, "READABLE")?;
            one = true
        }
        if self.is_writable() {
            if one {
                write!(fmt, " | ")?
            }
            write!(fmt, "WRITABLE")?;
            one = true
        }
        #[cfg(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        ))]
        {
            if self.is_aio() {
                if one {
                    write!(fmt, " | ")?
                }
                write!(fmt, "AIO")?;
                one = true
            }
        }
        #[cfg(any(target_os = "freebsd"))]
        {
            if self.is_lio() {
                if one {
                    write!(fmt, " | ")?
                }
                write!(fmt, "LIO")?;
                one = true
            }
        }
        debug_assert!(one, "printing empty interests");
        Ok(())
    }
}
<commit_msg>Prefer pass-by-value over pass-by-reference for public methods<commit_after>use std::num::NonZeroU8;
use std::{fmt, ops};

/// Interests used in registering.
///
/// Interests are used in [registering] [`event::Source`]s with [`Poll`], they
/// indicate what readiness should be monitored for. For example if a socket is
/// registered with [readable] interests and the socket becomes writable, no
/// event will be returned from a call to [`poll`].
///
/// The size of `Option<Interests>` should be identical to itself.
///
/// ```
/// use std::mem::size_of;
/// use mio::Interests;
///
/// assert_eq!(size_of::<Option<Interests>>(), size_of::<Interests>());
/// ```
///
/// [registering]: crate::Registry::register
/// [`event::Source`]: crate::event::Source
/// [`Poll`]: crate::Poll
/// [readable]: Interests::READABLE
/// [`poll`]: crate::Poll::poll
#[derive(Copy, PartialEq, Eq, Clone, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Interests(NonZeroU8);

// These must be unique.
const READABLE: u8 = 0b0_001;
const WRITABLE: u8 = 0b0_010;
// The following are not available on all platforms.
#[cfg_attr(
    not(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos"
    )),
    allow(dead_code)
)]
const AIO: u8 = 0b0_100;
#[cfg_attr(not(target_os = "freebsd"), allow(dead_code))]
const LIO: u8 = 0b1_000;

impl Interests {
    /// Returns a `Interests` set representing readable interests.
    pub const READABLE: Interests = Interests(unsafe { NonZeroU8::new_unchecked(READABLE) });

    /// Returns a `Interests` set representing writable interests.
    pub const WRITABLE: Interests = Interests(unsafe { NonZeroU8::new_unchecked(WRITABLE) });

    /// Returns a `Interests` set representing AIO completion interests.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos"
    ))]
    pub const AIO: Interests = Interests(unsafe { NonZeroU8::new_unchecked(AIO) });

    /// Returns a `Interests` set representing LIO completion interests.
    #[cfg(target_os = "freebsd")]
    pub const LIO: Interests = Interests(unsafe { NonZeroU8::new_unchecked(LIO) });

    /// Returns true if the value includes readable readiness.
    pub fn is_readable(self) -> bool {
        (self.0.get() & READABLE) != 0
    }

    /// Returns true if the value includes writable readiness.
    pub fn is_writable(self) -> bool {
        (self.0.get() & WRITABLE) != 0
    }

    /// Returns true if `Interests` contains AIO readiness
    pub fn is_aio(self) -> bool {
        (self.0.get() & AIO) != 0
    }

    /// Returns true if `Interests` contains LIO readiness
    pub fn is_lio(self) -> bool {
        (self.0.get() & LIO) != 0
    }

    #[cfg(windows)]
    pub(crate) fn as_u8(self) -> u8 {
        self.0.get()
    }
}

impl ops::BitOr for Interests {
    type Output = Self;

    #[inline]
    fn bitor(self, other: Self) -> Self {
        Interests(unsafe { NonZeroU8::new_unchecked(self.0.get() | other.0.get()) })
    }
}

impl ops::BitOrAssign for Interests {
    #[inline]
    fn bitor_assign(&mut self, other: Self) {
        self.0 = (*self | other).0;
    }
}

impl fmt::Debug for Interests {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut one = false;
        if self.is_readable() {
            if one {
                write!(fmt, " | ")?
            }
            write!(fmt, "READABLE")?;
            one = true
        }
        if self.is_writable() {
            if one {
                write!(fmt, " | ")?
            }
            write!(fmt, "WRITABLE")?;
            one = true
        }
        #[cfg(any(
            target_os = "dragonfly",
            target_os = "freebsd",
            target_os = "ios",
            target_os = "macos"
        ))]
        {
            if self.is_aio() {
                if one {
                    write!(fmt, " | ")?
                }
                write!(fmt, "AIO")?;
                one = true
            }
        }
        #[cfg(any(target_os = "freebsd"))]
        {
            if self.is_lio() {
                if one {
                    write!(fmt, " | ")?
                }
                write!(fmt, "LIO")?;
                one = true
            }
        }
        debug_assert!(one, "printing empty interests");
        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::frame::Ax25Frame;
use crate::kiss;
use crate::linux;
use std::collections::VecDeque;
use std::str::FromStr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use thiserror::Error;

/// Errors that can occur when interacting with a `Tnc`.
#[derive(Debug, Error)]
pub enum TncError {
    #[error("Unable to connect to TNC: {}", source)]
    OpenTnc { source: std::io::Error },
    #[error("Interface with specified callsign '{}' does not exist", callsign)]
    InterfaceNotFound { callsign: String },
    #[error("Unable to send frame: {}", source)]
    SendFrame { source: std::io::Error },
    #[error("Unable to receive frame: {}", source)]
    ReceiveFrame { source: std::io::Error },
    #[error("Unable to make configuration change: {}", source)]
    ConfigFailed { source: std::io::Error },
}

/// Errors that can occur when parsing a `TncAddress` from a string.
#[derive(Debug, Error, PartialEq)]
pub enum ParseError {
    #[error("TNC address '{}' is invalid - it should begin with 'tnc:'", string)]
    NoTncPrefix { string: String },
    #[error("Unknown TNC type {}", tnc_type)]
    UnknownType { tnc_type: String },
    #[error(
        "TNC type '{}' expects {} parameters to follow but there are {}",
        tnc_type,
        expected,
        actual
    )]
    WrongParameterCount {
        tnc_type: String,
        expected: usize,
        actual: usize,
    },
    #[error("Supplied port '{}' should be a number from 0 to 65535", input)]
    InvalidPort {
        input: String,
        source: std::num::ParseIntError,
    },
}

/// Configuration details for a TCP KISS TNC. This structure can be created directly
/// or indirectly by parsing a string into a `TncAddress`.
#[derive(PartialEq, Debug)]
pub struct TcpKissConfig {
    /// Hostname or IP address of the computer with the TNC
    pub host: String,
    /// Port number
    pub port: u16,
}

/// Configuration details for a TNC attached as a Linux network interface using
/// `kissattach`. This structure can be created directly or indirectly by parsing
/// a string into a `TncAddress`.
#[derive(PartialEq, Debug)]
pub struct LinuxIfConfig {
    /// The hardware address associated with the interface, e.g. "VK7NTK-2"
    pub callsign: String,
}

#[derive(PartialEq, Debug)]
pub(crate) enum ConnectConfig {
    TcpKiss(TcpKissConfig),
    LinuxIf(LinuxIfConfig),
}

/// A parsed TNC address that can be used to open a `Tnc`.
#[derive(PartialEq, Debug)]
pub struct TncAddress {
    pub(crate) config: ConnectConfig,
}

impl TncAddress {
    /// Programmatically create a `TncAddress` pointing to a Linux network interface.
    pub fn new_linuxif(linuxif: LinuxIfConfig) -> Self {
        TncAddress {
            config: ConnectConfig::LinuxIf(linuxif),
        }
    }

    /// Porgrammatically create a `TncAddress` pointing to a KISS TCP service.
    pub fn new_tcpkiss(tcpkiss: TcpKissConfig) -> Self {
        TncAddress {
            config: ConnectConfig::TcpKiss(tcpkiss),
        }
    }
}

impl FromStr for TncAddress {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if !s.starts_with("tnc:") {
            return Err(ParseError::NoTncPrefix {
                string: s.to_string(),
            });
        }
        let components: Vec<&str> = s.split(':').collect();
        let len = components.len();
        Ok(match components[1] {
            "tcpkiss" => {
                if len != 4 {
                    return Err(ParseError::WrongParameterCount {
                        tnc_type: components[1].to_string(),
                        expected: 2usize,
                        actual: len - 2,
                    });
                }
                TncAddress {
                    config: ConnectConfig::TcpKiss(TcpKissConfig {
                        host: components[2].to_string(),
                        port: components[3].parse().map_err(|e| ParseError::InvalidPort {
                            input: components[3].to_string(),
                            source: e,
                        })?,
                    }),
                }
            }
            "linuxif" => {
                if len != 3 {
                    return Err(ParseError::WrongParameterCount {
                        tnc_type: components[1].to_string(),
                        expected: 1usize,
                        actual: len - 2,
                    });
                }
                TncAddress {
                    config: ConnectConfig::LinuxIf(LinuxIfConfig {
                        callsign: components[2].to_string(),
                    }),
                }
            }
            unknown => {
                return Err(ParseError::UnknownType {
                    tnc_type: unknown.to_string(),
                })
            }
        })
    }
}

trait TncImpl: Send + Sync {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError>;
    fn receive_frame(&self) -> Result<Ax25Frame, TncError>;
    fn clone(&self) -> Box<dyn TncImpl>;
}

/// A local or remote TNC attached to a radio, which can send and receive frames.
pub struct Tnc {
    imp: Box<dyn TncImpl>,
    senders: Arc<Mutex<Vec<Sender<Ax25Frame>>>>,
    buffer: Arc<Mutex<VecDeque<Ax25Frame>>>,
}

impl Tnc {
    /// Attempt to obtain a `Tnc` connection using the provided address.
    pub fn open(address: &TncAddress) -> Result<Self, TncError> {
        let imp: Box<dyn TncImpl> = match &address.config {
            ConnectConfig::TcpKiss(config) => Box::new(TcpKissTnc::open(&config)?),
            ConnectConfig::LinuxIf(config) => Box::new(LinuxIfTnc::open(&config)?),
        };
        Ok(Tnc::new(imp))
    }

    fn new(imp: Box<dyn TncImpl>) -> Self {
        let tnc = Tnc {
            imp,
            senders: Arc::new(Mutex::new(Vec::new())),
            buffer: Arc::new(Mutex::new(VecDeque::new())),
        };
        {
            let tnc = tnc.clone();
            // TODO how to allow this thread to exit?
            thread::spawn(move || {
                loop {
                    // TODO what to do if this fails?
                    if let Ok(x) = tnc.imp.receive_frame() {
                        tnc.buffer.lock().unwrap().push_back(x.clone());
                        tnc.senders.lock().unwrap().retain(|s| {
                            let x = (&x).clone();
                            // If there's an error, remove sender from vec
                            s.send(x).is_ok()
                        });
                    }
                }
            });
        }
        tnc
    }

    /// Transmit a frame on the radio. Transmission is not guaranteed even if a
    /// `Ok` result is returned.
    pub fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.imp.send_frame(frame)
    }

    /// Block to receive a frame from the radio. If you want to do this on
    /// a separate thread from sending, clone the `Tnc`.
    pub fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        self.imp.receive_frame()
    }

    /// Create a new `Receiver<Ax25Frame>`
    /// This will receive a copy of all incoming frames.
    pub fn incoming(&mut self) -> Receiver<Ax25Frame> {
        let (sender, receiver) = channel();
        self.senders.lock().unwrap().push(sender);
        receiver
    }
}

impl Clone for Tnc {
    fn clone(&self) -> Self {
        Tnc {
            imp: self.imp.clone(),
            senders: self.senders.clone(),
            buffer: self.buffer.clone(),
        }
    }
}

struct LinuxIfTnc {
    socket: Arc<linux::Ax25RawSocket>,
    ifindex: i32,
}

impl LinuxIfTnc {
    fn open(config: &LinuxIfConfig) -> Result<Self, TncError> {
        let socket = linux::Ax25RawSocket::new().map_err(|e| TncError::OpenTnc { source: e })?;
        let ifindex = match socket
            .list_ax25_interfaces()
            .map_err(|e| TncError::OpenTnc { source: e })?
            .iter()
            .find(|nd| nd.name.to_uppercase() == config.callsign.to_uppercase())
        {
            Some(nd) => nd.ifindex,
            None => {
                return Err(TncError::InterfaceNotFound {
                    callsign: config.callsign.clone(),
                })
            }
        };
        Ok(Self {
            socket: Arc::new(socket),
            ifindex,
        })
    }
}

impl TncImpl for LinuxIfTnc {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.socket
            .send_frame(&frame.to_bytes(), self.ifindex)
            .map_err(|e| TncError::SendFrame { source: e })
    }

    fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        loop {
            let bytes = self
                .socket
                .receive_frame(self.ifindex)
                .map_err(|e| TncError::ReceiveFrame { source: e })?;
            if let Ok(parsed) = Ax25Frame::from_bytes(&bytes) {
                return Ok(parsed);
            }
        }
    }

    fn clone(&self) -> Box<dyn TncImpl> {
        Box::new(LinuxIfTnc {
            socket: self.socket.clone(),
            ifindex: self.ifindex,
        })
    }
}

struct TcpKissTnc {
    iface: Arc<kiss::TcpKissInterface>,
}

impl TcpKissTnc {
    fn open(config: &TcpKissConfig) -> Result<Self, TncError> {
        Ok(Self {
            iface: Arc::new(
                kiss::TcpKissInterface::new(format!("{}:{}", config.host, config.port))
                    .map_err(|e| TncError::OpenTnc { source: e })?,
            ),
        })
    }
}

impl TncImpl for TcpKissTnc {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.iface
            .send_frame(&frame.to_bytes())
            .map_err(|e| TncError::SendFrame { source: e })
    }

    fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        loop {
            let bytes = self
                .iface
                .receive_frame()
                .map_err(|e| TncError::ReceiveFrame { source: e })?;
            if let Ok(parsed) = Ax25Frame::from_bytes(&bytes) {
                return Ok(parsed);
            }
        }
    }

    fn clone(&self) -> Box<dyn TncImpl> {
        Box::new(TcpKissTnc {
            iface: self.iface.clone(),
        })
    }
}

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

    #[test]
    fn parse_tnc_addresses() {
        assert_eq!(
            "tnc:tcpkiss:192.168.0.1:8001".parse::<TncAddress>(),
            Ok(TncAddress {
                config: ConnectConfig::TcpKiss(TcpKissConfig {
                    host: "192.168.0.1".to_string(),
                    port: 8001_u16,
                })
            })
        );
        assert_eq!(
            "tnc:linuxif:VK7NTK-2".parse::<TncAddress>(),
            Ok(TncAddress {
                config: ConnectConfig::LinuxIf(LinuxIfConfig {
                    callsign: "VK7NTK-2".to_string(),
                })
            })
        );
        assert!(match "fish".parse::<TncAddress>() {
            Err(ParseError::NoTncPrefix { .. }) => true,
            _ => false,
        });
        assert!(match "tnc:".parse::<TncAddress>() {
            Err(ParseError::UnknownType { tnc_type }) => tnc_type == "",
            _ => false,
        });
        assert!(match "tnc:fish".parse::<TncAddress>() {
            Err(ParseError::UnknownType { tnc_type }) => tnc_type == "fish",
            _ => false,
        });
        assert!(match "tnc:tcpkiss".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 0
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 1
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:a:b:c".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 3
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:192.168.0.1".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 1
            }
            _ => false,
        });
        assert!(
            match "tnc:tcpkiss:192.168.0.1:hello".parse::<TncAddress>() {
                Err(ParseError::InvalidPort { input, .. }) => input == "hello",
                _ => false,
            }
        );
    }
}
<commit_msg>Code improvement<commit_after>use crate::frame::Ax25Frame;
use crate::kiss;
use crate::linux;
use std::collections::VecDeque;
use std::str::FromStr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
use thiserror::Error;

/// Errors that can occur when interacting with a `Tnc`.
#[derive(Debug, Error)]
pub enum TncError {
    #[error("Unable to connect to TNC: {}", source)]
    OpenTnc { source: std::io::Error },
    #[error("Interface with specified callsign '{}' does not exist", callsign)]
    InterfaceNotFound { callsign: String },
    #[error("Unable to send frame: {}", source)]
    SendFrame { source: std::io::Error },
    #[error("Unable to receive frame: {}", source)]
    ReceiveFrame { source: std::io::Error },
    #[error("Unable to make configuration change: {}", source)]
    ConfigFailed { source: std::io::Error },
}

/// Errors that can occur when parsing a `TncAddress` from a string.
#[derive(Debug, Error, PartialEq)]
pub enum ParseError {
    #[error("TNC address '{}' is invalid - it should begin with 'tnc:'", string)]
    NoTncPrefix { string: String },
    #[error("Unknown TNC type {}", tnc_type)]
    UnknownType { tnc_type: String },
    #[error(
        "TNC type '{}' expects {} parameters to follow but there are {}",
        tnc_type,
        expected,
        actual
    )]
    WrongParameterCount {
        tnc_type: String,
        expected: usize,
        actual: usize,
    },
    #[error("Supplied port '{}' should be a number from 0 to 65535", input)]
    InvalidPort {
        input: String,
        source: std::num::ParseIntError,
    },
}

/// Configuration details for a TCP KISS TNC. This structure can be created directly
/// or indirectly by parsing a string into a `TncAddress`.
#[derive(PartialEq, Debug)]
pub struct TcpKissConfig {
    /// Hostname or IP address of the computer with the TNC
    pub host: String,
    /// Port number
    pub port: u16,
}

/// Configuration details for a TNC attached as a Linux network interface using
/// `kissattach`. This structure can be created directly or indirectly by parsing
/// a string into a `TncAddress`.
#[derive(PartialEq, Debug)]
pub struct LinuxIfConfig {
    /// The hardware address associated with the interface, e.g. "VK7NTK-2"
    pub callsign: String,
}

#[derive(PartialEq, Debug)]
pub(crate) enum ConnectConfig {
    TcpKiss(TcpKissConfig),
    LinuxIf(LinuxIfConfig),
}

/// A parsed TNC address that can be used to open a `Tnc`.
#[derive(PartialEq, Debug)]
pub struct TncAddress {
    pub(crate) config: ConnectConfig,
}

impl TncAddress {
    /// Programmatically create a `TncAddress` pointing to a Linux network interface.
    pub fn new_linuxif(linuxif: LinuxIfConfig) -> Self {
        TncAddress {
            config: ConnectConfig::LinuxIf(linuxif),
        }
    }

    /// Porgrammatically create a `TncAddress` pointing to a KISS TCP service.
    pub fn new_tcpkiss(tcpkiss: TcpKissConfig) -> Self {
        TncAddress {
            config: ConnectConfig::TcpKiss(tcpkiss),
        }
    }
}

impl FromStr for TncAddress {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if !s.starts_with("tnc:") {
            return Err(ParseError::NoTncPrefix {
                string: s.to_string(),
            });
        }
        let components: Vec<&str> = s.split(':').collect();
        let len = components.len();
        Ok(match components[1] {
            "tcpkiss" => {
                if len != 4 {
                    return Err(ParseError::WrongParameterCount {
                        tnc_type: components[1].to_string(),
                        expected: 2usize,
                        actual: len - 2,
                    });
                }
                TncAddress {
                    config: ConnectConfig::TcpKiss(TcpKissConfig {
                        host: components[2].to_string(),
                        port: components[3].parse().map_err(|e| ParseError::InvalidPort {
                            input: components[3].to_string(),
                            source: e,
                        })?,
                    }),
                }
            }
            "linuxif" => {
                if len != 3 {
                    return Err(ParseError::WrongParameterCount {
                        tnc_type: components[1].to_string(),
                        expected: 1usize,
                        actual: len - 2,
                    });
                }
                TncAddress {
                    config: ConnectConfig::LinuxIf(LinuxIfConfig {
                        callsign: components[2].to_string(),
                    }),
                }
            }
            unknown => {
                return Err(ParseError::UnknownType {
                    tnc_type: unknown.to_string(),
                })
            }
        })
    }
}

trait TncImpl: Send + Sync {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError>;
    fn receive_frame(&self) -> Result<Ax25Frame, TncError>;
    fn clone(&self) -> Box<dyn TncImpl>;
}

/// A local or remote TNC attached to a radio, which can send and receive frames.
pub struct Tnc {
    imp: Box<dyn TncImpl>,
    senders: Arc<Mutex<Vec<Sender<Ax25Frame>>>>,
    buffer: Arc<Mutex<VecDeque<Ax25Frame>>>,
}

impl Tnc {
    /// Attempt to obtain a `Tnc` connection using the provided address.
    pub fn open(address: &TncAddress) -> Result<Self, TncError> {
        let imp: Box<dyn TncImpl> = match &address.config {
            ConnectConfig::TcpKiss(config) => Box::new(TcpKissTnc::open(&config)?),
            ConnectConfig::LinuxIf(config) => Box::new(LinuxIfTnc::open(&config)?),
        };
        Ok(Tnc::new(imp))
    }

    fn new(imp: Box<dyn TncImpl>) -> Self {
        let tnc = Tnc {
            imp,
            senders: Arc::new(Mutex::new(Vec::new())),
            buffer: Arc::new(Mutex::new(VecDeque::new())),
        };
        {
            let tnc = tnc.clone();
            // TODO how to allow this thread to exit?
            thread::spawn(move || {
                loop {
                    // TODO what to do if this fails?
                    if let Ok(x) = tnc.imp.receive_frame() {
                        tnc.buffer.lock().unwrap().push_back(x.clone());
                        tnc.senders.lock().unwrap().retain(|s| {
                            // If there's an error, remove sender from vec
                            s.send(x.clone()).is_ok()
                        });
                    }
                }
            });
        }
        tnc
    }

    /// Transmit a frame on the radio. Transmission is not guaranteed even if a
    /// `Ok` result is returned.
    pub fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.imp.send_frame(frame)
    }

    /// Block to receive a frame from the radio. If you want to do this on
    /// a separate thread from sending, clone the `Tnc`.
    pub fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        self.imp.receive_frame()
    }

    /// Create a new `Receiver<Ax25Frame>`
    /// This will receive a copy of all incoming frames.
    pub fn incoming(&mut self) -> Receiver<Ax25Frame> {
        let (sender, receiver) = channel();
        self.senders.lock().unwrap().push(sender);
        receiver
    }
}

impl Clone for Tnc {
    fn clone(&self) -> Self {
        Tnc {
            imp: self.imp.clone(),
            senders: self.senders.clone(),
            buffer: self.buffer.clone(),
        }
    }
}

struct LinuxIfTnc {
    socket: Arc<linux::Ax25RawSocket>,
    ifindex: i32,
}

impl LinuxIfTnc {
    fn open(config: &LinuxIfConfig) -> Result<Self, TncError> {
        let socket = linux::Ax25RawSocket::new().map_err(|e| TncError::OpenTnc { source: e })?;
        let ifindex = match socket
            .list_ax25_interfaces()
            .map_err(|e| TncError::OpenTnc { source: e })?
            .iter()
            .find(|nd| nd.name.to_uppercase() == config.callsign.to_uppercase())
        {
            Some(nd) => nd.ifindex,
            None => {
                return Err(TncError::InterfaceNotFound {
                    callsign: config.callsign.clone(),
                })
            }
        };
        Ok(Self {
            socket: Arc::new(socket),
            ifindex,
        })
    }
}

impl TncImpl for LinuxIfTnc {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.socket
            .send_frame(&frame.to_bytes(), self.ifindex)
            .map_err(|e| TncError::SendFrame { source: e })
    }

    fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        loop {
            let bytes = self
                .socket
                .receive_frame(self.ifindex)
                .map_err(|e| TncError::ReceiveFrame { source: e })?;
            if let Ok(parsed) = Ax25Frame::from_bytes(&bytes) {
                return Ok(parsed);
            }
        }
    }

    fn clone(&self) -> Box<dyn TncImpl> {
        Box::new(LinuxIfTnc {
            socket: self.socket.clone(),
            ifindex: self.ifindex,
        })
    }
}

struct TcpKissTnc {
    iface: Arc<kiss::TcpKissInterface>,
}

impl TcpKissTnc {
    fn open(config: &TcpKissConfig) -> Result<Self, TncError> {
        Ok(Self {
            iface: Arc::new(
                kiss::TcpKissInterface::new(format!("{}:{}", config.host, config.port))
                    .map_err(|e| TncError::OpenTnc { source: e })?,
            ),
        })
    }
}

impl TncImpl for TcpKissTnc {
    fn send_frame(&self, frame: &Ax25Frame) -> Result<(), TncError> {
        self.iface
            .send_frame(&frame.to_bytes())
            .map_err(|e| TncError::SendFrame { source: e })
    }

    fn receive_frame(&self) -> Result<Ax25Frame, TncError> {
        loop {
            let bytes = self
                .iface
                .receive_frame()
                .map_err(|e| TncError::ReceiveFrame { source: e })?;
            if let Ok(parsed) = Ax25Frame::from_bytes(&bytes) {
                return Ok(parsed);
            }
        }
    }

    fn clone(&self) -> Box<dyn TncImpl> {
        Box::new(TcpKissTnc {
            iface: self.iface.clone(),
        })
    }
}

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

    #[test]
    fn parse_tnc_addresses() {
        assert_eq!(
            "tnc:tcpkiss:192.168.0.1:8001".parse::<TncAddress>(),
            Ok(TncAddress {
                config: ConnectConfig::TcpKiss(TcpKissConfig {
                    host: "192.168.0.1".to_string(),
                    port: 8001_u16,
                })
            })
        );
        assert_eq!(
            "tnc:linuxif:VK7NTK-2".parse::<TncAddress>(),
            Ok(TncAddress {
                config: ConnectConfig::LinuxIf(LinuxIfConfig {
                    callsign: "VK7NTK-2".to_string(),
                })
            })
        );
        assert!(match "fish".parse::<TncAddress>() {
            Err(ParseError::NoTncPrefix { .. }) => true,
            _ => false,
        });
        assert!(match "tnc:".parse::<TncAddress>() {
            Err(ParseError::UnknownType { tnc_type }) => tnc_type == "",
            _ => false,
        });
        assert!(match "tnc:fish".parse::<TncAddress>() {
            Err(ParseError::UnknownType { tnc_type }) => tnc_type == "fish",
            _ => false,
        });
        assert!(match "tnc:tcpkiss".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 0
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 1
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:a:b:c".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 3
            }
            _ => false,
        });
        assert!(match "tnc:tcpkiss:192.168.0.1".parse::<TncAddress>() {
            Err(ParseError::WrongParameterCount {
                tnc_type,
                expected,
                actual,
            }) => {
                tnc_type == "tcpkiss" && expected == 2 && actual == 1
            }
            _ => false,
        });
        assert!(
            match "tnc:tcpkiss:192.168.0.1:hello".parse::<TncAddress>() {
                Err(ParseError::InvalidPort { input, .. }) => input == "hello",
                _ => false,
            }
        );
    }
}
<|endoftext|>


<commit_before>use check_gl;
use gl;
use gl::types::{GLint, GLuint, GLenum, GLsizeiptr};
use libc;
use libc::c_void;
use std::mem;

type IndexType = u16;

#[allow(raw_pointer_deriving)]
#[deriving(Clone)]
struct VertexAttribute {
    layout: GLuint,
    gl_type: GLenum,
    size: GLint,
    normalized: u8,
    offset: *const c_void
}

pub struct BufferBuilder<VertexType: Copy> {
    attributes: Vec<VertexAttribute>,
    used_layouts: Vec<bool>,
    vertex_size: uint
}
impl<VertexType: Copy>  BufferBuilder<VertexType> {
    pub fn new(capacity: uint) -> BufferBuilder<VertexType> {
        BufferBuilder {
            attributes: Vec::with_capacity(capacity),
            used_layouts: Vec::with_capacity(capacity),
            vertex_size: 0,
        }
    }

    pub fn max_attribute_size_left(&self) -> uint {
        let final_size = mem::size_of::<VertexType>();
        assert!(final_size >= self.vertex_size);
        final_size - self.vertex_size
    }

    pub fn attribute_f32(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 1, false, offset)
    }

    pub fn attribute_vec2f(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 2, false, offset)
    }

    pub fn attribute_vec3f(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 3, false, offset)
    }

    pub fn new_attribute(&mut self, layout: uint, gl_type: GLenum,
                         size: uint, normalized: bool, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        let attr_size = size * match gl_type {
            gl::FLOAT => mem::size_of::<f32>(),
            _ => fail!("Unsupported attribute type.")
        };
        assert!(self.max_attribute_size_left() >= attr_size);
        self.vertex_size += attr_size;
        if self.used_layouts.len() > layout {
            assert!(!self.used_layouts[layout]);
            *self.used_layouts.get_mut(layout) = true;
        } else {
            for _ in range(0, layout - self.used_layouts.len()) {
                self.used_layouts.push(false);
            }
            self.used_layouts.push(true);
        }
        self.attributes.push(VertexAttribute {
            layout: layout as GLuint,
            gl_type: gl_type,
            size: size as GLint,
            normalized: normalized as u8,
            offset: offset,
        });
        self
    }

    pub fn build(&self) -> VertexBuffer {
        assert_eq!(self.max_attribute_size_left(), 0);
        VertexBuffer::new(self.vertex_size, self.attributes.clone())
    }
}

fn bind_attributes(stride: uint, attributes: &[VertexAttribute]) {
    let stride = stride as i32;
    for attr in attributes.iter() {
        check_gl!(gl::EnableVertexAttribArray(attr.layout));
        check_gl_unsafe!(gl::VertexAttribPointer(attr.layout, attr.size,
                                                 attr.gl_type, attr.normalized,
                                                 stride, attr.offset));
    }
}

fn unbind_attributes(attributes: &[VertexAttribute]) {
    for attr in attributes.iter() {
        check_gl!(gl::DisableVertexAttribArray(attr.layout));
    }
}

struct VboId {
    id: GLuint
}
impl VboId {
    fn new_orphaned() -> VboId {
        VboId { id: 0 }
    }

    fn orphan(&mut self) -> &mut VboId {
        if self.id != 0 {
            check_gl_unsafe!(gl::DeleteBuffers(1, &self.id));
            self.id = 0;
        }
        self
    }

    fn reset(&mut self) -> &mut VboId {
        self.orphan();
        check_gl_unsafe!(gl::GenBuffers(1, &mut self.id));
        assert!(self.id != 0);
        self
    }

    fn bind(&self) -> &VboId {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, self.id));
        self
    }

    fn unbind(&self) -> &VboId {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, 0));
        self
    }

    fn id(&self) -> GLuint { self.id }
}

pub struct VertexBuffer {
    id: VboId,
    length: uint,
    attributes: Vec<VertexAttribute>,
    vertex_size: uint,
}

impl VertexBuffer {
    fn new(vertex_size: uint, attributes: Vec<VertexAttribute>)
            -> VertexBuffer {
        VertexBuffer {
            id: VboId::new_orphaned(),
            length: 0,
            vertex_size: vertex_size,
            attributes: attributes,
        }
    }

    pub fn draw_triangles(&self) -> &VertexBuffer {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, self.id.id()));
        bind_attributes(self.vertex_size, self.attributes.as_slice());
        check_gl!(gl::DrawArrays(gl::TRIANGLES, 0, self.length as i32));
        unbind_attributes(self.attributes.as_slice());
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, 0));
        self
    }

    pub fn len(&self) -> uint { self.length }

    pub fn set_data<V: Copy> (&mut self, usage: GLenum, data: &[V])
            -> &mut VertexBuffer {
        assert_eq!(self.vertex_size, mem::size_of::<V>());
        self.id.reset().bind();
        self.length = data.len();
        check_gl_unsafe!(gl::BufferData(
                gl::ARRAY_BUFFER, (data.len() * self.vertex_size) as GLsizeiptr,
                data.as_ptr() as *const libc::c_void, usage));
        self.id.unbind();
        self
    }
}
<commit_msg>Added u8 vertex attributes<commit_after>use check_gl;
use gl;
use gl::types::{GLint, GLuint, GLenum, GLsizeiptr};
use libc;
use libc::c_void;
use std::mem;

type IndexType = u16;

#[allow(raw_pointer_deriving)]
#[deriving(Clone)]
struct VertexAttribute {
    layout: GLuint,
    gl_type: GLenum,
    size: GLint,
    normalized: u8,
    offset: *const c_void
}

pub struct BufferBuilder<VertexType: Copy> {
    attributes: Vec<VertexAttribute>,
    used_layouts: Vec<bool>,
    vertex_size: uint
}
impl<VertexType: Copy>  BufferBuilder<VertexType> {
    pub fn new(capacity: uint) -> BufferBuilder<VertexType> {
        BufferBuilder {
            attributes: Vec::with_capacity(capacity),
            used_layouts: Vec::with_capacity(capacity),
            vertex_size: 0,
        }
    }

    pub fn max_attribute_size_left(&self) -> uint {
        let final_size = mem::size_of::<VertexType>();
        assert!(final_size >= self.vertex_size);
        final_size - self.vertex_size
    }

    pub fn attribute_u8(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::UNSIGNED_BYTE, 1, false, offset)
    }

    pub fn attribute_f32(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 1, false, offset)
    }

    pub fn attribute_vec2f(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 2, false, offset)
    }

    pub fn attribute_vec3f(&mut self, layout: uint, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        self.new_attribute(layout, gl::FLOAT, 3, false, offset)
    }

    pub fn new_attribute(&mut self, layout: uint, gl_type: GLenum,
                         size: uint, normalized: bool, offset: *const c_void)
            -> &mut BufferBuilder<VertexType> {
        let attr_size = size * match gl_type {
            gl::FLOAT => mem::size_of::<f32>(),
            gl::UNSIGNED_BYTE => mem::size_of::<u8>(),
            _ => fail!("Unsupported attribute type.")
        };
        assert!(self.max_attribute_size_left() >= attr_size);
        self.vertex_size += attr_size;
        if self.used_layouts.len() > layout {
            assert!(!self.used_layouts[layout]);
            *self.used_layouts.get_mut(layout) = true;
        } else {
            for _ in range(0, layout - self.used_layouts.len()) {
                self.used_layouts.push(false);
            }
            self.used_layouts.push(true);
        }
        self.attributes.push(VertexAttribute {
            layout: layout as GLuint,
            gl_type: gl_type,
            size: size as GLint,
            normalized: normalized as u8,
            offset: offset,
        });
        self
    }

    pub fn build(&self) -> VertexBuffer {
        assert_eq!(self.max_attribute_size_left(), 0);
        VertexBuffer::new(self.vertex_size, self.attributes.clone())
    }
}

fn bind_attributes(stride: uint, attributes: &[VertexAttribute]) {
    let stride = stride as i32;
    for attr in attributes.iter() {
        check_gl!(gl::EnableVertexAttribArray(attr.layout));
        match attr.gl_type {
            gl::FLOAT => check_gl_unsafe!(gl::VertexAttribPointer(
                attr.layout, attr.size, attr.gl_type, attr.normalized,
                stride, attr.offset)),
            gl::UNSIGNED_BYTE => check_gl_unsafe!(gl::VertexAttribIPointer(
                attr.layout, attr.size, attr.gl_type, stride, attr.offset)),
            _ => fail!("Missing attribute type from attrib ptr.")
        }
    }
}

fn unbind_attributes(attributes: &[VertexAttribute]) {
    for attr in attributes.iter() {
        check_gl!(gl::DisableVertexAttribArray(attr.layout));
    }
}

struct VboId {
    id: GLuint
}
impl VboId {
    fn new_orphaned() -> VboId {
        VboId { id: 0 }
    }

    fn orphan(&mut self) -> &mut VboId {
        if self.id != 0 {
            check_gl_unsafe!(gl::DeleteBuffers(1, &self.id));
            self.id = 0;
        }
        self
    }

    fn reset(&mut self) -> &mut VboId {
        self.orphan();
        check_gl_unsafe!(gl::GenBuffers(1, &mut self.id));
        assert!(self.id != 0);
        self
    }

    fn bind(&self) -> &VboId {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, self.id));
        self
    }

    fn unbind(&self) -> &VboId {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, 0));
        self
    }

    fn id(&self) -> GLuint { self.id }
}

pub struct VertexBuffer {
    id: VboId,
    length: uint,
    attributes: Vec<VertexAttribute>,
    vertex_size: uint,
}

impl VertexBuffer {
    fn new(vertex_size: uint, attributes: Vec<VertexAttribute>)
            -> VertexBuffer {
        VertexBuffer {
            id: VboId::new_orphaned(),
            length: 0,
            vertex_size: vertex_size,
            attributes: attributes,
        }
    }

    pub fn draw_triangles(&self) -> &VertexBuffer {
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, self.id.id()));
        bind_attributes(self.vertex_size, self.attributes.as_slice());
        check_gl!(gl::DrawArrays(gl::TRIANGLES, 0, self.length as i32));
        unbind_attributes(self.attributes.as_slice());
        check_gl!(gl::BindBuffer(gl::ARRAY_BUFFER, 0));
        self
    }

    pub fn len(&self) -> uint { self.length }

    pub fn set_data<V: Copy> (&mut self, usage: GLenum, data: &[V])
            -> &mut VertexBuffer {
        assert_eq!(self.vertex_size, mem::size_of::<V>());
        self.id.reset().bind();
        self.length = data.len();
        check_gl_unsafe!(gl::BufferData(
                gl::ARRAY_BUFFER, (data.len() * self.vertex_size) as GLsizeiptr,
                data.as_ptr() as *const libc::c_void, usage));
        self.id.unbind();
        self
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::iter;
use std::io::{self, Write};
use std::cmp;
use std::str;
use std::rt::heap;
use std::cell::Cell;
use std::marker::PhantomData;
use std::intrinsics;

// A chunk in the vector
struct Chunk<T> {
    // Next chunk
    next: *mut Chunk<T>,
    // Count of items
    len: usize,
    // Capacity
    cap: usize,
    // Items follow in memory
    items: [T; 0]
}

pub struct MonoVec<T> {
    head: *mut Chunk<T>,
    tail: Cell<*mut Chunk<T>>,
    _ph: PhantomData<T>
}

unsafe impl<T: Send> Send for MonoVec<T> {}

impl<T> Chunk<T> {
    fn array_size(len: usize) -> usize {
        len.checked_mul(mem::size_of::<T>()).unwrap()
    }

    fn mem_size(len: usize) -> usize {
        mem::size_of::<Self>().checked_add(Self::array_size(len)).unwrap()
    }

    fn new(cap: usize) -> *mut Self {
        unsafe {
            let res = heap::allocate(Self::mem_size(cap),
                                     mem::align_of::<Self>()) as *mut Self;
            ptr::write(&mut (*res).next, ptr::null_mut());
            ptr::write(&mut (*res).len, 0);
            ptr::write(&mut (*res).cap, cap);
            res
        }
    }
}

impl<T> MonoVec<T> {
    pub fn new() -> Self {
        Self::with_capacity(8)
    }

    pub fn with_capacity(cap: usize) -> Self {
        let head = Chunk::new(cap);
        MonoVec {
            head: head,
            tail: Cell::new(head),
            _ph: PhantomData
        }
    }

    // FIXME: track total len in header to make this O(1)?
    pub fn len(&self) -> usize {
        let mut len = 0;
        let mut cur = self.head;

        while !cur.is_null() {
            unsafe {
                len += (*cur).len;
                cur = (*cur).next;
            }
        }
        len
    }

    // Reserves space for at least `len` more contiguous elements, returning
    // a pointer to the space and the available capacity (which may be > `len`)
    #[inline(never)]
    pub fn reserve(&self, len: usize) -> (*mut T, usize) {
        unsafe {
            let tail = self.tail.get();
            let cap = (*tail).cap;
            if cap - (*tail).len < len {
                // Grow capacity exponentially to amortize cost of insertions
                let mut new_cap = cap.checked_mul(2).unwrap();
                while new_cap < len {
                    new_cap = new_cap.checked_mul(2).unwrap();
                }
                let new = Chunk::new(new_cap);

                (*tail).next = new;
                self.tail.set(new);
            }

            let tail = self.tail.get();
            let ptr = (*tail).items.as_mut_ptr().offset((*tail).len as isize);
            let cap = (*tail).cap - (*tail).len;
            (ptr, cap)
        }
    }

    // Adds to length of curent chunk.  Usually used after
    // writing into reserved space.
    pub unsafe fn add_len(&self, len: usize) {
        let tail = self.tail.get();
        (*tail).len = len.checked_add((*tail).len).unwrap();
    }

    #[inline]
    pub fn push(&self, elem: T) -> &T {
        let (ptr, _) = self.reserve(1);
        unsafe {
            ptr::write(ptr, elem);
            self.add_len(1);
            &*ptr
        }
    }

    pub fn push_as_slice<E: IntoIterator<Item=T>>(&self, elems: E) -> &[T]
            where E::IntoIter: ExactSizeIterator {
        let iter = elems.into_iter();
        let len = iter.len();
        let (ptr, _) = self.reserve(len);
        let mut cur = ptr;
        unsafe {
            for elem in iter {
                ptr::write(cur, elem);
                cur = cur.offset(1);
            }
            self.add_len(len);
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn chunks(&self) -> Chunks<T> {
        Chunks {
            chunk: self.head,
            _ph: PhantomData
        }
    }

    pub fn items(&self) -> Items<T> {
        #[inline(always)]
        fn id<T>(x: T) -> T { x }
        Items(self.chunks().flat_map(id))
    }
}

impl<'a, T: 'a> IntoIterator for &'a MonoVec<T> {
    type Item = &'a T;
    type IntoIter = Items<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.items()
    }
}

impl MonoVec<u8> {
    pub fn format(&self, args: fmt::Arguments) -> &str {
        let mut needed = 1;
        loop {
            let (ptr, len) = self.reserve(needed);
            let mut slice = unsafe { slice::from_raw_parts_mut(ptr, len) };
            let res = slice.write_fmt(args);
            match res {
                Ok(()) => unsafe {
                    let len = len - slice.len();
                    self.add_len(len);
                    return str::from_utf8_unchecked(slice::from_raw_parts(ptr, len))
                },
                Err(ref err) if err.kind() == io::ErrorKind::WriteZero => {
                    needed = len + 1;
                    continue
                }
                Err(err) => panic!(err)
            }
        }
    }
}

impl io::Write for MonoVec<u8> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let (ptr, len) = self.reserve(0);
        let len = cmp::min(len, buf.len());
        if len != 0 {
            unsafe {
                ptr::copy_nonoverlapping(buf.as_ptr(), ptr, len);
            }
        }
        Ok(len)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<T> Drop for MonoVec<T> {
    fn drop(&mut self) {
        unsafe {
            while !self.head.is_null() {
                let head = self.head;
                self.head = (*head).next;
                let start = (*head).items.as_mut_ptr();
                let end = start.offset((*head).len as isize);
                let mut cur = start;
                while cur < end {
                    intrinsics::drop_in_place(cur);
                    cur = cur.offset(1);
                }
            }
        }
    }
}

pub struct Chunks<'a, T: 'a> {
    chunk: *mut Chunk<T>,
    _ph: PhantomData<&'a [T]>
}

impl<'a, T> Iterator for Chunks<'a, T> {
    type Item = &'a [T];

    fn next(&mut self) -> Option<&'a [T]> {
        let chunk = self.chunk;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                self.chunk = (*chunk).next;
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

// Wrapper to hide ugly adapter type
pub struct Items<'a, T: 'a>(iter::FlatMap<Chunks<'a, T>, &'a [T], fn(&'a [T]) -> &'a [T]>);

impl<'a, T: 'a> Iterator for Items<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.0.next()
    }
}

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

    #[test]
    fn format() {
        let buffer = MonoVec::new();
        for i in 0..100 {
            assert_eq!(buffer.format(format_args!("hello {}", i)),
                       format!("hello {}", i));
        }
    }
}
<commit_msg>Fix Drop for MonoVec<commit_after>use std::fmt;
use std::mem;
use std::ptr;
use std::slice;
use std::iter;
use std::io::{self, Write};
use std::cmp;
use std::str;
use std::rt::heap;
use std::cell::Cell;
use std::marker::PhantomData;
use std::intrinsics;

// A chunk in the vector
struct Chunk<T> {
    // Next chunk
    next: *mut Chunk<T>,
    // Count of items
    len: usize,
    // Capacity
    cap: usize,
    // Items follow in memory
    items: [T; 0]
}

pub struct MonoVec<T> {
    head: *mut Chunk<T>,
    tail: Cell<*mut Chunk<T>>,
    _ph: PhantomData<T>
}

unsafe impl<T: Send> Send for MonoVec<T> {}

impl<T> Chunk<T> {
    fn array_size(len: usize) -> usize {
        len.checked_mul(mem::size_of::<T>()).unwrap()
    }

    fn mem_size(len: usize) -> usize {
        mem::size_of::<Self>().checked_add(Self::array_size(len)).unwrap()
    }

    fn new(cap: usize) -> *mut Self {
        unsafe {
            let res = heap::allocate(Self::mem_size(cap),
                                     mem::align_of::<Self>()) as *mut Self;
            ptr::write(&mut (*res).next, ptr::null_mut());
            ptr::write(&mut (*res).len, 0);
            ptr::write(&mut (*res).cap, cap);
            res
        }
    }
}

impl<T> MonoVec<T> {
    pub fn new() -> Self {
        Self::with_capacity(8)
    }

    pub fn with_capacity(cap: usize) -> Self {
        let head = Chunk::new(cap);
        MonoVec {
            head: head,
            tail: Cell::new(head),
            _ph: PhantomData
        }
    }

    // FIXME: track total len in header to make this O(1)?
    pub fn len(&self) -> usize {
        let mut len = 0;
        let mut cur = self.head;

        while !cur.is_null() {
            unsafe {
                len += (*cur).len;
                cur = (*cur).next;
            }
        }
        len
    }

    // Reserves space for at least `len` more contiguous elements, returning
    // a pointer to the space and the available capacity (which may be > `len`)
    #[inline(never)]
    pub fn reserve(&self, len: usize) -> (*mut T, usize) {
        unsafe {
            let tail = self.tail.get();
            let cap = (*tail).cap;
            if cap - (*tail).len < len {
                // Grow capacity exponentially to amortize cost of insertions
                let mut new_cap = cap.checked_mul(2).unwrap();
                while new_cap < len {
                    new_cap = new_cap.checked_mul(2).unwrap();
                }
                let new = Chunk::new(new_cap);

                (*tail).next = new;
                self.tail.set(new);
            }

            let tail = self.tail.get();
            let ptr = (*tail).items.as_mut_ptr().offset((*tail).len as isize);
            let cap = (*tail).cap - (*tail).len;
            (ptr, cap)
        }
    }

    // Adds to length of curent chunk.  Usually used after
    // writing into reserved space.
    pub unsafe fn add_len(&self, len: usize) {
        let tail = self.tail.get();
        (*tail).len = len.checked_add((*tail).len).unwrap();
    }

    #[inline]
    pub fn push(&self, elem: T) -> &T {
        let (ptr, _) = self.reserve(1);
        unsafe {
            ptr::write(ptr, elem);
            self.add_len(1);
            &*ptr
        }
    }

    pub fn push_as_slice<E: IntoIterator<Item=T>>(&self, elems: E) -> &[T]
            where E::IntoIter: ExactSizeIterator {
        let iter = elems.into_iter();
        let len = iter.len();
        let (ptr, _) = self.reserve(len);
        let mut cur = ptr;
        unsafe {
            for elem in iter {
                ptr::write(cur, elem);
                cur = cur.offset(1);
            }
            self.add_len(len);
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn chunks(&self) -> Chunks<T> {
        Chunks {
            chunk: self.head,
            _ph: PhantomData
        }
    }

    pub fn items(&self) -> Items<T> {
        #[inline(always)]
        fn id<T>(x: T) -> T { x }
        Items(self.chunks().flat_map(id))
    }
}

impl<'a, T: 'a> IntoIterator for &'a MonoVec<T> {
    type Item = &'a T;
    type IntoIter = Items<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.items()
    }
}

impl MonoVec<u8> {
    pub fn format(&self, args: fmt::Arguments) -> &str {
        let mut needed = 1;
        loop {
            let (ptr, len) = self.reserve(needed);
            let mut slice = unsafe { slice::from_raw_parts_mut(ptr, len) };
            let res = slice.write_fmt(args);
            match res {
                Ok(()) => unsafe {
                    let len = len - slice.len();
                    self.add_len(len);
                    return str::from_utf8_unchecked(slice::from_raw_parts(ptr, len))
                },
                Err(ref err) if err.kind() == io::ErrorKind::WriteZero => {
                    needed = len + 1;
                    continue
                }
                Err(err) => panic!(err)
            }
        }
    }
}

impl io::Write for MonoVec<u8> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let (ptr, len) = self.reserve(0);
        let len = cmp::min(len, buf.len());
        if len != 0 {
            unsafe {
                ptr::copy_nonoverlapping(buf.as_ptr(), ptr, len);
            }
        }
        Ok(len)
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<T> Drop for MonoVec<T> {
    fn drop(&mut self) {
        unsafe {
            while !self.head.is_null() {
                let head = self.head;
                self.head = (*head).next;
                if intrinsics::needs_drop::<T>() {
                    let mut cur = (*head).items.as_mut_ptr();
                    let end = cur.offset((*head).len as isize);
                    while cur < end {
                        intrinsics::drop_in_place(cur);
                        cur = cur.offset(1);
                    }
                }
                heap::deallocate(head as *mut u8,
                                 (*head).cap * mem::size_of::<T>(),
                                 mem::min_align_of::<T>());
            }
        }
    }
}

pub struct Chunks<'a, T: 'a> {
    chunk: *mut Chunk<T>,
    _ph: PhantomData<&'a [T]>
}

impl<'a, T> Iterator for Chunks<'a, T> {
    type Item = &'a [T];

    fn next(&mut self) -> Option<&'a [T]> {
        let chunk = self.chunk;
        if chunk.is_null() {
            None
        } else {
            unsafe {
                self.chunk = (*chunk).next;
                Some(slice::from_raw_parts((*chunk).items.as_ptr(), (*chunk).len))
            }
        }
    }
}

// Wrapper to hide ugly adapter type
pub struct Items<'a, T: 'a>(iter::FlatMap<Chunks<'a, T>, &'a [T], fn(&'a [T]) -> &'a [T]>);

impl<'a, T: 'a> Iterator for Items<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.0.next()
    }
}

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

    #[test]
    fn format() {
        let buffer = MonoVec::new();
        for i in 0..100 {
            assert_eq!(buffer.format(format_args!("hello {}", i)),
                       format!("hello {}", i));
        }
    }
}
<|endoftext|>


<commit_before>use rand::Rng;
use seckey::SecKey;
use kyber::{ params, kem };
use super::{ KeyExchange, CheckedExchange };
use ::Packing;

pub struct Kyber;
pub struct PrivateKey(pub SecKey<[u8; params::SECRETKEYBYTES]>);
pub struct PublicKey(pub [u8; params::PUBLICKEYBYTES]);
pub struct Message(pub [u8; params::CIPHERTEXTBYTES]);

impl KeyExchange for Kyber {
    type PrivateKey = PrivateKey;
    type PublicKey = PublicKey;
    type Message = Message;

    const SHARED_LENGTH: usize = params::SYMBYTES;

    fn kerpair<R: Rng>(mut r: R) -> (Self::PrivateKey, Self::PublicKey) {
        // TODO use `SecKey::with_default()`
        let mut sk = SecKey::new([0; params::SECRETKEYBYTES]).ok().expect("memsec malloc failed");

        let mut pk = [0; params::PUBLICKEYBYTES];
        kem::keypair(&mut r, &mut pk, &mut sk.write());
        (PrivateKey(sk), PublicKey(pk))
    }

    fn exchange_to<R: Rng>(mut r: R, sharedkey: &mut [u8], &PublicKey(ref pk): &Self::PublicKey) -> Self::Message {
        let sharedkey = array_mut_ref!(sharedkey, 0, params::SYMBYTES);
        let mut c = [0; params::CIPHERTEXTBYTES];
        kem::enc(&mut r, &mut c, sharedkey, pk);
        Message(c)
    }

    fn exchange_from(sharedkey: &mut [u8], sk: &Self::PrivateKey, m: &Self::Message) {
        <Kyber as CheckedExchange>::exchange_from(sharedkey, sk, m);
    }
}

impl CheckedExchange for Kyber {
    fn exchange_from(
        sharedkey: &mut [u8],
        &PrivateKey(ref sk): &Self::PrivateKey,
        &Message(ref m): &Self::Message
    ) -> bool {
        let sharedkey = array_mut_ref!(sharedkey, 0, params::SYMBYTES);
        kem::dec(sharedkey, m, &sk.read())
    }
}

impl Packing for PrivateKey {
    const BYTES_LENGTH: usize = params::SECRETKEYBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::SECRETKEYBYTES);
        buf.clone_from(&*self.0.read())
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::SECRETKEYBYTES);
        SecKey::from_ref(buf).map(PrivateKey)
            .expect("memsec malloc failed")
    }
}

impl Packing for PublicKey {
    const BYTES_LENGTH: usize = params::PUBLICKEYBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::PUBLICKEYBYTES);
        buf.clone_from(&self.0)
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::PUBLICKEYBYTES);
        let mut pk = [0; params::PUBLICKEYBYTES];
        pk.clone_from(buf);
        PublicKey(pk)
    }
}

impl Packing for Message {
    const BYTES_LENGTH: usize = params::CIPHERTEXTBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::CIPHERTEXTBYTES);
        buf.clone_from(&self.0)
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::CIPHERTEXTBYTES);
        let mut sig = [0; params::CIPHERTEXTBYTES];
        sig.clone_from(buf);
        Message(sig)
    }
}
<commit_msg>Fix typo in implementation<commit_after>use rand::Rng;
use seckey::SecKey;
use kyber::{ params, kem };
use super::{ KeyExchange, CheckedExchange };
use ::Packing;

pub struct Kyber;
pub struct PrivateKey(pub SecKey<[u8; params::SECRETKEYBYTES]>);
pub struct PublicKey(pub [u8; params::PUBLICKEYBYTES]);
pub struct Message(pub [u8; params::CIPHERTEXTBYTES]);

impl KeyExchange for Kyber {
    type PrivateKey = PrivateKey;
    type PublicKey = PublicKey;
    type Message = Message;

    const SHARED_LENGTH: usize = params::SYMBYTES;

    fn keypair<R: Rng>(mut r: R) -> (Self::PrivateKey, Self::PublicKey) {
        // TODO use `SecKey::with_default()`
        let mut sk = SecKey::new([0; params::SECRETKEYBYTES]).ok().expect("memsec malloc failed");

        let mut pk = [0; params::PUBLICKEYBYTES];
        kem::keypair(&mut r, &mut pk, &mut sk.write());
        (PrivateKey(sk), PublicKey(pk))
    }

    fn exchange_to<R: Rng>(mut r: R, sharedkey: &mut [u8], &PublicKey(ref pk): &Self::PublicKey) -> Self::Message {
        let sharedkey = array_mut_ref!(sharedkey, 0, params::SYMBYTES);
        let mut c = [0; params::CIPHERTEXTBYTES];
        kem::enc(&mut r, &mut c, sharedkey, pk);
        Message(c)
    }

    fn exchange_from(sharedkey: &mut [u8], sk: &Self::PrivateKey, m: &Self::Message) {
        <Kyber as CheckedExchange>::exchange_from(sharedkey, sk, m);
    }
}

impl CheckedExchange for Kyber {
    fn exchange_from(
        sharedkey: &mut [u8],
        &PrivateKey(ref sk): &Self::PrivateKey,
        &Message(ref m): &Self::Message
    ) -> bool {
        let sharedkey = array_mut_ref!(sharedkey, 0, params::SYMBYTES);
        kem::dec(sharedkey, m, &sk.read())
    }
}

impl Packing for PrivateKey {
    const BYTES_LENGTH: usize = params::SECRETKEYBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::SECRETKEYBYTES);
        buf.clone_from(&*self.0.read())
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::SECRETKEYBYTES);
        SecKey::from_ref(buf).map(PrivateKey)
            .expect("memsec malloc failed")
    }
}

impl Packing for PublicKey {
    const BYTES_LENGTH: usize = params::PUBLICKEYBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::PUBLICKEYBYTES);
        buf.clone_from(&self.0)
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::PUBLICKEYBYTES);
        let mut pk = [0; params::PUBLICKEYBYTES];
        pk.clone_from(buf);
        PublicKey(pk)
    }
}

impl Packing for Message {
    const BYTES_LENGTH: usize = params::CIPHERTEXTBYTES;

    fn read_bytes(&self, buf: &mut [u8]) {
        let buf = array_mut_ref!(buf, 0, params::CIPHERTEXTBYTES);
        buf.clone_from(&self.0)
    }

    fn from_bytes(buf: &[u8]) -> Self {
        let buf = array_ref!(buf, 0, params::CIPHERTEXTBYTES);
        let mut sig = [0; params::CIPHERTEXTBYTES];
        sig.clone_from(buf);
        Message(sig)
    }
}
<|endoftext|>