<commit_before>use std;
use super::{ScriptError, Parser, Parsed};
use super::opcode::*;
use super::flags::*;
use primitive::{Transaction, TxOut, pubkey, UInt256};
use ::serialize::{self, Serializable, SerializeParam, CompactSize};
use std::error::Error;

pub struct Checker<'a> {
   tx:     &'a Transaction,
   in_idx: usize,
}

impl <'a> Checker<'a> {
   pub fn new(tx:&'a Transaction, in_idx:usize) -> Checker {
      Checker { tx:tx, in_idx:in_idx }
   }
   pub fn verify(&self, target:&[u8], pk:&[u8], sig:&[u8], flags:u32) -> Result<(), ScriptError> {
      try!(check_signature_encoding(sig, flags));
      try!(check_pubkey_encoding(pk, flags));

      let (hash_type, _sig) = match sig.split_last() {
         None              => return Err(ScriptError::new("short sig")),
         Some((last, pre)) => (*last, pre)
      };
      let hash = hash_target(target, self.tx, self.in_idx, hash_type);

      match pubkey::verify(pk, &hash, sig) {
         Ok(_)  => Ok(()),
         Err(e) => Err(ScriptError::new(e.description())),
      }
   }
}

#[allow(dead_code)]
struct ByteBuf<'a>(&'a [u8]);
impl<'a> std::fmt::LowerHex for ByteBuf<'a> {
    fn fmt(&self, fmtr: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        for byte in self.0 {
            try!( fmtr.write_fmt(format_args!("{:02x}", byte)));
        }
        Ok(())
    }
}

fn is_compressed_or_uncompressed_pubkey(vch:&[u8]) -> bool {
   let len = vch.len();

   if len < 33 { return false; }
   match vch[0] {
      0x02 => len == 33,
      0x03 => len == 33,
      0x04 => len == 65,
      _ => false
   }
}

fn is_valid_signature_encoding(vch:&[u8]) -> bool {
   let len = vch.len();
   if len < 9 { return false; }
   if len >73 { return false; }

   if vch[0] != 0x30 { return false; }
   if (vch[1]+3) as usize != len { return false; }

   let len_r = vch[3] as usize;
   if 5 + len_r >= len { return false; }

   let len_s = vch[5 + len_r] as usize;
   if len_r + len_s + 7 != len { return false; }

   if vch[2] != 0x02 { return false; }
   if len_r == 0 { return false; }
   if (vch[4] & 0x80) != 0 { return false; }
   if (len_r > 1) && (vch[4] == 0x00) && ((vch[5] & 0x80) != 0) { return false; }

   if vch[len_r+4] != 0x02 { return false; }
   if len_s == 0 { return false; }
   if (vch[len_r+6] & 0x80) != 0 { return false; }
   if (len_s > 1) && (vch[len_r+6] == 0x00) && ((vch[len_r+7] & 0x80) != 0) { return false; }

   true
}

fn is_low_der_signature(vch:&[u8]) -> bool {
   if !is_valid_signature_encoding(vch) {
      return false;
   }
   match pubkey::check_low_s(vch) {
      Ok(_)  => true,
      Err(_) => false,
   }
}

fn is_defined_hashtype_signature(vch:&[u8]) -> bool {
   if vch.len() == 0 { return false; }

   let hash_type = vch[vch.len()-1] & !SIGHASH_ANYONECANPAY;
   if hash_type < SIGHASH_ALL || hash_type > SIGHASH_SINGLE {
      return false;
   }
   true
}

pub fn check_pubkey_encoding(vch:&[u8], flags:u32) -> Result<bool, ScriptError> {
   if (flags & SCRIPT_VERIFY_STRICTENC) != 0 {
      if !is_compressed_or_uncompressed_pubkey(vch) {
         return Err(ScriptError::new("pubkey encoding"));
      }
   }
   Ok(true)
}

pub fn check_signature_encoding(vch:&[u8], flags:u32) -> Result<bool, ScriptError> {
   if vch.len() == 0 { return Ok(true); }

   if (flags & (SCRIPT_VERIFY_DERSIG | SCRIPT_VERIFY_LOW_S | SCRIPT_VERIFY_STRICTENC)) != 0 {
      if !is_valid_signature_encoding(vch) {
         return Err(ScriptError::new("signature encoding"));
      }
   }

   if (flags & SCRIPT_VERIFY_LOW_S) != 0 {
      if !is_low_der_signature(vch) {
         return Err(ScriptError::new("not a low der signature"));
      }
   }

   if (flags & SCRIPT_VERIFY_STRICTENC) != 0 {
      if !is_defined_hashtype_signature(vch) {
         return Err(ScriptError::new("not a defined sig hashtype"));
      }
   }

   Ok(true)
}

fn hash_target(target: &[u8], tx:&Transaction, in_idx:usize, hash_type:u8) -> UInt256 {
   let one = UInt256::from_str("0000000000000000000000000000000000000000000000000000000000000001").unwrap();
   if in_idx >= tx.ins.len() {
      return one;
   }

   // Check for invalid use of SIGHASH_SINGLE
   if (hash_type & 0x1f) == SIGHASH_SINGLE {
      if in_idx >= tx.outs.len() { //??
         return one;
      }
   }

   {
      let tx_ser = TransactionSignatureSerializer::new(target, tx, in_idx, hash_type);
      let io = &mut serialize::HashWriter::new();
      let param  = SerializeParam::new_gethash_ver(0);
      Ok(0)
         .and_then(|_| tx_ser.serialize(io, &param))
         .and_then(|_| hash_type.serialize(io, &param))
         .map(|_| io.get_hash())
         .unwrap_or(one)
   }
}

struct TransactionSignatureSerializer<'a> {
   target: &'a [u8],
   tx: &'a Transaction,
   in_idx: usize,
   anyone_can_pay: bool,
   hash_single: bool,
   hash_none: bool,
}
impl <'a> TransactionSignatureSerializer<'a> {
   fn new(target: &'a [u8], tx:&'a Transaction, in_idx:usize, hash_type:u8) -> TransactionSignatureSerializer<'a> {
      let anyone_can_pay = (hash_type & SIGHASH_ANYONECANPAY) != 0;
      let hash_single    = (hash_type & 0x1f) == SIGHASH_SINGLE;
      let hash_none      = (hash_type & 0x1f) == SIGHASH_NONE;
      TransactionSignatureSerializer {
         target: target,
         tx: tx,
         in_idx: in_idx,
         anyone_can_pay: anyone_can_pay,
         hash_single: hash_single,
         hash_none: hash_none,
      }
   }

   fn serialize_script_code(&self, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      let mut parser = Parser::new(self.target);
      let mut num_sep = 0usize;
      while let Some(Parsed(_pos, code, _follow)) = parser.parse_op() {
         if code == OP_CODESEPARATOR { num_sep += 1; }
      }
      r += try!(CompactSize::Serialize((self.target.len() - num_sep) as u64, io, param));

      let mut parser = Parser::new(self.target);
      let mut pos0 = 0;
      while let Some(Parsed(pos, code, _follow)) = parser.parse_op() {
         if code == OP_CODESEPARATOR {
            r += try!(self.target[pos0..pos].serialize(io, param));
            pos0 = pos + 1;
         }
      }
      if pos0 != self.target.len() {
         r += try!(self.target[pos0..].serialize(io, param));
      }

      Ok(r)
   }

   fn serialize_input(&self, index:usize, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      let index = if self.anyone_can_pay { self.in_idx } else { index };
      r += try!(self.tx.ins[index].prevout.serialize(io, param));

      if index != self.in_idx {
         r += try!(CompactSize::Serialize(0u64, io, param)); // empty script = empty vector
      } else {
         r += try!(self.serialize_script_code(io, param));
      }

      if (index != self.in_idx) && (self.hash_single || self.hash_none) {
         r += try!(0u32.serialize(io, param));
      } else {
         r += try!(self.tx.ins[index].sequence.serialize(io, param));
      }

      Ok(r)
   }

   fn serialize_output(&self, index:usize, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;

      if self.hash_single && index != self.in_idx {
         let out = TxOut::new();
         r += try!(out.serialize(io, param));
      } else {
         r += try!(self.tx.outs[index].serialize(io, param));
      }
      Ok(r)
   }
}
impl <'a> Serializable for TransactionSignatureSerializer<'a> {
   fn get_serialize_size(&self, param:&SerializeParam) -> usize {
      let mut io = serialize::DummyWriter::new();
      let _ = self.serialize(&mut io, param);
      io.len()
   }
   fn serialize(&self, io: &mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      r += try!(self.tx.version.serialize(io, param));

      let num_in = if self.anyone_can_pay { 1usize } else { self.tx.ins.len() };
      r += try!(CompactSize::Serialize(num_in as u64, io, param));
      for i in 0..num_in {
         r += try!(self.serialize_input(i, io, param));
      }

      let num_out = match (self.hash_none, self.hash_single) {
         (true, _)     => 0usize,
         (false, true) => self.in_idx + 1,
         _             => self.tx.outs.len()
      };
      r += try!(CompactSize::Serialize(num_out as u64, io, param));
      for i in 0..num_in {
         r += try!(self.serialize_output(i, io, param));
      }

      r += try!(self.tx.locktime.serialize(io, param));
      Ok(r)
   }
   fn deserialize(&mut self, _io: &mut std::io::Read, _param:&SerializeParam) -> serialize::Result {
      Ok(0usize)
   }
}

<commit_msg>refine error handling<commit_after>use std;
use super::{ScriptError, Parser, Parsed};
use super::opcode::*;
use super::flags::*;
use primitive::{Transaction, TxOut, pubkey, UInt256};
use ::serialize::{self, Serializable, SerializeParam, CompactSize};
use std::error::Error;

pub struct Checker<'a> {
   tx:     &'a Transaction,
   in_idx: usize,
}

impl <'a> Checker<'a> {
   pub fn new(tx:&'a Transaction, in_idx:usize) -> Checker {
      Checker { tx:tx, in_idx:in_idx }
   }
   pub fn verify(&self, target:&[u8], pk:&[u8], sig:&[u8], flags:u32) -> Result<(), ScriptError> {
      try!(check_signature_encoding(sig, flags));
      try!(check_pubkey_encoding(pk, flags));

      let (hash_type, _sig) = match sig.split_last() {
         None              => return Err(ScriptError::new("short sig")),
         Some((last, pre)) => (*last, pre)
      };
      let hash = hash_target(target, self.tx, self.in_idx, hash_type);

      match pubkey::verify(pk, &hash, sig) {
         Ok(_)  => Ok(()),
         Err(e) => Err(ScriptError::new(e.description())),
      }
   }
}

#[allow(dead_code)]
struct ByteBuf<'a>(&'a [u8]);
impl<'a> std::fmt::LowerHex for ByteBuf<'a> {
    fn fmt(&self, fmtr: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        for byte in self.0 {
            try!( fmtr.write_fmt(format_args!("{:02x}", byte)));
        }
        Ok(())
    }
}

fn is_compressed_or_uncompressed_pubkey(vch:&[u8]) -> bool {
   let len = vch.len();

   if len < 33 { return false; }
   match vch[0] {
      0x02 => len == 33,
      0x03 => len == 33,
      0x04 => len == 65,
      _ => false
   }
}

fn is_valid_signature_encoding(vch:&[u8]) -> bool {
   let len = vch.len();
   if len < 9 { return false; }
   if len >73 { return false; }

   if vch[0] != 0x30 { return false; }
   if (vch[1]+3) as usize != len { return false; }

   let len_r = vch[3] as usize;
   if 5 + len_r >= len { return false; }

   let len_s = vch[5 + len_r] as usize;
   if len_r + len_s + 7 != len { return false; }

   if vch[2] != 0x02 { return false; }
   if len_r == 0 { return false; }
   if (vch[4] & 0x80) != 0 { return false; }
   if (len_r > 1) && (vch[4] == 0x00) && ((vch[5] & 0x80) != 0) { return false; }

   if vch[len_r+4] != 0x02 { return false; }
   if len_s == 0 { return false; }
   if (vch[len_r+6] & 0x80) != 0 { return false; }
   if (len_s > 1) && (vch[len_r+6] == 0x00) && ((vch[len_r+7] & 0x80) != 0) { return false; }

   true
}

fn is_low_der_signature(vch:&[u8]) -> bool {
   if !is_valid_signature_encoding(vch) {
      return false;
   }
   match pubkey::check_low_s(vch) {
      Ok(_)  => true,
      Err(_) => false,
   }
}

fn is_defined_hashtype_signature(vch:&[u8]) -> bool {
   if vch.len() == 0 { return false; }

   let hash_type = vch[vch.len()-1] & !SIGHASH_ANYONECANPAY;
   if hash_type < SIGHASH_ALL || hash_type > SIGHASH_SINGLE {
      return false;
   }
   true
}

pub fn check_pubkey_encoding(vch:&[u8], flags:u32) -> Result<bool, ScriptError> {
   if (flags & SCRIPT_VERIFY_STRICTENC) != 0 {
      if !is_compressed_or_uncompressed_pubkey(vch) {
         return Err(ScriptError::new("pubkey encoding"));
      }
   }
   Ok(true)
}

pub fn check_signature_encoding(vch:&[u8], flags:u32) -> Result<bool, ScriptError> {
   if vch.len() == 0 { return Ok(true); }

   if (flags & (SCRIPT_VERIFY_DERSIG | SCRIPT_VERIFY_LOW_S | SCRIPT_VERIFY_STRICTENC)) != 0 {
      if !is_valid_signature_encoding(vch) {
         return Err(ScriptError::new("signature encoding"));
      }
   }

   if (flags & SCRIPT_VERIFY_LOW_S) != 0 {
      if !is_low_der_signature(vch) {
         return Err(ScriptError::new("not a low der signature"));
      }
   }

   if (flags & SCRIPT_VERIFY_STRICTENC) != 0 {
      if !is_defined_hashtype_signature(vch) {
         return Err(ScriptError::new("not a defined sig hashtype"));
      }
   }

   Ok(true)
}

fn hash_target(target: &[u8], tx:&Transaction, in_idx:usize, hash_type:u8) -> UInt256 {
   let one = UInt256::from_str("0000000000000000000000000000000000000000000000000000000000000001").unwrap();
   if in_idx >= tx.ins.len() {
      return one;
   }

   // Check for invalid use of SIGHASH_SINGLE
   if (hash_type & 0x1f) == SIGHASH_SINGLE {
      if in_idx >= tx.outs.len() { //??
         return one;
      }
   }

   {
      let tx_ser = TransactionSignatureSerializer::new(target, tx, in_idx, hash_type);
      let io = &mut serialize::HashWriter::new();
      let param  = SerializeParam::new_gethash_ver(0);
      Ok(0)
         .and_then(|_| tx_ser.serialize(io, &param))
         .and_then(|_| (hash_type as u32).serialize(io, &param))
         .map(|_| io.get_hash())
         .unwrap_or(one)
   }
}

struct TransactionSignatureSerializer<'a> {
   target: &'a [u8],
   tx: &'a Transaction,
   in_idx: usize,
   anyone_can_pay: bool,
   hash_single: bool,
   hash_none: bool,
}
impl <'a> TransactionSignatureSerializer<'a> {
   fn new(target: &'a [u8], tx:&'a Transaction, in_idx:usize, hash_type:u8) -> TransactionSignatureSerializer<'a> {
      let anyone_can_pay = (hash_type & SIGHASH_ANYONECANPAY) != 0;
      let hash_single    = (hash_type & 0x1f) == SIGHASH_SINGLE;
      let hash_none      = (hash_type & 0x1f) == SIGHASH_NONE;
      TransactionSignatureSerializer {
         target: target,
         tx: tx,
         in_idx: in_idx,
         anyone_can_pay: anyone_can_pay,
         hash_single: hash_single,
         hash_none: hash_none,
      }
   }

   fn serialize_script_code(&self, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      let mut parser = Parser::new(self.target);
      let mut num_sep = 0usize;
      while let Some(Parsed(_pos, code, _follow)) = parser.parse_op() {
         if code == OP_CODESEPARATOR { num_sep += 1; }
      }
      r += try!(CompactSize::Serialize((self.target.len() - num_sep) as u64, io, param));

      let mut parser = Parser::new(self.target);
      let mut pos0 = 0;
      while let Some(Parsed(pos, code, _follow)) = parser.parse_op() {
         if code == OP_CODESEPARATOR {
            r += try!(self.target[pos0..pos].serialize(io, param));
            pos0 = pos + 1;
         }
      }
      if pos0 != self.target.len() {
         r += try!(self.target[pos0..].serialize(io, param));
      }

      Ok(r)
   }

   fn serialize_input(&self, index:usize, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      let index = if self.anyone_can_pay { self.in_idx } else { index };
      r += try!(self.tx.ins[index].prevout.serialize(io, param));

      if index != self.in_idx {
         r += try!(CompactSize::Serialize(0u64, io, param)); // empty script = empty vector
      } else {
         r += try!(self.serialize_script_code(io, param));
      }

      if (index != self.in_idx) && (self.hash_single || self.hash_none) {
         r += try!(0u32.serialize(io, param));
      } else {
         r += try!(self.tx.ins[index].sequence.serialize(io, param));
      }

      Ok(r)
   }

   fn serialize_output(&self, index:usize, io:&mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;

      if self.hash_single && index != self.in_idx {
         let out = TxOut::new();
         r += try!(out.serialize(io, param));
      } else {
         r += try!(self.tx.outs[index].serialize(io, param));
      }
      Ok(r)
   }
}
impl <'a> Serializable for TransactionSignatureSerializer<'a> {
   fn get_serialize_size(&self, param:&SerializeParam) -> usize {
      let mut io = serialize::DummyWriter::new();
      let _ = self.serialize(&mut io, param);
      io.len()
   }
   fn serialize(&self, io: &mut std::io::Write, param:&SerializeParam) -> serialize::Result {
      let mut r = 0usize;
      r += try!(self.tx.version.serialize(io, param));

      let num_in = if self.anyone_can_pay { 1usize } else { self.tx.ins.len() };
      r += try!(CompactSize::Serialize(num_in as u64, io, param));
      for i in 0..num_in {
         r += try!(self.serialize_input(i, io, param));
      }

      let num_out = match (self.hash_none, self.hash_single) {
         (true, _)     => 0usize,
         (false, true) => self.in_idx + 1,
         _             => self.tx.outs.len()
      };
      r += try!(CompactSize::Serialize(num_out as u64, io, param));
      for i in 0..num_in {
         r += try!(self.serialize_output(i, io, param));
      }

      r += try!(self.tx.locktime.serialize(io, param));
      Ok(r)
   }
   fn deserialize(&mut self, _io: &mut std::io::Read, _param:&SerializeParam) -> serialize::Result {
      Ok(0usize)
   }
}

<|endoftext|>


<commit_before>use super::Script;
use crate::Lang;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MultiLangScript {
    Latin,
    Cyrillic,
    Arabic,
    Devanagari,
    Hebrew,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScriptLangGroup {
    Multi(MultiLangScript),
    One(Lang),
}

impl Script {
    pub fn to_lang_group(&self) -> ScriptLangGroup {
        use ScriptLangGroup::{Multi, One};
        use MultiLangScript as MLS;

        match *self {
            Script::Latin => Multi(MLS::Latin),
            Script::Cyrillic => Multi(MLS::Cyrillic),
            Script::Arabic => Multi(MLS::Arabic),
            Script::Devanagari => Multi(MLS::Devanagari),
            Script::Hebrew => Multi(MLS::Hebrew),
            Script::Mandarin => One(Lang::Cmn),
            Script::Bengali => One(Lang::Ben),
            Script::Hangul => One(Lang::Kor),
            Script::Georgian => One(Lang::Kat),
            Script::Greek => One(Lang::Ell),
            Script::Kannada => One(Lang::Kan),
            Script::Tamil => One(Lang::Tam),
            Script::Thai => One(Lang::Tha),
            Script::Gujarati => One(Lang::Guj),
            Script::Gurmukhi => One(Lang::Pan),
            Script::Telugu => One(Lang::Tel),
            Script::Malayalam => One(Lang::Mal),
            Script::Oriya => One(Lang::Ori),
            Script::Myanmar => One(Lang::Mya),
            Script::Sinhala => One(Lang::Sin),
            Script::Khmer => One(Lang::Khm),
            Script::Ethiopic => One(Lang::Amh),
            Script::Katakana | Script::Hiragana => One(Lang::Jpn),
        }
    }
}

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

    #[test]
    fn test_to_lang_group() {
        assert_eq!(
            Script::Latin.to_lang_group(),
            ScriptLangGroup::Multi(MultiLangScript::Latin)
        );

        assert_eq!(
            Script::Georgian.to_lang_group(),
            ScriptLangGroup::One(Lang::Kat)
        );
    }
}
<commit_msg>Implement MultiLangScript.to_script()<commit_after>use super::Script;
use crate::Lang;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MultiLangScript {
    Latin,
    Cyrillic,
    Arabic,
    Devanagari,
    Hebrew,
}

impl MultiLangScript {
    pub fn to_script(&self) -> Script {
        match self {
            Self::Latin => Script::Latin,
            Self::Cyrillic => Script::Cyrillic,
            Self::Arabic => Script::Arabic,
            Self::Devanagari => Script::Devanagari,
            Self::Hebrew => Script::Hebrew
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScriptLangGroup {
    Multi(MultiLangScript),
    One(Lang),
}

impl Script {
    pub fn to_lang_group(&self) -> ScriptLangGroup {
        use ScriptLangGroup::{Multi, One};
        use MultiLangScript as MLS;

        match *self {
            Script::Latin => Multi(MLS::Latin),
            Script::Cyrillic => Multi(MLS::Cyrillic),
            Script::Arabic => Multi(MLS::Arabic),
            Script::Devanagari => Multi(MLS::Devanagari),
            Script::Hebrew => Multi(MLS::Hebrew),
            Script::Mandarin => One(Lang::Cmn),
            Script::Bengali => One(Lang::Ben),
            Script::Hangul => One(Lang::Kor),
            Script::Georgian => One(Lang::Kat),
            Script::Greek => One(Lang::Ell),
            Script::Kannada => One(Lang::Kan),
            Script::Tamil => One(Lang::Tam),
            Script::Thai => One(Lang::Tha),
            Script::Gujarati => One(Lang::Guj),
            Script::Gurmukhi => One(Lang::Pan),
            Script::Telugu => One(Lang::Tel),
            Script::Malayalam => One(Lang::Mal),
            Script::Oriya => One(Lang::Ori),
            Script::Myanmar => One(Lang::Mya),
            Script::Sinhala => One(Lang::Sin),
            Script::Khmer => One(Lang::Khm),
            Script::Ethiopic => One(Lang::Amh),
            Script::Katakana | Script::Hiragana => One(Lang::Jpn),
        }
    }
}

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

    #[test]
    fn test_script_to_lang_group() {
        assert_eq!(
            Script::Latin.to_lang_group(),
            ScriptLangGroup::Multi(MultiLangScript::Latin)
        );

        assert_eq!(
            Script::Georgian.to_lang_group(),
            ScriptLangGroup::One(Lang::Kat)
        );
    }

    #[test]
    fn test_multi_lang_script_to_script() {
        assert_eq!(MultiLangScript::Latin.to_script(), Script::Latin);
        assert_eq!(MultiLangScript::Hebrew.to_script(), Script::Hebrew);
    }
}
<|endoftext|>


<commit_before>//! Implements search parallelization facilities.

use std::thread;
use std::cmp::{min, max};
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use moves::*;
use position::Position;
use tt::*;
use search::INITIAL_ASPIRATION_WINDOW;
use search::alpha_beta::Search;

/// Represents a command to a search thread.
pub enum Command {
    /// Requests a new search.
    Search {
        /// A number identifying the new search.
        search_id: usize,

        /// The root position.
        position: Position,

        /// The requested search depth.
        depth: u8,

        /// The lower bound for the new search.
        lower_bound: Value,

        /// The upper bound for the new search.
        upper_bound: Value,

        /// The evaluation of the root position so far, or
        /// `VALUE_UNKNOWN` if not available.
        value: Value,
    },

    /// Stops the currently running search.
    Stop,

    /// Stops the currently running search and exits the search
    /// thread.
    Exit,
}

/// Represents progress report from a search thread.
pub struct Report {
    /// The ID passed with the search command.
    pub search_id: usize,

    /// The number of positions searched so far.
    pub searched_nodes: NodeCount,

    /// The search depth completed so far.
    pub depth: u8,

    /// The evaluation of the root position so far, or `VALUE_UNKNOWN`
    /// if not available.
    pub value: Value,

    /// `true` if the search is finished or has been stopped, `false`
    /// otherwise.
    pub done: bool,
}

/// Listens for commands, executes simple searches, sends reports
/// back.
///
/// This function will block and wait to receive commands on the
/// `commands` channel to start, stop, or exit searches. It is
/// intended to be called in a separate thread. While the search is
/// executed, regular `Report` messages will be send back to the
/// master thread via the `reports` channel. When the search is done,
/// the final `Report` message will have its `done` field set to
/// `true`.
///
/// # Example:
///
/// ```rust
/// // Spawn a slave thread:
/// let tt = Arc::new(tt);
/// let (commands_tx, commands_rx) = channel();
/// let (reports_tx, reports_rx) = channel();
/// thread::spawn(move || {
///     serve_simple(tt, commands_rx, reports_tx);
/// });
///
/// // Send a command to start a new search:
/// commands_tx.send(Command::Search {
///     search_id: 0,
///     position: Position::form_fen("8/8/8/8/8/8/7P/5k1K b - - 0 99"),
///     depth: 5,
///     lower_bound: -29999,
///     upper_bound: 29999,
/// }).unwrap();
/// ```
///
/// This function executes sequential (non-parallel) search to a fixed
/// depth.
pub fn serve_simple(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a command.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound, .. } => {
                    debug_assert!(lower_bound < upper_bound);
                    let mut report = |searched_nodes| {
                        reports.send(Report {
                                   search_id: search_id,
                                   searched_nodes: searched_nodes,
                                   depth: 0,
                                   value: VALUE_UNKNOWN,
                                   done: false,
                               })
                               .ok();
                        if let Ok(cmd) = commands.try_recv() {
                            pending_command = Some(cmd);
                            true
                        } else {
                            false
                        }
                    };
                    let mut search = Search::new(position, &tt, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth, Move::invalid())
                                      .unwrap_or(VALUE_UNKNOWN);
                    reports.send(Report {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: if value == VALUE_UNKNOWN {
                                   0
                               } else {
                                   depth
                               },
                               value: value,
                               done: true,
                           })
                           .ok();
                    search.reset();
                }

                Command::Stop => continue,

                Command::Exit => break,
            }
        }
    })
}

/// Listens for commands, executes deepening searches, sends reports
/// back.
///
/// This function will block and wait to receive commands on the
/// `commands` channel to start, stop, or exit searches. It is
/// intended to be called in a separate thread. While the search is
/// executed, regular `Report` messages will be send back to the
/// master thread via the `reports` channel. When the search is done,
/// the final `Report` message will have its `done` field set to
/// `true`.
///
/// # Example:
///
/// ```rust
/// // Spawn a slave thread:
/// let tt = Arc::new(tt);
/// let (commands_tx, commands_rx) = channel();
/// let (reports_tx, reports_rx) = channel();
/// thread::spawn(move || {
///     serve_deepening(tt, commands_rx, reports_tx);
/// });
///
/// // Send a command to start a new search:
/// commands_tx.send(Command::Search {
///     search_id: 0,
///     position: Position::form_fen("8/8/8/8/8/8/7P/5k1K b - - 0 99"),
///     depth: 5,
///     lower_bound: -29999,
///     upper_bound: 29999,
/// }).unwrap();
/// ```
///
/// This function executes a deepening search with aspiration
/// window. It starts at depth `1` and consequently increases it until
/// the specified final depth is reached.
pub fn serve_deepening(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>) {
    // Start a slave thread that will be commanded to run searches
    // with increasing depths (search deepening).
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        serve_simple(tt, slave_commands_rx, slave_reports_tx);
    });

    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound, .. } => {
                debug_assert!(lower_bound < upper_bound);
                let mut current_searched_nodes = 0;
                let mut current_value = VALUE_UNKNOWN;
                let mut current_depth = 1;

                'depthloop: while current_depth <= depth {
                    // First we set up the aspiration window. Aspiration windows are a way
                    // to reduce the search space in the search. We use `current_value`
                    // from the last iteration of `depth`, and calculate a window around
                    // this as the alpha-beta bounds. Because the window is narrower, more
                    // beta cutoffs are achieved, and the search takes a shorter time. The
                    // drawback is that if the true score is outside this window, then a
                    // costly re-search must be made. But then most probably the re-search
                    // will be much faster, because many positions will be remembered from
                    // the TT.
                    //
                    // Here `delta` is the initial half-width of the window, that will be
                    // exponentially increased each time the search failed. We use `isize`
                    // type to avoid overflows.
                    let mut delta = INITIAL_ASPIRATION_WINDOW as isize;
                    let (mut alpha, mut beta) = if current_depth < 5 {
                        (lower_bound, upper_bound)
                    } else {
                        let v = current_value as isize;
                        (max(lower_bound as isize, v - delta) as Value,
                         min(v + delta, upper_bound as isize) as Value)
                    };

                    'aspiration: loop {
                        if alpha >= beta {
                            // This may happen if (v - delta, v + delta) and
                            // (lower_bound, upper_bound) do not intersect.
                            alpha = lower_bound;
                            beta = upper_bound;
                        }

                        // Command the slave thread to run a search.
                        slave_commands_tx.send(Command::Search {
                                             search_id: current_depth as usize,
                                             position: position.clone(),
                                             depth: current_depth,
                                             lower_bound: alpha,
                                             upper_bound: beta,
                                             value: VALUE_UNKNOWN,
                                         })
                                         .unwrap();

                        'report: loop {
                            // In this loop we process the reports coming from the slave
                            // thread, but we also constantly check if there is a new
                            // pending command for us, in which case we have to terminate
                            // the search.
                            let Report { searched_nodes, value, done, .. } =
                                slave_reports_rx.recv().unwrap();
                            reports.send(Report {
                                       search_id: search_id,
                                       searched_nodes: current_searched_nodes + searched_nodes,
                                       depth: current_depth - 1,
                                       value: current_value,
                                       done: false,
                                   })
                                   .ok();
                            if pending_command.is_none() {
                                if done {
                                    current_searched_nodes += searched_nodes;
                                    current_value = value;
                                    break 'report;
                                }
                                if let Ok(cmd) = commands.try_recv() {
                                    slave_commands_tx.send(Command::Stop).unwrap();
                                    pending_command = Some(cmd);
                                }
                            }
                            if done {
                                current_searched_nodes += searched_nodes;
                                break 'depthloop;
                            }
                        } // end of 'report

                        // Check if the `current_value` is within the aspiration window
                        // (alpha, beta). If not so, we must consider running a re-search.
                        let v = current_value as isize;
                        if current_value <= alpha && lower_bound < alpha {
                            alpha = max(lower_bound as isize, v - delta) as Value;
                        } else if current_value >= beta && upper_bound > beta {
                            beta = min(v + delta, upper_bound as isize) as Value;
                        } else {
                            break 'aspiration;
                        }

                        // Increase the half-width of the aspiration window.
                        delta += 3 * delta / 8;
                        if delta > 1500 {
                            delta = 1_000_000;
                        }

                    } // end of 'aspiration

                    // Send a progress report with `current_value` for
                    // every completed depth.
                    reports.send(Report {
                               search_id: search_id,
                               searched_nodes: current_searched_nodes,
                               depth: current_depth,
                               value: current_value,
                               done: false,
                           })
                           .ok();
                    current_depth += 1;

                } // end of 'depthloop

                // The search is done -- send a final report.
                reports.send(Report {
                           search_id: search_id,
                           searched_nodes: current_searched_nodes,
                           depth: current_depth - 1,
                           value: current_value,
                           done: true,
                       })
                       .ok();
            }

            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }

            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}
<commit_msg>add SearchRefinement and serve()<commit_after>//! Implements search parallelization facilities.

use std::thread;
use std::cmp::{min, max};
use std::cell::UnsafeCell;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver, RecvError};
use basetypes::*;
use moves::*;
use position::Position;
use tt::*;
use search::INITIAL_ASPIRATION_WINDOW;
use search::alpha_beta::Search;

/// Represents a command to a search thread.
pub enum Command {
    /// Requests a new search.
    Search {
        /// A number identifying the new search.
        search_id: usize,

        /// The root position.
        position: Position,

        /// The requested search depth.
        depth: u8,

        /// The lower bound for the new search.
        lower_bound: Value,

        /// The upper bound for the new search.
        upper_bound: Value,

        /// The evaluation of the root position so far, or
        /// `VALUE_UNKNOWN` if not available.
        value: Value,
    },

    /// Stops the currently running search.
    Stop,

    /// Stops the currently running search and exits the search
    /// thread.
    Exit,
}

/// Represents progress report from a search thread.
pub struct Report {
    /// The ID passed with the search command.
    pub search_id: usize,

    /// The number of positions searched so far.
    pub searched_nodes: NodeCount,

    /// The search depth completed so far.
    pub depth: u8,

    /// The evaluation of the root position so far, or `VALUE_UNKNOWN`
    /// if not available.
    pub value: Value,

    /// `true` if the search is finished or has been stopped, `false`
    /// otherwise.
    pub done: bool,
}

/// Listens for commands, executes simple searches, sends reports
/// back.
///
/// This function will block and wait to receive commands on the
/// `commands` channel to start, stop, or exit searches. It is
/// intended to be called in a separate thread. While the search is
/// executed, regular `Report` messages will be send back to the
/// master thread via the `reports` channel. When the search is done,
/// the final `Report` message will have its `done` field set to
/// `true`.
///
/// # Example:
///
/// ```rust
/// // Spawn a slave thread:
/// let tt = Arc::new(tt);
/// let (commands_tx, commands_rx) = channel();
/// let (reports_tx, reports_rx) = channel();
/// thread::spawn(move || {
///     serve_simple(tt, commands_rx, reports_tx);
/// });
///
/// // Send a command to start a new search:
/// commands_tx.send(Command::Search {
///     search_id: 0,
///     position: Position::form_fen("8/8/8/8/8/8/7P/5k1K b - - 0 99"),
///     depth: 5,
///     lower_bound: -29999,
///     upper_bound: 29999,
/// }).unwrap();
/// ```
///
/// This function executes sequential (non-parallel) search to a fixed
/// depth.
pub fn serve_simple(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>) {
    thread_local!(
        static MOVE_STACK: UnsafeCell<MoveStack> = UnsafeCell::new(MoveStack::new())
    );
    MOVE_STACK.with(|s| {
        let mut move_stack = unsafe { &mut *s.get() };
        let mut pending_command = None;
        loop {
            // If there is a pending command, we take it, otherwise we
            // block and wait to receive a command.
            let command = match pending_command.take() {
                Some(cmd) => cmd,
                None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
            };

            match command {
                Command::Search { search_id, position, depth, lower_bound, upper_bound, .. } => {
                    debug_assert!(lower_bound < upper_bound);
                    let mut report = |searched_nodes| {
                        reports.send(Report {
                                   search_id: search_id,
                                   searched_nodes: searched_nodes,
                                   depth: 0,
                                   value: VALUE_UNKNOWN,
                                   done: false,
                               })
                               .ok();
                        if let Ok(cmd) = commands.try_recv() {
                            pending_command = Some(cmd);
                            true
                        } else {
                            false
                        }
                    };
                    let mut search = Search::new(position, &tt, move_stack, &mut report);
                    let value = search.run(lower_bound, upper_bound, depth, Move::invalid())
                                      .unwrap_or(VALUE_UNKNOWN);
                    reports.send(Report {
                               search_id: search_id,
                               searched_nodes: search.node_count(),
                               depth: if value == VALUE_UNKNOWN {
                                   0
                               } else {
                                   depth
                               },
                               value: value,
                               done: true,
                           })
                           .ok();
                    search.reset();
                }

                Command::Stop => continue,

                Command::Exit => break,
            }
        }
    })
}

/// Listens for commands, executes deepening searches, sends reports
/// back.
///
/// This function will block and wait to receive commands on the
/// `commands` channel to start, stop, or exit searches. It is
/// intended to be called in a separate thread. While the search is
/// executed, regular `Report` messages will be send back to the
/// master thread via the `reports` channel. When the search is done,
/// the final `Report` message will have its `done` field set to
/// `true`.
///
/// # Example:
///
/// ```rust
/// // Spawn a slave thread:
/// let tt = Arc::new(tt);
/// let (commands_tx, commands_rx) = channel();
/// let (reports_tx, reports_rx) = channel();
/// thread::spawn(move || {
///     serve_deepening(tt, commands_rx, reports_tx);
/// });
///
/// // Send a command to start a new search:
/// commands_tx.send(Command::Search {
///     search_id: 0,
///     position: Position::form_fen("8/8/8/8/8/8/7P/5k1K b - - 0 99"),
///     depth: 5,
///     lower_bound: -29999,
///     upper_bound: 29999,
/// }).unwrap();
/// ```
///
/// This function executes a deepening search with aspiration
/// window. It starts at depth `1` and consequently increases it until
/// the specified final depth is reached.
pub fn serve_deepening(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>) {
    // Start a slave thread that will be commanded to run searches
    // with increasing depths (search deepening).
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        serve_simple(tt, slave_commands_rx, slave_reports_tx);
    });

    let mut pending_command = None;
    loop {
        // If there is a pending command, we take it, otherwise we
        // block and wait to receive a new one.
        let command = match pending_command.take() {
            Some(cmd) => cmd,
            None => commands.recv().or::<RecvError>(Ok(Command::Exit)).unwrap(),
        };

        match command {
            Command::Search { search_id, position, depth, lower_bound, upper_bound, .. } => {
                debug_assert!(lower_bound < upper_bound);
                let mut current_searched_nodes = 0;
                let mut current_value = VALUE_UNKNOWN;
                let mut current_depth = 1;

                'depthloop: while current_depth <= depth {
                    // First we set up the aspiration window. Aspiration windows are a way
                    // to reduce the search space in the search. We use `current_value`
                    // from the last iteration of `depth`, and calculate a window around
                    // this as the alpha-beta bounds. Because the window is narrower, more
                    // beta cutoffs are achieved, and the search takes a shorter time. The
                    // drawback is that if the true score is outside this window, then a
                    // costly re-search must be made. But then most probably the re-search
                    // will be much faster, because many positions will be remembered from
                    // the TT.
                    //
                    // Here `delta` is the initial half-width of the window, that will be
                    // exponentially increased each time the search failed. We use `isize`
                    // type to avoid overflows.
                    let mut delta = INITIAL_ASPIRATION_WINDOW as isize;
                    let (mut alpha, mut beta) = if current_depth < 5 {
                        (lower_bound, upper_bound)
                    } else {
                        let v = current_value as isize;
                        (max(lower_bound as isize, v - delta) as Value,
                         min(v + delta, upper_bound as isize) as Value)
                    };

                    'aspiration: loop {
                        if alpha >= beta {
                            // This may happen if (v - delta, v + delta) and
                            // (lower_bound, upper_bound) do not intersect.
                            alpha = lower_bound;
                            beta = upper_bound;
                        }

                        // Command the slave thread to run a search.
                        slave_commands_tx.send(Command::Search {
                                             search_id: current_depth as usize,
                                             position: position.clone(),
                                             depth: current_depth,
                                             lower_bound: alpha,
                                             upper_bound: beta,
                                             value: VALUE_UNKNOWN,
                                         })
                                         .unwrap();

                        'report: loop {
                            // In this loop we process the reports coming from the slave
                            // thread, but we also constantly check if there is a new
                            // pending command for us, in which case we have to terminate
                            // the search.
                            let Report { searched_nodes, value, done, .. } =
                                slave_reports_rx.recv().unwrap();
                            reports.send(Report {
                                       search_id: search_id,
                                       searched_nodes: current_searched_nodes + searched_nodes,
                                       depth: current_depth - 1,
                                       value: current_value,
                                       done: false,
                                   })
                                   .ok();
                            if pending_command.is_none() {
                                if done {
                                    current_searched_nodes += searched_nodes;
                                    current_value = value;
                                    break 'report;
                                }
                                if let Ok(cmd) = commands.try_recv() {
                                    slave_commands_tx.send(Command::Stop).unwrap();
                                    pending_command = Some(cmd);
                                }
                            }
                            if done {
                                current_searched_nodes += searched_nodes;
                                break 'depthloop;
                            }
                        } // end of 'report

                        // Check if the `current_value` is within the aspiration window
                        // (alpha, beta). If not so, we must consider running a re-search.
                        let v = current_value as isize;
                        if current_value <= alpha && lower_bound < alpha {
                            alpha = max(lower_bound as isize, v - delta) as Value;
                        } else if current_value >= beta && upper_bound > beta {
                            beta = min(v + delta, upper_bound as isize) as Value;
                        } else {
                            break 'aspiration;
                        }

                        // Increase the half-width of the aspiration window.
                        delta += 3 * delta / 8;
                        if delta > 1500 {
                            delta = 1_000_000;
                        }

                    } // end of 'aspiration

                    // Send a progress report with `current_value` for
                    // every completed depth.
                    reports.send(Report {
                               search_id: search_id,
                               searched_nodes: current_searched_nodes,
                               depth: current_depth,
                               value: current_value,
                               done: false,
                           })
                           .ok();
                    current_depth += 1;

                } // end of 'depthloop

                // The search is done -- send a final report.
                reports.send(Report {
                           search_id: search_id,
                           searched_nodes: current_searched_nodes,
                           depth: current_depth - 1,
                           value: current_value,
                           done: true,
                       })
                       .ok();
            }

            Command::Stop => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }

            Command::Exit => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
        }
    }
    slave.join().unwrap();
}

trait SearchRefinement {
    fn run_slave(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>);
    fn new(tt: Arc<Tt>,
           slave_commands_tx: &Sender<Command>,
           slave_reports_rx: Receiver<Report>,
           reports: Sender<Report>)
           -> Self;
    fn start_search(&mut self,
                    search_id: usize,
                    position: Position,
                    depth: u8,
                    lower_bound: Value,
                    upper_bound: Value,
                    value: Value);
    fn progress(&mut self, search_is_terminated: bool) -> bool;
}

fn serve<T: SearchRefinement>(tt: Arc<Tt>, commands: Receiver<Command>, reports: Sender<Report>) {
    // Start a slave thread that we will send commands to.
    let slave_tt = tt.clone();
    let (slave_commands_tx, slave_commands_rx) = channel();
    let (slave_reports_tx, slave_reports_rx) = channel();
    let slave = thread::spawn(move || {
        T::run_slave(slave_tt, slave_commands_rx, slave_reports_tx);
    });

    // Create a master object that will send commands to the slave,
    // receive slave's reports and write to `reports`.
    let mut master = T::new(tt, &slave_commands_tx, slave_reports_rx, reports);

    // Orchestrate the work of the master and the slave.
    let mut pending_command = None;
    loop {
        match pending_command.take() {
            Some(Command::Search { search_id,
                                   position,
                                   depth,
                                   lower_bound,
                                   upper_bound,
                                   value }) => {
                master.start_search(search_id, position, depth, lower_bound, upper_bound, value);
                while !master.progress(pending_command.is_some()) {
                    if pending_command.is_none() {
                        if let Ok(cmd) = commands.try_recv() {
                            slave_commands_tx.send(Command::Stop).unwrap();
                            pending_command = Some(cmd);
                        }
                    }
                }
            }
            Some(Command::Stop) => {
                slave_commands_tx.send(Command::Stop).unwrap();
                continue;
            }
            Some(Command::Exit) => {
                slave_commands_tx.send(Command::Exit).unwrap();
                break;
            }
            None => pending_command = commands.recv().or::<RecvError>(Ok(Command::Exit)).ok(),
        }
    }
    slave.join().unwrap();
}
<|endoftext|>


<commit_before>use super::Hash256;
use super::TransitionError;

pub fn get_signed_parent_hashes(
    cycle_length: u8,
    block_slot: u64,
    attestation_slot: u64,
    current_hashes: Vec<Hash256>,
    oblique_hashes: Vec<Hash256>)
    -> Result<Vec<Hash256>, TransitionError>
{
    // This cast places a limit on cycle_length. If you change it, check math
    // for overflow.
    let cycle_length: u64 = cycle_length as u64;

    if current_hashes.len() as u64 != (cycle_length * 2) {
        return Err(TransitionError::InvalidInput(String::from(
                    "current_hashes.len() must equal cycle_length * 2")));
    }
    if attestation_slot >= block_slot {
        return Err(TransitionError::InvalidInput(String::from(
                    "attestation_slot must be less than block_slot")));
    }
    if oblique_hashes.len() as u64 > cycle_length {
        return Err(TransitionError::InvalidInput(String::from(
                    "oblique_hashes.len() must be <= cycle_length * 2")));
    }

    /*
     * Cannot underflow as block_slot cannot be less
     * than attestation_slot.
     */
    let attestation_distance = block_slot - attestation_slot;

    if attestation_distance > cycle_length {
        return Err(TransitionError::InvalidInput(String::from(
                    "attestation_slot must be withing one cycle of block_slot")));
    }

    /*
     * Cannot underflow because attestation_distance cannot
     * be larger than cycle_length.
     */
    let start = cycle_length - attestation_distance;

    /*
     * Overflow is potentially impossible, but proof is complicated
     * enough to just use checked math.
     *
     * Math is:
     * start + cycle_length - oblique_hashes.len()
     */
    let end = start.checked_add(cycle_length)
        .and_then(|x| x.checked_sub(oblique_hashes.len() as u64))
        .ok_or(TransitionError::IntWrapping)?;

    let mut hashes = Vec::new();
    hashes.extend_from_slice(
        &current_hashes[(start as usize)..(end as usize)]);
    hashes.append(&mut oblique_hashes.clone());

    Ok(hashes)
}

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

    fn get_range_of_hashes(from: usize, to: usize) -> Vec<Hash256> {
        (from..to).map(|i| get_hash(&vec![i as u8])).collect()
    }

    fn get_hash(value: &[u8]) -> Hash256 {
        Hash256::from_slice(value)
    }

    #[test]
    fn test_get_signed_hashes_no_oblique() {
        /*
         * Google Slides example.
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 19;
        let attestation_slot: u64 = 15;
        let current_hashes = get_range_of_hashes(3, 19);
        let oblique_hashes = vec![];
        let result = get_signed_parent_hashes(
            cycle_length,
            block_slot,
            attestation_slot,
            current_hashes,
            oblique_hashes);
        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.len(), cycle_length as usize);
        let expected_result = get_range_of_hashes(7, 15);
        assert_eq!(result, expected_result);
    }
}
<commit_msg>Use borrows for get_signed_parent_hases, add tests<commit_after>use super::Hash256;
use super::TransitionError;

pub fn get_signed_parent_hashes(
    cycle_length: &u8,
    block_slot: &u64,
    attestation_slot: &u64,
    current_hashes: &Vec<Hash256>,
    oblique_hashes: &Vec<Hash256>)
    -> Result<Vec<Hash256>, TransitionError>
{
    // This cast places a limit on cycle_length. If you change it, check math
    // for overflow.
    let cycle_length: u64 = *cycle_length as u64;

    if current_hashes.len() as u64 != (cycle_length * 2) {
        return Err(TransitionError::InvalidInput(String::from(
                    "current_hashes.len() must equal cycle_length * 2")));
    }
    if attestation_slot >= block_slot {
        return Err(TransitionError::InvalidInput(String::from(
                    "attestation_slot must be less than block_slot")));
    }
    if oblique_hashes.len() as u64 > cycle_length {
        return Err(TransitionError::InvalidInput(String::from(
                    "oblique_hashes.len() must be <= cycle_length * 2")));
    }

    /*
     * Cannot underflow as block_slot cannot be less
     * than attestation_slot.
     */
    let attestation_distance = block_slot - attestation_slot;

    if attestation_distance > cycle_length {
        return Err(TransitionError::InvalidInput(String::from(
                    "attestation_slot must be withing one cycle of block_slot")));
    }

    /*
     * Cannot underflow because attestation_distance cannot
     * be larger than cycle_length.
     */
    let start = cycle_length - attestation_distance;

    /*
     * Overflow is potentially impossible, but proof is complicated
     * enough to just use checked math.
     *
     * Arithmetic is:
     * start + cycle_length - oblique_hashes.len()
     */
    let end = start.checked_add(cycle_length)
        .and_then(|x| x.checked_sub(oblique_hashes.len() as u64))
        .ok_or(TransitionError::IntWrapping)?;

    let mut hashes = Vec::new();
    hashes.extend_from_slice(
        &current_hashes[(start as usize)..(end as usize)]);
    hashes.append(&mut oblique_hashes.clone());

    Ok(hashes)
}

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

    fn get_range_of_hashes(from: usize, to: usize) -> Vec<Hash256> {
        (from..to).map(|i| get_hash(&vec![i as u8])).collect()
    }

    fn get_hash(value: &[u8]) -> Hash256 {
        Hash256::from_slice(value)
    }

    #[test]
    fn test_get_signed_hashes_oblique_scenario_1() {
        /*
         * Two oblique hashes.
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 19;
        let attestation_slot: u64 = 15;
        let current_hashes = get_range_of_hashes(3, 19);
        let oblique_hashes = get_range_of_hashes(100, 102);
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.len(), cycle_length as usize);

        let mut expected_result = get_range_of_hashes(7, 13);
        expected_result.append(&mut get_range_of_hashes(100, 102));
        assert_eq!(result, expected_result);
    }

    #[test]
    fn test_get_signed_hashes_oblique_scenario_2() {
        /*
         * All oblique hashes.
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 19;
        let attestation_slot: u64 = 15;
        let current_hashes = get_range_of_hashes(3, 19);
        let oblique_hashes = get_range_of_hashes(100, 108);
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.len(), cycle_length as usize);

        let expected_result = get_range_of_hashes(100, 108);
        assert_eq!(result, expected_result);
    }

    #[test]
    fn test_get_signed_hashes_scenario_1() {
        /*
         * Google Slides example.
         * https://tinyurl.com/ybzn2spw
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 19;
        let attestation_slot: u64 = 15;
        let current_hashes = get_range_of_hashes(3, 19);
        let oblique_hashes = vec![];
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.len(), cycle_length as usize);
        let expected_result = get_range_of_hashes(7, 15);
        assert_eq!(result, expected_result);
    }

    #[test]
    fn test_get_signed_hashes_scenario_2() {
        /*
         * Block 1, attestation 0.
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 1;
        let attestation_slot: u64 = 0;
        let current_hashes = get_range_of_hashes(0, 16);
        let oblique_hashes = vec![];
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.len(), cycle_length as usize);
        let expected_result = get_range_of_hashes(7, 15);
        assert_eq!(result, expected_result);
    }

    #[test]
    fn test_get_signed_hashes_scenario_3() {
        /*
         * attestation_slot too large
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 100;
        let attestation_slot: u64 = 100;
        let current_hashes = get_range_of_hashes(0, 16);
        let oblique_hashes = vec![];
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_err());
    }

    #[test]
    fn test_get_signed_hashes_scenario_4() {
        /*
         * Current hashes too small
         */
        let cycle_length: u8 = 8;
        let block_slot: u64 = 100;
        let attestation_slot: u64 = 99;
        let current_hashes = get_range_of_hashes(0, 15);
        let oblique_hashes = vec![];
        let result = get_signed_parent_hashes(
            &cycle_length,
            &block_slot,
            &attestation_slot,
            &current_hashes,
            &oblique_hashes);
        assert!(result.is_err());
    }
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::cell::{Cell, UnsafeCell};
use std::{mem, ptr};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};

use futures::{Async, Poll};
use futures::stream::Stream;
use futures::task::{self, Task};

#[cfg(target_pointer_width = "32")]
const PAD: usize = 16;
#[cfg(target_pointer_width = "64")]
const PAD: usize = 8;

#[repr(C)]
struct RingBuffer<T, E> {
    read:          AtomicUsize,
    last_write:    Cell<usize>,
    _padding1:     [usize; PAD - 2],
    write:         AtomicUsize,
    last_read:     Cell<usize>,
    _padding2:     [usize; PAD - 2],
    buffer:        *mut T,
    capacity:      usize,
    size:          usize,
    pub error:     AtomicPtr<E>,
    pub is_closed: AtomicBool,
    task:          UnsafeCell<Option<Task>>,
    pub task_lock: AtomicBool
}

impl<T, E> RingBuffer<T, E> {
    #[inline]
    pub fn new(capacity: usize) -> RingBuffer<T, E> {
        let size = capacity.next_power_of_two();
        let mut vec = Vec::with_capacity(size);
        let ptr = vec.as_mut_ptr();

        mem::forget(vec);

        RingBuffer {
            read:       AtomicUsize::new(0),
            last_write: Cell::new(0),
            _padding1:  [0; PAD - 2],
            write:      AtomicUsize::new(0),
            last_read:  Cell::new(0),
            _padding2:  [0; PAD - 2],
            buffer:     ptr,
            capacity:   capacity,
            size:       size,
            error:      AtomicPtr::new(ptr::null_mut()),
            is_closed:  AtomicBool::new(false),
            task:       UnsafeCell::new(None),
            task_lock:  AtomicBool::new(false)
        }
    }

    #[inline]
    pub fn try_push(&self, value: T) -> Option<T> {
        let read = self.last_read.get();
        let write = self.write.load(Ordering::Relaxed);

        if read + self.capacity <= write {
            let read = self.read.load(Ordering::Relaxed);
            self.last_read.set(read);

            if read + self.capacity <= write {
                return Some(value);
            }
        }

        self.push(write, value);

        self.write.store(write.wrapping_add(1), Ordering::Relaxed);

        self.try_unpark();

        None
    }

    #[inline]
    pub fn try_pop(&self) -> Option<T> {
        let read = self.read.load(Ordering::Relaxed);
        let write = self.last_write.get();

        if read == write {
            let write = self.write.load(Ordering::Relaxed);
            self.last_write.set(write);

            if read == write {
                return None;
            }
        }

        let result = self.pop(read);

        self.read.store(read.wrapping_add(1), Ordering::Relaxed);

        Some(result)
    }

    #[inline]
    pub fn try_park(&self) -> bool {
        let locked = self.task_lock.swap(true, Ordering::Relaxed);

        if !locked {
            self.swap(Some(task::park()));

            self.task_lock.store(false, Ordering::Relaxed);
        }

        !locked
    }

    #[inline]
    pub fn try_unpark(&self) -> bool {
        let locked = self.task_lock.swap(true, Ordering::Relaxed);

        if !locked {
            let task = self.swap(None);

            if let Some(task) = task {
                task.unpark();
            }

            self.task_lock.store(false, Ordering::Relaxed);
        }

        !locked
    }

    #[inline]
    pub fn swap(&self, option: Option<Task>) -> Option<Task> {
        let mut task = unsafe { &mut *self.task.get()};
        mem::replace(task, option)
    }

    #[inline]
    fn push(&self, index: usize, value: T) {
        unsafe {
            ptr::write(&mut *self.buffer.offset(self.index(index)), value);
        }
    }

    #[inline]
    fn pop(&self, index: usize) -> T {
        unsafe {
            ptr::read(self.buffer.offset(self.index(index)))
        }
    }

    #[inline]
    fn index(&self, index: usize) -> isize {
        (index & (self.size - 1)) as isize
    }
}

impl<T, E> Drop for RingBuffer<T, E> {
    fn drop(&mut self) {
        unsafe  {
            Vec::from_raw_parts(self.buffer, 0, self.size);
        }
    }
}

unsafe impl<T: Send, E: Send> Send for RingBuffer<T, E> {}
unsafe impl<T: Sync, E: Sync> Sync for RingBuffer<T, E> {}

pub struct BufferedStream<T, E> {
    buffer: Arc<RingBuffer<T, E>>
}

impl<T, E> Stream for BufferedStream<T, E> {
    type Item  = T;
    type Error = E;

    fn poll(&mut self) -> Poll<Option<T>, E> {
        let error = self.buffer.error.load(Ordering::Acquire);

        if !error.is_null() {
            let error = unsafe { Err(*Box::from_raw(error)) };
            self.buffer.error.store(ptr::null_mut(), Ordering::Release);

            return error;
        }

        match self.buffer.try_pop() {
            Some(value) => Ok(Async::Ready(Some(value))),
            None => {
                if self.buffer.is_closed.load(Ordering::Relaxed) {
                    Ok(Async::Ready(None))
                } else {
                    loop {
                        if self.buffer.try_park() {
                            break;
                        }
                    }

                    Ok(Async::NotReady)
                }
            }
        }
    }
}

pub struct BufferedSender<T, E> {
    buffer: Arc<RingBuffer<T, E>>
}

impl<T, E> BufferedSender<T, E> {
    #[inline]
    pub fn send(&self, value: T) {
        let mut value = value;

        loop {
            value = match self.buffer.try_push(value) {
                Some(value) => value,
                None        => break
            };
        }
    }

    #[inline]
    pub fn fail(self, error: E) {
        let boxed = Box::new(error);
        self.buffer.error.store(Box::into_raw(boxed), Ordering::Release);
    }
}

impl<T, E> Drop for BufferedSender<T, E> {
    fn drop(&mut self) {
        self.buffer.is_closed.store(true, Ordering::Relaxed);

        loop {
            if self.buffer.try_unpark() {
                break;
            }
        }
    }
}

pub fn buffered<T, E>(capacity: usize) -> (BufferedSender<T, E>, BufferedStream<T, E>) {
    let buffer = Arc::new(RingBuffer::new(capacity));

    (BufferedSender { buffer: buffer.clone() }, BufferedStream { buffer: buffer })
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::thread;
    use std::time::Duration;

    use futures::future::Future;
    use futures::stream::Stream;

    use super::*;

    #[test]
    fn ring_push_pop() {
        let ring: RingBuffer<i32, ()> = RingBuffer::new(3);

        assert_eq!(ring.try_push(1), None);
        assert_eq!(ring.try_push(2), None);
        assert_eq!(ring.try_push(3), None);
        assert_eq!(ring.try_push(4), Some(4));

        assert_eq!(ring.try_pop(), Some(1));
        assert_eq!(ring.try_pop(), Some(2));
        assert_eq!(ring.try_pop(), Some(3));
        assert_eq!(ring.try_pop(), None);

        assert_eq!(ring.try_push(1), None);
        assert_eq!(ring.try_push(2), None);
        assert_eq!(ring.try_push(3), None);
        assert_eq!(ring.try_push(4), Some(4));

        assert_eq!(ring.try_pop(), Some(1));
        assert_eq!(ring.try_pop(), Some(2));
        assert_eq!(ring.try_pop(), Some(3));
        assert_eq!(ring.try_pop(), None);
    }

    #[test]
    fn buffer_fail() {
        let (sender, stream) = buffered::<(), String>(1);

        sender.fail("error".to_owned());

        assert_eq!(stream.collect().wait(), Err("error".to_owned()))
    }

    #[test]
    fn ring_async() {
        let ring: Arc<RingBuffer<i32, ()>> = Arc::new(RingBuffer::new(3));
        let clone = ring.clone();

        thread::spawn(move || {
            for _ in 0..10 {
                thread::sleep(Duration::from_millis(1));

                assert_eq!(clone.try_push(1), None);
            }
        });

        let mut count = 0;

        loop {
            if ring.try_pop().is_some() {
                count += 1;
            }

            if count == 10 {
                break;
            }
        }
    }

    #[test]
    fn sender_stream() {
        let (sender, stream) = buffered::<i32, ()>(3);

        thread::spawn(move || {
            thread::sleep(Duration::from_millis(10));

            for _ in 0..100 {
                sender.send(1);
            }
        });

        let sum = stream.fold(0, |sum, value| {
            Ok(sum + value)
        }).wait().unwrap();

        assert_eq!(sum, 100);
    }
}
<commit_msg>Fixed flaky test.<commit_after>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::cell::{Cell, UnsafeCell};
use std::{mem, ptr};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering};

use futures::{Async, Poll};
use futures::stream::Stream;
use futures::task::{self, Task};

#[cfg(target_pointer_width = "32")]
const PAD: usize = 16;
#[cfg(target_pointer_width = "64")]
const PAD: usize = 8;

#[repr(C)]
struct RingBuffer<T, E> {
    read:          AtomicUsize,
    last_write:    Cell<usize>,
    _padding1:     [usize; PAD - 2],
    write:         AtomicUsize,
    last_read:     Cell<usize>,
    _padding2:     [usize; PAD - 2],
    buffer:        *mut T,
    capacity:      usize,
    size:          usize,
    pub error:     AtomicPtr<E>,
    pub is_closed: AtomicBool,
    task:          UnsafeCell<Option<Task>>,
    pub task_lock: AtomicBool
}

impl<T, E> RingBuffer<T, E> {
    #[inline]
    pub fn new(capacity: usize) -> RingBuffer<T, E> {
        let size = capacity.next_power_of_two();
        let mut vec = Vec::with_capacity(size);
        let ptr = vec.as_mut_ptr();

        mem::forget(vec);

        RingBuffer {
            read:       AtomicUsize::new(0),
            last_write: Cell::new(0),
            _padding1:  [0; PAD - 2],
            write:      AtomicUsize::new(0),
            last_read:  Cell::new(0),
            _padding2:  [0; PAD - 2],
            buffer:     ptr,
            capacity:   capacity,
            size:       size,
            error:      AtomicPtr::new(ptr::null_mut()),
            is_closed:  AtomicBool::new(false),
            task:       UnsafeCell::new(None),
            task_lock:  AtomicBool::new(false)
        }
    }

    #[inline]
    pub fn try_push(&self, value: T) -> Option<T> {
        let read = self.last_read.get();
        let write = self.write.load(Ordering::Relaxed);

        if read + self.capacity <= write {
            let read = self.read.load(Ordering::Relaxed);
            self.last_read.set(read);

            if read + self.capacity <= write {
                return Some(value);
            }
        }

        self.push(write, value);

        self.write.store(write.wrapping_add(1), Ordering::Relaxed);

        self.try_unpark();

        None
    }

    #[inline]
    pub fn try_pop(&self) -> Option<T> {
        let read = self.read.load(Ordering::Relaxed);
        let write = self.last_write.get();

        if read == write {
            let write = self.write.load(Ordering::Relaxed);
            self.last_write.set(write);

            if read == write {
                return None;
            }
        }

        let result = self.pop(read);

        self.read.store(read.wrapping_add(1), Ordering::Relaxed);

        Some(result)
    }

    #[inline]
    pub fn try_park(&self) -> bool {
        let locked = self.task_lock.swap(true, Ordering::Relaxed);

        if !locked {
            self.swap(Some(task::park()));

            self.task_lock.store(false, Ordering::Relaxed);
        }

        !locked
    }

    #[inline]
    pub fn try_unpark(&self) -> bool {
        let locked = self.task_lock.swap(true, Ordering::Relaxed);

        if !locked {
            let task = self.swap(None);

            if let Some(task) = task {
                task.unpark();
            }

            self.task_lock.store(false, Ordering::Relaxed);
        }

        !locked
    }

    #[inline]
    pub fn swap(&self, option: Option<Task>) -> Option<Task> {
        let mut task = unsafe { &mut *self.task.get()};
        mem::replace(task, option)
    }

    #[inline]
    fn push(&self, index: usize, value: T) {
        unsafe {
            ptr::write(&mut *self.buffer.offset(self.index(index)), value);
        }
    }

    #[inline]
    fn pop(&self, index: usize) -> T {
        unsafe {
            ptr::read(self.buffer.offset(self.index(index)))
        }
    }

    #[inline]
    fn index(&self, index: usize) -> isize {
        (index & (self.size - 1)) as isize
    }
}

impl<T, E> Drop for RingBuffer<T, E> {
    fn drop(&mut self) {
        unsafe  {
            Vec::from_raw_parts(self.buffer, 0, self.size);
        }
    }
}

unsafe impl<T: Send, E: Send> Send for RingBuffer<T, E> {}
unsafe impl<T: Sync, E: Sync> Sync for RingBuffer<T, E> {}

pub struct BufferedStream<T, E> {
    buffer: Arc<RingBuffer<T, E>>
}

impl<T, E> Stream for BufferedStream<T, E> {
    type Item  = T;
    type Error = E;

    fn poll(&mut self) -> Poll<Option<T>, E> {
        let error = self.buffer.error.load(Ordering::Acquire);

        if !error.is_null() {
            let error = unsafe { Err(*Box::from_raw(error)) };
            self.buffer.error.store(ptr::null_mut(), Ordering::Release);

            return error;
        }

        match self.buffer.try_pop() {
            Some(value) => Ok(Async::Ready(Some(value))),
            None => {
                if self.buffer.is_closed.load(Ordering::Relaxed) {
                    Ok(Async::Ready(None))
                } else {
                    loop {
                        if self.buffer.try_park() {
                            break;
                        }
                    }

                    Ok(Async::NotReady)
                }
            }
        }
    }
}

pub struct BufferedSender<T, E> {
    buffer: Arc<RingBuffer<T, E>>
}

impl<T, E> BufferedSender<T, E> {
    #[inline]
    pub fn send(&self, value: T) {
        let mut value = value;

        loop {
            value = match self.buffer.try_push(value) {
                Some(value) => value,
                None        => break
            };
        }
    }

    #[inline]
    pub fn fail(self, error: E) {
        let boxed = Box::new(error);
        self.buffer.error.store(Box::into_raw(boxed), Ordering::Release);
    }
}

impl<T, E> Drop for BufferedSender<T, E> {
    fn drop(&mut self) {
        self.buffer.is_closed.store(true, Ordering::Relaxed);

        loop {
            if self.buffer.try_unpark() {
                break;
            }
        }
    }
}

pub fn buffered<T, E>(capacity: usize) -> (BufferedSender<T, E>, BufferedStream<T, E>) {
    let buffer = Arc::new(RingBuffer::new(capacity));

    (BufferedSender { buffer: buffer.clone() }, BufferedStream { buffer: buffer })
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::thread;
    use std::time::Duration;

    use futures::future::Future;
    use futures::stream::Stream;

    use super::*;

    #[test]
    fn ring_push_pop() {
        let ring: RingBuffer<i32, ()> = RingBuffer::new(3);

        assert_eq!(ring.try_push(1), None);
        assert_eq!(ring.try_push(2), None);
        assert_eq!(ring.try_push(3), None);
        assert_eq!(ring.try_push(4), Some(4));

        assert_eq!(ring.try_pop(), Some(1));
        assert_eq!(ring.try_pop(), Some(2));
        assert_eq!(ring.try_pop(), Some(3));
        assert_eq!(ring.try_pop(), None);

        assert_eq!(ring.try_push(1), None);
        assert_eq!(ring.try_push(2), None);
        assert_eq!(ring.try_push(3), None);
        assert_eq!(ring.try_push(4), Some(4));

        assert_eq!(ring.try_pop(), Some(1));
        assert_eq!(ring.try_pop(), Some(2));
        assert_eq!(ring.try_pop(), Some(3));
        assert_eq!(ring.try_pop(), None);
    }

    #[test]
    fn buffer_fail() {
        let (sender, stream) = buffered::<(), String>(1);

        sender.fail("error".to_owned());

        assert_eq!(stream.collect().wait(), Err("error".to_owned()))
    }

    #[test]
    fn ring_async() {
        let ring: Arc<RingBuffer<i32, ()>> = Arc::new(RingBuffer::new(10));
        let clone = ring.clone();

        thread::spawn(move || {
            for _ in 0..10 {
                thread::sleep(Duration::from_millis(1));

                assert_eq!(clone.try_push(1), None);
            }
        });

        let mut count = 0;

        loop {
            if ring.try_pop().is_some() {
                count += 1;
            }

            if count == 10 {
                break;
            }
        }
    }

    #[test]
    fn sender_stream() {
        let (sender, stream) = buffered::<i32, ()>(3);

        thread::spawn(move || {
            thread::sleep(Duration::from_millis(10));

            for _ in 0..100 {
                sender.send(1);
            }
        });

        let sum = stream.fold(0, |sum, value| {
            Ok(sum + value)
        }).wait().unwrap();

        assert_eq!(sum, 100);
    }
}
<|endoftext|>


<commit_before>#![allow(unstable)]
#![deny(warnings)]
#![feature(plugin_registrar)]

extern crate rustc;
extern crate syntax;

use rustc::plugin::registry::Registry;
use syntax::ast::{CompilerGenerated, TtToken, TokenTree, UnsafeBlock};
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacExpr, MacResult, NormalTT};
use syntax::ext::build::AstBuilder;
use syntax::parse::token::{self, Comma};

/// Executes several closures in parallel
///
/// - This macro is an expression, and will return a tuple containing the values returned by the
/// closures.
/// - One thread per closure (i.e. there is no "load balancing").
/// - This macro will block until all the spawned threads have finished. For this reason the
///   closures don't need to fulfill the `Send` trait (i.e. the closures may capture slices and
///   references).
///
/// # Expansion
///
/// Consider the following call:
///
/// ``` ignore
/// let (a, b, c) = execute!{
///     || job_0(),
///     || job_1(),
///     || job_2(),
/// };
/// ```
///
/// This macro uses the unsafe [`fork`](../parallel/fn.fork.html) function defined in the
/// [`parallel`](../parallel/index.html) crate, and its expansion looks (roughly) like this:
///
/// ``` ignore
/// let (a, b, c) = {
///     let __thread_0 = || job_0();
///     let __thread_1 = unsafe { ::parallel::fork(|| job_1()) };  // spawns a thread
///     let __thread_2 = unsafe { ::parallel::fork(|| job_2()) };  // spawns another thread
///
///     // the current thread takes care of the first closure
///     (__thread_0(), __thread_1.join().ok().unwrap(), __thread_2.join().ok().unwrap())
///     // ^~ then blocks until the other two threads finish
/// };
/// ```
///
/// # Panics
///
/// This macro will panic if any of the spawned threads panics.
///
/// # Example
///
/// I'll borrow the binary tree example from Niko's
/// [blog](http://smallcultfollowing.com/babysteps/blog/2013/06/11/data-parallelism-in-rust/).
///
/// ```
/// #![feature(plugin)]
///
/// extern crate parallel;
/// #[plugin]
/// extern crate parallel_macros;
///
/// struct Tree {
///     left: Option<Box<Tree>>,
///     right: Option<Box<Tree>>,
///     value: i32,
/// }
///
/// impl Tree {
///     fn sum(&self) -> i32 {
///         fn sum(subtree: &Option<Box<Tree>>) -> i32 {
///             match *subtree {
///                 None => 0,
///                 Some(ref tree) => tree.sum(),
///             }
///         }
///
///         let (left_sum, right_sum) = execute! {
///             // NB Each closure captures a reference and therefore doesn't fulfill `Send`
///             |:| sum(&self.left),
///             |:| sum(&self.right),
///         };
///
///         left_sum + self.value + right_sum
///     }
/// }
///
/// fn main() {
///     let tree = Tree {
///         value: 5,
///         left: Some(Box::new(Tree {
///             value: 3,
///             left: Some(Box::new(Tree {
///                 value: 1,
///                 left: None,
///                 right: Some(Box::new(Tree {
///                     value: 4,
///                     left: None,
///                     right: None,
///                 })),
///             })),
///             right: None,
///         })),
///         right: Some(Box::new(Tree {
///             value: 7,
///             left: None,
///             right: None,
///         })),
///     };
///
///     assert_eq!(tree.sum(), 20);
/// }
/// ```
#[macro_export]
macro_rules! execute {
    ($($closure:expr),+,) => ({ /* syntax extension */ });
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(r: &mut Registry) {
    r.register_syntax_extension(token::intern("execute"), NormalTT(Box::new(expand_execute), None));
}

fn expand_execute<'cx>(
    cx: &'cx mut ExtCtxt,
    sp: Span,
    tts: &[TokenTree],
) -> Box<MacResult + 'cx> {
    let parallel_fork_fn_path = {
        let segments = vec![
            cx.ident_of("parallel"),
            cx.ident_of("fork"),
        ];

        cx.expr_path(cx.path_global(sp, segments))
    };

    let mut stmts = vec![];
    let threads = tts.split(|tt| match *tt {
        TtToken(_, Comma) => true,
        _ => false,
    }).filter(|tts| {
        !tts.is_empty()
    }).enumerate().map(|(i, tts)|  {
        let closure = cx.new_parser_from_tts(tts).parse_expr();
        let ident = cx.ident_of(&*format!("__thread_{}", i));

        let expr = if i == 0 {
            closure
        } else {
            let fn_name = parallel_fork_fn_path.clone();
            let args = vec![closure];

            // XXX There has to be a simpler way to wrap an expression in `unsafe`
            let block = cx.block_expr(cx.expr_call(sp, fn_name, args)).map(|mut b| {
                b.rules = UnsafeBlock(CompilerGenerated);
                b
            });
            cx.expr_block(block)
        };

        stmts.push(cx.stmt_let(sp, false, ident, expr));

        ident
    }).collect::<Vec<_>>();

    let mut is_first = true;
    let expr = cx.expr_tuple(sp, threads.into_iter().map(|thread| {
        let thread = cx.expr_ident(sp, thread);

        if is_first {
            let args = vec![];
            is_first = false;

            cx.expr_call(sp, thread, args)
        } else {
            let args = vec![];
            let join_method = cx.ident_of("join");
            let ok_method = cx.ident_of("ok");
            let unwrap_method = cx.ident_of("unwrap");

            let join_call = cx.expr_method_call(sp, thread, join_method, args.clone());
            let ok_call = cx.expr_method_call(sp, join_call, ok_method, args.clone());
            cx.expr_method_call(sp, ok_call, unwrap_method, args)
        }
    }).collect());

    MacExpr::new(cx.expr_block(cx.block(sp, stmts, Some(expr))))
}
<commit_msg>style: split long line<commit_after>#![allow(unstable)]
#![deny(warnings)]
#![feature(plugin_registrar)]

extern crate rustc;
extern crate syntax;

use rustc::plugin::registry::Registry;
use syntax::ast::{CompilerGenerated, TtToken, TokenTree, UnsafeBlock};
use syntax::codemap::Span;
use syntax::ext::base::{ExtCtxt, MacExpr, MacResult, NormalTT};
use syntax::ext::build::AstBuilder;
use syntax::parse::token::{self, Comma};

/// Executes several closures in parallel
///
/// - This macro is an expression, and will return a tuple containing the values returned by the
/// closures.
/// - One thread per closure (i.e. there is no "load balancing").
/// - This macro will block until all the spawned threads have finished. For this reason the
///   closures don't need to fulfill the `Send` trait (i.e. the closures may capture slices and
///   references).
///
/// # Expansion
///
/// Consider the following call:
///
/// ``` ignore
/// let (a, b, c) = execute!{
///     || job_0(),
///     || job_1(),
///     || job_2(),
/// };
/// ```
///
/// This macro uses the unsafe [`fork`](../parallel/fn.fork.html) function defined in the
/// [`parallel`](../parallel/index.html) crate, and its expansion looks (roughly) like this:
///
/// ``` ignore
/// let (a, b, c) = {
///     let __thread_0 = || job_0();
///     let __thread_1 = unsafe { ::parallel::fork(|| job_1()) };  // spawns a thread
///     let __thread_2 = unsafe { ::parallel::fork(|| job_2()) };  // spawns another thread
///
///     // the current thread takes care of the first closure
///     (__thread_0(), __thread_1.join().ok().unwrap(), __thread_2.join().ok().unwrap())
///     // ^~ then blocks until the other two threads finish
/// };
/// ```
///
/// # Panics
///
/// This macro will panic if any of the spawned threads panics.
///
/// # Example
///
/// I'll borrow the binary tree example from Niko's
/// [blog](http://smallcultfollowing.com/babysteps/blog/2013/06/11/data-parallelism-in-rust/).
///
/// ```
/// #![feature(plugin)]
///
/// extern crate parallel;
/// #[plugin]
/// extern crate parallel_macros;
///
/// struct Tree {
///     left: Option<Box<Tree>>,
///     right: Option<Box<Tree>>,
///     value: i32,
/// }
///
/// impl Tree {
///     fn sum(&self) -> i32 {
///         fn sum(subtree: &Option<Box<Tree>>) -> i32 {
///             match *subtree {
///                 None => 0,
///                 Some(ref tree) => tree.sum(),
///             }
///         }
///
///         let (left_sum, right_sum) = execute! {
///             // NB Each closure captures a reference and therefore doesn't fulfill `Send`
///             |:| sum(&self.left),
///             |:| sum(&self.right),
///         };
///
///         left_sum + self.value + right_sum
///     }
/// }
///
/// fn main() {
///     let tree = Tree {
///         value: 5,
///         left: Some(Box::new(Tree {
///             value: 3,
///             left: Some(Box::new(Tree {
///                 value: 1,
///                 left: None,
///                 right: Some(Box::new(Tree {
///                     value: 4,
///                     left: None,
///                     right: None,
///                 })),
///             })),
///             right: None,
///         })),
///         right: Some(Box::new(Tree {
///             value: 7,
///             left: None,
///             right: None,
///         })),
///     };
///
///     assert_eq!(tree.sum(), 20);
/// }
/// ```
#[macro_export]
macro_rules! execute {
    ($($closure:expr),+,) => ({ /* syntax extension */ });
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(r: &mut Registry) {
    r.register_syntax_extension(
        token::intern("execute"),
        NormalTT(Box::new(expand_execute), None));
}

fn expand_execute<'cx>(
    cx: &'cx mut ExtCtxt,
    sp: Span,
    tts: &[TokenTree],
) -> Box<MacResult + 'cx> {
    let parallel_fork_fn_path = {
        let segments = vec![
            cx.ident_of("parallel"),
            cx.ident_of("fork"),
        ];

        cx.expr_path(cx.path_global(sp, segments))
    };

    let mut stmts = vec![];
    let threads = tts.split(|tt| match *tt {
        TtToken(_, Comma) => true,
        _ => false,
    }).filter(|tts| {
        !tts.is_empty()
    }).enumerate().map(|(i, tts)|  {
        let closure = cx.new_parser_from_tts(tts).parse_expr();
        let ident = cx.ident_of(&*format!("__thread_{}", i));

        let expr = if i == 0 {
            closure
        } else {
            let fn_name = parallel_fork_fn_path.clone();
            let args = vec![closure];

            // XXX There has to be a simpler way to wrap an expression in `unsafe`
            let block = cx.block_expr(cx.expr_call(sp, fn_name, args)).map(|mut b| {
                b.rules = UnsafeBlock(CompilerGenerated);
                b
            });
            cx.expr_block(block)
        };

        stmts.push(cx.stmt_let(sp, false, ident, expr));

        ident
    }).collect::<Vec<_>>();

    let mut is_first = true;
    let expr = cx.expr_tuple(sp, threads.into_iter().map(|thread| {
        let thread = cx.expr_ident(sp, thread);

        if is_first {
            let args = vec![];
            is_first = false;

            cx.expr_call(sp, thread, args)
        } else {
            let args = vec![];
            let join_method = cx.ident_of("join");
            let ok_method = cx.ident_of("ok");
            let unwrap_method = cx.ident_of("unwrap");

            let join_call = cx.expr_method_call(sp, thread, join_method, args.clone());
            let ok_call = cx.expr_method_call(sp, join_call, ok_method, args.clone());
            cx.expr_method_call(sp, ok_call, unwrap_method, args)
        }
    }).collect());

    MacExpr::new(cx.expr_block(cx.block(sp, stmts, Some(expr))))
}
<|endoftext|>