<commit_before>
// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! LookupHost result from a resolution with a Resolver

use std::io;
use std::mem;
use std::net::IpAddr;
use std::slice::Iter;

use futures::{Async, future, Future, Poll, task};

use trust_dns::client::ClientHandle;
use trust_dns::error::ClientError;
use trust_dns::op::Message;
use trust_dns::rr::{DNSClass, Name, RData, RecordType};

use config::LookupIpStrategy;

/// Result of a DNS query when querying for A or AAAA records.
#[derive(Debug)]
pub struct LookupIp {
    ips: Vec<IpAddr>,
}

impl LookupIp {
    fn new(ips: Vec<IpAddr>) -> Self {
        LookupIp { ips }
    }

    /// Returns a borrowed iterator of the returned IPs
    pub fn iter(&self) -> LookupIpIter {
        LookupIpIter(self.ips.iter())
    }
}

impl Iterator for LookupIp {
    type Item = IpAddr;

    fn next(&mut self) -> Option<Self::Item> {
        self.ips.pop()
    }
}

/// Borrowed view of set of IPs returned from a LookupIp
pub struct LookupIpIter<'a>(Iter<'a, IpAddr>);

impl<'a> Iterator for LookupIpIter<'a> {
    type Item = &'a IpAddr;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

/// The Future returned from ResolverFuture when performing an A or AAAA lookup.
pub struct LookupIpFuture {
    future: Box<Future<Item = LookupIp, Error = io::Error>>,
}

impl LookupIpFuture {
    pub(crate) fn lookup<C: ClientHandle>(
        host: &str,
        strategy: LookupIpStrategy,
        client: &mut C,
    ) -> Self {
        let name = match Name::parse(host, None) {
            Ok(name) => name,
            Err(err) => {
                return LookupIpFuture { future: Box::new(future::err(io::Error::from(err))) }
            }
        };

        let query = LookupIpState::lookup(name, strategy, client);
        LookupIpFuture { future: Box::new(query) }
    }
}

impl Future for LookupIpFuture {
    type Item = LookupIp;
    type Error = io::Error;

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

// TODO: maximum recursion on CNAME, etc, chains...
// struct LookupStack(Vec<Query>);

// impl LookupStack {
//     // pushes the Query onto the stack, and returns a reference. An error will be returned
//     fn push(&mut self, query: Query) -> io::Result<&Query> {
//         if self.0.contains(&query) {
//             return Err(io::Error::new(io::ErrorKind::Other, "circular CNAME or other recursion"));
//         }

//         self.0.push(query);
//         Ok(self.0.last().unwrap())
//     }
// }

enum LookupIpState {
    Query(Box<Future<Item = Vec<IpAddr>, Error = ClientError>>),
    Fin(Vec<IpAddr>),
}

impl LookupIpState {
    pub fn lookup<C: ClientHandle>(name: Name, strategy: LookupIpStrategy, client: &mut C) -> Self {
        let query_future = lookup(name, strategy, client);

        LookupIpState::Query(query_future)
    }

    fn transition_query(&mut self, ips: Vec<IpAddr>) {
        assert!(if let LookupIpState::Query(_) = *self {
            true
        } else {
            false
        });

        // transition
        mem::replace(self, LookupIpState::Fin(ips));
    }
}

impl Future for LookupIpState {
    type Item = LookupIp;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // first transition any polling that is needed (mutable refs...)
        let poll;
        match *self {
            LookupIpState::Query(ref mut query) => {
                poll = query.poll().map_err(|e| e.into());
                match poll {
                    Ok(Async::NotReady) => {
                        return Ok(Async::NotReady);
                    }
                    Ok(Async::Ready(_)) => (), // handled in next match
                    Err(e) => {
                        return Err(e);
                    }
                }
            }
            LookupIpState::Fin(ref mut ips) => {
                let ips = mem::replace(ips, Vec::<IpAddr>::new());
                let ips = LookupIp::new(ips);
                return Ok(Async::Ready(ips));
            }
        }

        // getting here means there are Aync::Ready available.
        match *self {
            LookupIpState::Query(_) => {
                match poll {
                    Ok(Async::Ready(ips)) => self.transition_query(ips),
                    _ => panic!("should have returned earlier"),
                }
            }
            LookupIpState::Fin(_) => panic!("should have returned earlier"),
        }

        task::current().notify(); // yield
        Ok(Async::NotReady)
    }
}

fn lookup<C: ClientHandle>(
    name: Name,
    strategy: LookupIpStrategy,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    match strategy {
        LookupIpStrategy::Ipv4Only => ipv4_only(name, client),
        LookupIpStrategy::Ipv6Only => ipv6_only(name, client),
        LookupIpStrategy::Ipv4AndIpv6 => ipv4_and_ipv6(name, client),
        LookupIpStrategy::Ipv6thenIpv4 => unimplemented!(),
        LookupIpStrategy::Ipv4thenIpv6 => unimplemented!(),
    }
}

fn map_message_to_ipaddr(mut message: Message) -> Vec<IpAddr> {
    message
        .take_answers()
        .iter()
        .filter_map(|r| match *r.rdata() {
            RData::A(ipaddr) => Some(IpAddr::V4(ipaddr)),
            RData::AAAA(ipaddr) => Some(IpAddr::V6(ipaddr)),
            _ => None,
        })
        .collect()
}

/// queries only for A records
fn ipv4_only<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(client.query(name, DNSClass::IN, RecordType::A).map(
        map_message_to_ipaddr,
    ))
}

fn ipv6_only<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(client.query(name, DNSClass::IN, RecordType::AAAA).map(
        map_message_to_ipaddr,
    ))
}

fn ipv4_and_ipv6<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(
        client
            .query(name.clone(), DNSClass::IN, RecordType::A)
            .map(map_message_to_ipaddr)
            .select(client.query(name, DNSClass::IN, RecordType::AAAA).map(
                map_message_to_ipaddr,
            ))
            .then(|sel_res| {
                match sel_res {
                    // A record returned, get the other record result, or else just return record
                    Ok((mut ips, remaining_query)) => {
                        Box::new(remaining_query.then(move |query_res| match query_res {
                            /// join AAAA and A results
                            Ok(mut rem_ips) => {
                                rem_ips.append(&mut ips);
                                Box::new(future::ok(rem_ips))
                            }
                            // AAAA failed, just return A
                            Err(_) => Box::new(future::ok(ips)),
                        })) as
                            // This cast is to resolve a comilation error, not sure of it's necessity
                            Box<Future<Item = Vec<IpAddr>, Error = ClientError>>
                    }

                    // A failed, just return the AAAA result
                    Err((_, remaining_query)) => Box::new(remaining_query),
                }
            }),
    )
}

// TODO: build some non-network tests which test all variants of the Strategies...<commit_msg>remove a couple unnecessary Boxen<commit_after>
// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! LookupHost result from a resolution with a Resolver

use std::io;
use std::mem;
use std::net::IpAddr;
use std::slice::Iter;

use futures::{Async, future, Future, Poll, task};

use trust_dns::client::ClientHandle;
use trust_dns::error::ClientError;
use trust_dns::op::Message;
use trust_dns::rr::{DNSClass, Name, RData, RecordType};

use config::LookupIpStrategy;

/// Result of a DNS query when querying for A or AAAA records.
#[derive(Debug)]
pub struct LookupIp {
    ips: Vec<IpAddr>,
}

impl LookupIp {
    fn new(ips: Vec<IpAddr>) -> Self {
        LookupIp { ips }
    }

    /// Returns a borrowed iterator of the returned IPs
    pub fn iter(&self) -> LookupIpIter {
        LookupIpIter(self.ips.iter())
    }
}

impl Iterator for LookupIp {
    type Item = IpAddr;

    fn next(&mut self) -> Option<Self::Item> {
        self.ips.pop()
    }
}

/// Borrowed view of set of IPs returned from a LookupIp
pub struct LookupIpIter<'a>(Iter<'a, IpAddr>);

impl<'a> Iterator for LookupIpIter<'a> {
    type Item = &'a IpAddr;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

/// The Future returned from ResolverFuture when performing an A or AAAA lookup.
pub struct LookupIpFuture {
    future: Box<Future<Item = LookupIp, Error = io::Error>>,
}

impl LookupIpFuture {
    pub(crate) fn lookup<C: ClientHandle>(
        host: &str,
        strategy: LookupIpStrategy,
        client: &mut C,
    ) -> Self {
        let name = match Name::parse(host, None) {
            Ok(name) => name,
            Err(err) => {
                return LookupIpFuture { future: Box::new(future::err(io::Error::from(err))) }
            }
        };

        let query = LookupIpState::lookup(name, strategy, client);
        LookupIpFuture { future: Box::new(query) }
    }
}

impl Future for LookupIpFuture {
    type Item = LookupIp;
    type Error = io::Error;

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

// TODO: maximum recursion on CNAME, etc, chains...
// struct LookupStack(Vec<Query>);

// impl LookupStack {
//     // pushes the Query onto the stack, and returns a reference. An error will be returned
//     fn push(&mut self, query: Query) -> io::Result<&Query> {
//         if self.0.contains(&query) {
//             return Err(io::Error::new(io::ErrorKind::Other, "circular CNAME or other recursion"));
//         }

//         self.0.push(query);
//         Ok(self.0.last().unwrap())
//     }
// }

enum LookupIpState {
    Query(Box<Future<Item = Vec<IpAddr>, Error = ClientError>>),
    Fin(Vec<IpAddr>),
}

impl LookupIpState {
    pub fn lookup<C: ClientHandle>(name: Name, strategy: LookupIpStrategy, client: &mut C) -> Self {
        let query_future = lookup(name, strategy, client);

        LookupIpState::Query(query_future)
    }

    fn transition_query(&mut self, ips: Vec<IpAddr>) {
        assert!(if let LookupIpState::Query(_) = *self {
            true
        } else {
            false
        });

        // transition
        mem::replace(self, LookupIpState::Fin(ips));
    }
}

impl Future for LookupIpState {
    type Item = LookupIp;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // first transition any polling that is needed (mutable refs...)
        let poll;
        match *self {
            LookupIpState::Query(ref mut query) => {
                poll = query.poll().map_err(|e| e.into());
                match poll {
                    Ok(Async::NotReady) => {
                        return Ok(Async::NotReady);
                    }
                    Ok(Async::Ready(_)) => (), // handled in next match
                    Err(e) => {
                        return Err(e);
                    }
                }
            }
            LookupIpState::Fin(ref mut ips) => {
                let ips = mem::replace(ips, Vec::<IpAddr>::new());
                let ips = LookupIp::new(ips);
                return Ok(Async::Ready(ips));
            }
        }

        // getting here means there are Aync::Ready available.
        match *self {
            LookupIpState::Query(_) => {
                match poll {
                    Ok(Async::Ready(ips)) => self.transition_query(ips),
                    _ => panic!("should have returned earlier"),
                }
            }
            LookupIpState::Fin(_) => panic!("should have returned earlier"),
        }

        task::current().notify(); // yield
        Ok(Async::NotReady)
    }
}

fn lookup<C: ClientHandle>(
    name: Name,
    strategy: LookupIpStrategy,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    match strategy {
        LookupIpStrategy::Ipv4Only => ipv4_only(name, client),
        LookupIpStrategy::Ipv6Only => ipv6_only(name, client),
        LookupIpStrategy::Ipv4AndIpv6 => ipv4_and_ipv6(name, client),
        LookupIpStrategy::Ipv6thenIpv4 => unimplemented!(),
        LookupIpStrategy::Ipv4thenIpv6 => unimplemented!(),
    }
}

fn map_message_to_ipaddr(mut message: Message) -> Vec<IpAddr> {
    message
        .take_answers()
        .iter()
        .filter_map(|r| match *r.rdata() {
            RData::A(ipaddr) => Some(IpAddr::V4(ipaddr)),
            RData::AAAA(ipaddr) => Some(IpAddr::V6(ipaddr)),
            _ => None,
        })
        .collect()
}

/// queries only for A records
fn ipv4_only<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(client.query(name, DNSClass::IN, RecordType::A).map(
        map_message_to_ipaddr,
    ))
}

fn ipv6_only<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(client.query(name, DNSClass::IN, RecordType::AAAA).map(
        map_message_to_ipaddr,
    ))
}

fn ipv4_and_ipv6<C: ClientHandle>(
    name: Name,
    client: &mut C,
) -> Box<Future<Item = Vec<IpAddr>, Error = ClientError>> {
    Box::new(
        client
            .query(name.clone(), DNSClass::IN, RecordType::A)
            .map(map_message_to_ipaddr)
            .select(client.query(name, DNSClass::IN, RecordType::AAAA).map(
                map_message_to_ipaddr,
            ))
            .then(|sel_res| {
                match sel_res {
                    // A record returned, get the other record result, or else just return record
                    Ok((mut ips, remaining_query)) => {
                        Box::new(remaining_query.then(move |query_res| match query_res {
                            /// join AAAA and A results
                            Ok(mut rem_ips) => {
                                rem_ips.append(&mut ips);
                                future::ok(rem_ips)
                            }
                            // AAAA failed, just return A
                            Err(_) => future::ok(ips),
                        })) as
                            // This cast is to resolve a comilation error, not sure of it's necessity
                            Box<Future<Item = Vec<IpAddr>, Error = ClientError>>
                    }

                    // A failed, just return the AAAA result
                    Err((_, remaining_query)) => Box::new(remaining_query),
                }
            }),
    )
}

// TODO: build some non-network tests which test all variants of the Strategies...<|endoftext|>


<commit_before>#![crate_name = "twobit"]
#![crate_type = "dylib"]
#![feature(macro_rules)]

//! Implements the TwoBit struct giving read access to 2bit files in the format
//! used by the UCSC Genome Browser (details at: http://genome.ucsc.edu/FAQ/FAQformat.html#format7).

extern crate rustrt;

use std::os::MemoryMap;
use std::collections::HashMap;
use std::io::{ IoResult, IoError };
use std::iter::FromIterator;
use std::os::unix::AsRawFd;

#[deriving(Show)]
struct Block { start: u32, length: u32 }

struct Sequence {
	n_dna_bases: u32,
	unk_blocks: Vec<Block>,
	#[allow(dead_code)]
	mask_blocks: Vec<Block>,
	dna_start: *mut u8
}

impl Sequence {
	fn range<'a>(&'a self, start: u32, end: u32) -> SeqRange<'a> {
		let mut rsize = (end - start + 1) as uint;
		let mut end = end;

		if end >= self.n_dna_bases {
			let rest = rsize - self.n_dna_bases as uint;
			end = self.n_dna_bases - 1;
			rsize = (end - start + 1) as uint;

			// seq range chained with an iterator that repeats the same value
			SeqRange {
				rsize: rsize,
				ptr: unsafe { self.dna_start.offset( (start / 4) as int) },
				idx: 0u,
				offset: (start % 4) as uint,
				unk_blocks: &self.unk_blocks,
				ub_exhausted: self.unk_blocks.len() == 0,
				ub_idx: 0,
				ub_start: if self.unk_blocks.len() > 0 { self.unk_blocks[0].start as uint } else { 0 },
				ub_end: if self.unk_blocks.len() > 0 { (self.unk_blocks[0].start + self.unk_blocks[0].length - 1) as uint } else { 0 },
				n_more: rest
			}
		} else {
			SeqRange {
				rsize: rsize,
				ptr: unsafe { self.dna_start.offset( (start / 4) as int) },
				idx: 0u,
				offset: (start % 4) as uint,
				unk_blocks: &self.unk_blocks,
				ub_exhausted: self.unk_blocks.len() == 0,
				ub_idx: 0,
				ub_start: if self.unk_blocks.len() > 0 { self.unk_blocks[0].start as uint } else { 0 },
				ub_end: if self.unk_blocks.len() > 0 { (self.unk_blocks[0].start + self.unk_blocks[0].length - 1) as uint } else { 0 },
				n_more: 0
			}
		}
	}

	fn string(&self, start: u32, end: u32) -> String {
		FromIterator::from_iter(self.range(start, end))
	}
}

/// Sequence range iterator
pub struct SeqRange<'a> {
	rsize: uint,
	ptr: * mut u8,
	idx: uint,
	offset: uint,
	unk_blocks: &'a Vec<Block>,
	ub_exhausted: bool,
	ub_idx: uint,
	ub_start: uint,
	ub_end: uint,
	n_more: uint
}

impl<'a> SeqRange<'a> {
	#[inline]
	fn increment_idx(&mut self) {
		self.idx = self.idx + 1;
		self.offset = self.offset + 1;
		if self.offset == 4 {
			self.offset = 0;
			self.ptr = unsafe { self.ptr.offset(1) };
		}
	}
}

impl<'a> Iterator<char> for SeqRange<'a> {
	fn next(&mut self) -> Option<char> {
		if self.idx == self.rsize {
			if self.n_more > 0 {
				self.n_more = self.n_more - 1;
				return Some('N');
			}
			return None;
		} else {
			unsafe {
				// are we within a block?
				if !self.ub_exhausted {
					loop {
						if self.idx > self.ub_end {
							self.ub_idx = self.ub_idx + 1;
							if self.ub_idx == self.unk_blocks.len() {
								self.ub_exhausted = true;
								break;
							} else {
								self.ub_start = self.unk_blocks[self.ub_idx].start as uint;
								self.ub_end = self.ub_start + self.unk_blocks[self.ub_idx].length as uint - 1;
							}
						} else if self.idx >= self.ub_start {
							self.increment_idx();

							return Some('N');
						} else {
							// outside block, so continue to data collection
							break;
						}
					}
				}

				// no, so collect data
				let result = Some(byte_to_base(*self.ptr, self.offset));

				self.increment_idx();

				return result;
			}
		}
	}

	fn size_hint(&self) -> (uint, Option<uint>) {
		// lower and upper bound (same)
		let bound = self.rsize - self.idx - self.n_more;
		(bound, Some(bound))
	}
}

/// TwoBit type, represents a 2bit object in the UCSC Genome Browser format
pub struct TwoBit {
	seqs: HashMap<String, Sequence>,

	#[allow(dead_code)]
	data: MemoryMap // this is needed to keep the memory map alive
}

macro_rules! try_rt(
    ($e:expr) => (match $e { Ok(e) => e, Err(rustrt::rtio::IoError{code: code, extra: _, detail: _}) => return Err(IoError::from_errno(code, true)) })
)

fn mmap_read_u32(ptr: * mut u8, offset: uint) -> u32 {
	return unsafe {
		let tmp : *const u32 = std::mem::transmute(ptr.offset(offset as int) as *const [u8, ..4]);
		*tmp };
}

fn mmap_read_u8(ptr: * mut u8, offset: uint) -> u8 {
	return unsafe {
		*ptr.offset(offset as int)
	};
}

fn read_blocks(data: * mut u8, offset: u32) -> (Vec<Block>, u32) {
	let len = mmap_read_u32(data, offset as uint);
	let mut vec = Vec::<Block>::new();

	if len > 0 {
		let off1 = offset + 4;
		let off2 = offset + 4 + len*4;

		for i in range(0, len) {
			let start = mmap_read_u32(data, (off1 + i*4) as uint);
			let size = mmap_read_u32(data, (off2 + i*4) as uint);

			vec.push(Block{ start: start, length: size });
		}
	}

	return (vec, offset + 4 + 2*len*4);
}

fn mmap_read_index(data: *mut u8, count: u32) -> HashMap<String, Sequence> {
		let mut index = HashMap::with_capacity(count as uint);

		let mut header_start = 16u;

		for _ in range(0, count) {
			let name_size = mmap_read_u8(data, header_start);
			let name = unsafe { std::string::String::from_raw_parts(data.offset((header_start + 1) as int), name_size as uint, name_size as uint) };
			let offset = mmap_read_u32(data, header_start + 1 + name_size as uint);

			// get actual info
			let dna_size = mmap_read_u32(data, offset as uint);

			// unknown blocks
			let (unk_blocks, offset) = read_blocks(data, offset + 4);

			// masked blocks
			let (mask_blocks, offset) = read_blocks(data, offset);

			// actual pointer to DNA data
			let dna_ptr = unsafe { data.offset((offset + 4) as int) }; // + reserved

			index.insert(name, Sequence { n_dna_bases: dna_size, unk_blocks: unk_blocks, mask_blocks: mask_blocks, dna_start: dna_ptr });

			header_start = header_start + 1 + name_size as uint + 4;
		}

		return index;
}

fn byte_to_base(value: u8, offset: uint) -> char {
	let bases = ['T', 'C', 'A', 'G'];
	let rev_offset = 3 - offset;
	let mask = 3 << (rev_offset * 2);
	let idx = (value & mask) >> (rev_offset * 2);

	return bases[idx as uint];
}

impl TwoBit {

	/// Create a new TwoBit object from the supplied filename
	///
	/// # Arguments
	///
	/// - filename - string slice with the path to the 2bit file
	pub fn new(filename: &str) -> IoResult<TwoBit> {
		// open file
		//let mut fh = try_rt!(native::io::file::open(&filename.to_c_str(), rustrt::rtio::Open, rustrt::rtio::Read));
		let fh = try!(std::io::File::open(&Path::new(filename)));
		let fs = try!(fh.stat());

		// build memory map
		let mmap = match MemoryMap::new(fs.size as uint, &[ std::os::MapReadable, std::os::MapFd(fh.as_raw_fd())]) {
			Ok(val) => val,
			Err(_) => return Err(IoError{kind: std::io::OtherIoError, desc: "Memory map failed!", detail: None})
		};

		// TODO: add madvise call

		// validate header
		let val = mmap_read_u32(mmap.data(), 0);

		if val != 0x1A412743 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Invalid signature or wrong architecture.", detail: None });
		}

		let val = mmap_read_u32(mmap.data(), 4);
		if val != 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Unknown file version.", detail: None });
		} // TODO: actually report version found

		let n_sequences = mmap_read_u32(mmap.data(), 8);
		if n_sequences == 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Zero sequence count.", detail: None });
		}

		let val = mmap_read_u32(mmap.data(), 12);
		if val != 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Reserved bytes not zero.", detail: None });
		} // TODO: actually report value found

		// parse index
		let index = mmap_read_index(mmap.data(), n_sequences);

		return Ok(TwoBit { seqs: index, data: mmap });
	}

	/// Collect the sequence for the range [start, end]
	///
	/// Note that when ranges exceed the recorded information contained in the 2bit file
	/// excess is padded with 'N's.
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	/// - start - zero based start coordinate for range
	/// - end - zero based end coordinate (inclusive) for range
	pub fn sequence(&self, chrom: &str, start: u32, end: u32) -> Option<String> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => Some(seq.string(start, end)),
			None => None
		}
	}

	/// Iterator on the sequence for the range [start, end]
	///
	/// Note that when ranges exceed the recorded information contained in the 2bit file
	/// excess is padded with 'N's.
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	/// - start - zero based start coordinate for range
	/// - end - zero based end coordinate (inclusive) for range
	pub fn sequence_iter<'a>(&'a self, chrom: &str, start: u32, end: u32) -> Option<SeqRange<'a>> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => Some(seq.range(start, end)),
			None => None
		}
	}

	/// Retrieve the length (number of bases) of a given sequence
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	pub fn sequence_len(&self, chrom: &str) -> Option<u32> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(&Sequence{ n_dna_bases: n, ..}) => Some(n),
			None => None
		}
	}

	/// Retrieve the names of the sequences contained in this file
	pub fn sequence_names<'a>(&'a self) -> Vec<&'a String> {
		self.seqs.keys().collect()
	}

	/// Compute the nucleotide frequencies (ACGT) of a given sequence
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	pub fn base_frequencies(&self, chrom: &str) -> Option<[f64, ..4]> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => {
				let mut counts = [0f64, 0.0, 0.0, 0.0];

				for c in seq.range(0, seq.n_dna_bases - 1) {
					match c {
						'A' => counts[0] = counts[0] + 1.0,
						'C' => counts[1] = counts[1] + 1.0,
						'G' => counts[2] = counts[2] + 1.0,
						'T' => counts[3] = counts[3] + 1.0,
						_ => {}
					}
				}

				let sum = counts[0] + counts[1] + counts[2] + counts[3];
				counts[0] = counts[0] / sum;
				counts[1] = counts[1] / sum;
				counts[2] = counts[2] / sum;
				counts[3] = counts[3] / sum;

				Some(counts)
			}
			None => None
		}
	}
}
<commit_msg>Useful DNA sequence manipulation operations<commit_after>#![crate_name = "twobit"]
#![crate_type = "dylib"]
#![feature(macro_rules)]

//! Implements the TwoBit struct giving read access to 2bit files in the format
//! used by the UCSC Genome Browser (details at: http://genome.ucsc.edu/FAQ/FAQformat.html#format7).

extern crate rustrt;

use std::os::MemoryMap;
use std::collections::HashMap;
use std::io::{ IoResult, IoError };
use std::iter::FromIterator;
use std::os::unix::AsRawFd;

#[deriving(Show)]
struct Block { start: u32, length: u32 }

struct Sequence {
	n_dna_bases: u32,
	unk_blocks: Vec<Block>,
	#[allow(dead_code)]
	mask_blocks: Vec<Block>,
	dna_start: *mut u8
}

impl Sequence {
	fn range<'a>(&'a self, start: u32, end: u32) -> SeqRange<'a> {
		let mut rsize = (end - start + 1) as uint;
		let mut end = end;

		if end >= self.n_dna_bases {
			let rest = rsize - self.n_dna_bases as uint;
			end = self.n_dna_bases - 1;
			rsize = (end - start + 1) as uint;

			// seq range chained with an iterator that repeats the same value
			SeqRange {
				rsize: rsize,
				ptr: unsafe { self.dna_start.offset( (start / 4) as int) },
				idx: 0u,
				offset: (start % 4) as uint,
				unk_blocks: &self.unk_blocks,
				ub_exhausted: self.unk_blocks.len() == 0,
				ub_idx: 0,
				ub_start: if self.unk_blocks.len() > 0 { self.unk_blocks[0].start as uint } else { 0 },
				ub_end: if self.unk_blocks.len() > 0 { (self.unk_blocks[0].start + self.unk_blocks[0].length - 1) as uint } else { 0 },
				n_more: rest
			}
		} else {
			SeqRange {
				rsize: rsize,
				ptr: unsafe { self.dna_start.offset( (start / 4) as int) },
				idx: 0u,
				offset: (start % 4) as uint,
				unk_blocks: &self.unk_blocks,
				ub_exhausted: self.unk_blocks.len() == 0,
				ub_idx: 0,
				ub_start: if self.unk_blocks.len() > 0 { self.unk_blocks[0].start as uint } else { 0 },
				ub_end: if self.unk_blocks.len() > 0 { (self.unk_blocks[0].start + self.unk_blocks[0].length - 1) as uint } else { 0 },
				n_more: 0
			}
		}
	}

	fn string(&self, start: u32, end: u32) -> String {
		FromIterator::from_iter(self.range(start, end))
	}
}

/// Sequence range iterator
pub struct SeqRange<'a> {
	rsize: uint,
	ptr: * mut u8,
	idx: uint,
	offset: uint,
	unk_blocks: &'a Vec<Block>,
	ub_exhausted: bool,
	ub_idx: uint,
	ub_start: uint,
	ub_end: uint,
	n_more: uint
}

impl<'a> SeqRange<'a> {
	#[inline]
	fn increment_idx(&mut self) {
		self.idx = self.idx + 1;
		self.offset = self.offset + 1;
		if self.offset == 4 {
			self.offset = 0;
			self.ptr = unsafe { self.ptr.offset(1) };
		}
	}
}

impl<'a> Iterator<char> for SeqRange<'a> {
	fn next(&mut self) -> Option<char> {
		if self.idx == self.rsize {
			if self.n_more > 0 {
				self.n_more = self.n_more - 1;
				return Some('N');
			}
			return None;
		} else {
			unsafe {
				// are we within a block?
				if !self.ub_exhausted {
					loop {
						if self.idx > self.ub_end {
							self.ub_idx = self.ub_idx + 1;
							if self.ub_idx == self.unk_blocks.len() {
								self.ub_exhausted = true;
								break;
							} else {
								self.ub_start = self.unk_blocks[self.ub_idx].start as uint;
								self.ub_end = self.ub_start + self.unk_blocks[self.ub_idx].length as uint - 1;
							}
						} else if self.idx >= self.ub_start {
							self.increment_idx();

							return Some('N');
						} else {
							// outside block, so continue to data collection
							break;
						}
					}
				}

				// no, so collect data
				let result = Some(byte_to_base(*self.ptr, self.offset));

				self.increment_idx();

				return result;
			}
		}
	}

	fn size_hint(&self) -> (uint, Option<uint>) {
		// lower and upper bound (same)
		let bound = self.rsize - self.idx - self.n_more;
		(bound, Some(bound))
	}
}

/// TwoBit type, represents a 2bit object in the UCSC Genome Browser format
pub struct TwoBit {
	seqs: HashMap<String, Sequence>,

	#[allow(dead_code)]
	data: MemoryMap // this is needed to keep the memory map alive
}

macro_rules! try_rt(
    ($e:expr) => (match $e { Ok(e) => e, Err(rustrt::rtio::IoError{code: code, extra: _, detail: _}) => return Err(IoError::from_errno(code, true)) })
)

fn mmap_read_u32(ptr: * mut u8, offset: uint) -> u32 {
	return unsafe {
		let tmp : *const u32 = std::mem::transmute(ptr.offset(offset as int) as *const [u8, ..4]);
		*tmp };
}

fn mmap_read_u8(ptr: * mut u8, offset: uint) -> u8 {
	return unsafe {
		*ptr.offset(offset as int)
	};
}

fn read_blocks(data: * mut u8, offset: u32) -> (Vec<Block>, u32) {
	let len = mmap_read_u32(data, offset as uint);
	let mut vec = Vec::<Block>::new();

	if len > 0 {
		let off1 = offset + 4;
		let off2 = offset + 4 + len*4;

		for i in range(0, len) {
			let start = mmap_read_u32(data, (off1 + i*4) as uint);
			let size = mmap_read_u32(data, (off2 + i*4) as uint);

			vec.push(Block{ start: start, length: size });
		}
	}

	return (vec, offset + 4 + 2*len*4);
}

fn mmap_read_index(data: *mut u8, count: u32) -> HashMap<String, Sequence> {
		let mut index = HashMap::with_capacity(count as uint);

		let mut header_start = 16u;

		for _ in range(0, count) {
			let name_size = mmap_read_u8(data, header_start);
			let name = unsafe { std::string::String::from_raw_parts(data.offset((header_start + 1) as int), name_size as uint, name_size as uint) };
			let offset = mmap_read_u32(data, header_start + 1 + name_size as uint);

			// get actual info
			let dna_size = mmap_read_u32(data, offset as uint);

			// unknown blocks
			let (unk_blocks, offset) = read_blocks(data, offset + 4);

			// masked blocks
			let (mask_blocks, offset) = read_blocks(data, offset);

			// actual pointer to DNA data
			let dna_ptr = unsafe { data.offset((offset + 4) as int) }; // + reserved

			index.insert(name, Sequence { n_dna_bases: dna_size, unk_blocks: unk_blocks, mask_blocks: mask_blocks, dna_start: dna_ptr });

			header_start = header_start + 1 + name_size as uint + 4;
		}

		return index;
}

fn byte_to_base(value: u8, offset: uint) -> char {
	let bases = ['T', 'C', 'A', 'G'];
	let rev_offset = 3 - offset;
	let mask = 3 << (rev_offset * 2);
	let idx = (value & mask) >> (rev_offset * 2);

	return bases[idx as uint];
}

impl TwoBit {

	/// Create a new TwoBit object from the supplied filename
	///
	/// # Arguments
	///
	/// - filename - string slice with the path to the 2bit file
	pub fn new(filename: &str) -> IoResult<TwoBit> {
		// TODO: revise interface to take a "File" instance instead of a filename

		// open file
		let fh = try!(std::io::File::open(&Path::new(filename)));
		let fs = try!(fh.stat());

		// build memory map
		let mmap = match MemoryMap::new(fs.size as uint, &[ std::os::MapReadable, std::os::MapFd(fh.as_raw_fd())]) {
			Ok(val) => val,
			Err(_) => return Err(IoError{kind: std::io::OtherIoError, desc: "Memory map failed!", detail: None})
		};

		// TODO: add madvise call

		// validate header
		let val = mmap_read_u32(mmap.data(), 0);

		if val != 0x1A412743 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Invalid signature or wrong architecture.", detail: None });
		}

		let val = mmap_read_u32(mmap.data(), 4);
		if val != 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Unknown file version.", detail: None });
		} // TODO: actually report version found

		let n_sequences = mmap_read_u32(mmap.data(), 8);
		if n_sequences == 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Zero sequence count.", detail: None });
		}

		let val = mmap_read_u32(mmap.data(), 12);
		if val != 0 {
			return Err(IoError { kind: std::io::OtherIoError, desc: "Reserved bytes not zero.", detail: None });
		} // TODO: actually report value found

		// parse index
		let index = mmap_read_index(mmap.data(), n_sequences);

		return Ok(TwoBit { seqs: index, data: mmap });
	}

	/// Collect the sequence for the range [start, end]
	///
	/// Note that when ranges exceed the recorded information contained in the 2bit file
	/// excess is padded with 'N's.
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	/// - start - zero based start coordinate for range
	/// - end - zero based end coordinate (inclusive) for range
	pub fn sequence(&self, chrom: &str, start: u32, end: u32) -> Option<String> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => Some(seq.string(start, end)),
			None => None
		}
	}

	/// Iterator on the sequence for the range [start, end]
	///
	/// Note that when ranges exceed the recorded information contained in the 2bit file
	/// excess is padded with 'N's.
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	/// - start - zero based start coordinate for range
	/// - end - zero based end coordinate (inclusive) for range
	pub fn sequence_iter<'a>(&'a self, chrom: &str, start: u32, end: u32) -> Option<SeqRange<'a>> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => Some(seq.range(start, end)),
			None => None
		}
	}

	/// Retrieve the length (number of bases) of a given sequence
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	pub fn sequence_len(&self, chrom: &str) -> Option<u32> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(&Sequence{ n_dna_bases: n, ..}) => Some(n),
			None => None
		}
	}

	/// Retrieve the names of the sequences contained in this file
	pub fn sequence_names<'a>(&'a self) -> Vec<&'a String> {
		self.seqs.keys().collect()
	}

	/// Compute the nucleotide frequencies (ACGT) of a given sequence
	///
	/// # Arguments
	///
	/// - chrom - sequence name, typically the chromosome name
	pub fn base_frequencies(&self, chrom: &str) -> Option<[f64, ..4]> {
		match self.seqs.get(&String::from_str(chrom)) {
			Some(ref seq) => {
				let mut counts = [0f64, 0.0, 0.0, 0.0];

				for c in seq.range(0, seq.n_dna_bases - 1) {
					match c {
						'A' => counts[0] = counts[0] + 1.0,
						'C' => counts[1] = counts[1] + 1.0,
						'G' => counts[2] = counts[2] + 1.0,
						'T' => counts[3] = counts[3] + 1.0,
						_ => {}
					}
				}

				let sum = counts[0] + counts[1] + counts[2] + counts[3];
				counts[0] = counts[0] / sum;
				counts[1] = counts[1] / sum;
				counts[2] = counts[2] / sum;
				counts[3] = counts[3] / sum;

				Some(counts)
			}
			None => None
		}
	}
}

/// Set of useful nucleotide sequence manipulation operations.
pub trait DNAOps {

	/// Obtain reverse complement of nucleotide sequence
	fn reverse_complement(&self) -> Self;

	/// Convert a DNA string into an integer vector
	///
	/// # Arguments
	///
	/// - offset - starting value for DNA nucleotides (A = offset, C = offset + 1, G = offset + 2, T = offset + 3, else offset + 4)
	fn into_numeric(&self, offset: u8) -> Vec<u8>;
}

impl DNAOps for String {

	fn reverse_complement(&self) -> String {
		// create a new string by iterating over the chars of the old one
		let mut result = String::with_capacity(self.len());

		for x in self.chars().map(|base|
		  match base {
		  	'a' => 't',
		  	'A' => 'T',
		  	'c' => 'g',
		  	'C' => 'G',
		  	'g' => 'c',
		  	'G' => 'C',
		  	't' => 'a',
		  	'T' => 'A',
		  	_ => 'N'
		  }) {
		  result.push(x);
		}

		// use as_mut_vec to reverse the resulting string inplace
		unsafe {
			let vec = result.as_mut_vec();
			vec.reverse();
		}
		return result;
	}

	fn into_numeric(&self, offset: u8) -> Vec<u8> {
		self.chars().map(|base: char|
		 match base {
		 	'a' => offset,
		 	'A' => offset,
		 	'c' => offset + 1,
		 	'C' => offset + 1,
		 	'g' => offset + 2,
		 	'G' => offset + 2,
		 	't' => offset + 3,
		 	'T' => offset + 3,
		 	_ => offset + 4
		 }).collect::<Vec<u8>>()
	}
}
<|endoftext|>


<commit_before>//! Michael-Scott lock-free queue.
//!
//! Usable with any number of producers and consumers.
//!
//! Michael and Scott.  Simple, Fast, and Practical Non-Blocking and Blocking Concurrent Queue
//! Algorithms.  PODC 1996.  http://dl.acm.org/citation.cfm?id=248106
//!
//! Simon Doherty, Lindsay Groves, Victor Luchangco, and Mark Moir. 2004b. Formal Verification of a
//! Practical Lock-Free Queue Algorithm. https://doi.org/10.1007/978-3-540-30232-2_7

use core::mem::MaybeUninit;
use core::sync::atomic::Ordering::{Acquire, Relaxed, Release};

use crossbeam_utils::CachePadded;

use {unprotected, Atomic, Guard, Owned, Shared};

// The representation here is a singly-linked list, with a sentinel node at the front. In general
// the `tail` pointer may lag behind the actual tail. Non-sentinel nodes are either all `Data` or
// all `Blocked` (requests for data from blocked threads).
#[derive(Debug)]
pub struct Queue<T> {
    head: CachePadded<Atomic<Node<T>>>,
    tail: CachePadded<Atomic<Node<T>>>,
}

#[derive(Debug)]
struct Node<T> {
    /// The slot in which a value of type `T` can be stored.
    ///
    /// The type of `data` is `MaybeUninit<T>` because a `Node<T>` doesn't always contain a `T`.
    /// For example, the sentinel node in a queue never contains a value: its slot is always empty.
    /// Other nodes start their life with a push operation and contain a value until it gets popped
    /// out. After that such empty nodes get added to the collector for destruction.
    data: MaybeUninit<T>,

    next: Atomic<Node<T>>,
}

// Any particular `T` should never be accessed concurrently, so no need for `Sync`.
unsafe impl<T: Send> Sync for Queue<T> {}
unsafe impl<T: Send> Send for Queue<T> {}

impl<T> Queue<T> {
    /// Create a new, empty queue.
    pub fn new() -> Queue<T> {
        let q = Queue {
            head: CachePadded::new(Atomic::null()),
            tail: CachePadded::new(Atomic::null()),
        };
        let sentinel = Owned::new(Node {
            data: MaybeUninit::uninit(),
            next: Atomic::null(),
        });
        unsafe {
            let guard = &unprotected();
            let sentinel = sentinel.into_shared(guard);
            q.head.store(sentinel, Relaxed);
            q.tail.store(sentinel, Relaxed);
            q
        }
    }

    /// Attempts to atomically place `n` into the `next` pointer of `onto`, and returns `true` on
    /// success. The queue's `tail` pointer may be updated.
    #[inline(always)]
    fn push_internal(&self, onto: Shared<Node<T>>, new: Shared<Node<T>>, guard: &Guard) -> bool {
        // is `onto` the actual tail?
        let o = unsafe { onto.deref() };
        let next = o.next.load(Acquire, guard);
        if unsafe { next.as_ref().is_some() } {
            // if not, try to "help" by moving the tail pointer forward
            let _ = self.tail.compare_and_set(onto, next, Release, guard);
            false
        } else {
            // looks like the actual tail; attempt to link in `n`
            let result = o
                .next
                .compare_and_set(Shared::null(), new, Release, guard)
                .is_ok();
            if result {
                // try to move the tail pointer forward
                let _ = self.tail.compare_and_set(onto, new, Release, guard);
            }
            result
        }
    }

    /// Adds `t` to the back of the queue, possibly waking up threads blocked on `pop`.
    pub fn push(&self, t: T, guard: &Guard) {
        let new = Owned::new(Node {
            data: MaybeUninit::new(t),
            next: Atomic::null(),
        });
        let new = Owned::into_shared(new, guard);

        loop {
            // We push onto the tail, so we'll start optimistically by looking there first.
            let tail = self.tail.load(Acquire, guard);

            // Attempt to push onto the `tail` snapshot; fails if `tail.next` has changed.
            if self.push_internal(tail, new, guard) {
                break;
            }
        }
    }

    /// Attempts to pop a data node. `Ok(None)` if queue is empty; `Err(())` if lost race to pop.
    #[inline(always)]
    fn pop_internal(&self, guard: &Guard) -> Result<Option<T>, ()> {
        let head = self.head.load(Acquire, guard);
        let h = unsafe { head.deref() };
        let next = h.next.load(Acquire, guard);
        match unsafe { next.as_ref() } {
            Some(n) => unsafe {
                self.head
                    .compare_and_set(head, next, Release, guard)
                    .map(|_| {
                        let tail = self.tail.load(Relaxed, guard);
                        // Advance the tail so that we don't retire a pointer to a reachable node.
                        if head == tail {
                            let _ = self.tail.compare_and_set(tail, next, Release, guard);
                        }
                        guard.defer_destroy(head);
                        // TODO: Replace with MaybeUninit::read when api is stable
                        Some(n.data.as_ptr().read())
                    })
                    .map_err(|_| ())
            },
            None => Ok(None),
        }
    }

    /// Attempts to pop a data node, if the data satisfies the given condition. `Ok(None)` if queue
    /// is empty or the data does not satisfy the condition; `Err(())` if lost race to pop.
    #[inline(always)]
    fn pop_if_internal<F>(&self, condition: F, guard: &Guard) -> Result<Option<T>, ()>
    where
        T: Sync,
        F: Fn(&T) -> bool,
    {
        let head = self.head.load(Acquire, guard);
        let h = unsafe { head.deref() };
        let next = h.next.load(Acquire, guard);
        unsafe {
            match next.as_ref() {
                Some(n) if condition(&*n.data.as_ptr()) => {
                    self.head
                        .compare_and_set(head, next, Release, guard)
                        .map(|_| {
                            let tail = self.tail.load(Relaxed, guard);
                            // Advance the tail so that we don't retire a pointer to a reachable node.
                            if head == tail {
                                let _ = self.tail.compare_and_set(tail, next, Release, guard);
                            }
                            guard.defer_destroy(head);
                            // TODO: Replace with MaybeUninit::read when api is stable
                            Some(n.data.as_ptr().read())
                        })
                        .map_err(|_| ())
                }
                None | Some(_) => Ok(None),
            }
        }
    }

    /// Attempts to dequeue from the front.
    ///
    /// Returns `None` if the queue is observed to be empty.
    pub fn try_pop(&self, guard: &Guard) -> Option<T> {
        loop {
            if let Ok(head) = self.pop_internal(guard) {
                return head;
            }
        }
    }

    /// Attempts to dequeue from the front, if the item satisfies the given condition.
    ///
    /// Returns `None` if the queue is observed to be empty, or the head does not satisfy the given
    /// condition.
    pub fn try_pop_if<F>(&self, condition: F, guard: &Guard) -> Option<T>
    where
        T: Sync,
        F: Fn(&T) -> bool,
    {
        loop {
            if let Ok(head) = self.pop_if_internal(&condition, guard) {
                return head;
            }
        }
    }
}

impl<T> Drop for Queue<T> {
    fn drop(&mut self) {
        unsafe {
            let guard = &unprotected();

            while let Some(_) = self.try_pop(guard) {}

            // Destroy the remaining sentinel node.
            let sentinel = self.head.load(Relaxed, guard);
            drop(sentinel.into_owned());
        }
    }
}

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

    struct Queue<T> {
        queue: super::Queue<T>,
    }

    impl<T> Queue<T> {
        pub fn new() -> Queue<T> {
            Queue {
                queue: super::Queue::new(),
            }
        }

        pub fn push(&self, t: T) {
            let guard = &pin();
            self.queue.push(t, guard);
        }

        pub fn is_empty(&self) -> bool {
            let guard = &pin();
            let head = self.queue.head.load(Acquire, guard);
            let h = unsafe { head.deref() };
            h.next.load(Acquire, guard).is_null()
        }

        pub fn try_pop(&self) -> Option<T> {
            let guard = &pin();
            self.queue.try_pop(guard)
        }

        pub fn pop(&self) -> T {
            loop {
                match self.try_pop() {
                    None => continue,
                    Some(t) => return t,
                }
            }
        }
    }

    const CONC_COUNT: i64 = 1000000;

    #[test]
    fn push_try_pop_1() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(37));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_2() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        q.push(48);
        assert_eq!(q.try_pop(), Some(37));
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(48));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_seq() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.try_pop(), Some(i));
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_1() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.pop(), 37);
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_2() {
        let q: Queue<i64> = Queue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), 37);
        assert_eq!(q.pop(), 48);
    }

    #[test]
    fn push_pop_many_seq() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.pop(), i);
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_spsc() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());

        thread::scope(|scope| {
            scope.spawn(|_| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.try_pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        })
        .unwrap();
    }

    #[test]
    fn push_try_pop_many_spmc() {
        fn recv(_t: i32, q: &Queue<i64>) {
            let mut cur = -1;
            for _i in 0..CONC_COUNT {
                if let Some(elem) = q.try_pop() {
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 {
                        break;
                    }
                }
            }
        }

        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        thread::scope(|scope| {
            for i in 0..3 {
                let q = &q;
                scope.spawn(move |_| recv(i, q));
            }

            scope.spawn(|_| {
                for i in 0..CONC_COUNT {
                    q.push(i);
                }
            });
        })
        .unwrap();
    }

    #[test]
    fn push_try_pop_many_mpmc() {
        enum LR {
            Left(i64),
            Right(i64),
        }

        let q: Queue<LR> = Queue::new();
        assert!(q.is_empty());

        thread::scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|_| {
                    for i in CONC_COUNT - 1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|_| {
                    for i in CONC_COUNT - 1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|_| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.try_pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        })
        .unwrap();
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: Queue<i64> = Queue::new();

        thread::scope(|scope| {
            scope.spawn(|_| {
                let mut next = 0;
                while next < CONC_COUNT {
                    assert_eq!(q.pop(), next);
                    next += 1;
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        })
        .unwrap();
        assert!(q.is_empty());
    }

    #[test]
    fn is_empty_dont_pop() {
        let q: Queue<i64> = Queue::new();
        q.push(20);
        q.push(20);
        assert!(!q.is_empty());
        assert!(!q.is_empty());
        assert!(q.try_pop().is_some());
    }
}
<commit_msg>epoch: remove Debug impl on Node<commit_after>//! Michael-Scott lock-free queue.
//!
//! Usable with any number of producers and consumers.
//!
//! Michael and Scott.  Simple, Fast, and Practical Non-Blocking and Blocking Concurrent Queue
//! Algorithms.  PODC 1996.  http://dl.acm.org/citation.cfm?id=248106
//!
//! Simon Doherty, Lindsay Groves, Victor Luchangco, and Mark Moir. 2004b. Formal Verification of a
//! Practical Lock-Free Queue Algorithm. https://doi.org/10.1007/978-3-540-30232-2_7

use core::mem::MaybeUninit;
use core::sync::atomic::Ordering::{Acquire, Relaxed, Release};

use crossbeam_utils::CachePadded;

use {unprotected, Atomic, Guard, Owned, Shared};

// The representation here is a singly-linked list, with a sentinel node at the front. In general
// the `tail` pointer may lag behind the actual tail. Non-sentinel nodes are either all `Data` or
// all `Blocked` (requests for data from blocked threads).
#[derive(Debug)]
pub struct Queue<T> {
    head: CachePadded<Atomic<Node<T>>>,
    tail: CachePadded<Atomic<Node<T>>>,
}

struct Node<T> {
    /// The slot in which a value of type `T` can be stored.
    ///
    /// The type of `data` is `MaybeUninit<T>` because a `Node<T>` doesn't always contain a `T`.
    /// For example, the sentinel node in a queue never contains a value: its slot is always empty.
    /// Other nodes start their life with a push operation and contain a value until it gets popped
    /// out. After that such empty nodes get added to the collector for destruction.
    data: MaybeUninit<T>,

    next: Atomic<Node<T>>,
}

// Any particular `T` should never be accessed concurrently, so no need for `Sync`.
unsafe impl<T: Send> Sync for Queue<T> {}
unsafe impl<T: Send> Send for Queue<T> {}

impl<T> Queue<T> {
    /// Create a new, empty queue.
    pub fn new() -> Queue<T> {
        let q = Queue {
            head: CachePadded::new(Atomic::null()),
            tail: CachePadded::new(Atomic::null()),
        };
        let sentinel = Owned::new(Node {
            data: MaybeUninit::uninit(),
            next: Atomic::null(),
        });
        unsafe {
            let guard = &unprotected();
            let sentinel = sentinel.into_shared(guard);
            q.head.store(sentinel, Relaxed);
            q.tail.store(sentinel, Relaxed);
            q
        }
    }

    /// Attempts to atomically place `n` into the `next` pointer of `onto`, and returns `true` on
    /// success. The queue's `tail` pointer may be updated.
    #[inline(always)]
    fn push_internal(&self, onto: Shared<Node<T>>, new: Shared<Node<T>>, guard: &Guard) -> bool {
        // is `onto` the actual tail?
        let o = unsafe { onto.deref() };
        let next = o.next.load(Acquire, guard);
        if unsafe { next.as_ref().is_some() } {
            // if not, try to "help" by moving the tail pointer forward
            let _ = self.tail.compare_and_set(onto, next, Release, guard);
            false
        } else {
            // looks like the actual tail; attempt to link in `n`
            let result = o
                .next
                .compare_and_set(Shared::null(), new, Release, guard)
                .is_ok();
            if result {
                // try to move the tail pointer forward
                let _ = self.tail.compare_and_set(onto, new, Release, guard);
            }
            result
        }
    }

    /// Adds `t` to the back of the queue, possibly waking up threads blocked on `pop`.
    pub fn push(&self, t: T, guard: &Guard) {
        let new = Owned::new(Node {
            data: MaybeUninit::new(t),
            next: Atomic::null(),
        });
        let new = Owned::into_shared(new, guard);

        loop {
            // We push onto the tail, so we'll start optimistically by looking there first.
            let tail = self.tail.load(Acquire, guard);

            // Attempt to push onto the `tail` snapshot; fails if `tail.next` has changed.
            if self.push_internal(tail, new, guard) {
                break;
            }
        }
    }

    /// Attempts to pop a data node. `Ok(None)` if queue is empty; `Err(())` if lost race to pop.
    #[inline(always)]
    fn pop_internal(&self, guard: &Guard) -> Result<Option<T>, ()> {
        let head = self.head.load(Acquire, guard);
        let h = unsafe { head.deref() };
        let next = h.next.load(Acquire, guard);
        match unsafe { next.as_ref() } {
            Some(n) => unsafe {
                self.head
                    .compare_and_set(head, next, Release, guard)
                    .map(|_| {
                        let tail = self.tail.load(Relaxed, guard);
                        // Advance the tail so that we don't retire a pointer to a reachable node.
                        if head == tail {
                            let _ = self.tail.compare_and_set(tail, next, Release, guard);
                        }
                        guard.defer_destroy(head);
                        // TODO: Replace with MaybeUninit::read when api is stable
                        Some(n.data.as_ptr().read())
                    })
                    .map_err(|_| ())
            },
            None => Ok(None),
        }
    }

    /// Attempts to pop a data node, if the data satisfies the given condition. `Ok(None)` if queue
    /// is empty or the data does not satisfy the condition; `Err(())` if lost race to pop.
    #[inline(always)]
    fn pop_if_internal<F>(&self, condition: F, guard: &Guard) -> Result<Option<T>, ()>
    where
        T: Sync,
        F: Fn(&T) -> bool,
    {
        let head = self.head.load(Acquire, guard);
        let h = unsafe { head.deref() };
        let next = h.next.load(Acquire, guard);
        unsafe {
            match next.as_ref() {
                Some(n) if condition(&*n.data.as_ptr()) => {
                    self.head
                        .compare_and_set(head, next, Release, guard)
                        .map(|_| {
                            let tail = self.tail.load(Relaxed, guard);
                            // Advance the tail so that we don't retire a pointer to a reachable node.
                            if head == tail {
                                let _ = self.tail.compare_and_set(tail, next, Release, guard);
                            }
                            guard.defer_destroy(head);
                            // TODO: Replace with MaybeUninit::read when api is stable
                            Some(n.data.as_ptr().read())
                        })
                        .map_err(|_| ())
                }
                None | Some(_) => Ok(None),
            }
        }
    }

    /// Attempts to dequeue from the front.
    ///
    /// Returns `None` if the queue is observed to be empty.
    pub fn try_pop(&self, guard: &Guard) -> Option<T> {
        loop {
            if let Ok(head) = self.pop_internal(guard) {
                return head;
            }
        }
    }

    /// Attempts to dequeue from the front, if the item satisfies the given condition.
    ///
    /// Returns `None` if the queue is observed to be empty, or the head does not satisfy the given
    /// condition.
    pub fn try_pop_if<F>(&self, condition: F, guard: &Guard) -> Option<T>
    where
        T: Sync,
        F: Fn(&T) -> bool,
    {
        loop {
            if let Ok(head) = self.pop_if_internal(&condition, guard) {
                return head;
            }
        }
    }
}

impl<T> Drop for Queue<T> {
    fn drop(&mut self) {
        unsafe {
            let guard = &unprotected();

            while let Some(_) = self.try_pop(guard) {}

            // Destroy the remaining sentinel node.
            let sentinel = self.head.load(Relaxed, guard);
            drop(sentinel.into_owned());
        }
    }
}

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

    struct Queue<T> {
        queue: super::Queue<T>,
    }

    impl<T> Queue<T> {
        pub fn new() -> Queue<T> {
            Queue {
                queue: super::Queue::new(),
            }
        }

        pub fn push(&self, t: T) {
            let guard = &pin();
            self.queue.push(t, guard);
        }

        pub fn is_empty(&self) -> bool {
            let guard = &pin();
            let head = self.queue.head.load(Acquire, guard);
            let h = unsafe { head.deref() };
            h.next.load(Acquire, guard).is_null()
        }

        pub fn try_pop(&self) -> Option<T> {
            let guard = &pin();
            self.queue.try_pop(guard)
        }

        pub fn pop(&self) -> T {
            loop {
                match self.try_pop() {
                    None => continue,
                    Some(t) => return t,
                }
            }
        }
    }

    const CONC_COUNT: i64 = 1000000;

    #[test]
    fn push_try_pop_1() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(37));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_2() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        q.push(48);
        assert_eq!(q.try_pop(), Some(37));
        assert!(!q.is_empty());
        assert_eq!(q.try_pop(), Some(48));
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_seq() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.try_pop(), Some(i));
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_1() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        q.push(37);
        assert!(!q.is_empty());
        assert_eq!(q.pop(), 37);
        assert!(q.is_empty());
    }

    #[test]
    fn push_pop_2() {
        let q: Queue<i64> = Queue::new();
        q.push(37);
        q.push(48);
        assert_eq!(q.pop(), 37);
        assert_eq!(q.pop(), 48);
    }

    #[test]
    fn push_pop_many_seq() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        for i in 0..200 {
            q.push(i)
        }
        assert!(!q.is_empty());
        for i in 0..200 {
            assert_eq!(q.pop(), i);
        }
        assert!(q.is_empty());
    }

    #[test]
    fn push_try_pop_many_spsc() {
        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());

        thread::scope(|scope| {
            scope.spawn(|_| {
                let mut next = 0;

                while next < CONC_COUNT {
                    if let Some(elem) = q.try_pop() {
                        assert_eq!(elem, next);
                        next += 1;
                    }
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        })
        .unwrap();
    }

    #[test]
    fn push_try_pop_many_spmc() {
        fn recv(_t: i32, q: &Queue<i64>) {
            let mut cur = -1;
            for _i in 0..CONC_COUNT {
                if let Some(elem) = q.try_pop() {
                    assert!(elem > cur);
                    cur = elem;

                    if cur == CONC_COUNT - 1 {
                        break;
                    }
                }
            }
        }

        let q: Queue<i64> = Queue::new();
        assert!(q.is_empty());
        thread::scope(|scope| {
            for i in 0..3 {
                let q = &q;
                scope.spawn(move |_| recv(i, q));
            }

            scope.spawn(|_| {
                for i in 0..CONC_COUNT {
                    q.push(i);
                }
            });
        })
        .unwrap();
    }

    #[test]
    fn push_try_pop_many_mpmc() {
        enum LR {
            Left(i64),
            Right(i64),
        }

        let q: Queue<LR> = Queue::new();
        assert!(q.is_empty());

        thread::scope(|scope| {
            for _t in 0..2 {
                scope.spawn(|_| {
                    for i in CONC_COUNT - 1..CONC_COUNT {
                        q.push(LR::Left(i))
                    }
                });
                scope.spawn(|_| {
                    for i in CONC_COUNT - 1..CONC_COUNT {
                        q.push(LR::Right(i))
                    }
                });
                scope.spawn(|_| {
                    let mut vl = vec![];
                    let mut vr = vec![];
                    for _i in 0..CONC_COUNT {
                        match q.try_pop() {
                            Some(LR::Left(x)) => vl.push(x),
                            Some(LR::Right(x)) => vr.push(x),
                            _ => {}
                        }
                    }

                    let mut vl2 = vl.clone();
                    let mut vr2 = vr.clone();
                    vl2.sort();
                    vr2.sort();

                    assert_eq!(vl, vl2);
                    assert_eq!(vr, vr2);
                });
            }
        })
        .unwrap();
    }

    #[test]
    fn push_pop_many_spsc() {
        let q: Queue<i64> = Queue::new();

        thread::scope(|scope| {
            scope.spawn(|_| {
                let mut next = 0;
                while next < CONC_COUNT {
                    assert_eq!(q.pop(), next);
                    next += 1;
                }
            });

            for i in 0..CONC_COUNT {
                q.push(i)
            }
        })
        .unwrap();
        assert!(q.is_empty());
    }

    #[test]
    fn is_empty_dont_pop() {
        let q: Queue<i64> = Queue::new();
        q.push(20);
        q.push(20);
        assert!(!q.is_empty());
        assert!(!q.is_empty());
        assert!(q.try_pop().is_some());
    }
}
<|endoftext|>


<commit_before>use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use std::time::Duration;

use crossbeam_deque::{Injector, Steal, Stealer, Worker};
use rand::distributions::{Distribution, Uniform};
use rand::thread_rng;
use scoped_threadpool::Pool;

use crate::gc::root::Slot;
use crate::gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, threadpool: &mut Pool) {
    let number_workers = threadpool.thread_count() as usize;
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);
    let injector = Injector::new();

    for _ in 0..number_workers {
        let w = Worker::new_lifo();
        let s = w.stealer();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                injector.push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let terminator = Terminator::new(number_workers);

    threadpool.scoped(|scoped| {
        for (task_id, worker) in workers.into_iter().enumerate() {
            let heap_region = heap.clone();
            let perm_region = perm.clone();

            let injector = &injector;
            let stealers = &stealers;
            let terminator = &terminator;

            scoped.execute(move || {
                let mut task = MarkingTask {
                    task_id,
                    local: Segment::new(),
                    worker,
                    injector,
                    stealers,
                    terminator,
                    heap_region,
                    perm_region,
                    marked: 0,
                };

                task.run();
            });
        }
    });
}

pub struct Terminator {
    const_nworkers: usize,
    nworkers: AtomicUsize,
}

impl Terminator {
    pub fn new(number_workers: usize) -> Terminator {
        Terminator {
            const_nworkers: number_workers,
            nworkers: AtomicUsize::new(number_workers),
        }
    }

    pub fn try_terminate(&self) -> bool {
        if self.const_nworkers == 1 {
            return true;
        }

        self.decrease_workers();
        thread::sleep(Duration::from_micros(1));
        self.zero_or_increase_workers()
    }

    fn decrease_workers(&self) -> bool {
        self.nworkers.fetch_sub(1, Ordering::SeqCst) == 1
    }

    fn zero_or_increase_workers(&self) -> bool {
        let mut nworkers = self.nworkers.load(Ordering::Relaxed);

        loop {
            if nworkers == 0 {
                return true;
            }

            let prev_nworkers =
                self.nworkers
                    .compare_and_swap(nworkers, nworkers + 1, Ordering::SeqCst);

            if nworkers == prev_nworkers {
                return false;
            }

            nworkers = prev_nworkers;
        }
    }
}

struct MarkingTask<'a> {
    task_id: usize,
    local: Segment,
    worker: Worker<Address>,
    injector: &'a Injector<Address>,
    stealers: &'a [Stealer<Address>],
    terminator: &'a Terminator,
    heap_region: Region,
    perm_region: Region,
    marked: usize,
}

impl<'a> MarkingTask<'a> {
    fn pop(&mut self) -> Option<Address> {
        self.pop_local()
            .or_else(|| self.pop_worker())
            .or_else(|| self.pop_global())
            .or_else(|| self.steal())
    }

    fn pop_local(&mut self) -> Option<Address> {
        if self.local.is_empty() {
            return None;
        }

        let obj = self.local.pop().expect("should be non-empty");
        Some(obj)
    }

    fn pop_worker(&mut self) -> Option<Address> {
        self.worker.pop()
    }

    fn pop_global(&mut self) -> Option<Address> {
        loop {
            let result = self.injector.steal_batch_and_pop(&mut self.worker);

            match result {
                Steal::Empty => break,
                Steal::Success(value) => return Some(value),
                Steal::Retry => continue,
            }
        }

        None
    }

    fn steal(&self) -> Option<Address> {
        if self.stealers.len() == 1 {
            return None;
        }

        let mut rng = thread_rng();
        let range = Uniform::new(0, self.stealers.len());

        for _ in 0..2 * self.stealers.len() {
            let mut stealer_id = self.task_id;

            while stealer_id == self.task_id {
                stealer_id = range.sample(&mut rng);
            }

            let stealer = &self.stealers[stealer_id];

            loop {
                match stealer.steal_batch_and_pop(&self.worker) {
                    Steal::Empty => break,
                    Steal::Success(address) => return Some(address),
                    Steal::Retry => continue,
                }
            }
        }

        None
    }

    fn run(&mut self) {
        loop {
            let object_addr = if let Some(object_addr) = self.pop() {
                object_addr
            } else if self.terminator.try_terminate() {
                break;
            } else {
                continue;
            };

            let object = object_addr.to_mut_obj();

            object.visit_reference_fields(|field| {
                self.trace(field);
            });
        }
    }

    fn trace(&mut self, slot: Slot) {
        let field_addr = slot.get();

        if self.heap_region.contains(field_addr) {
            let field_obj = field_addr.to_mut_obj();

            if field_obj.header().try_mark_non_atomic() {
                if self.local.has_capacity() {
                    self.local.push(field_addr);
                    self.defensive_push();
                } else {
                    self.worker.push(field_addr);
                }
            }
        } else {
            debug_assert!(field_addr.is_null() || self.perm_region.contains(field_addr));
        }
    }

    fn defensive_push(&mut self) {
        self.marked += 1;

        if self.marked > 256 {
            if self.local.len() > 4 {
                let target_len = self.local.len() / 2;

                while self.local.len() > target_len {
                    let val = self.local.pop().unwrap();
                    self.injector.push(val);
                }
            }

            self.marked = 0;
        }
    }
}

const SEGMENT_SIZE: usize = 64;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }

    fn len(&mut self) -> usize {
        self.data.len()
    }
}
<commit_msg>pmarking: use relaxed memory ordering for updating counter<commit_after>use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use std::time::Duration;

use crossbeam_deque::{Injector, Steal, Stealer, Worker};
use rand::distributions::{Distribution, Uniform};
use rand::thread_rng;
use scoped_threadpool::Pool;

use crate::gc::root::Slot;
use crate::gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, threadpool: &mut Pool) {
    let number_workers = threadpool.thread_count() as usize;
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);
    let injector = Injector::new();

    for _ in 0..number_workers {
        let w = Worker::new_lifo();
        let s = w.stealer();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                injector.push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let terminator = Terminator::new(number_workers);

    threadpool.scoped(|scoped| {
        for (task_id, worker) in workers.into_iter().enumerate() {
            let heap_region = heap.clone();
            let perm_region = perm.clone();

            let injector = &injector;
            let stealers = &stealers;
            let terminator = &terminator;

            scoped.execute(move || {
                let mut task = MarkingTask {
                    task_id,
                    local: Segment::new(),
                    worker,
                    injector,
                    stealers,
                    terminator,
                    heap_region,
                    perm_region,
                    marked: 0,
                };

                task.run();
            });
        }
    });
}

pub struct Terminator {
    const_nworkers: usize,
    nworkers: AtomicUsize,
}

impl Terminator {
    pub fn new(number_workers: usize) -> Terminator {
        Terminator {
            const_nworkers: number_workers,
            nworkers: AtomicUsize::new(number_workers),
        }
    }

    pub fn try_terminate(&self) -> bool {
        if self.const_nworkers == 1 {
            return true;
        }

        if self.decrease_workers() {
            // reached 0, no need to wait
            return true;
        }

        thread::sleep(Duration::from_micros(1));
        self.zero_or_increase_workers()
    }

    fn decrease_workers(&self) -> bool {
        self.nworkers.fetch_sub(1, Ordering::Relaxed) == 1
    }

    fn zero_or_increase_workers(&self) -> bool {
        let mut nworkers = self.nworkers.load(Ordering::Relaxed);

        loop {
            if nworkers == 0 {
                return true;
            }

            let prev_nworkers =
                self.nworkers
                    .compare_and_swap(nworkers, nworkers + 1, Ordering::Relaxed);

            if nworkers == prev_nworkers {
                // Value was successfully increased again, workers didn't terminate in time. There is still work left.
                return false;
            }

            nworkers = prev_nworkers;
        }
    }
}

struct MarkingTask<'a> {
    task_id: usize,
    local: Segment,
    worker: Worker<Address>,
    injector: &'a Injector<Address>,
    stealers: &'a [Stealer<Address>],
    terminator: &'a Terminator,
    heap_region: Region,
    perm_region: Region,
    marked: usize,
}

impl<'a> MarkingTask<'a> {
    fn pop(&mut self) -> Option<Address> {
        self.pop_local()
            .or_else(|| self.pop_worker())
            .or_else(|| self.pop_global())
            .or_else(|| self.steal())
    }

    fn pop_local(&mut self) -> Option<Address> {
        if self.local.is_empty() {
            return None;
        }

        let obj = self.local.pop().expect("should be non-empty");
        Some(obj)
    }

    fn pop_worker(&mut self) -> Option<Address> {
        self.worker.pop()
    }

    fn pop_global(&mut self) -> Option<Address> {
        loop {
            let result = self.injector.steal_batch_and_pop(&mut self.worker);

            match result {
                Steal::Empty => break,
                Steal::Success(value) => return Some(value),
                Steal::Retry => continue,
            }
        }

        None
    }

    fn steal(&self) -> Option<Address> {
        if self.stealers.len() == 1 {
            return None;
        }

        let mut rng = thread_rng();
        let range = Uniform::new(0, self.stealers.len());

        for _ in 0..2 * self.stealers.len() {
            let mut stealer_id = self.task_id;

            while stealer_id == self.task_id {
                stealer_id = range.sample(&mut rng);
            }

            let stealer = &self.stealers[stealer_id];

            loop {
                match stealer.steal_batch_and_pop(&self.worker) {
                    Steal::Empty => break,
                    Steal::Success(address) => return Some(address),
                    Steal::Retry => continue,
                }
            }
        }

        None
    }

    fn run(&mut self) {
        loop {
            let object_addr = if let Some(object_addr) = self.pop() {
                object_addr
            } else if self.terminator.try_terminate() {
                break;
            } else {
                continue;
            };

            let object = object_addr.to_mut_obj();

            object.visit_reference_fields(|field| {
                self.trace(field);
            });
        }
    }

    fn trace(&mut self, slot: Slot) {
        let field_addr = slot.get();

        if self.heap_region.contains(field_addr) {
            let field_obj = field_addr.to_mut_obj();

            if field_obj.header().try_mark_non_atomic() {
                if self.local.has_capacity() {
                    self.local.push(field_addr);
                    self.defensive_push();
                } else {
                    self.worker.push(field_addr);
                }
            }
        } else {
            debug_assert!(field_addr.is_null() || self.perm_region.contains(field_addr));
        }
    }

    fn defensive_push(&mut self) {
        self.marked += 1;

        if self.marked > 256 {
            if self.local.len() > 4 {
                let target_len = self.local.len() / 2;

                while self.local.len() > target_len {
                    let val = self.local.pop().unwrap();
                    self.injector.push(val);
                }
            }

            self.marked = 0;
        }
    }
}

const SEGMENT_SIZE: usize = 64;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }

    fn len(&mut self) -> usize {
        self.data.len()
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::firmware::tests::{
    create_fake_installation_set, create_fake_starup_callbacks, create_hook, device_attributes_dir,
    device_identity_dir, hardware_hook, product_uid_hook, version_hook,
};
use std::{any::Any, env, fs, io::Write, os::unix::fs::PermissionsExt, path::PathBuf};

pub use crate::{
    firmware::Metadata, runtime_settings::RuntimeSettings, settings::Settings,
    states::machine::Context,
};

pub struct TestEnvironment {
    pub firmware: Data<Metadata>,
    pub runtime_settings: Data<RuntimeSettings>,
    pub settings: Data<Settings>,
    pub binaries: Data<PathBuf>,
}

pub struct Data<T> {
    pub stored_path: PathBuf,
    #[allow(dead_code)]
    pub guard: Vec<Box<dyn Any>>,
    pub data: T,
}

#[derive(Default)]
pub struct TestEnvironmentBuilder {
    disable_polling: bool,
    invalid_hardware: bool,
    extra_binaries: Vec<String>,
    server_address: Option<String>,
    listen_socket: Option<String>,
}

impl TestEnvironment {
    pub fn build() -> TestEnvironmentBuilder {
        TestEnvironmentBuilder::default()
    }

    pub fn gen_context(&self) -> Context {
        Context::new(
            self.settings.data.clone(),
            self.runtime_settings.data.clone(),
            self.firmware.data.clone(),
        )
    }
}

impl TestEnvironmentBuilder {
    pub fn add_echo_binary(mut self, binary_name: &str) -> Self {
        self.extra_binaries.push(binary_name.to_owned());
        self
    }

    pub fn invalid_hardware(self) -> Self {
        TestEnvironmentBuilder { invalid_hardware: true, ..self }
    }

    pub fn disable_polling(self) -> Self {
        TestEnvironmentBuilder { disable_polling: true, ..self }
    }

    pub fn listen_socket(self, s: String) -> Self {
        TestEnvironmentBuilder { listen_socket: Some(s), ..self }
    }

    pub fn server_address(self, s: String) -> Self {
        TestEnvironmentBuilder { server_address: Some(s), ..self }
    }

    pub fn finish(self) -> TestEnvironment {
        let firmware = {
            let dir = tempfile::tempdir().unwrap();
            let dir_path = dir.path();

            // create fake hooks to be used to validate the load
            create_hook(
                product_uid_hook(dir_path),
                "#!/bin/sh\necho 229ffd7e08721d716163fc81a2dbaf6c90d449f0a3b009b6a2defe8a0b0d7381",
            );

            create_hook(version_hook(dir_path), "#!/bin/sh\necho 1.1");
            create_hook(
                hardware_hook(&dir_path),
                &format!(
                    "#!/bin/sh\necho {}",
                    match self.invalid_hardware {
                        false => "board",
                        true => "invalid",
                    }
                ),
            );
            create_hook(
                device_identity_dir(dir_path),
                "#!/bin/sh\necho id1=value1\necho id2=value2",
            );
            create_hook(
                device_attributes_dir(dir_path),
                "#!/bin/sh\necho attr1=attrvalue1\necho attr2=attrvalue2",
            );

            Data {
                data: Metadata::from_path(dir_path).unwrap(),
                stored_path: dir_path.to_owned(),
                guard: vec![Box::new(dir)],
            }
        };

        let binaries = {
            let bin_dir = tempfile::tempdir().unwrap();
            let bin_dir_path = bin_dir.path();
            let output_file = bin_dir_path.join("output");

            create_fake_installation_set(bin_dir_path, 0);
            // Startup callbacks will be stored in the firmware directory
            create_fake_starup_callbacks(&firmware.stored_path, &output_file);

            for bin in self.extra_binaries.into_iter() {
                let mut file = fs::File::create(&bin_dir_path.join(&bin)).unwrap();
                writeln!(file, "#!/bin/sh\necho $0 >> {}", output_file.to_string_lossy()).unwrap();
                let mut permissions = file.metadata().unwrap().permissions();
                permissions.set_mode(0o755);
                file.set_permissions(permissions).unwrap();
            }
            let curr_path = env::var("PATH").map(|s| ":".to_string() + &s).unwrap_or_default();
            env::set_var("PATH", format!("{}{}", bin_dir_path.to_string_lossy(), curr_path,));

            Data {
                data: output_file,
                stored_path: bin_dir_path.to_owned(),
                guard: vec![Box::new(bin_dir)],
            }
        };

        let runtime_settings = {
            let file = tempfile::NamedTempFile::new().unwrap();
            let file_path = file.path().to_owned();
            fs::remove_file(&file_path).unwrap();

            let mut runtime_settings = RuntimeSettings::default();
            runtime_settings.path = file_path.clone();

            Data { data: runtime_settings, stored_path: file_path, guard: vec![Box::new(file)] }
        };

        let settings = {
            let mut file = tempfile::NamedTempFile::new().unwrap();
            let file_path = file.path().to_owned();
            let download_dir = tempfile::tempdir().unwrap();

            write!(
                file,
                r#"[network]
server_address={}
listen_socket={}

[storage]
read_only=false
runtime_settings={runtime_settings}

[polling]
enabled={polling_enabled}
interval="1d"

[update]
download_dir={download_dir}
supported_install_modes=["copy", "tarball"]

[firmware]
metadata={metadata}"#,
                server_address = toml::to_string(
                    self.server_address.as_deref().unwrap_or("https://api.updatehub.io")
                )
                .unwrap(),
                listen_socket =
                    toml::to_string(self.listen_socket.as_deref().unwrap_or("localhost:8080"))
                        .unwrap(),
                runtime_settings = toml::to_string(&runtime_settings.stored_path).unwrap(),
                polling_enabled = toml::to_string(&!self.disable_polling).unwrap(),
                download_dir = toml::to_string(download_dir.path()).unwrap(),
                metadata = toml::to_string(&firmware.stored_path).unwrap()
            )
            .unwrap();

            Data {
                data: Settings::load(&file_path).unwrap(),
                stored_path: file_path,
                guard: vec![Box::new(file), Box::new(download_dir)],
            }
        };

        TestEnvironment { firmware, runtime_settings, settings, binaries }
    }
}
<commit_msg>updatehub: Add test object to supported_install_modes on TestEnviroment<commit_after>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

use crate::firmware::tests::{
    create_fake_installation_set, create_fake_starup_callbacks, create_hook, device_attributes_dir,
    device_identity_dir, hardware_hook, product_uid_hook, version_hook,
};
use std::{any::Any, env, fs, io::Write, os::unix::fs::PermissionsExt, path::PathBuf};

pub use crate::{
    firmware::Metadata, runtime_settings::RuntimeSettings, settings::Settings,
    states::machine::Context,
};

pub struct TestEnvironment {
    pub firmware: Data<Metadata>,
    pub runtime_settings: Data<RuntimeSettings>,
    pub settings: Data<Settings>,
    pub binaries: Data<PathBuf>,
}

pub struct Data<T> {
    pub stored_path: PathBuf,
    #[allow(dead_code)]
    pub guard: Vec<Box<dyn Any>>,
    pub data: T,
}

#[derive(Default)]
pub struct TestEnvironmentBuilder {
    disable_polling: bool,
    invalid_hardware: bool,
    extra_binaries: Vec<String>,
    server_address: Option<String>,
    listen_socket: Option<String>,
}

impl TestEnvironment {
    pub fn build() -> TestEnvironmentBuilder {
        TestEnvironmentBuilder::default()
    }

    pub fn gen_context(&self) -> Context {
        Context::new(
            self.settings.data.clone(),
            self.runtime_settings.data.clone(),
            self.firmware.data.clone(),
        )
    }
}

impl TestEnvironmentBuilder {
    pub fn add_echo_binary(mut self, binary_name: &str) -> Self {
        self.extra_binaries.push(binary_name.to_owned());
        self
    }

    pub fn invalid_hardware(self) -> Self {
        TestEnvironmentBuilder { invalid_hardware: true, ..self }
    }

    pub fn disable_polling(self) -> Self {
        TestEnvironmentBuilder { disable_polling: true, ..self }
    }

    pub fn listen_socket(self, s: String) -> Self {
        TestEnvironmentBuilder { listen_socket: Some(s), ..self }
    }

    pub fn server_address(self, s: String) -> Self {
        TestEnvironmentBuilder { server_address: Some(s), ..self }
    }

    pub fn finish(self) -> TestEnvironment {
        let firmware = {
            let dir = tempfile::tempdir().unwrap();
            let dir_path = dir.path();

            // create fake hooks to be used to validate the load
            create_hook(
                product_uid_hook(dir_path),
                "#!/bin/sh\necho 229ffd7e08721d716163fc81a2dbaf6c90d449f0a3b009b6a2defe8a0b0d7381",
            );

            create_hook(version_hook(dir_path), "#!/bin/sh\necho 1.1");
            create_hook(
                hardware_hook(&dir_path),
                &format!(
                    "#!/bin/sh\necho {}",
                    match self.invalid_hardware {
                        false => "board",
                        true => "invalid",
                    }
                ),
            );
            create_hook(
                device_identity_dir(dir_path),
                "#!/bin/sh\necho id1=value1\necho id2=value2",
            );
            create_hook(
                device_attributes_dir(dir_path),
                "#!/bin/sh\necho attr1=attrvalue1\necho attr2=attrvalue2",
            );

            Data {
                data: Metadata::from_path(dir_path).unwrap(),
                stored_path: dir_path.to_owned(),
                guard: vec![Box::new(dir)],
            }
        };

        let binaries = {
            let bin_dir = tempfile::tempdir().unwrap();
            let bin_dir_path = bin_dir.path();
            let output_file = bin_dir_path.join("output");

            create_fake_installation_set(bin_dir_path, 0);
            // Startup callbacks will be stored in the firmware directory
            create_fake_starup_callbacks(&firmware.stored_path, &output_file);

            for bin in self.extra_binaries.into_iter() {
                let mut file = fs::File::create(&bin_dir_path.join(&bin)).unwrap();
                writeln!(file, "#!/bin/sh\necho $0 >> {}", output_file.to_string_lossy()).unwrap();
                let mut permissions = file.metadata().unwrap().permissions();
                permissions.set_mode(0o755);
                file.set_permissions(permissions).unwrap();
            }
            let curr_path = env::var("PATH").map(|s| ":".to_string() + &s).unwrap_or_default();
            env::set_var("PATH", format!("{}{}", bin_dir_path.to_string_lossy(), curr_path,));

            Data {
                data: output_file,
                stored_path: bin_dir_path.to_owned(),
                guard: vec![Box::new(bin_dir)],
            }
        };

        let runtime_settings = {
            let file = tempfile::NamedTempFile::new().unwrap();
            let file_path = file.path().to_owned();
            fs::remove_file(&file_path).unwrap();

            let mut runtime_settings = RuntimeSettings::default();
            runtime_settings.path = file_path.clone();

            Data { data: runtime_settings, stored_path: file_path, guard: vec![Box::new(file)] }
        };

        let settings = {
            let mut file = tempfile::NamedTempFile::new().unwrap();
            let file_path = file.path().to_owned();
            let download_dir = tempfile::tempdir().unwrap();

            write!(
                file,
                r#"[network]
server_address={}
listen_socket={}

[storage]
read_only=false
runtime_settings={runtime_settings}

[polling]
enabled={polling_enabled}
interval="1d"

[update]
download_dir={download_dir}
supported_install_modes=["copy", "tarball", "test"]

[firmware]
metadata={metadata}"#,
                server_address = toml::to_string(
                    self.server_address.as_deref().unwrap_or("https://api.updatehub.io")
                )
                .unwrap(),
                listen_socket =
                    toml::to_string(self.listen_socket.as_deref().unwrap_or("localhost:8080"))
                        .unwrap(),
                runtime_settings = toml::to_string(&runtime_settings.stored_path).unwrap(),
                polling_enabled = toml::to_string(&!self.disable_polling).unwrap(),
                download_dir = toml::to_string(download_dir.path()).unwrap(),
                metadata = toml::to_string(&firmware.stored_path).unwrap()
            )
            .unwrap();

            Data {
                data: Settings::load(&file_path).unwrap(),
                stored_path: file_path,
                guard: vec![Box::new(file), Box::new(download_dir)],
            }
        };

        TestEnvironment { firmware, runtime_settings, settings, binaries }
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Write};
use std::path::Path;
use std::process;

const NOSYMLINK: &str = "
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When building `cxx` from a git clone, git's symlink support needs
to be enabled on platforms that have it off by default (Windows).
Either use:

   $ git config --global core.symlinks true

prior to cloning, or else use:

   $ git clone -c core.symlinks=true https://github.com/dtolnay/cxx

for the clone.

Symlinks are only required when compiling locally from a clone of
the git repository---they are NOT required when building `cxx` as
a Cargo-managed (possibly transitive) build dependency downloaded
through crates.io.
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
";

fn main() {
    if !Path::new("src/syntax/mod.rs").exists() {
        let _ = io::stderr().lock().write_all(NOSYMLINK.as_bytes());
        process::exit(1);
    }
}
<commit_msg>Try to detect symlinking disabled on Windows<commit_after>use std::io::{self, Write};
#[cfg(windows)]
use std::os::windows::fs as windows;
use std::path::Path;
use std::process;
#[cfg(windows)]
use std::{env, fs};

const MISSING: &str = "
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When building `cxx` from a git clone, git's symlink support needs
to be enabled on platforms that have it off by default (Windows).
Either use:

   $ git config --global core.symlinks true

prior to cloning, or else use:

   $ git clone -c core.symlinks=true https://github.com/dtolnay/cxx

for the clone.

Symlinks are only required when compiling locally from a clone of
the git repository---they are NOT required when building `cxx` as
a Cargo-managed (possibly transitive) build dependency downloaded
through crates.io.
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
";

#[cfg(windows)]
const DENIED: &str = "
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When building `cxx` from a git clone on Windows we need Developer
Mode enabled for symlink support.

To enable Developer Mode: go under Settings to Update & Security,
then 'For developers', and turn on the toggle for Developer Mode.

For more explanation of symlinks in Windows, see these resources:
> https://blogs.windows.com/windowsdeveloper/2016/12/02/symlinks-windows-10/
> https://docs.microsoft.com/windows/uwp/get-started/enable-your-device-for-development

Symlinks are only required when compiling locally from a clone of
the git repository---they are NOT required when building `cxx` as
a Cargo-managed (possibly transitive) build dependency downloaded
through crates.io.
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
";

fn main() {
    if Path::new("src/syntax/mod.rs").exists() {
        return;
    }

    #[allow(unused_mut)]
    let mut message = MISSING;

    #[cfg(windows)]
    if let Some(out_dir) = env::var_os("OUT_DIR") {
        let parent_dir = Path::new(&out_dir).join("symlink");
        let from_dir = parent_dir.join("from");
        let to_dir = parent_dir.join("to");
        if fs::create_dir_all(&from_dir).is_ok()
            && windows::symlink_dir(&from_dir, &to_dir).is_err()
        {
            message = DENIED;
        }
    }

    let _ = io::stderr().write_all(message.as_bytes());
    process::exit(1);
}
<|endoftext|>