<commit_before>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA, VA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
        let size = ws.read_u32(load_config_directory)?;

        // max offset into the config directory that we'll read.
        let max_config_directory_offset = match is_64 {
            true => 0x94,   // CFG flags
            false => 0x58,  // CFG flags
        };

        // in `d3d11sdklayers.dll` for example, the config size is 0x70,
        // which is much too small to read the CFG table options.
        if max_config_directory_offset > size {
            debug!("no CF Guard table: load config directory too small");
            return Ok(())
        }

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 0x90)?,
            false => ws.read_u32(load_config_directory + 0x58)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 0x80)?,
            false => ws.read_va(load_config_directory + 0x50)?,
        };
        debug!("CF Guard table: {}", cfg_table_va);

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 0x88)?,
            false => ws.read_u32(load_config_directory + 0x54)?,
        };
        debug!("CF Guard table count: {}", cfg_table_count);

        if cfg_table_va == VA(0x0) {
            debug!("CF guard table empty");
            return Ok(())
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        debug!("CF guard table: {}", cfg_table_rva);

        let mut offset = cfg_table_rva;
        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 112)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 72)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8, Permissions::R) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1, Permissions::X) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 120)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 76)?).unwrap(),
        };
        if ws.probe(guard_dispatch_icall_fptr, 8, Permissions::R) {
            let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
            if ws.probe(guard_dispatch_icall, 1, Permissions::X) {
                debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                ws.make_function(guard_dispatch_icall)?;
                ws.analyze()?;
            }
        };

        Ok(())
    }
}
<commit_msg>cfguard: use hex offsets<commit_after>/// this analyzer parses the Control Flow Guard function table.
/// it populates functions in the workspace.
/// the purpose of the table is to enumerate the functions that may be called indirectly.
///
/// references:
///   - https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory
///   - https://lucasg.github.io/2017/02/05/Control-Flow-Guard/
use log::{debug};
use goblin::{Object};
use failure::{Error};

use super::super::super::arch::{RVA, VA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct CFGuardTableAnalyzer {}

impl CFGuardTableAnalyzer {
    pub fn new() -> CFGuardTableAnalyzer {
        CFGuardTableAnalyzer {}
    }
}

const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK: u32 = 0xF0000000;
const IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT: u32 = 28;

impl Analyzer for CFGuardTableAnalyzer {
    fn get_name(&self) -> String {
        "CF Guard Table analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::CFGuardTableAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// CFGuardTableAnalyzer::new().analyze(&mut ws);
    ///
    /// // export: RtlVirtualUnwind
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x1010)).is_some());
    ///
    /// // __guard_check_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21960)).is_some());
    ///
    /// // __guard_dispatch_icall
    /// assert!(ws.get_functions().find(|&&rva| rva == RVA(0x21B40)).is_some());
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let (load_config_directory, is_64) = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let load_config_directory = match opt_header.data_directories.get_load_config_table() {
                Some(load_config_directory) => load_config_directory,
                _ => return Ok(()),
            };

            (RVA::from(load_config_directory.virtual_address as i64), pe.is_64)
        };

        debug!("load config directory: {}", load_config_directory);

        // offsets defined here:
        // https://docs.microsoft.com/en-us/windows/desktop/debug/pe-format#load-configuration-directory

        // according to IDA, the first DWORD is `Size` not `Characteristics` (unused).
        let size = ws.read_u32(load_config_directory)?;

        // max offset into the config directory that we'll read.
        let max_config_directory_offset = match is_64 {
            true => 0x94,   // CFG flags
            false => 0x58,  // CFG flags
        };

        // in `d3d11sdklayers.dll` for example, the config size is 0x70,
        // which is much too small to read the CFG table options.
        if max_config_directory_offset > size {
            debug!("no CF Guard table: load config directory too small");
            return Ok(())
        }

        let cfg_flags = match is_64 {
            true => ws.read_u32(load_config_directory + 0x90)?,
            false => ws.read_u32(load_config_directory + 0x58)?,
        };

        let stride = (cfg_flags & IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_MASK) >> IMAGE_GUARD_CF_FUNCTION_TABLE_SIZE_SHIFT;

        let cfg_table_va = match is_64 {
            true => ws.read_va(load_config_directory + 0x80)?,
            false => ws.read_va(load_config_directory + 0x50)?,
        };
        debug!("CF Guard table: {}", cfg_table_va);

        let cfg_table_count = match is_64 {
            true => ws.read_u32(load_config_directory + 0x88)?,
            false => ws.read_u32(load_config_directory + 0x54)?,
        };
        debug!("CF Guard table count: {}", cfg_table_count);

        if cfg_table_va == VA(0x0) {
            debug!("CF guard table empty");
            return Ok(())
        };

        let cfg_table_rva = ws.rva(cfg_table_va).unwrap();
        debug!("CF guard table: {}", cfg_table_rva);

        let mut offset = cfg_table_rva;
        for _ in 0..cfg_table_count {
            let function = RVA::from(ws.read_i32(offset)?);

            debug!("CF guard function: {}", function);
            ws.make_function(function)?;
            ws.analyze()?;

            // 4 == sizeof(32-bit RVA)
            offset = offset + (4 + stride as usize);
        }

        // add function pointed to by GuardCFCheckFunctionPointer
        let guard_check_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 0x70)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 0x48)?).unwrap(),
        };
        if ws.probe(guard_check_icall_fptr, 8, Permissions::R) {
            let guard_check_icall = ws.rva(ws.read_va(guard_check_icall_fptr)?).unwrap();
            if ws.probe(guard_check_icall, 1, Permissions::X) {
                debug!("CF guard check function: {:#x}", guard_check_icall);
                ws.make_function(guard_check_icall)?;
                ws.analyze()?;
            }
        };

        // add function pointed to by GuardCFDispatchFunctionPointer
        let guard_dispatch_icall_fptr = match is_64 {
            true => ws.rva(ws.read_va(load_config_directory + 120)?).unwrap(),
            false => ws.rva(ws.read_va(load_config_directory + 76)?).unwrap(),
        };
        if ws.probe(guard_dispatch_icall_fptr, 8, Permissions::R) {
            let guard_dispatch_icall = ws.rva(ws.read_va(guard_dispatch_icall_fptr)?).unwrap();
            if ws.probe(guard_dispatch_icall, 1, Permissions::X) {
                debug!("CF guard dispatch function: {:#x}", guard_dispatch_icall);
                ws.make_function(guard_dispatch_icall)?;
                ws.analyze()?;
            }
        };

        Ok(())
    }
}
<|endoftext|>


<commit_before>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![feature(test)]

extern crate protobuf;
extern crate test;

use protobuf::stream;

use self::test::Bencher;

#[inline]
fn buffer_write_byte(os: &mut stream::CodedOutputStream) {
    for i in 0..10 {
        os.write_raw_byte(test::black_box(i as u8)).unwrap();
    }
    os.flush().unwrap();
}

#[inline]
fn buffer_write_bytes(os: &mut stream::CodedOutputStream) {
    for _ in 0..10 {
        os.write_raw_bytes(test::black_box(b"1234567890")).unwrap();
    }
    os.flush().unwrap();
}

#[bench]
fn bench_buffer(b: &mut Bencher) {
    b.iter(|| {
        let mut v = Vec::new();
        {
            let mut os = stream::CodedOutputStream::new(&mut v);
            buffer_write_byte(&mut os);
        }
        v
    });
}

#[bench]
fn bench_buffer_bytes(b: &mut Bencher) {
    b.iter(|| {
        let mut v = Vec::new();
        {
            let mut os = stream::CodedOutputStream::new(&mut v);
            buffer_write_bytes(&mut os);
        }
        v
    });
}
<commit_msg>bench os.write_raw_varint32<commit_after>// `cargo test --benches` and `#[feature(test)]` work only in nightly
#![cfg(rustc_nightly)]
#![feature(test)]

extern crate protobuf;
extern crate test;

use protobuf::stream;

use self::test::Bencher;

#[inline]
fn buffer_write_byte(os: &mut stream::CodedOutputStream) {
    for i in 0..10 {
        os.write_raw_byte(test::black_box(i as u8)).unwrap();
    }
    os.flush().unwrap();
}

#[inline]
fn buffer_write_bytes(os: &mut stream::CodedOutputStream) {
    for _ in 0..10 {
        os.write_raw_bytes(test::black_box(b"1234567890")).unwrap();
    }
    os.flush().unwrap();
}

#[bench]
fn bench_buffer(b: &mut Bencher) {
    b.iter(|| {
        let mut v = Vec::new();
        {
            let mut os = stream::CodedOutputStream::new(&mut v);
            buffer_write_byte(&mut os);
        }
        v
    });
}

#[bench]
fn bench_buffer_bytes(b: &mut Bencher) {
    b.iter(|| {
        let mut v = Vec::new();
        {
            let mut os = stream::CodedOutputStream::new(&mut v);
            buffer_write_bytes(&mut os);
        }
        v
    });
}

#[bench]
fn bench_write_raw_varint_32(b: &mut Bencher) {
    let mut v = Vec::with_capacity(10_000);
    b.iter(|| {
        v.clear();
        {
            let mut os = stream::CodedOutputStream::new(&mut v);
            for i in 0..1000 {
                os.write_raw_varint32(i * 139 % 1000).unwrap();
            }
        }
        v.len()
    })
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::io::{Cursor, Read, Seek, SeekFrom};
use cbor::cbor::{CborError, CborType};

// We limit the length of any cbor byte array to 128MiB. This is a somewhat
// arbitrary limit that should work on all platforms and is large enough for
// any benign data.
pub const MAX_ARRAY_SIZE: usize = 134217728;

/// Struct holding a cursor and additional information for decoding.
#[derive(Debug)]
struct DecoderCursor {
    pub cursor: Cursor<Vec<u8>>,
}

/// Apply this mask (with &) to get the value part of the initial byte of a CBOR item.
const INITIAL_VALUE_MASK: u64 = 0b0001_1111;

impl DecoderCursor {
    /// Read and return the given number of bytes from the cursor. Advances the cursor.
    fn read_bytes(&mut self, len: usize) -> Result<Vec<u8>, CborError> {
        if len > MAX_ARRAY_SIZE {
            return Err(CborError::InputTooLarge);
        }
        let mut buf: Vec<u8> = vec![0; len];
        match self.cursor.read_exact(&mut buf) {
            Err(_) => return Err(CborError::TruncatedInput),
            Ok(()) => {}
        }
        Ok(buf)
    }

    /// Convert num bytes to a u64
    fn read_int_from_bytes(&mut self, num: usize) -> Result<u64, CborError> {
        let x = self.read_bytes(num)?;
        let mut result: u64 = 0;
        for i in (0..num).rev() {
            result += u64::from(x[num - 1 - i]) << (i * 8);
        }
        Ok(result)
    }

    /// Read an integer and return it as u64.
    fn read_int(&mut self) -> Result<u64, CborError> {
        let first_value = self.read_int_from_bytes(1)? & INITIAL_VALUE_MASK;
        let val = match first_value {
            0...23 => first_value,
            24 => self.read_int_from_bytes(1)?,
            25 => self.read_int_from_bytes(2)?,
            26 => self.read_int_from_bytes(4)?,
            27 => self.read_int_from_bytes(8)?,
            _ => return Err(CborError::MalformedInput),
        };
        Ok(val)
    }

    fn read_signed_int(&mut self) -> Result<CborType, CborError> {
        let uint = self.read_int()?;
        if uint > i64::max_value() as u64 {
            return Err(CborError::InputValueOutOfRange);
        }
        let result: i64 = -1 - uint as i64;
        Ok(CborType::SignedInteger(result))
    }

    /// Read an array of data items and return it.
    fn read_array(&mut self) -> Result<CborType, CborError> {
        // Create a new array.
        let mut array: Vec<CborType> = Vec::new();
        // Read the length of the array.
        let num_items = self.read_int()?;
        // Decode each of the num_items data items.
        for _ in 0..num_items {
            let new_item = self.decode_item()?;
            array.push(new_item);
        }
        Ok(CborType::Array(array))
    }

    /// Read a byte string and return it.
    fn read_byte_string(&mut self) -> Result<CborType, CborError> {
        let length = self.read_int()?;
        let limit = usize::max_value() as u64;
        if length > limit {
            return Err(CborError::InputTooLarge);
        }
        let byte_string = self.read_bytes(length as usize)?;
        Ok(CborType::Bytes(byte_string))
    }

    /// Read a map.
    fn read_map(&mut self) -> Result<CborType, CborError> {
        let num_items = self.read_int()?;
        // Create a new array.
        let mut map: BTreeMap<CborType, CborType> = BTreeMap::new();
        // Decode each of the num_items (key, data item) pairs.
        for _ in 0..num_items {
            let key_val = self.decode_item()?;
            let item_value = self.decode_item()?;
            if map.contains_key(&key_val) {
                return Err(CborError::DuplicateMapKey);
            }
            map.insert(key_val, item_value);
        }
        Ok(CborType::Map(map))
    }

    fn read_null(&mut self) -> Result<CborType, CborError> {
        let value = self.read_int_from_bytes(1)? & INITIAL_VALUE_MASK;
        if value != 22 {
            return Err(CborError::UnsupportedType);
        }
        Ok(CborType::Null)
    }

    /// Peeks at the next byte in the cursor, but does not change the position.
    fn peek_byte(&mut self) -> Result<u8, CborError> {
        let x = self.read_bytes(1)?;
        if self.cursor.seek(SeekFrom::Current(-1)).is_err() {
            return Err(CborError::LibraryError)
        };
        Ok(x[0])
    }

    /// Decodes the next CBOR item.
    pub fn decode_item(&mut self) -> Result<CborType, CborError> {
        let major_type = self.peek_byte()? >> 5;
        let result = match major_type {
            0 => {
                let value = self.read_int()?;
                CborType::Integer(value)
            }
            1 => self.read_signed_int()?,
            2 => self.read_byte_string()?,
            4 => self.read_array()?,
            5 => self.read_map()?,
            6 => {
                let tag = self.read_int()?;
                let item = self.decode_item()?;
                CborType::Tag(tag, Box::new(item))
            }
            7 => self.read_null()?,
            _ => unreachable!(),
        };
        Ok(result)
    }
}

/// Read the CBOR structure in bytes and return it as a `CborType`.
pub fn decode(bytes: Vec<u8>) -> Result<CborType, CborError> {
    let mut decoder_cursor = DecoderCursor { cursor: Cursor::new(bytes) };
    decoder_cursor.decode_item()
    // TODO: check cursor at end?
}
<commit_msg>use MAX_ARRAY_SIZE as limit in read_byte_string<commit_after>use std::collections::BTreeMap;
use std::io::{Cursor, Read, Seek, SeekFrom};
use cbor::cbor::{CborError, CborType};

// We limit the length of any cbor byte array to 128MiB. This is a somewhat
// arbitrary limit that should work on all platforms and is large enough for
// any benign data.
pub const MAX_ARRAY_SIZE: usize = 134_217_728;

/// Struct holding a cursor and additional information for decoding.
#[derive(Debug)]
struct DecoderCursor {
    pub cursor: Cursor<Vec<u8>>,
}

/// Apply this mask (with &) to get the value part of the initial byte of a CBOR item.
const INITIAL_VALUE_MASK: u64 = 0b0001_1111;

impl DecoderCursor {
    /// Read and return the given number of bytes from the cursor. Advances the cursor.
    fn read_bytes(&mut self, len: usize) -> Result<Vec<u8>, CborError> {
        if len > MAX_ARRAY_SIZE {
            return Err(CborError::InputTooLarge);
        }
        let mut buf: Vec<u8> = vec![0; len];
        match self.cursor.read_exact(&mut buf) {
            Err(_) => return Err(CborError::TruncatedInput),
            Ok(()) => {}
        }
        Ok(buf)
    }

    /// Convert num bytes to a u64
    fn read_int_from_bytes(&mut self, num: usize) -> Result<u64, CborError> {
        let x = self.read_bytes(num)?;
        let mut result: u64 = 0;
        for i in (0..num).rev() {
            result += u64::from(x[num - 1 - i]) << (i * 8);
        }
        Ok(result)
    }

    /// Read an integer and return it as u64.
    fn read_int(&mut self) -> Result<u64, CborError> {
        let first_value = self.read_int_from_bytes(1)? & INITIAL_VALUE_MASK;
        let val = match first_value {
            0...23 => first_value,
            24 => self.read_int_from_bytes(1)?,
            25 => self.read_int_from_bytes(2)?,
            26 => self.read_int_from_bytes(4)?,
            27 => self.read_int_from_bytes(8)?,
            _ => return Err(CborError::MalformedInput),
        };
        Ok(val)
    }

    fn read_signed_int(&mut self) -> Result<CborType, CborError> {
        let uint = self.read_int()?;
        if uint > i64::max_value() as u64 {
            return Err(CborError::InputValueOutOfRange);
        }
        let result: i64 = -1 - uint as i64;
        Ok(CborType::SignedInteger(result))
    }

    /// Read an array of data items and return it.
    fn read_array(&mut self) -> Result<CborType, CborError> {
        // Create a new array.
        let mut array: Vec<CborType> = Vec::new();
        // Read the length of the array.
        let num_items = self.read_int()?;
        // Decode each of the num_items data items.
        for _ in 0..num_items {
            let new_item = self.decode_item()?;
            array.push(new_item);
        }
        Ok(CborType::Array(array))
    }

    /// Read a byte string and return it.
    fn read_byte_string(&mut self) -> Result<CborType, CborError> {
        let length = self.read_int()?;
        if length > MAX_ARRAY_SIZE as u64 {
            return Err(CborError::InputTooLarge);
        }
        let byte_string = self.read_bytes(length as usize)?;
        Ok(CborType::Bytes(byte_string))
    }

    /// Read a map.
    fn read_map(&mut self) -> Result<CborType, CborError> {
        let num_items = self.read_int()?;
        // Create a new array.
        let mut map: BTreeMap<CborType, CborType> = BTreeMap::new();
        // Decode each of the num_items (key, data item) pairs.
        for _ in 0..num_items {
            let key_val = self.decode_item()?;
            let item_value = self.decode_item()?;
            if map.contains_key(&key_val) {
                return Err(CborError::DuplicateMapKey);
            }
            map.insert(key_val, item_value);
        }
        Ok(CborType::Map(map))
    }

    fn read_null(&mut self) -> Result<CborType, CborError> {
        let value = self.read_int_from_bytes(1)? & INITIAL_VALUE_MASK;
        if value != 22 {
            return Err(CborError::UnsupportedType);
        }
        Ok(CborType::Null)
    }

    /// Peeks at the next byte in the cursor, but does not change the position.
    fn peek_byte(&mut self) -> Result<u8, CborError> {
        let x = self.read_bytes(1)?;
        if self.cursor.seek(SeekFrom::Current(-1)).is_err() {
            return Err(CborError::LibraryError)
        };
        Ok(x[0])
    }

    /// Decodes the next CBOR item.
    pub fn decode_item(&mut self) -> Result<CborType, CborError> {
        let major_type = self.peek_byte()? >> 5;
        let result = match major_type {
            0 => {
                let value = self.read_int()?;
                CborType::Integer(value)
            }
            1 => self.read_signed_int()?,
            2 => self.read_byte_string()?,
            4 => self.read_array()?,
            5 => self.read_map()?,
            6 => {
                let tag = self.read_int()?;
                let item = self.decode_item()?;
                CborType::Tag(tag, Box::new(item))
            }
            7 => self.read_null()?,
            _ => unreachable!(),
        };
        Ok(result)
    }
}

/// Read the CBOR structure in bytes and return it as a `CborType`.
pub fn decode(bytes: Vec<u8>) -> Result<CborType, CborError> {
    let mut decoder_cursor = DecoderCursor { cursor: Cursor::new(bytes) };
    decoder_cursor.decode_item()
    // TODO: check cursor at end?
}
<|endoftext|>


<commit_before>use color::PuyoColor;
use field;
use field_bit::FieldBit;
use plain_field::PuyoPlainField;
use std::mem;
use x86intrin::*;

#[derive(Clone, Copy)]
pub struct BitField {
    m: [FieldBit; 3],
}

impl BitField {
    pub fn new() -> BitField {
        BitField {
            m: [FieldBit::new_empty(),
                FieldBit::from_values(0xFFFF, 0x8001, 0x8001, 0x8001, 0x8001, 0x8001, 0x8001, 0xFFFF),
                FieldBit::new_empty(), ]
        }
    }

    pub fn from_plain_field(pf: PuyoPlainField) -> BitField {
        let mut bf = BitField::new();

        // TODO(mayah): We have better algorithm here.
        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                bf.set_color(x, y, pf.color(x, y))
            }
        }

        bf
    }

    pub fn from_str(s: &str) -> BitField {
        BitField::from_plain_field(PuyoPlainField::from_str(s))
    }

    pub fn color(&self, x: usize, y: usize) -> PuyoColor {
        let b0: u8 = if self.m[0].get(x, y) { 1 } else { 0 };
        let b1: u8 = if self.m[1].get(x, y) { 2 } else { 0 };
        let b2: u8 = if self.m[2].get(x, y) { 4 } else { 0 };

        unsafe {
            mem::transmute(b0 | b1 | b2)
        }
    }

    pub fn is_color(&self, x: usize, y: usize, c: PuyoColor) -> bool {
        self.bits(c).get(x, y)
    }

    pub fn is_normal_color(&self, x: usize, y: usize) -> bool {
        self.m[2].get(x, y)
    }

    pub fn set_color(&mut self, x: usize, y: usize, c: PuyoColor) {
        let cc = c as u8;
        for i in 0 .. 3 {
            if (cc & (1 << i)) != 0 {
                self.m[i as usize].set(x, y);
            } else {
                self.m[i as usize].unset(x, y);
            }
        }
    }

    pub fn count_connected_puyos(&self, x: usize, y: usize) -> usize {
        if y > field::HEIGHT {
            return 0
        }

        let c = self.color(x, y);
        let color_bits = self.bits(c).masked_field_12();
        FieldBit::from_onebit(x, y).expand(&color_bits).popcount()
    }

    pub fn bits(&self, c: PuyoColor) -> FieldBit {
        let r0 = self.m[0].as_m128i();
        let r1 = self.m[1].as_m128i();
        let r2 = self.m[2].as_m128i();

        let v = match c {
            PuyoColor::EMPTY => {  // 0
                let x = mm_or_si128(mm_or_si128(r0, r1), r2);
                mm_xor_si128(x, mm_setr_epi32(!0, !0, !0, !0))
            },
            PuyoColor::OJAMA => {  // 1
                mm_andnot_si128(r2, mm_andnot_si128(r1, r0))
            },
            PuyoColor::WALL => {   // 2
                mm_andnot_si128(r2, mm_andnot_si128(r0, r1))
            },
            PuyoColor::IRON => {   // 3
                mm_andnot_si128(r2, mm_and_si128(r0, r1))
            },
            PuyoColor::RED => {    // 4
                mm_andnot_si128(r0, mm_andnot_si128(r1, r2))
            },
            PuyoColor::BLUE => {   // 5
                mm_and_si128(r0, mm_andnot_si128(r1, r2))
            },
            PuyoColor::YELLOW => { // 6
                mm_andnot_si128(r0, mm_and_si128(r1, r2))
            },
            PuyoColor::GREEN => {  // 7
                mm_and_si128(r0, mm_and_si128(r1, r2))
            },
        };

        FieldBit::new(v)
    }
}

#[cfg(test)]
mod tests {
    use super::BitField;
    use color;
    use color::PuyoColor;
    use field;

    #[test]
    fn test_initial() {
        let bf = BitField::new();
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                if x == 0 || x == 7 || y == 0 || y == 15 {
                    assert_eq!(bf.color(x, y), PuyoColor::WALL);
                } else {
                    assert_eq!(bf.color(x, y), PuyoColor::EMPTY);
                }
            }
        }
    }

    #[test]
    fn test_from_str() {
        let bf = BitField::from_str(concat!(
            "RGBRGB",
            "RGBRGB",
            "RGBRGB"));

        assert_eq!(bf.color(1, 1), PuyoColor::RED);
        assert_eq!(bf.color(2, 1), PuyoColor::GREEN);
        assert_eq!(bf.color(3, 1), PuyoColor::BLUE);
        assert_eq!(bf.color(1, 3), PuyoColor::RED);
        assert_eq!(bf.color(2, 3), PuyoColor::GREEN);
        assert_eq!(bf.color(3, 3), PuyoColor::BLUE);
    }

    #[test]
    fn test_set_color() {
        let colors = [
            PuyoColor::EMPTY, PuyoColor::OJAMA, PuyoColor::WALL, PuyoColor::IRON,
            PuyoColor::RED, PuyoColor::BLUE, PuyoColor::YELLOW, PuyoColor::GREEN,
        ];
        let mut bf = BitField::new();

        for c in colors.iter() {
            bf.set_color(1, 1, *c);
            assert_eq!(*c, bf.color(1, 1));
        }
    }

    #[test]
    fn test_each_cell() {
        let bf = BitField::from_str(concat!(
            "&&&&&&",
            "OOOOOO",
            "YYYYYY",
            "BBBBBB",
            "GGGGGG",
            "RRRRRR"));

        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                for c in color::color::ALL_PUYO_COLORS.iter() {
                    assert_eq!(bf.bits(*c).get(x, y), *c == bf.color(x, y));
                    assert_eq!(bf.is_color(x, y, *c), bf.is_color(x, y, *c));
                }

                assert_eq!(bf.is_normal_color(x, y), bf.is_normal_color(x, y));
            }
        }
    }

    #[test]
    fn test_count_connected_puyos() {
        let bf = BitField::from_str(concat!(
            "RRRRRR",
            "BYBRRY",
            "RRRBBB"));

        assert_eq!(bf.count_connected_puyos(1, 1), 3);
        assert_eq!(bf.count_connected_puyos(4, 1), 3);
        assert_eq!(bf.count_connected_puyos(1, 2), 1);
        assert_eq!(bf.count_connected_puyos(3, 2), 1);
        assert_eq!(bf.count_connected_puyos(6, 2), 1);
        assert_eq!(bf.count_connected_puyos(4, 2), 8);
    }
}
<commit_msg>Add BitField::is_empty<commit_after>use color::PuyoColor;
use field;
use field_bit::FieldBit;
use plain_field::PuyoPlainField;
use std::mem;
use x86intrin::*;

#[derive(Clone, Copy)]
pub struct BitField {
    m: [FieldBit; 3],
}

impl BitField {
    pub fn new() -> BitField {
        BitField {
            m: [FieldBit::new_empty(),
                FieldBit::from_values(0xFFFF, 0x8001, 0x8001, 0x8001, 0x8001, 0x8001, 0x8001, 0xFFFF),
                FieldBit::new_empty(), ]
        }
    }

    pub fn from_plain_field(pf: PuyoPlainField) -> BitField {
        let mut bf = BitField::new();

        // TODO(mayah): We have better algorithm here.
        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                bf.set_color(x, y, pf.color(x, y))
            }
        }

        bf
    }

    pub fn from_str(s: &str) -> BitField {
        BitField::from_plain_field(PuyoPlainField::from_str(s))
    }

    pub fn color(&self, x: usize, y: usize) -> PuyoColor {
        let b0: u8 = if self.m[0].get(x, y) { 1 } else { 0 };
        let b1: u8 = if self.m[1].get(x, y) { 2 } else { 0 };
        let b2: u8 = if self.m[2].get(x, y) { 4 } else { 0 };

        unsafe {
            mem::transmute(b0 | b1 | b2)
        }
    }

    pub fn is_color(&self, x: usize, y: usize, c: PuyoColor) -> bool {
        self.bits(c).get(x, y)
    }

    pub fn is_empty(&self, x: usize, y: usize) -> bool {
        let whole = self.m[0] | self.m[1] | self.m[2];
        return !(whole.get(x, y))
    }

    pub fn is_normal_color(&self, x: usize, y: usize) -> bool {
        self.m[2].get(x, y)
    }

    pub fn set_color(&mut self, x: usize, y: usize, c: PuyoColor) {
        let cc = c as u8;
        for i in 0 .. 3 {
            if (cc & (1 << i)) != 0 {
                self.m[i as usize].set(x, y);
            } else {
                self.m[i as usize].unset(x, y);
            }
        }
    }

    pub fn count_connected_puyos(&self, x: usize, y: usize) -> usize {
        if y > field::HEIGHT {
            return 0
        }

        let c = self.color(x, y);
        let color_bits = self.bits(c).masked_field_12();
        FieldBit::from_onebit(x, y).expand(&color_bits).popcount()
    }

    pub fn bits(&self, c: PuyoColor) -> FieldBit {
        let r0 = self.m[0].as_m128i();
        let r1 = self.m[1].as_m128i();
        let r2 = self.m[2].as_m128i();

        let v = match c {
            PuyoColor::EMPTY => {  // 0
                let x = mm_or_si128(mm_or_si128(r0, r1), r2);
                mm_xor_si128(x, mm_setr_epi32(!0, !0, !0, !0))
            },
            PuyoColor::OJAMA => {  // 1
                mm_andnot_si128(r2, mm_andnot_si128(r1, r0))
            },
            PuyoColor::WALL => {   // 2
                mm_andnot_si128(r2, mm_andnot_si128(r0, r1))
            },
            PuyoColor::IRON => {   // 3
                mm_andnot_si128(r2, mm_and_si128(r0, r1))
            },
            PuyoColor::RED => {    // 4
                mm_andnot_si128(r0, mm_andnot_si128(r1, r2))
            },
            PuyoColor::BLUE => {   // 5
                mm_and_si128(r0, mm_andnot_si128(r1, r2))
            },
            PuyoColor::YELLOW => { // 6
                mm_andnot_si128(r0, mm_and_si128(r1, r2))
            },
            PuyoColor::GREEN => {  // 7
                mm_and_si128(r0, mm_and_si128(r1, r2))
            },
        };

        FieldBit::new(v)
    }
}

#[cfg(test)]
mod tests {
    use super::BitField;
    use color;
    use color::PuyoColor;
    use field;

    #[test]
    fn test_initial() {
        let bf = BitField::new();
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                if x == 0 || x == 7 || y == 0 || y == 15 {
                    assert_eq!(bf.color(x, y), PuyoColor::WALL);
                } else {
                    assert_eq!(bf.color(x, y), PuyoColor::EMPTY);
                }
            }
        }
    }

    #[test]
    fn test_from_str() {
        let bf = BitField::from_str(concat!(
            "RGBRGB",
            "RGBRGB",
            "RGBRGB"));

        assert_eq!(bf.color(1, 1), PuyoColor::RED);
        assert_eq!(bf.color(2, 1), PuyoColor::GREEN);
        assert_eq!(bf.color(3, 1), PuyoColor::BLUE);
        assert_eq!(bf.color(1, 3), PuyoColor::RED);
        assert_eq!(bf.color(2, 3), PuyoColor::GREEN);
        assert_eq!(bf.color(3, 3), PuyoColor::BLUE);
    }

    #[test]
    fn test_set_color() {
        let colors = [
            PuyoColor::EMPTY, PuyoColor::OJAMA, PuyoColor::WALL, PuyoColor::IRON,
            PuyoColor::RED, PuyoColor::BLUE, PuyoColor::YELLOW, PuyoColor::GREEN,
        ];
        let mut bf = BitField::new();

        for c in colors.iter() {
            bf.set_color(1, 1, *c);
            assert_eq!(*c, bf.color(1, 1));
        }
    }

    #[test]
    fn test_is_empty() {
        let bf = BitField::from_str(concat!(
            "RRR..."));

        assert!(!bf.is_empty(1, 1));
        assert!(!bf.is_empty(2, 1));
        assert!(!bf.is_empty(3, 1));
        assert!(bf.is_empty(4, 1));
        assert!(bf.is_empty(5, 1));
        assert!(bf.is_empty(6, 1));
    }

    #[test]
    fn test_each_cell() {
        let bf = BitField::from_str(concat!(
            "&&&&&&",
            "OOOOOO",
            "YYYYYY",
            "BBBBBB",
            "GGGGGG",
            "RRRRRR"));

        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                for c in color::color::ALL_PUYO_COLORS.iter() {
                    assert_eq!(bf.bits(*c).get(x, y), *c == bf.color(x, y));
                    assert_eq!(bf.is_color(x, y, *c), bf.is_color(x, y, *c));
                }

                assert_eq!(bf.is_normal_color(x, y), bf.is_normal_color(x, y));
            }
        }
    }

    #[test]
    fn test_count_connected_puyos() {
        let bf = BitField::from_str(concat!(
            "RRRRRR",
            "BYBRRY",
            "RRRBBB"));

        assert_eq!(bf.count_connected_puyos(1, 1), 3);
        assert_eq!(bf.count_connected_puyos(4, 1), 3);
        assert_eq!(bf.count_connected_puyos(1, 2), 1);
        assert_eq!(bf.count_connected_puyos(3, 2), 1);
        assert_eq!(bf.count_connected_puyos(6, 2), 1);
        assert_eq!(bf.count_connected_puyos(4, 2), 8);
    }
}
<|endoftext|>


<commit_before>use std;
use x86intrin::*;

#[derive(Clone, Copy, Debug)]
pub struct FieldBit {
    m: m128i,
}

#[derive(Clone, Copy, Debug)]
pub struct FieldBit256 {
    m: m256i,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum LowHigh {
    LOW,
    HIGH,
}

impl FieldBit {
    pub fn new(m: m128i) -> FieldBit {
        FieldBit {
            m: m,
        }
    }

    pub fn empty() -> FieldBit {
        FieldBit {
            m: mm_setzero_si128()
        }
    }

    pub fn from_values(v1: u16, v2: u16, v3: u16, v4: u16, v5: u16, v6: u16, v7: u16, v8: u16) -> FieldBit {
        FieldBit {
            m: mm_setr_epi16(v1 as i16, v2 as i16, v3 as i16, v4 as i16,
                             v5 as i16, v6 as i16, v7 as i16, v8 as i16)
        }
    }

    pub fn from_onebit(x: usize, y: usize) -> FieldBit {
        FieldBit {
            m: FieldBit::onebit(x, y)
        }
    }

    pub fn from_str(s: &str) -> FieldBit {
        let mut f = FieldBit::empty();

        assert!(s.len() % 6 == 0);

        let mut cnt = 0;
        for b in s.bytes().rev() {
            let x = 6 - (cnt % 6);
            let y = (cnt / 6) + 1;
            if b == b'1' {
                f.set(x, y);
            }
            cnt += 1;
        }

        f
    }

    pub fn as_m128i(&self) -> m128i {
        self.m
    }

    pub fn get(&self, x: usize, y: usize) -> bool {
        debug_assert!(FieldBit::check_in_range(x, y));
        mm_testz_si128(FieldBit::onebit(x, y), self.m) == 0
    }

    pub fn set(&mut self, x: usize, y: usize) {
        debug_assert!(FieldBit::check_in_range(x, y));
        self.m = mm_or_si128(FieldBit::onebit(x, y), self.m)
    }

    pub fn unset(&mut self, x: usize, y: usize) {
        debug_assert!(FieldBit::check_in_range(x, y));
        self.m = mm_andnot_si128(FieldBit::onebit(x, y), self.m)
    }

    pub fn masked_field_12(&self) -> FieldBit {
        FieldBit {
            m: mm_and_si128(self.m, mm_setr_epi16(0, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0)),
        }
    }

    pub fn masked_field_13(&self) -> FieldBit {
        FieldBit {
            m: mm_and_si128(self.m, mm_setr_epi16(0, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0)),
        }
    }

    pub fn popcount(&self) -> usize {
        let x = self.m.as_u64x2();
        let low: u64 = x.extract(0);
        let high: u64 = x.extract(1);
        (low.count_ones() + high.count_ones()) as usize
    }

    pub fn expand(&self, mask: &FieldBit) -> FieldBit {
        let mut seed = self.m;
        loop {
            let mut expanded = seed;
            expanded = mm_or_si128(mm_slli_epi16(seed, 1), expanded);
            expanded = mm_or_si128(mm_srli_epi16(seed, 1), expanded);
            expanded = mm_or_si128(mm_slli_si128(seed, 2), expanded);
            expanded = mm_or_si128(mm_srli_si128(seed, 2), expanded);
            expanded = mm_and_si128(mask.m, expanded);

            if mm_testc_si128(seed, expanded) != 0 { // seed == expanded
                return FieldBit { m: expanded };
            }
            seed = expanded;
        }
    }

    pub fn expand1(&self, mask: &FieldBit) -> FieldBit {
        let seed = self.m;
        let v1 = mm_slli_epi16(seed, 1);
        let v2 = mm_srli_epi16(seed, 1);
        let v3 = mm_slli_si128(seed, 2);
        let v4 = mm_srli_si128(seed, 2);

        let m = ((seed | v1) | (v2 | v3) | v4) & mask.m;
        FieldBit { m: m }
    }

    fn check_in_range(x: usize, y: usize) -> bool {
        x < 8 && y < 16
    }

    fn onebit(x: usize, y: usize) -> m128i {
        debug_assert!(FieldBit::check_in_range(x, y));

        let shift = ((x << 4) | y) & 0x3F;
        let hi: i64 = (x as i64) >> 2;
        let lo: i64 = hi ^ 1;

        mm_set_epi64x(hi << shift, lo << shift)
    }
}

impl std::ops::BitOr for FieldBit {
    type Output = FieldBit;

    fn bitor(self, rhs: FieldBit) -> FieldBit {
        FieldBit::new(mm_or_si128(self.m, rhs.m))
    }
}

impl std::cmp::PartialEq<FieldBit> for FieldBit {
    fn eq(&self, other: &FieldBit) -> bool {
        let x = mm_xor_si128(self.m, other.m);
        mm_testz_si128(x, x) == 1
    }
}

impl std::fmt::Display for FieldBit {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let x = self.m.as_u16x8();
        write!(f, "({}, {}, {}, {}, {}, {}, {}, {})",
               x.extract(0), x.extract(1), x.extract(2), x.extract(3),
               x.extract(4), x.extract(5), x.extract(6), x.extract(7))
    }
}

impl FieldBit256 {
    pub fn new(m: m256i) -> FieldBit256 {
        FieldBit256 {
            m: m,
        }
    }

    pub fn empty() -> FieldBit256 {
        FieldBit256 {
            m: mm256_setzero_si256()
        }
    }

    pub fn from_low_high(low: FieldBit, high: FieldBit) -> FieldBit256 {
        let m = mm256_inserti128_si256(mm256_castsi128_si256(low.m), high.m, 1);
        FieldBit256 {
            m: m
        }
    }

    pub fn get(&self, lowhigh: LowHigh, x: usize, y: usize) -> bool {
        debug_assert!(FieldBit256::check_in_range(x, y));
        mm256_testz_si256(FieldBit256::onebit(lowhigh, x, y), self.m) == 0
    }

    fn check_in_range(x: usize, y: usize) -> bool {
        x < 8 && y < 16
    }

    fn onebit(lowhigh: LowHigh, x: usize, y: usize) -> m256i {
        debug_assert!(FieldBit256::check_in_range(x, y));

        // TODO(mayah): Maybe we have more good solution.

        let shift = ((x << 4) | y) & 0x3F;
        let zero = mm256_setzero_si256();
        if lowhigh == LowHigh::LOW {
            if x < 4 {
                return mm256_insert_epi64(zero, 1 << shift, 0)
            } else {
                return mm256_insert_epi64(zero, 1 << shift, 1)
            }
        } else {
            if x < 4 {
                return mm256_insert_epi64(zero, 1 << shift, 2)
            } else {
                return mm256_insert_epi64(zero, 1 << shift, 3)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use field;
    use field_bit::FieldBit;

    #[test]
    fn test_empty() {
        let fb = FieldBit::empty();
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(fb.get(x, y), false);
            }
        }
    }

    #[test]
    fn test_from_value() {
        let fb = FieldBit::from_values(1 << 0, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6, 1 << 7);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(fb.get(x, y), x == y, "failed at x={}, y={}, fb.get(x, y)={}", x, y, fb.get(x, y));
            }
        }
    }

    #[test]
    fn test_from_str() {
        let fb = FieldBit::from_str(concat!(
            "111...",
            "......",
            "111..."));

        for x in 0 .. 8 {
            for y in 0 .. 16 {
                let b = (y == 1 || y == 3) && (1 <= x && x <= 3);
                assert_eq!(fb.get(x, y), b, "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_set_get() {
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                let mut fb = FieldBit::empty();
                assert!(!fb.get(x, y));
                fb.set(x, y);
                assert!(fb.get(x, y));
                fb.unset(x, y);
                assert!(!fb.get(x, y));
            }
        }
    }

    #[test]
    fn test_masked_field() {
        let fb = FieldBit::from_values(0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF);
        let fb12 = fb.masked_field_12();
        let fb13 = fb.masked_field_13();

        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                assert!(fb.get(x, y), "x={}, y={}", x, y);
                assert_eq!(fb12.get(x, y), 1 <= x && x <= 6 && 1 <= y && y <= 12, "x={}, y={}", x, y);
                assert_eq!(fb13.get(x, y), 1 <= x && x <= 6 && 1 <= y && y <= 13, "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_popcount() {
        let fb = FieldBit::from_values(1, 2, 3, 4, 5, 6, 7, 8);
        assert_eq!(fb.popcount(), 1 + 1 + 2 + 1 + 2 + 2 + 3 + 1)
    }

    #[test]
    fn test_expand() {
        let mask = FieldBit::from_str(concat!(
            ".1....",
            "1.11..",
            "1.1...",
            "1.1...",
        ));

        let expected = FieldBit::from_str(concat!(
            "..11..",
            "..1...",
            "..1..."));

        let actual = FieldBit::from_onebit(3, 1).expand(&mask);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(actual.get(x, y), expected.get(x, y), "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_expand1() {
        let mask = FieldBit::from_str(concat!(
            "111111",
            "111111",
            "111111"));

        let seed = FieldBit::from_str(concat!(
            "......",
            "1...1.",
            "......"));

        let expected = FieldBit::from_str(concat!(
            "1...1.",
            "11.111",
            "1...1."));

        let actual = seed.expand1(&mask);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(actual.get(x, y), expected.get(x, y), "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_eq() {
        let fb1 = FieldBit::from_str(concat!(
            "1....1",
            "1....1",
            "111111"));
        let fb2 = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "1....1"));

        assert!(fb1 == fb1);
        assert!(fb2 == fb2);
        assert!(fb1 != fb2);
    }

    #[test]
    fn test_bitor() {
        let fb1 = FieldBit::from_str(concat!(
            "1....1",
            "1....1",
            "111111"));
        let fb2 = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "1....1"));
        let expected = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "111111"));

        assert_eq!(expected, fb1 | fb2);
    }
}

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

    #[test]
    fn test_constructor() {
        let fb256 = FieldBit256::empty();

        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert!(!fb256.get(LowHigh::LOW, x, y));
                assert!(!fb256.get(LowHigh::HIGH, x, y));
            }
        }
    }

    #[test]
    fn test_from_low_high() {
        let mut low = FieldBit::empty();
        let mut high = FieldBit::empty();
        low.set(1, 3);
        low.set(4, 8);
        high.set(2, 4);
        high.set(5, 9);

        let fb256 = FieldBit256::from_low_high(low, high);
        assert!(fb256.get(LowHigh::LOW, 1, 3));
        assert!(fb256.get(LowHigh::LOW, 4, 8));
        assert!(fb256.get(LowHigh::HIGH, 2, 4));
        assert!(fb256.get(LowHigh::HIGH, 5, 9));

        assert!(!fb256.get(LowHigh::HIGH, 1, 3));
        assert!(!fb256.get(LowHigh::HIGH, 4, 8));
        assert!(!fb256.get(LowHigh::LOW, 2, 4));
        assert!(!fb256.get(LowHigh::LOW, 5, 9));
    }
}
<commit_msg>Add FieldBit::uninitialized()<commit_after>use std;
use x86intrin::*;

#[derive(Clone, Copy, Debug)]
pub struct FieldBit {
    m: m128i,
}

#[derive(Clone, Copy, Debug)]
pub struct FieldBit256 {
    m: m256i,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum LowHigh {
    LOW,
    HIGH,
}

impl FieldBit {
    pub fn new(m: m128i) -> FieldBit {
        FieldBit {
            m: m,
        }
    }

    pub fn uninitialized() -> FieldBit {
        unsafe { std::mem::uninitialized::<FieldBit>() }
    }

    pub fn empty() -> FieldBit {
        FieldBit {
            m: mm_setzero_si128()
        }
    }

    pub fn from_values(v1: u16, v2: u16, v3: u16, v4: u16, v5: u16, v6: u16, v7: u16, v8: u16) -> FieldBit {
        FieldBit {
            m: mm_setr_epi16(v1 as i16, v2 as i16, v3 as i16, v4 as i16,
                             v5 as i16, v6 as i16, v7 as i16, v8 as i16)
        }
    }

    pub fn from_onebit(x: usize, y: usize) -> FieldBit {
        FieldBit {
            m: FieldBit::onebit(x, y)
        }
    }

    pub fn from_str(s: &str) -> FieldBit {
        let mut f = FieldBit::empty();

        assert!(s.len() % 6 == 0);

        let mut cnt = 0;
        for b in s.bytes().rev() {
            let x = 6 - (cnt % 6);
            let y = (cnt / 6) + 1;
            if b == b'1' {
                f.set(x, y);
            }
            cnt += 1;
        }

        f
    }

    pub fn as_m128i(&self) -> m128i {
        self.m
    }

    pub fn get(&self, x: usize, y: usize) -> bool {
        debug_assert!(FieldBit::check_in_range(x, y));
        mm_testz_si128(FieldBit::onebit(x, y), self.m) == 0
    }

    pub fn set(&mut self, x: usize, y: usize) {
        debug_assert!(FieldBit::check_in_range(x, y));
        self.m = mm_or_si128(FieldBit::onebit(x, y), self.m)
    }

    pub fn unset(&mut self, x: usize, y: usize) {
        debug_assert!(FieldBit::check_in_range(x, y));
        self.m = mm_andnot_si128(FieldBit::onebit(x, y), self.m)
    }

    pub fn masked_field_12(&self) -> FieldBit {
        FieldBit {
            m: mm_and_si128(self.m, mm_setr_epi16(0, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0x1FFE, 0)),
        }
    }

    pub fn masked_field_13(&self) -> FieldBit {
        FieldBit {
            m: mm_and_si128(self.m, mm_setr_epi16(0, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0x3FFE, 0)),
        }
    }

    pub fn popcount(&self) -> usize {
        let x = self.m.as_u64x2();
        let low: u64 = x.extract(0);
        let high: u64 = x.extract(1);
        (low.count_ones() + high.count_ones()) as usize
    }

    pub fn expand(&self, mask: &FieldBit) -> FieldBit {
        let mut seed = self.m;
        loop {
            let mut expanded = seed;
            expanded = mm_or_si128(mm_slli_epi16(seed, 1), expanded);
            expanded = mm_or_si128(mm_srli_epi16(seed, 1), expanded);
            expanded = mm_or_si128(mm_slli_si128(seed, 2), expanded);
            expanded = mm_or_si128(mm_srli_si128(seed, 2), expanded);
            expanded = mm_and_si128(mask.m, expanded);

            if mm_testc_si128(seed, expanded) != 0 { // seed == expanded
                return FieldBit { m: expanded };
            }
            seed = expanded;
        }
    }

    pub fn expand1(&self, mask: &FieldBit) -> FieldBit {
        let seed = self.m;
        let v1 = mm_slli_epi16(seed, 1);
        let v2 = mm_srli_epi16(seed, 1);
        let v3 = mm_slli_si128(seed, 2);
        let v4 = mm_srli_si128(seed, 2);

        let m = ((seed | v1) | (v2 | v3) | v4) & mask.m;
        FieldBit { m: m }
    }

    fn check_in_range(x: usize, y: usize) -> bool {
        x < 8 && y < 16
    }

    fn onebit(x: usize, y: usize) -> m128i {
        debug_assert!(FieldBit::check_in_range(x, y));

        let shift = ((x << 4) | y) & 0x3F;
        let hi: i64 = (x as i64) >> 2;
        let lo: i64 = hi ^ 1;

        mm_set_epi64x(hi << shift, lo << shift)
    }
}

impl std::ops::BitOr for FieldBit {
    type Output = FieldBit;

    fn bitor(self, rhs: FieldBit) -> FieldBit {
        FieldBit::new(mm_or_si128(self.m, rhs.m))
    }
}

impl std::cmp::PartialEq<FieldBit> for FieldBit {
    fn eq(&self, other: &FieldBit) -> bool {
        let x = mm_xor_si128(self.m, other.m);
        mm_testz_si128(x, x) == 1
    }
}

impl std::fmt::Display for FieldBit {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let x = self.m.as_u16x8();
        write!(f, "({}, {}, {}, {}, {}, {}, {}, {})",
               x.extract(0), x.extract(1), x.extract(2), x.extract(3),
               x.extract(4), x.extract(5), x.extract(6), x.extract(7))
    }
}

impl FieldBit256 {
    pub fn new(m: m256i) -> FieldBit256 {
        FieldBit256 {
            m: m,
        }
    }

    pub fn uninitialized() -> FieldBit256 {
        unsafe { std::mem::uninitialized::<FieldBit256>() }
    }

    pub fn empty() -> FieldBit256 {
        FieldBit256 {
            m: mm256_setzero_si256()
        }
    }

    pub fn from_low_high(low: FieldBit, high: FieldBit) -> FieldBit256 {
        let m = mm256_inserti128_si256(mm256_castsi128_si256(low.m), high.m, 1);
        FieldBit256 {
            m: m
        }
    }

    pub fn get(&self, lowhigh: LowHigh, x: usize, y: usize) -> bool {
        debug_assert!(FieldBit256::check_in_range(x, y));
        mm256_testz_si256(FieldBit256::onebit(lowhigh, x, y), self.m) == 0
    }

    fn check_in_range(x: usize, y: usize) -> bool {
        x < 8 && y < 16
    }

    fn onebit(lowhigh: LowHigh, x: usize, y: usize) -> m256i {
        debug_assert!(FieldBit256::check_in_range(x, y));

        // TODO(mayah): Maybe we have more good solution.

        let shift = ((x << 4) | y) & 0x3F;
        let zero = mm256_setzero_si256();
        if lowhigh == LowHigh::LOW {
            if x < 4 {
                return mm256_insert_epi64(zero, 1 << shift, 0)
            } else {
                return mm256_insert_epi64(zero, 1 << shift, 1)
            }
        } else {
            if x < 4 {
                return mm256_insert_epi64(zero, 1 << shift, 2)
            } else {
                return mm256_insert_epi64(zero, 1 << shift, 3)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use field;
    use field_bit::FieldBit;

    #[test]
    fn test_empty() {
        let fb = FieldBit::empty();
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(fb.get(x, y), false);
            }
        }
    }

    #[test]
    fn test_from_value() {
        let fb = FieldBit::from_values(1 << 0, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6, 1 << 7);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(fb.get(x, y), x == y, "failed at x={}, y={}, fb.get(x, y)={}", x, y, fb.get(x, y));
            }
        }
    }

    #[test]
    fn test_from_str() {
        let fb = FieldBit::from_str(concat!(
            "111...",
            "......",
            "111..."));

        for x in 0 .. 8 {
            for y in 0 .. 16 {
                let b = (y == 1 || y == 3) && (1 <= x && x <= 3);
                assert_eq!(fb.get(x, y), b, "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_set_get() {
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                let mut fb = FieldBit::empty();
                assert!(!fb.get(x, y));
                fb.set(x, y);
                assert!(fb.get(x, y));
                fb.unset(x, y);
                assert!(!fb.get(x, y));
            }
        }
    }

    #[test]
    fn test_masked_field() {
        let fb = FieldBit::from_values(0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF);
        let fb12 = fb.masked_field_12();
        let fb13 = fb.masked_field_13();

        for x in 0 .. field::MAP_WIDTH {
            for y in 0 .. field::MAP_HEIGHT {
                assert!(fb.get(x, y), "x={}, y={}", x, y);
                assert_eq!(fb12.get(x, y), 1 <= x && x <= 6 && 1 <= y && y <= 12, "x={}, y={}", x, y);
                assert_eq!(fb13.get(x, y), 1 <= x && x <= 6 && 1 <= y && y <= 13, "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_popcount() {
        let fb = FieldBit::from_values(1, 2, 3, 4, 5, 6, 7, 8);
        assert_eq!(fb.popcount(), 1 + 1 + 2 + 1 + 2 + 2 + 3 + 1)
    }

    #[test]
    fn test_expand() {
        let mask = FieldBit::from_str(concat!(
            ".1....",
            "1.11..",
            "1.1...",
            "1.1...",
        ));

        let expected = FieldBit::from_str(concat!(
            "..11..",
            "..1...",
            "..1..."));

        let actual = FieldBit::from_onebit(3, 1).expand(&mask);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(actual.get(x, y), expected.get(x, y), "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_expand1() {
        let mask = FieldBit::from_str(concat!(
            "111111",
            "111111",
            "111111"));

        let seed = FieldBit::from_str(concat!(
            "......",
            "1...1.",
            "......"));

        let expected = FieldBit::from_str(concat!(
            "1...1.",
            "11.111",
            "1...1."));

        let actual = seed.expand1(&mask);
        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert_eq!(actual.get(x, y), expected.get(x, y), "x={}, y={}", x, y);
            }
        }
    }

    #[test]
    fn test_eq() {
        let fb1 = FieldBit::from_str(concat!(
            "1....1",
            "1....1",
            "111111"));
        let fb2 = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "1....1"));

        assert!(fb1 == fb1);
        assert!(fb2 == fb2);
        assert!(fb1 != fb2);
    }

    #[test]
    fn test_bitor() {
        let fb1 = FieldBit::from_str(concat!(
            "1....1",
            "1....1",
            "111111"));
        let fb2 = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "1....1"));
        let expected = FieldBit::from_str(concat!(
            "111111",
            "1....1",
            "111111"));

        assert_eq!(expected, fb1 | fb2);
    }
}

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

    #[test]
    fn test_constructor() {
        let fb256 = FieldBit256::empty();

        for x in 0 .. 8 {
            for y in 0 .. 16 {
                assert!(!fb256.get(LowHigh::LOW, x, y));
                assert!(!fb256.get(LowHigh::HIGH, x, y));
            }
        }
    }

    #[test]
    fn test_from_low_high() {
        let mut low = FieldBit::empty();
        let mut high = FieldBit::empty();
        low.set(1, 3);
        low.set(4, 8);
        high.set(2, 4);
        high.set(5, 9);

        let fb256 = FieldBit256::from_low_high(low, high);
        assert!(fb256.get(LowHigh::LOW, 1, 3));
        assert!(fb256.get(LowHigh::LOW, 4, 8));
        assert!(fb256.get(LowHigh::HIGH, 2, 4));
        assert!(fb256.get(LowHigh::HIGH, 5, 9));

        assert!(!fb256.get(LowHigh::HIGH, 1, 3));
        assert!(!fb256.get(LowHigh::HIGH, 4, 8));
        assert!(!fb256.get(LowHigh::LOW, 2, 4));
        assert!(!fb256.get(LowHigh::LOW, 5, 9));
    }
}
<|endoftext|>


<commit_before>use uuid::Uuid;

use common::err::{Error, TResult};
use common::types::Type;

// let ctx = TokamakContext::new();
// ctx.from(RandomGenerator).select(...);

use super::TokamakContext;

pub struct DataFrame<'a> {
  pub ctx : &'a TokamakContext,
  pub decl: DataFrameDecl<'a>
}

pub enum DataFrameDecl<'a> {
  From(Box<DataSet>),
  Select(Box<DataFrame<'a>>, Vec<Expr>)
}

pub enum Expr {
  Plus,
  Field(String)
}

impl<'a> DataFrame<'a> {
  pub fn kind(&self) -> &'static str {
    match self.decl {
      DataFrameDecl::From  (_) => "from",
      DataFrameDecl::Select(_,_) => "select"
    }
  }

  pub fn select(self, exprs: Vec<Expr>) -> DataFrame<'a> {
    DataFrame {ctx: self.ctx, decl: DataFrameDecl::Select(Box::new(self), exprs)}
  }

  fn count(&self) -> TResult<usize> {
    Ok(0)
  }
}

pub trait DataSet {
  fn name(&self) -> &str;

  //fn schema(&self) -> &Vec<Box<Type>>;
}

pub struct CustomDataSource {
  name     : String,
  src_type : String,
  schema   : Vec<String>,
  props    : Vec<(String, String)>
}

impl DataSet for CustomDataSource {
  fn name(&self) -> &str {
    &self.name
  }
}

pub fn RandomGenerator(types: Vec<&str>) -> Box<DataSet>
{
  Box::new(CustomDataSource {
    name    :  Uuid::new_v4().to_hyphenated_string(),
    src_type: "random".to_string(),
    schema  : types.into_iter().map(|s| s.to_string()).collect::<Vec<String>>(),
    props   : Vec::new()
  })
}

fn typestr_to_schema(ctx: &TokamakContext, types: Vec<&str>) -> Vec<Box<Type>>
{
  types.iter()
    .map( |s| ctx.get_type(s).unwrap().clone_box() )
    .collect::<Vec<Box<Type>>>()
}<commit_msg>Add head, tail APIs.<commit_after>use uuid::Uuid;

use common::err::{Error, TResult};
use common::types::Type;

// let ctx = TokamakContext::new();
// ctx.from(RandomGenerator).select(...);

use super::TokamakContext;

pub struct DataFrame<'a> {
  pub ctx : &'a TokamakContext,
  pub decl: DataFrameDecl<'a>
}

pub enum DataFrameDecl<'a> {
  From(Box<DataSet>),
  Select(Box<DataFrame<'a>>, Vec<Expr>)
}

pub enum Expr {
  Plus,
  Field(String)
}

impl<'a> DataFrame<'a> {
  pub fn kind(&self) -> &'static str {
    match self.decl {
      DataFrameDecl::From  (_) => "from",
      DataFrameDecl::Select(_,_) => "select"
    }
  }

  pub fn select(self, exprs: Vec<Expr>) -> DataFrame<'a> {
    DataFrame {ctx: self.ctx, decl: DataFrameDecl::Select(Box::new(self), exprs)}
  }

  fn count(&self) -> TResult<usize> {
    Ok(0)
  }

  fn head(&self) -> TResult<Box<DataSet>> {
    Err(Error::InternalError)
  }

  fn head_with(&self, num: usize) -> TResult<Box<DataSet>> {
    Err(Error::InternalError)
  }

  fn tail_with(&self, num: usize) -> TResult<Box<DataSet>> {
    Err(Error::InternalError)
  }
}

pub trait DataSet {
  fn name(&self) -> &str;

  //fn schema(&self) -> &Vec<Box<Type>>;
}

pub struct CustomDataSource {
  name     : String,
  src_type : String,
  schema   : Vec<String>,
  props    : Vec<(String, String)>
}

impl DataSet for CustomDataSource {
  fn name(&self) -> &str {
    &self.name
  }
}

pub fn RandomGenerator(types: Vec<&str>) -> Box<DataSet>
{
  Box::new(CustomDataSource {
    name    :  Uuid::new_v4().to_hyphenated_string(),
    src_type: "random".to_string(),
    schema  : types.into_iter().map(|s| s.to_string()).collect::<Vec<String>>(),
    props   : Vec::new()
  })
}

fn typestr_to_schema(ctx: &TokamakContext, types: Vec<&str>) -> Vec<Box<Type>>
{
  types.iter()
    .map( |s| ctx.get_type(s).unwrap().clone_box() )
    .collect::<Vec<Box<Type>>>()
}<|endoftext|>


<commit_before>//! Runtime support needed for testing the stdsimd crate.
//!
//! This basically just disassembles the current executable and then parses the
//! output once globally and then provides the `assert` function which makes
//! assertions about the disassembly of a function.

#![cfg_attr(
    feature = "cargo-clippy",
    allow(clippy::missing_docs_in_private_items, clippy::print_stdout)
)]

extern crate assert_instr_macro;
extern crate backtrace;
extern crate cc;
#[macro_use]
extern crate lazy_static;
extern crate rustc_demangle;
extern crate simd_test_macro;
#[macro_use]
extern crate cfg_if;

pub use assert_instr_macro::*;
pub use simd_test_macro::*;
use std::{collections::HashMap, env, str};

// `println!` doesn't work on wasm32 right now, so shadow the compiler's `println!`
// macro with our own shim that redirects to `console.log`.
#[allow(unused)]
#[cfg(target_arch = "wasm32")]
#[macro_export]
macro_rules! println {
    ($($args:tt)*) => (wasm::js_console_log(&format!($($args)*)))
}

cfg_if! {
    if #[cfg(target_arch = "wasm32")] {
        extern crate wasm_bindgen;
        extern crate console_error_panic_hook;
        pub mod wasm;
        use wasm::disassemble_myself;
    } else {
        mod disassembly;
        use disassembly::disassemble_myself;
    }
}

lazy_static! {
    static ref DISASSEMBLY: HashMap<String, Vec<Function>> = disassemble_myself();
}

struct Function {
    addr: Option<usize>,
    instrs: Vec<Instruction>,
}

struct Instruction {
    parts: Vec<String>,
}

fn normalize(symbol: &str) -> String {
    let symbol = rustc_demangle::demangle(symbol).to_string();
    let mut ret = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };
    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while ret.starts_with("_") {
        ret.remove(0);
    }
    return ret;
}

/// Main entry point for this crate, called by the `#[assert_instr]` macro.
///
/// This asserts that the function at `fnptr` contains the instruction
/// `expected` provided.
pub fn assert(fnptr: usize, fnname: &str, expected: &str) {
    let mut fnname = fnname.to_string();
    let functions = get_functions(fnptr, &mut fnname);
    assert_eq!(functions.len(), 1);
    let function = &functions[0];

    let mut instrs = &function.instrs[..];
    while instrs.last().map_or(false, |s| s.parts == ["nop"]) {
        instrs = &instrs[..instrs.len() - 1];
    }

    // Look for `expected` as the first part of any instruction in this
    // function, returning if we do indeed find it.
    let mut found = false;
    for instr in instrs {
        // Get the first instruction, e.g., tzcntl in tzcntl %rax,%rax.
        if let Some(part) = instr.parts.get(0) {
            // Truncate the instruction with the length of the expected
            // instruction: tzcntl => tzcnt and compares that.
            if part.starts_with(expected) {
                found = true;
                break;
            }
        }
    }

    // Look for `call` instructions in the disassembly to detect whether
    // inlining failed: all intrinsics are `#[inline(always)]`, so
    // calling one intrinsic from another should not generate `call`
    // instructions.
    let mut inlining_failed = false;
    for (i, instr) in instrs.iter().enumerate() {
        let part = match instr.parts.get(0) {
            Some(part) => part,
            None => continue,
        };
        if !part.contains("call") {
            continue;
        }

        // On 32-bit x86 position independent code will call itself and be
        // immediately followed by a `pop` to learn about the current address.
        // Let's not take that into account when considering whether a function
        // failed inlining something.
        let followed_by_pop = function
            .instrs
            .get(i + 1)
            .and_then(|i| i.parts.get(0))
            .map_or(false, |s| s.contains("pop"));
        if followed_by_pop && cfg!(target_arch = "x86") {
            continue;
        }

        inlining_failed = true;
        break;
    }

    let instruction_limit = std::env::var("STDSIMD_ASSERT_INSTR_LIMIT")
        .ok()
        .map_or_else(
            || match expected {
                // `cpuid` returns a pretty big aggregate structure, so exempt
                // it from the slightly more restrictive 22 instructions below.
                "cpuid" => 30,

                // Apparently, on Windows, LLVM generates a bunch of
                // saves/restores of xmm registers around these intstructions,
                // which exceeds the limit of 20 below. As it seems dictated by
                // Windows's ABI (I believe?), we probably can't do much
                // about it.
                "vzeroall" | "vzeroupper" if cfg!(windows) => 30,

                // Intrinsics using `cvtpi2ps` are typically "composites" and
                // in some cases exceed the limit.
                "cvtpi2ps" => 25,

                // Original limit was 20 instructions, but ARM DSP Intrinsics
                // are exactly 20 instructions long. So, bump the limit to 22
                // instead of adding here a long list of exceptions.
                _ => 22,
            },
            |v| v.parse().unwrap(),
        );
    let probably_only_one_instruction = instrs.len() < instruction_limit;

    if found && probably_only_one_instruction && !inlining_failed {
        return;
    }

    // Help debug by printing out the found disassembly, and then panic as we
    // didn't find the instruction.
    println!("disassembly for {}: ", fnname,);
    for (i, instr) in instrs.iter().enumerate() {
        let mut s = format!("\t{:2}: ", i);
        for part in &instr.parts {
            s.push_str(part);
            s.push_str(" ");
        }
        println!("{}", s);
    }

    if !found {
        panic!(
            "failed to find instruction `{}` in the disassembly",
            expected
        );
    } else if !probably_only_one_instruction {
        panic!(
            "instruction found, but the disassembly contains too many \
             instructions: #instructions = {} >= {} (limit)",
            instrs.len(),
            instruction_limit
        );
    } else if inlining_failed {
        panic!(
            "instruction found, but the disassembly contains `call` \
             instructions, which hint that inlining failed"
        );
    }
}

fn get_functions(fnptr: usize, fnname: &mut String) -> &'static [Function] {
    // Translate this function pointer to a symbolic name that we'd have found
    // in the disassembly.
    let mut sym = None;
    backtrace::resolve(fnptr as *mut _, |name| {
        sym = name.name().and_then(|s| s.as_str()).map(normalize);
    });

    if let Some(sym) = &sym {
        if let Some(s) = DISASSEMBLY.get(sym) {
            *fnname = sym.to_string();
            return s;
        }
    }

    let exact_match = DISASSEMBLY
        .iter()
        .find(|(_, list)| list.iter().any(|f| f.addr == Some(fnptr)));
    if let Some((name, list)) = exact_match {
        *fnname = name.to_string();
        return list;
    }

    if let Some(sym) = sym {
        println!("assumed symbol name: `{}`", sym);
    }
    println!("maybe related functions");
    for f in DISASSEMBLY.keys().filter(|k| k.contains(&**fnname)) {
        println!("\t- {}", f);
    }
    panic!("failed to find disassembly of {:#x} ({})", fnptr, fnname);
}

pub fn assert_skip_test_ok(name: &str) {
    if env::var("STDSIMD_TEST_EVERYTHING").is_err() {
        return;
    }
    panic!("skipped test `{}` when it shouldn't be skipped", name);
}

// See comment in `assert-instr-macro` crate for why this exists
pub static mut _DONT_DEDUP: &'static str = "";
<commit_msg>assert_instr: bump instruction limit for simd32<commit_after>//! Runtime support needed for testing the stdsimd crate.
//!
//! This basically just disassembles the current executable and then parses the
//! output once globally and then provides the `assert` function which makes
//! assertions about the disassembly of a function.

#![cfg_attr(
    feature = "cargo-clippy",
    allow(clippy::missing_docs_in_private_items, clippy::print_stdout)
)]

extern crate assert_instr_macro;
extern crate backtrace;
extern crate cc;
#[macro_use]
extern crate lazy_static;
extern crate rustc_demangle;
extern crate simd_test_macro;
#[macro_use]
extern crate cfg_if;

pub use assert_instr_macro::*;
pub use simd_test_macro::*;
use std::{collections::HashMap, env, str};

// `println!` doesn't work on wasm32 right now, so shadow the compiler's `println!`
// macro with our own shim that redirects to `console.log`.
#[allow(unused)]
#[cfg(target_arch = "wasm32")]
#[macro_export]
macro_rules! println {
    ($($args:tt)*) => (wasm::js_console_log(&format!($($args)*)))
}

cfg_if! {
    if #[cfg(target_arch = "wasm32")] {
        extern crate wasm_bindgen;
        extern crate console_error_panic_hook;
        pub mod wasm;
        use wasm::disassemble_myself;
    } else {
        mod disassembly;
        use disassembly::disassemble_myself;
    }
}

lazy_static! {
    static ref DISASSEMBLY: HashMap<String, Vec<Function>> = disassemble_myself();
}

struct Function {
    addr: Option<usize>,
    instrs: Vec<Instruction>,
}

struct Instruction {
    parts: Vec<String>,
}

fn normalize(symbol: &str) -> String {
    let symbol = rustc_demangle::demangle(symbol).to_string();
    let mut ret = match symbol.rfind("::h") {
        Some(i) => symbol[..i].to_string(),
        None => symbol.to_string(),
    };
    // Normalize to no leading underscore to handle platforms that may
    // inject extra ones in symbol names.
    while ret.starts_with("_") {
        ret.remove(0);
    }
    return ret;
}

/// Main entry point for this crate, called by the `#[assert_instr]` macro.
///
/// This asserts that the function at `fnptr` contains the instruction
/// `expected` provided.
pub fn assert(fnptr: usize, fnname: &str, expected: &str) {
    let mut fnname = fnname.to_string();
    let functions = get_functions(fnptr, &mut fnname);
    assert_eq!(functions.len(), 1);
    let function = &functions[0];

    let mut instrs = &function.instrs[..];
    while instrs.last().map_or(false, |s| s.parts == ["nop"]) {
        instrs = &instrs[..instrs.len() - 1];
    }

    // Look for `expected` as the first part of any instruction in this
    // function, returning if we do indeed find it.
    let mut found = false;
    for instr in instrs {
        // Get the first instruction, e.g., tzcntl in tzcntl %rax,%rax.
        if let Some(part) = instr.parts.get(0) {
            // Truncate the instruction with the length of the expected
            // instruction: tzcntl => tzcnt and compares that.
            if part.starts_with(expected) {
                found = true;
                break;
            }
        }
    }

    // Look for `call` instructions in the disassembly to detect whether
    // inlining failed: all intrinsics are `#[inline(always)]`, so
    // calling one intrinsic from another should not generate `call`
    // instructions.
    let mut inlining_failed = false;
    for (i, instr) in instrs.iter().enumerate() {
        let part = match instr.parts.get(0) {
            Some(part) => part,
            None => continue,
        };
        if !part.contains("call") {
            continue;
        }

        // On 32-bit x86 position independent code will call itself and be
        // immediately followed by a `pop` to learn about the current address.
        // Let's not take that into account when considering whether a function
        // failed inlining something.
        let followed_by_pop = function
            .instrs
            .get(i + 1)
            .and_then(|i| i.parts.get(0))
            .map_or(false, |s| s.contains("pop"));
        if followed_by_pop && cfg!(target_arch = "x86") {
            continue;
        }

        inlining_failed = true;
        break;
    }

    let instruction_limit = std::env::var("STDSIMD_ASSERT_INSTR_LIMIT")
        .ok()
        .map_or_else(
            || match expected {
                // `cpuid` returns a pretty big aggregate structure, so exempt
                // it from the slightly more restrictive 22 instructions below.
                "cpuid" => 30,

                // Apparently, on Windows, LLVM generates a bunch of
                // saves/restores of xmm registers around these intstructions,
                // which exceeds the limit of 20 below. As it seems dictated by
                // Windows's ABI (I believe?), we probably can't do much
                // about it.
                "vzeroall" | "vzeroupper" if cfg!(windows) => 30,

                // Intrinsics using `cvtpi2ps` are typically "composites" and
                // in some cases exceed the limit.
                "cvtpi2ps" => 25,

                // core_arch/src/acle/simd32
                "usad8" => 27,
                "qadd8" | "qsub8" | "sadd8" | "sel" | "shadd8" | "shsub8" => 29,

                // Original limit was 20 instructions, but ARM DSP Intrinsics
                // are exactly 20 instructions long. So, bump the limit to 22
                // instead of adding here a long list of exceptions.
                _ => 22,
            },
            |v| v.parse().unwrap(),
        );
    let probably_only_one_instruction = instrs.len() < instruction_limit;

    if found && probably_only_one_instruction && !inlining_failed {
        return;
    }

    // Help debug by printing out the found disassembly, and then panic as we
    // didn't find the instruction.
    println!("disassembly for {}: ", fnname,);
    for (i, instr) in instrs.iter().enumerate() {
        let mut s = format!("\t{:2}: ", i);
        for part in &instr.parts {
            s.push_str(part);
            s.push_str(" ");
        }
        println!("{}", s);
    }

    if !found {
        panic!(
            "failed to find instruction `{}` in the disassembly",
            expected
        );
    } else if !probably_only_one_instruction {
        panic!(
            "instruction found, but the disassembly contains too many \
             instructions: #instructions = {} >= {} (limit)",
            instrs.len(),
            instruction_limit
        );
    } else if inlining_failed {
        panic!(
            "instruction found, but the disassembly contains `call` \
             instructions, which hint that inlining failed"
        );
    }
}

fn get_functions(fnptr: usize, fnname: &mut String) -> &'static [Function] {
    // Translate this function pointer to a symbolic name that we'd have found
    // in the disassembly.
    let mut sym = None;
    backtrace::resolve(fnptr as *mut _, |name| {
        sym = name.name().and_then(|s| s.as_str()).map(normalize);
    });

    if let Some(sym) = &sym {
        if let Some(s) = DISASSEMBLY.get(sym) {
            *fnname = sym.to_string();
            return s;
        }
    }

    let exact_match = DISASSEMBLY
        .iter()
        .find(|(_, list)| list.iter().any(|f| f.addr == Some(fnptr)));
    if let Some((name, list)) = exact_match {
        *fnname = name.to_string();
        return list;
    }

    if let Some(sym) = sym {
        println!("assumed symbol name: `{}`", sym);
    }
    println!("maybe related functions");
    for f in DISASSEMBLY.keys().filter(|k| k.contains(&**fnname)) {
        println!("\t- {}", f);
    }
    panic!("failed to find disassembly of {:#x} ({})", fnptr, fnname);
}

pub fn assert_skip_test_ok(name: &str) {
    if env::var("STDSIMD_TEST_EVERYTHING").is_err() {
        return;
    }
    panic!("skipped test `{}` when it shouldn't be skipped", name);
}

// See comment in `assert-instr-macro` crate for why this exists
pub static mut _DONT_DEDUP: &'static str = "";
<|endoftext|>


<commit_before>use core::cell::Cell;
use core::fmt;
use core::sync::atomic;

const SPIN_LIMIT: u32 = 6;
const YIELD_LIMIT: u32 = 10;

/// A counter for exponential backoff in spin loops.
///
/// Backing off in spin loops reduces contention and improves overall performance.
///
/// This primitive can execute "YIELD" and "PAUSE" instructions, yield the current thread to the OS
/// scheduler, and tell when is a good time to block the thread using a different synchronization
/// mechanism. Each step of the back off procedure takes roughly twice as long as the previous
/// step.
///
/// # Examples
///
/// Backing off in a lock-free loop:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicUsize;
/// use std::sync::atomic::Ordering::SeqCst;
///
/// fn fetch_mul(a: &AtomicUsize, b: usize) -> usize {
///     let backoff = Backoff::new();
///     loop {
///         let val = a.load(SeqCst);
///         if a.compare_and_swap(val, val.wrapping_mul(b), SeqCst) == val {
///             return val;
///         }
///         backoff.spin();
///     }
/// }
/// ```
///
/// Waiting for an [`AtomicBool`] to become `true`:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::SeqCst;
///
/// fn spin_wait(ready: &AtomicBool) {
///     let backoff = Backoff::new();
///     while !ready.load(SeqCst) {
///         backoff.snooze();
///     }
/// }
/// ```
///
/// Waiting for an [`AtomicBool`] to become `true` and parking the thread after a long wait.
/// Note that whoever sets the atomic variable to `true` must notify the parked thread by calling
/// [`unpark()`]:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::SeqCst;
/// use std::thread;
///
/// fn blocking_wait(ready: &AtomicBool) {
///     let backoff = Backoff::new();
///     while !ready.load(SeqCst) {
///         if backoff.is_complete() {
///             thread::park();
///         } else {
///             backoff.snooze();
///         }
///     }
/// }
/// ```
///
/// [`is_complete`]: struct.Backoff.html#method.is_complete
/// [`std::thread::park()`]: https://doc.rust-lang.org/std/thread/fn.park.html
/// [`Condvar`]: https://doc.rust-lang.org/std/sync/struct.Condvar.html
/// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
/// [`unpark()`]: https://doc.rust-lang.org/std/thread/struct.Thread.html#method.unpark
pub struct Backoff {
    step: Cell<u32>,
}

impl Backoff {
    /// Creates a new `Backoff`.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    ///
    /// let backoff = Backoff::new();
    /// ```
    #[inline]
    pub fn new() -> Self {
        Backoff {
            step: Cell::new(0),
        }
    }

    /// Resets the `Backoff`.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    ///
    /// let backoff = Backoff::new();
    /// backoff.reset();
    /// ```
    #[inline]
    pub fn reset(&self) {
        self.step.set(0);
    }

    /// Backs off in a lock-free loop.
    ///
    /// This method should be used when we need to retry an operation because another thread made
    /// progress.
    ///
    /// The processor may yield using the "YIELD" or "PAUSE" instruction.
    ///
    /// # Examples
    ///
    /// Backing off in a lock-free loop:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::atomic::AtomicUsize;
    /// use std::sync::atomic::Ordering::SeqCst;
    ///
    /// fn fetch_mul(a: &AtomicUsize, b: usize) -> usize {
    ///     let backoff = Backoff::new();
    ///     loop {
    ///         let val = a.load(SeqCst);
    ///         if a.compare_and_swap(val, val.wrapping_mul(b), SeqCst) == val {
    ///             return val;
    ///         }
    ///         backoff.spin();
    ///     }
    /// }
    ///
    /// let a = AtomicUsize::new(7);
    /// assert_eq!(fetch_mul(&a, 8), 7);
    /// assert_eq!(a.load(SeqCst), 56);
    /// ```
    #[inline]
    pub fn spin(&self) {
        for _ in 0..1 << self.step.get().min(SPIN_LIMIT) {
            atomic::spin_loop_hint();
        }

        if self.step.get() <= SPIN_LIMIT {
            self.step.set(self.step.get() + 1);
        }
    }

    /// Backs off in a blocking loop.
    ///
    /// This method should be used when we need to wait for another thread to make progress.
    ///
    /// The processor may yield using the "YIELD" or "PAUSE" instruction and the current thread
    /// may yield by giving up a timeslice to the OS scheduler.
    ///
    /// In `#[no_std]` environments, this method is equivalent to [`spin`].
    ///
    /// If possible, use [`is_complete`] to check when it is advised to stop using backoff and
    /// block the current thread using a different synchronization mechanism instead.
    ///
    /// [`spin`]: struct.Backoff.html#method.spin
    /// [`is_complete`]: struct.Backoff.html#method.is_complete
    ///
    /// # Examples
    ///
    /// Waiting for an [`AtomicBool`] to become `true`:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::Arc;
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::SeqCst;
    /// use std::thread;
    /// use std::time::Duration;
    ///
    /// fn spin_wait(ready: &AtomicBool) {
    ///     let backoff = Backoff::new();
    ///     while !ready.load(SeqCst) {
    ///         backoff.snooze();
    ///     }
    /// }
    ///
    /// let ready = Arc::new(AtomicBool::new(false));
    /// let ready2 = ready.clone();
    ///
    /// thread::spawn(move || {
    ///     thread::sleep(Duration::from_millis(100));
    ///     ready2.store(true, SeqCst);
    /// });
    ///
    /// assert_eq!(ready.load(SeqCst), false);
    /// spin_wait(&ready);
    /// assert_eq!(ready.load(SeqCst), true);
    /// ```
    ///
    /// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
    #[inline]
    pub fn snooze(&self) {
        if self.step.get() <= SPIN_LIMIT {
            for _ in 0..1 << self.step.get() {
                atomic::spin_loop_hint();
            }
        } else {
            #[cfg(not(feature = "std"))]
            for _ in 0..1 << self.step.get() {
                atomic::spin_loop_hint();
            }

            #[cfg(feature = "std")]
            ::std::thread::yield_now();
        }

        if self.step.get() <= YIELD_LIMIT {
            self.step.set(self.step.get() + 1);
        }
    }

    /// Returns `true` if exponential backoff is complete and blocking the thread is advised.
    ///
    /// # Examples
    ///
    /// Waiting for an [`AtomicBool`] to become `true` and parking the thread after a long wait:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::Arc;
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::SeqCst;
    /// use std::thread;
    /// use std::time::Duration;
    ///
    /// fn blocking_wait(ready: &AtomicBool) {
    ///     let backoff = Backoff::new();
    ///     while !ready.load(SeqCst) {
    ///         if backoff.is_complete() {
    ///             thread::park();
    ///         } else {
    ///             backoff.snooze();
    ///         }
    ///     }
    /// }
    ///
    /// let ready = Arc::new(AtomicBool::new(false));
    /// let ready2 = ready.clone();
    /// let waiter = thread::current();
    ///
    /// thread::spawn(move || {
    ///     thread::sleep(Duration::from_millis(100));
    ///     ready2.store(true, SeqCst);
    ///     waiter.unpark();
    /// });
    ///
    /// assert_eq!(ready.load(SeqCst), false);
    /// blocking_wait(&ready);
    /// assert_eq!(ready.load(SeqCst), true);
    /// ```
    ///
    /// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
    #[inline]
    pub fn is_complete(&self) -> bool {
        self.step.get() > YIELD_LIMIT
    }
}

impl fmt::Debug for Backoff {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Backoff")
            .field("step", &self.step)
            .field("is_complete", &self.is_complete())
            .finish()
    }
}

impl Default for Backoff {
    fn default() -> Backoff {
        Backoff::new()
    }
}
<commit_msg>Quotes to italics<commit_after>use core::cell::Cell;
use core::fmt;
use core::sync::atomic;

const SPIN_LIMIT: u32 = 6;
const YIELD_LIMIT: u32 = 10;

/// Performs exponential backoff in spin loops.
///
/// Backing off in spin loops reduces contention and improves overall performance.
///
/// This primitive can execute *YIELD* and *PAUSE* instructions, yield the current thread to the OS
/// scheduler, and tell when is a good time to block the thread using a different synchronization
/// mechanism. Each step of the back off procedure takes roughly twice as long as the previous
/// step.
///
/// # Examples
///
/// Backing off in a lock-free loop:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicUsize;
/// use std::sync::atomic::Ordering::SeqCst;
///
/// fn fetch_mul(a: &AtomicUsize, b: usize) -> usize {
///     let backoff = Backoff::new();
///     loop {
///         let val = a.load(SeqCst);
///         if a.compare_and_swap(val, val.wrapping_mul(b), SeqCst) == val {
///             return val;
///         }
///         backoff.spin();
///     }
/// }
/// ```
///
/// Waiting for an [`AtomicBool`] to become `true`:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::SeqCst;
///
/// fn spin_wait(ready: &AtomicBool) {
///     let backoff = Backoff::new();
///     while !ready.load(SeqCst) {
///         backoff.snooze();
///     }
/// }
/// ```
///
/// Waiting for an [`AtomicBool`] to become `true` and parking the thread after a long wait.
/// Note that whoever sets the atomic variable to `true` must notify the parked thread by calling
/// [`unpark()`]:
///
/// ```
/// use crossbeam_utils::Backoff;
/// use std::sync::atomic::AtomicBool;
/// use std::sync::atomic::Ordering::SeqCst;
/// use std::thread;
///
/// fn blocking_wait(ready: &AtomicBool) {
///     let backoff = Backoff::new();
///     while !ready.load(SeqCst) {
///         if backoff.is_complete() {
///             thread::park();
///         } else {
///             backoff.snooze();
///         }
///     }
/// }
/// ```
///
/// [`is_complete`]: struct.Backoff.html#method.is_complete
/// [`std::thread::park()`]: https://doc.rust-lang.org/std/thread/fn.park.html
/// [`Condvar`]: https://doc.rust-lang.org/std/sync/struct.Condvar.html
/// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
/// [`unpark()`]: https://doc.rust-lang.org/std/thread/struct.Thread.html#method.unpark
pub struct Backoff {
    step: Cell<u32>,
}

impl Backoff {
    /// Creates a new `Backoff`.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    ///
    /// let backoff = Backoff::new();
    /// ```
    #[inline]
    pub fn new() -> Self {
        Backoff {
            step: Cell::new(0),
        }
    }

    /// Resets the `Backoff`.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    ///
    /// let backoff = Backoff::new();
    /// backoff.reset();
    /// ```
    #[inline]
    pub fn reset(&self) {
        self.step.set(0);
    }

    /// Backs off in a lock-free loop.
    ///
    /// This method should be used when we need to retry an operation because another thread made
    /// progress.
    ///
    /// The processor may yield using the *YIELD* or *PAUSE* instruction.
    ///
    /// # Examples
    ///
    /// Backing off in a lock-free loop:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::atomic::AtomicUsize;
    /// use std::sync::atomic::Ordering::SeqCst;
    ///
    /// fn fetch_mul(a: &AtomicUsize, b: usize) -> usize {
    ///     let backoff = Backoff::new();
    ///     loop {
    ///         let val = a.load(SeqCst);
    ///         if a.compare_and_swap(val, val.wrapping_mul(b), SeqCst) == val {
    ///             return val;
    ///         }
    ///         backoff.spin();
    ///     }
    /// }
    ///
    /// let a = AtomicUsize::new(7);
    /// assert_eq!(fetch_mul(&a, 8), 7);
    /// assert_eq!(a.load(SeqCst), 56);
    /// ```
    #[inline]
    pub fn spin(&self) {
        for _ in 0..1 << self.step.get().min(SPIN_LIMIT) {
            atomic::spin_loop_hint();
        }

        if self.step.get() <= SPIN_LIMIT {
            self.step.set(self.step.get() + 1);
        }
    }

    /// Backs off in a blocking loop.
    ///
    /// This method should be used when we need to wait for another thread to make progress.
    ///
    /// The processor may yield using the *YIELD* or *PAUSE* instruction and the current thread
    /// may yield by giving up a timeslice to the OS scheduler.
    ///
    /// In `#[no_std]` environments, this method is equivalent to [`spin`].
    ///
    /// If possible, use [`is_complete`] to check when it is advised to stop using backoff and
    /// block the current thread using a different synchronization mechanism instead.
    ///
    /// [`spin`]: struct.Backoff.html#method.spin
    /// [`is_complete`]: struct.Backoff.html#method.is_complete
    ///
    /// # Examples
    ///
    /// Waiting for an [`AtomicBool`] to become `true`:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::Arc;
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::SeqCst;
    /// use std::thread;
    /// use std::time::Duration;
    ///
    /// fn spin_wait(ready: &AtomicBool) {
    ///     let backoff = Backoff::new();
    ///     while !ready.load(SeqCst) {
    ///         backoff.snooze();
    ///     }
    /// }
    ///
    /// let ready = Arc::new(AtomicBool::new(false));
    /// let ready2 = ready.clone();
    ///
    /// thread::spawn(move || {
    ///     thread::sleep(Duration::from_millis(100));
    ///     ready2.store(true, SeqCst);
    /// });
    ///
    /// assert_eq!(ready.load(SeqCst), false);
    /// spin_wait(&ready);
    /// assert_eq!(ready.load(SeqCst), true);
    /// ```
    ///
    /// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
    #[inline]
    pub fn snooze(&self) {
        if self.step.get() <= SPIN_LIMIT {
            for _ in 0..1 << self.step.get() {
                atomic::spin_loop_hint();
            }
        } else {
            #[cfg(not(feature = "std"))]
            for _ in 0..1 << self.step.get() {
                atomic::spin_loop_hint();
            }

            #[cfg(feature = "std")]
            ::std::thread::yield_now();
        }

        if self.step.get() <= YIELD_LIMIT {
            self.step.set(self.step.get() + 1);
        }
    }

    /// Returns `true` if exponential backoff is complete and blocking the thread is advised.
    ///
    /// # Examples
    ///
    /// Waiting for an [`AtomicBool`] to become `true` and parking the thread after a long wait:
    ///
    /// ```
    /// use crossbeam_utils::Backoff;
    /// use std::sync::Arc;
    /// use std::sync::atomic::AtomicBool;
    /// use std::sync::atomic::Ordering::SeqCst;
    /// use std::thread;
    /// use std::time::Duration;
    ///
    /// fn blocking_wait(ready: &AtomicBool) {
    ///     let backoff = Backoff::new();
    ///     while !ready.load(SeqCst) {
    ///         if backoff.is_complete() {
    ///             thread::park();
    ///         } else {
    ///             backoff.snooze();
    ///         }
    ///     }
    /// }
    ///
    /// let ready = Arc::new(AtomicBool::new(false));
    /// let ready2 = ready.clone();
    /// let waiter = thread::current();
    ///
    /// thread::spawn(move || {
    ///     thread::sleep(Duration::from_millis(100));
    ///     ready2.store(true, SeqCst);
    ///     waiter.unpark();
    /// });
    ///
    /// assert_eq!(ready.load(SeqCst), false);
    /// blocking_wait(&ready);
    /// assert_eq!(ready.load(SeqCst), true);
    /// ```
    ///
    /// [`AtomicBool`]: https://doc.rust-lang.org/std/sync/atomic/struct.AtomicBool.html
    #[inline]
    pub fn is_complete(&self) -> bool {
        self.step.get() > YIELD_LIMIT
    }
}

impl fmt::Debug for Backoff {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Backoff")
            .field("step", &self.step)
            .field("is_complete", &self.is_complete())
            .finish()
    }
}

impl Default for Backoff {
    fn default() -> Backoff {
        Backoff::new()
    }
}
<|endoftext|>