<commit_before>extern crate bincode;
extern crate rustc_serialize;
extern crate rand;

mod wal_file;

use wal_file::{KeyValuePair, WALFile, WALIterator};

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use rustc_serialize::{Encodable, Decodable};

use std::cmp::max;
use std::convert::From;
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::{Read, Write, Seek, SeekFrom, ErrorKind};
use std::mem::{size_of};
use std::str;

const NUM_CHILDREN: usize = 32;
const FILE_HEADER: &'static str = "B+Tree\0";
const CURRENT_VERSION: u8 = 0x01;

// specify the types for the keys & values
pub trait KeyType: Ord + Encodable + Decodable {}
pub trait ValueType: Ord + Encodable + Decodable {}

// provide generic implementations
impl<T> KeyType for T where T: Ord + Encodable + Decodable {}
impl<T> ValueType for T where T: Ord + Encodable + Decodable {}

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
enum Payload<K: KeyType, V: ValueType> {
        Value(V),
        Children([(K,u64); NUM_CHILDREN]),
    }

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
struct Node<K: KeyType, V: ValueType> {
    key: K,
    parent: u64,
    payload: Payload<K,V>, // either children, or actual values
}

/// This struct represents an on-disk B+Tree. There are NUM_CHILDREN keys at each
/// level in the tree. The on-disk format is as follows where VV is the version
/// number:
/// |-------------------------------------------|
/// | 0x42 0x2b 0x54 0x72 | 0x65 0x65 0x00 0xVV |
/// | B    +    T    r    | e    e    \0   0xVV |
/// |-------------------------------------------|
/// | smallest record in bincode format         |
/// |-------------------------------------------|
/// | ...                                       |
/// |-------------------------------------------|
/// | largest record in bincode format          |
/// |-------------------------------------------|
/// | internal nodes ...                        |
/// |-------------------------------------------|
/// | root node                                 |
/// |-------------------------------------------|
pub struct BTree<K: KeyType, V: ValueType> {
    tree_file_path: String,         // the path to the tree file
    tree_file: File,                // the file backing the whole thing
    wal_file: WALFile<K,V>,              // write-ahead log for in-memory items
    root: Option<Node<K,V>>,        // optional in-memory copy of the root node
    max_key_size: usize,            // the size of the key in bytes
    max_value_size: usize,          // the size of the value in bytes
    mem_tree: BTreeMap<K, BTreeSet<V>>,  // the in-memory BTree that gets merged with the on-disk one
}

impl <K: KeyType, V: ValueType> BTree<K, V> {
    pub fn new(tree_file_path: String, max_key_size: usize, max_value_size: usize) -> Result<BTree<K,V>, Box<Error>> {
        // create our mem_tree
        let mut mem_tree = BTreeMap::<K, BTreeSet<V>>::new();

        let wal_file_path = tree_file_path.to_owned() + ".wal";
        let mut wal_file = try!(WALFile::<K,V>::new(wal_file_path.to_owned(), max_key_size, max_value_size));

        let record_size = max_key_size + max_value_size;

        // if we have a WAL file, replay it into the mem_tree
        if try!(wal_file.is_new()) {
            let mut wal_it = wal_file.into_iter();

            loop {
                match wal_it.next() {
                    // add it to the in-memory table
                    Some(kv) => {mem_tree.entry(kv.key).or_insert(BTreeSet::<V>::new()).insert(kv.value);},
                    None => break,
                }
            }
        }

        // compute the size of a on-disk Node
        let node_size: usize = (max_key_size + size_of::<u64>() + max(max_value_size, (max_key_size + size_of::<u64>()) * NUM_CHILDREN)) as usize;

        // open the data file
        let mut tree_file = try!(OpenOptions::new().read(true).write(true).create(true).open(tree_file_path.to_owned()));

        let metadata = try!(tree_file.metadata());

        // check to see if this is a new file
        if metadata.len() == 0 {
            // write out our header
            try!(tree_file.write(FILE_HEADER.as_bytes()));

            // write out our version
            try!(tree_file.write(&[CURRENT_VERSION]));

            // construct and return our BTree object
            Ok(BTree{tree_file_path: tree_file_path,
                     tree_file: tree_file,
                     wal_file: wal_file,
                     root: None,
                     max_key_size: max_key_size,
                     max_value_size: max_value_size,
                     mem_tree: mem_tree
            })
        } else {
            let mut version_string = vec![0; 8];

            try!(tree_file.read_exact(&mut version_string));

            // make sure we've opened a proper file
            if try!(str::from_utf8(&version_string[0..FILE_HEADER.len()])) != FILE_HEADER ||
               version_string[FILE_HEADER.len()] != CURRENT_VERSION {
                return Err(From::from(std::io::Error::new(ErrorKind::InvalidData, "Invalid BTree file or BTree version")));
            }

            let mut buff = vec![0; node_size];

            // make sure we have a root node to read
            if metadata.len() < (version_string.len() + node_size) as u64 {
                // if we don't have a root node yet, just return
                return Ok(BTree{tree_file_path: tree_file_path,
                                tree_file: tree_file,
                                wal_file: wal_file,
                                root: None,
                                max_key_size: max_key_size,
                                max_value_size: max_value_size,
                                mem_tree: mem_tree
                });
            }

            // seek node_size in from the end of the file to read the root node
            try!(tree_file.seek(SeekFrom::End((node_size as isize * -1) as i64)));
            try!(tree_file.read_exact(&mut buff));

            let root_node: Node<K,V> = try!(decode(&buff[..]));

            Ok(BTree{tree_file_path: tree_file_path,
                     tree_file: tree_file,
                     wal_file: wal_file,
                     root: Some(root_node),
                     max_key_size: max_key_size,
                     max_value_size: max_value_size,
                     mem_tree: mem_tree
            })
        }
    }

    /// Inserts a key into the BTree
    pub fn insert(&mut self, key: K, value: V) -> Result<(), Box<Error>> {
        let record = KeyValuePair{key: key, value: value};

        try!(self.wal_file.write_record(&record));

        let KeyValuePair{key, value} = record;

        self.mem_tree.entry(key).or_insert(BTreeSet::<V>::new()).insert(value);

        Ok( () )
    }

/*
    /// Merges the records on disk with the records in memory
    fn compact(&mut self) -> Result<(), Box<Error>>{
        let mut new_tree_file = try!(OpenOptions::new().read(true).write(true).create(true).truncate(true).open(self.tree_file_path + ".new"));

        let mut mem_iter = self.mem_tree.iter().fuse();  // get an iterator that always returns None when done

        loop {
            let mem_item = mem_iter.next();

        }
    }
*/
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::fs::OpenOptions;
    use ::BTree;
    use rand::{thread_rng, Rng};

    pub fn gen_temp_name() -> String {
        let file_name: String = thread_rng().gen_ascii_chars().take(10).collect();

        return String::from("/tmp/") + &file_name + &String::from(".btr");
    }

    fn remove_files(file_path: String) {
        fs::remove_file(&file_path);
        fs::remove_file(file_path + ".wal");
    }
/*

    #[test]
    fn new_blank_file() {
        let file_path = gen_temp_name();

        BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        // make sure our two files were created
        let btf = OpenOptions::new().read(true).write(false).create(false).open(&file_path).unwrap();
        assert!(btf.metadata().unwrap().len() == 8);

        let wal = OpenOptions::new().read(true).write(false).create(false).open(file_path.to_owned() + ".wal").unwrap();
        assert!(wal.metadata().unwrap().len() == 0);

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn new_existing_file() {
        let file_path = gen_temp_name();

        {
            BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();
        }

        let btree = BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        // check our file lengths from the struct
        assert!(btree.tree_file.metadata().unwrap().len() == 8);
        assert!(btree.wal_file.metadata().unwrap().len() == 0);

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_new_u8() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        let len = btree.insert(2, 3).unwrap(); // insert into a new file

        assert!(btree.wal_file.metadata().unwrap().len() == 2);
        assert!(btree.mem_tree.contains_key(&2));

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_new_str() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<String, String>::new(file_path.to_owned(), 15, 15).unwrap();

        let size = btree.insert("Hello".to_owned(), "World".to_owned()).unwrap(); // insert into a new file

        assert!(btree.wal_file.metadata().unwrap().len() == size as u64);
        assert!(btree.mem_tree.contains_key(&"Hello"));

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_multiple() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<String, String>::new(file_path.to_owned(), 15, 15).unwrap();

        let size1 = btree.insert("Hello".to_owned(), "World".to_owned()).unwrap(); // insert into a new file
        assert!(btree.wal_file.metadata().unwrap().len() == size1 as u64);

        let size2 = btree.insert("Hello".to_owned(), "Everyone".to_owned()).unwrap();
        assert!(btree.wal_file.metadata().unwrap().len() == (size1 + size2) as u64);

        remove_files(file_path); // remove files assuming it all went well
    }
*/
}
<commit_msg>Updated iterator loop to use for<commit_after>extern crate bincode;
extern crate rustc_serialize;
extern crate rand;

mod wal_file;

use wal_file::{KeyValuePair, WALFile, WALIterator};

use bincode::SizeLimit;
use bincode::rustc_serialize::{encode, decode};
use rustc_serialize::{Encodable, Decodable};

use std::cmp::max;
use std::convert::From;
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fs::File;
use std::fs::OpenOptions;
use std::io::{Read, Write, Seek, SeekFrom, ErrorKind};
use std::mem::{size_of};
use std::str;

const NUM_CHILDREN: usize = 32;
const FILE_HEADER: &'static str = "B+Tree\0";
const CURRENT_VERSION: u8 = 0x01;

// specify the types for the keys & values
pub trait KeyType: Ord + Encodable + Decodable {}
pub trait ValueType: Ord + Encodable + Decodable {}

// provide generic implementations
impl<T> KeyType for T where T: Ord + Encodable + Decodable {}
impl<T> ValueType for T where T: Ord + Encodable + Decodable {}

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
enum Payload<K: KeyType, V: ValueType> {
        Value(V),
        Children([(K,u64); NUM_CHILDREN]),
    }

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
struct Node<K: KeyType, V: ValueType> {
    key: K,
    parent: u64,
    payload: Payload<K,V>, // either children, or actual values
}

/// This struct represents an on-disk B+Tree. There are NUM_CHILDREN keys at each
/// level in the tree. The on-disk format is as follows where VV is the version
/// number:
/// |-------------------------------------------|
/// | 0x42 0x2b 0x54 0x72 | 0x65 0x65 0x00 0xVV |
/// | B    +    T    r    | e    e    \0   0xVV |
/// |-------------------------------------------|
/// | smallest record in bincode format         |
/// |-------------------------------------------|
/// | ...                                       |
/// |-------------------------------------------|
/// | largest record in bincode format          |
/// |-------------------------------------------|
/// | internal nodes ...                        |
/// |-------------------------------------------|
/// | root node                                 |
/// |-------------------------------------------|
pub struct BTree<K: KeyType, V: ValueType> {
    tree_file_path: String,         // the path to the tree file
    tree_file: File,                // the file backing the whole thing
    wal_file: WALFile<K,V>,              // write-ahead log for in-memory items
    root: Option<Node<K,V>>,        // optional in-memory copy of the root node
    max_key_size: usize,            // the size of the key in bytes
    max_value_size: usize,          // the size of the value in bytes
    mem_tree: BTreeMap<K, BTreeSet<V>>,  // the in-memory BTree that gets merged with the on-disk one
}

impl <K: KeyType, V: ValueType> BTree<K, V> {
    pub fn new(tree_file_path: String, max_key_size: usize, max_value_size: usize) -> Result<BTree<K,V>, Box<Error>> {
        // create our mem_tree
        let mut mem_tree = BTreeMap::<K, BTreeSet<V>>::new();

        let wal_file_path = tree_file_path.to_owned() + ".wal";
        let mut wal_file = try!(WALFile::<K,V>::new(wal_file_path.to_owned(), max_key_size, max_value_size));

        let record_size = max_key_size + max_value_size;

        // if we have a WAL file, replay it into the mem_tree
        if try!(wal_file.is_new()) {
            for kv in wal_file {
                mem_tree.entry(kv.key).or_insert(BTreeSet::<V>::new()).insert(kv.value);
            }
        }

        // compute the size of a on-disk Node
        let node_size: usize = (max_key_size + size_of::<u64>() + max(max_value_size, (max_key_size + size_of::<u64>()) * NUM_CHILDREN)) as usize;

        // open the data file
        let mut tree_file = try!(OpenOptions::new().read(true).write(true).create(true).open(tree_file_path.to_owned()));

        let metadata = try!(tree_file.metadata());

        // check to see if this is a new file
        if metadata.len() == 0 {
            // write out our header
            try!(tree_file.write(FILE_HEADER.as_bytes()));

            // write out our version
            try!(tree_file.write(&[CURRENT_VERSION]));

            // construct and return our BTree object
            Ok(BTree{tree_file_path: tree_file_path,
                     tree_file: tree_file,
                     wal_file: wal_file,
                     root: None,
                     max_key_size: max_key_size,
                     max_value_size: max_value_size,
                     mem_tree: mem_tree
            })
        } else {
            let mut version_string = vec![0; 8];

            try!(tree_file.read_exact(&mut version_string));

            // make sure we've opened a proper file
            if try!(str::from_utf8(&version_string[0..FILE_HEADER.len()])) != FILE_HEADER ||
               version_string[FILE_HEADER.len()] != CURRENT_VERSION {
                return Err(From::from(std::io::Error::new(ErrorKind::InvalidData, "Invalid BTree file or BTree version")));
            }

            let mut buff = vec![0; node_size];

            // make sure we have a root node to read
            if metadata.len() < (version_string.len() + node_size) as u64 {
                // if we don't have a root node yet, just return
                return Ok(BTree{tree_file_path: tree_file_path,
                                tree_file: tree_file,
                                wal_file: wal_file,
                                root: None,
                                max_key_size: max_key_size,
                                max_value_size: max_value_size,
                                mem_tree: mem_tree
                });
            }

            // seek node_size in from the end of the file to read the root node
            try!(tree_file.seek(SeekFrom::End((node_size as isize * -1) as i64)));
            try!(tree_file.read_exact(&mut buff));

            let root_node: Node<K,V> = try!(decode(&buff[..]));

            Ok(BTree{tree_file_path: tree_file_path,
                     tree_file: tree_file,
                     wal_file: wal_file,
                     root: Some(root_node),
                     max_key_size: max_key_size,
                     max_value_size: max_value_size,
                     mem_tree: mem_tree
            })
        }
    }

    /// Inserts a key into the BTree
    pub fn insert(&mut self, key: K, value: V) -> Result<(), Box<Error>> {
        let record = KeyValuePair{key: key, value: value};

        try!(self.wal_file.write_record(&record));

        let KeyValuePair{key, value} = record;

        self.mem_tree.entry(key).or_insert(BTreeSet::<V>::new()).insert(value);

        Ok( () )
    }

/*
    /// Merges the records on disk with the records in memory
    fn compact(&mut self) -> Result<(), Box<Error>>{
        let mut new_tree_file = try!(OpenOptions::new().read(true).write(true).create(true).truncate(true).open(self.tree_file_path + ".new"));

        let mut mem_iter = self.mem_tree.iter().fuse();  // get an iterator that always returns None when done

        loop {
            let mem_item = mem_iter.next();

        }
    }
*/
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::fs::OpenOptions;
    use ::BTree;
    use rand::{thread_rng, Rng};

    pub fn gen_temp_name() -> String {
        let file_name: String = thread_rng().gen_ascii_chars().take(10).collect();

        return String::from("/tmp/") + &file_name + &String::from(".btr");
    }

    fn remove_files(file_path: String) {
        fs::remove_file(&file_path);
        fs::remove_file(file_path + ".wal");
    }
/*

    #[test]
    fn new_blank_file() {
        let file_path = gen_temp_name();

        BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        // make sure our two files were created
        let btf = OpenOptions::new().read(true).write(false).create(false).open(&file_path).unwrap();
        assert!(btf.metadata().unwrap().len() == 8);

        let wal = OpenOptions::new().read(true).write(false).create(false).open(file_path.to_owned() + ".wal").unwrap();
        assert!(wal.metadata().unwrap().len() == 0);

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn new_existing_file() {
        let file_path = gen_temp_name();

        {
            BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();
        }

        let btree = BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        // check our file lengths from the struct
        assert!(btree.tree_file.metadata().unwrap().len() == 8);
        assert!(btree.wal_file.metadata().unwrap().len() == 0);

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_new_u8() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<u8, u8>::new(file_path.to_owned(), 1, 1).unwrap();

        let len = btree.insert(2, 3).unwrap(); // insert into a new file

        assert!(btree.wal_file.metadata().unwrap().len() == 2);
        assert!(btree.mem_tree.contains_key(&2));

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_new_str() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<String, String>::new(file_path.to_owned(), 15, 15).unwrap();

        let size = btree.insert("Hello".to_owned(), "World".to_owned()).unwrap(); // insert into a new file

        assert!(btree.wal_file.metadata().unwrap().len() == size as u64);
        assert!(btree.mem_tree.contains_key(&"Hello"));

        remove_files(file_path); // remove files assuming it all went well
    }

    #[test]
    fn insert_multiple() {
        let file_path = gen_temp_name();

        let mut btree = BTree::<String, String>::new(file_path.to_owned(), 15, 15).unwrap();

        let size1 = btree.insert("Hello".to_owned(), "World".to_owned()).unwrap(); // insert into a new file
        assert!(btree.wal_file.metadata().unwrap().len() == size1 as u64);

        let size2 = btree.insert("Hello".to_owned(), "Everyone".to_owned()).unwrap();
        assert!(btree.wal_file.metadata().unwrap().len() == (size1 + size2) as u64);

        remove_files(file_path); // remove files assuming it all went well
    }
*/
}
<|endoftext|>


<commit_before>#![feature(iter_cmp)]
extern crate disque;

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread::{spawn, JoinHandle};

use disque::Disque;

#[derive(Clone)]
pub enum JobStatus {
    FastAck,
    AckJob,
    NAck,
}

pub trait Handler {
    fn process_job(&self, queue_name: &[u8], jobid: &String, body: Vec<u8>) -> JobStatus;
    fn process_error(&self, queue_name: &[u8], jobid: &String, nack: u32, additional_deliveries: u32) -> bool;
}

#[derive(Clone)]
struct HandlerWrapper<H: Handler> {
    handler: Arc<H>,
}
unsafe impl<H: Handler> Send for HandlerWrapper<H> {}
unsafe impl<H: Handler> Sync for HandlerWrapper<H> {}

fn create_worker<H: Handler + Clone + 'static>(position: usize,
        task_rx: Receiver<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>,
        completion_tx: Sender<(usize, String, JobStatus)>,
        handler_: HandlerWrapper<H>,
        ) -> JoinHandle<()> {
    let handlerw = handler_.clone();
    spawn(move || {
        let handler = handlerw.handler;
        loop {
            let (queue, jobid, job, nack,
                additional_deliveries) = match task_rx.recv() {
                Ok(o) => match o {
                    Some(v) => v,
                    None => break,
                },
                Err(e) => {
                    // TODO: better log
                    println!("Error in worker thread {:?}", e);
                    break;
                }
            };

            if nack > 0 || additional_deliveries > 0 {
                if !handler.process_error(&*queue, &jobid, nack,
                    additional_deliveries) {
                    return;
                }
            }
            let status = handler.process_job(&*queue, &jobid, job);

            completion_tx.send((position, jobid, status)).unwrap();
        }
    })
}

pub struct EventLoop {
    disque: Disque,
    workers: Vec<(JoinHandle<()>, Sender<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>)>,
    completion_rx: Receiver<(usize, String, JobStatus)>,
    free_workers: HashSet<usize>,
    queues: HashSet<Vec<u8>>,
    hello: (u8, String, Vec<(String, String, u16, u32)>),
    node_counter: HashMap<Vec<u8>, usize>,
}

impl EventLoop {
    pub fn new<H: Handler + Clone + 'static>(
            disque: Disque, numworkers: usize,
            handler: H) -> Self {
        let mut workers = Vec::with_capacity(numworkers);
        let mut free_workers = HashSet::with_capacity(numworkers);
        let (completion_tx, completion_rx) = channel();
        let ahandler = HandlerWrapper { handler: Arc::new(handler) };
        for i in 0..numworkers {
            let (task_tx, task_rx) = channel();
            let jg = create_worker(i, task_rx, completion_tx.clone(),
                    ahandler.clone());
            workers.push((jg, task_tx));
            free_workers.insert(i);
        }
        let hello = disque.hello().unwrap();
        EventLoop {
            disque: disque,
            completion_rx: completion_rx,
            workers: workers,
            hello: hello,
            free_workers: free_workers,
            queues: HashSet::new(),
            node_counter: HashMap::new(),
        }
    }

    pub fn watch_queue(&mut self, queue_name: Vec<u8>) {
        self.queues.insert(queue_name);
    }

    pub fn unwatch_queue(&mut self, queue_name: &Vec<u8>) {
        self.queues.remove(queue_name);
    }

    fn completed(&mut self, worker: usize, jobid: String, status: JobStatus) {
        self.free_workers.insert(worker);
        match status {
            JobStatus::FastAck => self.disque.fastackjob(jobid.as_bytes()),
            JobStatus::AckJob => self.disque.ackjob(jobid.as_bytes()),
            JobStatus::NAck => self.disque.nackjob(jobid.as_bytes()),
        }.unwrap();
    }

    fn mark_completed(&mut self) {
        loop {
            match self.completion_rx.try_recv() {
                Ok(c) => self.completed(c.0, c.1, c.2),
                Err(_) => break,
            }
        }
    }

    pub fn choose_favorite_node(&self) -> (Vec<u8>, usize) {
        let default = (&Vec::new(), &0);
        let r = self.node_counter.iter().max_by(|node| node.1).unwrap_or(default);
        (r.0.clone(), r.1.clone())
    }

    pub fn jobcount_current_node(&self) -> usize {
        let nodeid = self.hello.1.as_bytes()[0..8].to_vec();
        self.node_counter.get(&nodeid).unwrap_or(&0).clone()
    }

    pub fn current_node_id(&self) -> String {
        self.hello.1.clone()
    }

    fn run_once(&mut self) -> bool {
        self.mark_completed();
        let worker = match self.free_workers.iter().next() {
            Some(w) => w.clone(),
            None => return false,
        };

        let job = match self.disque.getjob_withcounters(false, None,
                &*self.queues.iter().map(|k| &**k).collect::<Vec<_>>()
                ).unwrap() {
            Some(j) => j,
            None => return false,
        };

        let nodeid = job.1.as_bytes()[2..10].to_vec();
        let v = self.node_counter.remove(&nodeid).unwrap_or(0);
        self.node_counter.insert(nodeid, v + 1);

        self.free_workers.remove(&worker);
        self.workers[worker].1.send(Some(job)).unwrap();
        true
    }

    fn connect_to_node(&mut self, new_master: Vec<u8>) -> bool {
        let mut hello = None;
        for node in self.hello.2.iter() {
            if node.0.as_bytes()[..new_master.len()] == *new_master {
                match Disque::open(&*format!("redis://{}:{}/", node.1, node.2)) {
                    Ok(disque) => {
                        hello = Some(match disque.hello() {
                            Ok(hello) => hello,
                            Err(_) => break,
                        });
                        self.disque = disque;
                        break;
                    },
                    Err(_) => (),
                }
                break;
            }
        }
        match hello {
            Some(h) => { self.hello = h; true }
            None => false,
        }
    }

    pub fn do_cycle(&mut self) {
        let (fav_node, fav_count) = self.choose_favorite_node();
        let current_count = self.jobcount_current_node();
        // only change if it is at least 20% better than the current node
        if fav_count as f64 / current_count as f64 > 1.2 {
            self.connect_to_node(fav_node);
        }
    }

    pub fn run(&mut self, cycle: usize) {
        let mut c = 0;
        loop {
            self.run_once();
            if cycle > 0 {
                c += 1;
                if c == cycle {
                    self.do_cycle();
                    c = 0;
                }
            }
        }
    }

    pub fn run_times(&mut self, mut times: usize) {
        while times > 0 {
            if self.run_once() {
                times -= 1;
            }
        }
    }

    pub fn stop(mut self) {
        for worker in std::mem::replace(&mut self.workers, vec![]).into_iter() {
            worker.1.send(None).unwrap();
            worker.0.join().unwrap();
        }
        self.mark_completed();
    }
}

#[test]
fn favorite() {
    #[derive(Clone)]
    struct MyHandler;
    impl Handler for MyHandler {
        fn process_job(&self, _: &[u8], _: &String, _: Vec<u8>) -> JobStatus {
            JobStatus::AckJob
        }
        fn process_error(&self, _: &[u8], _: &String, _: u32, _: u32) -> bool {
            false
        }
    }
    let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
    let mut el = EventLoop::new(disque, 1, MyHandler);
    el.node_counter.insert(vec![1, 2, 3], 123);
    el.node_counter.insert(vec![4, 5, 6], 456);
    el.node_counter.insert(vec![0, 0, 0], 0);
    assert_eq!(el.choose_favorite_node(), (vec![4, 5, 6], 456));
}
<commit_msg>Wait for a worker to complete before polling<commit_after>#![feature(iter_cmp)]
extern crate disque;

use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread::{spawn, JoinHandle};

use disque::Disque;

#[derive(Clone)]
pub enum JobStatus {
    FastAck,
    AckJob,
    NAck,
}

pub trait Handler {
    fn process_job(&self, queue_name: &[u8], jobid: &String, body: Vec<u8>) -> JobStatus;
    fn process_error(&self, queue_name: &[u8], jobid: &String, nack: u32, additional_deliveries: u32) -> bool;
}

#[derive(Clone)]
struct HandlerWrapper<H: Handler> {
    handler: Arc<H>,
}
unsafe impl<H: Handler> Send for HandlerWrapper<H> {}
unsafe impl<H: Handler> Sync for HandlerWrapper<H> {}

fn create_worker<H: Handler + Clone + 'static>(position: usize,
        task_rx: Receiver<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>,
        completion_tx: Sender<(usize, String, JobStatus)>,
        handler_: HandlerWrapper<H>,
        ) -> JoinHandle<()> {
    let handlerw = handler_.clone();
    spawn(move || {
        let handler = handlerw.handler;
        loop {
            let (queue, jobid, job, nack,
                additional_deliveries) = match task_rx.recv() {
                Ok(o) => match o {
                    Some(v) => v,
                    None => break,
                },
                Err(e) => {
                    // TODO: better log
                    println!("Error in worker thread {:?}", e);
                    break;
                }
            };

            if nack > 0 || additional_deliveries > 0 {
                if !handler.process_error(&*queue, &jobid, nack,
                    additional_deliveries) {
                    return;
                }
            }
            let status = handler.process_job(&*queue, &jobid, job);

            completion_tx.send((position, jobid, status)).unwrap();
        }
    })
}

pub struct EventLoop {
    disque: Disque,
    workers: Vec<(JoinHandle<()>, Sender<Option<(Vec<u8>, String, Vec<u8>, u32, u32)>>)>,
    completion_rx: Receiver<(usize, String, JobStatus)>,
    free_workers: HashSet<usize>,
    queues: HashSet<Vec<u8>>,
    hello: (u8, String, Vec<(String, String, u16, u32)>),
    node_counter: HashMap<Vec<u8>, usize>,
}

impl EventLoop {
    pub fn new<H: Handler + Clone + 'static>(
            disque: Disque, numworkers: usize,
            handler: H) -> Self {
        let mut workers = Vec::with_capacity(numworkers);
        let mut free_workers = HashSet::with_capacity(numworkers);
        let (completion_tx, completion_rx) = channel();
        let ahandler = HandlerWrapper { handler: Arc::new(handler) };
        for i in 0..numworkers {
            let (task_tx, task_rx) = channel();
            let jg = create_worker(i, task_rx, completion_tx.clone(),
                    ahandler.clone());
            workers.push((jg, task_tx));
            free_workers.insert(i);
        }
        let hello = disque.hello().unwrap();
        EventLoop {
            disque: disque,
            completion_rx: completion_rx,
            workers: workers,
            hello: hello,
            free_workers: free_workers,
            queues: HashSet::new(),
            node_counter: HashMap::new(),
        }
    }

    pub fn watch_queue(&mut self, queue_name: Vec<u8>) {
        self.queues.insert(queue_name);
    }

    pub fn unwatch_queue(&mut self, queue_name: &Vec<u8>) {
        self.queues.remove(queue_name);
    }

    fn completed(&mut self, worker: usize, jobid: String, status: JobStatus) {
        self.free_workers.insert(worker);
        match status {
            JobStatus::FastAck => self.disque.fastackjob(jobid.as_bytes()),
            JobStatus::AckJob => self.disque.ackjob(jobid.as_bytes()),
            JobStatus::NAck => self.disque.nackjob(jobid.as_bytes()),
        }.unwrap();
    }

    fn mark_completed(&mut self, blocking: bool) -> bool {
        if blocking {
            match self.completion_rx.recv() {
                Ok(c) => self.completed(c.0, c.1, c.2),
                Err(_) => return false,
            }
        }
        loop {
            match self.completion_rx.try_recv() {
                Ok(c) => self.completed(c.0, c.1, c.2),
                Err(_) => break,
            }
        }
        true
    }

    pub fn choose_favorite_node(&self) -> (Vec<u8>, usize) {
        let default = (&Vec::new(), &0);
        let r = self.node_counter.iter().max_by(|node| node.1).unwrap_or(default);
        (r.0.clone(), r.1.clone())
    }

    pub fn jobcount_current_node(&self) -> usize {
        let nodeid = self.hello.1.as_bytes()[0..8].to_vec();
        self.node_counter.get(&nodeid).unwrap_or(&0).clone()
    }

    pub fn current_node_id(&self) -> String {
        self.hello.1.clone()
    }

    fn run_once(&mut self) -> bool {
        self.mark_completed(false);
        let worker = match self.free_workers.iter().next() {
            Some(w) => w.clone(),
            None => return false,
        };

        let job = match self.disque.getjob_withcounters(false, None,
                &*self.queues.iter().map(|k| &**k).collect::<Vec<_>>()
                ).unwrap() {
            Some(j) => j,
            None => return false,
        };

        let nodeid = job.1.as_bytes()[2..10].to_vec();
        let v = self.node_counter.remove(&nodeid).unwrap_or(0);
        self.node_counter.insert(nodeid, v + 1);

        self.free_workers.remove(&worker);
        self.workers[worker].1.send(Some(job)).unwrap();
        true
    }

    fn connect_to_node(&mut self, new_master: Vec<u8>) -> bool {
        let mut hello = None;
        for node in self.hello.2.iter() {
            if node.0.as_bytes()[..new_master.len()] == *new_master {
                match Disque::open(&*format!("redis://{}:{}/", node.1, node.2)) {
                    Ok(disque) => {
                        hello = Some(match disque.hello() {
                            Ok(hello) => hello,
                            Err(_) => break,
                        });
                        self.disque = disque;
                        break;
                    },
                    Err(_) => (),
                }
                break;
            }
        }
        match hello {
            Some(h) => { self.hello = h; true }
            None => false,
        }
    }

    pub fn do_cycle(&mut self) {
        let (fav_node, fav_count) = self.choose_favorite_node();
        let current_count = self.jobcount_current_node();
        // only change if it is at least 20% better than the current node
        if fav_count as f64 / current_count as f64 > 1.2 {
            self.connect_to_node(fav_node);
        }
    }

    pub fn run(&mut self, cycle: usize) {
        self.run_times_cycle(0, cycle)
    }

    pub fn run_times(&mut self, times: usize) {
        self.run_times_cycle(times, 0)
    }

    pub fn run_times_cycle(&mut self, times: usize, cycle: usize) {
        let mut c = 0;
        let mut counter = 0;
        loop {
            let did_run = self.run_once();
            if did_run {
                if times > 0 {
                    counter += 1;
                    if counter == times {
                        break;
                    }
                }
                if cycle > 0 {
                    c += 1;
                    if c == cycle {
                        self.do_cycle();
                        c = 0;
                    }
                }
            } else {
                self.mark_completed(true);
            }
        }
        self.mark_completed(false);
    }

    pub fn stop(mut self) {
        for worker in std::mem::replace(&mut self.workers, vec![]).into_iter() {
            worker.1.send(None).unwrap();
            worker.0.join().unwrap();
        }
        self.mark_completed(false);
    }
}

#[test]
fn favorite() {
    #[derive(Clone)]
    struct MyHandler;
    impl Handler for MyHandler {
        fn process_job(&self, _: &[u8], _: &String, _: Vec<u8>) -> JobStatus {
            JobStatus::AckJob
        }
        fn process_error(&self, _: &[u8], _: &String, _: u32, _: u32) -> bool {
            false
        }
    }
    let disque = Disque::open("redis://127.0.0.1:7711/").unwrap();
    let mut el = EventLoop::new(disque, 1, MyHandler);
    el.node_counter.insert(vec![1, 2, 3], 123);
    el.node_counter.insert(vec![4, 5, 6], 456);
    el.node_counter.insert(vec![0, 0, 0], 0);
    assert_eq!(el.choose_favorite_node(), (vec![4, 5, 6], 456));
}
<|endoftext|>


<commit_before>// TODO(dsprenkels) Move the hazardous symbols to a separate module (with proper documentation
// and proper warnings)
/*!
This crate provides bindings to my [Shamir secret sharing library][sss].

The main functions to use are `create_shares` and `combine_shares`. **`create_keyshares` and
`combine_shares` are only for experts!** These latter functions miss some security guarantees, so
do not use them unless you really know what you are doing.

Encapsulated in the `SSSResult`, `combine_shares` will return an `Option<_>` which will be
`Some(data)` if the data could be restored. If the data could not be restored, `combine_shares`
will return `Ok(None)`. This means that could mean either of:

1. More shares were needed to reach the treshold.
2. Shares of different sets (corresponding to different secrets) were supplied or some of the
   shares were tampered with.

# Example

```rust
use shamirsecretsharing::*;

// Create a some shares over the secret data `[42, 42, 42, ...]`
let data = vec![42; DATA_SIZE];
let count = 5;
let treshold = 3;
let mut shares = create_shares(&data, count, treshold).unwrap();

// Lose some shares (for demonstrational purposes)
shares.remove(2);
shares.remove(0);

// We still have 3 shares, so we should still be able to restore the secret
let restored = combine_shares(&shares).unwrap();
assert_eq!(restored, Some(data));

// If we lose one more share the secret is lost
shares.remove(0);
let restored2 = combine_shares(&shares).unwrap();
assert_eq!(restored2, None);
```

This library supports can generate sets with at most `count` and a `treshold` shares.

[sss]: https://github.com/dsprenkels/sss
*/

extern crate libc;
#[link(name = "sss", kind = "static")]

use libc::{uint8_t, c_int};
use std::error;
use std::fmt;

/// Custom error types for errors originating from this crate
#[derive(Debug, PartialEq, Eq)]
pub enum SSSError {
    /// The `n` parameter was invalid
    InvalidN(u8),
    /// The `n` parameter was invalid
    InvalidK(u8),
    /// There was a (key)share that had an invalid length
    BadShareLen((usize, usize)),
    /// The input supplied to a function had an incorrect length
    BadInputLen(usize),
}

/// The size of the input data to `create_shares`
pub const DATA_SIZE: usize = 64;
/// The size of the input data to `create_keyshares`
#[doc(hidden)]
pub const KEY_SIZE: usize = 32;
/// Regular share size from shares produced by `create_shares`
pub const SHARE_SIZE: usize = 113;
/// Keyshare size from shares produced by `create_keyshares`
#[doc(hidden)]
pub const KEYSHARE_SIZE: usize = 33;

impl fmt::Display for SSSError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use SSSError::*;
        match *self {
            // Both underlying errors already impl `Display`, so we defer to
            // their implementations.
            InvalidN(n) => write!(f, "Error: invalid share count ({})", n),
            InvalidK(k) => write!(f, "Error: invalid treshold ({})", k),
            BadShareLen((i, x)) => write!(f, "Error: share {} has bad length ({})", i, x),
            BadInputLen(x) => write!(f, "Error: bad input length ({})", x),
        }
    }
}

impl error::Error for SSSError {
    fn description(&self) -> &str {
        use SSSError::*;
        match *self {
            InvalidN(_) => "invalid n",
            InvalidK(_) => "invalid k",
            BadShareLen(_) => "bad share length",
            BadInputLen(_) => "bad input length",
        }
    }
}

type SSSResult<T> = Result<T, SSSError>;

extern {
    fn sss_create_shares(out: *mut uint8_t, data: *const uint8_t, n: uint8_t, k: uint8_t);
    fn sss_combine_shares(data: *mut uint8_t, shares: *const uint8_t, k: uint8_t) -> c_int;
    fn sss_create_keyshares(out: *mut uint8_t, key: *const uint8_t, n: uint8_t, k: uint8_t);
    fn sss_combine_keyshares(key: *mut uint8_t, shares: *const uint8_t, k: uint8_t);
}

/// Check the parameters `n` and `k` and return `Ok(())` if they were valid
fn check_nk(n: u8, k: u8) -> SSSResult<()> {
    if n < 1 {
        return Err(SSSError::InvalidN(n));
    }
    if k < 1 || k > n {
        return Err(SSSError::InvalidK(k));
    }
    Ok(())
}

/// Check `data` and return `Ok(())` if its length is correct for being shared with
/// `create_shares`
fn check_data_len(data: &[u8]) -> SSSResult<()> {
    if data.len() != DATA_SIZE {
        Err(SSSError::BadInputLen(data.len()))
    } else {
        Ok(())
    }
}

/// Check `key` and return `Ok(())` if its length is correct for being shared with
/// `create_keyshares`
fn check_key_len(key: &[u8]) -> SSSResult<()> {
    if key.len() != KEY_SIZE {
        Err(SSSError::BadInputLen(key.len()))
    } else {
        Ok(())
    }
}

/**
Create a set of shares

- `data` must be a `&[u8]` slice of length `DATA_SIZE` (64)
- `n` is the number of shares that is to be generated
- `k` is the treshold value of how many shares are needed to restore the secret

The value that is returned is a newly allocated vector of vectors. Each of these vectors will
contain `SHARE_SIZE` `u8` items.

# Example
```
use shamirsecretsharing::*;

// Create a some shares over the secret data `[42, 42, 42, ...]`
let data = vec![42; DATA_SIZE];
let count = 5;
let treshold = 3;
let shares = create_shares(&data, count, treshold);
match shares {
 Ok(shares) => println!("Created some shares: {:?}", shares),
 Err(err) => panic!("Oops! Something went wrong: {}", err),
}
```
*/
pub fn create_shares(data: &[u8], n: u8, k: u8) -> SSSResult<Vec<Vec<u8>>> {
    // TODO(dsprenks) Currenlty this function uses the `group` function to group the shares from
    // the output buffer that `sss_create_shares` fills for us. While it is good that this uses
    // no unsafe code, it *does* use a 10 lines of helper function and it needs an extra `tmp`
    // vector. We should probably rewrite this to use `Vec::from_raw_parts` instead.
    try!(check_nk(n, k));
    try!(check_data_len(data));

    // Restore the shares into one buffer
    let mut tmp = Vec::with_capacity(SHARE_SIZE * (n as usize));
    unsafe {
        sss_create_shares(tmp.as_mut_ptr(), data.as_ptr(), n, k);
        tmp.set_len(SHARE_SIZE * (n as usize)); // `sss_create_shares` has written to `tmp`
    }

    // This function groups the elements in `tmp` into a new Vec `acc` in-place.
    let group = |mut acc: Vec<Vec<_>>, x| {
        if acc.last().map_or(false, |x| x.len() < SHARE_SIZE) {
            acc.last_mut().unwrap().push(x);
        } else {
            let mut new_group = Vec::with_capacity(SHARE_SIZE);
            new_group.push(x);
            acc.push(new_group);
        }
        acc
    };

    // Put each share in a separate Vec
    Ok(tmp.into_iter().fold(Vec::with_capacity(n as usize), group))
}

/**
Combine a set of shares and return the original secret

`shares` must be a slice of share vectors.

The return type will be a `Result` which will only be `Err(err)` of the input shares were
malformed. When the input shares are of the correct length, this function will always return
`Ok(())`.

Attempts at restoring a secret may fail. Then `combine_shares` will return `Ok(None)`. This only
cases in which this can happen are:

1. More shares were needed to reach the treshold.
2. Shares of different sets (corresponding to different secrets) were supplied or some of the
   shares were tampered with.

If the shares were correct---and a secret could be restored---this function will return
`Ok(Some(data))`, with `data` being a vector of `u8`s. This `data` will be the same length as When
it was shared, namely `DATA_SIZE` (64) bytes.

# Example

```rust
use shamirsecretsharing::*;

# let mut shares = create_shares(&vec![42; DATA_SIZE], 3, 3).unwrap();
// When `shares` contains a set of valid shares
let restored = combine_shares(&shares).unwrap();
let data = restored.expect("`shares` did not contain a valid set of shares");
println!("Restored some data: {:?}", data);

# // Remove a share s.t. the treshold is not reached
# shares.pop();
// When `shares` contains an invalid set of shares
let restored = combine_shares(&shares).unwrap();
assert_eq!(restored, None);
```
*/
pub fn combine_shares(shares: &[Vec<u8>]) -> SSSResult<Option<Vec<u8>>> {
    for (i, share) in shares.iter().enumerate() {
        if share.len() != SHARE_SIZE {
            return Err(SSSError::BadShareLen((i, share.len())));
        }
    }

    // Build a slice containing all the shares sequentially
    let mut tmp = Vec::with_capacity(SHARE_SIZE * shares.len());
    for share in shares {
        tmp.extend(share.iter());
    }

    // Combine the shares
    let mut data = Vec::with_capacity(DATA_SIZE);
    let ret = unsafe {
        let ret = sss_combine_shares(data.as_mut_ptr(), tmp.as_mut_ptr(), shares.len() as uint8_t);
        data.set_len(DATA_SIZE);
        ret
    };

    match ret {
        0 => Ok(Some(data)),
        _ => Ok(None),
    }
}

#[doc(hidden)]
pub fn create_keyshares(key: &[u8], n: u8, k: u8) -> SSSResult<Vec<Vec<u8>>> {
    try!(check_nk(n, k));
    try!(check_key_len(key));

    // Restore the keyshares into one buffer
    let mut tmp = Vec::with_capacity(KEYSHARE_SIZE * (n as usize));
    unsafe {
        sss_create_keyshares(tmp.as_mut_ptr(), key.as_ptr(), n, k);
        tmp.set_len(KEYSHARE_SIZE * (n as usize)); // `sss_create_shares` has written to `tmp`
    }

    // This function groups the elements in `tmp` into a new Vec `acc` in-place.
    let group = |mut acc: Vec<Vec<_>>, x| {
        if acc.last().map_or(false, |x| x.len() < KEYSHARE_SIZE) {
            acc.last_mut().unwrap().push(x);
        } else {
            let mut new_group = Vec::with_capacity(KEYSHARE_SIZE);
            new_group.push(x);
            acc.push(new_group);
        }
        acc
    };

    // Put each share in a separate Vec
    // TODO(dsprenks) Just as in `create_shares`, we should use `Vec::from_raw_parts` instead
    // of the complex fold.
    Ok(tmp.into_iter().fold(Vec::with_capacity(n as usize), group))
}

#[doc(hidden)]
pub fn combine_keyshares(keyshares: &Vec<Vec<u8>>) -> SSSResult<Vec<u8>> {
    for (i, keyshare) in keyshares.iter().enumerate() {
        if keyshare.len() != KEYSHARE_SIZE {
            return Err(SSSError::BadShareLen((i, keyshare.len())));
        }
    }

    // Build a slice containing all the keyshares sequentially
    let mut tmp = Vec::with_capacity(KEYSHARE_SIZE * keyshares.len());
    for keyshare in keyshares {
        tmp.extend(keyshare.iter());
    }

    // Combine the keyshares
    let mut key = Vec::with_capacity(KEY_SIZE);
    unsafe {
        sss_combine_keyshares(key.as_mut_ptr(), tmp.as_mut_ptr(), keyshares.len() as uint8_t);
        key.set_len(KEY_SIZE);
    };

    Ok(key)
}

#[cfg(test)]
mod tests {
    use super::*;
    const DATA: &[u8] = &[42; DATA_SIZE];
    const KEY: &[u8] = &[42; KEY_SIZE];

    #[test]
    fn create_shares_ok() {
        let shares = create_shares(DATA, 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        for share in shares {
            assert_eq!(share.len(), SHARE_SIZE);
        }
    }

    #[test]
    fn create_shares_err() {
        assert_eq!(create_shares(DATA, 0, 0), Err(SSSError::InvalidN(0)));
        assert_eq!(create_shares(DATA, 5, 0), Err(SSSError::InvalidK(0)));
        assert_eq!(create_shares(DATA, 5, 6), Err(SSSError::InvalidK(6)));
    }

    #[test]
    fn combine_shares_ok() {
        let mut shares = create_shares(DATA, 5, 3).unwrap();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
    }

    #[test]
    fn combine_shares_err() {
        let shares = vec![vec![]];
        assert_eq!(combine_shares(&shares), Err(SSSError::BadShareLen((0, 0))));
    }

    #[test]
    fn create_keyshares_ok() {
        let keyshares = create_keyshares(KEY, 5, 3).unwrap();
        assert_eq!(keyshares.len(), 5);
        for keyshare in keyshares {
            assert_eq!(keyshare.len(), KEYSHARE_SIZE);;
        }
    }

    #[test]
    fn combine_keyshares_ok() {
        let mut keyshares = create_keyshares(KEY, 5, 3).unwrap();
        assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
        keyshares.pop();
        assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
        keyshares.pop();
        assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
        keyshares.pop();
        assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
        keyshares.pop();
        assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
        keyshares.pop();
        assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
    }
}
<commit_msg>Move hazmat into separate module<commit_after>/*!
This crate provides bindings to my [Shamir secret sharing library][sss].

The main functions to use are `create_shares` and `combine_shares`. **`create_keyshares` and
`combine_shares` are only for experts!** These latter functions miss some security guarantees, so
do not use them unless you really know what you are doing.

Encapsulated in the `SSSResult`, `combine_shares` will return an `Option<_>` which will be
`Some(data)` if the data could be restored. If the data could not be restored, `combine_shares`
will return `Ok(None)`. This means that could mean either of:

1. More shares were needed to reach the treshold.
2. Shares of different sets (corresponding to different secrets) were supplied or some of the
   shares were tampered with.

# Example

```rust
use shamirsecretsharing::*;

// Create a some shares over the secret data `[42, 42, 42, ...]`
let data = vec![42; DATA_SIZE];
let count = 5;
let treshold = 3;
let mut shares = create_shares(&data, count, treshold).unwrap();

// Lose some shares (for demonstrational purposes)
shares.remove(2);
shares.remove(0);

// We still have 3 shares, so we should still be able to restore the secret
let restored = combine_shares(&shares).unwrap();
assert_eq!(restored, Some(data));

// If we lose one more share the secret is lost
shares.remove(0);
let restored2 = combine_shares(&shares).unwrap();
assert_eq!(restored2, None);
```

This library supports can generate sets with at most `count` and a `treshold` shares.

[sss]: https://github.com/dsprenkels/sss
*/

extern crate libc;
#[link(name = "sss", kind = "static")]

use libc::{uint8_t, c_int};
use std::error;
use std::fmt;

/// Custom error types for errors originating from this crate
#[derive(Debug, PartialEq, Eq)]
pub enum SSSError {
    /// The `n` parameter was invalid
    InvalidN(u8),
    /// The `n` parameter was invalid
    InvalidK(u8),
    /// There was a (key)share that had an invalid length
    BadShareLen((usize, usize)),
    /// The input supplied to a function had an incorrect length
    BadInputLen(usize),
}

/// The size of the input data to `create_shares`
pub const DATA_SIZE: usize = 64;
/// Regular share size from shares produced by `create_shares`
pub const SHARE_SIZE: usize = 113;

impl fmt::Display for SSSError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use SSSError::*;
        match *self {
            // Both underlying errors already impl `Display`, so we defer to
            // their implementations.
            InvalidN(n) => write!(f, "Error: invalid share count ({})", n),
            InvalidK(k) => write!(f, "Error: invalid treshold ({})", k),
            BadShareLen((i, x)) => write!(f, "Error: share {} has bad length ({})", i, x),
            BadInputLen(x) => write!(f, "Error: bad input length ({})", x),
        }
    }
}

impl error::Error for SSSError {
    fn description(&self) -> &str {
        use SSSError::*;
        match *self {
            InvalidN(_) => "invalid n",
            InvalidK(_) => "invalid k",
            BadShareLen(_) => "bad share length",
            BadInputLen(_) => "bad input length",
        }
    }
}

type SSSResult<T> = Result<T, SSSError>;

extern {
    fn sss_create_shares(out: *mut uint8_t, data: *const uint8_t, n: uint8_t, k: uint8_t);
    fn sss_combine_shares(data: *mut uint8_t, shares: *const uint8_t, k: uint8_t) -> c_int;
}

/// Check the parameters `n` and `k` and return `Ok(())` if they were valid
fn check_nk(n: u8, k: u8) -> SSSResult<()> {
    if n < 1 {
        return Err(SSSError::InvalidN(n));
    }
    if k < 1 || k > n {
        return Err(SSSError::InvalidK(k));
    }
    Ok(())
}

/// Check `data` and return `Ok(())` if its length is correct for being shared with
/// `create_shares`
fn check_data_len(data: &[u8]) -> SSSResult<()> {
    if data.len() != DATA_SIZE {
        Err(SSSError::BadInputLen(data.len()))
    } else {
        Ok(())
    }
}

/**
Create a set of shares

- `data` must be a `&[u8]` slice of length `DATA_SIZE` (64)
- `n` is the number of shares that is to be generated
- `k` is the treshold value of how many shares are needed to restore the secret

The value that is returned is a newly allocated vector of vectors. Each of these vectors will
contain `SHARE_SIZE` `u8` items.

# Example
```
use shamirsecretsharing::*;

// Create a some shares over the secret data `[42, 42, 42, ...]`
let data = vec![42; DATA_SIZE];
let count = 5;
let treshold = 3;
let shares = create_shares(&data, count, treshold);
match shares {
 Ok(shares) => println!("Created some shares: {:?}", shares),
 Err(err) => panic!("Oops! Something went wrong: {}", err),
}
```
*/
pub fn create_shares(data: &[u8], n: u8, k: u8) -> SSSResult<Vec<Vec<u8>>> {
    // TODO(dsprenks) Currenlty this function uses the `group` function to group the shares from
    // the output buffer that `sss_create_shares` fills for us. While it is good that this uses
    // no unsafe code, it *does* use a 10 lines of helper function and it needs an extra `tmp`
    // vector. We should probably rewrite this to use `Vec::from_raw_parts` instead.
    try!(check_nk(n, k));
    try!(check_data_len(data));

    // Restore the shares into one buffer
    let mut tmp = Vec::with_capacity(SHARE_SIZE * (n as usize));
    unsafe {
        sss_create_shares(tmp.as_mut_ptr(), data.as_ptr(), n, k);
        tmp.set_len(SHARE_SIZE * (n as usize)); // `sss_create_shares` has written to `tmp`
    }

    // This function groups the elements in `tmp` into a new Vec `acc` in-place.
    let group = |mut acc: Vec<Vec<_>>, x| {
        if acc.last().map_or(false, |x| x.len() < SHARE_SIZE) {
            acc.last_mut().unwrap().push(x);
        } else {
            let mut new_group = Vec::with_capacity(SHARE_SIZE);
            new_group.push(x);
            acc.push(new_group);
        }
        acc
    };

    // Put each share in a separate Vec
    Ok(tmp.into_iter().fold(Vec::with_capacity(n as usize), group))
}

/**
Combine a set of shares and return the original secret

`shares` must be a slice of share vectors.

The return type will be a `Result` which will only be `Err(err)` of the input shares were
malformed. When the input shares are of the correct length, this function will always return
`Ok(())`.

Attempts at restoring a secret may fail. Then `combine_shares` will return `Ok(None)`. This only
cases in which this can happen are:

1. More shares were needed to reach the treshold.
2. Shares of different sets (corresponding to different secrets) were supplied or some of the
   shares were tampered with.

If the shares were correct---and a secret could be restored---this function will return
`Ok(Some(data))`, with `data` being a vector of `u8`s. This `data` will be the same length as When
it was shared, namely `DATA_SIZE` (64) bytes.

# Example

```rust
use shamirsecretsharing::*;

# let mut shares = create_shares(&vec![42; DATA_SIZE], 3, 3).unwrap();
// When `shares` contains a set of valid shares
let restored = combine_shares(&shares).unwrap();
let data = restored.expect("`shares` did not contain a valid set of shares");
println!("Restored some data: {:?}", data);

# // Remove a share s.t. the treshold is not reached
# shares.pop();
// When `shares` contains an invalid set of shares
let restored = combine_shares(&shares).unwrap();
assert_eq!(restored, None);
```
*/
pub fn combine_shares(shares: &[Vec<u8>]) -> SSSResult<Option<Vec<u8>>> {
    for (i, share) in shares.iter().enumerate() {
        if share.len() != SHARE_SIZE {
            return Err(SSSError::BadShareLen((i, share.len())));
        }
    }

    // Build a slice containing all the shares sequentially
    let mut tmp = Vec::with_capacity(SHARE_SIZE * shares.len());
    for share in shares {
        tmp.extend(share.iter());
    }

    // Combine the shares
    let mut data = Vec::with_capacity(DATA_SIZE);
    let ret = unsafe {
        let ret = sss_combine_shares(data.as_mut_ptr(), tmp.as_mut_ptr(), shares.len() as uint8_t);
        data.set_len(DATA_SIZE);
        ret
    };

    match ret {
        0 => Ok(Some(data)),
        _ => Ok(None),
    }
}

pub mod hazmat {
    use libc::uint8_t;
    use super::{SSSResult, SSSError};

    extern {
        fn sss_create_keyshares(out: *mut uint8_t, key: *const uint8_t, n: uint8_t, k: uint8_t);
        fn sss_combine_keyshares(key: *mut uint8_t, shares: *const uint8_t, k: uint8_t);
    }

    /// The size of the input data to `create_keyshares`
    #[doc(hidden)]
    pub const KEY_SIZE: usize = 32;

    /// Keyshare size from shares produced by `create_keyshares`
    #[doc(hidden)]
    pub const KEYSHARE_SIZE: usize = 33;

    /// Check `key` and return `Ok(())` if its length is correct for being shared with
    /// `create_keyshares`
    fn check_key_len(key: &[u8]) -> SSSResult<()> {
        if key.len() != KEY_SIZE {
            Err(SSSError::BadInputLen(key.len()))
        } else {
            Ok(())
        }
    }

    #[doc(hidden)]
    pub fn create_keyshares(key: &[u8], n: u8, k: u8) -> SSSResult<Vec<Vec<u8>>> {
        try!(super::check_nk(n, k));
        try!(check_key_len(key));

        // Restore the keyshares into one buffer
        let mut tmp = Vec::with_capacity(KEYSHARE_SIZE * (n as usize));
        unsafe {
            sss_create_keyshares(tmp.as_mut_ptr(), key.as_ptr(), n, k);
            tmp.set_len(KEYSHARE_SIZE * (n as usize)); // `sss_create_shares` has written to `tmp`
        }

        // This function groups the elements in `tmp` into a new Vec `acc` in-place.
        let group = |mut acc: Vec<Vec<_>>, x| {
            if acc.last().map_or(false, |x| x.len() < KEYSHARE_SIZE) {
                acc.last_mut().unwrap().push(x);
            } else {
                let mut new_group = Vec::with_capacity(KEYSHARE_SIZE);
                new_group.push(x);
                acc.push(new_group);
            }
            acc
        };

        // Put each share in a separate Vec
        // TODO(dsprenks) Just as in `create_shares`, we should use `Vec::from_raw_parts` instead
        // of the complex fold.
        Ok(tmp.into_iter().fold(Vec::with_capacity(n as usize), group))
    }

    #[doc(hidden)]
    pub fn combine_keyshares(keyshares: &Vec<Vec<u8>>) -> SSSResult<Vec<u8>> {
        for (i, keyshare) in keyshares.iter().enumerate() {
            if keyshare.len() != KEYSHARE_SIZE {
                return Err(SSSError::BadShareLen((i, keyshare.len())));
            }
        }

        // Build a slice containing all the keyshares sequentially
        let mut tmp = Vec::with_capacity(KEYSHARE_SIZE * keyshares.len());
        for keyshare in keyshares {
            tmp.extend(keyshare.iter());
        }

        // Combine the keyshares
        let mut key = Vec::with_capacity(KEY_SIZE);
        unsafe {
            sss_combine_keyshares(key.as_mut_ptr(), tmp.as_mut_ptr(), keyshares.len() as uint8_t);
            key.set_len(KEY_SIZE);
        };

        Ok(key)
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        const KEY: &[u8] = &[42; KEY_SIZE];

        #[test]
        fn create_keyshares_ok() {
            let keyshares = create_keyshares(KEY, 5, 3).unwrap();
            assert_eq!(keyshares.len(), 5);
            for keyshare in keyshares {
                assert_eq!(keyshare.len(), KEYSHARE_SIZE);;
            }
        }

        #[test]
        fn combine_keyshares_ok() {
            let mut keyshares = create_keyshares(KEY, 5, 3).unwrap();
            assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
            keyshares.pop();
            assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
            keyshares.pop();
            assert_eq!(combine_keyshares(&keyshares).unwrap(), KEY);
            keyshares.pop();
            assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
            keyshares.pop();
            assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
            keyshares.pop();
            assert_ne!(combine_keyshares(&keyshares).unwrap(), KEY);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    const DATA: &[u8] = &[42; DATA_SIZE];

    #[test]
    fn create_shares_ok() {
        let shares = create_shares(DATA, 5, 3).unwrap();
        assert_eq!(shares.len(), 5);
        for share in shares {
            assert_eq!(share.len(), SHARE_SIZE);
        }
    }

    #[test]
    fn create_shares_err() {
        assert_eq!(create_shares(DATA, 0, 0), Err(SSSError::InvalidN(0)));
        assert_eq!(create_shares(DATA, 5, 0), Err(SSSError::InvalidK(0)));
        assert_eq!(create_shares(DATA, 5, 6), Err(SSSError::InvalidK(6)));
    }

    #[test]
    fn combine_shares_ok() {
        let mut shares = create_shares(DATA, 5, 3).unwrap();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap().unwrap(), DATA);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
        shares.pop();
        assert_eq!(combine_shares(&shares).unwrap(), None);
    }

    #[test]
    fn combine_shares_err() {
        let shares = vec![vec![]];
        assert_eq!(combine_shares(&shares), Err(SSSError::BadShareLen((0, 0))));
    }
}
<|endoftext|>


<commit_before>#![feature(phase)]
extern crate collections;
extern crate url;
#[phase(plugin, link)] extern crate log;

pub mod connection;
pub mod channel;
pub mod exchange;
pub mod queue;
pub mod framing;
pub mod table;
pub mod protocol;
pub mod session;
<commit_msg>Fix build.<commit_after>#![feature(phase)]
extern crate collections;
extern crate url;
#[phase(plugin, link)] extern crate log;

pub mod connection;
pub mod channel;
pub mod framing;
pub mod table;
pub mod protocol;
pub mod session;
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter { inner: self.hashmap1.iter().chain(self.hashmap2.iter()) }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<commit_msg>Add keys()<commit_after>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter { inner: self.hashmap1.iter().chain(self.hashmap2.iter()) }
    }

    pub fn keys(&self) -> Keys<K, V> {
        Keys { inner: self.hashmap1.keys().chain(self.hashmap2.keys()) }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Keys<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Keys<'a, K, V>, hash_map::Keys<'a, K, V>>
}

impl<'a, K, V> Iterator for Keys<'a, K, V> {
    type Item = &'a K;

    #[inline] fn next(&mut self) -> Option<&'a K> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn keys() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.keys() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<|endoftext|>


<commit_before>use std::env;

#[derive(PartialEq, Debug)]
pub enum CiService {
    Travis,
    Circle,
    Unknown,
}

macro_rules! err {
    ($expr:expr) => {
        match $expr {
            Ok(val) => Some(val),
            Err(e) => match e {
                ::std::env::VarError::NotPresent => None,
                _ => panic!(e),
            }
        }
    }
}

macro_rules! lang_version {
    ($lang:expr, $name:ident) => {
        /// **Travis only**: Returns the $lang version that is used.
        pub fn $name(&self) -> Option<String> {
            match self.service {
                CiService::Travis => err!(env::var(format!("TRAVIS_{}_VERSION", $lang))),
                _ => None,
            }
        }
    }
}

pub struct Ci {
    service: CiService,
}
impl Ci {
    pub fn new() -> Option<Self> {
        let s = Ci::which_ci();
        match s {
            CiService::Travis | CiService::Circle => Some(Ci { service: s }),
            _ => None,
        }
    }

    pub fn which_ci() -> CiService {
        match (err!(env::var("TRAVIS")), err!(env::var("CIRCLECI"))) {
            (Some(_), None) => CiService::Travis,
            (None, Some(_)) => CiService::Circle,
            _ => CiService::Unknown,
        }
    }

    /// Returns the locale setting, g.e. `en_US.UTF-8`.
    pub fn lang() -> Option<String> {
        err!(env::var("LANG"))
    }

    /// Returns the search path.
    pub fn path() -> Option<String> {
        err!(env::var("PATH"))
    }

    /// Returns the path to the users home directory.
    pub fn home() -> Option<String> {
        err!(env::var("HOME"))
    }

    pub fn is_travis(&self) -> bool {
        match self.service {
            CiService::Travis => true,
            _ => false,
        }
    }

    pub fn is_circle(&self) -> bool {
        match self.service {
            CiService::Circle => true,
            _ => false,
        }
    }

    pub fn branch(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BRANCH")),
            CiService::Travis => err!(env::var("TRAVIS_BRANCH")),
            _ => None,
        }
    }

    /// **Circle only**
    /// A permanent link to the current build, for example,
    /// https://circleci.com/gh/circleci/frontend/933
    pub fn build_url(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BUILD_URL")),
            _ => None,
        }
    }

    /// Returns the build number.
    /// TODO: convert this to a number.
    pub fn build_id(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BUILD_NUM")),
            CiService::Travis => err!(env::var("TRAVIS_BUILD_NUMBER")),
            _ => None,
        }
    }

    /// **Travis only**: The absolute path to the directory where the repository
    /// being built has been copied on the worker.
    /// TODO: Return a filesystem path instead?
    pub fn build_dir(&self) -> Option<String> {
        match self.service {
            CiService::Travis => err!(env::var("TRAVIS_BUILD_DIR")),
            _ => None,
        }
    }

    /// The sha1 hash of the commit being tested.
    pub fn commit(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_SHA1")),
            CiService::Travis => err!(env::var("TRAVIS_COMMIT")),
            _ => None,
        }
    }

    /// The number of the pull request this build forms part of.
    /// If this build is not part of a pull request, `None` is returned.
    /// TODO: convert this to a number.
    pub fn pull_request(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_PR_NUMBER")),
            CiService::Travis => {
                let pr = err!(env::var("TRAVIS_PULL_REQUEST")).unwrap_or("false".to_string());
                if pr == "false" {
                    None
                } else {
                    Some(pr)
                }
            }
            _ => None,
        }
    }

    lang_version!("DART", dart);
    lang_version!("GO", go);
    lang_version!("HAXE", haxe);
    lang_version!("JDK", java);
    lang_version!("JULIA", julia);
    lang_version!("NODE", node);
    lang_version!("OTP", otp);
    lang_version!("PERL", perl);
    lang_version!("PHP", php);
    lang_version!("PYTHON", python);
    lang_version!("R", r);
    lang_version!("RUBY", ruby);
    lang_version!("RUST", rust);
    lang_version!("SCALA", scala);
}
<commit_msg>Implement `Display` for `Ci`<commit_after>use std::env;
use std::fmt;

#[derive(PartialEq, Debug)]
pub enum CiService {
    Travis,
    Circle,
    Unknown,
}

macro_rules! err {
    ($expr:expr) => {
        match $expr {
            Ok(val) => Some(val),
            Err(e) => match e {
                ::std::env::VarError::NotPresent => None,
                _ => panic!(e),
            }
        }
    }
}

macro_rules! lang_version {
    ($lang:expr, $name:ident) => {
        /// **Travis only**: Returns the $lang version that is used.
        pub fn $name(&self) -> Option<String> {
            match self.service {
                CiService::Travis => err!(env::var(format!("TRAVIS_{}_VERSION", $lang))),
                _ => None,
            }
        }
    }
}

pub struct Ci {
    service: CiService,
}
impl Ci {
    pub fn new() -> Option<Self> {
        let s = Ci::which_ci();
        match s {
            CiService::Travis | CiService::Circle => Some(Ci { service: s }),
            _ => None,
        }
    }

    pub fn which_ci() -> CiService {
        match (err!(env::var("TRAVIS")), err!(env::var("CIRCLECI"))) {
            (Some(_), None) => CiService::Travis,
            (None, Some(_)) => CiService::Circle,
            _ => CiService::Unknown,
        }
    }

    /// Returns the locale setting, g.e. `en_US.UTF-8`.
    pub fn lang() -> Option<String> {
        err!(env::var("LANG"))
    }

    /// Returns the search path.
    pub fn path() -> Option<String> {
        err!(env::var("PATH"))
    }

    /// Returns the path to the users home directory.
    pub fn home() -> Option<String> {
        err!(env::var("HOME"))
    }

    pub fn is_travis(&self) -> bool {
        match self.service {
            CiService::Travis => true,
            _ => false,
        }
    }

    pub fn is_circle(&self) -> bool {
        match self.service {
            CiService::Circle => true,
            _ => false,
        }
    }

    pub fn branch(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BRANCH")),
            CiService::Travis => err!(env::var("TRAVIS_BRANCH")),
            _ => None,
        }
    }

    /// **Circle only**
    /// A permanent link to the current build, for example,
    /// https://circleci.com/gh/circleci/frontend/933
    pub fn build_url(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BUILD_URL")),
            _ => None,
        }
    }

    /// Returns the build number.
    /// TODO: convert this to a number.
    pub fn build_id(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_BUILD_NUM")),
            CiService::Travis => err!(env::var("TRAVIS_BUILD_NUMBER")),
            _ => None,
        }
    }

    /// **Travis only**: The absolute path to the directory where the repository
    /// being built has been copied on the worker.
    /// TODO: Return a filesystem path instead?
    pub fn build_dir(&self) -> Option<String> {
        match self.service {
            CiService::Travis => err!(env::var("TRAVIS_BUILD_DIR")),
            _ => None,
        }
    }

    /// The sha1 hash of the commit being tested.
    pub fn commit(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_SHA1")),
            CiService::Travis => err!(env::var("TRAVIS_COMMIT")),
            _ => None,
        }
    }

    /// The number of the pull request this build forms part of.
    /// If this build is not part of a pull request, `None` is returned.
    /// TODO: convert this to a number.
    pub fn pull_request(&self) -> Option<String> {
        match self.service {
            CiService::Circle => err!(env::var("CIRCLE_PR_NUMBER")),
            CiService::Travis => {
                let pr = err!(env::var("TRAVIS_PULL_REQUEST")).unwrap_or("false".to_string());
                if pr == "false" {
                    None
                } else {
                    Some(pr)
                }
            }
            _ => None,
        }
    }

    lang_version!("DART", dart);
    lang_version!("GO", go);
    lang_version!("HAXE", haxe);
    lang_version!("JDK", java);
    lang_version!("JULIA", julia);
    lang_version!("NODE", node);
    lang_version!("OTP", otp);
    lang_version!("PERL", perl);
    lang_version!("PHP", php);
    lang_version!("PYTHON", python);
    lang_version!("R", r);
    lang_version!("RUBY", ruby);
    lang_version!("RUST", rust);
    lang_version!("SCALA", scala);
}
impl fmt::Display for Ci {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Continuous Integration Service: {:?}", self.service)
    }
}
<|endoftext|>