<commit_before>use primal_bit::BitVec;
use std::cmp;

use crate::wheel;

pub mod primes;
mod presieve;

/// A heavily optimised prime sieve.
///
/// This is a streaming segmented sieve, meaning it sieves numbers in
/// intervals, extracting whatever it needs and discarding it. See
/// `Sieve` for a wrapper that caches the information to allow for
/// repeated queries, at the cost of *O(limit)* memory use.
///
/// This uses *O(sqrt(limit))* memory, and is designed to be as
/// cache-friendly as possible. `StreamingSieve` should be used for
/// one-off calls, or simple linear iteration.
///
/// The design is *heavily* inspired/adopted from Kim Walisch's
/// [primesieve](http://primesieve.org/), and has similar speed
/// (around 5-20% slower).
///
/// # Examples
///
/// ```rust
/// let count = primal::StreamingSieve::prime_pi(123456);
/// println!("𝜋(123456) = {}", count);
/// ```
#[derive(Debug)]
pub struct StreamingSieve {
    small: Option<crate::Sieve>,
    sieve: BitVec,
    primes: Vec<wheel::State<wheel::Wheel210>>,
    small_primes: Vec<wheel::State<wheel::Wheel30>>,
    large_primes: Vec<wheel::State<wheel::Wheel210>>,
    presieve: presieve::Presieve,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = 32 << 10;
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::BYTE_MODULO / wheel::BYTE_SIZE;

fn isqrt(x: usize) -> usize {
    (x as f64).sqrt() as usize
}

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub(crate) fn new(limit: usize) -> StreamingSieve {
        let low = 0;

        let elems = cmp::min(wheel::bits_for(limit), SEG_ELEMS);
        let presieve = presieve::Presieve::new(elems);
        let current = presieve.smallest_unincluded_prime();

        let small = if limit < current * current {
            None
        } else {
            Some(crate::Sieve::new(isqrt(limit) + 1))
        };

        StreamingSieve {
            small,
            sieve: BitVec::from_elem(elems, true),
            primes: vec![],
            small_primes: vec![],
            large_primes: vec![],
            presieve,

            low,
            current,
            limit
        }
    }
    fn split_index(&self, idx: usize) -> (usize, usize) {
        let len = SEG_ELEMS;
        (idx / len,idx % len)
    }
    fn index_for(&self, n: usize) -> (bool, usize, usize) {
        let (b, idx) = wheel::bit_index(n);
        let (base, tweak) = self.split_index(idx);
        (b, base, tweak)
    }

    /// Count the number of primes upto and including `n`, that is, 𝜋,
    /// the [prime counting
    /// function](https://en.wikipedia.org/wiki/Prime-counting_function).
    ///
    /// # Examples
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::prime_pi(10), 4);
    /// // the endpoint is included
    /// assert_eq!(primal::StreamingSieve::prime_pi(11), 5);
    ///
    /// assert_eq!(primal::StreamingSieve::prime_pi(100), 25);
    /// assert_eq!(primal::StreamingSieve::prime_pi(1000), 168);
    /// ```
    pub fn prime_pi(n: usize) -> usize {
        match n {
            0..=1 => 0,
            2 => 1,
            3..=4 => 2,
            5..=6 => 3,
            7..=10 => 4,
            _ => {
                let mut sieve = StreamingSieve::new(n);
                let (includes, base, tweak) = sieve.index_for(n);
                let mut count = match wheel::BYTE_MODULO {
                    30 => 3,
                    _ => unimplemented!()
                };

                for _ in 0..base {
                    let (_, bitv) = sieve.next().unwrap();
                    count += bitv.count_ones();
                }

                let (_, last) = sieve.next().unwrap();
                count += last.count_ones_before(tweak + includes as usize);
                count
            }
        }
    }

    /// Compute *pn*, the `n` prime number, 1-indexed
    /// (i.e. *p1* = 2, *p2* = 3).
    ///
    /// # Panics
    ///
    /// `n` must be larger than 0 and less than the total number of
    /// primes in this sieve (that is,
    /// `self.prime_pi(self.upper_bound())`).
    ///
    /// # Example
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::nth_prime(1_000), 7919);
    /// ```
    pub fn nth_prime(n: usize) -> usize {
        assert!(n > 0);
        match n {
            1 => 2,
            2 => 3,
            3 => 5,
            _ => {
                let mut bit_n = n - 3;
                let (_, hi) = primal_estimate::nth_prime(n as u64);
                let mut sieve = StreamingSieve::new(hi as usize);
                while let Some((low, bits)) = sieve.next() {
                    let count = bits.count_ones();
                    if count >= bit_n {
                        let bit_idx = bits.find_nth_bit(bit_n - 1).unwrap();
                        return low + wheel::from_bit_index(bit_idx)
                    }

                    bit_n -= count
                }
                unreachable!()
            }
        }
    }

    fn add_sieving_prime(&mut self, p: usize, low: usize) {
        if p <= CACHE / 2 {
            self.small_primes.push(wheel::State::new(wheel::Wheel30, p, low));
        } else {
            let elem = wheel::State::new(wheel::Wheel210, p, low);
            if p < CACHE * 5 / 2 {
                self.primes.push(elem)
            } else {
                self.large_primes.push(elem)
            }
        }
    }

    fn find_new_sieving_primes(&mut self, low: usize, high: usize) {
        if let Some(small) = self.small.take() {
            for p in small.primes_from(self.current) {
                if p * p > high {
                    self.current = p;
                    break
                }
                self.add_sieving_prime(p, low);
            }
            self.small = Some(small);
        }
    }

    fn small_primes_sieve<W: wheel::Wheel>(sieve: &mut BitVec,
                                           small_primes: &mut [wheel::State<W>]) {
        let bytes = sieve.as_bytes_mut();
        for wi in small_primes {
            wi.sieve_hardcoded(bytes);
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.primes.chunks_exact_mut(3);

        while let Some([wi1, wi2, wi3]) = chunks.next() {
            wi1.sieve_triple(wi2, wi3, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    fn large_primes_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.large_primes.chunks_exact_mut(2);

        while let Some([wi1, wi2]) = chunks.next() {
            wi1.sieve_pair(wi2, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub(crate) fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low = self.low.saturating_add(SEG_LEN);
        let high = cmp::min(low.saturating_add(SEG_LEN - 1), self.limit);

        self.find_new_sieving_primes(low, high);

        self.presieve.apply(&mut self.sieve, low);
        StreamingSieve::small_primes_sieve(&mut self.sieve, &mut self.small_primes);
        self.direct_sieve();
        self.large_primes_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, false);
            self.presieve.mark_small_primes(&mut self.sieve);
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use primal_slowsieve::Primes;
    use crate::wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::BYTE_MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_slowsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::BYTE_MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::BYTE_SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
    #[test]
    fn prime_pi() {
        let (limit, mult) = if cfg!(feature = "slow_tests") {
            (2_000_000, 19_998)
        } else {
            (200_000, 1_998)
        };
        let real = Primes::sieve(limit);

        for i in (0..20).chain((0..100).map(|n| n * mult + 1)) {
            let val = StreamingSieve::prime_pi(i);
            let true_ = real.primes().take_while(|p| *p <= i).count();
            assert!(val == true_, "failed for {}, true {}, computed {}",
                    i, true_, val)
        }
    }

    #[test]
    fn prime_pi_issue_48() {
        // Expected values independently confirmed on Wolfram|Alpha.
        let limits = [49, 121, 289, 961, 11_047, 32_611, 230_907, 455_166_135];
        let expected = [15, 30, 61, 162, 1_338, 3_501, 20_513, 24_112_077];
        for (&limit, &expected) in limits.iter().zip(&expected) {
            let val = StreamingSieve::prime_pi(limit);
            assert_eq!(val, expected, "failed for limit {}", limit);
        }
    }

    #[test]
    fn nth_prime() {
        let primes = Sieve::new(2_000_000);

        for (i, p) in primes.primes_from(0).enumerate() {
            let n = i + 1;
            if n < 2000 || n % 1000 == 0 {
                assert_eq!(StreamingSieve::nth_prime(n), p);
            }
        }
    }

    // These are designed to specifically test the medium sized and
    // large prime sieving.
    #[test]
    fn prime_pi_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (10_000_000_000, 455_052_511);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (4_294_000_000, 203_236_859);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (500_000_000, 26_355_867);

        assert_eq!(StreamingSieve::prime_pi(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }

    #[test]
    fn nth_prime_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (455_052_512, 10_000_000_019);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (203_236_860, 4_294_000_079);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (26_355_868, 500_000_003);

        assert_eq!(StreamingSieve::nth_prime(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }
}
<commit_msg>Fix prime_pi on multiples of SEG_LEN<commit_after>use primal_bit::BitVec;
use std::cmp;

use crate::wheel;

pub mod primes;
mod presieve;

/// A heavily optimised prime sieve.
///
/// This is a streaming segmented sieve, meaning it sieves numbers in
/// intervals, extracting whatever it needs and discarding it. See
/// `Sieve` for a wrapper that caches the information to allow for
/// repeated queries, at the cost of *O(limit)* memory use.
///
/// This uses *O(sqrt(limit))* memory, and is designed to be as
/// cache-friendly as possible. `StreamingSieve` should be used for
/// one-off calls, or simple linear iteration.
///
/// The design is *heavily* inspired/adopted from Kim Walisch's
/// [primesieve](http://primesieve.org/), and has similar speed
/// (around 5-20% slower).
///
/// # Examples
///
/// ```rust
/// let count = primal::StreamingSieve::prime_pi(123456);
/// println!("𝜋(123456) = {}", count);
/// ```
#[derive(Debug)]
pub struct StreamingSieve {
    small: Option<crate::Sieve>,
    sieve: BitVec,
    primes: Vec<wheel::State<wheel::Wheel210>>,
    small_primes: Vec<wheel::State<wheel::Wheel30>>,
    large_primes: Vec<wheel::State<wheel::Wheel210>>,
    presieve: presieve::Presieve,

    low: usize,
    current: usize,
    limit: usize,
}

const CACHE: usize = 32 << 10;
const SEG_ELEMS: usize = 8 * CACHE;
const SEG_LEN: usize = SEG_ELEMS * wheel::BYTE_MODULO / wheel::BYTE_SIZE;

fn isqrt(x: usize) -> usize {
    (x as f64).sqrt() as usize
}

impl StreamingSieve {
    /// Create a new instance of the streaming sieve that will
    /// correctly progressively filter primes up to `limit`.
    pub(crate) fn new(limit: usize) -> StreamingSieve {
        let low = 0;

        let elems = cmp::min(wheel::bits_for(limit), SEG_ELEMS);
        let presieve = presieve::Presieve::new(elems);
        let current = presieve.smallest_unincluded_prime();

        let small = if limit < current * current {
            None
        } else {
            Some(crate::Sieve::new(isqrt(limit) + 1))
        };

        StreamingSieve {
            small,
            sieve: BitVec::from_elem(elems, true),
            primes: vec![],
            small_primes: vec![],
            large_primes: vec![],
            presieve,

            low,
            current,
            limit
        }
    }
    fn split_index(&self, idx: usize) -> (usize, usize) {
        let len = SEG_ELEMS;
        (idx / len,idx % len)
    }
    fn index_for(&self, n: usize) -> (bool, usize, usize) {
        let (b, idx) = wheel::bit_index(n);
        let (base, tweak) = self.split_index(idx);
        (b, base, tweak)
    }

    /// Count the number of primes upto and including `n`, that is, 𝜋,
    /// the [prime counting
    /// function](https://en.wikipedia.org/wiki/Prime-counting_function).
    ///
    /// # Examples
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::prime_pi(10), 4);
    /// // the endpoint is included
    /// assert_eq!(primal::StreamingSieve::prime_pi(11), 5);
    ///
    /// assert_eq!(primal::StreamingSieve::prime_pi(100), 25);
    /// assert_eq!(primal::StreamingSieve::prime_pi(1000), 168);
    /// ```
    pub fn prime_pi(n: usize) -> usize {
        match n {
            0..=1 => 0,
            2 => 1,
            3..=4 => 2,
            5..=6 => 3,
            7..=10 => 4,
            _ => {
                let mut sieve = StreamingSieve::new(n);
                let (includes, base, tweak) = sieve.index_for(n);
                let mut count = match wheel::BYTE_MODULO {
                    30 => 3,
                    _ => unimplemented!()
                };

                for _ in 0..base {
                    let (_, bitv) = sieve.next().unwrap();
                    count += bitv.count_ones();
                }

                let tail = tweak + includes as usize;
                if tail != 0 {
                    let (_, last) = sieve.next().unwrap();
                    count += last.count_ones_before(tail);
                }
                count
            }
        }
    }

    /// Compute *pn*, the `n` prime number, 1-indexed
    /// (i.e. *p1* = 2, *p2* = 3).
    ///
    /// # Panics
    ///
    /// `n` must be larger than 0 and less than the total number of
    /// primes in this sieve (that is,
    /// `self.prime_pi(self.upper_bound())`).
    ///
    /// # Example
    ///
    /// ```rust
    /// assert_eq!(primal::StreamingSieve::nth_prime(1_000), 7919);
    /// ```
    pub fn nth_prime(n: usize) -> usize {
        assert!(n > 0);
        match n {
            1 => 2,
            2 => 3,
            3 => 5,
            _ => {
                let mut bit_n = n - 3;
                let (_, hi) = primal_estimate::nth_prime(n as u64);
                let mut sieve = StreamingSieve::new(hi as usize);
                while let Some((low, bits)) = sieve.next() {
                    let count = bits.count_ones();
                    if count >= bit_n {
                        let bit_idx = bits.find_nth_bit(bit_n - 1).unwrap();
                        return low + wheel::from_bit_index(bit_idx)
                    }

                    bit_n -= count
                }
                unreachable!()
            }
        }
    }

    fn add_sieving_prime(&mut self, p: usize, low: usize) {
        if p <= CACHE / 2 {
            self.small_primes.push(wheel::State::new(wheel::Wheel30, p, low));
        } else {
            let elem = wheel::State::new(wheel::Wheel210, p, low);
            if p < CACHE * 5 / 2 {
                self.primes.push(elem)
            } else {
                self.large_primes.push(elem)
            }
        }
    }

    fn find_new_sieving_primes(&mut self, low: usize, high: usize) {
        if let Some(small) = self.small.take() {
            for p in small.primes_from(self.current) {
                if p * p > high {
                    self.current = p;
                    break
                }
                self.add_sieving_prime(p, low);
            }
            self.small = Some(small);
        }
    }

    fn small_primes_sieve<W: wheel::Wheel>(sieve: &mut BitVec,
                                           small_primes: &mut [wheel::State<W>]) {
        let bytes = sieve.as_bytes_mut();
        for wi in small_primes {
            wi.sieve_hardcoded(bytes);
        }
    }

    fn direct_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.primes.chunks_exact_mut(3);

        while let Some([wi1, wi2, wi3]) = chunks.next() {
            wi1.sieve_triple(wi2, wi3, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    fn large_primes_sieve(&mut self) {
        let bytes = self.sieve.as_bytes_mut();
        let mut chunks = self.large_primes.chunks_exact_mut(2);

        while let Some([wi1, wi2]) = chunks.next() {
            wi1.sieve_pair(wi2, bytes);
        }

        for wi in chunks.into_remainder() {
            wi.sieve(bytes);
        }
    }

    /// Extract the next chunk of filtered primes, the return value is
    /// `Some((low, v))` or `None` if the sieve has reached the limit.
    ///
    /// The vector stores bits for each odd number starting at `low`.
    /// Bit `n` of `v` is set if and only if `low + 2 * n + 1` is
    /// prime.
    ///
    /// NB. the prime 2 is not included in any of these sieves and so
    /// needs special handling.
    pub(crate) fn next(&mut self) -> Option<(usize, &BitVec)> {
        if self.low >= self.limit {
            return None
        }

        let low = self.low;
        self.low = self.low.saturating_add(SEG_LEN);
        let high = cmp::min(low.saturating_add(SEG_LEN - 1), self.limit);

        self.find_new_sieving_primes(low, high);

        self.presieve.apply(&mut self.sieve, low);
        StreamingSieve::small_primes_sieve(&mut self.sieve, &mut self.small_primes);
        self.direct_sieve();
        self.large_primes_sieve();

        if low == 0 {
            // 1 is not prime.
            self.sieve.set(0, false);
            self.presieve.mark_small_primes(&mut self.sieve);
        }

        Some((low, &self.sieve))
    }
}

#[cfg(test)]
mod tests {
    use crate::Sieve;
    use primal_slowsieve::Primes;
    use crate::wheel;
    use super::StreamingSieve;
    fn gcd(x: usize, y: usize) -> usize {
        if y == 0 { x }
        else { gcd(y, x % y) }
    }
    fn coprime_to(x: usize) -> Vec<usize> {
        (1..x).filter(|&n| gcd(n, x) == 1).collect()
    }
    #[test]
    fn test() {
        let coprime = coprime_to(wheel::BYTE_MODULO);
        const LIMIT: usize = 2_000_000;
        let mut sieve = StreamingSieve::new(LIMIT);
        let primes = ::primal_slowsieve::Primes::sieve(LIMIT);

        let mut base = 0;
        let mut index = 0;

        while let Some((_low, next)) = sieve.next() {
            for val in next {
                let i = wheel::BYTE_MODULO * base + coprime[index];
                if i >= LIMIT { break }
                assert!(primes.is_prime(i) == val,
                        "failed for {} (is prime = {})", i, primes.is_prime(i));

                index += 1;
                if index == wheel::BYTE_SIZE {
                    index = 0;
                    base += 1
                }
            }
        }
    }
    #[test]
    fn prime_pi() {
        let (limit, mult) = if cfg!(feature = "slow_tests") {
            (2_000_000, 19_998)
        } else {
            (200_000, 1_998)
        };
        let real = Primes::sieve(limit);

        for i in (0..20).chain((0..100).map(|n| n * mult + 1)) {
            let val = StreamingSieve::prime_pi(i);
            let true_ = real.primes().take_while(|p| *p <= i).count();
            assert!(val == true_, "failed for {}, true {}, computed {}",
                    i, true_, val)
        }
    }

    #[test]
    fn prime_pi_issue_48() {
        // Expected values independently confirmed on Wolfram|Alpha.
        let limits = [49, 121, 289, 961, 11_047, 32_611, 230_907, 455_166_135];
        let expected = [15, 30, 61, 162, 1_338, 3_501, 20_513, 24_112_077];
        for (&limit, &expected) in limits.iter().zip(&expected) {
            let val = StreamingSieve::prime_pi(limit);
            assert_eq!(val, expected, "failed for limit {}", limit);
        }
    }

    #[test]
    fn prime_pi_seg_len() {
        const PRIME1: usize = 982_981;
        const PRIME2: usize = 983_063;
        const EXPECTED: usize = 77_279;
        assert!((PRIME1..PRIME2).contains(&super::SEG_LEN));
        assert_eq!(StreamingSieve::prime_pi(PRIME1 - 1), EXPECTED - 1);
        for limit in PRIME1..PRIME2 {
            assert_eq!(StreamingSieve::prime_pi(limit), EXPECTED);
        }
        assert_eq!(StreamingSieve::prime_pi(PRIME2), EXPECTED + 1);
    }

    #[test]
    fn nth_prime() {
        let primes = Sieve::new(2_000_000);

        for (i, p) in primes.primes_from(0).enumerate() {
            let n = i + 1;
            if n < 2000 || n % 1000 == 0 {
                assert_eq!(StreamingSieve::nth_prime(n), p);
            }
        }
    }

    // These are designed to specifically test the medium sized and
    // large prime sieving.
    #[test]
    fn prime_pi_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (10_000_000_000, 455_052_511);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (4_294_000_000, 203_236_859);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (500_000_000, 26_355_867);

        assert_eq!(StreamingSieve::prime_pi(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }

    #[test]
    fn nth_prime_huge() {
        #[cfg(all(feature = "slow_tests", target_pointer_width = "64"))]
        const LIMIT_RESULT: (usize, usize) = (455_052_512, 10_000_000_019);
        #[cfg(all(feature = "slow_tests", target_pointer_width = "32"))]
        const LIMIT_RESULT: (usize, usize) = (203_236_860, 4_294_000_079);
        #[cfg(not(feature = "slow_tests"))]
        const LIMIT_RESULT: (usize, usize) = (26_355_868, 500_000_003);

        assert_eq!(StreamingSieve::nth_prime(LIMIT_RESULT.0), LIMIT_RESULT.1);
    }
}
<|endoftext|>


<commit_before>use crate::context::Context;
use crate::instance::{
    FaultDetails, Instance, State, TerminationDetails, CURRENT_INSTANCE, HOST_CTX,
};
use crate::sysdeps::UContextPtr;
use crate::trapcode::{TrapCode, TrapCodeType};
use failure::Error;
use lazy_static::lazy_static;
use libc::{c_int, c_void, siginfo_t};
use nix::sys::signal::{
    pthread_sigmask, raise, sigaction, SaFlags, SigAction, SigHandler, SigSet, SigmaskHow, Signal,
};
use std::sync::Mutex;

lazy_static! {
    // TODO: work out an alternative to this that is signal-safe for `reraise_host_signal_in_handler`
    static ref LUCET_SIGNAL_STATE: Mutex<Option<SignalState>> = Mutex::new(None);
}

/// The value returned by
/// [`Instance.signal_handler`](struct.Instance.html#structfield.signal_handler) to determine the
/// outcome of a handled signal.
pub enum SignalBehavior {
    /// Use default behavior, which switches back to the host with `State::Fault` populated.
    Default,
    /// Override default behavior and cause the instance to continue.
    Continue,
    /// Override default behavior and cause the instance to terminate.
    Terminate,
}

pub type SignalHandler =
    dyn Fn(&Instance, &TrapCode, libc::c_int, *const siginfo_t, *const c_void) -> SignalBehavior;

pub fn signal_handler_none(
    _inst: &Instance,
    _trapcode: &TrapCode,
    _signum: libc::c_int,
    _siginfo_ptr: *const siginfo_t,
    _ucontext_ptr: *const c_void,
) -> SignalBehavior {
    SignalBehavior::Default
}

impl Instance {
    pub(crate) fn with_signals_on<F, R>(&mut self, f: F) -> Result<R, Error>
    where
        F: FnOnce(&mut Instance) -> Result<R, Error>,
    {
        // setup signal stack for this thread
        let guest_sigstack = SigStack::new(
            self.alloc.slot().sigstack,
            SigStackFlags::empty(),
            libc::SIGSTKSZ,
        );
        let saved_sigstack = unsafe { sigaltstack(&guest_sigstack).expect("sigaltstack succeeds") };

        let mut ostate = LUCET_SIGNAL_STATE.lock().unwrap();
        if let Some(ref mut state) = *ostate {
            state.counter += 1;
        } else {
            unsafe {
                setup_guest_signal_state(&mut ostate);
            }
        }
        drop(ostate);

        // run the body
        let res = f(self);

        let mut ostate = LUCET_SIGNAL_STATE.lock().unwrap();
        let counter_zero = if let Some(ref mut state) = *ostate {
            state.counter -= 1;
            if state.counter == 0 {
                unsafe {
                    // restore the host signal stack
                    sigaltstack(&saved_sigstack).expect("sigaltstack succeeds");
                    restore_host_signal_state(state);
                }
                true
            } else {
                false
            }
        } else {
            panic!("signal handlers weren't installed at instance exit");
        };
        if counter_zero {
            *ostate = None;
        }

        res
    }
}

/// Signal handler installed during instance execution.
///
/// This function is only designed to handle signals that are the direct result of execution of a
/// hardware instruction from the faulting WASM thread. It thus safely assumes the signal is
/// directed specifically at this thread (i.e. not a different thread or the process as a whole).
extern "C" fn handle_signal(signum: c_int, siginfo_ptr: *mut siginfo_t, ucontext_ptr: *mut c_void) {
    let signal = Signal::from_c_int(signum).expect("signum is a valid signal");
    if !(signal == Signal::SIGBUS
        || signal == Signal::SIGSEGV
        || signal == Signal::SIGILL
        || signal == Signal::SIGFPE)
    {
        panic!("unexpected signal in guest signal handler: {:?}", signal);
    }
    assert!(!siginfo_ptr.is_null(), "siginfo must not be null");

    // Safety: when using a SA_SIGINFO sigaction, the third argument can be cast to a `ucontext_t`
    // pointer per the manpage
    assert!(!ucontext_ptr.is_null());
    let ctx = UContextPtr::new(ucontext_ptr);
    let rip = ctx.get_ip();

    let switch_to_host = CURRENT_INSTANCE.with(|current_instance| {
        let mut current_instance = current_instance.borrow_mut();

        if current_instance.is_none() {
            // If there is no current instance, we've caught a signal raised by a thread that's not
            // running a lucet instance. Restore the host signal handler and reraise the signal,
            // then return if the host handler returns
            unsafe {
                reraise_host_signal_in_handler(signal, signum, siginfo_ptr, ucontext_ptr);
            }
            // don't try context-switching
            return false;
        }

        // Safety: the memory pointed to by CURRENT_INSTANCE should be a valid instance. This is not
        // a trivial property, but relies on the compiler not emitting guest programs that can
        // overwrite the instance.
        let inst = unsafe {
            current_instance
                .as_mut()
                .expect("current instance exists")
                .as_mut()
        };

        let trapcode = inst
            .module
            .lookup_trapcode(rip)
            // if we couldn't find a code in the manifest, return an unknown trapcode that will be
            // converted to a fatal trap when we switch back to the host
            .unwrap_or(TrapCode {
                ty: TrapCodeType::Unknown,
                tag: 0,
            });

        let behavior = (inst.signal_handler)(inst, &trapcode, signum, siginfo_ptr, ucontext_ptr);
        match behavior {
            SignalBehavior::Continue => {
                // return to the guest context without making any modifications to the instance
                false
            }
            SignalBehavior::Terminate => {
                // set the state before jumping back to the host context
                inst.state = State::Terminated {
                    details: TerminationDetails::Signal,
                };
                true
            }
            SignalBehavior::Default => {
                // record the fault and jump back to the host context
                inst.state = State::Fault {
                    details: FaultDetails {
                        // fatal field is set false here by default - have to wait until
                        // `verify_trap_safety` to have enough information to determine whether trap was
                        // fatal. It is not signal safe to access some of the required information.
                        fatal: false,
                        trapcode: trapcode,
                        rip_addr: rip as usize,
                        rip_addr_details: None,
                    },
                    // safety: pointer is checked for null at the top of the function, and the
                    // manpage guarantees that a siginfo_t will be passed as the second argument
                    siginfo: unsafe { *siginfo_ptr },
                    context: ctx.into(),
                };
                true
            }
        }
    });

    if switch_to_host {
        HOST_CTX.with(|host_ctx| unsafe {
            Context::set_from_signal(&*host_ctx.get())
                .expect("can successfully switch back to the host context");
        });
        unreachable!()
    }
}

struct SignalState {
    counter: usize,
    saved_sigbus: SigAction,
    saved_sigfpe: SigAction,
    saved_sigill: SigAction,
    saved_sigsegv: SigAction,
}

// raw pointers in the saved types
unsafe impl Send for SignalState {}

unsafe fn setup_guest_signal_state(ostate: &mut Option<SignalState>) {
    let mut masked_signals = SigSet::empty();
    masked_signals.add(Signal::SIGBUS);
    masked_signals.add(Signal::SIGFPE);
    masked_signals.add(Signal::SIGILL);
    masked_signals.add(Signal::SIGSEGV);

    // setup signal handlers
    let sa = SigAction::new(
        SigHandler::SigAction(handle_signal),
        SaFlags::SA_RESTART | SaFlags::SA_SIGINFO | SaFlags::SA_ONSTACK,
        masked_signals,
    );
    let saved_sigbus = sigaction(Signal::SIGBUS, &sa).expect("sigaction succeeds");
    let saved_sigfpe = sigaction(Signal::SIGFPE, &sa).expect("sigaction succeeds");
    let saved_sigill = sigaction(Signal::SIGILL, &sa).expect("sigaction succeeds");
    let saved_sigsegv = sigaction(Signal::SIGSEGV, &sa).expect("sigaction succeeds");

    *ostate = Some(SignalState {
        counter: 1,
        saved_sigbus,
        saved_sigfpe,
        saved_sigill,
        saved_sigsegv,
    });
}

unsafe fn restore_host_signal_state(state: &mut SignalState) {
    // restore signal handlers
    sigaction(Signal::SIGBUS, &state.saved_sigbus).expect("sigaction succeeds");
    sigaction(Signal::SIGFPE, &state.saved_sigfpe).expect("sigaction succeeds");
    sigaction(Signal::SIGILL, &state.saved_sigill).expect("sigaction succeeds");
    sigaction(Signal::SIGSEGV, &state.saved_sigsegv).expect("sigaction succeeds");
}

unsafe fn reraise_host_signal_in_handler(
    sig: Signal,
    signum: libc::c_int,
    siginfo_ptr: *mut libc::siginfo_t,
    ucontext_ptr: *mut c_void,
) {
    let saved_handler = {
        // TODO: avoid taking a mutex here, probably by having some static muts just for this
        // function
        if let Some(ref state) = *LUCET_SIGNAL_STATE.lock().unwrap() {
            match sig {
                Signal::SIGBUS => state.saved_sigbus.clone(),
                Signal::SIGFPE => state.saved_sigfpe.clone(),
                Signal::SIGILL => state.saved_sigill.clone(),
                Signal::SIGSEGV => state.saved_sigsegv.clone(),
                sig => panic!(
                    "unexpected signal in reraise_host_signal_in_handler: {:?}",
                    sig
                ),
            }
        } else {
            // this case is very fishy; it can arise when the last lucet instance spins down and
            // uninstalls the lucet handlers while a signal handler is running on this thread, but
            // before taking the mutex above. The theory is that if this has happened, the host
            // handler has been reinstalled, so we shouldn't end up back here if we reraise

            // unmask the signal to reraise; we don't have to restore it because the handler will return
            // after this. If it signals again between here and now, that's a double fault and the
            // process is going to die anyway
            let mut unmask = SigSet::empty();
            unmask.add(sig);
            pthread_sigmask(SigmaskHow::SIG_UNBLOCK, Some(&unmask), None)
                .expect("pthread_sigmask succeeds");
            // if there's no current signal state, just re-raise and hope for the best
            raise(sig).expect("raise succeeds");
            return;
        }
    };

    match saved_handler.handler() {
        SigHandler::SigDfl => {
            // reinstall default signal handler and reraise the signal; this should terminate the
            // program
            sigaction(sig, &saved_handler).expect("sigaction succeeds");
            let mut unmask = SigSet::empty();
            unmask.add(sig);
            pthread_sigmask(SigmaskHow::SIG_UNBLOCK, Some(&unmask), None)
                .expect("pthread_sigmask succeeds");
            raise(sig).expect("raise succeeds");
        }
        SigHandler::SigIgn => {
            // don't do anything; if we hit this case, whatever program is hosting us is almost
            // certainly doing something wrong, because our set of signals requires intervention to
            // proceed
            return;
        }
        SigHandler::Handler(f) => {
            // call the saved handler directly so there is no altstack confusion
            f(signum)
        }
        SigHandler::SigAction(f) => {
            // call the saved handler directly so there is no altstack confusion
            f(signum, siginfo_ptr, ucontext_ptr)
        }
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// A collection of wrappers that will be upstreamed to the `nix` crate eventually.
////////////////////////////////////////////////////////////////////////////////////////////////////

use bitflags::bitflags;

#[derive(Copy, Clone)]
pub struct SigStack {
    stack: libc::stack_t,
}

impl SigStack {
    pub fn new(sp: *mut libc::c_void, flags: SigStackFlags, size: libc::size_t) -> SigStack {
        let mut stack = unsafe { std::mem::uninitialized::<libc::stack_t>() };
        stack.ss_sp = sp;
        stack.ss_flags = flags.bits();
        stack.ss_size = size;
        SigStack { stack }
    }

    pub fn flags(&self) -> SigStackFlags {
        SigStackFlags::from_bits_truncate(self.stack.ss_flags)
    }
}

impl AsRef<libc::stack_t> for SigStack {
    fn as_ref(&self) -> &libc::stack_t {
        &self.stack
    }
}

impl AsMut<libc::stack_t> for SigStack {
    fn as_mut(&mut self) -> &mut libc::stack_t {
        &mut self.stack
    }
}

bitflags! {
    pub struct SigStackFlags: libc::c_int {
        const SS_ONSTACK = libc::SS_ONSTACK;
        const SS_DISABLE = libc::SS_DISABLE;
    }
}

pub unsafe fn sigaltstack(ss: &SigStack) -> nix::Result<SigStack> {
    let mut oldstack = std::mem::uninitialized::<libc::stack_t>();

    let res = libc::sigaltstack(
        &ss.stack as *const libc::stack_t,
        &mut oldstack as *mut libc::stack_t,
    );

    nix::errno::Errno::result(res).map(|_| SigStack { stack: oldstack })
}
<commit_msg>Make expect() messages more explicit<commit_after>use crate::context::Context;
use crate::instance::{
    FaultDetails, Instance, State, TerminationDetails, CURRENT_INSTANCE, HOST_CTX,
};
use crate::sysdeps::UContextPtr;
use crate::trapcode::{TrapCode, TrapCodeType};
use failure::Error;
use lazy_static::lazy_static;
use libc::{c_int, c_void, siginfo_t};
use nix::sys::signal::{
    pthread_sigmask, raise, sigaction, SaFlags, SigAction, SigHandler, SigSet, SigmaskHow, Signal,
};
use std::sync::Mutex;

lazy_static! {
    // TODO: work out an alternative to this that is signal-safe for `reraise_host_signal_in_handler`
    static ref LUCET_SIGNAL_STATE: Mutex<Option<SignalState>> = Mutex::new(None);
}

/// The value returned by
/// [`Instance.signal_handler`](struct.Instance.html#structfield.signal_handler) to determine the
/// outcome of a handled signal.
pub enum SignalBehavior {
    /// Use default behavior, which switches back to the host with `State::Fault` populated.
    Default,
    /// Override default behavior and cause the instance to continue.
    Continue,
    /// Override default behavior and cause the instance to terminate.
    Terminate,
}

pub type SignalHandler =
    dyn Fn(&Instance, &TrapCode, libc::c_int, *const siginfo_t, *const c_void) -> SignalBehavior;

pub fn signal_handler_none(
    _inst: &Instance,
    _trapcode: &TrapCode,
    _signum: libc::c_int,
    _siginfo_ptr: *const siginfo_t,
    _ucontext_ptr: *const c_void,
) -> SignalBehavior {
    SignalBehavior::Default
}

impl Instance {
    pub(crate) fn with_signals_on<F, R>(&mut self, f: F) -> Result<R, Error>
    where
        F: FnOnce(&mut Instance) -> Result<R, Error>,
    {
        // setup signal stack for this thread
        let guest_sigstack = SigStack::new(
            self.alloc.slot().sigstack,
            SigStackFlags::empty(),
            libc::SIGSTKSZ,
        );
        let saved_sigstack =
            unsafe { sigaltstack(&guest_sigstack).expect("saving sigaltstack succeeds") };

        let mut ostate = LUCET_SIGNAL_STATE.lock().unwrap();
        if let Some(ref mut state) = *ostate {
            state.counter += 1;
        } else {
            unsafe {
                setup_guest_signal_state(&mut ostate);
            }
        }
        drop(ostate);

        // run the body
        let res = f(self);

        let mut ostate = LUCET_SIGNAL_STATE.lock().unwrap();
        let counter_zero = if let Some(ref mut state) = *ostate {
            state.counter -= 1;
            if state.counter == 0 {
                unsafe {
                    // restore the host signal stack
                    sigaltstack(&saved_sigstack).expect("sigaltstack restoration succeeds");
                    restore_host_signal_state(state);
                }
                true
            } else {
                false
            }
        } else {
            panic!("signal handlers weren't installed at instance exit");
        };
        if counter_zero {
            *ostate = None;
        }

        res
    }
}

/// Signal handler installed during instance execution.
///
/// This function is only designed to handle signals that are the direct result of execution of a
/// hardware instruction from the faulting WASM thread. It thus safely assumes the signal is
/// directed specifically at this thread (i.e. not a different thread or the process as a whole).
extern "C" fn handle_signal(signum: c_int, siginfo_ptr: *mut siginfo_t, ucontext_ptr: *mut c_void) {
    let signal = Signal::from_c_int(signum).expect("signum is a valid signal");
    if !(signal == Signal::SIGBUS
        || signal == Signal::SIGSEGV
        || signal == Signal::SIGILL
        || signal == Signal::SIGFPE)
    {
        panic!("unexpected signal in guest signal handler: {:?}", signal);
    }
    assert!(!siginfo_ptr.is_null(), "siginfo must not be null");

    // Safety: when using a SA_SIGINFO sigaction, the third argument can be cast to a `ucontext_t`
    // pointer per the manpage
    assert!(!ucontext_ptr.is_null());
    let ctx = UContextPtr::new(ucontext_ptr);
    let rip = ctx.get_ip();

    let switch_to_host = CURRENT_INSTANCE.with(|current_instance| {
        let mut current_instance = current_instance.borrow_mut();

        if current_instance.is_none() {
            // If there is no current instance, we've caught a signal raised by a thread that's not
            // running a lucet instance. Restore the host signal handler and reraise the signal,
            // then return if the host handler returns
            unsafe {
                reraise_host_signal_in_handler(signal, signum, siginfo_ptr, ucontext_ptr);
            }
            // don't try context-switching
            return false;
        }

        // Safety: the memory pointed to by CURRENT_INSTANCE should be a valid instance. This is not
        // a trivial property, but relies on the compiler not emitting guest programs that can
        // overwrite the instance.
        let inst = unsafe {
            current_instance
                .as_mut()
                .expect("current instance exists")
                .as_mut()
        };

        let trapcode = inst
            .module
            .lookup_trapcode(rip)
            // if we couldn't find a code in the manifest, return an unknown trapcode that will be
            // converted to a fatal trap when we switch back to the host
            .unwrap_or(TrapCode {
                ty: TrapCodeType::Unknown,
                tag: 0,
            });

        let behavior = (inst.signal_handler)(inst, &trapcode, signum, siginfo_ptr, ucontext_ptr);
        match behavior {
            SignalBehavior::Continue => {
                // return to the guest context without making any modifications to the instance
                false
            }
            SignalBehavior::Terminate => {
                // set the state before jumping back to the host context
                inst.state = State::Terminated {
                    details: TerminationDetails::Signal,
                };
                true
            }
            SignalBehavior::Default => {
                // record the fault and jump back to the host context
                inst.state = State::Fault {
                    details: FaultDetails {
                        // fatal field is set false here by default - have to wait until
                        // `verify_trap_safety` to have enough information to determine whether trap was
                        // fatal. It is not signal safe to access some of the required information.
                        fatal: false,
                        trapcode: trapcode,
                        rip_addr: rip as usize,
                        rip_addr_details: None,
                    },
                    // safety: pointer is checked for null at the top of the function, and the
                    // manpage guarantees that a siginfo_t will be passed as the second argument
                    siginfo: unsafe { *siginfo_ptr },
                    context: ctx.into(),
                };
                true
            }
        }
    });

    if switch_to_host {
        HOST_CTX.with(|host_ctx| unsafe {
            Context::set_from_signal(&*host_ctx.get())
                .expect("can successfully switch back to the host context");
        });
        unreachable!()
    }
}

struct SignalState {
    counter: usize,
    saved_sigbus: SigAction,
    saved_sigfpe: SigAction,
    saved_sigill: SigAction,
    saved_sigsegv: SigAction,
}

// raw pointers in the saved types
unsafe impl Send for SignalState {}

unsafe fn setup_guest_signal_state(ostate: &mut Option<SignalState>) {
    let mut masked_signals = SigSet::empty();
    masked_signals.add(Signal::SIGBUS);
    masked_signals.add(Signal::SIGFPE);
    masked_signals.add(Signal::SIGILL);
    masked_signals.add(Signal::SIGSEGV);

    // setup signal handlers
    let sa = SigAction::new(
        SigHandler::SigAction(handle_signal),
        SaFlags::SA_RESTART | SaFlags::SA_SIGINFO | SaFlags::SA_ONSTACK,
        masked_signals,
    );
    let saved_sigbus = sigaction(Signal::SIGBUS, &sa).expect("sigaction succeeds");
    let saved_sigfpe = sigaction(Signal::SIGFPE, &sa).expect("sigaction succeeds");
    let saved_sigill = sigaction(Signal::SIGILL, &sa).expect("sigaction succeeds");
    let saved_sigsegv = sigaction(Signal::SIGSEGV, &sa).expect("sigaction succeeds");

    *ostate = Some(SignalState {
        counter: 1,
        saved_sigbus,
        saved_sigfpe,
        saved_sigill,
        saved_sigsegv,
    });
}

unsafe fn restore_host_signal_state(state: &mut SignalState) {
    // restore signal handlers
    sigaction(Signal::SIGBUS, &state.saved_sigbus).expect("sigaction succeeds");
    sigaction(Signal::SIGFPE, &state.saved_sigfpe).expect("sigaction succeeds");
    sigaction(Signal::SIGILL, &state.saved_sigill).expect("sigaction succeeds");
    sigaction(Signal::SIGSEGV, &state.saved_sigsegv).expect("sigaction succeeds");
}

unsafe fn reraise_host_signal_in_handler(
    sig: Signal,
    signum: libc::c_int,
    siginfo_ptr: *mut libc::siginfo_t,
    ucontext_ptr: *mut c_void,
) {
    let saved_handler = {
        // TODO: avoid taking a mutex here, probably by having some static muts just for this
        // function
        if let Some(ref state) = *LUCET_SIGNAL_STATE.lock().unwrap() {
            match sig {
                Signal::SIGBUS => state.saved_sigbus.clone(),
                Signal::SIGFPE => state.saved_sigfpe.clone(),
                Signal::SIGILL => state.saved_sigill.clone(),
                Signal::SIGSEGV => state.saved_sigsegv.clone(),
                sig => panic!(
                    "unexpected signal in reraise_host_signal_in_handler: {:?}",
                    sig
                ),
            }
        } else {
            // this case is very fishy; it can arise when the last lucet instance spins down and
            // uninstalls the lucet handlers while a signal handler is running on this thread, but
            // before taking the mutex above. The theory is that if this has happened, the host
            // handler has been reinstalled, so we shouldn't end up back here if we reraise

            // unmask the signal to reraise; we don't have to restore it because the handler will return
            // after this. If it signals again between here and now, that's a double fault and the
            // process is going to die anyway
            let mut unmask = SigSet::empty();
            unmask.add(sig);
            pthread_sigmask(SigmaskHow::SIG_UNBLOCK, Some(&unmask), None)
                .expect("pthread_sigmask succeeds");
            // if there's no current signal state, just re-raise and hope for the best
            raise(sig).expect("raise succeeds");
            return;
        }
    };

    match saved_handler.handler() {
        SigHandler::SigDfl => {
            // reinstall default signal handler and reraise the signal; this should terminate the
            // program
            sigaction(sig, &saved_handler).expect("sigaction succeeds");
            let mut unmask = SigSet::empty();
            unmask.add(sig);
            pthread_sigmask(SigmaskHow::SIG_UNBLOCK, Some(&unmask), None)
                .expect("pthread_sigmask succeeds");
            raise(sig).expect("raise succeeds");
        }
        SigHandler::SigIgn => {
            // don't do anything; if we hit this case, whatever program is hosting us is almost
            // certainly doing something wrong, because our set of signals requires intervention to
            // proceed
            return;
        }
        SigHandler::Handler(f) => {
            // call the saved handler directly so there is no altstack confusion
            f(signum)
        }
        SigHandler::SigAction(f) => {
            // call the saved handler directly so there is no altstack confusion
            f(signum, siginfo_ptr, ucontext_ptr)
        }
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// A collection of wrappers that will be upstreamed to the `nix` crate eventually.
////////////////////////////////////////////////////////////////////////////////////////////////////

use bitflags::bitflags;

#[derive(Copy, Clone)]
pub struct SigStack {
    stack: libc::stack_t,
}

impl SigStack {
    pub fn new(sp: *mut libc::c_void, flags: SigStackFlags, size: libc::size_t) -> SigStack {
        let mut stack = unsafe { std::mem::uninitialized::<libc::stack_t>() };
        stack.ss_sp = sp;
        stack.ss_flags = flags.bits();
        stack.ss_size = size;
        SigStack { stack }
    }

    pub fn flags(&self) -> SigStackFlags {
        SigStackFlags::from_bits_truncate(self.stack.ss_flags)
    }
}

impl AsRef<libc::stack_t> for SigStack {
    fn as_ref(&self) -> &libc::stack_t {
        &self.stack
    }
}

impl AsMut<libc::stack_t> for SigStack {
    fn as_mut(&mut self) -> &mut libc::stack_t {
        &mut self.stack
    }
}

bitflags! {
    pub struct SigStackFlags: libc::c_int {
        const SS_ONSTACK = libc::SS_ONSTACK;
        const SS_DISABLE = libc::SS_DISABLE;
    }
}

pub unsafe fn sigaltstack(ss: &SigStack) -> nix::Result<SigStack> {
    let mut oldstack = std::mem::uninitialized::<libc::stack_t>();

    let res = libc::sigaltstack(
        &ss.stack as *const libc::stack_t,
        &mut oldstack as *mut libc::stack_t,
    );

    nix::errno::Errno::result(res).map(|_| SigStack { stack: oldstack })
}
<|endoftext|>


<commit_before>use crate::proto::grpc_status::GrpcStatus;
use crate::Metadata;
use bytes::Bytes;
use httpbis::Header;
use httpbis::Headers;

pub(crate) static HEADER_GRPC_STATUS: &'static str = "grpc-status";
pub(crate) static HEADER_GRPC_MESSAGE: &'static str = "grpc-message";

pub(crate) fn headers_500(grpc_status: GrpcStatus, message: String) -> Headers {
    Headers::from_vec(vec![
        Header::new(":status", "500"),
        Header::new(HEADER_GRPC_STATUS, format!("{}", grpc_status as i32)),
        Header::new(HEADER_GRPC_MESSAGE, message),
    ])
}

pub(crate) fn headers_200(metadata: Metadata) -> Headers {
    let mut headers = Headers::from_vec(vec![
        // TODO: do not allocate
        Header::new(":status", "200"),
        Header::new("content-type", "application/grpc"),
        Header::new(HEADER_GRPC_STATUS, "0"),
    ]);
    headers.extend(metadata.into_headers());
    headers
}

/// Create HTTP response for gRPC error
pub(crate) fn grpc_error_message(message: &str) -> httpbis::SimpleHttpMessage {
    let headers = Headers::from_vec(vec![
        Header::new(":status", "200"),
        // TODO: alloc
        Header::new(
            HEADER_GRPC_STATUS,
            format!("{}", GrpcStatus::Internal.code()),
        ),
        Header::new(HEADER_GRPC_MESSAGE, message.to_owned()),
    ]);
    httpbis::SimpleHttpMessage {
        headers,
        body: Bytes::new(),
    }
}

// Trailers -> Status [Status-Message] *Custom-Metadata
pub(crate) fn trailers(
    grpc_status: GrpcStatus,
    message: Option<String>,
    metadata: Metadata,
) -> Headers {
    let mut headers = Headers::from_vec(vec![Header::new(
        HEADER_GRPC_STATUS,
        format!("{}", grpc_status as i32),
    )]);
    if let Some(message) = message {
        headers.add_header(Header::new(HEADER_GRPC_MESSAGE, message));
    }
    headers.extend(metadata.into_headers());
    headers
}
<commit_msg>Update to latest rust-http2<commit_after>use crate::proto::grpc_status::GrpcStatus;
use crate::Metadata;
use bytes::Bytes;
use httpbis::BytesDeque;
use httpbis::Header;
use httpbis::Headers;

pub(crate) static HEADER_GRPC_STATUS: &'static str = "grpc-status";
pub(crate) static HEADER_GRPC_MESSAGE: &'static str = "grpc-message";

pub(crate) fn headers_500(grpc_status: GrpcStatus, message: String) -> Headers {
    Headers::from_vec(vec![
        Header::new(":status", "500"),
        Header::new(HEADER_GRPC_STATUS, format!("{}", grpc_status as i32)),
        Header::new(HEADER_GRPC_MESSAGE, message),
    ])
}

pub(crate) fn headers_200(metadata: Metadata) -> Headers {
    let mut headers = Headers::from_vec(vec![
        // TODO: do not allocate
        Header::new(":status", "200"),
        Header::new("content-type", "application/grpc"),
        Header::new(HEADER_GRPC_STATUS, "0"),
    ]);
    headers.extend(metadata.into_headers());
    headers
}

/// Create HTTP response for gRPC error
pub(crate) fn grpc_error_message(message: &str) -> httpbis::SimpleHttpMessage {
    let headers = Headers::from_vec(vec![
        Header::new(":status", "200"),
        // TODO: alloc
        Header::new(
            HEADER_GRPC_STATUS,
            format!("{}", GrpcStatus::Internal.code()),
        ),
        Header::new(HEADER_GRPC_MESSAGE, message.to_owned()),
    ]);
    httpbis::SimpleHttpMessage {
        headers,
        body: BytesDeque::new(),
    }
}

// Trailers -> Status [Status-Message] *Custom-Metadata
pub(crate) fn trailers(
    grpc_status: GrpcStatus,
    message: Option<String>,
    metadata: Metadata,
) -> Headers {
    let mut headers = Headers::from_vec(vec![Header::new(
        HEADER_GRPC_STATUS,
        format!("{}", grpc_status as i32),
    )]);
    if let Some(message) = message {
        headers.add_header(Header::new(HEADER_GRPC_MESSAGE, message));
    }
    headers.extend(metadata.into_headers());
    headers
}
<|endoftext|>


<commit_before>use core::ptr::{self};
use core::marker::PhantomData;
use core::cmp;

use matrix::{Scalar,Mat,Hierarch,Matrix,ResizableBuffer,HierarchyNode,RoCM};
use typenum::Unsigned;
use thread_comm::ThreadInfo;
use composables::{GemmNode,AlgorithmStep};

pub trait Adder <T: Scalar, At: Mat<T>, Apt: Mat<T>> {
    fn add(a: &mut At, a_pack: &mut Apt, thr: &ThreadInfo<T>);
}

//Adds Apt to At.
pub struct Unpacker<T: Scalar, At: Mat<T>, Apt: Mat<T>> {
    _t: PhantomData<T>,
    _at: PhantomData<At>,
    _apt: PhantomData<Apt>,
}

//Default implementation of Unpacker. Uses the getters and setters of Mat<T>
impl<T: Scalar, At: Mat<T>, Apt: Mat<T>> Adder<T, At, Apt>
    for Unpacker<T, At, Apt> {
    default fn add(a: &mut At, a_pack: &mut Apt, thr: &ThreadInfo<T>) {
        if a_pack.width() <= 0 || a_pack.height() <= 0 {
            return;
        }
        let cols_per_thread = (a.width()-1) / thr.num_threads() + 1;
        let start = cols_per_thread * thr.thread_id();
        let end = cmp::min(a.width(), start+cols_per_thread);

        for x in start..end {
            for y in 0..a.height() {
                let alpha = a_pack.get(y,x) + a.get(y,x);
                a.set(y,x,alpha);
            }
        }
    }
}

//returns the depth and score of the level with best parallelizability
fn score_parallelizability(m: usize, y_hier: &[HierarchyNode]) -> (usize, f64)  {
    let mut best_depth = 0;
    let mut best_score = 0.0;
    let mut m_tracker = m;

    for i in 0..y_hier.len() {
        let blksz = y_hier[i].blksz;
        let n_iter = (m_tracker as f64) / (blksz as f64);
        let score = n_iter;
        if score > best_score {
            best_score = score;
            best_depth = i;
        }
        m_tracker = blksz;
    }
    (best_depth, best_score)
}

fn unpack_hier_leaf<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
	(a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>,
	 x_parallelize_level: isize, x_threads: usize, x_id: usize,
     y_parallelize_level: isize, y_threads: usize, y_id: usize) {
    //Parallelize Y direction
    let mut ystart = 0;
    let mut yend = a.height();
    if y_parallelize_level == 0 {
		let rows_per_thread = (a.height()-1) / y_threads + 1; // micro-panels per thread
		ystart = rows_per_thread*y_id;
		yend   = ystart+rows_per_thread;
    }
    //Parallelize X direction
    let mut xstart = 0;
    let mut xend = a.width();
    if x_parallelize_level == 0 {
		let rows_per_thread = (a.width()-1) / x_threads + 1; // micro-panels per thread
		xstart = rows_per_thread*x_id;
		xend   = xstart+rows_per_thread;
    }
    unsafe{
        let cs_a = a.get_column_stride();
        let rs_a = a.get_row_stride();
        let ap = a.get_mut_buffer();

        let a_pack_p = a_pack.get_mut_buffer();

        for y in ystart..yend {
            for x in xstart..xend {
                let alpha_a = ptr::read(ap.offset((y*rs_a + x*cs_a) as isize));
                let alpha_ap = ptr::read(a_pack_p.offset((y*LRS::to_usize() + x*LCS::to_usize()) as isize));
                ptr::write(ap.offset((y*rs_a + x*cs_a) as isize), alpha_a + alpha_ap);
            }
        }
    }
}
fn unpack_hier_y<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    (a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, y_hier: &[HierarchyNode],
    x_parallelize_level: isize, x_threads: usize, x_id: usize,
    y_parallelize_level: isize, y_threads: usize, y_id: usize) {
    if y_hier.len()-1 == 0 {
        unpack_hier_leaf(a, a_pack, x_parallelize_level, x_threads, x_id, y_parallelize_level, y_threads, y_id);
    } else {
        let blksz = y_hier[0].blksz;
        let m_save = a.push_y_view(blksz);
        a_pack.push_y_view(blksz);

        let mut start = 0;
        let mut end = m_save;
        if y_parallelize_level == 0 {
			let range = m_save;
			let n_blocks = (range-1) / blksz + 1;
			let blocks_per_thread = (n_blocks-1) / y_threads + 1; // micro-panels per thread
			start = blksz*blocks_per_thread*y_id;
			end   = start+blksz*blocks_per_thread;
        }
        let mut i = start;
        while i < end {
            a.slide_y_view_to(i, blksz);
            a_pack.slide_y_view_to(i, blksz);

            unpack_hier_y(a, a_pack, &y_hier[1..y_hier.len()],
                x_parallelize_level, x_threads, x_id,
                y_parallelize_level-1, y_threads, y_id);
            i += blksz;
        }
        a.pop_y_view();
        a_pack.pop_y_view();
    }
}
fn unpack_hier_x<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    (a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, x_hier: &[HierarchyNode], y_hier: &[HierarchyNode],
	 x_parallelize_level: isize, x_threads: usize, x_id: usize,
     y_parallelize_level: isize, y_threads: usize, y_id: usize)
{
    if x_hier.len() - 1 == 0 {
        unpack_hier_y(a, a_pack, y_hier, x_parallelize_level, x_threads, x_id, y_parallelize_level, y_threads, y_id);
    } else {
        let blksz = x_hier[0].blksz;
        let n_save = a.push_x_view(blksz);
        a_pack.push_x_view(blksz);

        let mut start = 0;
        let mut end = n_save;
        if x_parallelize_level == 0 {
			let range = n_save;
			let n_blocks = (range-1) / blksz + 1;
			let blocks_per_thread = (n_blocks-1) / x_threads + 1; // micro-panels per thread
			start = blksz*blocks_per_thread*x_id;
			end   = start+blksz*blocks_per_thread;
        }
        let mut j = start;
        while j < end  {
            a.slide_x_view_to(j, blksz);
            a_pack.slide_x_view_to(j, blksz);
            unpack_hier_x(a, a_pack, &x_hier[1..x_hier.len()], y_hier,
                x_parallelize_level-1, x_threads, x_id,
                y_parallelize_level, y_threads, y_id);
            j += blksz;
        }
        a.pop_x_view();
        a_pack.pop_x_view();
    }
}

//Specialized implementation of Unpacker for adding to Matrix<T> from Hierarch<T>
impl<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    Adder<T, Matrix<T>, Hierarch<T, LH, LW, LRS, LCS>>
    for Unpacker<T, Matrix<T>, Hierarch<T, LH, LW, LRS, LCS>> {
    default fn add(a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, thr: &ThreadInfo<T>) {
        if a_pack.width() <= 0 || a_pack.height() <= 0 {
            return;
        }

        //Get copies of the x and y hierarchy.
        //Since we borrow a_pack as mutable during pack_hier_x,
        //we can't borrow x_hier and y_hier immutably so we must copy
        let x_hier : Vec<HierarchyNode> =
        {
            let tmp1 = a_pack.get_x_hierarchy();
            let mut tmp2 : Vec<HierarchyNode> = Vec::with_capacity(tmp1.len());
            tmp2.extend_from_slice(tmp1);
            tmp2
        };
        let y_hier : Vec<HierarchyNode> =
        {
            let tmp1 = a_pack.get_y_hierarchy();
            let mut tmp2 : Vec<HierarchyNode> = Vec::with_capacity(tmp1.len());
            tmp2.extend_from_slice(tmp1);
            tmp2
        };
        let (y_depth, y_score) = score_parallelizability(a.height(), &y_hier);
        let (x_depth, x_score) = score_parallelizability(a.width(), &x_hier);

		//Figure out x and y num threads
		let mut index = f64::sqrt(thr.num_threads() as f64) as usize;
        while (thr.num_threads() % index) != 0 {
            index = index - 1;
        }
        let (y_nt, x_nt) = (4,1);
            if y_score < x_score {
                (index, thr.num_threads() / index)
            } else {
                (thr.num_threads() / index, index)
            };

        let x_tid = thr.thread_id() / y_nt;
        let y_tid = thr.thread_id() % y_nt;

        unpack_hier_x(a, a_pack, &x_hier, &y_hier, x_depth as isize, x_nt, x_tid, y_depth as isize, y_nt, y_tid);
    }
}

pub struct UnpackC<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>,
    S: GemmNode<T, At, Bt, CPt>> {
    child: S,
    c_pack: CPt,
    _t: PhantomData<T>,
    _at: PhantomData<At>,
    _bt: PhantomData<Bt>,
    _ct: PhantomData<Ct>,
}
impl<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>, S: GemmNode<T, At, Bt, CPt>> UnpackC<T,At,Bt,Ct,CPt,S>
    where CPt: ResizableBuffer<T> {
}
impl<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>, S: GemmNode<T, At, Bt, CPt>>
    GemmNode<T, At, Bt, Ct> for UnpackC<T, At, Bt, Ct, CPt, S>
    where CPt: ResizableBuffer<T> {
    #[inline(always)]
    unsafe fn run(&mut self, a: &mut At, b: &mut Bt, c:&mut Ct, thr: &ThreadInfo<T>) -> () {
        let algo_desc = S::hierarchy_description();
        let y_marker = AlgorithmStep::M{bsz: 0};
        let x_marker = AlgorithmStep::N{bsz: 0};
        let capacity_for_cpt = CPt:: capacity_for(c, y_marker, x_marker, &algo_desc);

        thr.barrier();

        //Check if we need to resize packing buffer
        if self.c_pack.capacity() < capacity_for_cpt {
            if thr.thread_id() == 0 {
                self.c_pack.aquire_buffer_for(capacity_for_cpt);
            }
            else {
                self.c_pack.set_capacity(capacity_for_cpt);
            }
            self.c_pack.send_alias(thr);
        }

        //Logically resize the c_pack matrix
        self.c_pack.resize_to(c, y_marker, x_marker, &algo_desc);
        //thr.barrier();
        c_pack.set_scalar(T::zero());
        self.child.run(a, b, &mut self.c_pack, thr);
        thr.barrier();
        <Unpacker<T,Ct,CPt>>::add(c, &mut self.c_pack, thr);
        thr.barrier();
    }
    fn new() -> UnpackC<T, At, Bt, Ct, CPt, S>{
        let algo_desc = S::hierarchy_description();
        let y_marker = AlgorithmStep::M{bsz: 0};
        let x_marker = AlgorithmStep::N{bsz: 0};

        UnpackC{ child: S::new(),
               c_pack: CPt::empty(y_marker, x_marker, &algo_desc),
               _t: PhantomData, _at: PhantomData, _bt: PhantomData, _ct: PhantomData }
    }
    fn hierarchy_description() -> Vec<AlgorithmStep> {
        S::hierarchy_description()
    }
}
<commit_msg>fix typo<commit_after>use core::ptr::{self};
use core::marker::PhantomData;
use core::cmp;

use matrix::{Scalar,Mat,Hierarch,Matrix,ResizableBuffer,HierarchyNode,RoCM};
use typenum::Unsigned;
use thread_comm::ThreadInfo;
use composables::{GemmNode,AlgorithmStep};

pub trait Adder <T: Scalar, At: Mat<T>, Apt: Mat<T>> {
    fn add(a: &mut At, a_pack: &mut Apt, thr: &ThreadInfo<T>);
}

//Adds Apt to At.
pub struct Unpacker<T: Scalar, At: Mat<T>, Apt: Mat<T>> {
    _t: PhantomData<T>,
    _at: PhantomData<At>,
    _apt: PhantomData<Apt>,
}

//Default implementation of Unpacker. Uses the getters and setters of Mat<T>
impl<T: Scalar, At: Mat<T>, Apt: Mat<T>> Adder<T, At, Apt>
    for Unpacker<T, At, Apt> {
    default fn add(a: &mut At, a_pack: &mut Apt, thr: &ThreadInfo<T>) {
        if a_pack.width() <= 0 || a_pack.height() <= 0 {
            return;
        }
        let cols_per_thread = (a.width()-1) / thr.num_threads() + 1;
        let start = cols_per_thread * thr.thread_id();
        let end = cmp::min(a.width(), start+cols_per_thread);

        for x in start..end {
            for y in 0..a.height() {
                let alpha = a_pack.get(y,x) + a.get(y,x);
                a.set(y,x,alpha);
            }
        }
    }
}

//returns the depth and score of the level with best parallelizability
fn score_parallelizability(m: usize, y_hier: &[HierarchyNode]) -> (usize, f64)  {
    let mut best_depth = 0;
    let mut best_score = 0.0;
    let mut m_tracker = m;

    for i in 0..y_hier.len() {
        let blksz = y_hier[i].blksz;
        let n_iter = (m_tracker as f64) / (blksz as f64);
        let score = n_iter;
        if score > best_score {
            best_score = score;
            best_depth = i;
        }
        m_tracker = blksz;
    }
    (best_depth, best_score)
}

fn unpack_hier_leaf<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
	(a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>,
	 x_parallelize_level: isize, x_threads: usize, x_id: usize,
     y_parallelize_level: isize, y_threads: usize, y_id: usize) {
    //Parallelize Y direction
    let mut ystart = 0;
    let mut yend = a.height();
    if y_parallelize_level == 0 {
		let rows_per_thread = (a.height()-1) / y_threads + 1; // micro-panels per thread
		ystart = rows_per_thread*y_id;
		yend   = ystart+rows_per_thread;
    }
    //Parallelize X direction
    let mut xstart = 0;
    let mut xend = a.width();
    if x_parallelize_level == 0 {
		let rows_per_thread = (a.width()-1) / x_threads + 1; // micro-panels per thread
		xstart = rows_per_thread*x_id;
		xend   = xstart+rows_per_thread;
    }
    unsafe{
        let cs_a = a.get_column_stride();
        let rs_a = a.get_row_stride();
        let ap = a.get_mut_buffer();

        let a_pack_p = a_pack.get_mut_buffer();

        for y in ystart..yend {
            for x in xstart..xend {
                let alpha_a = ptr::read(ap.offset((y*rs_a + x*cs_a) as isize));
                let alpha_ap = ptr::read(a_pack_p.offset((y*LRS::to_usize() + x*LCS::to_usize()) as isize));
                ptr::write(ap.offset((y*rs_a + x*cs_a) as isize), alpha_a + alpha_ap);
            }
        }
    }
}
fn unpack_hier_y<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    (a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, y_hier: &[HierarchyNode],
    x_parallelize_level: isize, x_threads: usize, x_id: usize,
    y_parallelize_level: isize, y_threads: usize, y_id: usize) {
    if y_hier.len()-1 == 0 {
        unpack_hier_leaf(a, a_pack, x_parallelize_level, x_threads, x_id, y_parallelize_level, y_threads, y_id);
    } else {
        let blksz = y_hier[0].blksz;
        let m_save = a.push_y_view(blksz);
        a_pack.push_y_view(blksz);

        let mut start = 0;
        let mut end = m_save;
        if y_parallelize_level == 0 {
			let range = m_save;
			let n_blocks = (range-1) / blksz + 1;
			let blocks_per_thread = (n_blocks-1) / y_threads + 1; // micro-panels per thread
			start = blksz*blocks_per_thread*y_id;
			end   = start+blksz*blocks_per_thread;
        }
        let mut i = start;
        while i < end {
            a.slide_y_view_to(i, blksz);
            a_pack.slide_y_view_to(i, blksz);

            unpack_hier_y(a, a_pack, &y_hier[1..y_hier.len()],
                x_parallelize_level, x_threads, x_id,
                y_parallelize_level-1, y_threads, y_id);
            i += blksz;
        }
        a.pop_y_view();
        a_pack.pop_y_view();
    }
}
fn unpack_hier_x<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    (a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, x_hier: &[HierarchyNode], y_hier: &[HierarchyNode],
	 x_parallelize_level: isize, x_threads: usize, x_id: usize,
     y_parallelize_level: isize, y_threads: usize, y_id: usize)
{
    if x_hier.len() - 1 == 0 {
        unpack_hier_y(a, a_pack, y_hier, x_parallelize_level, x_threads, x_id, y_parallelize_level, y_threads, y_id);
    } else {
        let blksz = x_hier[0].blksz;
        let n_save = a.push_x_view(blksz);
        a_pack.push_x_view(blksz);

        let mut start = 0;
        let mut end = n_save;
        if x_parallelize_level == 0 {
			let range = n_save;
			let n_blocks = (range-1) / blksz + 1;
			let blocks_per_thread = (n_blocks-1) / x_threads + 1; // micro-panels per thread
			start = blksz*blocks_per_thread*x_id;
			end   = start+blksz*blocks_per_thread;
        }
        let mut j = start;
        while j < end  {
            a.slide_x_view_to(j, blksz);
            a_pack.slide_x_view_to(j, blksz);
            unpack_hier_x(a, a_pack, &x_hier[1..x_hier.len()], y_hier,
                x_parallelize_level-1, x_threads, x_id,
                y_parallelize_level, y_threads, y_id);
            j += blksz;
        }
        a.pop_x_view();
        a_pack.pop_x_view();
    }
}

//Specialized implementation of Unpacker for adding to Matrix<T> from Hierarch<T>
impl<T: Scalar, LH: Unsigned, LW: Unsigned, LRS: Unsigned, LCS: Unsigned>
    Adder<T, Matrix<T>, Hierarch<T, LH, LW, LRS, LCS>>
    for Unpacker<T, Matrix<T>, Hierarch<T, LH, LW, LRS, LCS>> {
    default fn add(a: &mut Matrix<T>, a_pack: &mut Hierarch<T, LH, LW, LRS, LCS>, thr: &ThreadInfo<T>) {
        if a_pack.width() <= 0 || a_pack.height() <= 0 {
            return;
        }

        //Get copies of the x and y hierarchy.
        //Since we borrow a_pack as mutable during pack_hier_x,
        //we can't borrow x_hier and y_hier immutably so we must copy
        let x_hier : Vec<HierarchyNode> =
        {
            let tmp1 = a_pack.get_x_hierarchy();
            let mut tmp2 : Vec<HierarchyNode> = Vec::with_capacity(tmp1.len());
            tmp2.extend_from_slice(tmp1);
            tmp2
        };
        let y_hier : Vec<HierarchyNode> =
        {
            let tmp1 = a_pack.get_y_hierarchy();
            let mut tmp2 : Vec<HierarchyNode> = Vec::with_capacity(tmp1.len());
            tmp2.extend_from_slice(tmp1);
            tmp2
        };
        let (y_depth, y_score) = score_parallelizability(a.height(), &y_hier);
        let (x_depth, x_score) = score_parallelizability(a.width(), &x_hier);

		//Figure out x and y num threads
		let mut index = f64::sqrt(thr.num_threads() as f64) as usize;
        while (thr.num_threads() % index) != 0 {
            index = index - 1;
        }
        let (y_nt, x_nt) = (4,1);
            if y_score < x_score {
                (index, thr.num_threads() / index)
            } else {
                (thr.num_threads() / index, index)
            };

        let x_tid = thr.thread_id() / y_nt;
        let y_tid = thr.thread_id() % y_nt;

        unpack_hier_x(a, a_pack, &x_hier, &y_hier, x_depth as isize, x_nt, x_tid, y_depth as isize, y_nt, y_tid);
    }
}

pub struct UnpackC<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>,
    S: GemmNode<T, At, Bt, CPt>> {
    child: S,
    c_pack: CPt,
    _t: PhantomData<T>,
    _at: PhantomData<At>,
    _bt: PhantomData<Bt>,
    _ct: PhantomData<Ct>,
}
impl<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>, S: GemmNode<T, At, Bt, CPt>> UnpackC<T,At,Bt,Ct,CPt,S>
    where CPt: ResizableBuffer<T> {
}
impl<T: Scalar, At: Mat<T>, Bt: Mat<T>, Ct: Mat<T>, CPt: Mat<T>, S: GemmNode<T, At, Bt, CPt>>
    GemmNode<T, At, Bt, Ct> for UnpackC<T, At, Bt, Ct, CPt, S>
    where CPt: ResizableBuffer<T> {
    #[inline(always)]
    unsafe fn run(&mut self, a: &mut At, b: &mut Bt, c:&mut Ct, thr: &ThreadInfo<T>) -> () {
        let algo_desc = S::hierarchy_description();
        let y_marker = AlgorithmStep::M{bsz: 0};
        let x_marker = AlgorithmStep::N{bsz: 0};
        let capacity_for_cpt = CPt:: capacity_for(c, y_marker, x_marker, &algo_desc);

        thr.barrier();

        //Check if we need to resize packing buffer
        if self.c_pack.capacity() < capacity_for_cpt {
            if thr.thread_id() == 0 {
                self.c_pack.aquire_buffer_for(capacity_for_cpt);
            }
            else {
                self.c_pack.set_capacity(capacity_for_cpt);
            }
            self.c_pack.send_alias(thr);
        }

        //Logically resize the c_pack matrix
        self.c_pack.resize_to(c, y_marker, x_marker, &algo_desc);
        //thr.barrier();
        self.c_pack.set_scalar(T::zero());
        self.child.run(a, b, &mut self.c_pack, thr);
        thr.barrier();
        <Unpacker<T,Ct,CPt>>::add(c, &mut self.c_pack, thr);
        thr.barrier();
    }
    fn new() -> UnpackC<T, At, Bt, Ct, CPt, S>{
        let algo_desc = S::hierarchy_description();
        let y_marker = AlgorithmStep::M{bsz: 0};
        let x_marker = AlgorithmStep::N{bsz: 0};

        UnpackC{ child: S::new(),
               c_pack: CPt::empty(y_marker, x_marker, &algo_desc),
               _t: PhantomData, _at: PhantomData, _bt: PhantomData, _ct: PhantomData }
    }
    fn hierarchy_description() -> Vec<AlgorithmStep> {
        S::hierarchy_description()
    }
}
<|endoftext|>