<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::write_unchecked;
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_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::BTreeSet;
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),
}

enum PruningSearch {
    DeepestInnode,
    Innode(NodeCellRef),
}

struct NodeRemoval<'a> {
    empty_pages: Vec<&'a EntryKey>,
    split: Vec<(NodeCellRef, EntryKey)>,
}

impl<'a> NodeRemoval<'a> {
    fn new() -> Self {
        Self {
            empty_pages: vec![],
            split: vec![],
        }
    }
}

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 mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn merge_innode_remnant<'a, KS, PS>(
    current_node: &mut NodeWriteGuard<KS, PS>,
    prev_key: &'a EntryKey,
    removal: &mut NodeRemoval<'a>,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let curr_right_ref = current_node.right_ref_mut_no_empty().unwrap().clone();
    let curr_innode = current_node.innode_mut();
    let curr_right_bound = &curr_innode.right_bound;
    let curr_last_child = mem::replace(&mut curr_innode.ptrs.as_slice()[0], NodeCellRef::default());
    debug_assert_eq!(curr_innode.len, 0);
    removal.empty_pages.push(prev_key);
    if curr_last_child.is_default() {
        return;
    }
    let mut next_node = write_targeted(write_node::<KS, PS>(&curr_right_ref), curr_right_bound);
    let pos = next_node.search(curr_right_bound);
    let next_innode = next_node.innode_mut();
    next_innode.debug_check_integrity();
    if next_innode.len == KS::slice_len() {
        removal.split.push(next_innode.split_insert(
            curr_right_bound.clone(),
            curr_last_child,
            pos,
            false,
        ))
    } else {
        next_innode.insert_in_place(curr_right_bound.clone(), curr_last_child, pos, false);
    }
}

fn apply_removal<'a, KS, PS>(
    cursor_guard: &mut NodeWriteGuard<KS, PS>,
    poses: &mut BTreeSet<usize>,
    node_removal: &mut NodeRemoval<'a>,
    prev_key: &Option<&'a EntryKey>,
    remove_children_right_nodes: bool,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    if poses.is_empty() {
        return;
    }
    debug!(
        "Applying removal, guard key {:?}, poses {:?}",
        cursor_guard.first_key(),
        poses
    );
    {
        let innode = cursor_guard.innode_mut();
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            if remove_children_right_nodes {
                // A node will be reclaimed if and only if it have zero references
                // If there is any sequential empty nodes, they have to be unlinked each other
                // The upper level will do the rest to unlink it from the tree
                innode.ptrs.as_slice()[..innode.len + 1]
                    .iter()
                    .enumerate()
                    .filter(|(i, _)| poses.contains(i))
                    .for_each(|(_, r)| {
                        write_node::<KS, PS>(r).right_ref_mut_no_empty();
                    });
            }
            let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, p)| p)
                .collect();

            let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .take(if ptrs.len() == 0 { 0 } else { ptrs.len() - 1 })
                .map(|(_, k)| k)
                .collect();
            innode.len = keys.len();
            debug!("Prune filtered page have keys {:?}", &keys);
            for (i, key) in keys.into_iter().enumerate() {
                debug_assert!(key > &mut smallvec!(0));
                mem::swap(&mut new_keys.as_slice()[i], key);
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                debug_assert!(!ptr.is_default());
                mem::swap(&mut new_ptrs.as_slice()[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;
        innode.debug_check_integrity();

        debug!(
            "Pruned page have keys {:?}",
            &innode.keys.as_slice_immute()[..innode.len]
        );
    }

    if cursor_guard.is_empty() {
        if let &Some(k) = prev_key {
            debug_assert!(!cursor_guard.is_ext());
            merge_innode_remnant(cursor_guard, k, node_removal);
        }
        debug!("Pruned page is empty: {:?}", prev_key);
        cursor_guard.make_empty_node();
    }
    poses.clear();
}

fn prune_selected<'a, KS, PS>(
    node: &NodeCellRef,
    mut removal: Box<NodeRemoval<'a>>,
    level: usize,
) -> Box<NodeRemoval<'a>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, removal.empty_pages.first().unwrap());
    let pruning = match first_search {
        MutSearchResult::Internal(sub_node) => {
            if read_unchecked::<KS, PS>(&NodeData::<KS, PS>::get_non_empty_node(&sub_node)).is_ext()
            {
                PruningSearch::DeepestInnode
            } else {
                PruningSearch::Innode(sub_node)
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut deepest = false;
    match pruning {
        PruningSearch::DeepestInnode => {
            debug!(
                "Removing in deepest nodes keys {:?}, level {}",
                &removal.empty_pages, level
            );
            deepest = true;
        }
        PruningSearch::Innode(sub_node) => {
            removal = prune_selected::<KS, PS>(&sub_node, removal, level + 1);
        }
    }
    // empty page references that will dealt with by upper level
    let mut upper_removal = NodeRemoval::new();
    if !removal.empty_pages.is_empty() {
        debug!(
            "Pruning page containing keys {:?}, level {}",
            &removal.empty_pages, level
        );
        // start delete
        let mut cursor_guard = write_non_empty(write_node::<KS, PS>(node));
        let mut guard_removing_poses = BTreeSet::new();
        let mut prev_key = None;
        debug!(
            "Prune deepest node starting key: {:?}, level {}",
            cursor_guard.first_key(),
            level
        );
        // iterate over all removal empty nodes
        // in the beginning cursor_guard contains the first key
        // if the empty node go out of the right bound of current cursor guard,
        // it will `apply_removal` and switch to the next right page
        for (_i, &key_to_del) in removal.empty_pages.iter().enumerate() {
            if key_to_del >= cursor_guard.right_bound() {
                debug!(
                    "Applying removal for overflow current page ({}/{}) key: {:?} >= bound: {:?}. guard keys: {:?}, level {}",
                    guard_removing_poses.len(),
                    cursor_guard.len() + 1,
                    key_to_del,
                    cursor_guard.right_bound(),
                    &cursor_guard.keys()[..cursor_guard.len()],
                    level
                );
                apply_removal(
                    &mut cursor_guard,
                    &mut guard_removing_poses,
                    &mut upper_removal,
                    &prev_key,
                    !deepest,
                );
                cursor_guard = write_targeted(cursor_guard, key_to_del);
                debug_assert!(!cursor_guard.is_empty_node());
                debug!(
                    "Applied removal for overflow current page ({}), level {}",
                    cursor_guard.len(),
                    level
                );
            }
            let pos = cursor_guard.search(key_to_del);
            debug!(
                "Key to delete have position {}, key: {:?}, level {}",
                pos, key_to_del, level
            );
            guard_removing_poses.insert(pos);
            prev_key = Some(key_to_del)
        }
        if !guard_removing_poses.is_empty() {
            debug!(
                "Applying removal for last keys {:?}, level {}",
                &guard_removing_poses, level
            );
            apply_removal(
                &mut cursor_guard,
                &mut guard_removing_poses,
                &mut upper_removal,
                &prev_key,
                !deepest,
            );
        }
    }
    if !removal.split.is_empty() {
        let mut cursor_guard = write_node::<KS, PS>(node);
        for (node, pivot) in &removal.split {
            cursor_guard = write_targeted(cursor_guard, &pivot);
            let pos = cursor_guard.search(&pivot);
            let innode = cursor_guard.innode_mut();
            let pivot = pivot.clone();
            let node = node.clone();
            if innode.len >= KS::slice_len() {
                upper_removal
                    .split
                    .push(innode.split_insert(pivot, node, pos, true));
            } else {
                innode.insert_in_place(pivot, node, pos, true)
            }
        }
    }
    debug!(
        "Have empty nodes {:?}, level {:?}",
        &upper_removal.empty_pages, level
    );
    box upper_removal
}

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;
    debug!("Merge selected {} pages", left_most_leaf_guards.len());

    // 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);
        }
    }

    // cleanup upper level references
    {
        let page_keys = left_most_leaf_guards
            .iter()
            .filter(|g| !g.is_empty())
            .map(|g| g.first_key())
            .collect_vec();
        let mut removal = NodeRemoval::new();
        removal.empty_pages = page_keys;
        prune_selected::<KS, PS>(&src_tree.get_root(), box removal, 0);
    }

    // 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());

        let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
        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 = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

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

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<commit_msg>consider none left most node<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::write_unchecked;
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_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::BTreeSet;
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),
}

enum PruningSearch {
    DeepestInnode,
    Innode(NodeCellRef),
}

struct NodeRemoval<'a> {
    empty_pages: Vec<&'a EntryKey>,
    split: Vec<(NodeCellRef, EntryKey)>,
}

impl<'a> NodeRemoval<'a> {
    fn new() -> Self {
        Self {
            empty_pages: vec![],
            split: vec![],
        }
    }
}

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 mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn merge_innode_remnant<'a, KS, PS>(
    current_node: &mut NodeWriteGuard<KS, PS>,
    prev_key: &'a EntryKey,
    removal: &mut NodeRemoval<'a>,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let curr_right_ref = current_node.right_ref_mut_no_empty().unwrap().clone();
    let curr_innode = current_node.innode_mut();
    let curr_right_bound = &curr_innode.right_bound;
    let curr_last_child = mem::replace(&mut curr_innode.ptrs.as_slice()[0], NodeCellRef::default());
    debug_assert_eq!(curr_innode.len, 0);
    removal.empty_pages.push(prev_key);
    if curr_last_child.is_default() {
        return;
    }
    let mut next_node = write_targeted(write_node::<KS, PS>(&curr_right_ref), curr_right_bound);
    let pos = next_node.search(curr_right_bound);
    let next_innode = next_node.innode_mut();
    next_innode.debug_check_integrity();
    if next_innode.len == KS::slice_len() {
        removal.split.push(next_innode.split_insert(
            curr_right_bound.clone(),
            curr_last_child,
            pos,
            false,
        ))
    } else {
        next_innode.insert_in_place(curr_right_bound.clone(), curr_last_child, pos, false);
    }
}

fn apply_removal<'a, KS, PS>(
    cursor_guard: &mut NodeWriteGuard<KS, PS>,
    poses: &mut BTreeSet<usize>,
    node_removal: &mut NodeRemoval<'a>,
    prev_key: &Option<&'a EntryKey>,
    remove_children_right_nodes: bool,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    if poses.is_empty() {
        return;
    }
    debug!(
        "Applying removal, guard key {:?}, poses {:?}",
        cursor_guard.first_key(),
        poses
    );
    {
        let innode = cursor_guard.innode_mut();
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            if remove_children_right_nodes {
                // A node will be reclaimed if and only if it have zero references
                // If there is any sequential empty nodes, they have to be unlinked each other
                // The upper level will do the rest to unlink it from the tree
                innode.ptrs.as_slice()[..innode.len + 1]
                    .iter()
                    .enumerate()
                    .filter(|(i, _)| poses.contains(i))
                    .for_each(|(_, r)| {
                        write_node::<KS, PS>(r).right_ref_mut_no_empty();
                    });
            }
            let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, p)| p)
                .collect();

            let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .take(if ptrs.len() == 0 { 0 } else { ptrs.len() - 1 })
                .map(|(_, k)| k)
                .collect();
            innode.len = keys.len();
            debug!("Prune filtered page have keys {:?}", &keys);
            for (i, key) in keys.into_iter().enumerate() {
                debug_assert!(key > &mut smallvec!(0));
                mem::swap(&mut new_keys.as_slice()[i], key);
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                debug_assert!(!ptr.is_default());
                mem::swap(&mut new_ptrs.as_slice()[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;
        innode.debug_check_integrity();

        debug!(
            "Pruned page have keys {:?}",
            &innode.keys.as_slice_immute()[..innode.len]
        );
    }

    if cursor_guard.is_empty() {
        if let &Some(k) = prev_key {
            debug_assert!(!cursor_guard.is_ext());
            merge_innode_remnant(cursor_guard, k, node_removal);
        }
        debug!("Pruned page is empty: {:?}", prev_key);
        cursor_guard.make_empty_node();
    }
    poses.clear();
}

fn prune_selected<'a, KS, PS>(
    node: &NodeCellRef,
    mut removal: Box<NodeRemoval<'a>>,
    level: usize,
) -> Box<NodeRemoval<'a>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, removal.empty_pages.first().unwrap());
    let pruning = match first_search {
        MutSearchResult::Internal(sub_node) => {
            if read_unchecked::<KS, PS>(&NodeData::<KS, PS>::get_non_empty_node(&sub_node)).is_ext()
            {
                PruningSearch::DeepestInnode
            } else {
                PruningSearch::Innode(sub_node)
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut deepest = false;
    match pruning {
        PruningSearch::DeepestInnode => {
            debug!(
                "Removing in deepest nodes keys {:?}, level {}",
                &removal.empty_pages, level
            );
            deepest = true;
        }
        PruningSearch::Innode(sub_node) => {
            removal = prune_selected::<KS, PS>(&sub_node, removal, level + 1);
        }
    }
    // empty page references that will dealt with by upper level
    let mut upper_removal = NodeRemoval::new();
    if !removal.empty_pages.is_empty() {
        debug!(
            "Pruning page containing keys {:?}, level {}",
            &removal.empty_pages, level
        );
        // start delete
        let mut cursor_guard = write_non_empty(write_node::<KS, PS>(node));
        let mut guard_removing_poses = BTreeSet::new();
        let mut prev_key = None;
        debug!(
            "Prune deepest node starting key: {:?}, level {}",
            cursor_guard.first_key(),
            level
        );
        // iterate over all removal empty nodes
        // in the beginning cursor_guard contains the first key
        // if the empty node go out of the right bound of current cursor guard,
        // it will `apply_removal` and switch to the next right page
        for (_i, &key_to_del) in removal.empty_pages.iter().enumerate() {
            if key_to_del >= cursor_guard.right_bound() {
                debug!(
                    "Applying removal for overflow current page ({}/{}) key: {:?} >= bound: {:?}. guard keys: {:?}, level {}",
                    guard_removing_poses.len(),
                    cursor_guard.len() + 1,
                    key_to_del,
                    cursor_guard.right_bound(),
                    &cursor_guard.keys()[..cursor_guard.len()],
                    level
                );
                apply_removal(
                    &mut cursor_guard,
                    &mut guard_removing_poses,
                    &mut upper_removal,
                    &prev_key,
                    !deepest,
                );
                cursor_guard = write_targeted(cursor_guard, key_to_del);
                debug_assert!(!cursor_guard.is_empty_node());
                debug!(
                    "Applied removal for overflow current page ({}), level {}",
                    cursor_guard.len(),
                    level
                );
            }
            let pos = cursor_guard.search(key_to_del);
            debug!(
                "Key to delete have position {}, key: {:?}, level {}",
                pos, key_to_del, level
            );
            guard_removing_poses.insert(pos);
            prev_key = Some(key_to_del)
        }
        if !guard_removing_poses.is_empty() {
            debug!(
                "Applying removal for last keys {:?}, level {}",
                &guard_removing_poses, level
            );
            apply_removal(
                &mut cursor_guard,
                &mut guard_removing_poses,
                &mut upper_removal,
                &prev_key,
                !deepest,
            );
        }
    }
    if !removal.split.is_empty() {
        let mut cursor_guard = write_node::<KS, PS>(node);
        for (node, pivot) in &removal.split {
            cursor_guard = write_targeted(cursor_guard, &pivot);
            let pos = cursor_guard.search(&pivot);
            let innode = cursor_guard.innode_mut();
            let pivot = pivot.clone();
            let node = node.clone();
            if innode.len >= KS::slice_len() {
                upper_removal
                    .split
                    .push(innode.split_insert(pivot, node, pos, true));
            } else {
                innode.insert_in_place(pivot, node, pos, true)
            }
        }
    }
    debug!(
        "Have empty nodes {:?}, level {:?}",
        &upper_removal.empty_pages, level
    );
    box upper_removal
}

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;
    debug!("Merge selected {} pages", left_most_leaf_guards.len());

    // 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);
        }
    }

    // cleanup upper level references
    {
        let page_keys = left_most_leaf_guards
            .iter()
            .filter(|g| !g.is_empty())
            .map(|g| g.first_key())
            .collect_vec();
        let mut removal = NodeRemoval::new();
        removal.empty_pages = page_keys;
        prune_selected::<KS, PS>(&src_tree.get_root(), box removal, 0);
    }

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

        let left_left_most = {
            let left_most = left_most_leaf_guards.first().unwrap();
            if left_most.is_none() {
                left_most.node_ref().clone()
            } else {
                left_most.left_ref().unwrap().clone()
            }
        };

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

        let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
        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 = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

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

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 fd developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::process::Command;
use std::sync::{Arc, Mutex};
use std::io;

/// Executes a command.
#[cfg(not(unix))]
pub fn execute_command(mut cmd: Command, out_perm: Arc<Mutex<()>>) {
    use std::process::Stdio;
    use std::io::Write;

    // Spawn the supplied command.
    let output = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).output();

    // Then wait for the command to exit, if it was spawned.
    match output {
        Ok(output) => {
            // While this lock is active, this thread will be the only thread allowed
            // to write it's outputs.
            let _lock = out_perm.lock().unwrap();

            let stdout = io::stdout();
            let stderr = io::stderr();

            let _ = stdout.lock().write_all(&output.stdout);
            let _ = stderr.lock().write_all(&output.stderr);
        }
        Err(why) => eprintln!("fd: exec error: {}", why),
    }
}

/// Executes a command.
#[cfg(all(unix))]
pub fn execute_command(mut cmd: Command, out_perm: Arc<Mutex<()>>) {
    use libc::{close, dup2, pipe, STDERR_FILENO, STDOUT_FILENO};
    use std::fs::File;
    use std::os::unix::process::CommandExt;
    use std::os::unix::io::FromRawFd;

    // Initial a pair of pipes that will be used to
    // pipe the std{out,err} of the spawned process.
    let mut stdout_fds = [0; 2];
    let mut stderr_fds = [0; 2];

    unsafe {
        pipe(stdout_fds.as_mut_ptr());
        pipe(stderr_fds.as_mut_ptr());
    }

    // Configure the pipes accordingly in the child.
    let child = cmd.before_exec(move || unsafe {
        // Redirect the child's std{out,err} to the write ends of our pipe.
        dup2(stdout_fds[1], STDOUT_FILENO);
        dup2(stderr_fds[1], STDERR_FILENO);

        // Close all the fds we created here, so EOF will be sent when the program exits.
        close(stdout_fds[0]);
        close(stdout_fds[1]);
        close(stderr_fds[0]);
        close(stderr_fds[1]);
        Ok(())
    }).spawn();

    // Open the read end of the pipes as `File`s.
    let (mut pout, mut perr) = unsafe {
        // Close the write ends of the pipes in the parent
        close(stdout_fds[1]);
        close(stderr_fds[1]);
        (
            // But create files from the read ends.
            File::from_raw_fd(stdout_fds[0]),
            File::from_raw_fd(stderr_fds[0]),
        )
    };

    match child {
        Ok(mut child) => {
            let _ = child.wait();

            // Create a lock to ensure that this thread has exclusive access to writing.
            let _lock = out_perm.lock().unwrap();

            // And then write the outputs of the process until EOF is sent to each file.
            let stdout = io::stdout();
            let stderr = io::stderr();
            let _ = io::copy(&mut pout, &mut stdout.lock());
            let _ = io::copy(&mut perr, &mut stderr.lock());
        }
        Err(why) => eprintln!("fd: exec error: {}", why),
    }
}
<commit_msg>Update error message if command could not be found<commit_after>// Copyright (c) 2017 fd developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::process::Command;
use std::sync::{Arc, Mutex};
use std::io;

/// Executes a command.
#[cfg(not(unix))]
pub fn execute_command(mut cmd: Command, out_perm: Arc<Mutex<()>>) {
    use std::process::Stdio;
    use std::io::Write;

    // Spawn the supplied command.
    let output = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).output();

    // Then wait for the command to exit, if it was spawned.
    match output {
        Ok(output) => {
            // While this lock is active, this thread will be the only thread allowed
            // to write it's outputs.
            let _lock = out_perm.lock().unwrap();

            let stdout = io::stdout();
            let stderr = io::stderr();

            let _ = stdout.lock().write_all(&output.stdout);
            let _ = stderr.lock().write_all(&output.stderr);
        }
        Err(why) => {
            if why.kind() == io::ErrorKind::NotFound {
                eprintln!("fd: execution error: command not found");
            } else {
                eprintln!("fd: execution error: {}", why);
            }
        }
    }
}

/// Executes a command.
#[cfg(all(unix))]
pub fn execute_command(mut cmd: Command, out_perm: Arc<Mutex<()>>) {
    use libc::{close, dup2, pipe, STDERR_FILENO, STDOUT_FILENO};
    use std::fs::File;
    use std::os::unix::process::CommandExt;
    use std::os::unix::io::FromRawFd;

    // Initial a pair of pipes that will be used to
    // pipe the std{out,err} of the spawned process.
    let mut stdout_fds = [0; 2];
    let mut stderr_fds = [0; 2];

    unsafe {
        pipe(stdout_fds.as_mut_ptr());
        pipe(stderr_fds.as_mut_ptr());
    }

    // Configure the pipes accordingly in the child.
    let child = cmd.before_exec(move || unsafe {
        // Redirect the child's std{out,err} to the write ends of our pipe.
        dup2(stdout_fds[1], STDOUT_FILENO);
        dup2(stderr_fds[1], STDERR_FILENO);

        // Close all the fds we created here, so EOF will be sent when the program exits.
        close(stdout_fds[0]);
        close(stdout_fds[1]);
        close(stderr_fds[0]);
        close(stderr_fds[1]);
        Ok(())
    }).spawn();

    // Open the read end of the pipes as `File`s.
    let (mut pout, mut perr) = unsafe {
        // Close the write ends of the pipes in the parent
        close(stdout_fds[1]);
        close(stderr_fds[1]);
        (
            // But create files from the read ends.
            File::from_raw_fd(stdout_fds[0]),
            File::from_raw_fd(stderr_fds[0]),
        )
    };

    match child {
        Ok(mut child) => {
            let _ = child.wait();

            // Create a lock to ensure that this thread has exclusive access to writing.
            let _lock = out_perm.lock().unwrap();

            // And then write the outputs of the process until EOF is sent to each file.
            let stdout = io::stdout();
            let stderr = io::stderr();
            let _ = io::copy(&mut pout, &mut stdout.lock());
            let _ = io::copy(&mut perr, &mut stderr.lock());
        }
        Err(why) => {
            if why.kind() == io::ErrorKind::NotFound {
                eprintln!("fd: execution error: command not found");
            } else {
                eprintln!("fd: execution error: {}", why);
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::path::Path;
use std::ffi::OsStr;

use iron::prelude::*;
use iron_sessionstorage::Value as SessionValue;
use iron_sessionstorage::traits::SessionRequestExt;
use rand::*;
use crypto::md5::Md5;
use crypto::digest::Digest;  // used for input_str, result_str
use chrono::{Local, NaiveDateTime};
use urlencoded::{UrlEncodedQuery, UrlEncodedBody};
use serde::Serialize;
use serde_json::{self, Value};
use hbs::handlebars::{Helper, Handlebars, RenderContext, RenderError};
use router::{Router, Params};
use pulldown_cmark::{Parser, html};
use toml::value::Value::String as Toml_String;

use common::http::SessionData;
use common::lazy_static::{RECORDS_COUNT_PER_PAGE, ADMINS};

pub fn parse_to_html(text: &str) -> String {

    let mut temp = String::new();
    let parser = Parser::new(text);

    html::push_html(&mut temp, parser);

    temp
}

pub fn gen_salt() -> String {

    thread_rng()
        .gen_ascii_chars()
        .take(32)
        .collect::<String>()
}

pub fn gen_md5(str: &str) -> String {

    let mut sh = Md5::new();

    sh.input_str(str);
    sh.result_str().to_string()
}

pub fn get_file_ext(filename: &str) -> Option<&str>{

    Path::new(filename)
        .extension()
        .and_then(OsStr::to_str)
}

pub fn gen_gravatar_url(email: &str) -> String {

    "http://www.gravatar.com/avatar/".to_string() + &*gen_md5(email) + "?s=150"
}

pub fn gen_datetime() -> NaiveDateTime {

    Local::now().naive_utc()
}

pub fn check_and_get_string(value: &Value) -> String {

    if value.is_null() {

        "".to_string()
    } else {

        value.as_str().unwrap().to_string()
    }
}

pub fn is_login(req: &mut Request) -> bool {

    let session_wrapper = req.session().get::<SessionData>().unwrap();

    if session_wrapper.is_some() {
        true
    } else {
        false
    }
}

pub fn get_session_obj(req: &mut Request) -> Value {

    let session_wrapper = req.session().get::<SessionData>().unwrap();

    json_parse(&*session_wrapper.unwrap().into_raw())
}

pub fn get_router_params(req: &mut Request) -> Params {

    req.extensions.get::<Router>().unwrap().clone()
}

pub fn get_request_body(req: &mut Request) -> HashMap<String, Vec<String>> {

    req.get::<UrlEncodedBody>().unwrap()
}

pub fn get_request_query(req: &mut Request) -> HashMap<String, Vec<String>> {

    req.get::<UrlEncodedQuery>().unwrap()
}

pub fn has_request_query(req: &mut Request) -> bool {

    if req.get::<UrlEncodedQuery>().is_err() {

        false
    } else {

        true
    }
}

pub fn get_query_page(req: &mut Request) -> u32 {

    let has_query_params = has_request_query(req);
    let page: u32;

    if has_query_params {

        let query = get_request_query(req);

        if query.get("page").is_none()  {
            page = 1;
        } else {

            let page_wrapper = query.get("page").unwrap()[0].parse::<u32>();

            if page_wrapper.is_err() {

                page = 1;
            } else {

                page = page_wrapper.unwrap();
            }
        }
    } else {
        page = 1;
    }

    page
}

pub fn json_stringify<T: Serialize>(data: &T) -> String {

    serde_json::to_string(data).unwrap()
}

pub fn json_parse(data: &str) -> Value {

    serde_json::from_str(data).unwrap()
}

pub fn mount_template_var(helper: &Helper, _: &Handlebars, context: &mut RenderContext) -> Result<(), RenderError> {

    let param_key = helper.param(0);
    let param_value = helper.param(1);

    if param_key.is_none() || param_value.is_none() {

        return Ok(());
    }

    let key = param_key.unwrap().value().as_str().unwrap().to_string();
    let value = param_value.unwrap().value();

    let mut view_data = context.context_mut().data_mut().as_object_mut().unwrap();

    view_data.insert(key, json!(value));

    Ok(())
}

pub fn build_pagination(cur_page: u32, total: u32, base_url: &str) -> Value {

    let mut delta = 1;

    if total % RECORDS_COUNT_PER_PAGE == 0 {
        delta = 0;
    }

    let page_count = total / RECORDS_COUNT_PER_PAGE + delta;
    let mut is_show_prev_ellipsis = true;
    let mut is_show_next_ellipsis = true;
    let mut is_first_page_disabled = false;
    let mut is_last_page_disabled = false;
    let mut page_list;

    if page_count < 6 {

        is_show_prev_ellipsis = false;
        is_show_next_ellipsis = false;
    } else {

        if cur_page < 4  {

            is_show_prev_ellipsis = false;
        }

        if cur_page > page_count - 3 {

            is_show_next_ellipsis = false;
        }
    }

    if cur_page == 1 {

        is_first_page_disabled = true;
    }

    if cur_page == page_count {

        is_last_page_disabled = true;
    }

    page_list = vec![];

    if page_count < 6 {  // 总页数小于等于5时

        for i in 1..(page_count + 1) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    } else if cur_page < 4 {  // 总页数大于5，当前页码小于等于3时，隐藏左侧ellipsis

        for i in 1..6 {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    } else if cur_page > page_count - 3 {  // 总页数大于5，当前页码距离总页数小于等于3时，隐藏右侧ellipsis

        for i in (page_count - 4)..(page_count + 1) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }

    } else {  // 当前页码的左右两侧各放置两个页码

        for i in (cur_page - 2)..(cur_page + 3) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    }

    json!({
        "base_url": base_url.to_owned(),
        "is_show_prev_ellipsis": is_show_prev_ellipsis,
        "is_show_next_ellipsis": is_show_next_ellipsis,
        "page_list": page_list,
        "is_first_page_disabled": is_first_page_disabled,
        "is_last_page_disabled": is_last_page_disabled,
        "first_page": 1,
        "last_page": page_count
    })
}

pub fn is_admin(username: &str) -> bool {

    if ADMINS.contains(&Toml_String(username.to_string())) {
        true
    } else {
        false
    }
}<commit_msg>增加test<commit_after>use std::collections::HashMap;
use std::path::Path;
use std::ffi::OsStr;

use iron::prelude::*;
use iron_sessionstorage::Value as SessionValue;
use iron_sessionstorage::traits::SessionRequestExt;
use rand::*;
use crypto::md5::Md5;
use crypto::digest::Digest;  // used for input_str, result_str
use chrono::{Local, NaiveDateTime};
use urlencoded::{UrlEncodedQuery, UrlEncodedBody};
use serde::Serialize;
use serde_json::{self, Value};
use hbs::handlebars::{Helper, Handlebars, RenderContext, RenderError};
use router::{Router, Params};
use pulldown_cmark::{Parser, html};
use toml::value::Value::String as Toml_String;

use common::http::SessionData;
use common::lazy_static::{RECORDS_COUNT_PER_PAGE, ADMINS};

pub fn parse_to_html(text: &str) -> String {

    let mut temp = String::new();
    let parser = Parser::new(text);

    html::push_html(&mut temp, parser);

    temp
}

pub fn gen_salt() -> String {

    thread_rng()
        .gen_ascii_chars()
        .take(32)
        .collect::<String>()
}

pub fn gen_md5(str: &str) -> String {

    let mut sh = Md5::new();

    sh.input_str(str);
    sh.result_str().to_string()
}

pub fn get_file_ext(filename: &str) -> Option<&str>{

    Path::new(filename)
        .extension()
        .and_then(OsStr::to_str)
}

pub fn gen_gravatar_url(email: &str) -> String {

    "http://www.gravatar.com/avatar/".to_string() + &*gen_md5(email) + "?s=150"
}

pub fn gen_datetime() -> NaiveDateTime {

    Local::now().naive_utc()
}

pub fn check_and_get_string(value: &Value) -> String {

    if value.is_null() {

        "".to_string()
    } else {

        value.as_str().unwrap().to_string()
    }
}

pub fn is_login(req: &mut Request) -> bool {

    let session_wrapper = req.session().get::<SessionData>().unwrap();

    if session_wrapper.is_some() {
        true
    } else {
        false
    }
}

pub fn get_session_obj(req: &mut Request) -> Value {

    let session_wrapper = req.session().get::<SessionData>().unwrap();

    json_parse(&*session_wrapper.unwrap().into_raw())
}

pub fn get_router_params(req: &mut Request) -> Params {

    req.extensions.get::<Router>().unwrap().clone()
}

pub fn get_request_body(req: &mut Request) -> HashMap<String, Vec<String>> {

    req.get::<UrlEncodedBody>().unwrap()
}

pub fn get_request_query(req: &mut Request) -> HashMap<String, Vec<String>> {

    req.get::<UrlEncodedQuery>().unwrap()
}

pub fn has_request_query(req: &mut Request) -> bool {

    if req.get::<UrlEncodedQuery>().is_err() {

        false
    } else {

        true
    }
}

pub fn get_query_page(req: &mut Request) -> u32 {

    let has_query_params = has_request_query(req);
    let page: u32;

    if has_query_params {

        let query = get_request_query(req);

        if query.get("page").is_none()  {
            page = 1;
        } else {

            let page_wrapper = query.get("page").unwrap()[0].parse::<u32>();

            if page_wrapper.is_err() {

                page = 1;
            } else {

                page = page_wrapper.unwrap();
            }
        }
    } else {
        page = 1;
    }

    page
}

pub fn json_stringify<T: Serialize>(data: &T) -> String {

    serde_json::to_string(data).unwrap()
}

pub fn json_parse(data: &str) -> Value {

    serde_json::from_str(data).unwrap()
}

pub fn mount_template_var(helper: &Helper, _: &Handlebars, context: &mut RenderContext) -> Result<(), RenderError> {

    let param_key = helper.param(0);
    let param_value = helper.param(1);

    if param_key.is_none() || param_value.is_none() {

        return Ok(());
    }

    let key = param_key.unwrap().value().as_str().unwrap().to_string();
    let value = param_value.unwrap().value();

    let mut view_data = context.context_mut().data_mut().as_object_mut().unwrap();

    view_data.insert(key, json!(value));

    Ok(())
}

pub fn build_pagination(cur_page: u32, total: u32, base_url: &str) -> Value {

    let mut delta = 1;

    if total % RECORDS_COUNT_PER_PAGE == 0 {
        delta = 0;
    }

    let page_count = total / RECORDS_COUNT_PER_PAGE + delta;
    let mut is_show_prev_ellipsis = true;
    let mut is_show_next_ellipsis = true;
    let mut is_first_page_disabled = false;
    let mut is_last_page_disabled = false;
    let mut page_list;

    if page_count < 6 {

        is_show_prev_ellipsis = false;
        is_show_next_ellipsis = false;
    } else {

        if cur_page < 4  {

            is_show_prev_ellipsis = false;
        }

        if cur_page > page_count - 3 {

            is_show_next_ellipsis = false;
        }
    }

    if cur_page == 1 {

        is_first_page_disabled = true;
    }

    if cur_page == page_count {

        is_last_page_disabled = true;
    }

    page_list = vec![];

    if page_count < 6 {  // 总页数小于等于5时

        for i in 1..(page_count + 1) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    } else if cur_page < 4 {  // 总页数大于5，当前页码小于等于3时，隐藏左侧ellipsis

        for i in 1..6 {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    } else if cur_page > page_count - 3 {  // 总页数大于5，当前页码距离总页数小于等于3时，隐藏右侧ellipsis

        for i in (page_count - 4)..(page_count + 1) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }

    } else {  // 当前页码的左右两侧各放置两个页码

        for i in (cur_page - 2)..(cur_page + 3) {

            if i == cur_page {

                page_list.push(json!({
                    "page": i,
                    "is_active": true
                }));
            } else {

                page_list.push(json!({
                    "page": i
                }));
            }
        }
    }

    json!({
        "base_url": base_url.to_owned(),
        "is_show_prev_ellipsis": is_show_prev_ellipsis,
        "is_show_next_ellipsis": is_show_next_ellipsis,
        "page_list": page_list,
        "is_first_page_disabled": is_first_page_disabled,
        "is_last_page_disabled": is_last_page_disabled,
        "first_page": 1,
        "last_page": page_count
    })
}

pub fn is_admin(username: &str) -> bool {

    if ADMINS.contains(&Toml_String(username.to_string())) {
        true
    } else {
        false
    }
}

#[test]
fn test_gen_salt() {

    assert_ne!(gen_salt(), "runner".to_owned());
}

#[test]
fn test_gen_md5() {

    assert_ne!(gen_md5("runner"), "runner".to_owned());
}

#[test]
fn test_get_file_ext() {

    assert_eq!(get_file_ext("abc.txt").unwrap(), "txt");
}<|endoftext|>


<commit_before>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})?").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    let mut date: Option<String> = None;

    if let Some(val) = RE.find(url) {
        date = Some(val.as_str().to_string());
    }

    date
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let decoded_ldjson = Json::from_str(ldjson.as_str()).unwrap();

    if let Some(date_published) = decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    let mut date: Option<String> = None;

    for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            date = Some(dt.to_string())
        } else if let Some("timestamp") = time.attr("class") {
            date = Some(time.text())
        }
    }

    if date.is_none() {
        for tag in html.find(Name("meta")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    date = Some(v.to_string())
                } else {
                    date = Some(tag.text())
                }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use super::extract_article_published_date;
    use chrono::NaiveDate;
    use reqwest;
    use std::string::String;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-11-29T00:44:59Z".to_string()), extract_from_html_tag(&document));
    }

    #[test]
    fn extracting_article_published_date() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();

        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, None).unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, Some(body)).unwrap());

        assert!((extract_article_published_date("", None)).is_err());
    }
}
<commit_msg>Add new date formats<commit_after>use regex::Regex;
use chrono::NaiveDate;
use reqwest;
use std::io::Read;
use select::document::Document;
use select::predicate::{Name, Attr};
use rustc_serialize::json::Json;
use errors::*;

// Some formats borrowed from https://github.com/amir/article-date-extractor
static FMTS: &'static [&str] = &["%A, %B %e, %Y", "%Y-%m-%dT%H:%M:%S%:z", "/%Y/%m/%d/", "/%Y/%d/%m/", "%Y-%m-%d", "%B %e, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%SZ", "%B %k, %Y, %H:%M %p", "%Y-%m-%d %H:%M:%S.000000"];

// Use lazy_static to ensure we only compile the regex once
lazy_static! {
    // Regex by Newspaper3k  - https://github.com/codelucas/newspaper/blob/master/newspaper/urls.py
    static ref RE: Regex = Regex::new(r"([\./\-_]{0,1}(19|20)\d{2})[\./\-_]{0,1}(([0-3]{0,1}[0-9][\./\-_])|(\w{3,5}[\./\-_]))([0-3]{0,1}[0-9][\./\-]{0,1})?").unwrap();
}

// Parse the date, trying out each format
fn parse_date(input: &str) -> Result<NaiveDate> {
    let mut result: Result<NaiveDate> = Err("None of the formats matched the date".into());

    'outer: for fmt in FMTS {
        if let Ok(v) = NaiveDate::parse_from_str(input, fmt) {
            { result = Ok(v);
              break 'outer; }
        }
    }

    result
}

// Extract date from a URL
fn extract_from_url(url: &str) -> Option<String> {
    let mut date: Option<String> = None;

    if let Some(val) = RE.find(url) {
        date = Some(val.as_str().to_string());
    }

    date
}

// Extract date from JSON-LD
fn extract_from_ldjson<'a>(html: &'a Document) -> Option<String> {
    let mut json_date: Option<String> = None;
    let mut ldjson = String::new();
    if let Some(ldj) = html.find(Attr("type", "application/ld+json")).next() {
        ldjson = ldj.text();
    }

    let decoded_ldjson = Json::from_str(ldjson.as_str()).unwrap();

    if let Some(date_published) = decoded_ldjson.search("datePublished") {
        if let Some(date) = date_published.as_string() {
            json_date = Some(date.to_string())
        }
    } else if let Some(date_created) = decoded_ldjson.search("dateCreated") {
        if let Some(date) = date_created.as_string() {
            json_date = Some(date.to_string())
        }
    }

    json_date
}

// Extract date from meta tags
fn extract_from_meta<'a>(html: &'a Document) -> Option<String> {
    let mut meta_date: Option<String> = None;

    'outer: for meta in html.find(Name("meta")) {
        let meta_name     = meta.attr("name").unwrap_or("").to_lowercase();
        let item_prop     = meta.attr("itemprop").unwrap_or("").to_lowercase();
        let http_equiv    = meta.attr("http-equiv").unwrap_or("").to_lowercase();
        let meta_property = meta.attr("property").unwrap_or("").to_lowercase();

        match meta_name.as_ref() {
            "pubdate"               | "publishdate"                  | "timestamp"       |
            "dc.date.issued"        | "date"                         | "sailthru.date"   |
            "article.published"     | "published-date"               | "article.created" |
            "article_date_original" | "cxenseparse:recs:publishtime" | "date_published"  => { if let Some(ct) = meta.attr("content") {
                                                                                meta_date = Some(ct.trim().to_string())
                                                                                }
                                                                                break 'outer; },
            _ => {},
        }

        match item_prop.as_ref() {
            "datepublished" | "datecreated" => { if let Some(ct) = meta.attr("content") {
                                                   meta_date = Some(ct.trim().to_string())
                                                 }
                                                 break 'outer; },
            _ => {},
        }

        match http_equiv.as_ref() {
            "date" =>  { if let Some(ct) = meta.attr("content") {
                           meta_date = Some(ct.trim().to_string())
                         }
                         break 'outer; },
            _ => {},
        }

        match meta_property.as_ref() {
            "article:published_time" | "bt:pubdate" => { if let Some(ct) = meta.attr("content") {
                                                           meta_date = Some(ct.trim().to_string())
                                                         }
                                                         break 'outer; },
            "og:image"                              => { if let Some(url) = meta.attr("content") {
                                                           meta_date = extract_from_url(url.trim())
                                                         }
                                                         break 'outer; },

            _ => {},
        }

    }

    meta_date
}

// Extract from html tags
fn extract_from_html_tag<'a>(html: &'a Document) -> Option<String> {
    let mut date: Option<String> = None;

    for time in html.find(Name("time")) {
        if let Some(dt) = time.attr("datetime") {
            date = Some(dt.to_string())
        } else if let Some("timestamp") = time.attr("class") {
            date = Some(time.text())
        }
    }

    if date.is_none() {
        for tag in html.find(Name("meta")) {
            if let Some("datePublished") = tag.attr("itemprop") {
                if let Some(v) = tag.attr("content") {
                    date = Some(v.to_string())
                } else {
                    date = Some(tag.text())
                }
            }
        }
    }

    date
}

// Try to extract the date by using each function one by one
pub fn extract_article_published_date(link: &str, html: Option<String>) -> Result<NaiveDate> {
    let mut body: String = String::new();
    let mut parsed_body: Option<Document> = None;

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    }

    if html.is_none() {
        if let Ok(mut response) = reqwest::get(link) {
            response.read_to_string(&mut body).unwrap();
            let doc = Document::from(body.as_str());
            parsed_body = Some(doc);
        } else {
            return Err("Couldn't open the link".into())
        }
    } else {
        parsed_body = Some(Document::from(html.unwrap().as_str()))
    }

    if let Some(v) = extract_from_url(link) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_ldjson(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_meta(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else if let Some(v) = extract_from_html_tag(parsed_body.as_ref().unwrap()) {
        return parse_date(v.as_str())
    } else {
        return Err("Couldn't find the date to parse".into())
    }
}

// Unit tests
#[cfg(test)]
mod test {
    use super::extract_from_url;
    use super::parse_date;
    use super::extract_from_meta;
    use super::extract_from_ldjson;
    use super::extract_from_html_tag;
    use super::extract_article_published_date;
    use chrono::NaiveDate;
    use reqwest;
    use std::string::String;
    use std::io::Read;
    use select::document::Document;

    #[test]
    fn parsing_date() {
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/11/30/").unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,30), parse_date("/2015/30/11/").unwrap());

        assert!(parse_date("bad_format").is_err());
    }

    #[test]
    fn extracting_from_url() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        assert_eq!(Some("/2015/11/28/".to_string()), extract_from_url(link));

        let link = "";
        assert_eq!(None, extract_from_url(link));
    }

    #[test]
    fn extracting_from_meta() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some(("2015-11-30 23:50:48".to_string())), extract_from_meta(&document));
    }

    #[test]
    fn extracting_from_ldjson() {
        let mut response = reqwest::get("https://techcrunch.com/2015/11/30/atlassian-share-price/").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-12-01T07:50:48Z".to_string()), extract_from_ldjson(&document));
    }

    #[test]
    fn extracting_from_html_tag() {
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();
        let document = Document::from(body.as_str());

        assert_eq!(Some("2015-11-29T00:44:59Z".to_string()), extract_from_html_tag(&document));
    }

    #[test]
    fn extracting_article_published_date() {
        let link = "http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html";
        let mut response = reqwest::get("http://edition.cnn.com/2015/11/28/opinions/sutter-cop21-paris-preview-two-degrees/index.html").unwrap();
        let mut body = String::new();
        response.read_to_string(&mut body).unwrap();

        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, None).unwrap());
        assert_eq!(NaiveDate::from_ymd(2015,11,28), extract_article_published_date(&link, Some(body)).unwrap());

        assert!((extract_article_published_date("", None)).is_err());
    }
}
<|endoftext|>


<commit_before>//! Simple parsing of ASN.1 DER

use timmy::tup::*;

use num::bigint::{BigInt, Sign};

#[derive(Debug)]
pub enum ASN1Error {
    MissingData(String),
}

pub type ASN1Result<T> = Result<T, ASN1Error>;

#[derive(Debug)]
pub enum ASN1Type {
    Unknown,
    Cont(u8, Box<ASN1Type>),
    Null,
    Boolean(bool),
    Sequence(Vec<ASN1Type>),
    Set(Vec<ASN1Type>),
    Object(TupT<u32>),
    Integer(BigInt),
    BitString(Vec<u8>),
    OctetString(Vec<u8>),
    PrintableString(String),
    UTF8String(String),
    UTCTime(String),
}

pub fn asn1_to_raw_string(obj: &ASN1Type) -> Option<String> {
    match obj {
        &ASN1Type::PrintableString(ref string) => Some(string.clone()),
        &ASN1Type::UTF8String(ref string) => Some(string.clone()),
        _ => None,
    }
}

pub fn asn1_to_raw_integer(obj: &ASN1Type) -> Option<BigInt> {
    match obj {
        &ASN1Type::Integer(ref bigint) => Some(bigint.clone()),
        _ => None,
    }
}

pub struct DerParser {
    buf: Vec<u8>,
    idx: usize,
}

impl DerParser {
    pub fn new(der: &Vec<u8>) -> DerParser {
        DerParser {
            buf: der.clone(),
            idx: 0,
        }
    }

    fn check_buffer(&mut self) -> ASN1Result<()> {
        if self.idx >= self.buf.len() {
            return Err(ASN1Error::MissingData("length".to_string()));
        }
        Ok(())
    }

    fn get_length(&mut self) -> ASN1Result<usize> {
        try!(self.check_buffer());

        let len_tmp = self.buf[self.idx];
        if len_tmp <= 127 {
            self.idx += 1;
            return Ok(len_tmp as usize);
        }

        let len_octets = len_tmp & 127;
        assert!(len_octets <= 3);
        let mut length = 0 as usize;
        self.idx += 1;
        for _ in 0..len_octets {
            length <<= 8;
            try!(self.check_buffer());
            length |= self.buf[self.idx] as usize;
            self.idx += 1;
        }
        Ok(length)
    }

    fn parse_header(&mut self) -> ASN1Result<(u8, bool, u8, usize)> {
        try!(self.check_buffer());

        let c = self.buf[self.idx];
        let tag = c & 31;
        let p_c = (c >> 5) & 1;
        let class_ = (c >> 6) & 3;
        assert!(tag != 31);
        self.idx += 1;
        let length = try!(self.get_length());

        Ok((class_ as u8, p_c == 1, tag as u8, length as usize))
    }

    pub fn parse_entry(&mut self) -> ASN1Result<ASN1Type> {
        try!(self.check_buffer());

        let (class_bits, constructed, tag, length) = try!(self.parse_header());

        assert!(class_bits == 0 || class_bits == 2);

        match class_bits {
            0 => {
            }
            2 => {
                let pos = self.idx;
                self.idx += length;

                assert!(pos + length <= self.buf.len());

                let sub_buf = self.buf[pos..pos+length].iter().cloned().collect();
                let mut sub_parser = DerParser::new(&sub_buf);
                let sub_entry = try!(sub_parser.parse_entry());
                return Ok(ASN1Type::Cont(tag, Box::new(sub_entry)));
            }
            _ => {
                panic!("Parsing this not implemented.");
            }
        }

        let pos = self.idx;
        self.idx += length;

        assert!(pos + length <= self.buf.len());

        let raw = self.buf[pos..pos+length].iter().cloned().collect();

        let entry = match tag {
            5 => ASN1Type::Null,
            1 => {
                let buf: Vec<u8> = raw;
                ASN1Type::Boolean(buf[0] == 0xff)
            }
            16 => {
                let mut sub_parser = DerParser::new(&raw);
                let mut sequence = Vec::<ASN1Type>::new();
                loop {
                    let sub_ent = sub_parser.parse_entry();
                    match sub_ent {
                        Err(ASN1Error::MissingData(_)) => { break },
                        Ok(ent) => { sequence.push(ent) }
                    }
                }
                ASN1Type::Sequence(sequence)
            },
            17 => {
                let mut sub_parser = DerParser::new(&raw);
                let mut sequence = Vec::<ASN1Type>::new();
                loop {
                    let sub_ent = sub_parser.parse_entry();
                    match sub_ent {
                        Err(ASN1Error::MissingData(_)) => { break },
                        Ok(ent) => { sequence.push(ent) }
                    }
                }
                ASN1Type::Set(sequence)
            },
            2 => ASN1Type::Integer(BigInt::from_bytes_be(Sign::Plus, &raw)),
            6 => {
                let mut oi = Vec::<u32>::new();
                let obj_bytes = raw;
                oi.push((obj_bytes[0] / 40) as u32);
                oi.push((obj_bytes[0] % 40) as u32);
                let mut i = 1;
                while i < obj_bytes.len() {
                    if obj_bytes[i] <= 127 {
                        oi.push(obj_bytes[i] as u32);
                        i += 1;
                    } else {
                        let mut next_sub_id: u32 = obj_bytes[i] as u32;
                        let mut sub_id: u32 = 0;
                        i += 1;
                        while next_sub_id > 127 {
                            sub_id <<= 7;
                            sub_id |= next_sub_id & 127;
                            next_sub_id = obj_bytes[i] as u32;
                            i += 1;
                        }
                        oi.push((sub_id << 7) | next_sub_id);
                    }
                }
                ASN1Type::Object(vec_to_tup(&oi).expect("..."))
            },
            3 => ASN1Type::BitString(raw),
            4 => ASN1Type::OctetString(raw),
            19 => ASN1Type::PrintableString(String::from_utf8(raw).unwrap()),
            23 => ASN1Type::UTCTime(String::from_utf8(raw).unwrap()),
            12 => ASN1Type::UTF8String(String::from_utf8(raw).unwrap()),
            _ => {
                println!("Got Unknown ASN1 Type {} {} {} {}", class_bits, constructed, tag, length);
                ASN1Type::Unknown
            }
        };

        Ok(entry)
    }
}
<commit_msg>Hide some ASN1 constants.<commit_after>//! Simple parsing of ASN.1 DER

use timmy::tup::*;

use num::bigint::{BigInt, Sign};

#[derive(Debug)]
pub enum ASN1Error {
    MissingData(String),
    DecodeError(String),
}

pub type ASN1Result<T> = Result<T, ASN1Error>;

#[derive(Debug)]
pub enum ASN1Type {
    Unknown,
    Cont(u8, Box<ASN1Type>),
    Null,
    Boolean(bool),
    Sequence(Vec<ASN1Type>),
    Set(Vec<ASN1Type>),
    Object(TupT<u32>),
    Integer(BigInt),
    BitString(Vec<u8>),
    OctetString(Vec<u8>),
    PrintableString(String),
    UTF8String(String),
    UTCTime(String),
}

enum TypeId {
    Null = 5,
    Boolean = 1,
    Sequence = 16,
    Set = 17,
    Integer = 2,
    Object = 6,
    BitString = 3,
    OctetString = 4,
    PrintableString = 19,
    UTCTime = 23,
    UTF8String = 12,
}

pub fn asn1_to_raw_string(obj: &ASN1Type) -> Option<String> {
    match obj {
        &ASN1Type::PrintableString(ref string) => Some(string.clone()),
        &ASN1Type::UTF8String(ref string) => Some(string.clone()),
        _ => None,
    }
}

pub fn asn1_to_raw_integer(obj: &ASN1Type) -> Option<BigInt> {
    match obj {
        &ASN1Type::Integer(ref bigint) => Some(bigint.clone()),
        _ => None,
    }
}

pub struct DerParser {
    buf: Vec<u8>,
    idx: usize,
}

impl DerParser {
    pub fn new(der: &Vec<u8>) -> DerParser {
        DerParser {
            buf: der.clone(),
            idx: 0,
        }
    }

    fn check_buffer(&mut self) -> ASN1Result<()> {
        if self.idx >= self.buf.len() {
            return Err(ASN1Error::MissingData("length".to_string()));
        }
        Ok(())
    }

    fn get_length(&mut self) -> ASN1Result<usize> {
        try!(self.check_buffer());

        let len_tmp = self.buf[self.idx];
        if len_tmp <= 127 {
            self.idx += 1;
            return Ok(len_tmp as usize);
        }

        let len_octets = len_tmp & 127;
        assert!(len_octets <= 3);
        let mut length = 0 as usize;
        self.idx += 1;
        for _ in 0..len_octets {
            length <<= 8;
            try!(self.check_buffer());
            length |= self.buf[self.idx] as usize;
            self.idx += 1;
        }
        Ok(length)
    }

    fn parse_header(&mut self) -> ASN1Result<(u8, bool, u8, usize)> {
        try!(self.check_buffer());

        let c = self.buf[self.idx];
        let tag = c & 31;
        let p_c = (c >> 5) & 1;
        let class_ = (c >> 6) & 3;
        assert!(tag != 31);
        self.idx += 1;
        let length = try!(self.get_length());

        Ok((class_ as u8, p_c == 1, tag as u8, length as usize))
    }

    fn parse_utf8_type(&self, raw: Vec<u8>) -> ASN1Result<String> {
        if let Ok(string) = String::from_utf8(raw) {
            Ok(string)
        } else {
            Err(ASN1Error::DecodeError("utf8 decode error".to_string()))
        }
    }

    pub fn parse_entry(&mut self) -> ASN1Result<ASN1Type> {
        try!(self.check_buffer());

        let (class_bits, constructed, tag, length) = try!(self.parse_header());

        assert!(class_bits == 0 || class_bits == 2);

        match class_bits {
            0 => {
            }
            2 => {
                let pos = self.idx;
                self.idx += length;

                assert!(pos + length <= self.buf.len());

                let sub_buf = self.buf[pos..pos+length].iter().cloned().collect();
                let mut sub_parser = DerParser::new(&sub_buf);
                let sub_entry = try!(sub_parser.parse_entry());
                return Ok(ASN1Type::Cont(tag, Box::new(sub_entry)));
            }
            _ => {
                panic!("Parsing this not implemented.");
            }
        }

        let pos = self.idx;
        self.idx += length;

        assert!(pos + length <= self.buf.len());

        let raw = self.buf[pos..pos+length].iter().cloned().collect();

        match tag {
            tag if tag == TypeId::Null as u8 => Ok(ASN1Type::Null),
            tag if tag == TypeId::Boolean as u8 => {
                let buf: Vec<u8> = raw;
                Ok(ASN1Type::Boolean(buf[0] == 0xff))
            }
            tag if tag == TypeId::Sequence as u8 => {
                let mut sub_parser = DerParser::new(&raw);
                let mut sequence = Vec::<ASN1Type>::new();
                loop {
                    let sub_ent = sub_parser.parse_entry();
                    match sub_ent {
                        Err(ASN1Error::MissingData(_)) => { break },
                        Err(err) => { return Err(err) }
                        Ok(ent) => { sequence.push(ent) }
                    }
                }
                Ok(ASN1Type::Sequence(sequence))
            },
            tag if tag == TypeId::Set as u8 => {
                let mut sub_parser = DerParser::new(&raw);
                let mut sequence = Vec::<ASN1Type>::new();
                loop {
                    let sub_ent = sub_parser.parse_entry();
                    match sub_ent {
                        Err(ASN1Error::MissingData(_)) => { break },
                        Err(err) => { return Err(err) }
                        Ok(ent) => { sequence.push(ent) }
                    }
                }
                Ok(ASN1Type::Set(sequence))
            },
            tag if tag == TypeId::Integer as u8 => Ok(ASN1Type::Integer(BigInt::from_bytes_be(Sign::Plus, &raw))),
            tag if tag == TypeId::Object as u8 => {
                let mut oi = Vec::<u32>::new();
                let obj_bytes = raw;
                oi.push((obj_bytes[0] / 40) as u32);
                oi.push((obj_bytes[0] % 40) as u32);
                let mut i = 1;
                while i < obj_bytes.len() {
                    if obj_bytes[i] <= 127 {
                        oi.push(obj_bytes[i] as u32);
                        i += 1;
                    } else {
                        let mut next_sub_id: u32 = obj_bytes[i] as u32;
                        let mut sub_id: u32 = 0;
                        i += 1;
                        while next_sub_id > 127 {
                            sub_id <<= 7;
                            sub_id |= next_sub_id & 127;
                            next_sub_id = obj_bytes[i] as u32;
                            i += 1;
                        }
                        oi.push((sub_id << 7) | next_sub_id);
                    }
                }
                Ok(ASN1Type::Object(vec_to_tup(&oi).expect("...")))
            },
            tag if tag == TypeId::BitString as u8 => Ok(ASN1Type::BitString(raw)),
            tag if tag == TypeId::OctetString as u8 => Ok(ASN1Type::OctetString(raw)),
            tag if tag == TypeId::PrintableString as u8 => {
                let string = try!(self.parse_utf8_type(raw));
                Ok(ASN1Type::PrintableString(string))
            },
            tag if tag == TypeId::UTCTime as u8 => {
                let string = try!(self.parse_utf8_type(raw));
                Ok(ASN1Type::UTCTime(string))
            },
            tag if tag == TypeId::UTF8String as u8 => {
                let string = try!(self.parse_utf8_type(raw));
                Ok(ASN1Type::UTF8String(string))
            },
            _ => {
                println!("Got Unknown ASN1 Type {} {} {} {}", class_bits, constructed, tag, length);
                Ok(ASN1Type::Unknown)
            }
        }
    }
}
<|endoftext|>