<commit_before>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bv;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bv::BitVec;
//! use bv::BitsMut;
//!
//! let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
//! bits.set_bit(5, true);
//! bits.set_bit(32, true);
//! let rs = RankSelect::new(bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use std::cmp;

use bv::BitVec;
use bv::Bits;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: BitVec<u8>,
    superblocks_1: Vec<u64>,
    superblocks_0: Vec<u64>,
    s: usize,
    k: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: BitVec<u8>, k: usize) -> RankSelect {
        let n = bits.len() as usize;
        let s = k * 32;

        RankSelect {
            n,
            s,
            k,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Return the used k (see `RankSelect::new()`).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &BitVec<u8> {
        &self.bits
    }

    /// Return i-th bit.
    pub fn get(&self, i: u64) -> bool {
        self.bits.get_bit(i)
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: u64) -> Option<u64> {
        if i >= self.n as u64 {
            None
        } else {
            let s = i / self.s as u64; // the superblock
            let b = i / 8; // the block
            let j = i % 8; // the bit in the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s as usize];
            // add the rank within the block
            let mask = ((2u16 << j) - 1) as u8;
            rank += (self.bits.get_block(b as usize) & mask).count_ones() as u64;
            // add the popcounts of blocks from the beginning of the current superblock
            // up to the current block
            for block in (s * self.s as u64 / 8)..b {
                let b = self.bits.get_block(block as usize);
                rank += b.count_ones() as u64;
            }

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: u64) -> Option<u64> {
        self.rank_1(i).map(|r| (i + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: u64) -> Option<u64> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u64) -> Option<u64> {
        self.select_x(
            j,
            &self.superblocks_1,
            |bit| bit != 0,
            |block| block.count_ones(),
        )
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u64) -> Option<u64> {
        self.select_x(
            j,
            &self.superblocks_0,
            |bit| bit == 0,
            |block| block.count_zeros(),
        )
    }

    fn select_x<F: Fn(u8) -> bool, C: Fn(u8) -> u32>(
        &self,
        j: u64,
        superblocks: &[u64],
        is_match: F,
        count_all: C,
    ) -> Option<u64> {
        let mut superblock = match superblocks.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = superblocks[superblock];

        let first_block = superblock * self.s / 8;
        for block in first_block..cmp::min(first_block + self.s / 8, self.bits.block_len()) {
            let b = self.bits.get_block(block);
            let p = count_all(b) as u64;
            if rank + p >= j {
                let mut bit = 0b1;
                // do not look at unused bits of the last block
                let max_bit = cmp::min(8, self.bits.len() - block as u64 * 8);
                for i in 0..max_bit {
                    rank += is_match(b & bit) as u64;
                    if rank == j {
                        return Some((first_block + block) as u64 * 8 + i);
                    }
                    bit <<= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u64) -> Option<u64> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, bits: &BitVec<u8>) -> Vec<u64> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u64 = 0;
    let mut i = 0;
    let nblocks = (bits.len() as f64 / 8.0).ceil() as usize;
    for block in 0..nblocks {
        let b = bits.get_block(block);
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank += if t {
            b.count_ones() as u64
        } else {
            b.count_zeros() as u64
        };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bv::bit_vec;
    use bv::BitVec;
    use bv::BitsMut;

    #[test]
    fn test_rank_select() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
        bits.set_bit(5, true);
        bits.set_bit(32, true);
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(7).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
        assert_eq!(rs.get(5), true);
        assert_eq!(rs.get(1), false);
        assert_eq!(rs.get(32), true);
    }

    #[test]
    fn test_select() {
        let bits: BitVec<u8> = bit_vec![true, false];
        let rs = RankSelect::new(bits, 1);

        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_1(0), None);

        assert_eq!(rs.select_0(1), Some(1));
        assert_eq!(rs.select_1(1), Some(0));

        assert_eq!(rs.select_0(2), None);
        assert_eq!(rs.select_1(2), None);
    }

    #[test]
    fn test_single_select() {
        let bits: BitVec<u8> = bit_vec![true];
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.select_1(0), None);
        assert_eq!(rs.select_1(1), Some(0));
        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_0(1), None);

        let bits: BitVec<u8> = bit_vec![false];
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.select_1(1), None);
        assert_eq!(rs.select_1(0), Some(0));
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1), Some(0));
        assert_eq!(rs.rank_0(0), Some(1));
        assert_eq!(rs.rank_1(0), Some(0));
    }

    #[test]
    fn test_rank_k() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 72);
        bits.set_bit(63, true);
        let rs = RankSelect::new(bits, 2);
        assert_eq!(rs.rank_1(63), Some(1));
        assert_eq!(rs.rank_1(64), Some(1));
        assert_eq!(rs.rank_1(71), Some(1));
    }
}
<commit_msg>fix select<commit_after>// Copyright 2014-2016 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Rank/Select data structure based on Gonzalez, Grabowski, Mäkinen, Navarro (2005).
//! This implementation uses only a single level of blocks, and performs well for large n.
//!
//! Example
//!
//! ```
//! extern crate bv;
//! # extern crate bio;
//! # fn main() {
//! use bio::data_structures::rank_select::RankSelect;
//! use bv::BitVec;
//! use bv::BitsMut;
//!
//! let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
//! bits.set_bit(5, true);
//! bits.set_bit(32, true);
//! let rs = RankSelect::new(bits, 1);
//! assert!(rs.rank(6).unwrap() == 1);
//! # }
//! ```

use std::cmp;

use bv::BitVec;
use bv::Bits;

/// A rank/select data structure.
#[derive(Serialize, Deserialize)]
pub struct RankSelect {
    n: usize,
    bits: BitVec<u8>,
    superblocks_1: Vec<u64>,
    superblocks_0: Vec<u64>,
    /// superblock size in bits
    s: usize,
    /// superblock size in 32 bits
    k: usize,
}

impl RankSelect {
    /// Create a new instance.
    ///
    /// # Arguments
    ///
    /// * `bits` - A bit vector.
    /// * `k` - Determines the size (k * 32 bits) of the superblocks.
    ///   A small k means faster rank query times at the expense of using more
    ///   space and slower select query times.
    ///   The data structure needs O(n + n log n / (k * 32)) bits with n being the bits of the given bitvector.
    ///   The data structure is succinct if k is chosen as a sublinear function of n
    ///   (e.g. k = (log n)² / 32).
    pub fn new(bits: BitVec<u8>, k: usize) -> RankSelect {
        let n = bits.len() as usize;
        let s = k * 32;

        RankSelect {
            n,
            s,
            k,
            superblocks_1: superblocks(true, n, s, &bits),
            superblocks_0: superblocks(false, n, s, &bits),
            bits,
        }
    }

    /// Return the used k (see `RankSelect::new()`).
    pub fn k(&self) -> usize {
        self.k
    }

    /// Get internal representation of bit vector.
    pub fn bits(&self) -> &BitVec<u8> {
        &self.bits
    }

    /// Return i-th bit.
    pub fn get(&self, i: u64) -> bool {
        self.bits.get_bit(i)
    }

    /// Get the 1-rank of a given bit, i.e. the number of 1-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_1(&self, i: u64) -> Option<u64> {
        if i >= self.n as u64 {
            None
        } else {
            let s = i / self.s as u64; // the superblock
            let b = i / 8; // the block
            let j = i % 8; // the bit in the block
                           // take the superblock rank
            let mut rank = self.superblocks_1[s as usize];
            // add the rank within the block
            let mask = ((2u16 << j) - 1) as u8;
            rank += (self.bits.get_block(b as usize) & mask).count_ones() as u64;
            // add the popcounts of blocks from the beginning of the current superblock
            // up to the current block
            for block in (s * self.s as u64 / 8)..b {
                let b = self.bits.get_block(block as usize);
                rank += b.count_ones() as u64;
            }

            Some(rank)
        }
    }

    /// Get the 0-rank of a given bit, i.e. the number of 0-bits in the bitvector up to i (inclusive).
    /// Complexity: O(k).
    ///
    /// # Arguments
    ///
    /// * `i` - Position of the bit to determine the rank for.
    pub fn rank_0(&self, i: u64) -> Option<u64> {
        self.rank_1(i).map(|r| (i + 1) - r)
    }

    /// Alias for `RankSelect::rank_1`.
    pub fn rank(&self, i: u64) -> Option<u64> {
        self.rank_1(i)
    }

    /// Get the smallest bit with a given 1-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_1(&self, j: u64) -> Option<u64> {
        self.select_x(
            j,
            &self.superblocks_1,
            |bit| bit != 0,
            |block| block.count_ones(),
        )
    }

    /// Get the smallest bit with a given 0-rank.
    /// Complexity: O(log (n / k) + k).
    ///
    /// # Arguments
    ///
    /// * `j` - The rank to find the smallest bit for.
    pub fn select_0(&self, j: u64) -> Option<u64> {
        self.select_x(
            j,
            &self.superblocks_0,
            |bit| bit == 0,
            |block| block.count_zeros(),
        )
    }

    fn select_x<F: Fn(u8) -> bool, C: Fn(u8) -> u32>(
        &self,
        j: u64,
        superblocks: &[u64],
        is_match: F,
        count_all: C,
    ) -> Option<u64> {
        let mut superblock = match superblocks.binary_search(&j) {
            Ok(i) | Err(i) => i, // superblock with same rank exists
        };
        if superblock > 0 {
            superblock -= 1;
        }
        let mut rank = superblocks[superblock];

        let first_block = superblock * self.s / 8;
        for block in first_block..cmp::min(first_block + self.s / 8, self.bits.block_len()) {
            let b = self.bits.get_block(block);
            let p = count_all(b) as u64;
            if rank + p >= j {
                let mut bit = 0b1;
                // do not look at unused bits of the last block
                let max_bit = cmp::min(8, self.bits.len() - block as u64 * 8);
                for i in 0..max_bit {
                    rank += is_match(b & bit) as u64;
                    if rank == j {
                        return Some(block as u64 * 8 + i);
                    }
                    bit <<= 1;
                }
            }
            rank += p;
        }

        None
    }

    /// Alias for `RankSelect::select_1`.
    pub fn select(&self, j: u64) -> Option<u64> {
        self.select_1(j)
    }
}

/// Create `n` superblocks of size `s` from a given bitvector.
fn superblocks(t: bool, n: usize, s: usize, bits: &BitVec<u8>) -> Vec<u64> {
    let mut superblocks = Vec::with_capacity(n / s + 1);
    let mut rank: u64 = 0;
    let mut i = 0;
    let nblocks = (bits.len() as f64 / 8.0).ceil() as usize;
    for block in 0..nblocks {
        let b = bits.get_block(block);
        if i % s == 0 {
            superblocks.push(rank);
        }
        rank += if t {
            b.count_ones() as u64
        } else {
            b.count_zeros() as u64
        };
        i += 8;
    }

    superblocks
}

#[cfg(test)]
mod tests {
    use super::*;
    use bv::bit_vec;
    use bv::BitVec;
    use bv::BitsMut;

    #[test]
    fn test_rank_select() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
        bits.set_bit(5, true);
        bits.set_bit(32, true);
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.rank_1(1).unwrap(), 0);
        assert_eq!(rs.rank_1(5).unwrap(), 1);
        assert_eq!(rs.rank_1(6).unwrap(), 1);
        assert_eq!(rs.rank_1(7).unwrap(), 1);
        assert_eq!(rs.rank_1(32).unwrap(), 2);
        assert_eq!(rs.rank_1(33).unwrap(), 2);
        assert_eq!(rs.rank_1(64), None);
        assert_eq!(rs.select_1(0).unwrap(), 0);
        assert_eq!(rs.select_1(1).unwrap(), 5);
        assert_eq!(rs.select_1(2).unwrap(), 32);
        assert_eq!(rs.rank_0(1).unwrap(), 2);
        assert_eq!(rs.rank_0(4).unwrap(), 5);
        assert_eq!(rs.rank_0(5).unwrap(), 5);
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1).unwrap(), 0);
        assert_eq!(rs.get(5), true);
        assert_eq!(rs.get(1), false);
        assert_eq!(rs.get(32), true);
    }

    #[test]
    fn test_rank_select2() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 64);
        bits.set_bit(5, true);
        bits.set_bit(32, true);
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.select_1(2).unwrap(), 32);
    }

    #[test]
    fn test_select() {
        let bits: BitVec<u8> = bit_vec![true, false];
        let rs = RankSelect::new(bits, 1);

        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_1(0), None);

        assert_eq!(rs.select_0(1), Some(1));
        assert_eq!(rs.select_1(1), Some(0));

        assert_eq!(rs.select_0(2), None);
        assert_eq!(rs.select_1(2), None);
    }

    #[test]
    fn test_single_select() {
        let bits: BitVec<u8> = bit_vec![true];
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.select_1(0), None);
        assert_eq!(rs.select_1(1), Some(0));
        assert_eq!(rs.select_0(0), Some(0));
        assert_eq!(rs.select_0(1), None);

        let bits: BitVec<u8> = bit_vec![false];
        let rs = RankSelect::new(bits, 1);
        assert_eq!(rs.select_1(1), None);
        assert_eq!(rs.select_1(0), Some(0));
        assert_eq!(rs.select_0(0), None);
        assert_eq!(rs.select_0(1), Some(0));
        assert_eq!(rs.rank_0(0), Some(1));
        assert_eq!(rs.rank_1(0), Some(0));
    }

    #[test]
    fn test_rank_k() {
        let mut bits: BitVec<u8> = BitVec::new_fill(false, 72);
        bits.set_bit(63, true);
        let rs = RankSelect::new(bits, 2);
        assert_eq!(rs.rank_1(63), Some(1));
        assert_eq!(rs.rank_1(64), Some(1));
        assert_eq!(rs.rank_1(71), Some(1));
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)]
#[template(path = "base.html")]
struct BaseTemplate<'a> {
    title: &'a str,
}

#[derive(Template)]
#[template(path = "child.html")]
struct ChildTemplate<'a> {
    _parent: BaseTemplate<'a>,
}

#[test]
fn test_use_base_directly() {
    let t = BaseTemplate { title: "Foo" };
    assert_eq!(t.render().unwrap(), "Foo\n\nFoo\nCopyright 2017");
}

#[test]
fn test_simple_extends() {
    let t = ChildTemplate { _parent: BaseTemplate { title: "Bar" } };
    assert_eq!(
        t.render().unwrap(),
        "Bar\n(Bar) Content goes here\nFoo\nCopyright 2017"
    );
}

pub mod parent {
    use askama::Template;
    #[derive(Template)]
    #[template(path = "base.html")]
    pub struct BaseTemplate<'a> {
        pub title: &'a str,
    }
}

pub mod child {
    use askama::Template;
    use super::parent::*;
    #[derive(Template)]
    #[template(path = "child.html")]
    pub struct ChildTemplate<'a> {
        pub _parent: BaseTemplate<'a>,
    }
}

#[test]
fn test_different_module() {
    let t = child::ChildTemplate { _parent: parent::BaseTemplate { title: "a" } };
    assert_eq!(
        t.render().unwrap(),
        "a\n(a) Content goes here\nFoo\nCopyright 2017"
    );
}

#[derive(Template)]
#[template(path = "nested-base.html", print = "code")]
struct NestedBaseTemplate {}

#[derive(Template)]
#[template(path = "nested-child.html", print = "code")]
struct NestedChildTemplate {
    _parent: NestedBaseTemplate,
}

#[test]
fn test_nested_blocks() {
    let t = NestedChildTemplate { _parent: NestedBaseTemplate {} };
    assert_eq!(t.render().unwrap(), "\ndurpy\n");
}
<commit_msg>Remove forgotten debugging cruft<commit_after>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)]
#[template(path = "base.html")]
struct BaseTemplate<'a> {
    title: &'a str,
}

#[derive(Template)]
#[template(path = "child.html")]
struct ChildTemplate<'a> {
    _parent: BaseTemplate<'a>,
}

#[test]
fn test_use_base_directly() {
    let t = BaseTemplate { title: "Foo" };
    assert_eq!(t.render().unwrap(), "Foo\n\nFoo\nCopyright 2017");
}

#[test]
fn test_simple_extends() {
    let t = ChildTemplate { _parent: BaseTemplate { title: "Bar" } };
    assert_eq!(
        t.render().unwrap(),
        "Bar\n(Bar) Content goes here\nFoo\nCopyright 2017"
    );
}

pub mod parent {
    use askama::Template;
    #[derive(Template)]
    #[template(path = "base.html")]
    pub struct BaseTemplate<'a> {
        pub title: &'a str,
    }
}

pub mod child {
    use askama::Template;
    use super::parent::*;
    #[derive(Template)]
    #[template(path = "child.html")]
    pub struct ChildTemplate<'a> {
        pub _parent: BaseTemplate<'a>,
    }
}

#[test]
fn test_different_module() {
    let t = child::ChildTemplate { _parent: parent::BaseTemplate { title: "a" } };
    assert_eq!(
        t.render().unwrap(),
        "a\n(a) Content goes here\nFoo\nCopyright 2017"
    );
}

#[derive(Template)]
#[template(path = "nested-base.html")]
struct NestedBaseTemplate {}

#[derive(Template)]
#[template(path = "nested-child.html")]
struct NestedChildTemplate {
    _parent: NestedBaseTemplate,
}

#[test]
fn test_nested_blocks() {
    let t = NestedChildTemplate { _parent: NestedBaseTemplate {} };
    assert_eq!(t.render().unwrap(), "\ndurpy\n");
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::error::Error;
use std::future::Future;
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{fmt, mem, thread};

/// Declares a new task-local key of type [`tokio::task::LocalKey`].
///
/// # Syntax
///
/// The macro wraps any number of static declarations and makes them local to the current task.
/// Publicity and attributes for each static is preserved. For example:
///
/// # Examples
///
/// ```
/// # use tokio::task_local;
/// task_local! {
///     pub static ONE: u32;
///
///     #[allow(unused)]
///     static TWO: f32;
/// }
/// # fn main() {}
/// ```
///
/// See [LocalKey documentation][`tokio::task::LocalKey`] for more
/// information.
///
/// [`tokio::task::LocalKey`]: struct@crate::task::LocalKey
#[macro_export]
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
macro_rules! task_local {
     // empty (base case for the recursion)
    () => {};

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty; $($rest:tt)*) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
        $crate::task_local!($($rest)*);
    };

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
    }
}

#[doc(hidden)]
#[cfg(tokio_const_thread_local)]
#[macro_export]
macro_rules! __task_local_inner {
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
        $vis static $name: $crate::task::LocalKey<$t> = {
            std::thread_local! {
                static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) };
            }

            $crate::task::LocalKey { inner: __KEY }
        };
    };
}

#[doc(hidden)]
#[cfg(not(tokio_const_thread_local))]
#[macro_export]
macro_rules! __task_local_inner {
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
        $vis static $name: $crate::task::LocalKey<$t> = {
            std::thread_local! {
                static __KEY: std::cell::RefCell<Option<$t>> = std::cell::RefCell::new(None);
            }

            $crate::task::LocalKey { inner: __KEY }
        };
    };
}

/// A key for task-local data.
///
/// This type is generated by the [`task_local!`] macro.
///
/// Unlike [`std::thread::LocalKey`], `tokio::task::LocalKey` will
/// _not_ lazily initialize the value on first access. Instead, the
/// value is first initialized when the future containing
/// the task-local is first polled by a futures executor, like Tokio.
///
/// # Examples
///
/// ```
/// # async fn dox() {
/// tokio::task_local! {
///     static NUMBER: u32;
/// }
///
/// NUMBER.scope(1, async move {
///     assert_eq!(NUMBER.get(), 1);
/// }).await;
///
/// NUMBER.scope(2, async move {
///     assert_eq!(NUMBER.get(), 2);
///
///     NUMBER.scope(3, async move {
///         assert_eq!(NUMBER.get(), 3);
///     }).await;
/// }).await;
/// # }
/// ```
///
/// [`std::thread::LocalKey`]: struct@std::thread::LocalKey
/// [`task_local!`]: ../macro.task_local.html
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub struct LocalKey<T: 'static> {
    #[doc(hidden)]
    pub inner: thread::LocalKey<RefCell<Option<T>>>,
}

impl<T: 'static> LocalKey<T> {
    /// Sets a value `T` as the task-local value for the future `F`.
    ///
    /// On completion of `scope`, the task-local will be dropped.
    ///
    /// ### Panics
    ///
    /// If you poll the returned future inside a call to [`with`] or
    /// [`try_with`] on the same `LocalKey`, then the call to `poll` will panic.
    ///
    /// ### Examples
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.scope(1, async move {
    ///     println!("task local value: {}", NUMBER.get());
    /// }).await;
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    pub fn scope<F>(&'static self, value: T, f: F) -> TaskLocalFuture<T, F>
    where
        F: Future,
    {
        TaskLocalFuture {
            local: self,
            slot: Some(value),
            future: Some(f),
            _pinned: PhantomPinned,
        }
    }

    /// Sets a value `T` as the task-local value for the closure `F`.
    ///
    /// On completion of `scope`, the task-local will be dropped.
    ///
    /// ### Panics
    ///
    /// This method panics if called inside a call to [`with`] or [`try_with`]
    /// on the same `LocalKey`.
    ///
    /// ### Examples
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.sync_scope(1, || {
    ///     println!("task local value: {}", NUMBER.get());
    /// });
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    #[track_caller]
    pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        let mut value = Some(value);
        match self.scope_inner(&mut value, f) {
            Ok(res) => res,
            Err(err) => err.panic(),
        }
    }

    fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
    where
        F: FnOnce() -> R,
    {
        struct Guard<'a, T: 'static> {
            local: &'static LocalKey<T>,
            slot: &'a mut Option<T>,
        }

        impl<'a, T: 'static> Drop for Guard<'a, T> {
            fn drop(&mut self) {
                // This should not panic.
                //
                // We know that the RefCell was not borrowed before the call to
                // `scope_inner`, so the only way for this to panic is if the
                // closure has created but not destroyed a RefCell guard.
                // However, we never give user-code access to the guards, so
                // there's no way for user-code to forget to destroy a guard.
                //
                // The call to `with` also should not panic, since the
                // thread-local wasn't destroyed when we first called
                // `scope_inner`, and it shouldn't have gotten destroyed since
                // then.
                self.local.inner.with(|inner| {
                    let mut ref_mut = inner.borrow_mut();
                    mem::swap(self.slot, &mut *ref_mut);
                });
            }
        }

        self.inner.try_with(|inner| {
            inner
                .try_borrow_mut()
                .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut))
        })??;

        let guard = Guard { local: self, slot };

        let res = f();

        drop(guard);

        Ok(res)
    }

    /// Accesses the current task-local and runs the provided closure.
    ///
    /// # Panics
    ///
    /// This function will panic if the task local doesn't have a value set.
    #[track_caller]
    pub fn with<F, R>(&'static self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        match self.try_with(f) {
            Ok(res) => res,
            Err(_) => panic!("cannot access a task-local storage value without setting it first"),
        }
    }

    /// Accesses the current task-local and runs the provided closure.
    ///
    /// If the task-local with the associated key is not present, this
    /// method will return an `AccessError`. For a panicking variant,
    /// see `with`.
    pub fn try_with<F, R>(&'static self, f: F) -> Result<R, AccessError>
    where
        F: FnOnce(&T) -> R,
    {
        // If called after the thread-local storing the task-local is destroyed,
        // then we are outside of a closure where the task-local is set.
        //
        // Therefore, it is correct to return an AccessError if `try_with`
        // returns an error.
        let try_with_res = self.inner.try_with(|v| {
            // This call to `borrow` cannot panic because no user-defined code
            // runs while a `borrow_mut` call is active.
            v.borrow().as_ref().map(f)
        });

        match try_with_res {
            Ok(Some(res)) => Ok(res),
            Ok(None) | Err(_) => Err(AccessError { _private: () }),
        }
    }
}

impl<T: Copy + 'static> LocalKey<T> {
    /// Returns a copy of the task-local value
    /// if the task-local value implements `Copy`.
    ///
    /// # Panics
    ///
    /// This function will panic if the task local doesn't have a value set.
    pub fn get(&'static self) -> T {
        self.with(|v| *v)
    }
}

impl<T: 'static> fmt::Debug for LocalKey<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("LocalKey { .. }")
    }
}

/// A future that sets a value `T` of a task local for the future `F` during
/// its execution.
///
/// The value of the task-local must be `'static` and will be dropped on the
/// completion of the future.
///
/// Created by the function [`LocalKey::scope`](self::LocalKey::scope).
///
/// ### Examples
///
/// ```
/// # async fn dox() {
/// tokio::task_local! {
///     static NUMBER: u32;
/// }
///
/// NUMBER.scope(1, async move {
///     println!("task local value: {}", NUMBER.get());
/// }).await;
/// # }
/// ```
// Doesn't use pin_project due to custom Drop.
pub struct TaskLocalFuture<T, F>
where
    T: 'static,
{
    local: &'static LocalKey<T>,
    slot: Option<T>,
    future: Option<F>,
    _pinned: PhantomPinned,
}

impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
    type Output = F::Output;

    #[track_caller]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // safety: The TaskLocalFuture struct is `!Unpin` so there is no way to
        // move `self.future` from now on.
        let this = unsafe { Pin::into_inner_unchecked(self) };
        let mut future_opt = unsafe { Pin::new_unchecked(&mut this.future) };

        let res =
            this.local
                .scope_inner(&mut this.slot, || match future_opt.as_mut().as_pin_mut() {
                    Some(fut) => {
                        let res = fut.poll(cx);
                        if res.is_ready() {
                            future_opt.set(None);
                        }
                        Some(res)
                    }
                    None => None,
                });

        match res {
            Ok(Some(res)) => res,
            Ok(None) => panic!("`TaskLocalFuture` polled after completion"),
            Err(err) => err.panic(),
        }
    }
}

impl<T: 'static, F> Drop for TaskLocalFuture<T, F> {
    fn drop(&mut self) {
        if mem::needs_drop::<F>() && self.future.is_some() {
            // Drop the future while the task-local is set, if possible. Otherwise
            // the future is dropped normally when the `Option<F>` field drops.
            let future = &mut self.future;
            let _ = self.local.scope_inner(&mut self.slot, || {
                *future = None;
            });
        }
    }
}

impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        /// Format the Option without Some.
        struct TransparentOption<'a, T> {
            value: &'a Option<T>,
        }
        impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                match self.value.as_ref() {
                    Some(value) => value.fmt(f),
                    // Hitting the None branch should not be possible.
                    None => f.pad("<missing>"),
                }
            }
        }

        f.debug_struct("TaskLocalFuture")
            .field("value", &TransparentOption { value: &self.slot })
            .finish()
    }
}

/// An error returned by [`LocalKey::try_with`](method@LocalKey::try_with).
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct AccessError {
    _private: (),
}

impl fmt::Debug for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AccessError").finish()
    }
}

impl fmt::Display for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt("task-local value not set", f)
    }
}

impl Error for AccessError {}

enum ScopeInnerErr {
    BorrowError,
    AccessError,
}

impl ScopeInnerErr {
    fn panic(&self) -> ! {
        match self {
            Self::BorrowError => panic!("cannot enter a task-local scope while the task-local storage is borrowed"),
            Self::AccessError => panic!("cannot enter a task-local scope during or after destruction of the underlying thread-local"),
        }
    }
}

impl From<std::cell::BorrowMutError> for ScopeInnerErr {
    fn from(_: std::cell::BorrowMutError) -> Self {
        Self::BorrowError
    }
}

impl From<std::thread::AccessError> for ScopeInnerErr {
    fn from(_: std::thread::AccessError) -> Self {
        Self::AccessError
    }
}
<commit_msg>task: propagate attributes on task-locals (#4837)<commit_after>use std::cell::RefCell;
use std::error::Error;
use std::future::Future;
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{fmt, mem, thread};

/// Declares a new task-local key of type [`tokio::task::LocalKey`].
///
/// # Syntax
///
/// The macro wraps any number of static declarations and makes them local to the current task.
/// Publicity and attributes for each static is preserved. For example:
///
/// # Examples
///
/// ```
/// # use tokio::task_local;
/// task_local! {
///     pub static ONE: u32;
///
///     #[allow(unused)]
///     static TWO: f32;
/// }
/// # fn main() {}
/// ```
///
/// See [LocalKey documentation][`tokio::task::LocalKey`] for more
/// information.
///
/// [`tokio::task::LocalKey`]: struct@crate::task::LocalKey
#[macro_export]
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
macro_rules! task_local {
     // empty (base case for the recursion)
    () => {};

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty; $($rest:tt)*) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
        $crate::task_local!($($rest)*);
    };

    ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty) => {
        $crate::__task_local_inner!($(#[$attr])* $vis $name, $t);
    }
}

#[doc(hidden)]
#[cfg(tokio_const_thread_local)]
#[macro_export]
macro_rules! __task_local_inner {
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
        $(#[$attr])*
        $vis static $name: $crate::task::LocalKey<$t> = {
            std::thread_local! {
                static __KEY: std::cell::RefCell<Option<$t>> = const { std::cell::RefCell::new(None) };
            }

            $crate::task::LocalKey { inner: __KEY }
        };
    };
}

#[doc(hidden)]
#[cfg(not(tokio_const_thread_local))]
#[macro_export]
macro_rules! __task_local_inner {
    ($(#[$attr:meta])* $vis:vis $name:ident, $t:ty) => {
        $(#[$attr])*
        $vis static $name: $crate::task::LocalKey<$t> = {
            std::thread_local! {
                static __KEY: std::cell::RefCell<Option<$t>> = std::cell::RefCell::new(None);
            }

            $crate::task::LocalKey { inner: __KEY }
        };
    };
}

/// A key for task-local data.
///
/// This type is generated by the [`task_local!`] macro.
///
/// Unlike [`std::thread::LocalKey`], `tokio::task::LocalKey` will
/// _not_ lazily initialize the value on first access. Instead, the
/// value is first initialized when the future containing
/// the task-local is first polled by a futures executor, like Tokio.
///
/// # Examples
///
/// ```
/// # async fn dox() {
/// tokio::task_local! {
///     static NUMBER: u32;
/// }
///
/// NUMBER.scope(1, async move {
///     assert_eq!(NUMBER.get(), 1);
/// }).await;
///
/// NUMBER.scope(2, async move {
///     assert_eq!(NUMBER.get(), 2);
///
///     NUMBER.scope(3, async move {
///         assert_eq!(NUMBER.get(), 3);
///     }).await;
/// }).await;
/// # }
/// ```
///
/// [`std::thread::LocalKey`]: struct@std::thread::LocalKey
/// [`task_local!`]: ../macro.task_local.html
#[cfg_attr(docsrs, doc(cfg(feature = "rt")))]
pub struct LocalKey<T: 'static> {
    #[doc(hidden)]
    pub inner: thread::LocalKey<RefCell<Option<T>>>,
}

impl<T: 'static> LocalKey<T> {
    /// Sets a value `T` as the task-local value for the future `F`.
    ///
    /// On completion of `scope`, the task-local will be dropped.
    ///
    /// ### Panics
    ///
    /// If you poll the returned future inside a call to [`with`] or
    /// [`try_with`] on the same `LocalKey`, then the call to `poll` will panic.
    ///
    /// ### Examples
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.scope(1, async move {
    ///     println!("task local value: {}", NUMBER.get());
    /// }).await;
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    pub fn scope<F>(&'static self, value: T, f: F) -> TaskLocalFuture<T, F>
    where
        F: Future,
    {
        TaskLocalFuture {
            local: self,
            slot: Some(value),
            future: Some(f),
            _pinned: PhantomPinned,
        }
    }

    /// Sets a value `T` as the task-local value for the closure `F`.
    ///
    /// On completion of `scope`, the task-local will be dropped.
    ///
    /// ### Panics
    ///
    /// This method panics if called inside a call to [`with`] or [`try_with`]
    /// on the same `LocalKey`.
    ///
    /// ### Examples
    ///
    /// ```
    /// # async fn dox() {
    /// tokio::task_local! {
    ///     static NUMBER: u32;
    /// }
    ///
    /// NUMBER.sync_scope(1, || {
    ///     println!("task local value: {}", NUMBER.get());
    /// });
    /// # }
    /// ```
    ///
    /// [`with`]: fn@Self::with
    /// [`try_with`]: fn@Self::try_with
    #[track_caller]
    pub fn sync_scope<F, R>(&'static self, value: T, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        let mut value = Some(value);
        match self.scope_inner(&mut value, f) {
            Ok(res) => res,
            Err(err) => err.panic(),
        }
    }

    fn scope_inner<F, R>(&'static self, slot: &mut Option<T>, f: F) -> Result<R, ScopeInnerErr>
    where
        F: FnOnce() -> R,
    {
        struct Guard<'a, T: 'static> {
            local: &'static LocalKey<T>,
            slot: &'a mut Option<T>,
        }

        impl<'a, T: 'static> Drop for Guard<'a, T> {
            fn drop(&mut self) {
                // This should not panic.
                //
                // We know that the RefCell was not borrowed before the call to
                // `scope_inner`, so the only way for this to panic is if the
                // closure has created but not destroyed a RefCell guard.
                // However, we never give user-code access to the guards, so
                // there's no way for user-code to forget to destroy a guard.
                //
                // The call to `with` also should not panic, since the
                // thread-local wasn't destroyed when we first called
                // `scope_inner`, and it shouldn't have gotten destroyed since
                // then.
                self.local.inner.with(|inner| {
                    let mut ref_mut = inner.borrow_mut();
                    mem::swap(self.slot, &mut *ref_mut);
                });
            }
        }

        self.inner.try_with(|inner| {
            inner
                .try_borrow_mut()
                .map(|mut ref_mut| mem::swap(slot, &mut *ref_mut))
        })??;

        let guard = Guard { local: self, slot };

        let res = f();

        drop(guard);

        Ok(res)
    }

    /// Accesses the current task-local and runs the provided closure.
    ///
    /// # Panics
    ///
    /// This function will panic if the task local doesn't have a value set.
    #[track_caller]
    pub fn with<F, R>(&'static self, f: F) -> R
    where
        F: FnOnce(&T) -> R,
    {
        match self.try_with(f) {
            Ok(res) => res,
            Err(_) => panic!("cannot access a task-local storage value without setting it first"),
        }
    }

    /// Accesses the current task-local and runs the provided closure.
    ///
    /// If the task-local with the associated key is not present, this
    /// method will return an `AccessError`. For a panicking variant,
    /// see `with`.
    pub fn try_with<F, R>(&'static self, f: F) -> Result<R, AccessError>
    where
        F: FnOnce(&T) -> R,
    {
        // If called after the thread-local storing the task-local is destroyed,
        // then we are outside of a closure where the task-local is set.
        //
        // Therefore, it is correct to return an AccessError if `try_with`
        // returns an error.
        let try_with_res = self.inner.try_with(|v| {
            // This call to `borrow` cannot panic because no user-defined code
            // runs while a `borrow_mut` call is active.
            v.borrow().as_ref().map(f)
        });

        match try_with_res {
            Ok(Some(res)) => Ok(res),
            Ok(None) | Err(_) => Err(AccessError { _private: () }),
        }
    }
}

impl<T: Copy + 'static> LocalKey<T> {
    /// Returns a copy of the task-local value
    /// if the task-local value implements `Copy`.
    ///
    /// # Panics
    ///
    /// This function will panic if the task local doesn't have a value set.
    pub fn get(&'static self) -> T {
        self.with(|v| *v)
    }
}

impl<T: 'static> fmt::Debug for LocalKey<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("LocalKey { .. }")
    }
}

/// A future that sets a value `T` of a task local for the future `F` during
/// its execution.
///
/// The value of the task-local must be `'static` and will be dropped on the
/// completion of the future.
///
/// Created by the function [`LocalKey::scope`](self::LocalKey::scope).
///
/// ### Examples
///
/// ```
/// # async fn dox() {
/// tokio::task_local! {
///     static NUMBER: u32;
/// }
///
/// NUMBER.scope(1, async move {
///     println!("task local value: {}", NUMBER.get());
/// }).await;
/// # }
/// ```
// Doesn't use pin_project due to custom Drop.
pub struct TaskLocalFuture<T, F>
where
    T: 'static,
{
    local: &'static LocalKey<T>,
    slot: Option<T>,
    future: Option<F>,
    _pinned: PhantomPinned,
}

impl<T: 'static, F: Future> Future for TaskLocalFuture<T, F> {
    type Output = F::Output;

    #[track_caller]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // safety: The TaskLocalFuture struct is `!Unpin` so there is no way to
        // move `self.future` from now on.
        let this = unsafe { Pin::into_inner_unchecked(self) };
        let mut future_opt = unsafe { Pin::new_unchecked(&mut this.future) };

        let res =
            this.local
                .scope_inner(&mut this.slot, || match future_opt.as_mut().as_pin_mut() {
                    Some(fut) => {
                        let res = fut.poll(cx);
                        if res.is_ready() {
                            future_opt.set(None);
                        }
                        Some(res)
                    }
                    None => None,
                });

        match res {
            Ok(Some(res)) => res,
            Ok(None) => panic!("`TaskLocalFuture` polled after completion"),
            Err(err) => err.panic(),
        }
    }
}

impl<T: 'static, F> Drop for TaskLocalFuture<T, F> {
    fn drop(&mut self) {
        if mem::needs_drop::<F>() && self.future.is_some() {
            // Drop the future while the task-local is set, if possible. Otherwise
            // the future is dropped normally when the `Option<F>` field drops.
            let future = &mut self.future;
            let _ = self.local.scope_inner(&mut self.slot, || {
                *future = None;
            });
        }
    }
}

impl<T: 'static, F> fmt::Debug for TaskLocalFuture<T, F>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        /// Format the Option without Some.
        struct TransparentOption<'a, T> {
            value: &'a Option<T>,
        }
        impl<'a, T: fmt::Debug> fmt::Debug for TransparentOption<'a, T> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                match self.value.as_ref() {
                    Some(value) => value.fmt(f),
                    // Hitting the None branch should not be possible.
                    None => f.pad("<missing>"),
                }
            }
        }

        f.debug_struct("TaskLocalFuture")
            .field("value", &TransparentOption { value: &self.slot })
            .finish()
    }
}

/// An error returned by [`LocalKey::try_with`](method@LocalKey::try_with).
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct AccessError {
    _private: (),
}

impl fmt::Debug for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AccessError").finish()
    }
}

impl fmt::Display for AccessError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt("task-local value not set", f)
    }
}

impl Error for AccessError {}

enum ScopeInnerErr {
    BorrowError,
    AccessError,
}

impl ScopeInnerErr {
    fn panic(&self) -> ! {
        match self {
            Self::BorrowError => panic!("cannot enter a task-local scope while the task-local storage is borrowed"),
            Self::AccessError => panic!("cannot enter a task-local scope during or after destruction of the underlying thread-local"),
        }
    }
}

impl From<std::cell::BorrowMutError> for ScopeInnerErr {
    fn from(_: std::cell::BorrowMutError) -> Self {
        Self::BorrowError
    }
}

impl From<std::thread::AccessError> for ScopeInnerErr {
    fn from(_: std::thread::AccessError) -> Self {
        Self::AccessError
    }
}
<|endoftext|>


<commit_before>use std::convert::TryFrom;

use rustc_apfloat::ieee::{Double, Single};
use rustc_apfloat::{Float, FloatConvert};
use rustc_ast::ast::FloatTy;
use rustc_middle::mir::interpret::{InterpResult, PointerArithmetic, Scalar};
use rustc_middle::mir::CastKind;
use rustc_middle::ty::adjustment::PointerCast;
use rustc_middle::ty::layout::TyAndLayout;
use rustc_middle::ty::{self, Ty, TypeAndMut, TypeFoldable};
use rustc_span::symbol::sym;
use rustc_target::abi::{LayoutOf, Size, Variants};

use super::{truncate, FnVal, ImmTy, Immediate, InterpCx, Machine, OpTy, PlaceTy};

impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
    pub fn cast(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        cast_kind: CastKind,
        cast_ty: Ty<'tcx>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        use rustc_middle::mir::CastKind::*;
        // FIXME: In which cases should we trigger UB when the source is uninit?
        match cast_kind {
            Pointer(PointerCast::Unsize) => {
                let cast_ty = self.layout_of(cast_ty)?;
                self.unsize_into(src, cast_ty, dest)?;
            }

            Misc => {
                let src = self.read_immediate(src)?;
                let res = self.misc_cast(src, cast_ty)?;
                self.write_immediate(res, dest)?;
            }

            Pointer(PointerCast::MutToConstPointer | PointerCast::ArrayToPointer) => {
                // These are NOPs, but can be wide pointers.
                let v = self.read_immediate(src)?;
                self.write_immediate(*v, dest)?;
            }

            Pointer(PointerCast::ReifyFnPointer) => {
                // The src operand does not matter, just its type
                match src.layout.ty.kind {
                    ty::FnDef(def_id, substs) => {
                        // All reifications must be monomorphic, bail out otherwise.
                        if src.layout.ty.needs_subst() {
                            throw_inval!(TooGeneric);
                        }

                        if self.tcx.has_attr(def_id, sym::rustc_args_required_const) {
                            bug!("reifying a fn ptr that requires const arguments");
                        }

                        let instance = ty::Instance::resolve_for_fn_ptr(
                            *self.tcx,
                            self.param_env,
                            def_id,
                            substs,
                        )
                        .ok_or_else(|| err_inval!(TooGeneric))?;

                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(fn_ptr, dest)?;
                    }
                    _ => bug!("reify fn pointer on {:?}", src.layout.ty),
                }
            }

            Pointer(PointerCast::UnsafeFnPointer) => {
                let src = self.read_immediate(src)?;
                match cast_ty.kind {
                    ty::FnPtr(_) => {
                        // No change to value
                        self.write_immediate(*src, dest)?;
                    }
                    _ => bug!("fn to unsafe fn cast on {:?}", cast_ty),
                }
            }

            Pointer(PointerCast::ClosureFnPointer(_)) => {
                // The src operand does not matter, just its type
                match src.layout.ty.kind {
                    ty::Closure(def_id, substs) => {
                        // All reifications must be monomorphic, bail out otherwise.
                        if src.layout.ty.needs_subst() {
                            throw_inval!(TooGeneric);
                        }

                        let instance = ty::Instance::resolve_closure(
                            *self.tcx,
                            def_id,
                            substs,
                            ty::ClosureKind::FnOnce,
                        );
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(fn_ptr, dest)?;
                    }
                    _ => bug!("closure fn pointer on {:?}", src.layout.ty),
                }
            }
        }
        Ok(())
    }

    fn misc_cast(
        &self,
        src: ImmTy<'tcx, M::PointerTag>,
        cast_ty: Ty<'tcx>,
    ) -> InterpResult<'tcx, Immediate<M::PointerTag>> {
        use rustc_middle::ty::TyKind::*;
        trace!("Casting {:?}: {:?} to {:?}", *src, src.layout.ty, cast_ty);

        match src.layout.ty.kind {
            // Floating point
            Float(FloatTy::F32) => {
                return Ok(self.cast_from_float(src.to_scalar()?.to_f32()?, cast_ty).into());
            }
            Float(FloatTy::F64) => {
                return Ok(self.cast_from_float(src.to_scalar()?.to_f64()?, cast_ty).into());
            }
            // The rest is integer/pointer-"like", including fn ptr casts and casts from enums that
            // are represented as integers.
            _ => assert!(
                src.layout.ty.is_bool()
                    || src.layout.ty.is_char()
                    || src.layout.ty.is_enum()
                    || src.layout.ty.is_integral()
                    || src.layout.ty.is_any_ptr(),
                "Unexpected cast from type {:?}",
                src.layout.ty
            ),
        }

        // # First handle non-scalar source values.

        // Handle cast from a univariant (ZST) enum.
        match src.layout.variants {
            Variants::Single { index } => {
                if let Some(discr) = src.layout.ty.discriminant_for_variant(*self.tcx, index) {
                    assert!(src.layout.is_zst());
                    let discr_layout = self.layout_of(discr.ty)?;
                    return Ok(self.cast_from_scalar(discr.val, discr_layout, cast_ty).into());
                }
            }
            Variants::Multiple { .. } => {}
        }

        // Handle casting any ptr to raw ptr (might be a fat ptr).
        if src.layout.ty.is_any_ptr() && cast_ty.is_unsafe_ptr() {
            let dest_layout = self.layout_of(cast_ty)?;
            if dest_layout.size == src.layout.size {
                // Thin or fat pointer that just hast the ptr kind of target type changed.
                return Ok(*src);
            } else {
                // Casting the metadata away from a fat ptr.
                assert_eq!(src.layout.size, 2 * self.memory.pointer_size());
                assert_eq!(dest_layout.size, self.memory.pointer_size());
                assert!(src.layout.ty.is_unsafe_ptr());
                return match *src {
                    Immediate::ScalarPair(data, _) => Ok(data.into()),
                    Immediate::Scalar(..) => bug!(
                        "{:?} input to a fat-to-thin cast ({:?} -> {:?})",
                        *src,
                        src.layout.ty,
                        cast_ty
                    ),
                };
            }
        }

        // # The remaining source values are scalar.

        // For all remaining casts, we either
        // (a) cast a raw ptr to usize, or
        // (b) cast from an integer-like (including bool, char, enums).
        // In both cases we want the bits.
        let bits = self.force_bits(src.to_scalar()?, src.layout.size)?;
        Ok(self.cast_from_scalar(bits, src.layout, cast_ty).into())
    }

    pub(super) fn cast_from_scalar(
        &self,
        v: u128, // raw bits (there is no ScalarTy so we separate data+layout)
        src_layout: TyAndLayout<'tcx>,
        cast_ty: Ty<'tcx>,
    ) -> Scalar<M::PointerTag> {
        // Let's make sure v is sign-extended *if* it has a signed type.
        let signed = src_layout.abi.is_signed(); // also checks that abi is `Scalar`.
        let v = if signed { self.sign_extend(v, src_layout) } else { v };
        trace!("cast_from_scalar: {}, {} -> {}", v, src_layout.ty, cast_ty);
        use rustc_middle::ty::TyKind::*;
        match cast_ty.kind {
            Int(_) | Uint(_) | RawPtr(_) => {
                let size = match cast_ty.kind {
                    // FIXME: Isn't there a helper for this? The same pattern occurs below.
                    Int(t) => {
                        t.bit_width().map(Size::from_bits).unwrap_or_else(|| self.pointer_size())
                    }
                    Uint(t) => {
                        t.bit_width().map(Size::from_bits).unwrap_or_else(|| self.pointer_size())
                    }
                    RawPtr(_) => self.pointer_size(),
                    _ => bug!(),
                };
                let v = truncate(v, size);
                Scalar::from_uint(v, size)
            }

            Float(FloatTy::F32) if signed => Scalar::from_f32(Single::from_i128(v as i128).value),
            Float(FloatTy::F64) if signed => Scalar::from_f64(Double::from_i128(v as i128).value),
            Float(FloatTy::F32) => Scalar::from_f32(Single::from_u128(v).value),
            Float(FloatTy::F64) => Scalar::from_f64(Double::from_u128(v).value),

            Char => {
                // `u8` to `char` cast
                Scalar::from_u32(u8::try_from(v).unwrap().into())
            }

            // Casts to bool are not permitted by rustc, no need to handle them here.
            _ => bug!("invalid int to {:?} cast", cast_ty),
        }
    }

    fn cast_from_float<F>(&self, f: F, dest_ty: Ty<'tcx>) -> Scalar<M::PointerTag>
    where
        F: Float + Into<Scalar<M::PointerTag>> + FloatConvert<Single> + FloatConvert<Double>,
    {
        use rustc_middle::ty::TyKind::*;
        match dest_ty.kind {
            // float -> uint
            Uint(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits());
                // `to_u128` is a saturating cast, which is what we need
                // (https://doc.rust-lang.org/nightly/nightly-rustc/rustc_apfloat/trait.Float.html#method.to_i128_r).
                let v = f.to_u128(usize::try_from(width).unwrap()).value;
                // This should already fit the bit width
                Scalar::from_uint(v, Size::from_bits(width))
            }
            // float -> int
            Int(t) => {
                let width = t.bit_width().unwrap_or_else(|| self.pointer_size().bits());
                // `to_i128` is a saturating cast, which is what we need
                // (https://doc.rust-lang.org/nightly/nightly-rustc/rustc_apfloat/trait.Float.html#method.to_i128_r).
                let v = f.to_i128(usize::try_from(width).unwrap()).value;
                Scalar::from_int(v, Size::from_bits(width))
            }
            // float -> f32
            Float(FloatTy::F32) => Scalar::from_f32(f.convert(&mut false).value),
            // float -> f64
            Float(FloatTy::F64) => Scalar::from_f64(f.convert(&mut false).value),
            // That's it.
            _ => bug!("invalid float to {:?} cast", dest_ty),
        }
    }

    fn unsize_into_ptr(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
        // The pointee types
        source_ty: Ty<'tcx>,
        cast_ty: Ty<'tcx>,
    ) -> InterpResult<'tcx> {
        // A<Struct> -> A<Trait> conversion
        let (src_pointee_ty, dest_pointee_ty) =
            self.tcx.struct_lockstep_tails_erasing_lifetimes(source_ty, cast_ty, self.param_env);

        match (&src_pointee_ty.kind, &dest_pointee_ty.kind) {
            (&ty::Array(_, length), &ty::Slice(_)) => {
                let ptr = self.read_immediate(src)?.to_scalar()?;
                // u64 cast is from usize to u64, which is always good
                let val = Immediate::new_slice(
                    ptr,
                    length.eval_usize(self.tcx.tcx, self.param_env),
                    self,
                );
                self.write_immediate(val, dest)
            }
            (&ty::Dynamic(..), &ty::Dynamic(..)) => {
                // For now, upcasts are limited to changes in marker
                // traits, and hence never actually require an actual
                // change to the vtable.
                let val = self.read_immediate(src)?;
                self.write_immediate(*val, dest)
            }
            (_, &ty::Dynamic(ref data, _)) => {
                // Initial cast from sized to dyn trait
                let vtable = self.get_vtable(src_pointee_ty, data.principal())?;
                let ptr = self.read_immediate(src)?.to_scalar()?;
                let val = Immediate::new_dyn_trait(ptr, vtable);
                self.write_immediate(val, dest)
            }

            _ => bug!("invalid unsizing {:?} -> {:?}", src.layout.ty, cast_ty),
        }
    }

    fn unsize_into(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        cast_ty: TyAndLayout<'tcx>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        trace!("Unsizing {:?} of type {} into {:?}", *src, src.layout.ty, cast_ty.ty);
        match (&src.layout.ty.kind, &cast_ty.ty.kind) {
            (&ty::Ref(_, s, _), &ty::Ref(_, c, _) | &ty::RawPtr(TypeAndMut { ty: c, .. }))
            | (&ty::RawPtr(TypeAndMut { ty: s, .. }), &ty::RawPtr(TypeAndMut { ty: c, .. })) => {
                self.unsize_into_ptr(src, dest, s, c)
            }
            (&ty::Adt(def_a, _), &ty::Adt(def_b, _)) => {
                assert_eq!(def_a, def_b);
                if def_a.is_box() || def_b.is_box() {
                    if !def_a.is_box() || !def_b.is_box() {
                        bug!("invalid unsizing between {:?} -> {:?}", src.layout.ty, cast_ty.ty);
                    }
                    return self.unsize_into_ptr(
                        src,
                        dest,
                        src.layout.ty.boxed_ty(),
                        cast_ty.ty.boxed_ty(),
                    );
                }

                // unsizing of generic struct with pointer fields
                // Example: `Arc<T>` -> `Arc<Trait>`
                // here we need to increase the size of every &T thin ptr field to a fat ptr
                for i in 0..src.layout.fields.count() {
                    let cast_ty_field = cast_ty.field(self, i)?;
                    if cast_ty_field.is_zst() {
                        continue;
                    }
                    let src_field = self.operand_field(src, i)?;
                    let dst_field = self.place_field(dest, i)?;
                    if src_field.layout.ty == cast_ty_field.ty {
                        self.copy_op(src_field, dst_field)?;
                    } else {
                        self.unsize_into(src_field, cast_ty_field, dst_field)?;
                    }
                }
                Ok(())
            }
            _ => bug!("unsize_into: invalid conversion: {:?} -> {:?}", src.layout, dest.layout),
        }
    }
}
<commit_msg>use helper method for determining size of int type<commit_after>use std::convert::TryFrom;

use rustc_apfloat::ieee::{Double, Single};
use rustc_apfloat::{Float, FloatConvert};
use rustc_ast::ast::FloatTy;
use rustc_attr as attr;
use rustc_middle::mir::interpret::{InterpResult, PointerArithmetic, Scalar};
use rustc_middle::mir::CastKind;
use rustc_middle::ty::adjustment::PointerCast;
use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
use rustc_middle::ty::{self, Ty, TypeAndMut, TypeFoldable};
use rustc_span::symbol::sym;
use rustc_target::abi::{Integer, LayoutOf, Variants};

use super::{truncate, FnVal, ImmTy, Immediate, InterpCx, Machine, OpTy, PlaceTy};

impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
    pub fn cast(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        cast_kind: CastKind,
        cast_ty: Ty<'tcx>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        use rustc_middle::mir::CastKind::*;
        // FIXME: In which cases should we trigger UB when the source is uninit?
        match cast_kind {
            Pointer(PointerCast::Unsize) => {
                let cast_ty = self.layout_of(cast_ty)?;
                self.unsize_into(src, cast_ty, dest)?;
            }

            Misc => {
                let src = self.read_immediate(src)?;
                let res = self.misc_cast(src, cast_ty)?;
                self.write_immediate(res, dest)?;
            }

            Pointer(PointerCast::MutToConstPointer | PointerCast::ArrayToPointer) => {
                // These are NOPs, but can be wide pointers.
                let v = self.read_immediate(src)?;
                self.write_immediate(*v, dest)?;
            }

            Pointer(PointerCast::ReifyFnPointer) => {
                // The src operand does not matter, just its type
                match src.layout.ty.kind {
                    ty::FnDef(def_id, substs) => {
                        // All reifications must be monomorphic, bail out otherwise.
                        if src.layout.ty.needs_subst() {
                            throw_inval!(TooGeneric);
                        }

                        if self.tcx.has_attr(def_id, sym::rustc_args_required_const) {
                            bug!("reifying a fn ptr that requires const arguments");
                        }

                        let instance = ty::Instance::resolve_for_fn_ptr(
                            *self.tcx,
                            self.param_env,
                            def_id,
                            substs,
                        )
                        .ok_or_else(|| err_inval!(TooGeneric))?;

                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(fn_ptr, dest)?;
                    }
                    _ => bug!("reify fn pointer on {:?}", src.layout.ty),
                }
            }

            Pointer(PointerCast::UnsafeFnPointer) => {
                let src = self.read_immediate(src)?;
                match cast_ty.kind {
                    ty::FnPtr(_) => {
                        // No change to value
                        self.write_immediate(*src, dest)?;
                    }
                    _ => bug!("fn to unsafe fn cast on {:?}", cast_ty),
                }
            }

            Pointer(PointerCast::ClosureFnPointer(_)) => {
                // The src operand does not matter, just its type
                match src.layout.ty.kind {
                    ty::Closure(def_id, substs) => {
                        // All reifications must be monomorphic, bail out otherwise.
                        if src.layout.ty.needs_subst() {
                            throw_inval!(TooGeneric);
                        }

                        let instance = ty::Instance::resolve_closure(
                            *self.tcx,
                            def_id,
                            substs,
                            ty::ClosureKind::FnOnce,
                        );
                        let fn_ptr = self.memory.create_fn_alloc(FnVal::Instance(instance));
                        self.write_scalar(fn_ptr, dest)?;
                    }
                    _ => bug!("closure fn pointer on {:?}", src.layout.ty),
                }
            }
        }
        Ok(())
    }

    fn misc_cast(
        &self,
        src: ImmTy<'tcx, M::PointerTag>,
        cast_ty: Ty<'tcx>,
    ) -> InterpResult<'tcx, Immediate<M::PointerTag>> {
        use rustc_middle::ty::TyKind::*;
        trace!("Casting {:?}: {:?} to {:?}", *src, src.layout.ty, cast_ty);

        match src.layout.ty.kind {
            // Floating point
            Float(FloatTy::F32) => {
                return Ok(self.cast_from_float(src.to_scalar()?.to_f32()?, cast_ty).into());
            }
            Float(FloatTy::F64) => {
                return Ok(self.cast_from_float(src.to_scalar()?.to_f64()?, cast_ty).into());
            }
            // The rest is integer/pointer-"like", including fn ptr casts and casts from enums that
            // are represented as integers.
            _ => assert!(
                src.layout.ty.is_bool()
                    || src.layout.ty.is_char()
                    || src.layout.ty.is_enum()
                    || src.layout.ty.is_integral()
                    || src.layout.ty.is_any_ptr(),
                "Unexpected cast from type {:?}",
                src.layout.ty
            ),
        }

        // # First handle non-scalar source values.

        // Handle cast from a univariant (ZST) enum.
        match src.layout.variants {
            Variants::Single { index } => {
                if let Some(discr) = src.layout.ty.discriminant_for_variant(*self.tcx, index) {
                    assert!(src.layout.is_zst());
                    let discr_layout = self.layout_of(discr.ty)?;
                    return Ok(self.cast_from_scalar(discr.val, discr_layout, cast_ty).into());
                }
            }
            Variants::Multiple { .. } => {}
        }

        // Handle casting any ptr to raw ptr (might be a fat ptr).
        if src.layout.ty.is_any_ptr() && cast_ty.is_unsafe_ptr() {
            let dest_layout = self.layout_of(cast_ty)?;
            if dest_layout.size == src.layout.size {
                // Thin or fat pointer that just hast the ptr kind of target type changed.
                return Ok(*src);
            } else {
                // Casting the metadata away from a fat ptr.
                assert_eq!(src.layout.size, 2 * self.memory.pointer_size());
                assert_eq!(dest_layout.size, self.memory.pointer_size());
                assert!(src.layout.ty.is_unsafe_ptr());
                return match *src {
                    Immediate::ScalarPair(data, _) => Ok(data.into()),
                    Immediate::Scalar(..) => bug!(
                        "{:?} input to a fat-to-thin cast ({:?} -> {:?})",
                        *src,
                        src.layout.ty,
                        cast_ty
                    ),
                };
            }
        }

        // # The remaining source values are scalar.

        // For all remaining casts, we either
        // (a) cast a raw ptr to usize, or
        // (b) cast from an integer-like (including bool, char, enums).
        // In both cases we want the bits.
        let bits = self.force_bits(src.to_scalar()?, src.layout.size)?;
        Ok(self.cast_from_scalar(bits, src.layout, cast_ty).into())
    }

    pub(super) fn cast_from_scalar(
        &self,
        v: u128, // raw bits (there is no ScalarTy so we separate data+layout)
        src_layout: TyAndLayout<'tcx>,
        cast_ty: Ty<'tcx>,
    ) -> Scalar<M::PointerTag> {
        // Let's make sure v is sign-extended *if* it has a signed type.
        let signed = src_layout.abi.is_signed(); // also checks that abi is `Scalar`.
        let v = if signed { self.sign_extend(v, src_layout) } else { v };
        trace!("cast_from_scalar: {}, {} -> {}", v, src_layout.ty, cast_ty);
        use rustc_middle::ty::TyKind::*;
        match cast_ty.kind {
            Int(_) | Uint(_) | RawPtr(_) => {
                let size = match cast_ty.kind {
                    Int(t) => Integer::from_attr(self, attr::IntType::SignedInt(t)).size(),
                    Uint(t) => Integer::from_attr(self, attr::IntType::UnsignedInt(t)).size(),
                    RawPtr(_) => self.pointer_size(),
                    _ => bug!(),
                };
                let v = truncate(v, size);
                Scalar::from_uint(v, size)
            }

            Float(FloatTy::F32) if signed => Scalar::from_f32(Single::from_i128(v as i128).value),
            Float(FloatTy::F64) if signed => Scalar::from_f64(Double::from_i128(v as i128).value),
            Float(FloatTy::F32) => Scalar::from_f32(Single::from_u128(v).value),
            Float(FloatTy::F64) => Scalar::from_f64(Double::from_u128(v).value),

            Char => {
                // `u8` to `char` cast
                Scalar::from_u32(u8::try_from(v).unwrap().into())
            }

            // Casts to bool are not permitted by rustc, no need to handle them here.
            _ => bug!("invalid int to {:?} cast", cast_ty),
        }
    }

    fn cast_from_float<F>(&self, f: F, dest_ty: Ty<'tcx>) -> Scalar<M::PointerTag>
    where
        F: Float + Into<Scalar<M::PointerTag>> + FloatConvert<Single> + FloatConvert<Double>,
    {
        use rustc_middle::ty::TyKind::*;
        match dest_ty.kind {
            // float -> uint
            Uint(t) => {
                let size = Integer::from_attr(self, attr::IntType::UnsignedInt(t)).size();
                // `to_u128` is a saturating cast, which is what we need
                // (https://doc.rust-lang.org/nightly/nightly-rustc/rustc_apfloat/trait.Float.html#method.to_i128_r).
                let v = f.to_u128(size.bits_usize()).value;
                // This should already fit the bit width
                Scalar::from_uint(v, size)
            }
            // float -> int
            Int(t) => {
                let size = Integer::from_attr(self, attr::IntType::SignedInt(t)).size();
                // `to_i128` is a saturating cast, which is what we need
                // (https://doc.rust-lang.org/nightly/nightly-rustc/rustc_apfloat/trait.Float.html#method.to_i128_r).
                let v = f.to_i128(size.bits_usize()).value;
                Scalar::from_int(v, size)
            }
            // float -> f32
            Float(FloatTy::F32) => Scalar::from_f32(f.convert(&mut false).value),
            // float -> f64
            Float(FloatTy::F64) => Scalar::from_f64(f.convert(&mut false).value),
            // That's it.
            _ => bug!("invalid float to {:?} cast", dest_ty),
        }
    }

    fn unsize_into_ptr(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        dest: PlaceTy<'tcx, M::PointerTag>,
        // The pointee types
        source_ty: Ty<'tcx>,
        cast_ty: Ty<'tcx>,
    ) -> InterpResult<'tcx> {
        // A<Struct> -> A<Trait> conversion
        let (src_pointee_ty, dest_pointee_ty) =
            self.tcx.struct_lockstep_tails_erasing_lifetimes(source_ty, cast_ty, self.param_env);

        match (&src_pointee_ty.kind, &dest_pointee_ty.kind) {
            (&ty::Array(_, length), &ty::Slice(_)) => {
                let ptr = self.read_immediate(src)?.to_scalar()?;
                // u64 cast is from usize to u64, which is always good
                let val = Immediate::new_slice(
                    ptr,
                    length.eval_usize(self.tcx.tcx, self.param_env),
                    self,
                );
                self.write_immediate(val, dest)
            }
            (&ty::Dynamic(..), &ty::Dynamic(..)) => {
                // For now, upcasts are limited to changes in marker
                // traits, and hence never actually require an actual
                // change to the vtable.
                let val = self.read_immediate(src)?;
                self.write_immediate(*val, dest)
            }
            (_, &ty::Dynamic(ref data, _)) => {
                // Initial cast from sized to dyn trait
                let vtable = self.get_vtable(src_pointee_ty, data.principal())?;
                let ptr = self.read_immediate(src)?.to_scalar()?;
                let val = Immediate::new_dyn_trait(ptr, vtable);
                self.write_immediate(val, dest)
            }

            _ => bug!("invalid unsizing {:?} -> {:?}", src.layout.ty, cast_ty),
        }
    }

    fn unsize_into(
        &mut self,
        src: OpTy<'tcx, M::PointerTag>,
        cast_ty: TyAndLayout<'tcx>,
        dest: PlaceTy<'tcx, M::PointerTag>,
    ) -> InterpResult<'tcx> {
        trace!("Unsizing {:?} of type {} into {:?}", *src, src.layout.ty, cast_ty.ty);
        match (&src.layout.ty.kind, &cast_ty.ty.kind) {
            (&ty::Ref(_, s, _), &ty::Ref(_, c, _) | &ty::RawPtr(TypeAndMut { ty: c, .. }))
            | (&ty::RawPtr(TypeAndMut { ty: s, .. }), &ty::RawPtr(TypeAndMut { ty: c, .. })) => {
                self.unsize_into_ptr(src, dest, s, c)
            }
            (&ty::Adt(def_a, _), &ty::Adt(def_b, _)) => {
                assert_eq!(def_a, def_b);
                if def_a.is_box() || def_b.is_box() {
                    if !def_a.is_box() || !def_b.is_box() {
                        bug!("invalid unsizing between {:?} -> {:?}", src.layout.ty, cast_ty.ty);
                    }
                    return self.unsize_into_ptr(
                        src,
                        dest,
                        src.layout.ty.boxed_ty(),
                        cast_ty.ty.boxed_ty(),
                    );
                }

                // unsizing of generic struct with pointer fields
                // Example: `Arc<T>` -> `Arc<Trait>`
                // here we need to increase the size of every &T thin ptr field to a fat ptr
                for i in 0..src.layout.fields.count() {
                    let cast_ty_field = cast_ty.field(self, i)?;
                    if cast_ty_field.is_zst() {
                        continue;
                    }
                    let src_field = self.operand_field(src, i)?;
                    let dst_field = self.place_field(dest, i)?;
                    if src_field.layout.ty == cast_ty_field.ty {
                        self.copy_op(src_field, dst_field)?;
                    } else {
                        self.unsize_into(src_field, cast_ty_field, dst_field)?;
                    }
                }
                Ok(())
            }
            _ => bug!("unsize_into: invalid conversion: {:?} -> {:?}", src.layout, dest.layout),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Error types

use core::fmt;
use core::num::NonZeroU32;

/// Error type of random number generators
///
/// In order to be compatible with `std` and `no_std`, this type has two
/// possible implementations: with `std` a boxed `Error` trait object is stored,
/// while with `no_std` we merely store an error code.
pub struct Error {
    #[cfg(feature="std")]
    inner: Box<dyn std::error::Error + Send + Sync + 'static>,
    #[cfg(not(feature="std"))]
    code: NonZeroU32,
}

impl Error {
    /// Construct from any type supporting `std::error::Error`
    ///
    /// Available only when configured with `std`.
    ///
    /// See also `From<NonZeroU32>`, which is available with and without `std`.
    #[cfg(feature="std")]
    #[inline]
    pub fn new<E>(err: E) -> Self
    where E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>
    {
        Error { inner: err.into() }
    }

    /// Reference the inner error (`std` only)
    ///
    /// When configured with `std`, this is a trivial operation and never
    /// panics. Without `std`, this method is simply unavailable.
    #[cfg(feature="std")]
    #[inline]
    pub fn inner(&self) -> &(dyn std::error::Error + Send + Sync + 'static) {
        &*self.inner
    }

    /// Unwrap the inner error (`std` only)
    ///
    /// When configured with `std`, this is a trivial operation and never
    /// panics. Without `std`, this method is simply unavailable.
    #[cfg(feature="std")]
    #[inline]
    pub fn take_inner(self) -> Box<dyn std::error::Error + Send + Sync + 'static> {
        self.inner
    }

    /// Codes below this point represent OS Errors (i.e. positive i32 values).
    /// Codes at or above this point, but below [`Error::CUSTOM_START`] are
    /// reserved for use by the `rand` and `getrandom` crates.
    pub const INTERNAL_START: u32 = 1 << 31;

    /// Codes at or above this point can be used by users to define their own
    /// custom errors.
    pub const CUSTOM_START: u32 = (1 << 31) + (1 << 30);

    /// Extract the raw OS error code (if this error came from the OS)
    ///
    /// This method is identical to `std::io::Error::raw_os_error()`, except
    /// that it works in `no_std` contexts. If this method returns `None`, the
    /// error value can still be formatted via the `Diplay` implementation.
    #[inline]
    pub fn raw_os_error(&self) -> Option<i32> {
        #[cfg(feature="std")] {
            if let Some(e) = self.inner.downcast_ref::<std::io::Error>() {
                return e.raw_os_error();
            }
        }
        match self.code() {
            Some(code) if u32::from(code) < Self::INTERNAL_START =>
                Some(u32::from(code) as i32),
            _ => None,
        }
    }

    /// Retrieve the error code, if any.
    ///
    /// If this `Error` was constructed via `From<NonZeroU32>`, then this method
    /// will return this `NonZeroU32` code (for `no_std` this is always the
    /// case). Otherwise, this method will return `None`.
    #[inline]
    pub fn code(&self) -> Option<NonZeroU32> {
        #[cfg(feature="std")] {
            self.inner.downcast_ref::<ErrorCode>().map(|c| c.0)
        }
        #[cfg(not(feature="std"))] {
            Some(self.code)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        #[cfg(feature="std")] {
            write!(f, "Error {{ inner: {:?} }}", self.inner)
        }
        #[cfg(all(feature="getrandom", not(feature="std")))] {
            getrandom::Error::from(self.code).fmt(f)
        }
        #[cfg(not(feature="getrandom"))] {
            write!(f, "Error {{ code: {} }}", self.code)
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        #[cfg(feature="std")] {
            write!(f, "{}", self.inner)
        }
        #[cfg(all(feature="getrandom", not(feature="std")))] {
            getrandom::Error::from(self.code).fmt(f)
        }
        #[cfg(not(feature="getrandom"))] {
            write!(f, "error code {}", self.code)
        }
    }
}

impl From<NonZeroU32> for Error {
    #[inline]
    fn from(code: NonZeroU32) -> Self {
        #[cfg(feature="std")] {
            Error { inner: Box::new(ErrorCode(code)) }
        }
        #[cfg(not(feature="std"))] {
            Error { code }
        }
    }
}

#[cfg(feature="getrandom")]
impl From<getrandom::Error> for Error {
    #[inline]
    fn from(error: getrandom::Error) -> Self {
        #[cfg(feature="std")] {
            Error { inner: Box::new(error) }
        }
        #[cfg(not(feature="std"))] {
            Error { code: error.code() }
        }
    }
}

#[cfg(feature="std")]
impl std::error::Error for Error {
    #[inline]
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.inner.source()
    }
}

#[cfg(feature="std")]
impl From<Error> for std::io::Error {
    #[inline]
    fn from(error: Error) -> Self {
        std::io::Error::new(std::io::ErrorKind::Other, error)
    }
}

#[cfg(feature="std")]
#[derive(Debug, Copy, Clone)]
struct ErrorCode(NonZeroU32);

#[cfg(feature="std")]
impl fmt::Display for ErrorCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "error code {}", self.0)
    }
}

#[cfg(feature="std")]
impl std::error::Error for ErrorCode {}
<commit_msg>Reconstruct io::Error from raw code if possible in From impl<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Error types

use core::fmt;
use core::num::NonZeroU32;

/// Error type of random number generators
///
/// In order to be compatible with `std` and `no_std`, this type has two
/// possible implementations: with `std` a boxed `Error` trait object is stored,
/// while with `no_std` we merely store an error code.
pub struct Error {
    #[cfg(feature="std")]
    inner: Box<dyn std::error::Error + Send + Sync + 'static>,
    #[cfg(not(feature="std"))]
    code: NonZeroU32,
}

impl Error {
    /// Construct from any type supporting `std::error::Error`
    ///
    /// Available only when configured with `std`.
    ///
    /// See also `From<NonZeroU32>`, which is available with and without `std`.
    #[cfg(feature="std")]
    #[inline]
    pub fn new<E>(err: E) -> Self
    where E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>
    {
        Error { inner: err.into() }
    }

    /// Reference the inner error (`std` only)
    ///
    /// When configured with `std`, this is a trivial operation and never
    /// panics. Without `std`, this method is simply unavailable.
    #[cfg(feature="std")]
    #[inline]
    pub fn inner(&self) -> &(dyn std::error::Error + Send + Sync + 'static) {
        &*self.inner
    }

    /// Unwrap the inner error (`std` only)
    ///
    /// When configured with `std`, this is a trivial operation and never
    /// panics. Without `std`, this method is simply unavailable.
    #[cfg(feature="std")]
    #[inline]
    pub fn take_inner(self) -> Box<dyn std::error::Error + Send + Sync + 'static> {
        self.inner
    }

    /// Codes below this point represent OS Errors (i.e. positive i32 values).
    /// Codes at or above this point, but below [`Error::CUSTOM_START`] are
    /// reserved for use by the `rand` and `getrandom` crates.
    pub const INTERNAL_START: u32 = 1 << 31;

    /// Codes at or above this point can be used by users to define their own
    /// custom errors.
    pub const CUSTOM_START: u32 = (1 << 31) + (1 << 30);

    /// Extract the raw OS error code (if this error came from the OS)
    ///
    /// This method is identical to `std::io::Error::raw_os_error()`, except
    /// that it works in `no_std` contexts. If this method returns `None`, the
    /// error value can still be formatted via the `Diplay` implementation.
    #[inline]
    pub fn raw_os_error(&self) -> Option<i32> {
        #[cfg(feature="std")] {
            if let Some(e) = self.inner.downcast_ref::<std::io::Error>() {
                return e.raw_os_error();
            }
        }
        match self.code() {
            Some(code) if u32::from(code) < Self::INTERNAL_START =>
                Some(u32::from(code) as i32),
            _ => None,
        }
    }

    /// Retrieve the error code, if any.
    ///
    /// If this `Error` was constructed via `From<NonZeroU32>`, then this method
    /// will return this `NonZeroU32` code (for `no_std` this is always the
    /// case). Otherwise, this method will return `None`.
    #[inline]
    pub fn code(&self) -> Option<NonZeroU32> {
        #[cfg(feature="std")] {
            self.inner.downcast_ref::<ErrorCode>().map(|c| c.0)
        }
        #[cfg(not(feature="std"))] {
            Some(self.code)
        }
    }
}

impl fmt::Debug for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        #[cfg(feature="std")] {
            write!(f, "Error {{ inner: {:?} }}", self.inner)
        }
        #[cfg(all(feature="getrandom", not(feature="std")))] {
            getrandom::Error::from(self.code).fmt(f)
        }
        #[cfg(not(feature="getrandom"))] {
            write!(f, "Error {{ code: {} }}", self.code)
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        #[cfg(feature="std")] {
            write!(f, "{}", self.inner)
        }
        #[cfg(all(feature="getrandom", not(feature="std")))] {
            getrandom::Error::from(self.code).fmt(f)
        }
        #[cfg(not(feature="getrandom"))] {
            write!(f, "error code {}", self.code)
        }
    }
}

impl From<NonZeroU32> for Error {
    #[inline]
    fn from(code: NonZeroU32) -> Self {
        #[cfg(feature="std")] {
            Error { inner: Box::new(ErrorCode(code)) }
        }
        #[cfg(not(feature="std"))] {
            Error { code }
        }
    }
}

#[cfg(feature="getrandom")]
impl From<getrandom::Error> for Error {
    #[inline]
    fn from(error: getrandom::Error) -> Self {
        #[cfg(feature="std")] {
            Error { inner: Box::new(error) }
        }
        #[cfg(not(feature="std"))] {
            Error { code: error.code() }
        }
    }
}

#[cfg(feature="std")]
impl std::error::Error for Error {
    #[inline]
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.inner.source()
    }
}

#[cfg(feature="std")]
impl From<Error> for std::io::Error {
    #[inline]
    fn from(error: Error) -> Self {
        if let Some(code) = error.raw_os_error() {
            std::io::Error::from_raw_os_error(code)
        } else {
            std::io::Error::new(std::io::ErrorKind::Other, error)
        }
    }
}

#[cfg(feature="std")]
#[derive(Debug, Copy, Clone)]
struct ErrorCode(NonZeroU32);

#[cfg(feature="std")]
impl fmt::Display for ErrorCode {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "error code {}", self.0)
    }
}

#[cfg(feature="std")]
impl std::error::Error for ErrorCode {}
<|endoftext|>