<commit_before>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::{LEVEL_2, LEVEL_3, LEVEL_PAGE_DIFF_MULTIPLIER};
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::cmp::min;
use std::collections::{BTreeSet, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let first_node = write_node(node);
            let mut collected_keys = first_node.len();
            let mut collected = vec![first_node];
            let target_keys = min(KS::slice_len() * LEVEL_PAGE_DIFF_MULTIPLIER, LEVEL_3);
            let target_guards = LEVEL_2;
            while collected_keys < target_keys && collected.len() < target_guards {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected_keys += right.len();
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound);
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut node_ref = node.clone();
    loop {
        let mut page = write_node::<KS, PS>(&node_ref);
        let page_len = page.len();
        if !page.is_empty_node() {
            let mut ptr_indices: HashSet<_> = page.innode().ptrs.as_slice_immute()[..page_len + 1]
                .iter()
                .enumerate()
                .filter_map(|(i, sub_level)| {
                    if read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                        Some(i)
                    } else {
                        None
                    }
                })
                .collect();
            let page_right_bound = page.right_bound().clone();
            if ptr_indices.len() == page.len() + 1 {
                // all sub nodes are empty
                // will set current node empty either
                let next_page_ref = page.innode().right.clone();
                *page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: next_page_ref,
                });
            } else {
                let mut innode = page.innode_mut();
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                {
                    let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..page_len + 1]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, p)| p)
                        .collect();

                    let keys: Vec<&mut _> = innode.keys.as_slice()[..page_len]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, k)| k)
                        .collect();
                    debug!("Prune filtered page have keys {:?}", &keys);
                    debug_assert_eq!(ptrs.len(), keys.len() + 1);
                    innode.len = keys.len();
                    let new_keys = new_keys.as_slice();
                    let new_ptrs = new_ptrs.as_slice();
                    for (i, key) in keys.into_iter().enumerate() {
                        debug_assert!(key > &mut smallvec!(0));
                        mem::swap(&mut new_keys[i], key);
                    }
                    for (i, ptr) in ptrs.into_iter().enumerate() {
                        debug_assert!(!ptr.is_default());
                        mem::swap(&mut new_ptrs[i], ptr);
                    }
                }
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
            if &page_right_bound >= bound {
                break;
            }
        }
        node_ref = page.right_ref().unwrap().clone();
    }
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = src_tree.head_page_id;
        new_first_node_ext.prev = left_left_most;

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<commit_msg>dump on assertion failure<commit_after>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::{LEVEL_2, LEVEL_3, LEVEL_PAGE_DIFF_MULTIPLIER};
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::cmp::min;
use std::collections::{BTreeSet, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;
use index::btree::dump::dump_tree;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let first_node = write_node(node);
            let mut collected_keys = first_node.len();
            let mut collected = vec![first_node];
            let target_keys = min(KS::slice_len() * LEVEL_PAGE_DIFF_MULTIPLIER, LEVEL_3);
            let target_guards = LEVEL_2;
            while collected_keys < target_keys && collected.len() < target_guards {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected_keys += right.len();
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound);
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut node_ref = node.clone();
    loop {
        let mut page = write_node::<KS, PS>(&node_ref);
        let page_len = page.len();
        if !page.is_empty_node() {
            let mut ptr_indices: HashSet<_> = page.innode().ptrs.as_slice_immute()[..page_len + 1]
                .iter()
                .enumerate()
                .filter_map(|(i, sub_level)| {
                    if read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                        Some(i)
                    } else {
                        None
                    }
                })
                .collect();
            let page_right_bound = page.right_bound().clone();
            if ptr_indices.len() == page.len() + 1 {
                // all sub nodes are empty
                // will set current node empty either
                let next_page_ref = page.innode().right.clone();
                *page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: next_page_ref,
                });
            } else {
                let mut innode = page.innode_mut();
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                {
                    let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..page_len + 1]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, p)| p)
                        .collect();

                    let keys: Vec<&mut _> = innode.keys.as_slice()[..page_len]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, k)| k)
                        .collect();
                    debug!("Prune filtered page have keys {:?}", &keys);
                    debug_assert_eq!(ptrs.len(), keys.len() + 1);
                    innode.len = keys.len();
                    let new_keys = new_keys.as_slice();
                    let new_ptrs = new_ptrs.as_slice();
                    for (i, key) in keys.into_iter().enumerate() {
                        debug_assert!(key > &mut smallvec!(0));
                        mem::swap(&mut new_keys[i], key);
                    }
                    for (i, ptr) in ptrs.into_iter().enumerate() {
                        debug_assert!(!ptr.is_default());
                        mem::swap(&mut new_ptrs[i], ptr);
                    }
                }
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
            if &page_right_bound >= bound {
                break;
            }
        }
        node_ref = page.right_ref().unwrap().clone();
    }
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    if cfg!(debug_assertions) {
        if left_most_id != src_tree.head_page_id {
            dump_tree(src_tree, "level_lsm_merge_failure_dump.json");
        }
    }
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = src_tree.head_page_id;
        new_first_node_ext.prev = left_left_most;

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    prune_selected::<KS, PS>(&src_tree.get_root(), &prune_bound);

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

#[derive(Debug, PartialEq)]
pub enum CacheError {
    CapacityExceeded,
    KeyNotFound,
    KeyTooLong,
    ValueTooLong,
}

pub type CacheResult<T> = Result<T, CacheError>;

#[derive(PartialEq, Eq, Hash, Debug, Clone)]
pub struct Key {
    item: Vec<u8>,
}

impl Key {
    fn new(item: Vec<u8>) -> Key {
        Key { item: item }
    }

    fn len(&self) -> usize {
        self.item.len()
    }
}

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Value {
    item: Vec<u8>,
}

impl Value {
    fn new(item: Vec<u8>) -> Value {
        Value { item: item }
    }

    fn len(&self) -> usize {
        self.item.len()
    }
}

pub struct Cache {
    capacity: u64,
    key_maxlen: u64,
    value_maxlen: u64,
    storage: HashMap<Key, Value>,
}

impl Cache {
    fn new(capacity: u64, key_maxlen: u64, value_maxlen: u64) -> Cache {
        Cache {
            capacity: capacity,
            key_maxlen: key_maxlen,
            value_maxlen: value_maxlen,
            storage: HashMap::new(),
        }
    }

    fn with_defaults(capacity: u64) -> Cache {
        Cache::new(
            capacity,
            1024, // key_maxlen = 1kb
            1048576 // value_maxlen = 1mb
            )
    }

    fn check_key_len(&self, key: &Key) -> bool {
        key.len() as u64 <= self.key_maxlen
    }

    fn check_value_len(&self, value: &Value) -> bool {
        value.len() as u64 <= self.value_maxlen
    }

    fn get(&self, key: Key) -> CacheResult<&Value> {
        // Check key size
        if (!self.check_key_len(&key)) {
            return Err(CacheError::KeyTooLong);
        }

        match self.storage.get(&key) {
            Some(value) => Ok(value),
            None => Err(CacheError::KeyNotFound),
        }
    }

    fn set(&mut self, key: Key, value: Value) -> CacheResult<()> {
        // Check key & value sizes
        if (!self.check_key_len(&key)) {
            return Err(CacheError::KeyTooLong);
        }
        if (!self.check_value_len(&value)) {
            return Err(CacheError::ValueTooLong);
        }

        // Check capacity if adding new key
        if (!self.storage.contains_key(&key)) {
            if (self.storage.len() as u64 == self.capacity) {
                return Err(CacheError::CapacityExceeded);
            }
        }

        self.storage.insert(key, value);
        Ok(())
    }
}

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

    // helper func since assert_eq!(rv.unwrap(), err) does not work
    fn assert_rv_eq<T>(rv: CacheResult<T>, err: CacheError) {
        // First of all it's supposed to be an error
        assert!(rv.is_err());

        // Now check if the error constructor is the right one
        match rv {
            Ok(_) => (),
            Err(e) => { assert_eq!(e, err); }
        };
    }

    #[test]
    fn test_set_one_key() {
        let mut cache = Cache::with_defaults(1);

        let key = Key::new(vec![1, 2, 3]);
        let value = Value::new(vec![4, 5, 6]);

        cache.set(key.clone(), value.clone());
        let value_found = cache.get(key).unwrap();

        assert_eq!(&value, value_found);
    }

    #[test]
    fn test_key_not_found() {
        let mut cache = Cache::with_defaults(1);

        // Retrieve a different key to the one set
        cache.set(Key::new(vec![1]), Value::new(vec![9]));
        let rv = cache.get(Key::new(vec![2]));

        assert_rv_eq(rv, CacheError::KeyNotFound);
    }

    #[test]
    fn test_store_beyond_capacity() {
        let mut cache = Cache::with_defaults(1);

        // we've now reached capacity
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9]));
        assert!(rv.is_ok());

        // overwriting is ok
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9]));
        assert!(rv.is_ok());

        // but we cannot store a new key
        let rv = cache.set(Key::new(vec![2]), Value::new(vec![9]));
        assert_rv_eq(rv, CacheError::CapacityExceeded);
    }

    #[test]
    fn test_exceed_item_size_limits() {
        let mut cache = Cache::new(1, 1, 1);

        // use a key that is too long
        let rv = cache.set(Key::new(vec![1, 2]), Value::new(vec![9]));
        assert_rv_eq(rv, CacheError::KeyTooLong);

        // use a value that is too long
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9, 8]));
        assert_rv_eq(rv, CacheError::ValueTooLong);

        // use a key that is too long
        let rv = cache.get(Key::new(vec![1, 2]));
        assert_rv_eq(rv, CacheError::KeyTooLong);
    }
}
<commit_msg>also implement contains_key<commit_after>use std::collections::HashMap;

#[derive(Debug, PartialEq)]
pub enum CacheError {
    CapacityExceeded,
    KeyNotFound,
    KeyTooLong,
    ValueTooLong,
}

pub type CacheResult<T> = Result<T, CacheError>;

#[derive(PartialEq, Eq, Hash, Debug, Clone)]
pub struct Key {
    item: Vec<u8>,
}

impl Key {
    fn new(item: Vec<u8>) -> Key {
        Key { item: item }
    }

    fn len(&self) -> usize {
        self.item.len()
    }
}

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Value {
    item: Vec<u8>,
}

impl Value {
    fn new(item: Vec<u8>) -> Value {
        Value { item: item }
    }

    fn len(&self) -> usize {
        self.item.len()
    }
}

pub struct Cache {
    capacity: u64,
    key_maxlen: u64,
    value_maxlen: u64,
    storage: HashMap<Key, Value>,
}

impl Cache {
    fn new(capacity: u64, key_maxlen: u64, value_maxlen: u64) -> Cache {
        Cache {
            capacity: capacity,
            key_maxlen: key_maxlen,
            value_maxlen: value_maxlen,
            storage: HashMap::new(),
        }
    }

    fn with_defaults(capacity: u64) -> Cache {
        Cache::new(
            capacity,
            1024, // key_maxlen = 1kb
            1048576 // value_maxlen = 1mb
            )
    }

    fn check_key_len(&self, key: &Key) -> bool {
        key.len() as u64 <= self.key_maxlen
    }

    fn check_value_len(&self, value: &Value) -> bool {
        value.len() as u64 <= self.value_maxlen
    }

    fn contains_key(&self, key: Key) -> CacheResult<bool> {
        // Check key size
        if (!self.check_key_len(&key)) {
            return Err(CacheError::KeyTooLong);
        }

        Ok(self.storage.contains_key(&key))
    }

    fn get(&self, key: Key) -> CacheResult<&Value> {
        // Check key size
        if (!self.check_key_len(&key)) {
            return Err(CacheError::KeyTooLong);
        }

        match self.storage.get(&key) {
            Some(value) => Ok(value),
            None => Err(CacheError::KeyNotFound),
        }
    }

    fn set(&mut self, key: Key, value: Value) -> CacheResult<()> {
        // Check key & value sizes
        if (!self.check_key_len(&key)) {
            return Err(CacheError::KeyTooLong);
        }
        if (!self.check_value_len(&value)) {
            return Err(CacheError::ValueTooLong);
        }

        // Check capacity if adding new key
        if (!self.storage.contains_key(&key)) {
            if (self.storage.len() as u64 == self.capacity) {
                return Err(CacheError::CapacityExceeded);
            }
        }

        self.storage.insert(key, value);
        Ok(())
    }
}

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

    // helper func since assert_eq!(rv.unwrap(), err) does not work
    fn assert_rv_eq<T>(rv: CacheResult<T>, err: CacheError) {
        // First of all it's supposed to be an error
        assert!(rv.is_err());

        // Now check if the error constructor is the right one
        match rv {
            Ok(_) => (),
            Err(e) => { assert_eq!(e, err); }
        };
    }

    #[test]
    fn test_set_one_key() {
        let mut cache = Cache::with_defaults(1);

        let key = Key::new(vec![1, 2, 3]);
        let value = Value::new(vec![4, 5, 6]);

        // First set it
        cache.set(key.clone(), value.clone());

        // Then test for it
        let rv = cache.contains_key(key.clone());
        assert_eq!(rv.unwrap(), true);

        // Test for a key that was not set
        let rv = cache.contains_key(Key::new(vec![9, 8]));
        assert_eq!(rv.unwrap(), false);

        // Now fetch it
        let value_found = cache.get(key).unwrap();

        assert_eq!(&value, value_found);
    }

    #[test]
    fn test_key_not_found() {
        let mut cache = Cache::with_defaults(1);

        // Retrieve a different key to the one set
        cache.set(Key::new(vec![1]), Value::new(vec![9]));
        let rv = cache.get(Key::new(vec![2]));

        assert_rv_eq(rv, CacheError::KeyNotFound);
    }

    #[test]
    fn test_store_beyond_capacity() {
        let mut cache = Cache::with_defaults(1);

        // we've now reached capacity
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9]));
        assert!(rv.is_ok());

        // overwriting is ok
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9]));
        assert!(rv.is_ok());

        // but we cannot store a new key
        let rv = cache.set(Key::new(vec![2]), Value::new(vec![9]));
        assert_rv_eq(rv, CacheError::CapacityExceeded);
    }

    #[test]
    fn test_exceed_item_size_limits() {
        let mut cache = Cache::new(1, 1, 1);

        // set: use a key that is too long
        let rv = cache.set(Key::new(vec![1, 2]), Value::new(vec![9]));
        assert_rv_eq(rv, CacheError::KeyTooLong);

        // set: use a value that is too long
        let rv = cache.set(Key::new(vec![1]), Value::new(vec![9, 8]));
        assert_rv_eq(rv, CacheError::ValueTooLong);

        // get: use a key that is too long
        let rv = cache.get(Key::new(vec![1, 2]));
        assert_rv_eq(rv, CacheError::KeyTooLong);

        // contains_key: use a key that is too long
        let rv = cache.contains_key(Key::new(vec![1, 2]));
        assert_rv_eq(rv, CacheError::KeyTooLong);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate prodbg_api;
extern crate libc;
use prodbg_api::*;
use libc::c_void;
pub mod gdb;
use gdb::{GdbRemote, Memory, NeedsAck};
//use std::net::{SocketAddr, SocketAddrV4, Ipv4Addr};
//use std::io::Read;
//use std::{slice, str};

//use std::ptr;

//static CODE: &'static [u8] = b"\xd4\x40\x87\x5a\x4e\x71\x02\xb4\xc0\xde\xc0\xde\x5c\x00\x1d\x80\x71\x12\x01\x23\xf2\x3c\x44\x22\x40\x49\x0e\x56\x54\xc5\xf2\x3c\x44\x00\x44\x7a\x00\x00\xf2\x00\x0a\x28\x4E\xB9\x00\x00\x00\x12\x4E\x75";

struct AmigaUaeBackend {
    capstone: Capstone,
    conn: gdb::GdbRemote,
    exception_location: u32,
}

impl AmigaUaeBackend {
    fn from_hex(ch: u8) -> u8 {
        if (ch >= b'a') && (ch <= b'f') {
            return ch - b'a' + 10;
        }

        if (ch >= b'0') && (ch <= b'9') {
            return ch - b'0';
        }

        if (ch >= b'A') && (ch <= b'F') {
            return ch - b'A' + 10;
        }

        return 0;
    }

    fn get_u32(data: &[u8]) -> u32 {
        let t0 = ((Self::from_hex(data[0]) << 4) | Self::from_hex(data[1])) as u32;
        let t1 = ((Self::from_hex(data[2]) << 4) | Self::from_hex(data[3])) as u32;
        let t2 = ((Self::from_hex(data[4]) << 4) | Self::from_hex(data[5])) as u32;
        let t3 = ((Self::from_hex(data[6]) << 4) | Self::from_hex(data[7])) as u32;
        (t0 << 24) | (t1 << 16) | (t2 << 8) | t3
    }

    fn write_register(writer: &mut Writer, name: &str, data: u32, read_only: bool) {
        writer.array_entry_begin();
        writer.write_string("name", name);
        writer.write_u8("size", 4);

        if read_only {
            writer.write_u8("read_only", 1);
        }

        writer.write_u32("register", data);

        writer.array_entry_end();
    }

    fn write_registers(&mut self, writer: &mut Writer, data: &Vec<u8>) {
        let mut index = 1;  // 1 to skip starting $

        writer.event_begin(EventType::SetRegisters as u16);
        writer.array_begin("registers");

        // a registers

        for i in 0..8 {
            let name = format!("a{}", i);
            let reg = Self::get_u32(&data[index..]);
            Self::write_register(writer, &name, reg, false);
            index += 8;
        }

        // d registers

        for i in 0..8 {
            let name = format!("d{}", i);
            let reg = Self::get_u32(&data[index..]);
            Self::write_register(writer, &name, reg, false);
            index += 8;
        }

        // Status registers & pc

        let sr = Self::get_u32(&data[index..]);
        let pc = Self::get_u32(&data[index + 8..]);

        self.exception_location = pc;

        Self::write_register(writer, "sr", sr, true);
        Self::write_register(writer, "pc", pc, true);

        writer.array_end();
        writer.event_end();
    }

    fn write_disassembly(&mut self, writer: &mut Writer, mem: &Memory) {
        match self.capstone.open(Arch::M68K, CS_MODE_M68K_000) {
            Err(e) => {
                println!("Unable to open Capstone {}", e as i32);
                return;
            }
            _ => (),
        }

        if let Ok(insns) = self.capstone.disasm(&mem.data, mem.address, 0) {
            writer.event_begin(EventType::SetDisassembly as u16);
            writer.array_begin("disassembly");

            for i in insns.iter() {
                let text = format!("{}\t{}", i.mnemonic().unwrap(), i.op_str().unwrap_or(""));
                writer.array_entry_begin();
                writer.write_u32("address", i.address as u32);
                writer.write_string("line", &text);

                //println!("{:?}", i);

                writer.array_entry_end();
            }

            writer.array_end();
            writer.event_end();
        } else {
            println!("No instructions :(");
        }
    }

    fn write_exception_location(&mut self, writer: &mut Writer) {
        writer.event_begin(EventType::SetExceptionLocation as u16);
        writer.write_u64("address", self.exception_location as u64);
        writer.write_u8("size", 4);
        writer.event_end();
    }

    /*
    fn test_capstone(&mut self) {
        match self.capstone.open(Arch::M68K, CS_MODE_M68K_000) {
            Err(e) => {
                println!("Unable to open Capstone {}", e as i32);
                return;
            }
            _ => (),
        }

        if let Ok(insns) = self.capstone.disasm(CODE, 0x1000, 0) {
            println!("Got {} instructions", insns.len());

            for i in insns.iter() {
                println!("{:?}", i);
            }
        } else {
            println!("No instructions :(");
        }
    }
    */
}

impl Backend for AmigaUaeBackend {
    fn new(service: &Service) -> Self {
        AmigaUaeBackend {
            capstone: service.get_capstone(),
            conn: GdbRemote::new(NeedsAck::No),
            exception_location: 0,
        }
    }

    // TODO: Something about action action as i32
    fn update(&mut self, action: i32, reader: &mut Reader, writer: &mut Writer) {
        for event in reader.get_event() {
            match event {
                 35 => {
                    if self.conn.connect("127.0.0.1:6860").is_ok() {
                        println!("Connected!");
                    }
                }
                _ => (),
            }
        }

        match action {
            ACTION_BREAK => {
                println!("Break");
            }

            ACTION_STEP => {
                self.conn.step_sync();

                if let Some(data) = self.conn.get_registers_sync() {
                    self.write_registers(writer, &data.clone());
                }

                if let Some(memory) = self.conn.get_memory_sync(self.exception_location - 128, 256) {
                    self.write_disassembly(writer, &memory);
                }

                self.write_exception_location(writer);
            }
            _ => (),
        }

        //self.conn.update();
    }

    fn register_menu(&mut self, menu_funcs: &mut MenuFuncs) -> *mut c_void {
        let menu = menu_funcs.create_menu("Amiga UAE Debugger");
        menu_funcs.add_menu_item(menu, "Connect to UAE...", 0, 0, 0);
        menu
    }
}

#[no_mangle]
pub fn init_plugin(plugin_handler: &mut PluginHandler) {
    define_backend_plugin!(PLUGIN, b"Amiga UAE Debugger\0", AmigaUaeBackend);
    plugin_handler.register_backend(&PLUGIN);
}

<commit_msg>align disassembly view<commit_after>#[macro_use]
extern crate prodbg_api;
extern crate libc;
use prodbg_api::*;
use libc::c_void;
pub mod gdb;
use gdb::{GdbRemote, Memory, NeedsAck};
//use std::net::{SocketAddr, SocketAddrV4, Ipv4Addr};
//use std::io::Read;
//use std::{slice, str};

//use std::ptr;

//static CODE: &'static [u8] = b"\xd4\x40\x87\x5a\x4e\x71\x02\xb4\xc0\xde\xc0\xde\x5c\x00\x1d\x80\x71\x12\x01\x23\xf2\x3c\x44\x22\x40\x49\x0e\x56\x54\xc5\xf2\x3c\x44\x00\x44\x7a\x00\x00\xf2\x00\x0a\x28\x4E\xB9\x00\x00\x00\x12\x4E\x75";

struct AmigaUaeBackend {
    capstone: Capstone,
    conn: gdb::GdbRemote,
    exception_location: u32,
}

impl AmigaUaeBackend {
    fn from_hex(ch: u8) -> u8 {
        if (ch >= b'a') && (ch <= b'f') {
            return ch - b'a' + 10;
        }

        if (ch >= b'0') && (ch <= b'9') {
            return ch - b'0';
        }

        if (ch >= b'A') && (ch <= b'F') {
            return ch - b'A' + 10;
        }

        return 0;
    }

    fn get_u32(data: &[u8]) -> u32 {
        let t0 = ((Self::from_hex(data[0]) << 4) | Self::from_hex(data[1])) as u32;
        let t1 = ((Self::from_hex(data[2]) << 4) | Self::from_hex(data[3])) as u32;
        let t2 = ((Self::from_hex(data[4]) << 4) | Self::from_hex(data[5])) as u32;
        let t3 = ((Self::from_hex(data[6]) << 4) | Self::from_hex(data[7])) as u32;
        (t0 << 24) | (t1 << 16) | (t2 << 8) | t3
    }

    fn write_register(writer: &mut Writer, name: &str, data: u32, read_only: bool) {
        writer.array_entry_begin();
        writer.write_string("name", name);
        writer.write_u8("size", 4);

        if read_only {
            writer.write_u8("read_only", 1);
        }

        writer.write_u32("register", data);

        writer.array_entry_end();
    }

    fn write_registers(&mut self, writer: &mut Writer, data: &Vec<u8>) {
        let mut index = 1;  // 1 to skip starting $

        writer.event_begin(EventType::SetRegisters as u16);
        writer.array_begin("registers");

        // d registers

        for i in 0..8 {
            let name = format!("d{}", i);
            let reg = Self::get_u32(&data[index..]);
            Self::write_register(writer, &name, reg, false);
            index += 8;
        }

        // a registers

        for i in 0..8 {
            let name = format!("a{}", i);
            let reg = Self::get_u32(&data[index..]);
            Self::write_register(writer, &name, reg, false);
            index += 8;
        }

        // Status registers & pc

        let sr = Self::get_u32(&data[index..]);
        let pc = Self::get_u32(&data[index + 8..]);

        self.exception_location = pc;

        Self::write_register(writer, "sr", sr, true);
        Self::write_register(writer, "pc", pc, true);

        writer.array_end();
        writer.event_end();
    }

    fn write_disassembly(&mut self, writer: &mut Writer, mem: &Memory) {
        match self.capstone.open(Arch::M68K, CS_MODE_M68K_000) {
            Err(e) => {
                println!("Unable to open Capstone {}", e as i32);
                return;
            }
            _ => (),
        }

        if let Ok(insns) = self.capstone.disasm(&mem.data, mem.address, 0) {
            writer.event_begin(EventType::SetDisassembly as u16);
            writer.array_begin("disassembly");

            for i in insns.iter() {
                let text = format!("{0: <10} {1: <10}", i.mnemonic().unwrap(), i.op_str().unwrap_or(""));
                writer.array_entry_begin();
                writer.write_u32("address", i.address as u32);
                writer.write_string("line", &text);

                //println!("{:?}", i);

                writer.array_entry_end();
            }

            writer.array_end();
            writer.event_end();
        } else {
            println!("No instructions :(");
        }
    }

    fn write_exception_location(&mut self, writer: &mut Writer) {
        writer.event_begin(EventType::SetExceptionLocation as u16);
        writer.write_u64("address", self.exception_location as u64);
        writer.write_u8("size", 4);
        writer.event_end();
    }

    /*
    fn test_capstone(&mut self) {
        match self.capstone.open(Arch::M68K, CS_MODE_M68K_000) {
            Err(e) => {
                println!("Unable to open Capstone {}", e as i32);
                return;
            }
            _ => (),
        }

        if let Ok(insns) = self.capstone.disasm(CODE, 0x1000, 0) {
            println!("Got {} instructions", insns.len());

            for i in insns.iter() {
                println!("{:?}", i);
            }
        } else {
            println!("No instructions :(");
        }
    }
    */
}

impl Backend for AmigaUaeBackend {
    fn new(service: &Service) -> Self {
        AmigaUaeBackend {
            capstone: service.get_capstone(),
            conn: GdbRemote::new(NeedsAck::No),
            exception_location: 0,
        }
    }

    // TODO: Something about action action as i32
    fn update(&mut self, action: i32, reader: &mut Reader, writer: &mut Writer) {
        for event in reader.get_event() {
            match event {
                 35 => {
                    if self.conn.connect("127.0.0.1:6860").is_ok() {
                        println!("Connected!");
                    }
                }
                _ => (),
            }
        }

        match action {
            ACTION_BREAK => {
                println!("Break");
            }

            ACTION_STEP => {
                self.conn.step_sync();

                if let Some(data) = self.conn.get_registers_sync() {
                    self.write_registers(writer, &data.clone());
                }

                if let Some(memory) = self.conn.get_memory_sync(self.exception_location - 128, 256) {
                    self.write_disassembly(writer, &memory);
                }

                self.write_exception_location(writer);
            }
            _ => (),
        }

        //self.conn.update();
    }

    fn register_menu(&mut self, menu_funcs: &mut MenuFuncs) -> *mut c_void {
        let menu = menu_funcs.create_menu("Amiga UAE Debugger");
        menu_funcs.add_menu_item(menu, "Connect to UAE...", 0, 0, 0);
        menu
    }
}

#[no_mangle]
pub fn init_plugin(plugin_handler: &mut PluginHandler) {
    define_backend_plugin!(PLUGIN, b"Amiga UAE Debugger\0", AmigaUaeBackend);
    plugin_handler.register_backend(&PLUGIN);
}

<|endoftext|>


<commit_before>//! Layout handling

use rustwlc::handle::{WlcView, WlcOutput};
use std::rc::{Rc, Weak};

#[derive(PartialEq, Clone, Copy)]
pub enum ContainerType {
    Root,        /* Root container, only one exists */
    Output,      /* Output, like a monitor or head */
    Workspace,   /* A workspace */
    View         /* A view (aka a window) */
}

enum Layout {
    None,
    Horizontal,
    Vertical,
    Stacked,
    Tabbed,
    Floating
}

/// Like i3, everything (workspaces, containers, views) are containable.
pub trait Containable {
    /// Gets the parent that this container sits in.
    ///
    /// If the container is the root, it returns None
    fn get_parent(&self) -> Option<&Containable>;

    /// Gets the children of this container.
    ///
    /// Views never have children
    fn get_children(&self) -> Vec<Weak<Box<Containable>>>;

    /// Gets the type of the container
    fn get_type(&self) -> ContainerType;

    /// Returns true if this container is focused.
    fn is_focused(&self) -> bool;

    /// Removes this container and all of its children
    fn remove_container(self);

    /// Sets this container (and everything in it) to given visibility
    fn set_visibility(&mut self, visibilty: bool);

    /// Gets the X and Y dimensions of the container
    fn get_dimensions(&self) -> (u64, u64);

    /// Gets the position of this container on the screen
    fn get_position(&self) -> (i64, i64);

    /// Finds a parent container with the given type
    fn get_parent_by_type(&self, type_: ContainerType) -> Option<&Containable>;
}

/// View specific functions
pub trait Viewable {
    /// Determines if the view is full screen
    fn is_fullscreen(&self) -> bool;

    /// Figures out if the view is focused
    fn is_active(&self) -> bool;

    /// Returns true if this view is a parent is an ancestor of the child
    fn is_parent_of<T: Containable>(&self, child: T) -> bool;

    /// Returns true if this view is a child is an decedent of the parent
    fn is_child_of<T: Containable>(&self, parent: T) -> bool;

    /// Gets the active workspace of the view
    fn active_workspace<T: Containable>(&self) -> T;
}

struct Container<T: Containable> {
    handle: Option<WlcOutput>,

    parent: Box<T>,
    children: Vec<Rc<Box<Containable>>>,
    type_: ContainerType,
    layout: Layout,

    width: u64,
    height: u64,

    x: i64,
    y: i64,

    visible: bool,
    is_focused: bool,
    is_floating: bool,
}

impl<C: Containable> Containable for Container<C> {

    /// Gets the parent that this container sits in.
    ///
    /// If the container is the root, it returns None
    fn get_parent(&self) -> Option<&Containable> {
        match self.type_ {
            ContainerType::Root => None,
            _ => Some(&*self.parent as &Containable),
        }
    }
    /// Gets the children of this container.
    ///
    /// Views never have children
    fn get_children(&self) -> Vec<Weak<Box<Containable>>> {
        self.children.iter().map(|child| Rc::downgrade(&child)).collect()
    }

    fn get_type(&self) -> ContainerType {
        self.type_
    }

    /// Returns true if this container is focused.
    fn is_focused(&self) -> bool {
        self.is_focused
    }

    /// Removes this container and all of its children
    fn remove_container(self) {
        unimplemented!();
    }

    /// Sets this container (and everything in it) to given visibility
    fn set_visibility(&mut self, visibilty: bool) {
        self.visible = visibilty
    }

    /// Gets the X (width) and Y (height) dimensions of the container
    fn get_dimensions(&self) -> (u64, u64) {
        (self.width, self.height)
    }

    /// Gets the position of this container on the screen
    fn get_position(&self) -> (i64, i64) {
        (self.x, self.y)
    }

    /// Finds a parent container with the given type, if there is any
    fn get_parent_by_type(&self, type_: ContainerType) -> Option<&Containable> {
        let mut container = self.get_parent();
        while container.is_some() && container.unwrap().get_type() != type_ {
            container = container.unwrap().get_parent();
        }
        container
    }
}

struct View {
    handle: Option<WlcView>,

    width: u64,
    height: u64,

    x: i64,
    y: i64,

    visible: bool,
    is_focused: bool,
    is_floating: bool,
}
<commit_msg>Implemented remove_container. Should definitely be tested to make sure it works<commit_after>//! Layout handling

use rustwlc::handle::{WlcView, WlcOutput};
use std::rc::{Rc, Weak};

#[derive(PartialEq, Clone, Copy)]
pub enum ContainerType {
    Root,        /* Root container, only one exists */
    Output,      /* Output, like a monitor or head */
    Workspace,   /* A workspace */
    View         /* A view (aka a window) */
}

enum Layout {
    None,
    Horizontal,
    Vertical,
    Stacked,
    Tabbed,
    Floating
}

/// Like i3, everything (workspaces, containers, views) are containable.
pub trait Containable {
    /// Gets the parent that this container sits in.
    ///
    /// If the container is the root, it returns None
    fn get_parent(&self) -> Option<&Containable>;

    /// Gets the children of this container.
    ///
    /// Views never have children
    fn get_children(&self) -> Vec<Weak<Box<Containable>>>;

    /// Gets the type of the container
    fn get_type(&self) -> ContainerType;

    /// Returns true if this container is focused.
    fn is_focused(&self) -> bool;

    /// Removes this container and all of its children
    fn remove_container(&self);

    /// Sets this container (and everything in it) to given visibility
    fn set_visibility(&mut self, visibilty: bool);

    /// Gets the X and Y dimensions of the container
    fn get_dimensions(&self) -> (u64, u64);

    /// Gets the position of this container on the screen
    fn get_position(&self) -> (i64, i64);

    /// Finds a parent container with the given type
    fn get_parent_by_type(&self, type_: ContainerType) -> Option<&Containable>;
}

/// View specific functions
pub trait Viewable {
    /// Determines if the view is full screen
    fn is_fullscreen(&self) -> bool;

    /// Figures out if the view is focused
    fn is_active(&self) -> bool;

    /// Returns true if this view is a parent is an ancestor of the child
    fn is_parent_of<T: Containable>(&self, child: T) -> bool;

    /// Returns true if this view is a child is an decedent of the parent
    fn is_child_of<T: Containable>(&self, parent: T) -> bool;

    /// Gets the active workspace of the view
    fn active_workspace<T: Containable>(&self) -> T;
}

struct Container<T: Containable> {
    handle: Option<WlcOutput>,

    parent: Box<T>,
    children: Vec<Rc<Box<Containable>>>,
    type_: ContainerType,
    layout: Layout,

    width: u64,
    height: u64,

    x: i64,
    y: i64,

    visible: bool,
    is_focused: bool,
    is_floating: bool,
}

impl<C: Containable> Containable for Container<C> {

    /// Gets the parent that this container sits in.
    ///
    /// If the container is the root, it returns None
    fn get_parent(&self) -> Option<&Containable> {
        match self.type_ {
            ContainerType::Root => None,
            _ => Some(&*self.parent as &Containable),
        }
    }
    /// Gets the children of this container.
    ///
    /// Views never have children
    fn get_children(&self) -> Vec<Weak<Box<Containable>>> {
        self.children.iter().map(|child| Rc::downgrade(&child)).collect()
    }

    fn get_type(&self) -> ContainerType {
        self.type_
    }

    /// Returns true if this container is focused.
    fn is_focused(&self) -> bool {
        self.is_focused
    }

    /// Removes this container and all of its children
    fn remove_container(&self) {
        for child in self.get_children() {
            if let Some(child) = child.upgrade() {
                if let Ok(child) = Rc::try_unwrap(child) {
                    child.remove_container();
                    drop(child);
                }
            }
        }
    }

    /// Sets this container (and everything in it) to given visibility
    fn set_visibility(&mut self, visibilty: bool) {
        self.visible = visibilty
    }

    /// Gets the X (width) and Y (height) dimensions of the container
    fn get_dimensions(&self) -> (u64, u64) {
        (self.width, self.height)
    }

    /// Gets the position of this container on the screen
    fn get_position(&self) -> (i64, i64) {
        (self.x, self.y)
    }

    /// Finds a parent container with the given type, if there is any
    fn get_parent_by_type(&self, type_: ContainerType) -> Option<&Containable> {
        let mut container = self.get_parent();
        while container.is_some() && container.unwrap().get_type() != type_ {
            container = container.unwrap().get_parent();
        }
        container
    }
}

struct View {
    handle: Option<WlcView>,

    width: u64,
    height: u64,

    x: i64,
    y: i64,

    visible: bool,
    is_focused: bool,
    is_floating: bool,
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cmp;
use collections::HashSet;
use std::local_data;
use std::uint;
use syntax::ast;
use rustc::util::nodemap::NodeSet;

use clean;
use clean::Item;
use plugins;
use fold;
use fold::DocFolder;

/// Strip items marked `#[doc(hidden)]`
pub fn strip_hidden(krate: clean::Crate) -> plugins::PluginResult {
    let mut stripped = HashSet::new();

    // strip all #[doc(hidden)] items
    let krate = {
        struct Stripper<'a> {
            stripped: &'a mut HashSet<ast::NodeId>
        };
        impl<'a> fold::DocFolder for Stripper<'a> {
            fn fold_item(&mut self, i: Item) -> Option<Item> {
                for attr in i.attrs.iter() {
                    match attr {
                        &clean::List(ref x, ref l) if "doc" == *x => {
                            for innerattr in l.iter() {
                                match innerattr {
                                    &clean::Word(ref s) if "hidden" == *s => {
                                        debug!("found one in strip_hidden; removing");
                                        self.stripped.insert(i.id);
                                        return None;
                                    },
                                    _ => (),
                                }
                            }
                        },
                        _ => ()
                    }
                }
                self.fold_item_recur(i)
            }
        }
        let mut stripper = Stripper{ stripped: &mut stripped };
        stripper.fold_crate(krate)
    };

    // strip any traits implemented on stripped items
    let krate = {
        struct ImplStripper<'a> {
            stripped: &'a mut HashSet<ast::NodeId>
        };
        impl<'a> fold::DocFolder for ImplStripper<'a> {
            fn fold_item(&mut self, i: Item) -> Option<Item> {
                match i.inner {
                    clean::ImplItem(clean::Impl{ for_: clean::ResolvedPath{ id: for_id, .. },
                                                 .. }) => {
                        if self.stripped.contains(&for_id) {
                            return None;
                        }
                    }
                    _ => {}
                }
                self.fold_item_recur(i)
            }
        }
        let mut stripper = ImplStripper{ stripped: &mut stripped };
        stripper.fold_crate(krate)
    };

    (krate, None)
}

/// Strip private items from the point of view of a crate or externally from a
/// crate, specified by the `xcrate` flag.
pub fn strip_private(krate: clean::Crate) -> plugins::PluginResult {
    // This stripper collects all *retained* nodes.
    let mut retained = HashSet::new();
    let exported_items = local_data::get(super::analysiskey, |analysis| {
        analysis.unwrap().exported_items.clone()
    });
    let mut krate = krate;

    // strip all private items
    {
        let mut stripper = Stripper {
            retained: &mut retained,
            exported_items: &exported_items,
        };
        krate = stripper.fold_crate(krate);
    }

    // strip all private implementations of traits
    {
        let mut stripper = ImplStripper(&retained);
        krate = stripper.fold_crate(krate);
    }
    (krate, None)
}

struct Stripper<'a> {
    retained: &'a mut HashSet<ast::NodeId>,
    exported_items: &'a NodeSet,
}

impl<'a> fold::DocFolder for Stripper<'a> {
    fn fold_item(&mut self, i: Item) -> Option<Item> {
        match i.inner {
            // These items can all get re-exported
            clean::TypedefItem(..) | clean::StaticItem(..) |
            clean::StructItem(..) | clean::EnumItem(..) |
            clean::TraitItem(..) | clean::FunctionItem(..) |
            clean::VariantItem(..) | clean::MethodItem(..) |
            clean::ForeignFunctionItem(..) | clean::ForeignStaticItem(..) => {
                if !self.exported_items.contains(&i.id) {
                    return None;
                }
            }

            clean::ViewItemItem(..) => {
                if i.visibility != Some(ast::Public) {
                    return None
                }
            }

            clean::StructFieldItem(..) => {
                if i.visibility == Some(ast::Private) {
                    return None;
                }
            }

            // handled below
            clean::ModuleItem(..) => {}

            // trait impls for private items should be stripped
            clean::ImplItem(clean::Impl{ for_: clean::ResolvedPath{ id: ref for_id, .. }, .. }) => {
                if !self.exported_items.contains(for_id) {
                    return None;
                }
            }
            clean::ImplItem(..) => {}

            // tymethods/macros have no control over privacy
            clean::MacroItem(..) | clean::TyMethodItem(..) => {}
        }

        let fastreturn = match i.inner {
            // nothing left to do for traits (don't want to filter their
            // methods out, visibility controlled by the trait)
            clean::TraitItem(..) => true,

            // implementations of traits are always public.
            clean::ImplItem(ref imp) if imp.trait_.is_some() => true,

            _ => false,
        };

        let i = if fastreturn {
            self.retained.insert(i.id);
            return Some(i);
        } else {
            self.fold_item_recur(i)
        };

        match i {
            Some(i) => {
                match i.inner {
                    // emptied modules/impls have no need to exist
                    clean::ModuleItem(ref m)
                        if m.items.len() == 0 &&
                           i.doc_value().is_none() => None,
                    clean::ImplItem(ref i) if i.methods.len() == 0 => None,
                    _ => {
                        self.retained.insert(i.id);
                        Some(i)
                    }
                }
            }
            None => None,
        }
    }
}

// This stripper discards all private impls of traits
struct ImplStripper<'a>(&'a HashSet<ast::NodeId>);
impl<'a> fold::DocFolder for ImplStripper<'a> {
    fn fold_item(&mut self, i: Item) -> Option<Item> {
        match i.inner {
            clean::ImplItem(ref imp) => {
                match imp.trait_ {
                    Some(clean::ResolvedPath{ id, .. }) => {
                        let ImplStripper(s) = *self;
                        if !s.contains(&id) {
                            return None;
                        }
                    }
                    Some(..) | None => {}
                }
            }
            _ => {}
        }
        self.fold_item_recur(i)
    }
}

pub fn unindent_comments(krate: clean::Crate) -> plugins::PluginResult {
    struct CommentCleaner;
    impl fold::DocFolder for CommentCleaner {
        fn fold_item(&mut self, i: Item) -> Option<Item> {
            let mut i = i;
            let mut avec: Vec<clean::Attribute> = Vec::new();
            for attr in i.attrs.iter() {
                match attr {
                    &clean::NameValue(ref x, ref s) if "doc" == *x => avec.push(
                        clean::NameValue(~"doc", unindent(*s))),
                    x => avec.push(x.clone())
                }
            }
            i.attrs = avec;
            self.fold_item_recur(i)
        }
    }
    let mut cleaner = CommentCleaner;
    let krate = cleaner.fold_crate(krate);
    (krate, None)
}

pub fn collapse_docs(krate: clean::Crate) -> plugins::PluginResult {
    struct Collapser;
    impl fold::DocFolder for Collapser {
        fn fold_item(&mut self, i: Item) -> Option<Item> {
            let mut docstr = ~"";
            let mut i = i;
            for attr in i.attrs.iter() {
                match *attr {
                    clean::NameValue(ref x, ref s) if "doc" == *x => {
                        docstr.push_str(s.clone());
                        docstr.push_char('\n');
                    },
                    _ => ()
                }
            }
            let mut a: Vec<clean::Attribute> = i.attrs.iter().filter(|&a| match a {
                &clean::NameValue(ref x, _) if "doc" == *x => false,
                _ => true
            }).map(|x| x.clone()).collect();
            if "" != docstr {
                a.push(clean::NameValue(~"doc", docstr));
            }
            i.attrs = a;
            self.fold_item_recur(i)
        }
    }
    let mut collapser = Collapser;
    let krate = collapser.fold_crate(krate);
    (krate, None)
}

pub fn unindent(s: &str) -> ~str {
    let lines = s.lines_any().collect::<Vec<&str> >();
    let mut saw_first_line = false;
    let mut saw_second_line = false;
    let min_indent = lines.iter().fold(uint::MAX, |min_indent, line| {

        // After we see the first non-whitespace line, look at
        // the line we have. If it is not whitespace, and therefore
        // part of the first paragraph, then ignore the indentation
        // level of the first line
        let ignore_previous_indents =
            saw_first_line &&
            !saw_second_line &&
            !line.is_whitespace();

        let min_indent = if ignore_previous_indents {
            uint::MAX
        } else {
            min_indent
        };

        if saw_first_line {
            saw_second_line = true;
        }

        if line.is_whitespace() {
            min_indent
        } else {
            saw_first_line = true;
            let mut spaces = 0;
            line.chars().all(|char| {
                // Only comparing against space because I wouldn't
                // know what to do with mixed whitespace chars
                if char == ' ' {
                    spaces += 1;
                    true
                } else {
                    false
                }
            });
            cmp::min(min_indent, spaces)
        }
    });

    if lines.len() >= 1 {
        let mut unindented = vec!( lines.get(0).trim() );
        unindented.push_all(lines.tail().iter().map(|&line| {
            if line.is_whitespace() {
                line
            } else {
                assert!(line.len() >= min_indent);
                line.slice_from(min_indent)
            }
        }).collect::<Vec<_>>().as_slice());
        unindented.connect("\n")
    } else {
        s.to_owned()
    }
}

#[cfg(test)]
mod unindent_tests {
    use super::unindent;

    #[test]
    fn should_unindent() {
        let s = ~"    line1\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\nline2");
    }

    #[test]
    fn should_unindent_multiple_paragraphs() {
        let s = ~"    line1\n\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\nline2");
    }

    #[test]
    fn should_leave_multiple_indent_levels() {
        // Line 2 is indented another level beyond the
        // base indentation and should be preserved
        let s = ~"    line1\n\n        line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\n    line2");
    }

    #[test]
    fn should_ignore_first_line_indent() {
        // Thi first line of the first paragraph may not be indented as
        // far due to the way the doc string was written:
        //
        // #[doc = "Start way over here
        //          and continue here"]
        let s = ~"line1\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\nline2");
    }

    #[test]
    fn should_not_ignore_first_line_indent_in_a_single_line_para() {
        let s = ~"line1\n\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\n    line2");
    }
}
<commit_msg>auto merge of #13319 : alexcrichton/rust/rustdoc-fields, r=brson<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cmp;
use collections::HashSet;
use std::local_data;
use std::uint;
use syntax::ast;
use rustc::util::nodemap::NodeSet;

use clean;
use clean::Item;
use plugins;
use fold;
use fold::DocFolder;

/// Strip items marked `#[doc(hidden)]`
pub fn strip_hidden(krate: clean::Crate) -> plugins::PluginResult {
    let mut stripped = HashSet::new();

    // strip all #[doc(hidden)] items
    let krate = {
        struct Stripper<'a> {
            stripped: &'a mut HashSet<ast::NodeId>
        };
        impl<'a> fold::DocFolder for Stripper<'a> {
            fn fold_item(&mut self, i: Item) -> Option<Item> {
                for attr in i.attrs.iter() {
                    match attr {
                        &clean::List(ref x, ref l) if "doc" == *x => {
                            for innerattr in l.iter() {
                                match innerattr {
                                    &clean::Word(ref s) if "hidden" == *s => {
                                        debug!("found one in strip_hidden; removing");
                                        self.stripped.insert(i.id);
                                        return None;
                                    },
                                    _ => (),
                                }
                            }
                        },
                        _ => ()
                    }
                }
                self.fold_item_recur(i)
            }
        }
        let mut stripper = Stripper{ stripped: &mut stripped };
        stripper.fold_crate(krate)
    };

    // strip any traits implemented on stripped items
    let krate = {
        struct ImplStripper<'a> {
            stripped: &'a mut HashSet<ast::NodeId>
        };
        impl<'a> fold::DocFolder for ImplStripper<'a> {
            fn fold_item(&mut self, i: Item) -> Option<Item> {
                match i.inner {
                    clean::ImplItem(clean::Impl{ for_: clean::ResolvedPath{ id: for_id, .. },
                                                 .. }) => {
                        if self.stripped.contains(&for_id) {
                            return None;
                        }
                    }
                    _ => {}
                }
                self.fold_item_recur(i)
            }
        }
        let mut stripper = ImplStripper{ stripped: &mut stripped };
        stripper.fold_crate(krate)
    };

    (krate, None)
}

/// Strip private items from the point of view of a crate or externally from a
/// crate, specified by the `xcrate` flag.
pub fn strip_private(krate: clean::Crate) -> plugins::PluginResult {
    // This stripper collects all *retained* nodes.
    let mut retained = HashSet::new();
    let exported_items = local_data::get(super::analysiskey, |analysis| {
        analysis.unwrap().exported_items.clone()
    });
    let mut krate = krate;

    // strip all private items
    {
        let mut stripper = Stripper {
            retained: &mut retained,
            exported_items: &exported_items,
        };
        krate = stripper.fold_crate(krate);
    }

    // strip all private implementations of traits
    {
        let mut stripper = ImplStripper(&retained);
        krate = stripper.fold_crate(krate);
    }
    (krate, None)
}

struct Stripper<'a> {
    retained: &'a mut HashSet<ast::NodeId>,
    exported_items: &'a NodeSet,
}

impl<'a> fold::DocFolder for Stripper<'a> {
    fn fold_item(&mut self, i: Item) -> Option<Item> {
        match i.inner {
            // These items can all get re-exported
            clean::TypedefItem(..) | clean::StaticItem(..) |
            clean::StructItem(..) | clean::EnumItem(..) |
            clean::TraitItem(..) | clean::FunctionItem(..) |
            clean::VariantItem(..) | clean::MethodItem(..) |
            clean::ForeignFunctionItem(..) | clean::ForeignStaticItem(..) => {
                if !self.exported_items.contains(&i.id) {
                    return None;
                }
            }

            clean::ViewItemItem(..) | clean::StructFieldItem(..) => {
                if i.visibility != Some(ast::Public) {
                    return None
                }
            }

            // handled below
            clean::ModuleItem(..) => {}

            // trait impls for private items should be stripped
            clean::ImplItem(clean::Impl{ for_: clean::ResolvedPath{ id: ref for_id, .. }, .. }) => {
                if !self.exported_items.contains(for_id) {
                    return None;
                }
            }
            clean::ImplItem(..) => {}

            // tymethods/macros have no control over privacy
            clean::MacroItem(..) | clean::TyMethodItem(..) => {}
        }

        let fastreturn = match i.inner {
            // nothing left to do for traits (don't want to filter their
            // methods out, visibility controlled by the trait)
            clean::TraitItem(..) => true,

            // implementations of traits are always public.
            clean::ImplItem(ref imp) if imp.trait_.is_some() => true,

            _ => false,
        };

        let i = if fastreturn {
            self.retained.insert(i.id);
            return Some(i);
        } else {
            self.fold_item_recur(i)
        };

        match i {
            Some(i) => {
                match i.inner {
                    // emptied modules/impls have no need to exist
                    clean::ModuleItem(ref m)
                        if m.items.len() == 0 &&
                           i.doc_value().is_none() => None,
                    clean::ImplItem(ref i) if i.methods.len() == 0 => None,
                    _ => {
                        self.retained.insert(i.id);
                        Some(i)
                    }
                }
            }
            None => None,
        }
    }
}

// This stripper discards all private impls of traits
struct ImplStripper<'a>(&'a HashSet<ast::NodeId>);
impl<'a> fold::DocFolder for ImplStripper<'a> {
    fn fold_item(&mut self, i: Item) -> Option<Item> {
        match i.inner {
            clean::ImplItem(ref imp) => {
                match imp.trait_ {
                    Some(clean::ResolvedPath{ id, .. }) => {
                        let ImplStripper(s) = *self;
                        if !s.contains(&id) {
                            return None;
                        }
                    }
                    Some(..) | None => {}
                }
            }
            _ => {}
        }
        self.fold_item_recur(i)
    }
}

pub fn unindent_comments(krate: clean::Crate) -> plugins::PluginResult {
    struct CommentCleaner;
    impl fold::DocFolder for CommentCleaner {
        fn fold_item(&mut self, i: Item) -> Option<Item> {
            let mut i = i;
            let mut avec: Vec<clean::Attribute> = Vec::new();
            for attr in i.attrs.iter() {
                match attr {
                    &clean::NameValue(ref x, ref s) if "doc" == *x => avec.push(
                        clean::NameValue(~"doc", unindent(*s))),
                    x => avec.push(x.clone())
                }
            }
            i.attrs = avec;
            self.fold_item_recur(i)
        }
    }
    let mut cleaner = CommentCleaner;
    let krate = cleaner.fold_crate(krate);
    (krate, None)
}

pub fn collapse_docs(krate: clean::Crate) -> plugins::PluginResult {
    struct Collapser;
    impl fold::DocFolder for Collapser {
        fn fold_item(&mut self, i: Item) -> Option<Item> {
            let mut docstr = ~"";
            let mut i = i;
            for attr in i.attrs.iter() {
                match *attr {
                    clean::NameValue(ref x, ref s) if "doc" == *x => {
                        docstr.push_str(s.clone());
                        docstr.push_char('\n');
                    },
                    _ => ()
                }
            }
            let mut a: Vec<clean::Attribute> = i.attrs.iter().filter(|&a| match a {
                &clean::NameValue(ref x, _) if "doc" == *x => false,
                _ => true
            }).map(|x| x.clone()).collect();
            if "" != docstr {
                a.push(clean::NameValue(~"doc", docstr));
            }
            i.attrs = a;
            self.fold_item_recur(i)
        }
    }
    let mut collapser = Collapser;
    let krate = collapser.fold_crate(krate);
    (krate, None)
}

pub fn unindent(s: &str) -> ~str {
    let lines = s.lines_any().collect::<Vec<&str> >();
    let mut saw_first_line = false;
    let mut saw_second_line = false;
    let min_indent = lines.iter().fold(uint::MAX, |min_indent, line| {

        // After we see the first non-whitespace line, look at
        // the line we have. If it is not whitespace, and therefore
        // part of the first paragraph, then ignore the indentation
        // level of the first line
        let ignore_previous_indents =
            saw_first_line &&
            !saw_second_line &&
            !line.is_whitespace();

        let min_indent = if ignore_previous_indents {
            uint::MAX
        } else {
            min_indent
        };

        if saw_first_line {
            saw_second_line = true;
        }

        if line.is_whitespace() {
            min_indent
        } else {
            saw_first_line = true;
            let mut spaces = 0;
            line.chars().all(|char| {
                // Only comparing against space because I wouldn't
                // know what to do with mixed whitespace chars
                if char == ' ' {
                    spaces += 1;
                    true
                } else {
                    false
                }
            });
            cmp::min(min_indent, spaces)
        }
    });

    if lines.len() >= 1 {
        let mut unindented = vec!( lines.get(0).trim() );
        unindented.push_all(lines.tail().iter().map(|&line| {
            if line.is_whitespace() {
                line
            } else {
                assert!(line.len() >= min_indent);
                line.slice_from(min_indent)
            }
        }).collect::<Vec<_>>().as_slice());
        unindented.connect("\n")
    } else {
        s.to_owned()
    }
}

#[cfg(test)]
mod unindent_tests {
    use super::unindent;

    #[test]
    fn should_unindent() {
        let s = ~"    line1\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\nline2");
    }

    #[test]
    fn should_unindent_multiple_paragraphs() {
        let s = ~"    line1\n\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\nline2");
    }

    #[test]
    fn should_leave_multiple_indent_levels() {
        // Line 2 is indented another level beyond the
        // base indentation and should be preserved
        let s = ~"    line1\n\n        line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\n    line2");
    }

    #[test]
    fn should_ignore_first_line_indent() {
        // Thi first line of the first paragraph may not be indented as
        // far due to the way the doc string was written:
        //
        // #[doc = "Start way over here
        //          and continue here"]
        let s = ~"line1\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\nline2");
    }

    #[test]
    fn should_not_ignore_first_line_indent_in_a_single_line_para() {
        let s = ~"line1\n\n    line2";
        let r = unindent(s);
        assert_eq!(r, ~"line1\n\n    line2");
    }
}
<|endoftext|>


<commit_before>extern crate cbor;
extern crate rustc_serialize;
extern crate secstr;
extern crate rusterpassword;
extern crate libc;
extern crate libsodium_sys as sodium;
extern crate sodiumoxide;
extern crate chrono;

use secstr::SecStr;
use rusterpassword::gen_site_seed;
use chrono::{DateTime, UTC};
use cbor::{Encoder, Decoder, CborError};
use libc::size_t;
use std::collections::btree_map::{BTreeMap, Keys};
use sodiumoxide::crypto::secretbox::xsalsa20poly1305 as secbox;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Vault {
    version: u16,
    padding: Vec<u8>,
    entries: BTreeMap<String, EncryptedEntry>
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub struct EncryptedEntry {
    nonce: Vec<u8>,
    counter: u32,
    ciphertext: Vec<u8>,
    metadata: EntryMetadata
}

#[derive(PartialEq, Clone, Debug, RustcDecodable, RustcEncodable)]
pub struct EntryMetadata {
    created_at: DateTime<UTC>,
    updated_at: DateTime<UTC>
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub struct Entry {
    fields: BTreeMap<String, Field>
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum Field {
    Derived { counter: u32, site_name: Option<String>, usage: DerivedUsage },
    Stored { data: SecStr, usage: StoredUsage }
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum DerivedUsage {
    Password(PasswordTemplate)
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum StoredUsage {
    Password
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum PasswordTemplate {
    Maximum, Long, Medium, Short, Basic, Pin
}

#[derive(Debug)]
pub enum EntryError {
    WrongEntriesKeyLength,
    WrongNonceLength,
    SeedGenerationError,
    DecryptionError,
    CodecError(CborError),
    DataError,
    EntryNotFound
}

impl From<CborError> for EntryError {
    fn from(err: CborError) -> EntryError {
        EntryError::CodecError(err)
    }
}

pub type EntryResult<T> = Result<T, EntryError>;

impl Vault {

    pub fn entry_names(&self) -> Keys<String, EncryptedEntry> {
        self.entries.keys()
    }

    pub fn get_entry(&self, entries_key: &SecStr, name: &str) -> EntryResult<(Entry, EntryMetadata)> {
        if let Some(ee) = self.entries.get(name) {
            let nonce_wrapped = try!(secbox::Nonce::from_slice(&ee.nonce).ok_or(EntryError::WrongNonceLength));
            let entry_key_wrapped = try!(gen_entry_key(entries_key, name, ee.counter));
            let plaintext = SecStr::new(try!(secbox::open(&ee.ciphertext, &nonce_wrapped, &entry_key_wrapped).map_err(|_| EntryError::DecryptionError)));
            let entry = try!(try!(Decoder::from_bytes(plaintext.unsecure()).decode::<Entry>().next().ok_or(EntryError::DataError)));
            Ok((entry, ee.metadata.clone()))
        } else {
            Err(EntryError::EntryNotFound)
        }
    }

    pub fn put_entry(&mut self, entries_key: &SecStr, name: &str, entry: &Entry, metadata: &mut EntryMetadata) -> EntryResult<()> {
        metadata.updated_at = UTC::now();
        let counter = self.entries.get(name).map(|ee| ee.counter + 1).unwrap_or(1);
        let nonce_wrapped = secbox::gen_nonce();
        let secbox::Nonce(nonce) = nonce_wrapped;
        let entry_key_wrapped = try!(gen_entry_key(entries_key, name, counter));
        let mut e = Encoder::from_memory();
        try!(e.encode(&[&*entry]));
        let plaintext = SecStr::new(e.into_bytes());
        let ciphertext = secbox::seal(plaintext.unsecure(), &nonce_wrapped, &entry_key_wrapped);
        self.entries.insert(name.to_owned(), EncryptedEntry {
            nonce: nonce.to_vec(), counter: counter, ciphertext: ciphertext, metadata: metadata.clone()
        });
        Ok(())
    }

}

fn gen_entry_key(entries_key: &SecStr, name: &str, counter: u32) -> EntryResult<secbox::Key> {
    let entry_key = try!(gen_site_seed(entries_key, name, counter).map_err(|_| EntryError::SeedGenerationError));
    Ok(try!(secbox::Key::from_slice(entry_key.unsecure()).ok_or(EntryError::WrongEntriesKeyLength)))
}

pub fn gen_entries_key(master_key: &SecStr) -> SecStr {
    let mut msg = vec![];
    msg.extend(b"technology.unrelenting.freepass");
    let mut dst = Vec::<u8>::with_capacity(64);
    unsafe {
        sodium::crypto_generichash_blake2b(
            dst.as_mut_ptr() as *mut u8, 64,
            msg.as_ptr(), msg.len() as u64,
            master_key.unsecure().as_ptr() as *const u8,
            master_key.unsecure().len() as size_t);
        dst.set_len(64);
    }
    SecStr::new(dst)
}

#[cfg(test)]
mod tests {
    use super::*;
    use secstr::*;
    use rusterpassword::*;
    use std::collections::btree_map::BTreeMap;
    use chrono::*;

    #[test]
    fn test_roundtrip_entry() {
        let mut fs = BTreeMap::new();
        fs.insert("password".to_owned(), Field::Derived { counter: 4, site_name: Some("twitter.com".to_owned()), usage: DerivedUsage::Password(PasswordTemplate::Maximum) });
        fs.insert("old_password".to_owned(), Field::Stored { data: SecStr::from("h0rse"), usage: StoredUsage::Password });
        let twitter = Entry { fields: fs };
        let mut metadata = EntryMetadata { created_at: UTC::now(), updated_at: UTC::now() };
        let mut vault = Vault { version: 0, padding: b"".to_vec(), entries: BTreeMap::new() };
        let master_key = gen_master_key(SecStr::from("Correct Horse Battery Staple"), "Clarke Griffin").unwrap();
        let entries_key = gen_entries_key(&master_key);
        vault.put_entry(&entries_key, "twitter", &twitter, &mut metadata).unwrap();
        assert!(vault.get_entry(&entries_key, "twitter").unwrap() == (twitter, metadata));
    }
}
<commit_msg>public fields, constant enum for password templates<commit_after>extern crate cbor;
extern crate rustc_serialize;
extern crate secstr;
extern crate rusterpassword;
extern crate libc;
extern crate libsodium_sys as sodium;
extern crate sodiumoxide;
extern crate chrono;

use secstr::SecStr;
use rusterpassword::gen_site_seed;
use chrono::{DateTime, UTC};
use cbor::{Encoder, Decoder, CborError};
use libc::size_t;
use std::collections::btree_map::{BTreeMap, Keys};
use sodiumoxide::crypto::secretbox::xsalsa20poly1305 as secbox;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Vault {
    pub version: u16,
    pub padding: Vec<u8>,
    pub entries: BTreeMap<String, EncryptedEntry>,
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub struct EncryptedEntry {
    pub nonce: Vec<u8>,
    pub counter: u32,
    pub ciphertext: Vec<u8>,
    pub metadata: EntryMetadata,
}

#[derive(PartialEq, Clone, Debug, RustcDecodable, RustcEncodable)]
pub struct EntryMetadata {
    pub created_at: DateTime<UTC>,
    pub updated_at: DateTime<UTC>,
    pub tags: Vec<String>,
}

impl EntryMetadata {
    pub fn new() -> EntryMetadata {
        EntryMetadata {
            created_at: UTC::now(),
            updated_at: UTC::now(),
            tags: Vec::new(),
        }
    }
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub struct Entry {
    pub fields: BTreeMap<String, Field>
}

impl Entry {
    pub fn new() -> Entry {
        Entry { fields: BTreeMap::new() }
    }
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum Field {
    Derived { counter: u32, site_name: Option<String>, usage: DerivedUsage },
    Stored { data: SecStr, usage: StoredUsage }
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum DerivedUsage {
    Password(PasswordTemplate)
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum StoredUsage {
    Password
}

#[derive(PartialEq, Debug, RustcDecodable, RustcEncodable)]
pub enum PasswordTemplate {
    // Same numbers as in the rusterpassword C API
    Maximum = 60,
    Long    = 50,
    Medium  = 40,
    Short   = 30,
    Basic   = 20,
    Pin     = 10,
}

#[derive(Debug)]
pub enum EntryError {
    WrongEntriesKeyLength,
    WrongNonceLength,
    SeedGenerationError,
    DecryptionError,
    CodecError(CborError),
    DataError,
    EntryNotFound,
}

impl From<CborError> for EntryError {
    fn from(err: CborError) -> EntryError {
        EntryError::CodecError(err)
    }
}

pub type EntryResult<T> = Result<T, EntryError>;

impl Vault {

    pub fn entry_names(&self) -> Keys<String, EncryptedEntry> {
        self.entries.keys()
    }

    pub fn get_entry(&self, entries_key: &SecStr, name: &str) -> EntryResult<(Entry, EntryMetadata)> {
        if let Some(ee) = self.entries.get(name) {
            let nonce_wrapped = try!(secbox::Nonce::from_slice(&ee.nonce).ok_or(EntryError::WrongNonceLength));
            let entry_key_wrapped = try!(gen_entry_key(entries_key, name, ee.counter));
            let plaintext = SecStr::new(try!(secbox::open(&ee.ciphertext, &nonce_wrapped, &entry_key_wrapped).map_err(|_| EntryError::DecryptionError)));
            let entry = try!(try!(Decoder::from_bytes(plaintext.unsecure()).decode::<Entry>().next().ok_or(EntryError::DataError)));
            Ok((entry, ee.metadata.clone()))
        } else {
            Err(EntryError::EntryNotFound)
        }
    }

    pub fn put_entry(&mut self, entries_key: &SecStr, name: &str, entry: &Entry, metadata: &mut EntryMetadata) -> EntryResult<()> {
        let counter = self.entries.get(name).map(|ee| ee.counter + 1).unwrap_or(1);
        let nonce_wrapped = secbox::gen_nonce();
        let secbox::Nonce(nonce) = nonce_wrapped;
        let entry_key_wrapped = try!(gen_entry_key(entries_key, name, counter));
        let mut e = Encoder::from_memory();
        try!(e.encode(&[&*entry]));
        let plaintext = SecStr::new(e.into_bytes());
        let ciphertext = secbox::seal(plaintext.unsecure(), &nonce_wrapped, &entry_key_wrapped);
        metadata.updated_at = UTC::now();
        self.entries.insert(name.to_owned(), EncryptedEntry {
            nonce: nonce.to_vec(), counter: counter, ciphertext: ciphertext, metadata: metadata.clone()
        });
        Ok(())
    }

}

fn gen_entry_key(entries_key: &SecStr, name: &str, counter: u32) -> EntryResult<secbox::Key> {
    let entry_key = try!(gen_site_seed(entries_key, name, counter).map_err(|_| EntryError::SeedGenerationError));
    Ok(try!(secbox::Key::from_slice(entry_key.unsecure()).ok_or(EntryError::WrongEntriesKeyLength)))
}

pub fn gen_entries_key(master_key: &SecStr) -> SecStr {
    let mut msg = vec![];
    msg.extend(b"technology.unrelenting.freepass");
    let mut dst = Vec::<u8>::with_capacity(64);
    unsafe {
        sodium::crypto_generichash_blake2b(
            dst.as_mut_ptr() as *mut u8, 64,
            msg.as_ptr(), msg.len() as u64,
            master_key.unsecure().as_ptr() as *const u8,
            master_key.unsecure().len() as size_t);
        dst.set_len(64);
    }
    SecStr::new(dst)
}

#[cfg(test)]
mod tests {
    use super::*;
    use secstr::*;
    use rusterpassword::*;
    use std::collections::btree_map::BTreeMap;

    #[test]
    fn test_roundtrip_entry() {
        let mut fs = BTreeMap::new();
        fs.insert("password".to_owned(), Field::Derived { counter: 4, site_name: Some("twitter.com".to_owned()), usage: DerivedUsage::Password(PasswordTemplate::Maximum) });
        fs.insert("old_password".to_owned(), Field::Stored { data: SecStr::from("h0rse"), usage: StoredUsage::Password });
        let twitter = Entry { fields: fs };
        let mut metadata = EntryMetadata::new();
        let mut vault = Vault { version: 0, padding: b"".to_vec(), entries: BTreeMap::new() };
        let master_key = gen_master_key(SecStr::from("Correct Horse Battery Staple"), "Clarke Griffin").unwrap();
        let entries_key = gen_entries_key(&master_key);
        vault.put_entry(&entries_key, "twitter", &twitter, &mut metadata).unwrap();
        assert!(vault.get_entry(&entries_key, "twitter").unwrap() == (twitter, metadata));
    }
}
<|endoftext|>