<commit_before>// Rust Bitcoin Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Bitcoin Blockchain
//!
//! This module provides the structures and functions to maintain the
//! blockchain.
//!

use alloc::rc::Rc;
use std::cell::RefCell;
use std::io::{IoResult, IoError, OtherIoError};

use blockdata::block::BlockHeader;
use blockdata::constants::{DIFFCHANGE_INTERVAL, DIFFCHANGE_TIMESPAN, max_target};
use network::serialize::{Serializable, SerializeIter};
use util::uint256::Uint256;
use util::hash::Sha256dHash;
use util::misc::prepend_err;
use util::patricia_tree::PatriciaTree;

/// A link in the blockchain
struct BlockchainNode {
  /// The blockheader
  header: BlockHeader,
  /// Total work from genesis to this point
  total_work: Uint256,
  /// Expected value of `header.bits` for this block; only changes every
  /// `blockdata::constants::DIFFCHANGE_INTERVAL;` blocks
  required_difficulty: Uint256,
  /// Height above genesis
  height: u32,
  /// Pointer to block's parent
  prev: RefCell<Option<Rc<BlockchainNode>>>
}

impl BlockchainNode {
  /// Look up the previous link, caching the result
  fn prev(&self, tree: &PatriciaTree<Rc<BlockchainNode>>) -> Option<Rc<BlockchainNode>> {
    let mut cache = self.prev.borrow_mut();
    if cache.is_some() {
      return Some(cache.get_ref().clone())
    }
    match tree.lookup(&self.header.prev_blockhash.as_bitv()) {
      Some(prev) => { *cache = Some(prev.clone()); return Some(prev.clone()); }
      None => { return None; }
    }
  }
}

impl Serializable for Rc<BlockchainNode> {
  fn serialize(&self) -> Vec<u8> {
    let mut ret = vec![];
    ret.extend(self.header.serialize().move_iter());
    ret.extend(self.total_work.serialize().move_iter());
    ret.extend(self.required_difficulty.serialize().move_iter());
    ret.extend(self.height.serialize().move_iter());
    // Don't serialize the prev pointer
    ret
  }

  fn deserialize<I: Iterator<u8>>(mut iter: I) -> IoResult<Rc<BlockchainNode>> {
    Ok(Rc::new(BlockchainNode {
      header: try!(prepend_err("header", Serializable::deserialize(iter.by_ref()))),
      total_work: try!(prepend_err("total_work", Serializable::deserialize(iter.by_ref()))),
      required_difficulty: try!(prepend_err("req_difficulty", Serializable::deserialize(iter.by_ref()))),
      height: try!(prepend_err("height", Serializable::deserialize(iter.by_ref()))),
      prev: RefCell::new(None)
    }))
  }

  // Override Serialize::hash to return the blockheader hash, since the
  // hash of the node itself is pretty much meaningless.
  fn hash(&self) -> Sha256dHash {
    self.header.hash()
  }
}

/// The blockchain
pub struct Blockchain {
  tree: PatriciaTree<Rc<BlockchainNode>>,
  best_tip: Rc<BlockchainNode>,
  best_hash: Sha256dHash
}

impl Serializable for Blockchain {
  fn serialize(&self) -> Vec<u8> {
    let mut ret = vec![];
    ret.extend(self.tree.serialize().move_iter());
    ret.extend(self.best_hash.serialize().move_iter());
    ret
  }

  fn serialize_iter<'a>(&'a self) -> SerializeIter<'a> {
    SerializeIter {
      data_iter: None,
      sub_iter_iter: box vec![ &self.tree as &Serializable,
                               &self.best_hash as &Serializable ].move_iter(),
      sub_iter: None,
      sub_started: false
    }
  }

  fn deserialize<I: Iterator<u8>>(mut iter: I) -> IoResult<Blockchain> {
    let tree: PatriciaTree<Rc<BlockchainNode>> = try!(prepend_err("tree", Serializable::deserialize(iter.by_ref())));
    let hash: Sha256dHash = try!(prepend_err("besthash", Serializable::deserialize(iter.by_ref())));
    let best = match tree.lookup(&hash.as_bitv()) {
      Some(rc) => rc.clone(),
      None => { return Err(IoError {
          kind: OtherIoError,
          desc: "best tip reference not found in tree",
          detail: Some(format!("best tip {:x} not found", hash))
        });
      }
    };
    Ok(Blockchain {
      tree: tree,
      best_tip: best.clone(),
      best_hash: best.hash()
    })
  }
}

struct LocatorHashIter<'tree> {
  index: Option<Rc<BlockchainNode>>,
  tree: &'tree PatriciaTree<Rc<BlockchainNode>>,
  count: uint,
  skip: uint
}

impl<'tree> LocatorHashIter<'tree> {
  fn new<'tree>(init: Rc<BlockchainNode>, tree: &'tree PatriciaTree<Rc<BlockchainNode>>) -> LocatorHashIter<'tree> {
    LocatorHashIter { index: Some(init), tree: tree, count: 0, skip: 1 }
  }
}

impl<'tree> Iterator<Sha256dHash> for LocatorHashIter<'tree> {
  fn next(&mut self) -> Option<Sha256dHash> {
    let ret = match self.index {
      Some(ref node) => Some(node.hash()),
      None => None
    };

    for _ in range(0, self.skip) {
      self.index = match self.index {
        Some(ref rc) => rc.prev(self.tree),
        None => None
      }
    }

    self.count += 1;
    if self.count > 10 {
      self.skip *= 2;
    }
    ret
  }
}

/// This function emulates the GetCompact(SetCompact(n)) in the satoshi code,
/// which drops the precision to something that can be encoded precisely in
/// the nBits block header field. Savour the perversity. This is in Bitcoin
/// consensus code. What. The. Fuck.
fn satoshi_the_precision(n: &Uint256) -> Uint256 {
  // Shift by B bits right then left to turn the low bits to zero
  let bits = 8 * ((n.bits() + 7) / 8 - 3);
  let mut ret = n.shr(bits);
  // Oh, did I say B was that fucked up formula? I meant sometimes also + 8.
  if ret.bit_value(23) {
    ret = ret.shr(8).shl(8);
  }
  ret.shl(bits)
}

impl Blockchain {
  /// Constructs a new blockchain
  pub fn new(genesis: BlockHeader) -> Blockchain {
    let genhash = genesis.hash().as_bitv();
    let rc_gen = Rc::new(BlockchainNode {
      header: genesis,
      total_work: Uint256::from_u64(0),
      required_difficulty: genesis.target(),
      height: 0,
      prev: RefCell::new(None)
    });
    Blockchain {
      tree: {
        let mut pat = PatriciaTree::new();
        pat.insert(&genhash, rc_gen.clone());
        pat
      },
      best_hash: rc_gen.hash(),
      best_tip: rc_gen
    }
  }

  /// Adds a block header to the chain
  pub fn add_header(&mut self, header: BlockHeader) -> bool {
    // Construct node, if possible
    let rc_header = match self.tree.lookup(&header.prev_blockhash.as_bitv()) {
      Some(prev) => {
        let difficulty =
          // Compute required difficulty if this is a diffchange block
          if (prev.height + 1) % DIFFCHANGE_INTERVAL == 0 {
            // Scan back DIFFCHANGE_INTERVAL blocks
            let mut scan = prev.clone();
            for _ in range(0, DIFFCHANGE_INTERVAL - 1) {
              scan = scan.prev(&self.tree).unwrap();
            }
            // Get clamped timespan between first and last blocks
            let timespan = match prev.header.time - scan.header.time {
              n if n < DIFFCHANGE_TIMESPAN / 4 => DIFFCHANGE_TIMESPAN / 4,
              n if n > DIFFCHANGE_TIMESPAN * 4 => DIFFCHANGE_TIMESPAN * 4,
              n => n
            };
            // Compute new target
            let mut target = prev.header.target();
            target = target.mul_u32(timespan);
            target = target.div(&Uint256::from_u64(DIFFCHANGE_TIMESPAN as u64));
            // Clamp below MAX_TARGET (difficulty 1)
            let max = max_target();
            if target > max { target = max };
            // Compactify (make expressible in the 8+24 nBits float format
            satoshi_the_precision(&target)
          } else {
          // Otherwise just use the last block's difficulty
             prev.required_difficulty
          };
        // Create node
        Rc::new(BlockchainNode {
          header: header,
          total_work: header.work().add(&prev.total_work),
          required_difficulty: difficulty,
          height: prev.height + 1,
          prev: RefCell::new(Some(prev.clone()))
        })
      },
      None => {
        println!("TODO: couldn't add blockheader");
        return false;
      }
    };

    // spv validate the block
    if !header.spv_validate(&rc_header.required_difficulty) {
      return false;
    }

    // Insert the new block
    self.tree.insert(&header.hash().as_bitv(), rc_header.clone());
    // Replace the best tip if necessary
    if rc_header.total_work > self.best_tip.total_work {
      self.set_best_tip(rc_header);
    }
    return true;
  }

  /// Sets the best tip (not public)
  fn set_best_tip(&mut self, tip: Rc<BlockchainNode>) {
    self.best_hash = tip.hash();
    self.best_tip = tip;
  }

  /// Returns the best tip
  pub fn best_tip<'a>(&'a self) -> &'a BlockHeader {
    &self.best_tip.header
  }

  /// Returns an array of locator hashes used in `getheaders` messages
  pub fn locator_hashes(&self) -> Vec<Sha256dHash> {
    LocatorHashIter::new(self.best_tip.clone(), &self.tree).collect()
  }
}

#[cfg(test)]
mod tests {
  use std::prelude::*;
  use std::io::IoResult;

  use blockdata::blockchain::Blockchain;
  use blockdata::constants::genesis_block;
  use network::serialize::Serializable;

  #[test]
  fn blockchain_serialize_test() {
    let empty_chain = Blockchain::new(genesis_block().header);
    assert_eq!(empty_chain.best_tip.hash().serialize(), genesis_block().header.hash().serialize());

    let serial = empty_chain.serialize();
    assert_eq!(serial, empty_chain.serialize_iter().collect());

    let deserial: IoResult<Blockchain> = Serializable::deserialize(serial.iter().map(|n| *n));
    assert!(deserial.is_ok());
    let read_chain = deserial.unwrap();
    assert_eq!(read_chain.best_tip.hash().serialize(), genesis_block().header.hash().serialize());
  }
}

<commit_msg>Add next pointers to blockchain for easier scanning<commit_after>// Rust Bitcoin Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Bitcoin Blockchain
//!
//! This module provides the structures and functions to maintain the
//! blockchain.
//!

use alloc::rc::Rc;
use std::cell::{Ref, RefCell};
use std::io::{IoResult, IoError, OtherIoError};

use blockdata::block::BlockHeader;
use blockdata::constants::{DIFFCHANGE_INTERVAL, DIFFCHANGE_TIMESPAN, max_target};
use network::serialize::{Serializable, SerializeIter};
use util::uint256::Uint256;
use util::hash::Sha256dHash;
use util::misc::prepend_err;
use util::patricia_tree::PatriciaTree;

/// A link in the blockchain
struct BlockchainNode {
  /// The blockheader
  header: BlockHeader,
  /// Total work from genesis to this point
  total_work: Uint256,
  /// Expected value of `header.bits` for this block; only changes every
  /// `blockdata::constants::DIFFCHANGE_INTERVAL;` blocks
  required_difficulty: Uint256,
  /// Height above genesis
  height: u32,
  /// Pointer to block's parent
  prev: RefCell<Option<Rc<BlockchainNode>>>,
  /// Pointer to block's child
  next: RefCell<Option<Rc<BlockchainNode>>>
}

impl BlockchainNode {
  /// Look up the previous link, caching the result
  fn prev(&self, tree: &PatriciaTree<Rc<BlockchainNode>>) -> Option<Rc<BlockchainNode>> {
    let mut cache = self.prev.borrow_mut();
    if cache.is_some() {
      return Some(cache.get_ref().clone())
    }
    match tree.lookup(&self.header.prev_blockhash.as_bitv()) {
      Some(prev) => { *cache = Some(prev.clone()); return Some(prev.clone()); }
      None => { return None; }
    }
  }

  /// Look up the next link
  fn next<'a>(&'a self) -> Ref<'a, Option<Rc<BlockchainNode>>> {
    self.next.borrow()
  }

  /// Set the next link
  fn set_next(&self, next: Rc<BlockchainNode>) {
    let mut cache = self.next.borrow_mut();
    *cache = Some(next);
  }
}

impl Serializable for Rc<BlockchainNode> {
  fn serialize(&self) -> Vec<u8> {
    let mut ret = vec![];
    ret.extend(self.header.serialize().move_iter());
    ret.extend(self.total_work.serialize().move_iter());
    ret.extend(self.required_difficulty.serialize().move_iter());
    ret.extend(self.height.serialize().move_iter());
    // Don't serialize the prev pointer
    ret
  }

  fn deserialize<I: Iterator<u8>>(mut iter: I) -> IoResult<Rc<BlockchainNode>> {
    Ok(Rc::new(BlockchainNode {
      header: try!(prepend_err("header", Serializable::deserialize(iter.by_ref()))),
      total_work: try!(prepend_err("total_work", Serializable::deserialize(iter.by_ref()))),
      required_difficulty: try!(prepend_err("req_difficulty", Serializable::deserialize(iter.by_ref()))),
      height: try!(prepend_err("height", Serializable::deserialize(iter.by_ref()))),
      prev: RefCell::new(None),
      next: RefCell::new(None)
    }))
  }

  // Override Serialize::hash to return the blockheader hash, since the
  // hash of the node itself is pretty much meaningless.
  fn hash(&self) -> Sha256dHash {
    self.header.hash()
  }
}

/// The blockchain
pub struct Blockchain {
  tree: PatriciaTree<Rc<BlockchainNode>>,
  best_tip: Rc<BlockchainNode>,
  best_hash: Sha256dHash,
  genesis_hash: Sha256dHash
}

impl Serializable for Blockchain {
  fn serialize(&self) -> Vec<u8> {
    let mut ret = vec![];
    ret.extend(self.tree.serialize().move_iter());
    ret.extend(self.best_hash.serialize().move_iter());
    ret.extend(self.genesis_hash.serialize().move_iter());
    ret
  }

  fn serialize_iter<'a>(&'a self) -> SerializeIter<'a> {
    SerializeIter {
      data_iter: None,
      sub_iter_iter: box vec![ &self.tree as &Serializable,
                               &self.best_hash as &Serializable,
                               &self.genesis_hash as &Serializable ].move_iter(),
      sub_iter: None,
      sub_started: false
    }
  }

  fn deserialize<I: Iterator<u8>>(mut iter: I) -> IoResult<Blockchain> {
    let tree: PatriciaTree<Rc<BlockchainNode>> = try!(prepend_err("tree", Serializable::deserialize(iter.by_ref())));
    let best_hash: Sha256dHash = try!(prepend_err("best_hash", Serializable::deserialize(iter.by_ref())));
    let genesis_hash: Sha256dHash = try!(prepend_err("genesis_hash", Serializable::deserialize(iter.by_ref())));
    // Lookup best tip
    let best = match tree.lookup(&best_hash.as_bitv()) {
      Some(rc) => rc.clone(),
      None => { return Err(IoError {
          kind: OtherIoError,
          desc: "best tip reference not found in tree",
          detail: Some(format!("best tip {:x} not found", best_hash))
        });
      }
    };
    // Lookup genesis
    if tree.lookup(&genesis_hash.as_bitv()).is_none() {
      return Err(IoError {
        kind: OtherIoError,
        desc: "genesis header not found in tree",
        detail: Some(format!("genesis header {:x} not found", genesis_hash))
      });
    }
    // Reconnect next and prev pointers back to "genesis", the first node
    // with no prev pointer.
    let mut scan = best.clone();
    let mut prev = best.prev(&tree);
    while prev.is_some() {
      prev.get_mut_ref().set_next(scan);
      scan = prev.get_ref().clone();
      prev = prev.get_ref().prev(&tree);
    }
    // Check that "genesis" is the genesis
    if scan.header.hash() != genesis_hash {
      Err(IoError {
          kind: OtherIoError,
          desc: "best tip did not link back to genesis",
          detail: Some(format!("no path from tip {:x} to genesis {:x}", best_hash, genesis_hash))
      })
    } else {
      // Return the chain
      Ok(Blockchain {
        tree: tree,
        best_tip: best.clone(),
        best_hash: best_hash,
        genesis_hash: genesis_hash
      })
    }
  }
}

struct LocatorHashIter<'tree> {
  index: Option<Rc<BlockchainNode>>,
  tree: &'tree PatriciaTree<Rc<BlockchainNode>>,
  count: uint,
  skip: uint
}

impl<'tree> LocatorHashIter<'tree> {
  fn new<'tree>(init: Rc<BlockchainNode>, tree: &'tree PatriciaTree<Rc<BlockchainNode>>) -> LocatorHashIter<'tree> {
    LocatorHashIter { index: Some(init), tree: tree, count: 0, skip: 1 }
  }
}

impl<'tree> Iterator<Sha256dHash> for LocatorHashIter<'tree> {
  fn next(&mut self) -> Option<Sha256dHash> {
    let ret = match self.index {
      Some(ref node) => Some(node.hash()),
      None => None
    };

    for _ in range(0, self.skip) {
      self.index = match self.index {
        Some(ref rc) => rc.prev(self.tree),
        None => { break; }
      }
    }

    self.count += 1;
    if self.count > 10 {
      self.skip *= 2;
    }
    ret
  }
}

/// This function emulates the GetCompact(SetCompact(n)) in the satoshi code,
/// which drops the precision to something that can be encoded precisely in
/// the nBits block header field. Savour the perversity. This is in Bitcoin
/// consensus code. What. The. Fuck.
fn satoshi_the_precision(n: &Uint256) -> Uint256 {
  // Shift by B bits right then left to turn the low bits to zero
  let bits = 8 * ((n.bits() + 7) / 8 - 3);
  let mut ret = n.shr(bits);
  // Oh, did I say B was that fucked up formula? I meant sometimes also + 8.
  if ret.bit_value(23) {
    ret = ret.shr(8).shl(8);
  }
  ret.shl(bits)
}

impl Blockchain {
  /// Constructs a new blockchain
  pub fn new(genesis: BlockHeader) -> Blockchain {
    let genhash = genesis.hash();
    let rc_gen = Rc::new(BlockchainNode {
      header: genesis,
      total_work: Uint256::from_u64(0),
      required_difficulty: genesis.target(),
      height: 0,
      prev: RefCell::new(None),
      next: RefCell::new(None)
    });
    Blockchain {
      tree: {
        let mut pat = PatriciaTree::new();
        pat.insert(&genhash.as_bitv(), rc_gen.clone());
        pat
      },
      best_hash: genhash,
      genesis_hash: genhash,
      best_tip: rc_gen,
    }
  }

  /// Adds a block header to the chain
  pub fn add_header(&mut self, header: BlockHeader) -> bool {
    // Construct node, if possible
    let rc_header = match self.tree.lookup(&header.prev_blockhash.as_bitv()) {
      Some(prev) => {
        let difficulty =
          // Compute required difficulty if this is a diffchange block
          if (prev.height + 1) % DIFFCHANGE_INTERVAL == 0 {
            // Scan back DIFFCHANGE_INTERVAL blocks
            let mut scan = prev.clone();
            for _ in range(0, DIFFCHANGE_INTERVAL - 1) {
              scan = scan.prev(&self.tree).unwrap();
            }
            // Get clamped timespan between first and last blocks
            let timespan = match prev.header.time - scan.header.time {
              n if n < DIFFCHANGE_TIMESPAN / 4 => DIFFCHANGE_TIMESPAN / 4,
              n if n > DIFFCHANGE_TIMESPAN * 4 => DIFFCHANGE_TIMESPAN * 4,
              n => n
            };
            // Compute new target
            let mut target = prev.header.target();
            target = target.mul_u32(timespan);
            target = target.div(&Uint256::from_u64(DIFFCHANGE_TIMESPAN as u64));
            // Clamp below MAX_TARGET (difficulty 1)
            let max = max_target();
            if target > max { target = max };
            // Compactify (make expressible in the 8+24 nBits float format
            satoshi_the_precision(&target)
          } else {
          // Otherwise just use the last block's difficulty
             prev.required_difficulty
          };
        // Create node
        let ret = Rc::new(BlockchainNode {
          header: header,
          total_work: header.work().add(&prev.total_work),
          required_difficulty: difficulty,
          height: prev.height + 1,
          prev: RefCell::new(Some(prev.clone())),
          next: RefCell::new(None)
        });
        prev.set_next(ret.clone());
        ret
      },
      None => {
        println!("TODO: couldn't add blockheader");
        return false;
      }
    };

    // spv validate the block
    if !header.spv_validate(&rc_header.required_difficulty) {
      return false;
    }

    // Insert the new block
    self.tree.insert(&header.hash().as_bitv(), rc_header.clone());
    // Replace the best tip if necessary
    if rc_header.total_work > self.best_tip.total_work {
      self.set_best_tip(rc_header);
    }
    return true;
  }

  /// Sets the best tip (not public)
  fn set_best_tip(&mut self, tip: Rc<BlockchainNode>) {
    let old_best = self.best_tip.clone();
    // Set best
    self.best_hash = tip.hash();
    self.best_tip = tip;
    // Fix next links
    let mut scan = self.best_tip.clone();
    let mut prev = self.best_tip.prev(&self.tree);
    // Scan backward
    loop {
      // If we hit the old best, there is no need to reorg
      if scan.header == old_best.header {
        break;
      }
      // If we hit the genesis, stop
      if prev.is_none() {
        println!("Warning: reorg past the genesis. This is a bug.");
        break;
      }
      // If we hit something pointing along the wrong chain, this is
      // a branch point at which we are reorg'ing
      if prev.get_ref().next().is_none() ||
         prev.get_ref().next().get_ref().header != scan.header {
        prev.get_mut_ref().set_next(scan);
      }
      scan = prev.clone().unwrap();
      prev = prev.unwrap().prev(&self.tree);
    }
  }

  /// Returns the best tip
  pub fn best_tip<'a>(&'a self) -> &'a BlockHeader {
    &self.best_tip.header
  }

  /// Returns an array of locator hashes used in `getheaders` messages
  pub fn locator_hashes(&self) -> Vec<Sha256dHash> {
    LocatorHashIter::new(self.best_tip.clone(), &self.tree).collect()
  }
}

#[cfg(test)]
mod tests {
  use std::prelude::*;
  use std::io::IoResult;

  use blockdata::blockchain::Blockchain;
  use blockdata::constants::genesis_block;
  use network::serialize::Serializable;

  #[test]
  fn blockchain_serialize_test() {
    let empty_chain = Blockchain::new(genesis_block().header);
    assert_eq!(empty_chain.best_tip.hash().serialize(), genesis_block().header.hash().serialize());

    let serial = empty_chain.serialize();
    assert_eq!(serial, empty_chain.serialize_iter().collect());

    let deserial: IoResult<Blockchain> = Serializable::deserialize(serial.iter().map(|n| *n));
    assert!(deserial.is_ok());
    let read_chain = deserial.unwrap();
    assert_eq!(read_chain.best_tip.hash().serialize(), genesis_block().header.hash().serialize());
  }
}

<|endoftext|>


<commit_before>#![cfg(all(not(target_os = "fuchsia"), not(target_os = "android")))]

extern crate libc;
extern crate pagecache;
extern crate rand;
extern crate sled;
extern crate tests;

use std::fs;
use std::mem::size_of;
use std::path::Path;
use std::thread;
use std::time::Duration;

use rand::Rng;

use pagecache::{Config, ConfigBuilder};

const CYCLE: usize = 16; // 65536;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (mut k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            k.reverse();
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    let lowest_vec = u32_to_vec(lowest);

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.scan(&*low_beginning) {
        let (mut k, v) = res.unwrap();
        if v != lowest_vec {
            k.reverse();
        }
        assert_eq!(
            v,
            lowest_vec,
            "expected key {} to have value {}, instead it had value {}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v)
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] =
        unsafe { std::mem::transmute(u) };
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    unsafe { std::mem::transmute(buf) }
}

fn run(config: Config) {
    // tests::setup_logger();
    let tree = sled::Tree::start(config).unwrap();

    // flush to ensure the initial root is stable.
    // TODO is this necessary or voodoo?
    tree.flush().unwrap();

    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        unsafe {
            libc::raise(9);
        }
    });

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let mut key = u32_to_vec((hu % CYCLE) as u32);
        key.reverse();

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 100_000);
        value.append(&mut vec![0u8; additional_len]);

        tree.set(key, value).unwrap();
    }
}

fn run_without_snapshot() {
    let config = ConfigBuilder::new()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes".to_string())
        .snapshot_after_ops(1 << 56)
        .build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot() {
    let config = ConfigBuilder::new()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes_with_snapshot".to_string())
        .snapshot_after_ops(1 << 10)
        .build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    cleanup_with_snapshots();
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(
                    child,
                    &mut status as *mut libc::c_int,
                    0,
                );
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup_with_snapshots();
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    cleanup();
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(
                    child,
                    &mut status as *mut libc::c_int,
                    0,
                );
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup();
}

fn cleanup_with_snapshots() {
    let dir = Path::new("test_crashes_with_snapshot");
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}

fn cleanup() {
    let dir = Path::new("test_crashes");
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}
<commit_msg>fix compilation error of tests on windows<commit_after>#![cfg(all(not(target_os = "fuchsia"), not(target_os = "android"), not(target_os = "windows")))]

extern crate libc;
extern crate pagecache;
extern crate rand;
extern crate sled;
extern crate tests;

use std::fs;
use std::mem::size_of;
use std::path::Path;
use std::thread;
use std::time::Duration;

use rand::Rng;

use pagecache::{Config, ConfigBuilder};

const CYCLE: usize = 16; // 65536;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        Some(Ok((_k, v))) => slice_to_u32(&*v),
        Some(Err(e)) => panic!("{:?}", e),
        None => return (0, 0),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for res in iter {
        let (mut k, v) = res.unwrap();
        if &v[..4] == &highest_vec[..4] {
            contiguous += 1;
        } else {
            k.reverse();
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    let lowest_vec = u32_to_vec(lowest);

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for res in tree.scan(&*low_beginning) {
        let (mut k, v) = res.unwrap();
        if v != lowest_vec {
            k.reverse();
        }
        assert_eq!(
            v,
            lowest_vec,
            "expected key {} to have value {}, instead it had value {}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v)
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; size_of::<u32>()] =
        unsafe { std::mem::transmute(u) };
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; size_of::<u32>()];
    buf.copy_from_slice(&b[..size_of::<u32>()]);

    unsafe { std::mem::transmute(buf) }
}

fn run(config: Config) {
    // tests::setup_logger();
    let tree = sled::Tree::start(config).unwrap();

    // flush to ensure the initial root is stable.
    // TODO is this necessary or voodoo?
    tree.flush().unwrap();

    thread::spawn(|| {
        let runtime = rand::thread_rng().gen_range(0, 200);
        thread::sleep(Duration::from_millis(runtime));
        unsafe {
            libc::raise(9);
        }
    });

    let (key, highest) = verify(&tree);

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let mut key = u32_to_vec((hu % CYCLE) as u32);
        key.reverse();

        let mut value = u32_to_vec((hu / CYCLE) as u32);
        let additional_len = rand::thread_rng().gen_range(0, 100_000);
        value.append(&mut vec![0u8; additional_len]);

        tree.set(key, value).unwrap();
    }
}

fn run_without_snapshot() {
    let config = ConfigBuilder::new()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes".to_string())
        .snapshot_after_ops(1 << 56)
        .build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

fn run_with_snapshot() {
    let config = ConfigBuilder::new()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes_with_snapshot".to_string())
        .snapshot_after_ops(1 << 10)
        .build();

    match thread::spawn(|| run(config)).join() {
        Err(e) => {
            println!("worker thread failed: {:?}", e);
            std::process::exit(15);
        }
        _ => {}
    }
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    cleanup_with_snapshots();
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(
                    child,
                    &mut status as *mut libc::c_int,
                    0,
                );
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup_with_snapshots();
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    cleanup();
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(
                    child,
                    &mut status as *mut libc::c_int,
                    0,
                );
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup();
}

fn cleanup_with_snapshots() {
    let dir = Path::new("test_crashes_with_snapshot");
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}

fn cleanup() {
    let dir = Path::new("test_crashes");
    if dir.exists() {
        fs::remove_dir_all(dir).unwrap();
    }
}
<|endoftext|>


<commit_before><commit_msg> Use raw identifiers for identifiers which collide with Rust keywords (#476)<commit_after><|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use num_format::{Locale, ToFormattedString};
use rand::{thread_rng, Rng};

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
    global_allocator
)]
#[cfg(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--flush-every=<ms>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --flush-every=<m>  Flush and sync the database every ms [default: 200].
";

#[derive(Clone, Copy)]
struct Args {
    threads: usize,
    burn_in: bool,
    duration: u64,
    key_len: usize,
    val_len: usize,
    get_prop: usize,
    set_prop: usize,
    del_prop: usize,
    cas_prop: usize,
    scan_prop: usize,
    merge_prop: usize,
    entries: usize,
    sequential: bool,
    total_ops: Option<usize>,
    flush_every: u64,
}

impl Default for Args {
    fn default() -> Args {
        Args {
            threads: 4,
            burn_in: false,
            duration: 10,
            key_len: 10,
            val_len: 100,
            get_prop: 94,
            set_prop: 2,
            del_prop: 1,
            cas_prop: 1,
            scan_prop: 1,
            merge_prop: 1,
            entries: 100000,
            sequential: false,
            total_ops: None,
            flush_every: 200,
        }
    }
}

fn parse<'a, I, T>(mut iter: I) -> T
where
    I: Iterator<Item = &'a str>,
    T: std::str::FromStr,
    <T as std::str::FromStr>::Err: std::fmt::Debug,
{
    iter.next().expect(USAGE).parse().expect(USAGE)
}

impl Args {
    fn parse() -> Args {
        let mut args = Args::default();
        for raw_arg in std::env::args().skip(1) {
            let mut splits = raw_arg[2..].split('=');
            match splits.next().unwrap() {
                "threads" => args.threads = parse(&mut splits),
                "burn-in" => args.burn_in = true,
                "duration" => args.duration = parse(&mut splits),
                "key-len" => args.key_len = parse(&mut splits),
                "val-len" => args.val_len = parse(&mut splits),
                "get-prop" => args.get_prop = parse(&mut splits),
                "set-prop" => args.set_prop = parse(&mut splits),
                "del-prop" => args.del_prop = parse(&mut splits),
                "cas-prop" => args.cas_prop = parse(&mut splits),
                "scan-prop" => args.scan_prop = parse(&mut splits),
                "merge-prop" => args.merge_prop = parse(&mut splits),
                "entries" => args.entries = parse(&mut splits),
                "sequential" => args.sequential = true,
                "total-ops" => args.total_ops = Some(parse(&mut splits)),
                "flush-every" => args.flush_every = parse(&mut splits),
                other => panic!("unknown option: {}, {}", other, USAGE),
            }
        }
        args
    }
}

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!(
            "did {} ops, {}mb RSS",
            (total - last).to_formatted_string(&Locale::en),
            rss() / (1024 * 1024)
        );

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(Vec::new);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(args: Args, tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let get_max = args.get_prop;
    let set_max = get_max + args.set_prop;
    let del_max = set_max + args.del_prop;
    let cas_max = del_max + args.cas_prop;
    let merge_max = cas_max + args.merge_prop;
    let scan_max = merge_max + args.scan_prop;

    let keygen = |len| -> sled::IVec {
        static SEQ: AtomicUsize = AtomicUsize::new(0);
        let i = if args.sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.entries;

        let start = if len < 8 { 8 - len } else { 0 };

        let i_keygen = &i.to_be_bytes()[start..];

        i_keygen.iter().cycle().take(len).copied().collect()
    };

    let valgen = |len| -> sled::IVec {
        if len == 0 {
            return vec![].into();
        }

        let i: usize = thread_rng().gen::<usize>() % (len * 8);

        let i_keygen = i.to_be_bytes();

        i_keygen
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };

    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = keygen(args.key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get(&key).unwrap();
            }
            v if v > get_max && v <= set_max => {
                let value = valgen(args.val_len);
                tree.insert(&key, value).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                let value = valgen(args.val_len);
                tree.merge(&key, value).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = Args::parse();

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(args.flush_every))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(args, tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.total_ops {
        assert!(!args.burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.burn_in {
        thread::sleep(std::time::Duration::from_secs(args.duration));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }
    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops.to_formatted_string(&Locale::en),
        time,
        ((ops * 1_000) / (time * 1_000)).to_formatted_string(&Locale::en)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if let Ok(env) = std::env::var("RUST_LOG") {
        builder.parse_filters(&env);
    }

    let _r = builder.try_init();
}
<commit_msg>Use zero-copy get in stress2 get bench<commit_after>use std::{
    sync::{
        atomic::{AtomicBool, AtomicUsize, Ordering},
        Arc,
    },
    thread,
};

use num_format::{Locale, ToFormattedString};
use rand::{thread_rng, Rng};

#[cfg_attr(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
    global_allocator
)]
#[cfg(
    // only enable jemalloc on linux and macos by default
    all(
        any(target_os = "linux", target_os = "macos"),
        feature = "jemalloc",
    ),
)]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

static TOTAL: AtomicUsize = AtomicUsize::new(0);

const USAGE: &str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>] \
    [--key-len=<l>] [--val-len=<l>] \
    [--get-prop=<p>] \
    [--set-prop=<p>] \
    [--del-prop=<p>] \
    [--cas-prop=<p>] \
    [--scan-prop=<p>] \
    [--merge-prop=<p>] \
    [--entries=<n>] \
    [--sequential] \
    [--total-ops=<n>] \
    [--flush-every=<ms>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
    --key-len=<l>      The length of keys [default: 10].
    --val-len=<l>      The length of values [default: 100].
    --get-prop=<p>     The relative proportion of get requests [default: 94].
    --set-prop=<p>     The relative proportion of set requests [default: 2].
    --del-prop=<p>     The relative proportion of del requests [default: 1].
    --cas-prop=<p>     The relative proportion of cas requests [default: 1].
    --scan-prop=<p>    The relative proportion of scan requests [default: 1].
    --merge-prop=<p>   The relative proportion of merge requests [default: 1].
    --entries=<n>      The total keyspace [default: 100000].
    --sequential       Run the test in sequential mode instead of random.
    --total-ops=<n>    Stop test after executing a total number of operations.
    --flush-every=<m>  Flush and sync the database every ms [default: 200].
";

#[derive(Clone, Copy)]
struct Args {
    threads: usize,
    burn_in: bool,
    duration: u64,
    key_len: usize,
    val_len: usize,
    get_prop: usize,
    set_prop: usize,
    del_prop: usize,
    cas_prop: usize,
    scan_prop: usize,
    merge_prop: usize,
    entries: usize,
    sequential: bool,
    total_ops: Option<usize>,
    flush_every: u64,
}

impl Default for Args {
    fn default() -> Args {
        Args {
            threads: 4,
            burn_in: false,
            duration: 10,
            key_len: 10,
            val_len: 100,
            get_prop: 94,
            set_prop: 2,
            del_prop: 1,
            cas_prop: 1,
            scan_prop: 1,
            merge_prop: 1,
            entries: 100000,
            sequential: false,
            total_ops: None,
            flush_every: 200,
        }
    }
}

fn parse<'a, I, T>(mut iter: I) -> T
where
    I: Iterator<Item = &'a str>,
    T: std::str::FromStr,
    <T as std::str::FromStr>::Err: std::fmt::Debug,
{
    iter.next().expect(USAGE).parse().expect(USAGE)
}

impl Args {
    fn parse() -> Args {
        let mut args = Args::default();
        for raw_arg in std::env::args().skip(1) {
            let mut splits = raw_arg[2..].split('=');
            match splits.next().unwrap() {
                "threads" => args.threads = parse(&mut splits),
                "burn-in" => args.burn_in = true,
                "duration" => args.duration = parse(&mut splits),
                "key-len" => args.key_len = parse(&mut splits),
                "val-len" => args.val_len = parse(&mut splits),
                "get-prop" => args.get_prop = parse(&mut splits),
                "set-prop" => args.set_prop = parse(&mut splits),
                "del-prop" => args.del_prop = parse(&mut splits),
                "cas-prop" => args.cas_prop = parse(&mut splits),
                "scan-prop" => args.scan_prop = parse(&mut splits),
                "merge-prop" => args.merge_prop = parse(&mut splits),
                "entries" => args.entries = parse(&mut splits),
                "sequential" => args.sequential = true,
                "total-ops" => args.total_ops = Some(parse(&mut splits)),
                "flush-every" => args.flush_every = parse(&mut splits),
                other => panic!("unknown option: {}, {}", other, USAGE),
            }
        }
        args
    }
}

fn report(shutdown: Arc<AtomicBool>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = TOTAL.load(Ordering::Acquire);

        println!(
            "did {} ops, {}mb RSS",
            (total - last).to_formatted_string(&Locale::en),
            rss() / (1024 * 1024)
        );

        last = total;
    }
}

fn concatenate_merge(
    _key: &[u8],              // the key being merged
    old_value: Option<&[u8]>, // the previous value, if one existed
    merged_bytes: &[u8],      // the new bytes being merged in
) -> Option<Vec<u8>> {
    // set the new value, return None to delete
    let mut ret = old_value.map(|ov| ov.to_vec()).unwrap_or_else(Vec::new);

    ret.extend_from_slice(merged_bytes);

    Some(ret)
}

fn run(args: Args, tree: Arc<sled::Db>, shutdown: Arc<AtomicBool>) {
    let get_max = args.get_prop;
    let set_max = get_max + args.set_prop;
    let del_max = set_max + args.del_prop;
    let cas_max = del_max + args.cas_prop;
    let merge_max = cas_max + args.merge_prop;
    let scan_max = merge_max + args.scan_prop;

    let keygen = |len| -> sled::IVec {
        static SEQ: AtomicUsize = AtomicUsize::new(0);
        let i = if args.sequential {
            SEQ.fetch_add(1, Ordering::Relaxed)
        } else {
            thread_rng().gen::<usize>()
        } % args.entries;

        let start = if len < 8 { 8 - len } else { 0 };

        let i_keygen = &i.to_be_bytes()[start..];

        i_keygen.iter().cycle().take(len).copied().collect()
    };

    let valgen = |len| -> sled::IVec {
        if len == 0 {
            return vec![].into();
        }

        let i: usize = thread_rng().gen::<usize>() % (len * 8);

        let i_keygen = i.to_be_bytes();

        i_keygen
            .iter()
            .skip_while(|v| **v == 0)
            .cycle()
            .take(len)
            .copied()
            .collect()
    };

    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        let op = TOTAL.fetch_add(1, Ordering::Release);
        let key = keygen(args.key_len);
        let choice = rng.gen_range(0, scan_max + 1);

        match choice {
            v if v <= get_max => {
                tree.get_zero_copy(&key, |_| {}).unwrap();
            }
            v if v > get_max && v <= set_max => {
                let value = valgen(args.val_len);
                tree.insert(&key, value).unwrap();
            }
            v if v > set_max && v <= del_max => {
                tree.remove(&key).unwrap();
            }
            v if v > del_max && v <= cas_max => {
                let old = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                let new = if rng.gen::<bool>() {
                    let value = valgen(args.val_len);
                    Some(value)
                } else {
                    None
                };

                if let Err(e) = tree.compare_and_swap(&key, old, new) {
                    panic!("operational error: {:?}", e);
                }
            }
            v if v > cas_max && v <= merge_max => {
                let value = valgen(args.val_len);
                tree.merge(&key, value).unwrap();
            }
            _ => {
                let iter = tree.range(key..).map(|res| res.unwrap());

                if op % 2 == 0 {
                    let _ = iter.take(rng.gen_range(0, 15)).collect::<Vec<_>>();
                } else {
                    let _ = iter
                        .rev()
                        .take(rng.gen_range(0, 15))
                        .collect::<Vec<_>>();
                }
            }
        }
    }
}

fn rss() -> usize {
    #[cfg(target_os = "linux")]
    {
        use std::io::prelude::*;
        use std::io::BufReader;

        let mut buf = String::new();
        let mut f =
            BufReader::new(std::fs::File::open("/proc/self/statm").unwrap());
        f.read_line(&mut buf).unwrap();
        let mut parts = buf.split_whitespace();
        let rss_pages = parts.nth(1).unwrap().parse::<usize>().unwrap();
        rss_pages * 4096
    }
    #[cfg(not(target_os = "linux"))]
    {
        0
    }
}

fn main() {
    #[cfg(feature = "logging")]
    setup_logger();

    let args = Args::parse();

    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::Config::new()
        .cache_capacity(256 * 1024 * 1024)
        .flush_every_ms(Some(args.flush_every))
        .print_profile_on_drop(true);

    let tree = Arc::new(config.open().unwrap());
    tree.set_merge_operator(concatenate_merge);

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.threads;

    for i in 0..=n_threads {
        let tree = tree.clone();
        let shutdown = shutdown.clone();

        let t = if i == 0 {
            thread::Builder::new()
                .name("reporter".into())
                .spawn(move || report(shutdown))
                .unwrap()
        } else {
            thread::spawn(move || run(args, tree, shutdown))
        };

        threads.push(t);
    }

    if let Some(ops) = args.total_ops {
        assert!(!args.burn_in, "don't set both --burn-in and --total-ops");
        while TOTAL.load(Ordering::Relaxed) < ops {
            thread::sleep(std::time::Duration::from_millis(50));
        }
        shutdown.store(true, Ordering::SeqCst);
    } else if !args.burn_in {
        thread::sleep(std::time::Duration::from_secs(args.duration));
        shutdown.store(true, Ordering::SeqCst);
    }

    for t in threads.into_iter() {
        t.join().unwrap();
    }
    let ops = TOTAL.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops.to_formatted_string(&Locale::en),
        time,
        ((ops * 1_000) / (time * 1_000)).to_formatted_string(&Locale::en)
    );
}

#[cfg(feature = "logging")]
pub fn setup_logger() {
    use std::io::Write;

    color_backtrace::install();

    fn tn() -> String {
        std::thread::current().name().unwrap_or("unknown").to_owned()
    }

    let mut builder = env_logger::Builder::new();
    builder
        .format(|buf, record| {
            writeln!(
                buf,
                "{:05} {:25} {:10} {}",
                record.level(),
                tn(),
                record.module_path().unwrap().split("::").last().unwrap(),
                record.args()
            )
        })
        .filter(None, log::LevelFilter::Info);

    if let Ok(env) = std::env::var("RUST_LOG") {
        builder.parse_filters(&env);
    }

    let _r = builder.try_init();
}
<|endoftext|>


<commit_before>use std::fmt::Debug;
use std::marker::PhantomData;
use super::Alphabet;
use expression::{Expression, Condition, ExpressionError};

#[derive(Debug, Eq, PartialEq)]
pub enum RuleError {
    SymbolMismatch,
    RuleArityMismatch,
    ArityMismatch,
    ConditionFalse,
    ConditionFailed,
    ExprFailed(ExpressionError),
}

pub trait ParametricSymbol: Clone + PartialEq + Debug
{
    type Sym: Alphabet;
    type Param: Clone + Debug + PartialEq;

    fn symbol(&self) -> &Self::Sym;
    fn symbol_mut(&mut self) -> &mut Self::Sym;

    fn params(&self) -> &[Self::Param];
    fn params_mut(&mut self) -> &mut [Self::Param];

    /// Construct a new ParametricSymbol. If the iterator contains a wrong
    /// number of parameters, return None.
    fn new_from_result_iter<I, E>(symbol: Self::Sym, iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>;

    fn new_from_iter<I>(symbol: Self::Sym, iter: I) -> Option<Self>
        where I: Iterator<Item = Self::Param>
    {
        match Self::new_from_result_iter::<_, ()>(symbol, iter.map(|i| Ok(i))) {
            Some(Ok(res)) => Some(res),
            Some(Err(())) => panic!(),
            None => None,
        }
    }

    fn new_from_vec(symbol: Self::Sym, vec: Vec<Self::Param>) -> Option<Self> {
        Self::new_from_iter(symbol, vec.into_iter())
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct PSym<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: Vec<Param>,
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let mut params = Vec::with_capacity(iter.size_hint().0);
        for p in iter {
            match p {
                Ok(p) => params.push(p),
                Err(e) => return Some(Err(e)),
            }
        }
        Some(Ok(PSym {
            symbol: symbol,
            params: params,
        }))
    }
}

#[derive(Debug, PartialEq)]
pub struct PSym1<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: [Param; 1],
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> Clone for PSym1<Sym, Param> {
    fn clone(&self) -> Self {
        PSym1 {
            symbol: self.symbol.clone(),
            params: [self.params[0].clone()],
        }
    }
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym1<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, mut iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let p1 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        if let Some(_) = iter.next() {
            return None;
        }

        Some(Ok(PSym1 {
            symbol: symbol,
            params: [p1],
        }))
    }
}

#[derive(Debug, PartialEq)]
pub struct PSym2<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: [Param; 2],
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> Clone for PSym2<Sym, Param> {
    fn clone(&self) -> Self {
        PSym2 {
            symbol: self.symbol.clone(),
            params: [self.params[0].clone(), self.params[1].clone()],
        }
    }
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym2<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, mut iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let p1 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        let p2 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        if let Some(_) = iter.next() {
            return None;
        }

        Some(Ok(PSym2 {
            symbol: symbol,
            params: [p1, p2],
        }))
    }
}

pub trait ParametricRule: Clone + Debug
{
    type InSym: ParametricSymbol;
    type OutSym: ParametricSymbol;

    fn apply(&self, psym: &Self::OutSym) -> Result<Vec<Self::OutSym>, RuleError>;
}

#[derive(Debug, Clone)]
pub struct PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    pub symbol: Sym,
    pub condition: C,
    pub production: Vec<PS>,
    pub arity: usize,
    _marker: PhantomData<PS2>,
}

impl<Sym, PS, PS2, C> PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    pub fn new(sym: Sym, cond: C, prod: Vec<PS>, arity: usize) -> PRule<Sym, PS, PS2, C> {
        PRule {
            symbol: sym,
            condition: cond,
            production: prod,
            arity: arity,
            _marker: PhantomData,
        }
    }
}

impl<Sym, PS, PS2, C> ParametricRule for PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    type InSym = PS;
    type OutSym = PS2;

    /// Tries to apply the rule and if applicable, produces a successor.
    fn apply(&self, psym: &PS2) -> Result<Vec<PS2>, RuleError> {
        if self.arity != psym.params().len() {
            Err(RuleError::RuleArityMismatch)
        } else if self.symbol.eq(psym.symbol()) {
            match self.condition.evaluate(psym.params()) {
                Ok(true) => {
                    let mut new_symbol_string = Vec::with_capacity(self.production.len());

                    // evaluate all parameters of all symbols in the production
                    for prod_sym in self.production.iter() {
                        match PS2::new_from_result_iter(prod_sym.symbol().clone(),
                                                        prod_sym.params().iter().map(|expr| {
                                                            expr.evaluate(psym.params())
                                                        })) {
                            Some(Ok(new_sym)) => {
                                new_symbol_string.push(new_sym);
                            }
                            Some(Err(e)) => {
                                return Err(RuleError::ExprFailed(e));
                            }
                            None => {
                                return Err(RuleError::ArityMismatch);
                            }
                        }
                    }

                    Ok(new_symbol_string)
                }
                Ok(false) => Err(RuleError::ConditionFalse),
                _ => Err(RuleError::ConditionFailed),
            }
        } else {
            Err(RuleError::SymbolMismatch)
        }
    }
}

#[test]
fn test_rule_apply() {
    use expression_num::NumExpr;
    use expression::cond::Cond;
    let expr_s = PSym::new_from_vec('P', vec![NumExpr::Const(123u32)]).unwrap();

    let rule = PRule::<_, PSym<_, NumExpr<u32>>, PSym<_, u32>, _>::new('A',
                                                                       Cond::True,
                                                                       vec![expr_s.clone()],
                                                                       1);

    let param_s = PSym::new_from_vec('P', vec![123u32]).unwrap();
    assert_eq!(Err(RuleError::SymbolMismatch), rule.apply(&param_s));

    let param_s = PSym::new_from_vec('A', vec![123u32]).unwrap();
    let result_s = PSym::new_from_vec('P', vec![123u32]).unwrap();
    assert_eq!(Ok(vec![result_s]), rule.apply(&param_s));

    let rule = PRule::<_, PSym<_, NumExpr<u32>>, PSym<_, u32>, _>::new('A',
                                                                       Cond::False,
                                                                       vec![expr_s.clone()],
                                                                       1);
    assert_eq!(Err(RuleError::ConditionFalse), rule.apply(&param_s));
}

pub trait ParametricSystem {
    type Rule: ParametricRule;

    fn apply_first_rule(&self, sym: &<Self::Rule as ParametricRule>::OutSym) -> Option<Vec<<Self::Rule as ParametricRule>::OutSym>>;

    /// Apply in parallel the first matching rule to each symbol in the string.
    /// Returns the total number of rule applications.
    fn develop_next(&self, axiom: &Vec<<Self::Rule as ParametricRule>::OutSym>) -> (Vec<<Self::Rule as ParametricRule>::OutSym>, usize) {
        let mut expanded = Vec::new();
        let mut rule_applications = 0;

        for sym in axiom.iter() {
            match self.apply_first_rule(sym) {
                Some(successor) => {
                    expanded.extend(successor);
                    rule_applications += 1;
                    // XXX: Count rule applications of the matching rule.
                }

                None => {
                    expanded.push(sym.clone());
                }
            }
        }

        (expanded, rule_applications)
    }

    /// Develop the system starting with `axiom` up to `max_iterations`. Return iteration count.
    fn develop(&self,
               axiom: Vec<<Self::Rule as ParametricRule>::OutSym>,
               max_iterations: usize)
               -> (Vec<<Self::Rule as ParametricRule>::OutSym>, usize) {
        let mut current = axiom;

        for iter in 0..max_iterations {
            let (next, rule_applications) = self.develop_next(&current);
            if rule_applications == 0 {
                return (current, iter);
            }
            current = next;
        }
        return (current, max_iterations);
    }
}

/// Apply first matching rule and return expanded successor.
pub fn apply_first_rule<R>(rules: &[R],
                           sym: &<R as ParametricRule>::OutSym)
                           -> Option<Vec<<R as ParametricRule>::OutSym>>
    where R: ParametricRule
{
    for rule in rules {
        if let Ok(successor) = rule.apply(sym) {
            return Some(successor);
        }
    }
    return None;
}

#[derive(Debug, Clone)]
pub struct PSystem<R>
    where R: ParametricRule
{
    rules: Vec<R>,
}

impl<R> PSystem<R> where R: ParametricRule
{
    pub fn new() -> PSystem<R> {
        PSystem { rules: vec![] }
    }

    pub fn add_rule(&mut self, rule: R) {
        self.rules.push(rule);
    }
}

impl<R> ParametricSystem for PSystem<R> where R: ParametricRule
{
    type Rule = R;

    /// Apply first matching rule and return expanded successor.
    fn apply_first_rule(&self, sym: &<Self::Rule as ParametricRule>::OutSym) -> Option<Vec<<Self::Rule as ParametricRule>::OutSym>> {
        apply_first_rule(&self.rules, sym)
    }
}
<commit_msg>Add a way to retrieve the symbol of a ParametricRule<commit_after>use std::fmt::Debug;
use std::marker::PhantomData;
use super::Alphabet;
use expression::{Expression, Condition, ExpressionError};

#[derive(Debug, Eq, PartialEq)]
pub enum RuleError {
    SymbolMismatch,
    RuleArityMismatch,
    ArityMismatch,
    ConditionFalse,
    ConditionFailed,
    ExprFailed(ExpressionError),
}

pub trait ParametricSymbol: Clone + PartialEq + Debug
{
    type Sym: Alphabet;
    type Param: Clone + Debug + PartialEq;

    fn symbol(&self) -> &Self::Sym;
    fn symbol_mut(&mut self) -> &mut Self::Sym;

    fn params(&self) -> &[Self::Param];
    fn params_mut(&mut self) -> &mut [Self::Param];

    /// Construct a new ParametricSymbol. If the iterator contains a wrong
    /// number of parameters, return None.
    fn new_from_result_iter<I, E>(symbol: Self::Sym, iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>;

    fn new_from_iter<I>(symbol: Self::Sym, iter: I) -> Option<Self>
        where I: Iterator<Item = Self::Param>
    {
        match Self::new_from_result_iter::<_, ()>(symbol, iter.map(|i| Ok(i))) {
            Some(Ok(res)) => Some(res),
            Some(Err(())) => panic!(),
            None => None,
        }
    }

    fn new_from_vec(symbol: Self::Sym, vec: Vec<Self::Param>) -> Option<Self> {
        Self::new_from_iter(symbol, vec.into_iter())
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct PSym<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: Vec<Param>,
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let mut params = Vec::with_capacity(iter.size_hint().0);
        for p in iter {
            match p {
                Ok(p) => params.push(p),
                Err(e) => return Some(Err(e)),
            }
        }
        Some(Ok(PSym {
            symbol: symbol,
            params: params,
        }))
    }
}

#[derive(Debug, PartialEq)]
pub struct PSym1<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: [Param; 1],
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> Clone for PSym1<Sym, Param> {
    fn clone(&self) -> Self {
        PSym1 {
            symbol: self.symbol.clone(),
            params: [self.params[0].clone()],
        }
    }
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym1<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, mut iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let p1 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        if let Some(_) = iter.next() {
            return None;
        }

        Some(Ok(PSym1 {
            symbol: symbol,
            params: [p1],
        }))
    }
}

#[derive(Debug, PartialEq)]
pub struct PSym2<Sym: Alphabet, Param: Clone + Debug + PartialEq> {
    symbol: Sym,
    params: [Param; 2],
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> Clone for PSym2<Sym, Param> {
    fn clone(&self) -> Self {
        PSym2 {
            symbol: self.symbol.clone(),
            params: [self.params[0].clone(), self.params[1].clone()],
        }
    }
}

impl<Sym: Alphabet, Param: Clone + Debug + PartialEq> ParametricSymbol for PSym2<Sym, Param> {
    type Sym = Sym;
    type Param = Param;

    fn symbol(&self) -> &Self::Sym {
        &self.symbol
    }
    fn symbol_mut(&mut self) -> &mut Self::Sym {
        &mut self.symbol
    }
    fn params(&self) -> &[Self::Param] {
        &self.params
    }
    fn params_mut(&mut self) -> &mut [Self::Param] {
        &mut self.params
    }

    fn new_from_result_iter<I, E>(symbol: Self::Sym, mut iter: I) -> Option<Result<Self, E>>
        where I: Iterator<Item = Result<Self::Param, E>>
    {
        let p1 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        let p2 = match iter.next() {
            Some(Ok(p)) => p,
            Some(Err(e)) => return Some(Err(e)),
            None => return None,
        };
        if let Some(_) = iter.next() {
            return None;
        }

        Some(Ok(PSym2 {
            symbol: symbol,
            params: [p1, p2],
        }))
    }
}

pub trait ParametricRule: Clone + Debug
{
    type InSym: ParametricSymbol;
    type OutSym: ParametricSymbol;

    fn apply(&self, psym: &Self::OutSym) -> Result<Vec<Self::OutSym>, RuleError>;

    fn symbol(&self) -> &<Self::InSym as ParametricSymbol>::Sym;
}

#[derive(Debug, Clone)]
pub struct PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    pub symbol: Sym,
    pub condition: C,
    pub production: Vec<PS>,
    pub arity: usize,
    _marker: PhantomData<PS2>,
}

impl<Sym, PS, PS2, C> PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    pub fn new(sym: Sym, cond: C, prod: Vec<PS>, arity: usize) -> PRule<Sym, PS, PS2, C> {
        PRule {
            symbol: sym,
            condition: cond,
            production: prod,
            arity: arity,
            _marker: PhantomData,
        }
    }
}

impl<Sym, PS, PS2, C> ParametricRule for PRule<Sym, PS, PS2, C>
    where Sym: Alphabet,
          PS: ParametricSymbol<Sym = Sym, Param = C::Expr>,
          PS2: ParametricSymbol<Sym = Sym, Param = <C::Expr as Expression>::Element>,
          C: Condition
{
    type InSym = PS;
    type OutSym = PS2;

    /// Tries to apply the rule and if applicable, produces a successor.
    fn apply(&self, psym: &PS2) -> Result<Vec<PS2>, RuleError> {
        if self.arity != psym.params().len() {
            Err(RuleError::RuleArityMismatch)
        } else if self.symbol.eq(psym.symbol()) {
            match self.condition.evaluate(psym.params()) {
                Ok(true) => {
                    let mut new_symbol_string = Vec::with_capacity(self.production.len());

                    // evaluate all parameters of all symbols in the production
                    for prod_sym in self.production.iter() {
                        match PS2::new_from_result_iter(prod_sym.symbol().clone(),
                                                        prod_sym.params().iter().map(|expr| {
                                                            expr.evaluate(psym.params())
                                                        })) {
                            Some(Ok(new_sym)) => {
                                new_symbol_string.push(new_sym);
                            }
                            Some(Err(e)) => {
                                return Err(RuleError::ExprFailed(e));
                            }
                            None => {
                                return Err(RuleError::ArityMismatch);
                            }
                        }
                    }

                    Ok(new_symbol_string)
                }
                Ok(false) => Err(RuleError::ConditionFalse),
                _ => Err(RuleError::ConditionFailed),
            }
        } else {
            Err(RuleError::SymbolMismatch)
        }
    }

    fn symbol(&self) -> &<Self::InSym as ParametricSymbol>::Sym {
        &self.symbol
    }
}

#[test]
fn test_rule_apply() {
    use expression_num::NumExpr;
    use expression::cond::Cond;
    let expr_s = PSym::new_from_vec('P', vec![NumExpr::Const(123u32)]).unwrap();

    let rule = PRule::<_, PSym<_, NumExpr<u32>>, PSym<_, u32>, _>::new('A',
                                                                       Cond::True,
                                                                       vec![expr_s.clone()],
                                                                       1);

    let param_s = PSym::new_from_vec('P', vec![123u32]).unwrap();
    assert_eq!(Err(RuleError::SymbolMismatch), rule.apply(&param_s));

    let param_s = PSym::new_from_vec('A', vec![123u32]).unwrap();
    let result_s = PSym::new_from_vec('P', vec![123u32]).unwrap();
    assert_eq!(Ok(vec![result_s]), rule.apply(&param_s));

    let rule = PRule::<_, PSym<_, NumExpr<u32>>, PSym<_, u32>, _>::new('A',
                                                                       Cond::False,
                                                                       vec![expr_s.clone()],
                                                                       1);
    assert_eq!(Err(RuleError::ConditionFalse), rule.apply(&param_s));
}

pub trait ParametricSystem {
    type Rule: ParametricRule;

    fn apply_first_rule(&self, sym: &<Self::Rule as ParametricRule>::OutSym) -> Option<Vec<<Self::Rule as ParametricRule>::OutSym>>;

    /// Apply in parallel the first matching rule to each symbol in the string.
    /// Returns the total number of rule applications.
    fn develop_next(&self, axiom: &Vec<<Self::Rule as ParametricRule>::OutSym>) -> (Vec<<Self::Rule as ParametricRule>::OutSym>, usize) {
        let mut expanded = Vec::new();
        let mut rule_applications = 0;

        for sym in axiom.iter() {
            match self.apply_first_rule(sym) {
                Some(successor) => {
                    expanded.extend(successor);
                    rule_applications += 1;
                    // XXX: Count rule applications of the matching rule.
                }

                None => {
                    expanded.push(sym.clone());
                }
            }
        }

        (expanded, rule_applications)
    }

    /// Develop the system starting with `axiom` up to `max_iterations`. Return iteration count.
    fn develop(&self,
               axiom: Vec<<Self::Rule as ParametricRule>::OutSym>,
               max_iterations: usize)
               -> (Vec<<Self::Rule as ParametricRule>::OutSym>, usize) {
        let mut current = axiom;

        for iter in 0..max_iterations {
            let (next, rule_applications) = self.develop_next(&current);
            if rule_applications == 0 {
                return (current, iter);
            }
            current = next;
        }
        return (current, max_iterations);
    }
}

/// Apply first matching rule and return expanded successor.
pub fn apply_first_rule<R>(rules: &[R],
                           sym: &<R as ParametricRule>::OutSym)
                           -> Option<Vec<<R as ParametricRule>::OutSym>>
    where R: ParametricRule
{
    for rule in rules {
        if let Ok(successor) = rule.apply(sym) {
            return Some(successor);
        }
    }
    return None;
}

#[derive(Debug, Clone)]
pub struct PSystem<R>
    where R: ParametricRule
{
    rules: Vec<R>,
}

impl<R> PSystem<R> where R: ParametricRule
{
    pub fn new() -> PSystem<R> {
        PSystem { rules: vec![] }
    }

    pub fn add_rule(&mut self, rule: R) {
        self.rules.push(rule);
    }
}

impl<R> ParametricSystem for PSystem<R> where R: ParametricRule
{
    type Rule = R;

    /// Apply first matching rule and return expanded successor.
    fn apply_first_rule(&self, sym: &<Self::Rule as ParametricRule>::OutSym) -> Option<Vec<<Self::Rule as ParametricRule>::OutSym>> {
        apply_first_rule(&self.rules, sym)
    }
}
<|endoftext|>


<commit_before>//! Parser module is responsible for parsing FITS files.

named!(fits<&[u8], (Vec<&[u8]>, Vec<&[u8]>) >,
       pair!(primary_header, many0!( take!(2880) )));

named!(primary_header<&[u8], Vec<&[u8]> >, count!(take!(2880), 2));

named!(a_header<&[u8], &[u8]>, take!(80));

named!(end_header<&[u8],(&[u8], Vec<&[u8]>)>, pair!(tag!("END"), count!(tag!(" "), 77)));

named!(blank_header<&[u8],Vec<&[u8]> >, count!(tag!(" "), 80));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::{fits, end_header, blank_header, a_header};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, (header, blocks)) => {
                assert_eq!(header.len(), 2);
                assert_eq!(blocks.len(), 3675);
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn a_header_should_parse_a_header(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = a_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn end_header_should_parse_an_END_header(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn blank_header_should_parse_a_BLANK_header(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<commit_msg>Parse known amounts headers of primary header<commit_after>//! Parser module is responsible for parsing FITS files.

named!(fits<&[u8], ((Vec<&[u8]>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >), Vec<&[u8]>) >,
       pair!(primary_header, many0!( take!(2880) )));

named!(primary_header<&[u8], (Vec<&[u8]>, (&[u8], Vec<&[u8]>), Vec<Vec<&[u8]> >)>,
       tuple!(
           count!(a_header, 55),
           end_header,
           count!(blank_header, 16)
       ));

named!(a_header<&[u8], &[u8]>, take!(80));

named!(end_header<&[u8],(&[u8], Vec<&[u8]>)>, pair!(tag!("END"), count!(tag!(" "), 77)));

named!(blank_header<&[u8],Vec<&[u8]> >, count!(tag!(" "), 80));

#[cfg(test)]
mod tests {
    use nom::{IResult};
    use super::{fits, primary_header, end_header, blank_header, a_header};

    #[test]
    fn it_should_parse_a_fits_file(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = fits(data);

        match result {
            IResult::Done(_, (header, blocks)) => {
                assert_eq!(blocks.len(), 3675);
            },
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn primary_header_should_parse_a_primary_header(){
        let data = include_bytes!("../../assets/images/k2-trappist1-unofficial-tpf-long-cadence.fits");

        let result = primary_header(&data[0..(2*2880)]);

        match result {
            IResult::Done(_, _) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn a_header_should_parse_a_header(){
        let data = "OBJECT  = 'EPIC 200164267'     / string version of target id                    "
            .as_bytes();

        let result = a_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn end_header_should_parse_an_END_header(){
        let data = "END                                                                             "
            .as_bytes();

        let result = end_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }

    #[test]
    fn blank_header_should_parse_a_BLANK_header(){
        let data = "                                                                                "
            .as_bytes();

        let result = blank_header(data);

        match result {
            IResult::Done(_,_) => assert!(true),
            IResult::Error(_) => panic!("Did not expect an error"),
            IResult::Incomplete(_) => panic!("Did not expect to be incomplete")
        }
    }
}
<|endoftext|>


<commit_before>use rustc::lint::*;
use rustc::middle::const_val::ConstVal;
use rustc::ty;
use rustc_const_eval::ConstContext;
use rustc_const_math::{ConstUsize,ConstInt};
use rustc::hir;
use syntax::ast::RangeLimits;
use utils::{self, higher};

/// **What it does:** Checks for out of bounds array indexing with a constant index.
///
/// **Why is this bad?** This will always panic at runtime.
///
/// **Known problems:** Hopefully none.
///
/// **Example:**
/// ```rust
/// let x = [1,2,3,4];
/// ...
/// x[9];
/// &x[2..9];
/// ```
declare_lint! {
    pub OUT_OF_BOUNDS_INDEXING,
    Deny,
    "out of bounds constant indexing"
}

/// **What it does:** Checks for usage of indexing or slicing.
///
/// **Why is this bad?** Usually, this can be safely allowed. However, in some
/// domains such as kernel development, a panic can cause the whole operating
/// system to crash.
///
/// **Known problems:** Hopefully none.
///
/// **Example:**
/// ```rust
/// ...
/// x[2];
/// &x[0..2];
/// ```
declare_restriction_lint! {
    pub INDEXING_SLICING,
    "indexing/slicing usage"
}

#[derive(Copy,Clone)]
pub struct ArrayIndexing;

impl LintPass for ArrayIndexing {
    fn get_lints(&self) -> LintArray {
        lint_array!(INDEXING_SLICING, OUT_OF_BOUNDS_INDEXING)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for ArrayIndexing {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx hir::Expr) {
        if let hir::ExprIndex(ref array, ref index) = e.node {
            // Array with known size can be checked statically
            let ty = cx.tables.expr_ty(array);
            if let ty::TyArray(_, size) = ty.sty {
                let size = ConstInt::Usize(ConstUsize::new(size as u64, cx.sess().target.uint_type).unwrap());
                let constcx = ConstContext::with_tables(cx.tcx, cx.tables);

                // Index is a constant uint
                let const_index = constcx.eval(index);
                if let Ok(ConstVal::Integral(const_index)) = const_index {
                    if size <= const_index {
                        utils::span_lint(cx, OUT_OF_BOUNDS_INDEXING, e.span, "const index is out of bounds");
                    }

                    return;
                }

                // Index is a constant range
                if let Some(range) = higher::range(index) {
                    let start = range.start
                        .map(|start| constcx.eval(start))
                        .map(|v| v.ok());
                    let end = range.end
                        .map(|end| constcx.eval(end))
                        .map(|v| v.ok());

                    if let Some((start, end)) = to_const_range(&start, &end, range.limits, size) {
                        if start > size || end > size {
                            utils::span_lint(cx, OUT_OF_BOUNDS_INDEXING, e.span, "range is out of bounds");
                        }
                        return;
                    }
                }
            }

            if let Some(range) = higher::range(index) {
                // Full ranges are always valid
                if range.start.is_none() && range.end.is_none() {
                    return;
                }

                // Impossible to know if indexing or slicing is correct
                utils::span_lint(cx, INDEXING_SLICING, e.span, "slicing may panic");
            } else {
                utils::span_lint(cx, INDEXING_SLICING, e.span, "indexing may panic");
            }
        }
    }
}

/// Returns an option containing a tuple with the start and end (exclusive) of the range.
fn to_const_range(
    start: &Option<Option<ConstVal>>,
    end: &Option<Option<ConstVal>>,
    limits: RangeLimits,
    array_size: ConstInt
) -> Option<(ConstInt, ConstInt)> {
    let start = match *start {
        Some(Some(ConstVal::Integral(x))) => x,
        Some(_) => return None,
        None => ConstInt::U8(0),
    };

    let end = match *end {
        Some(Some(ConstVal::Integral(x))) => {
            if limits == RangeLimits::Closed {
                (x + ConstInt::U8(1)).expect("such a big array is not realistic")
            } else {
                x
            }
        },
        Some(_) => return None,
        None => array_size,
    };

    Some((start, end))
}
<commit_msg>Use expect instead of unwrap to make clippy happy<commit_after>use rustc::lint::*;
use rustc::middle::const_val::ConstVal;
use rustc::ty;
use rustc_const_eval::ConstContext;
use rustc_const_math::{ConstUsize,ConstInt};
use rustc::hir;
use syntax::ast::RangeLimits;
use utils::{self, higher};

/// **What it does:** Checks for out of bounds array indexing with a constant index.
///
/// **Why is this bad?** This will always panic at runtime.
///
/// **Known problems:** Hopefully none.
///
/// **Example:**
/// ```rust
/// let x = [1,2,3,4];
/// ...
/// x[9];
/// &x[2..9];
/// ```
declare_lint! {
    pub OUT_OF_BOUNDS_INDEXING,
    Deny,
    "out of bounds constant indexing"
}

/// **What it does:** Checks for usage of indexing or slicing.
///
/// **Why is this bad?** Usually, this can be safely allowed. However, in some
/// domains such as kernel development, a panic can cause the whole operating
/// system to crash.
///
/// **Known problems:** Hopefully none.
///
/// **Example:**
/// ```rust
/// ...
/// x[2];
/// &x[0..2];
/// ```
declare_restriction_lint! {
    pub INDEXING_SLICING,
    "indexing/slicing usage"
}

#[derive(Copy,Clone)]
pub struct ArrayIndexing;

impl LintPass for ArrayIndexing {
    fn get_lints(&self) -> LintArray {
        lint_array!(INDEXING_SLICING, OUT_OF_BOUNDS_INDEXING)
    }
}

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for ArrayIndexing {
    fn check_expr(&mut self, cx: &LateContext<'a, 'tcx>, e: &'tcx hir::Expr) {
        if let hir::ExprIndex(ref array, ref index) = e.node {
            // Array with known size can be checked statically
            let ty = cx.tables.expr_ty(array);
            if let ty::TyArray(_, size) = ty.sty {
                let size = ConstInt::Usize(ConstUsize::new(size as u64, cx.sess().target.uint_type).expect("array size is invalid"));
                let constcx = ConstContext::with_tables(cx.tcx, cx.tables);

                // Index is a constant uint
                let const_index = constcx.eval(index);
                if let Ok(ConstVal::Integral(const_index)) = const_index {
                    if size <= const_index {
                        utils::span_lint(cx, OUT_OF_BOUNDS_INDEXING, e.span, "const index is out of bounds");
                    }

                    return;
                }

                // Index is a constant range
                if let Some(range) = higher::range(index) {
                    let start = range.start
                        .map(|start| constcx.eval(start))
                        .map(|v| v.ok());
                    let end = range.end
                        .map(|end| constcx.eval(end))
                        .map(|v| v.ok());

                    if let Some((start, end)) = to_const_range(&start, &end, range.limits, size) {
                        if start > size || end > size {
                            utils::span_lint(cx, OUT_OF_BOUNDS_INDEXING, e.span, "range is out of bounds");
                        }
                        return;
                    }
                }
            }

            if let Some(range) = higher::range(index) {
                // Full ranges are always valid
                if range.start.is_none() && range.end.is_none() {
                    return;
                }

                // Impossible to know if indexing or slicing is correct
                utils::span_lint(cx, INDEXING_SLICING, e.span, "slicing may panic");
            } else {
                utils::span_lint(cx, INDEXING_SLICING, e.span, "indexing may panic");
            }
        }
    }
}

/// Returns an option containing a tuple with the start and end (exclusive) of the range.
fn to_const_range(
    start: &Option<Option<ConstVal>>,
    end: &Option<Option<ConstVal>>,
    limits: RangeLimits,
    array_size: ConstInt
) -> Option<(ConstInt, ConstInt)> {
    let start = match *start {
        Some(Some(ConstVal::Integral(x))) => x,
        Some(_) => return None,
        None => ConstInt::U8(0),
    };

    let end = match *end {
        Some(Some(ConstVal::Integral(x))) => {
            if limits == RangeLimits::Closed {
                (x + ConstInt::U8(1)).expect("such a big array is not realistic")
            } else {
                x
            }
        },
        Some(_) => return None,
        None => array_size,
    };

    Some((start, end))
}
<|endoftext|>