<commit_before>//! QuotientFilter implementation.
use std::collections::hash_map::DefaultHasher;
use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};
use std::mem::size_of;

use fixedbitset::FixedBitSet;
use succinct::{IntVec, IntVecMut, IntVector};

use filters::Filter;

/// Error that signals that the QuotientFilter is full.
#[derive(Debug)]
pub struct QuotientFilterFull;

/// TODO
/// Run: List of fingerprints with same quotient.
/// Cluster: set of runs
///
/// https://en.wikipedia.org/wiki/Quotient_filter
/// http://static.usenix.org/events/hotstorage11/tech/final_files/Bender.pdf
/// https://www.vldb.org/pvldb/vol5/p1627_michaelabender_vldb2012.pdf
#[derive(Debug)]
pub struct QuotientFilter<B = BuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    is_occupied: FixedBitSet,
    is_continuation: FixedBitSet,
    is_shifted: FixedBitSet,
    remainders: IntVector,
    bits_quotient: usize,
    buildhasher: B,
    n_elements: usize,
}

impl QuotientFilter {
    /// Create new quotient filter with:
    ///
    /// - `bits_quotient`: number of bits used for a quotient, aka `2^bits_quotient` slots will be
    ///   allocated
    /// - `bits_remainder`: number of bits used for the remainder, so every slot will require
    ///   `bits_remainder + 3` bits of storage
    ///
    /// and a default hasher.
    pub fn with_params(bits_quotient: usize, bits_remainder: usize) -> Self {
        let buildhasher = BuildHasherDefault::<DefaultHasher>::default();
        QuotientFilter::with_params_and_hash(bits_quotient, bits_remainder, buildhasher)
    }
}

impl<B> QuotientFilter<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Create new quotient filter with:
    ///
    /// - `bits_quotient`: number of bits used for a quotient, aka `2^bits_quotient` slots will be
    ///   allocated
    /// - `bits_remainder`: number of bits used for the remainder, so every slot will require
    ///   `bits_remainder + 3` bits of storage
    /// - `buildhasher`: hash implementation
    pub fn with_params_and_hash(
        bits_quotient: usize,
        bits_remainder: usize,
        buildhasher: B,
    ) -> Self {
        assert!(
            (bits_remainder > 0) && (bits_remainder <= size_of::<usize>() * 8),
            "bits_remainder ({}) must be greater than 0 and smaller or equal than {}",
            bits_remainder,
            size_of::<usize>() * 8,
        );
        assert!(
            bits_quotient > 0,
            "bits_quotient ({}) must be greater than 0",
            bits_quotient,
        );
        assert!(
            bits_remainder + bits_quotient <= 64,
            "bits_remainder ({}) + bits_quotient ({}) must be smaller or equal than 64",
            bits_remainder,
            bits_quotient,
        );

        let len = 1 << bits_quotient;
        Self {
            is_occupied: FixedBitSet::with_capacity(len),
            is_continuation: FixedBitSet::with_capacity(len),
            is_shifted: FixedBitSet::with_capacity(len),
            remainders: IntVector::with_fill(bits_remainder, len as u64, 0),
            bits_quotient,
            buildhasher,
            n_elements: 0,
        }
    }

    /// Number of bits used for addressing slots.
    pub fn bits_quotient(&self) -> usize {
        self.bits_quotient
    }

    /// Number of bits stored as fingeprint information.
    pub fn bits_remainder(&self) -> usize {
        self.remainders.element_bits()
    }

    fn calc_quotient_remainder<T>(&self, obj: &T) -> (usize, usize)
    where
        T: Hash,
    {
        let bits_remainder = self.bits_remainder();
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let fingerprint = hasher.finish();
        let bits_trash = 64 - bits_remainder - self.bits_quotient;
        let trash = if bits_trash > 0 {
            (fingerprint >> (64 - bits_trash)) << (64 - bits_trash)
        } else {
            0
        };
        let fingerprint_clean = fingerprint - trash;
        let quotient = fingerprint_clean >> bits_remainder;
        let remainder = fingerprint_clean - (quotient << bits_remainder);
        (quotient as usize, remainder as usize)
    }

    fn decr(&self, pos: &mut usize) {
        *pos = if *pos == 0 {
            self.is_occupied.len() - 1
        } else {
            *pos - 1
        };
    }

    fn incr(&self, pos: &mut usize) {
        *pos = if *pos == self.is_occupied.len() - 1 {
            0
        } else {
            *pos + 1
        }
    }

    fn scan(
        &self,
        quotient: usize,
        remainder: usize,
        on_insert: bool,
    ) -> (bool, usize, bool, usize) {
        let run_exists = self.is_occupied[quotient];
        if (!run_exists) && (!on_insert) {
            // fast-path for query, since we don't need to find the correct position for the
            // insertion process
            return (run_exists, quotient, run_exists, quotient);
        }

        // walk back to find the beginning of the cluster
        let mut b = quotient;
        while self.is_shifted[b] {
            self.decr(&mut b);
        }

        // walk forward to find the actual start of the run
        let mut s = b;
        while b != quotient {
            // invariant: `s` poins to first slot of bucket `b`

            // skip all elements in the current run
            loop {
                self.incr(&mut s);
                if !self.is_continuation[s] {
                    break;
                }
            }

            // find the next occupied bucket
            loop {
                self.incr(&mut b);
                if self.is_occupied[b] || ((b == quotient) && on_insert) {
                    break;
                }
            }
        }
        // `s` now points to the first remainder in bucket at `quotient`

        // search of remainder within the run
        let start_of_run = s;
        if run_exists {
            loop {
                let r = self.remainders.get(s as u64);
                if r == remainder {
                    return (run_exists, s, run_exists, start_of_run);
                }
                if r > remainder {
                    // remainders are sorted within run
                    break;
                }
                self.incr(&mut s);
                if !self.is_continuation[s] {
                    break;
                }
            }
        }
        (false, s, run_exists, start_of_run)
    }
}

impl<B> Filter for QuotientFilter<B>
where
    B: BuildHasher + Clone + Eq,
{
    type InsertErr = QuotientFilterFull;

    fn clear(&mut self) {
        self.is_occupied.clear();
        self.is_continuation.clear();
        self.is_shifted.clear();
        self.remainders =
            IntVector::with_fill(self.remainders.element_bits(), self.remainders.len(), 0);
        self.n_elements = 0;
    }

    fn insert<T>(&mut self, t: &T) -> Result<(), Self::InsertErr>
    where
        T: Hash,
    {
        let (quotient, remainder) = self.calc_quotient_remainder(t);
        let (present, mut position, run_exists, start_of_run) =
            self.scan(quotient, remainder, true);

        // early exit if the element is already present
        if present {
            return Ok(());
        }
        // we need to insert the element into the filter

        // error out if there is no space left
        if self.n_elements == self.is_occupied.len() {
            return Err(QuotientFilterFull);
        }

        // set up swap chain
        let mut current_is_continuation =
            self.is_continuation[position] || (run_exists && (position == start_of_run));
        let mut current_remainder = self.remainders.get(position as u64);
        let mut current_used = self.is_occupied[position] || self.is_shifted[position];

        // set current state
        self.remainders.set(position as u64, remainder);
        if position != start_of_run {
            // might be an append operation, ensure is_continuation and is_shifted are set
            self.is_continuation.set(position, true);
        }
        if position != quotient {
            // not at canonical slot
            self.is_shifted.set(position, true);
        }

        // run swap chain until nothing to do
        let start = position;
        while current_used {
            self.incr(&mut position);
            let next_is_continuation = self.is_continuation[position];
            let next_remainder = self.remainders.get(position as u64);
            let next_used = self.is_occupied[position] || self.is_shifted[position];

            self.is_shifted.set(position, true);
            self.is_continuation.set(position, current_is_continuation);
            self.remainders.set(position as u64, current_remainder);

            current_is_continuation = next_is_continuation;
            current_remainder = next_remainder;
            current_used = next_used;

            if position == start {
                panic!("infinite loop detected");
            }
        }

        // mark canonical slot as occupied
        self.is_occupied.set(quotient, true);

        // done
        self.n_elements += 1;
        Ok(())
    }

    fn is_empty(&self) -> bool {
        self.n_elements == 0
    }

    fn len(&self) -> usize {
        self.n_elements
    }

    fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        let (quotient, remainder) = self.calc_quotient_remainder(obj);
        let (present, _position, _run_exists, _start_of_run) =
            self.scan(quotient, remainder, false);
        present
    }
}

#[cfg(test)]
mod tests {
    use super::QuotientFilter;
    use filters::Filter;

    #[test]
    #[should_panic(expected = "bits_quotient (0) must be greater than 0")]
    fn new_bits_quotient_0() {
        QuotientFilter::with_params(0, 16);
    }

    #[cfg(target_pointer_width = "32")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (0) must be greater than 0 and smaller or equal than 32"
    )]
    fn new_bits_remainder_0() {
        QuotientFilter::with_params(3, 0);
    }

    #[cfg(target_pointer_width = "64")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (0) must be greater than 0 and smaller or equal than 64"
    )]
    fn new_bits_remainder_0() {
        QuotientFilter::with_params(3, 0);
    }

    #[cfg(target_pointer_width = "32")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (33) must be greater than 0 and smaller or equal than 32"
    )]
    fn new_bits_remainder_too_large() {
        QuotientFilter::with_params(3, 33);
    }

    #[cfg(target_pointer_width = "64")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (65) must be greater than 0 and smaller or equal than 64"
    )]
    fn new_bits_remainder_too_large() {
        QuotientFilter::with_params(3, 65);
    }

    #[test]
    #[should_panic(
        expected = "bits_remainder (5) + bits_quotient (60) must be smaller or equal than 64"
    )]
    fn new_too_many_bits() {
        QuotientFilter::with_params(60, 5);
    }

    #[test]
    fn new() {
        let qf = QuotientFilter::with_params(3, 16);
        assert!(qf.is_empty());
        assert_eq!(qf.len(), 0);
        assert!(!qf.query(&13));
        assert_eq!(qf.bits_quotient(), 3);
        assert_eq!(qf.bits_remainder(), 16);
    }

    #[test]
    fn insert() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        assert!(!qf.is_empty());
        assert_eq!(qf.len(), 1);
        assert!(qf.query(&13));
        assert!(!qf.query(&42));
    }

    #[test]
    fn double_insert() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        qf.insert(&13).unwrap();
        assert!(!qf.is_empty());
        assert_eq!(qf.len(), 1);
        assert!(qf.query(&13));
        assert!(!qf.query(&42));
    }

    #[test]
    fn full() {
        let mut qf = QuotientFilter::with_params(3, 16);
        for i in 0..8 {
            qf.insert(&i).unwrap();
            for j in 0..i {
                assert!(qf.query(&j), "Cannot find {} after inserting {}", j, i);
            }
        }
        assert!(qf.insert(&1000).is_err());
    }

    #[test]
    fn clear() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        qf.clear();
        assert!(qf.is_empty());
        assert_eq!(qf.len(), 0);
        assert!(!qf.query(&13));
        assert_eq!(qf.bits_quotient(), 3);
        assert_eq!(qf.bits_remainder(), 16);
    }
}
<commit_msg>quotientfilter: impl Clone<commit_after>//! QuotientFilter implementation.
use std::collections::hash_map::DefaultHasher;
use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};
use std::mem::size_of;

use fixedbitset::FixedBitSet;
use succinct::{IntVec, IntVecMut, IntVector};

use filters::Filter;

/// Error that signals that the QuotientFilter is full.
#[derive(Debug)]
pub struct QuotientFilterFull;

/// TODO
/// Run: List of fingerprints with same quotient.
/// Cluster: set of runs
///
/// https://en.wikipedia.org/wiki/Quotient_filter
/// http://static.usenix.org/events/hotstorage11/tech/final_files/Bender.pdf
/// https://www.vldb.org/pvldb/vol5/p1627_michaelabender_vldb2012.pdf
#[derive(Clone, Debug)]
pub struct QuotientFilter<B = BuildHasherDefault<DefaultHasher>>
where
    B: BuildHasher + Clone + Eq,
{
    is_occupied: FixedBitSet,
    is_continuation: FixedBitSet,
    is_shifted: FixedBitSet,
    remainders: IntVector,
    bits_quotient: usize,
    buildhasher: B,
    n_elements: usize,
}

impl QuotientFilter {
    /// Create new quotient filter with:
    ///
    /// - `bits_quotient`: number of bits used for a quotient, aka `2^bits_quotient` slots will be
    ///   allocated
    /// - `bits_remainder`: number of bits used for the remainder, so every slot will require
    ///   `bits_remainder + 3` bits of storage
    ///
    /// and a default hasher.
    pub fn with_params(bits_quotient: usize, bits_remainder: usize) -> Self {
        let buildhasher = BuildHasherDefault::<DefaultHasher>::default();
        QuotientFilter::with_params_and_hash(bits_quotient, bits_remainder, buildhasher)
    }
}

impl<B> QuotientFilter<B>
where
    B: BuildHasher + Clone + Eq,
{
    /// Create new quotient filter with:
    ///
    /// - `bits_quotient`: number of bits used for a quotient, aka `2^bits_quotient` slots will be
    ///   allocated
    /// - `bits_remainder`: number of bits used for the remainder, so every slot will require
    ///   `bits_remainder + 3` bits of storage
    /// - `buildhasher`: hash implementation
    pub fn with_params_and_hash(
        bits_quotient: usize,
        bits_remainder: usize,
        buildhasher: B,
    ) -> Self {
        assert!(
            (bits_remainder > 0) && (bits_remainder <= size_of::<usize>() * 8),
            "bits_remainder ({}) must be greater than 0 and smaller or equal than {}",
            bits_remainder,
            size_of::<usize>() * 8,
        );
        assert!(
            bits_quotient > 0,
            "bits_quotient ({}) must be greater than 0",
            bits_quotient,
        );
        assert!(
            bits_remainder + bits_quotient <= 64,
            "bits_remainder ({}) + bits_quotient ({}) must be smaller or equal than 64",
            bits_remainder,
            bits_quotient,
        );

        let len = 1 << bits_quotient;
        Self {
            is_occupied: FixedBitSet::with_capacity(len),
            is_continuation: FixedBitSet::with_capacity(len),
            is_shifted: FixedBitSet::with_capacity(len),
            remainders: IntVector::with_fill(bits_remainder, len as u64, 0),
            bits_quotient,
            buildhasher,
            n_elements: 0,
        }
    }

    /// Number of bits used for addressing slots.
    pub fn bits_quotient(&self) -> usize {
        self.bits_quotient
    }

    /// Number of bits stored as fingeprint information.
    pub fn bits_remainder(&self) -> usize {
        self.remainders.element_bits()
    }

    fn calc_quotient_remainder<T>(&self, obj: &T) -> (usize, usize)
    where
        T: Hash,
    {
        let bits_remainder = self.bits_remainder();
        let mut hasher = self.buildhasher.build_hasher();
        obj.hash(&mut hasher);
        let fingerprint = hasher.finish();
        let bits_trash = 64 - bits_remainder - self.bits_quotient;
        let trash = if bits_trash > 0 {
            (fingerprint >> (64 - bits_trash)) << (64 - bits_trash)
        } else {
            0
        };
        let fingerprint_clean = fingerprint - trash;
        let quotient = fingerprint_clean >> bits_remainder;
        let remainder = fingerprint_clean - (quotient << bits_remainder);
        (quotient as usize, remainder as usize)
    }

    fn decr(&self, pos: &mut usize) {
        *pos = if *pos == 0 {
            self.is_occupied.len() - 1
        } else {
            *pos - 1
        };
    }

    fn incr(&self, pos: &mut usize) {
        *pos = if *pos == self.is_occupied.len() - 1 {
            0
        } else {
            *pos + 1
        }
    }

    fn scan(
        &self,
        quotient: usize,
        remainder: usize,
        on_insert: bool,
    ) -> (bool, usize, bool, usize) {
        let run_exists = self.is_occupied[quotient];
        if (!run_exists) && (!on_insert) {
            // fast-path for query, since we don't need to find the correct position for the
            // insertion process
            return (run_exists, quotient, run_exists, quotient);
        }

        // walk back to find the beginning of the cluster
        let mut b = quotient;
        while self.is_shifted[b] {
            self.decr(&mut b);
        }

        // walk forward to find the actual start of the run
        let mut s = b;
        while b != quotient {
            // invariant: `s` poins to first slot of bucket `b`

            // skip all elements in the current run
            loop {
                self.incr(&mut s);
                if !self.is_continuation[s] {
                    break;
                }
            }

            // find the next occupied bucket
            loop {
                self.incr(&mut b);
                if self.is_occupied[b] || ((b == quotient) && on_insert) {
                    break;
                }
            }
        }
        // `s` now points to the first remainder in bucket at `quotient`

        // search of remainder within the run
        let start_of_run = s;
        if run_exists {
            loop {
                let r = self.remainders.get(s as u64);
                if r == remainder {
                    return (run_exists, s, run_exists, start_of_run);
                }
                if r > remainder {
                    // remainders are sorted within run
                    break;
                }
                self.incr(&mut s);
                if !self.is_continuation[s] {
                    break;
                }
            }
        }
        (false, s, run_exists, start_of_run)
    }
}

impl<B> Filter for QuotientFilter<B>
where
    B: BuildHasher + Clone + Eq,
{
    type InsertErr = QuotientFilterFull;

    fn clear(&mut self) {
        self.is_occupied.clear();
        self.is_continuation.clear();
        self.is_shifted.clear();
        self.remainders =
            IntVector::with_fill(self.remainders.element_bits(), self.remainders.len(), 0);
        self.n_elements = 0;
    }

    fn insert<T>(&mut self, t: &T) -> Result<(), Self::InsertErr>
    where
        T: Hash,
    {
        let (quotient, remainder) = self.calc_quotient_remainder(t);
        let (present, mut position, run_exists, start_of_run) =
            self.scan(quotient, remainder, true);

        // early exit if the element is already present
        if present {
            return Ok(());
        }
        // we need to insert the element into the filter

        // error out if there is no space left
        if self.n_elements == self.is_occupied.len() {
            return Err(QuotientFilterFull);
        }

        // set up swap chain
        let mut current_is_continuation =
            self.is_continuation[position] || (run_exists && (position == start_of_run));
        let mut current_remainder = self.remainders.get(position as u64);
        let mut current_used = self.is_occupied[position] || self.is_shifted[position];

        // set current state
        self.remainders.set(position as u64, remainder);
        if position != start_of_run {
            // might be an append operation, ensure is_continuation and is_shifted are set
            self.is_continuation.set(position, true);
        }
        if position != quotient {
            // not at canonical slot
            self.is_shifted.set(position, true);
        }

        // run swap chain until nothing to do
        let start = position;
        while current_used {
            self.incr(&mut position);
            let next_is_continuation = self.is_continuation[position];
            let next_remainder = self.remainders.get(position as u64);
            let next_used = self.is_occupied[position] || self.is_shifted[position];

            self.is_shifted.set(position, true);
            self.is_continuation.set(position, current_is_continuation);
            self.remainders.set(position as u64, current_remainder);

            current_is_continuation = next_is_continuation;
            current_remainder = next_remainder;
            current_used = next_used;

            if position == start {
                panic!("infinite loop detected");
            }
        }

        // mark canonical slot as occupied
        self.is_occupied.set(quotient, true);

        // done
        self.n_elements += 1;
        Ok(())
    }

    fn is_empty(&self) -> bool {
        self.n_elements == 0
    }

    fn len(&self) -> usize {
        self.n_elements
    }

    fn query<T>(&self, obj: &T) -> bool
    where
        T: Hash,
    {
        let (quotient, remainder) = self.calc_quotient_remainder(obj);
        let (present, _position, _run_exists, _start_of_run) =
            self.scan(quotient, remainder, false);
        present
    }
}

#[cfg(test)]
mod tests {
    use super::QuotientFilter;
    use filters::Filter;

    #[test]
    #[should_panic(expected = "bits_quotient (0) must be greater than 0")]
    fn new_bits_quotient_0() {
        QuotientFilter::with_params(0, 16);
    }

    #[cfg(target_pointer_width = "32")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (0) must be greater than 0 and smaller or equal than 32"
    )]
    fn new_bits_remainder_0() {
        QuotientFilter::with_params(3, 0);
    }

    #[cfg(target_pointer_width = "64")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (0) must be greater than 0 and smaller or equal than 64"
    )]
    fn new_bits_remainder_0() {
        QuotientFilter::with_params(3, 0);
    }

    #[cfg(target_pointer_width = "32")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (33) must be greater than 0 and smaller or equal than 32"
    )]
    fn new_bits_remainder_too_large() {
        QuotientFilter::with_params(3, 33);
    }

    #[cfg(target_pointer_width = "64")]
    #[test]
    #[should_panic(
        expected = "bits_remainder (65) must be greater than 0 and smaller or equal than 64"
    )]
    fn new_bits_remainder_too_large() {
        QuotientFilter::with_params(3, 65);
    }

    #[test]
    #[should_panic(
        expected = "bits_remainder (5) + bits_quotient (60) must be smaller or equal than 64"
    )]
    fn new_too_many_bits() {
        QuotientFilter::with_params(60, 5);
    }

    #[test]
    fn new() {
        let qf = QuotientFilter::with_params(3, 16);
        assert!(qf.is_empty());
        assert_eq!(qf.len(), 0);
        assert!(!qf.query(&13));
        assert_eq!(qf.bits_quotient(), 3);
        assert_eq!(qf.bits_remainder(), 16);
    }

    #[test]
    fn insert() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        assert!(!qf.is_empty());
        assert_eq!(qf.len(), 1);
        assert!(qf.query(&13));
        assert!(!qf.query(&42));
    }

    #[test]
    fn double_insert() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        qf.insert(&13).unwrap();
        assert!(!qf.is_empty());
        assert_eq!(qf.len(), 1);
        assert!(qf.query(&13));
        assert!(!qf.query(&42));
    }

    #[test]
    fn full() {
        let mut qf = QuotientFilter::with_params(3, 16);
        for i in 0..8 {
            qf.insert(&i).unwrap();
            for j in 0..i {
                assert!(qf.query(&j), "Cannot find {} after inserting {}", j, i);
            }
        }
        assert!(qf.insert(&1000).is_err());
    }

    #[test]
    fn clear() {
        let mut qf = QuotientFilter::with_params(3, 16);
        qf.insert(&13).unwrap();
        qf.clear();
        assert!(qf.is_empty());
        assert_eq!(qf.len(), 0);
        assert!(!qf.query(&13));
        assert_eq!(qf.bits_quotient(), 3);
        assert_eq!(qf.bits_remainder(), 16);
    }

    #[test]
    fn clone() {
        let mut qf1 = QuotientFilter::with_params(3, 16);
        qf1.insert(&13).unwrap();

        let mut qf2 = qf1.clone();
        qf2.insert(&42).unwrap();

        assert_eq!(qf1.len(), 1);
        assert!(qf1.query(&13));
        assert!(!qf1.query(&42));

        assert_eq!(qf2.len(), 2);
        assert!(qf2.query(&13));
        assert!(qf2.query(&42));
    }
}
<|endoftext|>


<commit_before>//! The module exposes the `H2Client` struct, which implements a futures-based API for performing
//! HTTP/2 requests, based on Tokio.

use super::{HttpRequestHeaders, HttpRequestBody, HttpResponseHeaders, HttpResponseBody};
use client::tokio_layer::{H2ClientTokioProto};

use std::io::{self};
use std::net::SocketAddr;
use std::iter::{self, IntoIterator};

use futures::{Async, Future, Poll};
use futures::future::{BoxFuture};
use futures::stream::{self, Stream};

use tokio_core::reactor::{Handle};
use tokio_proto::{Connect, TcpClient};
use tokio_proto::streaming::{Message, Body};
use tokio_proto::streaming::multiplex::{StreamingMultiplex};
use tokio_proto::util::client_proxy::ClientProxy;

use tokio_service::{Service};

use solicit::http::{Header, StaticHeader};

/// A type alias for the request body stream. For now, we assume that all requests are a
/// single-chunk stream.
type RequestBodyStream = stream::Once<HttpRequestBody, io::Error>;

/// A type alias for the `ClientProxy` that we end up building after an `Io` is bound to a
/// `H2ClientTokioTransport` by `H2ClientTokioProto`.
type TokioClient =
    ClientProxy<
        Message<HttpRequestHeaders, RequestBodyStream>,
        Message<HttpResponseHeaders, Body<HttpResponseBody, io::Error>>,
        io::Error>;

/// A `futures::Stream` impl that represents the body of the response. The `Future` returned
/// by various `H2Client` methods returns an instance of this type, along with the response
/// headers.
pub struct ResponseBodyStream {
    /// The type simply hides away the Tokio `Body`, which will be returned by Tokio client
    /// Service.
    inner: Body<HttpResponseBody, io::Error>,
}

impl ResponseBodyStream {
    fn new(inner: Body<HttpResponseBody, io::Error>) -> ResponseBodyStream {
        ResponseBodyStream {
            inner: inner,
        }
    }
}

impl Stream for ResponseBodyStream {
    type Item = HttpResponseBody;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.inner.poll()
    }
}

/// A type alias for the future that is produced by the `H2Client`'s various `request` methods.
/// (`request`, `get`, `post`, ...)
type FutureH2Response = BoxFuture<(HttpResponseHeaders, ResponseBodyStream), io::Error>;

/// A struct that implements a futures-based API for an HTTP/2 client.
pub struct H2Client {
    /// The inner ClientProxy that hooks into the whole Tokio infrastructure.
    inner: TokioClient,
    /// The authority header (nee Host). Specifies the host name that the HTTP requests will
    /// be directed at. This is distinct from the socket address.
    authority: Vec<u8>,
}

impl H2Client {
    /// Creates a new `H2Client` from the given `TokioClient`.
    fn new(inner: TokioClient, authority: Vec<u8>) -> H2Client {
        H2Client {
            inner: inner,
            authority: authority,
        }
    }

    /// Connect to the given socket and yield a new `H2Client` that can be used to send HTTP/2
    /// requests to this socket.
    ///
    /// Returns a future that will resolve to the `H2Client`.
    pub fn connect(authority: &str, socket_addr: &SocketAddr, handle: &Handle) -> H2ClientNew {
        let proto = H2ClientTokioProto;

        let client = TcpClient::<StreamingMultiplex<RequestBodyStream>, _>::new(proto);
        let connect = client.connect(&socket_addr, &handle);

        H2ClientNew::new(connect, authority.as_bytes().to_vec())
    }

    /// Issues a GET request to the server.
    ///
    /// Yields a future that resolves to an `HttpRequestHeaders` struct. This struct will carry
    /// both the response headers, as well as the response body.
    ///
    /// TODO: WTB `-> impl` to avoid boxing this (or manually implementing a Future)
    pub fn get(&mut self, path: &[u8]) -> FutureH2Response {
        self.request(b"GET", path, iter::empty(), None)
    }

    /// Issues a POST request, carrying the given body.
    pub fn post(&mut self, path: &[u8], body: Vec<u8>) -> FutureH2Response {
        self.request(b"POST", path, iter::empty(), Some(body))
    }

    /// Perform a request, providing manually the request method, headers, and body.
    pub fn request<I>(&mut self,
                      method: &[u8],
                      path: &[u8],
                      user_headers: I,
                      body: Option<Vec<u8>>)
                      -> FutureH2Response
                      where I: IntoIterator<Item=StaticHeader> {
        let mut headers = Vec::new();
        headers.extend(vec![
            Header::new(b":method", method.to_vec()),
            Header::new(b":path", path.to_vec()),
            Header::new(b":authority", self.authority.clone()),
            Header::new(b":scheme", b"http"),
        ].into_iter());
        headers.extend(user_headers.into_iter());

        self.request_with_vec(headers, body)
    }

    /// Actually performs the full request. Avoids monomorphizing the entire code, but rather only
    /// the bit that requires the use of the IntoIterator trait, before passing off to this.
    fn request_with_vec(&mut self,
                        headers: Vec<StaticHeader>,
                        body: Option<Vec<u8>>)
                        -> FutureH2Response {

        let request_headers = HttpRequestHeaders::with_headers(headers);
        let tokio_req = match body {
            None => Message::WithoutBody(request_headers),
            Some(body) => {
                let body_stream = stream::once::<_, io::Error>(Ok(HttpRequestBody::new(body)));
                Message::WithBody(request_headers, body_stream)
            },
        };

        let response_future = Service::call(&self.inner, tokio_req).map(|response| {
            debug!("resolved response message");

            match response {
                Message::WithoutBody(resp @ HttpResponseHeaders { .. }) => {
                    // If there's no body, just yield an empty body stream.
                    (resp, ResponseBodyStream::new(Body::empty()))
                },
                Message::WithBody(resp @ HttpResponseHeaders { .. }, body) => {
                    (resp, ResponseBodyStream::new(body))
                },
            }
        });

        response_future.boxed()
    }
}

/// A simple `Future` implementation that resolves once the HTTP/2 client connection is
/// established.
pub struct H2ClientNew {
    /// The future that resolves to a new Tokio ClientProxy.
    inner: Connect<StreamingMultiplex<RequestBodyStream>, H2ClientTokioProto>,
    /// The authority that the new client will send requests to.
    authority: Option<Vec<u8>>,
}

impl H2ClientNew {
    fn new(connect: Connect<StreamingMultiplex<RequestBodyStream>, H2ClientTokioProto>,
           authority: Vec<u8>)
           -> H2ClientNew {
        H2ClientNew {
            inner: connect,
            authority: Some(authority),
        }
    }
}

impl Future for H2ClientNew {
    type Item = H2Client;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        trace!("waiting for client connection...");

        match self.inner.poll() {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready(client_proxy)) => {
                trace!("client connected");
                let authority = self.authority.take().expect("H2ClientNew future polled again");
                Ok(Async::Ready(H2Client::new(client_proxy, authority)))
            },
            Err(e) => Err(e),
        }
    }
}
<commit_msg>client: Refactor returned Future to be a wrapper struct<commit_after>//! The module exposes the `H2Client` struct, which implements a futures-based API for performing
//! HTTP/2 requests, based on Tokio.

use super::{HttpRequestHeaders, HttpRequestBody, HttpResponseHeaders, HttpResponseBody};
use client::tokio_layer::{H2ClientTokioProto};

use std::io::{self};
use std::net::SocketAddr;
use std::iter::{self, IntoIterator};

use futures::{Async, Future, Poll};
use futures::future::{BoxFuture};
use futures::stream::{self, Stream};

use tokio_core::reactor::{Handle};
use tokio_proto::{Connect, TcpClient};
use tokio_proto::streaming::{Message, Body};
use tokio_proto::streaming::multiplex::{StreamingMultiplex};
use tokio_proto::util::client_proxy::ClientProxy;

use tokio_service::{Service};

use solicit::http::{Header, StaticHeader};

/// A type alias for the request body stream. For now, we assume that all requests are a
/// single-chunk stream.
type RequestBodyStream = stream::Once<HttpRequestBody, io::Error>;

/// A type alias for the `ClientProxy` that we end up building after an `Io` is bound to a
/// `H2ClientTokioTransport` by `H2ClientTokioProto`.
type TokioClient =
    ClientProxy<
        Message<HttpRequestHeaders, RequestBodyStream>,
        Message<HttpResponseHeaders, Body<HttpResponseBody, io::Error>>,
        io::Error>;

/// A `futures::Stream` impl that represents the body of the response. The `Future` returned
/// by various `H2Client` methods returns an instance of this type, along with the response
/// headers.
pub struct ResponseBodyStream {
    /// The type simply hides away the Tokio `Body`, which will be returned by Tokio client
    /// Service.
    inner: Body<HttpResponseBody, io::Error>,
}

impl ResponseBodyStream {
    fn new(inner: Body<HttpResponseBody, io::Error>) -> ResponseBodyStream {
        ResponseBodyStream {
            inner: inner,
        }
    }
}

impl Stream for ResponseBodyStream {
    type Item = HttpResponseBody;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.inner.poll()
    }
}

/// A `Future` produced by the `H2Client`'s various `request` methods.
/// (`request`, `get`, `post`, ...)
pub struct FutureH2Response {
    /// Simply wraps a boxed future
    inner: BoxFuture<(HttpResponseHeaders, ResponseBodyStream), io::Error>,
}

impl Future for FutureH2Response {
    type Item = (HttpResponseHeaders, ResponseBodyStream);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.inner.poll()
    }
}

impl FutureH2Response {
    /// Creates a new `FutureH2Response` wrapping the given boxed future.
    fn new(inner: BoxFuture<(HttpResponseHeaders, ResponseBodyStream), io::Error>)
            -> FutureH2Response {
        FutureH2Response {
            inner: inner,
        }
    }
}

/// A struct that implements a futures-based API for an HTTP/2 client.
pub struct H2Client {
    /// The inner ClientProxy that hooks into the whole Tokio infrastructure.
    inner: TokioClient,
    /// The authority header (nee Host). Specifies the host name that the HTTP requests will
    /// be directed at. This is distinct from the socket address.
    authority: Vec<u8>,
}

impl H2Client {
    /// Creates a new `H2Client` from the given `TokioClient`.
    fn new(inner: TokioClient, authority: Vec<u8>) -> H2Client {
        H2Client {
            inner: inner,
            authority: authority,
        }
    }

    /// Connect to the given socket and yield a new `H2Client` that can be used to send HTTP/2
    /// requests to this socket.
    ///
    /// Returns a future that will resolve to the `H2Client`.
    pub fn connect(authority: &str, socket_addr: &SocketAddr, handle: &Handle) -> H2ClientNew {
        let proto = H2ClientTokioProto;

        let client = TcpClient::<StreamingMultiplex<RequestBodyStream>, _>::new(proto);
        let connect = client.connect(&socket_addr, &handle);

        H2ClientNew::new(connect, authority.as_bytes().to_vec())
    }

    /// Issues a GET request to the server.
    ///
    /// Yields a future that resolves to an `HttpRequestHeaders` struct. This struct will carry
    /// both the response headers, as well as the response body.
    pub fn get(&mut self, path: &[u8]) -> FutureH2Response {
        self.request(b"GET", path, iter::empty(), None)
    }

    /// Issues a POST request, carrying the given body.
    pub fn post(&mut self, path: &[u8], body: Vec<u8>) -> FutureH2Response {
        self.request(b"POST", path, iter::empty(), Some(body))
    }

    /// Perform a request, providing manually the request method, headers, and body.
    pub fn request<I>(&mut self,
                      method: &[u8],
                      path: &[u8],
                      user_headers: I,
                      body: Option<Vec<u8>>)
                      -> FutureH2Response
                      where I: IntoIterator<Item=StaticHeader> {
        let mut headers = Vec::new();
        headers.extend(vec![
            Header::new(b":method", method.to_vec()),
            Header::new(b":path", path.to_vec()),
            Header::new(b":authority", self.authority.clone()),
            Header::new(b":scheme", b"http"),
        ].into_iter());
        headers.extend(user_headers.into_iter());

        self.request_with_vec(headers, body)
    }

    /// Actually performs the full request. Avoids monomorphizing the entire code, but rather only
    /// the bit that requires the use of the IntoIterator trait, before passing off to this.
    fn request_with_vec(&mut self,
                        headers: Vec<StaticHeader>,
                        body: Option<Vec<u8>>)
                        -> FutureH2Response {

        let request_headers = HttpRequestHeaders::with_headers(headers);
        let tokio_req = match body {
            None => Message::WithoutBody(request_headers),
            Some(body) => {
                let body_stream = stream::once::<_, io::Error>(Ok(HttpRequestBody::new(body)));
                Message::WithBody(request_headers, body_stream)
            },
        };

        let response_future = Service::call(&self.inner, tokio_req).map(|response| {
            debug!("resolved response message");

            match response {
                Message::WithoutBody(resp @ HttpResponseHeaders { .. }) => {
                    // If there's no body, just yield an empty body stream.
                    (resp, ResponseBodyStream::new(Body::empty()))
                },
                Message::WithBody(resp @ HttpResponseHeaders { .. }, body) => {
                    (resp, ResponseBodyStream::new(body))
                },
            }
        });

        FutureH2Response::new(response_future.boxed())
    }
}

/// A simple `Future` implementation that resolves once the HTTP/2 client connection is
/// established.
pub struct H2ClientNew {
    /// The future that resolves to a new Tokio ClientProxy.
    inner: Connect<StreamingMultiplex<RequestBodyStream>, H2ClientTokioProto>,
    /// The authority that the new client will send requests to.
    authority: Option<Vec<u8>>,
}

impl H2ClientNew {
    fn new(connect: Connect<StreamingMultiplex<RequestBodyStream>, H2ClientTokioProto>,
           authority: Vec<u8>)
           -> H2ClientNew {
        H2ClientNew {
            inner: connect,
            authority: Some(authority),
        }
    }
}

impl Future for H2ClientNew {
    type Item = H2Client;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        trace!("waiting for client connection...");

        match self.inner.poll() {
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Ok(Async::Ready(client_proxy)) => {
                trace!("client connected");
                let authority = self.authority.take().expect("H2ClientNew future polled again");
                Ok(Async::Ready(H2Client::new(client_proxy, authority)))
            },
            Err(e) => Err(e),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Provides an implementation of the audio_streams::shm_streams::ShmStream trait using the VioS
//! client.
//! Given that the VioS server doesn't emit an event when the next buffer is expected, this
//! implementation uses thread::sleep to drive the frame timings.

use super::shm_vios::{VioSClient, VioSStreamParams};

use crate::virtio::snd::constants::*;

use audio_streams::shm_streams::{BufferSet, ServerRequest, ShmStream, ShmStreamSource};
use audio_streams::{BoxError, SampleFormat, StreamDirection, StreamEffect};

use base::{error, SharedMemory, SharedMemoryUnix};

use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use sys_util::{Error as SysError, SharedMemory as SysSharedMemory};

use super::shm_vios::{Error, Result};

// This is the error type used in audio_streams::shm_streams. Unfortunately, it's not declared
// public there so it needs to be re-declared here. It also prevents the usage of anyhow::Error.
type GenericResult<T> = std::result::Result<T, BoxError>;

/// Adapter that provides the ShmStreamSource trait around the VioS backend.
pub struct VioSShmStreamSource {
    vios_client: Arc<VioSClient>,
}

impl VioSShmStreamSource {
    /// Creates a new stream source given the path to the audio server's socket.
    pub fn new<P: AsRef<Path>>(server: P) -> Result<VioSShmStreamSource> {
        Ok(Self {
            vios_client: Arc::new(VioSClient::try_new(server)?),
        })
    }
}

fn from_sample_format(format: SampleFormat) -> u8 {
    match format {
        SampleFormat::U8 => VIRTIO_SND_PCM_FMT_U8,
        SampleFormat::S16LE => VIRTIO_SND_PCM_FMT_U16,
        SampleFormat::S24LE => VIRTIO_SND_PCM_FMT_U24,
        SampleFormat::S32LE => VIRTIO_SND_PCM_FMT_U32,
    }
}

fn virtio_frame_rate(frame_rate: u32) -> GenericResult<u8> {
    Ok(match frame_rate {
        5512u32 => VIRTIO_SND_PCM_RATE_5512,
        8000u32 => VIRTIO_SND_PCM_RATE_8000,
        11025u32 => VIRTIO_SND_PCM_RATE_11025,
        16000u32 => VIRTIO_SND_PCM_RATE_16000,
        22050u32 => VIRTIO_SND_PCM_RATE_22050,
        32000u32 => VIRTIO_SND_PCM_RATE_32000,
        44100u32 => VIRTIO_SND_PCM_RATE_44100,
        48000u32 => VIRTIO_SND_PCM_RATE_48000,
        64000u32 => VIRTIO_SND_PCM_RATE_64000,
        88200u32 => VIRTIO_SND_PCM_RATE_88200,
        96000u32 => VIRTIO_SND_PCM_RATE_96000,
        176400u32 => VIRTIO_SND_PCM_RATE_176400,
        192000u32 => VIRTIO_SND_PCM_RATE_192000,
        384000u32 => VIRTIO_SND_PCM_RATE_384000,
        _ => {
            return Err(Box::new(Error::UnsupportedFrameRate(frame_rate)));
        }
    })
}

impl ShmStreamSource for VioSShmStreamSource {
    /// Creates a new stream
    #[allow(clippy::too_many_arguments)]
    fn new_stream(
        &mut self,
        direction: StreamDirection,
        num_channels: usize,
        format: SampleFormat,
        frame_rate: u32,
        buffer_size: usize,
        _effects: &[StreamEffect],
        client_shm: &SysSharedMemory,
        _buffer_offsets: [u64; 2],
    ) -> GenericResult<Box<dyn ShmStream>> {
        self.vios_client.ensure_bg_thread_started()?;
        let virtio_dir = match direction {
            StreamDirection::Playback => VIRTIO_SND_D_OUTPUT,
            StreamDirection::Capture => VIRTIO_SND_D_INPUT,
        };
        let frame_size = num_channels * format.sample_bytes();
        let period_bytes = (frame_size * buffer_size) as u32;
        let stream_id = self
            .vios_client
            .get_unused_stream_id(virtio_dir)
            .ok_or(Box::new(Error::NoStreamsAvailable))?;
        // Create the stream object before any errors can be returned to guarantee the stream will
        // be released in all cases
        let stream_box = VioSndShmStream::new(
            buffer_size,
            num_channels,
            format,
            frame_rate,
            stream_id,
            direction,
            self.vios_client.clone(),
            client_shm,
        );
        self.vios_client.prepare_stream(stream_id)?;
        let params = VioSStreamParams {
            buffer_bytes: 2 * period_bytes,
            period_bytes,
            features: 0u32,
            channels: num_channels as u8,
            format: from_sample_format(format),
            rate: virtio_frame_rate(frame_rate)?,
        };
        self.vios_client.set_stream_parameters(stream_id, params)?;
        self.vios_client.start_stream(stream_id)?;
        stream_box
    }

    /// Get a list of file descriptors used by the implementation.
    ///
    /// Returns any open file descriptors needed by the implementation.
    /// This list helps users of the ShmStreamSource enter Linux jails without
    /// closing needed file descriptors.
    fn keep_fds(&self) -> Vec<RawFd> {
        self.vios_client.keep_fds()
    }
}

/// Adapter around a VioS stream that implements the ShmStream trait.
pub struct VioSndShmStream {
    num_channels: usize,
    frame_rate: u32,
    buffer_size: usize,
    frame_size: usize,
    interval: Duration,
    next_frame: Duration,
    start_time: Instant,
    stream_id: u32,
    direction: StreamDirection,
    vios_client: Arc<VioSClient>,
    client_shm: SharedMemory,
}

impl VioSndShmStream {
    /// Creates a new shm stream.
    fn new(
        buffer_size: usize,
        num_channels: usize,
        format: SampleFormat,
        frame_rate: u32,
        stream_id: u32,
        direction: StreamDirection,
        vios_client: Arc<VioSClient>,
        client_shm: &SysSharedMemory,
    ) -> GenericResult<Box<dyn ShmStream>> {
        let interval = Duration::from_millis(buffer_size as u64 * 1000 / frame_rate as u64);

        let dup_fd = unsafe {
            // Safe because fcntl doesn't affect memory and client_shm should wrap a known valid
            // file descriptor.
            libc::fcntl(client_shm.as_raw_fd(), libc::F_DUPFD_CLOEXEC, 0)
        };
        if dup_fd < 0 {
            return Err(Box::new(Error::DupError(SysError::last())));
        }
        let file = unsafe {
            // safe because we checked the result of libc::fcntl()
            File::from_raw_fd(dup_fd)
        };
        let client_shm_clone =
            SharedMemory::from_file(file).map_err(|e| Error::BaseMmapError(e))?;

        Ok(Box::new(Self {
            num_channels,
            frame_rate,
            buffer_size,
            frame_size: format.sample_bytes() * num_channels,
            interval,
            next_frame: interval,
            start_time: Instant::now(),
            stream_id,
            direction,
            vios_client,
            client_shm: client_shm_clone,
        }))
    }
}

impl ShmStream for VioSndShmStream {
    fn frame_size(&self) -> usize {
        self.frame_size
    }

    fn num_channels(&self) -> usize {
        self.num_channels
    }

    fn frame_rate(&self) -> u32 {
        self.frame_rate
    }

    /// Waits until the next time a frame should be sent to the server. The server may release the
    /// previous buffer much sooner than it needs the next one, so this function may sleep to wait
    /// for the right time.
    fn wait_for_next_action_with_timeout(
        &mut self,
        timeout: Duration,
    ) -> GenericResult<Option<ServerRequest>> {
        let elapsed = self.start_time.elapsed();
        if elapsed < self.next_frame {
            if timeout < self.next_frame - elapsed {
                std::thread::sleep(timeout);
                return Ok(None);
            } else {
                std::thread::sleep(self.next_frame - elapsed);
            }
        }
        self.next_frame += self.interval;
        Ok(Some(ServerRequest::new(self.buffer_size, self)))
    }
}

impl BufferSet for VioSndShmStream {
    fn callback(&mut self, offset: usize, frames: usize) -> GenericResult<()> {
        match self.direction {
            StreamDirection::Playback => {
                self.vios_client.inject_audio_data(
                    self.stream_id,
                    &mut self.client_shm,
                    offset,
                    frames * self.frame_size,
                )?;
            }
            StreamDirection::Capture => {
                self.vios_client.request_audio_data(
                    self.stream_id,
                    &mut self.client_shm,
                    offset,
                    frames * self.frame_size,
                )?;
            }
        }
        Ok(())
    }

    fn ignore(&mut self) -> GenericResult<()> {
        Ok(())
    }
}

impl Drop for VioSndShmStream {
    fn drop(&mut self) {
        let stream_id = self.stream_id;
        if let Err(e) = self
            .vios_client
            .stop_stream(stream_id)
            .and_then(|_| self.vios_client.release_stream(stream_id))
        {
            error!("Failed to stop and release stream {}: {}", stream_id, e);
        }
    }
}
<commit_msg>Error on stream creation no longer blocks stream ids<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Provides an implementation of the audio_streams::shm_streams::ShmStream trait using the VioS
//! client.
//! Given that the VioS server doesn't emit an event when the next buffer is expected, this
//! implementation uses thread::sleep to drive the frame timings.

use super::shm_vios::{VioSClient, VioSStreamParams};

use crate::virtio::snd::constants::*;

use audio_streams::shm_streams::{BufferSet, ServerRequest, ShmStream, ShmStreamSource};
use audio_streams::{BoxError, SampleFormat, StreamDirection, StreamEffect};

use base::{error, SharedMemory, SharedMemoryUnix};

use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use sys_util::{Error as SysError, SharedMemory as SysSharedMemory};

use super::shm_vios::{Error, Result};

// This is the error type used in audio_streams::shm_streams. Unfortunately, it's not declared
// public there so it needs to be re-declared here. It also prevents the usage of anyhow::Error.
type GenericResult<T> = std::result::Result<T, BoxError>;

/// Adapter that provides the ShmStreamSource trait around the VioS backend.
pub struct VioSShmStreamSource {
    vios_client: Arc<VioSClient>,
}

impl VioSShmStreamSource {
    /// Creates a new stream source given the path to the audio server's socket.
    pub fn new<P: AsRef<Path>>(server: P) -> Result<VioSShmStreamSource> {
        Ok(Self {
            vios_client: Arc::new(VioSClient::try_new(server)?),
        })
    }
}

fn from_sample_format(format: SampleFormat) -> u8 {
    match format {
        SampleFormat::U8 => VIRTIO_SND_PCM_FMT_U8,
        SampleFormat::S16LE => VIRTIO_SND_PCM_FMT_U16,
        SampleFormat::S24LE => VIRTIO_SND_PCM_FMT_U24,
        SampleFormat::S32LE => VIRTIO_SND_PCM_FMT_U32,
    }
}

fn virtio_frame_rate(frame_rate: u32) -> GenericResult<u8> {
    Ok(match frame_rate {
        5512u32 => VIRTIO_SND_PCM_RATE_5512,
        8000u32 => VIRTIO_SND_PCM_RATE_8000,
        11025u32 => VIRTIO_SND_PCM_RATE_11025,
        16000u32 => VIRTIO_SND_PCM_RATE_16000,
        22050u32 => VIRTIO_SND_PCM_RATE_22050,
        32000u32 => VIRTIO_SND_PCM_RATE_32000,
        44100u32 => VIRTIO_SND_PCM_RATE_44100,
        48000u32 => VIRTIO_SND_PCM_RATE_48000,
        64000u32 => VIRTIO_SND_PCM_RATE_64000,
        88200u32 => VIRTIO_SND_PCM_RATE_88200,
        96000u32 => VIRTIO_SND_PCM_RATE_96000,
        176400u32 => VIRTIO_SND_PCM_RATE_176400,
        192000u32 => VIRTIO_SND_PCM_RATE_192000,
        384000u32 => VIRTIO_SND_PCM_RATE_384000,
        _ => {
            return Err(Box::new(Error::UnsupportedFrameRate(frame_rate)));
        }
    })
}

impl VioSShmStreamSource {
    fn new_stream_inner(
        &mut self,
        stream_id: u32,
        direction: StreamDirection,
        num_channels: usize,
        format: SampleFormat,
        frame_rate: u32,
        buffer_size: usize,
        _effects: &[StreamEffect],
        client_shm: &SysSharedMemory,
        _buffer_offsets: [u64; 2],
    ) -> GenericResult<Box<dyn ShmStream>> {
        let frame_size = num_channels * format.sample_bytes();
        let period_bytes = (frame_size * buffer_size) as u32;
        self.vios_client.prepare_stream(stream_id)?;
        let params = VioSStreamParams {
            buffer_bytes: 2 * period_bytes,
            period_bytes,
            features: 0u32,
            channels: num_channels as u8,
            format: from_sample_format(format),
            rate: virtio_frame_rate(frame_rate)?,
        };
        self.vios_client.set_stream_parameters(stream_id, params)?;
        self.vios_client.start_stream(stream_id)?;
        VioSndShmStream::new(
            buffer_size,
            num_channels,
            format,
            frame_rate,
            stream_id,
            direction,
            self.vios_client.clone(),
            client_shm,
        )
    }
}

impl ShmStreamSource for VioSShmStreamSource {
    /// Creates a new stream
    #[allow(clippy::too_many_arguments)]
    fn new_stream(
        &mut self,
        direction: StreamDirection,
        num_channels: usize,
        format: SampleFormat,
        frame_rate: u32,
        buffer_size: usize,
        effects: &[StreamEffect],
        client_shm: &SysSharedMemory,
        buffer_offsets: [u64; 2],
    ) -> GenericResult<Box<dyn ShmStream>> {
        self.vios_client.ensure_bg_thread_started()?;
        let virtio_dir = match direction {
            StreamDirection::Playback => VIRTIO_SND_D_OUTPUT,
            StreamDirection::Capture => VIRTIO_SND_D_INPUT,
        };
        let stream_id = self
            .vios_client
            .get_unused_stream_id(virtio_dir)
            .ok_or(Box::new(Error::NoStreamsAvailable))?;
        self.new_stream_inner(
            stream_id,
            direction,
            num_channels,
            format,
            frame_rate,
            buffer_size,
            effects,
            client_shm,
            buffer_offsets,
        )
        .map_err(|e| {
            // Attempt to release the stream so that it can be used later. This is a best effort
            // attempt, so we ignore any error it may return.
            let _ = self.vios_client.release_stream(stream_id);
            e
        })
    }

    /// Get a list of file descriptors used by the implementation.
    ///
    /// Returns any open file descriptors needed by the implementation.
    /// This list helps users of the ShmStreamSource enter Linux jails without
    /// closing needed file descriptors.
    fn keep_fds(&self) -> Vec<RawFd> {
        self.vios_client.keep_fds()
    }
}

/// Adapter around a VioS stream that implements the ShmStream trait.
pub struct VioSndShmStream {
    num_channels: usize,
    frame_rate: u32,
    buffer_size: usize,
    frame_size: usize,
    interval: Duration,
    next_frame: Duration,
    start_time: Instant,
    stream_id: u32,
    direction: StreamDirection,
    vios_client: Arc<VioSClient>,
    client_shm: SharedMemory,
}

impl VioSndShmStream {
    /// Creates a new shm stream.
    fn new(
        buffer_size: usize,
        num_channels: usize,
        format: SampleFormat,
        frame_rate: u32,
        stream_id: u32,
        direction: StreamDirection,
        vios_client: Arc<VioSClient>,
        client_shm: &SysSharedMemory,
    ) -> GenericResult<Box<dyn ShmStream>> {
        let interval = Duration::from_millis(buffer_size as u64 * 1000 / frame_rate as u64);

        let dup_fd = unsafe {
            // Safe because fcntl doesn't affect memory and client_shm should wrap a known valid
            // file descriptor.
            libc::fcntl(client_shm.as_raw_fd(), libc::F_DUPFD_CLOEXEC, 0)
        };
        if dup_fd < 0 {
            return Err(Box::new(Error::DupError(SysError::last())));
        }
        let file = unsafe {
            // safe because we checked the result of libc::fcntl()
            File::from_raw_fd(dup_fd)
        };
        let client_shm_clone =
            SharedMemory::from_file(file).map_err(|e| Error::BaseMmapError(e))?;

        Ok(Box::new(Self {
            num_channels,
            frame_rate,
            buffer_size,
            frame_size: format.sample_bytes() * num_channels,
            interval,
            next_frame: interval,
            start_time: Instant::now(),
            stream_id,
            direction,
            vios_client,
            client_shm: client_shm_clone,
        }))
    }
}

impl ShmStream for VioSndShmStream {
    fn frame_size(&self) -> usize {
        self.frame_size
    }

    fn num_channels(&self) -> usize {
        self.num_channels
    }

    fn frame_rate(&self) -> u32 {
        self.frame_rate
    }

    /// Waits until the next time a frame should be sent to the server. The server may release the
    /// previous buffer much sooner than it needs the next one, so this function may sleep to wait
    /// for the right time.
    fn wait_for_next_action_with_timeout(
        &mut self,
        timeout: Duration,
    ) -> GenericResult<Option<ServerRequest>> {
        let elapsed = self.start_time.elapsed();
        if elapsed < self.next_frame {
            if timeout < self.next_frame - elapsed {
                std::thread::sleep(timeout);
                return Ok(None);
            } else {
                std::thread::sleep(self.next_frame - elapsed);
            }
        }
        self.next_frame += self.interval;
        Ok(Some(ServerRequest::new(self.buffer_size, self)))
    }
}

impl BufferSet for VioSndShmStream {
    fn callback(&mut self, offset: usize, frames: usize) -> GenericResult<()> {
        match self.direction {
            StreamDirection::Playback => {
                self.vios_client.inject_audio_data(
                    self.stream_id,
                    &mut self.client_shm,
                    offset,
                    frames * self.frame_size,
                )?;
            }
            StreamDirection::Capture => {
                self.vios_client.request_audio_data(
                    self.stream_id,
                    &mut self.client_shm,
                    offset,
                    frames * self.frame_size,
                )?;
            }
        }
        Ok(())
    }

    fn ignore(&mut self) -> GenericResult<()> {
        Ok(())
    }
}

impl Drop for VioSndShmStream {
    fn drop(&mut self) {
        let stream_id = self.stream_id;
        if let Err(e) = self
            .vios_client
            .stop_stream(stream_id)
            .and_then(|_| self.vios_client.release_stream(stream_id))
        {
            error!("Failed to stop and release stream {}: {}", stream_id, e);
        }
    }
}
<|endoftext|>


<commit_before>
use std::env;
use std::process::{Command, Stdio};
use std::io::{Result, Write};

macro_rules! try_log {
    ($exp:expr) => {
        match $exp {
            Ok(v) => v,
            Err(e) => { println!("<3>{}", e); return; }
        }
    }
}

fn get_systemd_unit_property(unit: &str, prop: &str) -> Result<String> {
    Command::new("systemctl")
        .arg("show")
        .arg(unit)
        .arg("--property")
        .arg(prop)
        .output()
        .map(|out| {
            String::from_utf8_lossy(&out.stdout[prop.len() + 1..])
                .trim_right_matches('\n')
                .to_owned()
        })
}

fn main() {
    let unit = match env::args().nth(1) {
        Some(unit) => unit,
        None => {
            println!("Usage: mail-on-failure <unit>");
            return;
        }
    };

    let mut user = try_log!(get_systemd_unit_property(&*unit, "User"));
    if user.len() == 0 {
        user = "root".to_owned();
    }

    let job_env = try_log!(get_systemd_unit_property(&*unit, "Environment"));
    for pair in job_env.split(' ') {
        let mut p = pair.splitn(2, '=');
        if let (Some(name), Some(value)) = (p.next(), p.next()) {
            if name == "MAILTO" {
                user = value.to_owned();
                break;
            }
        }
    }

    if user.len() == 0 {
        return;
    }

    let mut hostname = String::from_utf8_lossy(&try_log!(Command::new("uname")
                .arg("-n")
                .output())
            .stdout[..])
        .trim_right_matches('\n')
        .to_owned();

    if hostname.len() == 0 {
        hostname = "localhost".to_owned();
    }

    let mut head = String::new();
    head.push_str("From: root (systemd-cron)\nTo: ");
    head.push_str(&*user);
    head.push_str("\nSubject: [");
    head.push_str(&*hostname);
    head.push_str("] job ");
    head.push_str(&*unit);
    head.push_str(r###" failed
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Auto-Submitted: auto-generated

"###);

    let status = Command::new("systemctl")
        .arg("status")
        .arg(&*unit)
        .output()
        .unwrap();

    let mut mailer = try_log!(Command::new("sendmail")
        .arg("-i")
        .arg("-B8BITMIME")
        .arg(&*user)
        .stdin(Stdio::piped())
        .spawn());

    if let Some(ref mut stdin) = mailer.stdin {
        try_log!(stdin.write_all(head.as_bytes()).and_then(|_| stdin.write_all(&*status.stdout)));
    }

    mailer.wait().unwrap();
}
<commit_msg>replace deprecated trim_right_matches with trim_end_matches<commit_after>
use std::env;
use std::process::{Command, Stdio};
use std::io::{Result, Write};

macro_rules! try_log {
    ($exp:expr) => {
        match $exp {
            Ok(v) => v,
            Err(e) => { println!("<3>{}", e); return; }
        }
    }
}

fn get_systemd_unit_property(unit: &str, prop: &str) -> Result<String> {
    Command::new("systemctl")
        .arg("show")
        .arg(unit)
        .arg("--property")
        .arg(prop)
        .output()
        .map(|out| {
            String::from_utf8_lossy(&out.stdout[prop.len() + 1..])
                .trim_end_matches('\n')
                .to_owned()
        })
}

fn main() {
    let unit = match env::args().nth(1) {
        Some(unit) => unit,
        None => {
            println!("Usage: mail-on-failure <unit>");
            return;
        }
    };

    let mut user = try_log!(get_systemd_unit_property(&*unit, "User"));
    if user.len() == 0 {
        user = "root".to_owned();
    }

    let job_env = try_log!(get_systemd_unit_property(&*unit, "Environment"));
    for pair in job_env.split(' ') {
        let mut p = pair.splitn(2, '=');
        if let (Some(name), Some(value)) = (p.next(), p.next()) {
            if name == "MAILTO" {
                user = value.to_owned();
                break;
            }
        }
    }

    if user.len() == 0 {
        return;
    }

    let mut hostname = String::from_utf8_lossy(&try_log!(Command::new("uname")
                .arg("-n")
                .output())
            .stdout[..])
        .trim_end_matches('\n')
        .to_owned();

    if hostname.len() == 0 {
        hostname = "localhost".to_owned();
    }

    let mut head = String::new();
    head.push_str("From: root (systemd-cron)\nTo: ");
    head.push_str(&*user);
    head.push_str("\nSubject: [");
    head.push_str(&*hostname);
    head.push_str("] job ");
    head.push_str(&*unit);
    head.push_str(r###" failed
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Auto-Submitted: auto-generated

"###);

    let status = Command::new("systemctl")
        .arg("status")
        .arg(&*unit)
        .output()
        .unwrap();

    let mut mailer = try_log!(Command::new("sendmail")
        .arg("-i")
        .arg("-B8BITMIME")
        .arg(&*user)
        .stdin(Stdio::piped())
        .spawn());

    if let Some(ref mut stdin) = mailer.stdin {
        try_log!(stdin.write_all(head.as_bytes()).and_then(|_| stdin.write_all(&*status.stdout)));
    }

    mailer.wait().unwrap();
}
<|endoftext|>


<commit_before>use super::*;

impl Editor {
    /// Get pos of next char
    pub fn next_pos(&self) -> (usize, usize) {
        // TODO: Add numerals
        if self.x() == self.text[self.y()].len() {
            if self.y() < self.text.len() - 1 {
                (0, self.y() + 1)
            } else {
                (self.x(), self.y())
            }
        } else {
            (self.x() + 1, self.y())
        }
    }

    /// Get pos of previous char
    pub fn previous_pos(&self) -> (usize, usize) {
        if self.x() == 0 {
            if self.y() > 0 {
                (self.text[self.y()].len(), self.y() - 1)
            } else {
                (self.x(), self.y())
            }
        } else {
            (self.x() - 1, self.y())
        }
    }

    pub fn next(&mut self) {
        let (x, y) = self.next_pos();
        self.cursor_mut().x = x;
        self.cursor_mut().y = y;
    }
    pub fn previous(&mut self) {
        let (x, y) = self.previous_pos();
        self.cursor_mut().x = x;
        self.cursor_mut().y = y;
    }

    /// Go right
    pub fn right(&mut self, n: usize) {
        let x = self.x() + n;
        let y = self.y();

        let text = self.text.clone();
        let curs = self.cursor_mut();

        curs.x += n;

        if x > text[y].len() {
            curs.x = text[y].len();
        }
    }

    /// Go left
    pub fn left(&mut self, n: usize) {
        let x = self.x();
        let y = self.y();

        let text = self.text.clone();
        let curs = self.cursor_mut();

        if n <= x {
            curs.x = x - n;
        } else {
            curs.x = 0;
        }

    }

    /// Go up
    pub fn up(&mut self, n: usize) {
        let y = self.y();
        let curs = self.cursor_mut();
        if n <= y {
            curs.y -= n;
        } else {
            curs.y = 0;
        }
    }

    /// Go down
    pub fn down(&mut self, n: usize) {
        let x = self.x();
        let y = self.y() + n;

        let text = self.text.clone();
        let curs = self.cursor_mut();

        curs.y += n;

        if y >= text.len() {
            curs.y = text.len() - 1;
        }
    }
}
<commit_msg>Refactor in sodium<commit_after>use super::*;

impl Editor {
    /// Get pos of next char
    pub fn next_pos(&self) -> (usize, usize) {
        // TODO: Add numerals
        if self.x() == self.text[self.y()].len() {
            if self.y() < self.text.len() - 1 {
                (0, self.y() + 1)
            } else {
                (self.x(), self.y())
            }
        } else {
            (self.x() + 1, self.y())
        }
    }

    /// Get pos of previous char
    pub fn previous_pos(&self) -> (usize, usize) {
        if self.x() == 0 {
            if self.y() > 0 {
                (self.text[self.y()].len(), self.y() - 1)
            } else {
                (self.x(), self.y())
            }
        } else {
            (self.x() - 1, self.y())
        }
    }

    pub fn next(&mut self) {
        let (x, y) = self.next_pos();
        self.cursor_mut().x = x;
        self.cursor_mut().y = y;
    }
    pub fn previous(&mut self) {
        let (x, y) = self.previous_pos();
        self.cursor_mut().x = x;
        self.cursor_mut().y = y;
    }

    /// Get right pos
    pub fn right_pos(&self, n: usize) -> (usize, usize) {
        let x = self.x() + n;
        let y = self.y();

        if x > self.text[y].len() {
            (self.text[y].len(), y)
        } else {
            (x, y)
        }
    }
    pub fn right(&mut self, n: usize) {
        self.cursor_mut().x = self.right_pos(n).0;
    }

    /// Get left pos
    pub fn left_pos(&self, n: usize) -> (usize, usize) {
        let x = self.x();
        let y = self.y();

        if n <= x {
            (x - n, self.y())
        } else {
            (0, self.y())
        }

    }
    pub fn left(&mut self, n: usize) {
        self.cursor_mut().x = self.left_pos(n).0;
    }

    /// Go up
    pub fn up(&mut self, n: usize) {
        let y = self.y();
        let curs = self.cursor_mut();
        if n <= y {
            curs.y -= n;
        } else {
            curs.y = 0;
        }
    }

    /// Go down
    pub fn down(&mut self, n: usize) {
        let x = self.x();
        let y = self.y() + n;

        let text = self.text.clone();
        let curs = self.cursor_mut();

        curs.y += n;

        if y >= text.len() {
            curs.y = text.len() - 1;
        }
    }
}
<|endoftext|>


<commit_before>//! A Robin Hood hash map.

/// A Robin Hood hash map.
///
/// The hash map is parameterized by a supplied function acting as the hash function,
/// allowing the user more flexibility in defining the behavior of the map.
pub struct RHHMap<K: Eq, V, F: Fn(&K)->usize> {
    bucket: F,
    buffer: Vec<Option<(K,V)>>,
    shift:  usize,
    slop:   usize,
}

impl<K: Eq, V, F: Fn(&K)->usize> RHHMap<K, V, F> {
    /// Allocates a new Robin Hood hash map.
    pub fn new(function: F) -> RHHMap<K, V, F> {
        let slop = 4;
        let mut buffer = Vec::with_capacity(2 + slop);
        for _ in 0..buffer.capacity() { buffer.push(None); }
        RHHMap {
            bucket: function,
            buffer: buffer,
            shift: 63,
            slop: slop,
        }
    }
    /// The capacity of the underlying storage.
    pub fn capacity(&self) -> usize { self.buffer.capacity() }

    /// Get a reference to the value associated with a queried key.
    #[inline]
    pub fn get_ref<'a>(&'a self, query: &K) -> Option<&'a V> {
        let shift = self.shift;
        let bucket = &self.bucket;
        get_ref(&self.buffer[..], query, &|x| bucket(x), bucket(query) >> shift)
    }
    /// Get a mutable reference to the value associated with a queried key.
    #[inline]
    pub fn get_mut<'a>(&'a mut self, query: &K) -> Option<&'a mut V> {
        let shift = self.shift;
        let bucket = &self.bucket;
        get_mut(&mut self.buffer[..], query, &|x| bucket(x), bucket(query) >> shift)
    }
    /// Insert a key, value pair, recovering any pre-existing key, value pair.
    #[inline]
    pub fn insert(&mut self, key: K, val: V) -> Option<(K,V)> {

        let result = {
            let shift = self.shift;
            let bucket = &self.bucket;
            let location = bucket(&key) >> shift;
            insert(&mut self.buffer[..], key, val, &|x| bucket(x), location)
        };
        match result {
            Ok(result)      => result,
            Err((key, val)) => {

                let old_length = self.buffer.len() - self.slop;
                let mut new_buffer = Vec::with_capacity(2 * old_length + self.slop);
                for _ in 0..new_buffer.capacity() {
                    new_buffer.push(None);
                }
                let old_buffer = ::std::mem::replace(&mut self.buffer, new_buffer);
                self.shift -= 1;

                // let mut count = 0;
                let mut cursor = 0;
                for oldkeyval in old_buffer.into_iter() {
                    if let Some((oldkey, oldval)) = oldkeyval {
                        // count += 1;
                        let target = (self.bucket)(&oldkey) >> self.shift;
                        cursor = ::std::cmp::max(cursor, target);
                        self.buffer[cursor as usize] = Some((oldkey, oldval));
                        cursor += 1;
                    }
                }

                // println!("doubled length to {}; occupancy: {}", self.capacity() - self.slop, (count as f64) / ((self.capacity() - self.slop) as f64));

                self.insert(key, val)
            }
        }
    }
    /// Removes a key, recovering the key, value pair.
    pub fn remove(&mut self, key: &K) -> Option<(K,V)> {
        let shift = self.shift;
        let bucket = &self.bucket;
        remove(&mut self.buffer[..], key, &|x| bucket(x), bucket(key) >> shift)
    }
}

/// Returns either a shared reference to the value associated with `query`, or `None` if no value exists.
#[inline]
pub fn get_ref<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a [Option<(K,V)>], query: &K, function: &F, location: usize) -> Option<&'a V> {

    let target = function(query);
    let mut iterator = slice[location..].iter().map(|x| x.as_ref());
    while let Some(Some(&(ref key, ref val))) = iterator.next() {
        let found = function(key);
        if found >= target && key == query { return Some(val); }
        if found > target { return None; }
    }

    return None;
}

/// Returns either an exclusive reference to the value associated with `query`, or `None` if no value exists.
#[inline]
pub fn get_mut<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a mut [Option<(K,V)>], query: &K, function: &F, location: usize) -> Option<&'a mut V> {

    let target = function(query);
    let mut iterator = slice[location..].iter_mut().map(|x| x.as_mut());
    while let Some(Some(&mut (ref key, ref mut val))) = iterator.next() {
        let found = function(key);
        if found >= target && key == query { return Some(val); }
        if found > target { return None; }
    }

    return None;
}

#[derive(Eq, PartialEq)]
enum SearchStatus {
    Searching,
    Found(usize),
    Missing(usize),
}

/// Returns either the previous `(key,val)` pair for `key`, or `Err((key,val))` if `slice` does not have enough room to insert it.
#[inline]
pub fn insert<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a mut [Option<(K,V)>], key: K, val: V, function: &F, location: usize) -> Result<Option<(K,V)>,(K,V)> {

    let slop = 32;

    if location >= slice.len() {
        panic!("location {} > slice.len: {}", location, slice.len());
    }

    let target = function(&key);

    let mut position = location;
    let mut status = SearchStatus::Searching;
    let limit = ::std::cmp::min(location + slop, slice.len());
    while status == SearchStatus::Searching && position < limit {
        if let Some(ref mut kv) = slice[position].as_mut() {
            let found = function(&kv.0);
            if found == target && key == kv.0 {
                status = SearchStatus::Found(position);
            }
            if found > target { status = SearchStatus::Missing(position); }
        }
        else { status = SearchStatus::Missing(position); }

        position += 1;
    }

    match status {
        // if we bailed out of the search
        SearchStatus::Searching => {
            // println!("bailed out looking for element");
            Err((key, val))
        }
        // we found the key at `index`
        SearchStatus::Found(index) => {
            Ok(::std::mem::replace(&mut slice[index], Some((key, val))))
        }
        // we should have found the key by `index`, and should put the data there now
        SearchStatus::Missing(mut index) => {
            let begin = index;
            while index < slice.len() && slice[index].is_some() {
                index += 1;
            }

            // if we need to slide folks too far, is a bug
            if index > begin + 4 * slop {
                // println!("very long run: {}", index - begin);
                return Err((key, val));
            }

            if index < slice.len() {

                for i in 0..(index - begin) {
                    slice.swap(index - i - 1, index - i);
                }

                assert!(slice[begin].is_none());
                slice[begin] = Some((key, val));
                Ok(None)
            }
            else { Err((key, val)) }
        }
    }
}

/// Returns either the `(key,val)` pair associated with `key`, or `None` if it does not exist.
///
/// The `remove` method operates over a slice of `Option<(K,V)>`, given a target key and a hash
/// function `function`. It also takes an initial `location` from which to start its search, for
/// some reason.
#[inline]
pub fn remove<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &mut [Option<(K,V)>], key: &K, function: &F, location: usize) -> Option<(K,V)> {

    if true {
        panic!("rhhmap::remove is buggy; see code for details");
    }

    let target = function(&key);

    let mut success = false;
    let mut position = location;
    while position < slice.len() {
        if let Some(ref mut kv) = slice[position].as_mut() {
            let found = function(&kv.0);
            if found == target && key == &kv.0 {
                success = true;
                break;
            }
            if found > target { break; }
        }
        else { break; }

        position += 1;
    }

    if success {
        let result = slice[position].take();

        // now propagate the None forward as long as records are past their preferred location
        // BUG : `function` isn't the desired position, it is the hash. it should be shifted or something.
        while position + 1 < slice.len()
           && slice[position + 1].is_some()
           && function(&slice[position + 1].as_ref().unwrap().0) <= position {
            slice.swap(position, position + 1);
            position += 1;
        }

        result
    }
    else {
        None
    }
}
<commit_msg>moved SearchStatus around<commit_after>//! A Robin Hood hash map.

#[derive(Eq, PartialEq)]
enum SearchStatus {
    Searching,
    Found(usize),
    Missing(usize),
}

/// A Robin Hood hash map.
///
/// The hash map is parameterized by a supplied function acting as the hash function,
/// allowing the user more flexibility in defining the behavior of the map.
pub struct RHHMap<K: Eq, V, F: Fn(&K)->usize> {
    bucket: F,
    buffer: Vec<Option<(K,V)>>,
    shift:  usize,
    slop:   usize,
}

impl<K: Eq, V, F: Fn(&K)->usize> RHHMap<K, V, F> {
    /// Allocates a new Robin Hood hash map.
    pub fn new(function: F) -> RHHMap<K, V, F> {
        let slop = 4;
        let mut buffer = Vec::with_capacity(2 + slop);
        for _ in 0..buffer.capacity() { buffer.push(None); }
        RHHMap {
            bucket: function,
            buffer: buffer,
            shift: 63,
            slop: slop,
        }
    }
    /// The capacity of the underlying storage.
    pub fn capacity(&self) -> usize { self.buffer.capacity() }

    /// Get a reference to the value associated with a queried key.
    #[inline]
    pub fn get_ref<'a>(&'a self, query: &K) -> Option<&'a V> {
        let shift = self.shift;
        let bucket = &self.bucket;
        get_ref(&self.buffer[..], query, &|x| bucket(x), bucket(query) >> shift)
    }
    /// Get a mutable reference to the value associated with a queried key.
    #[inline]
    pub fn get_mut<'a>(&'a mut self, query: &K) -> Option<&'a mut V> {
        let shift = self.shift;
        let bucket = &self.bucket;
        get_mut(&mut self.buffer[..], query, &|x| bucket(x), bucket(query) >> shift)
    }
    /// Insert a key, value pair, recovering any pre-existing key, value pair.
    #[inline]
    pub fn insert(&mut self, key: K, val: V) -> Option<(K,V)> {

        let result = {
            let shift = self.shift;
            let bucket = &self.bucket;
            let location = bucket(&key) >> shift;
            insert(&mut self.buffer[..], key, val, &|x| bucket(x), location)
        };
        match result {
            Ok(result)      => result,
            Err((key, val)) => {

                let old_length = self.buffer.len() - self.slop;
                let mut new_buffer = Vec::with_capacity(2 * old_length + self.slop);
                for _ in 0..new_buffer.capacity() {
                    new_buffer.push(None);
                }
                let old_buffer = ::std::mem::replace(&mut self.buffer, new_buffer);
                self.shift -= 1;

                // let mut count = 0;
                let mut cursor = 0;
                for oldkeyval in old_buffer.into_iter() {
                    if let Some((oldkey, oldval)) = oldkeyval {
                        // count += 1;
                        let target = (self.bucket)(&oldkey) >> self.shift;
                        cursor = ::std::cmp::max(cursor, target);
                        self.buffer[cursor as usize] = Some((oldkey, oldval));
                        cursor += 1;
                    }
                }

                // println!("doubled length to {}; occupancy: {}", self.capacity() - self.slop, (count as f64) / ((self.capacity() - self.slop) as f64));

                self.insert(key, val)
            }
        }
    }
    /// Removes a key, recovering the key, value pair.
    pub fn remove(&mut self, key: &K) -> Option<(K,V)> {
        let shift = self.shift;
        let bucket = &self.bucket;
        remove(&mut self.buffer[..], key, &|x| bucket(x), bucket(key) >> shift)
    }
}

/// Returns either a shared reference to the value associated with `query`, or `None` if no value exists.
#[inline]
pub fn get_ref<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a [Option<(K,V)>], query: &K, function: &F, location: usize) -> Option<&'a V> {

    let target = function(query);
    let mut iterator = slice[location..].iter().map(|x| x.as_ref());
    while let Some(Some(&(ref key, ref val))) = iterator.next() {
        let found = function(key);
        if found >= target && key == query { return Some(val); }
        if found > target { return None; }
    }

    return None;
}

/// Returns either an exclusive reference to the value associated with `query`, or `None` if no value exists.
#[inline]
pub fn get_mut<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a mut [Option<(K,V)>], query: &K, function: &F, location: usize) -> Option<&'a mut V> {

    let target = function(query);
    let mut iterator = slice[location..].iter_mut().map(|x| x.as_mut());
    while let Some(Some(&mut (ref key, ref mut val))) = iterator.next() {
        let found = function(key);
        if found >= target && key == query { return Some(val); }
        if found > target { return None; }
    }

    return None;
}

/// Returns either the previous `(key,val)` pair for `key`, or `Err((key,val))` if `slice` does not have enough room to insert it.
#[inline]
pub fn insert<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &'a mut [Option<(K,V)>], key: K, val: V, function: &F, location: usize) -> Result<Option<(K,V)>,(K,V)> {

    let slop = 32;

    if location >= slice.len() {
        panic!("location {} > slice.len: {}", location, slice.len());
    }

    let target = function(&key);

    let mut position = location;
    let mut status = SearchStatus::Searching;
    let limit = ::std::cmp::min(location + slop, slice.len());
    while status == SearchStatus::Searching && position < limit {
        if let Some(ref mut kv) = slice[position].as_mut() {
            let found = function(&kv.0);
            if found == target && key == kv.0 {
                status = SearchStatus::Found(position);
            }
            if found > target { status = SearchStatus::Missing(position); }
        }
        else { status = SearchStatus::Missing(position); }

        position += 1;
    }

    match status {
        // if we bailed out of the search
        SearchStatus::Searching => {
            // println!("bailed out looking for element");
            Err((key, val))
        }
        // we found the key at `index`
        SearchStatus::Found(index) => {
            Ok(::std::mem::replace(&mut slice[index], Some((key, val))))
        }
        // we should have found the key by `index`, and should put the data there now
        SearchStatus::Missing(mut index) => {
            let begin = index;
            while index < slice.len() && slice[index].is_some() {
                index += 1;
            }

            // if we need to slide folks too far, is a bug
            if index > begin + 4 * slop {
                // println!("very long run: {}", index - begin);
                return Err((key, val));
            }

            if index < slice.len() {

                for i in 0..(index - begin) {
                    slice.swap(index - i - 1, index - i);
                }

                assert!(slice[begin].is_none());
                slice[begin] = Some((key, val));
                Ok(None)
            }
            else { Err((key, val)) }
        }
    }
}

/// Returns either the `(key,val)` pair associated with `key`, or `None` if it does not exist.
///
/// The `remove` method operates over a slice of `Option<(K,V)>`, given a target key and a hash
/// function `function`. It also takes an initial `location` from which to start its search, for
/// some reason.
#[inline]
pub fn remove<'a, K: Eq, V, F: Fn(&K)->usize>(slice: &mut [Option<(K,V)>], key: &K, function: &F, location: usize) -> Option<(K,V)> {

    if true {
        panic!("rhhmap::remove is buggy; see code for details");
    }

    let target = function(&key);

    let mut success = false;
    let mut position = location;
    while position < slice.len() {
        if let Some(ref mut kv) = slice[position].as_mut() {
            let found = function(&kv.0);
            if found == target && key == &kv.0 {
                success = true;
                break;
            }
            if found > target { break; }
        }
        else { break; }

        position += 1;
    }

    if success {
        let result = slice[position].take();

        // now propagate the None forward as long as records are past their preferred location
        // BUG : `function` isn't the desired position, it is the hash. it should be shifted or something.
        while position + 1 < slice.len()
           && slice[position + 1].is_some()
           && function(&slice[position + 1].as_ref().unwrap().0) <= position {
            slice.swap(position, position + 1);
            position += 1;
        }

        result
    }
    else {
        None
    }
}
<|endoftext|>