<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>will not consider empty 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 = 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
}
<|endoftext|>


<commit_before>use index::btree::dump::dump_tree;
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, HashMap, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;

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

struct KeyAltered {
    new_key: EntryKey,
    search_key: EntryKey,
    ptr: 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) -> Vec<KeyAltered>
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 prev_node_guard: Option<NodeWriteGuard<KS, PS>> = None;
    let mut page = write_node::<KS, PS>(&node);
    let mut next_live_ptrs = None;
    let mut altered = vec![];
    let mut empty_nodes = vec![];
    let select_live = |page: &NodeWriteGuard<KS, PS>| {
        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, sub_level.clone()))
                } else {
                    None
                }
            })
            .collect()
    };
    let mut next_non_empty = |page: &NodeWriteGuard<KS, PS>, empty_nodes: &mut Vec<_>| {
        let mut next = page.right_ref().cloned();
        while let Some(next_ref) = next {
            let guard = write_node::<KS, PS>(&next_ref);
            next = guard.right_ref().cloned();
            let live_ptrs: HashMap<_, _> = select_live(&guard);
            if live_ptrs.is_empty() {
                empty_nodes.push(guard);
            } else {
                return (guard, live_ptrs);
            }
        };
        unreachable!()
    };
    let preserve_live = |innode: &mut InNode<KS, PS>, live_ptrs: HashMap<usize, NodeCellRef>| {
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| live_ptrs.contains_key(&i))
                .map(|(_, k)| k)
                .collect();
            debug!("Prune filtered page have keys {:?}", &keys);
            debug_assert_eq!(live_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 live_ptrs.into_iter().enumerate() {
                debug_assert!(!ptr.is_default());
                mem::replace(&mut new_ptrs[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;
        innode.len == 0
    };
    loop {
        let page_right_ref = page.right_ref().unwrap().clone();
        if !page.is_empty_node() {
            let mut live_ptrs: HashMap<_, _> = next_live_ptrs.unwrap_or_else(|| select_live(&page));
            let is_empty = live_ptrs.is_empty();
            if !is_empty {
                let emptying = {
                    let mut innode = page.innode_mut();
                    preserve_live(&mut innode, live_ptrs)
                };
                if emptying {
                    // current innode have one ptr and no keys
                    // need to merge with the right page
                    // if the right page is full, partial of the right page will be moved to the current page
                    // merging right page will also been cleaned
                    let (mut next_guard, next_ptrs) = next_non_empty(&page, &mut empty_nodes);
                    let mid = next_guard.len() / 2;
                    let next_ref = next_guard.node_ref().clone();
                    let page_ref = page.node_ref().clone();
                    let mut page_innode = page.innode_mut();
                    let next_node_emptying = preserve_live(&mut next_guard.innode_mut(), next_ptrs);
                    page_innode.keys.as_slice()[0] = page_innode.right_bound.clone();
                    if next_guard.len() >= KS::slice_len() {
                        // will overflow if merged, have to rebalance right node contents with current node
                        let mut next_innode = next_guard.innode_mut();
                        let next_search_key = next_innode.keys.as_slice()[0].clone();
                        // move next node keys and ptrs to current node
                        for i in 0..mid {
                            mem::swap(&mut page_innode.keys.as_slice()[i + 1], &mut next_innode.keys.as_slice()[i]);
                            mem::swap(&mut page_innode.ptrs.as_slice()[i + 1], &mut next_innode.ptrs.as_slice()[i]);
                        }
                        for i in mid..next_innode.len {
                            next_innode.keys.as_slice().swap(i, i - mid);
                            next_innode.ptrs.as_slice().swap(i, i - mid);
                        }
                        next_innode.len -= mid;
                        page_innode.len += mid;
                        page_innode.right_bound = mem::replace(&mut page_innode.keys.as_slice()[mid], smallvec!());
                        page_innode.len = 1 + mid;
                        altered.push(KeyAltered {
                            new_key: page_innode.right_bound.clone(),
                            search_key: next_search_key,
                            ptr: next_ref
                        });

                    } else {
                        // Can move right page contents to current node and empty next node
                        {
                            let mut next_innode = next_guard.innode_mut();
                            let next_node_len = next_innode.len;
                            for i in 0..next_innode.len {
                                mem::swap(&mut page_innode.keys.as_slice()[i + 1], &mut next_innode.keys.as_slice()[i]);
                                mem::swap(&mut page_innode.ptrs.as_slice()[i + 1], &mut next_innode.ptrs.as_slice()[i]);
                            }
                            mem::swap(&mut page_innode.ptrs.as_slice()[next_node_len + 1], &mut next_innode.ptrs.as_slice()[next_node_len]);
                            page_innode.len = 1 + next_node_len;
                            page_innode.right = next_innode.right.clone();
                            page_innode.right_bound = next_innode.right_bound.clone();
                        }
                        *next_guard = NodeData::Empty(box EmptyNode {
                            left: None,
                            right: page_ref
                        });
                    }
                }
            }
            let (next_page, ptrs) = next_non_empty(&page, &mut empty_nodes);
            if !is_empty {
                if page.right_bound() >= bound {
                    break;
                } else {
                    prev_node_guard = Some(page);
                }
            } else {
                empty_nodes.push(page);
            }
            page = next_page;
            next_live_ptrs = Some(ptrs);
        } else {
            unreachable!();
        }
    }
    let last_page_ref = page.node_ref();
    for empty in &mut empty_nodes {
        // make nodes empty and set its next ptr to last non empty node
        **empty = NodeData::Empty(box EmptyNode {
            left: None,
            right: last_page_ref.clone(),
        });
    }
    altered
}

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());
        debug_assert!(read_unchecked::<KS, PS>(&right_right_most).is_ext());
        debug_assert!(!Arc::ptr_eq(
            &right_right_most.inner,
            &left_most_leaf_guards.last().unwrap().node_ref().inner
        ));

        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;

        debug_assert!(&new_first_node_ext.right_bound > &prune_bound);

        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>taken care of cornor cases<commit_after>use index::btree::dump::dump_tree;
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, HashMap, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::Arc;

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

struct KeyAltered {
    new_key: EntryKey,
    ptr: 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) -> Vec<KeyAltered>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    let extra_works = 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)
            } else {
                vec![]
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut prev_node_guard: Option<NodeWriteGuard<KS, PS>> = None;
    let mut page = write_node::<KS, PS>(&node);
    let mut next_live_ptrs = None;
    let mut altered_list = vec![];
    let mut empty_nodes = vec![];
    let is_altered = |ptr: &NodeCellRef| extra_works.iter().find(|a| a.ptr.ptr_eq(ptr));
    let select_live = |page: &NodeWriteGuard<KS, PS>| {
        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, sub_level.clone()))
                } else {
                    None
                }
            })
            .collect()
    };
    let mut next_non_empty = |page: &NodeWriteGuard<KS, PS>, empty_nodes: &mut Vec<_>| {
        let mut next = page.right_ref().cloned();
        while let Some(next_ref) = next {
            let guard = write_node::<KS, PS>(&next_ref);
            if guard.is_none() {
                return None;
            }
            next = guard.right_ref().cloned();
            let live_ptrs: HashMap<_, _> = select_live(&guard);
            if live_ptrs.is_empty() {
                empty_nodes.push(guard);
            } else {
                return Some((guard, live_ptrs));
            }
        }
        unreachable!()
    };
    let preserve_live = |innode: &mut InNode<KS, PS>,
                         node_ref: &NodeCellRef,
                         live_ptrs: HashMap<usize, NodeCellRef>,
                         altered_list: &mut Vec<KeyAltered>| {
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| live_ptrs.contains_key(&i))
                .map(|(_, k)| k)
                .collect();
            debug!("Prune filtered page have keys {:?}", &keys);
            debug_assert_eq!(live_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 live_ptrs.into_iter().enumerate() {
                debug_assert!(!ptr.is_default());
                if let Some(altered) = is_altered(&ptr) {
                    new_keys[i] = altered.new_key.clone();
                    if i == 0 {
                        altered_list.push(KeyAltered {
                            new_key: altered.new_key.clone(),
                            ptr: node_ref.clone(),
                        })
                    }
                }
                mem::replace(&mut new_ptrs[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;
        innode.len == 0
    };
    loop {
        let page_right_ref = page.right_ref().unwrap().clone();
        if !page.is_empty_node() {
            let mut live_ptrs: HashMap<_, _> = next_live_ptrs.unwrap_or_else(|| select_live(&page));
            let is_empty = live_ptrs.is_empty();
            if !is_empty {
                let emptying = {
                    let page_ref = page.node_ref().clone();
                    let mut innode = page.innode_mut();
                    preserve_live(&mut innode, &page_ref, live_ptrs, &mut altered_list)
                };
                if emptying {
                    // current innode have one ptr and no keys
                    // need to merge with the right page
                    // if the right page is full, partial of the right page will be moved to the current page
                    // merging right page will also been cleaned
                    if let Some((mut next_guard, next_ptrs)) =
                        next_non_empty(&page, &mut empty_nodes)
                    {
                        let mid = next_guard.len() / 2;
                        let next_ref = next_guard.node_ref().clone();
                        let page_ref = page.node_ref().clone();
                        let mut page_innode = page.innode_mut();
                        let next_node_emptying = preserve_live(
                            &mut next_guard.innode_mut(),
                            &next_ref,
                            next_ptrs,
                            &mut altered_list,
                        );
                        page_innode.keys.as_slice()[0] = page_innode.right_bound.clone();
                        if next_guard.len() >= KS::slice_len() {
                            // will overflow if merged, have to rebalance right node contents with current node
                            let mut next_innode = next_guard.innode_mut();
                            let next_search_key = next_innode.keys.as_slice()[0].clone();
                            // move next node keys and ptrs to current node
                            for i in 0..mid {
                                mem::swap(
                                    &mut page_innode.keys.as_slice()[i + 1],
                                    &mut next_innode.keys.as_slice()[i],
                                );
                                mem::swap(
                                    &mut page_innode.ptrs.as_slice()[i + 1],
                                    &mut next_innode.ptrs.as_slice()[i],
                                );
                            }
                            for i in mid..next_innode.len {
                                next_innode.keys.as_slice().swap(i, i - mid);
                                next_innode.ptrs.as_slice().swap(i, i - mid);
                            }
                            next_innode.len -= mid;
                            page_innode.len += mid;
                            page_innode.right_bound =
                                mem::replace(&mut page_innode.keys.as_slice()[mid], smallvec!());
                            page_innode.len = 1 + mid;
                            altered_list.push(KeyAltered {
                                new_key: page_innode.right_bound.clone(),
                                ptr: next_ref,
                            });
                        } else {
                            // Can move right page contents to current node and empty next node
                            {
                                let mut next_innode = next_guard.innode_mut();
                                let next_node_len = next_innode.len;
                                for i in 0..next_innode.len {
                                    mem::swap(
                                        &mut page_innode.keys.as_slice()[i + 1],
                                        &mut next_innode.keys.as_slice()[i],
                                    );
                                    mem::swap(
                                        &mut page_innode.ptrs.as_slice()[i + 1],
                                        &mut next_innode.ptrs.as_slice()[i],
                                    );
                                }
                                mem::swap(
                                    &mut page_innode.ptrs.as_slice()[next_node_len + 1],
                                    &mut next_innode.ptrs.as_slice()[next_node_len],
                                );
                                page_innode.len = 1 + next_node_len;
                                page_innode.right = next_innode.right.clone();
                                page_innode.right_bound = next_innode.right_bound.clone();
                            }
                            *next_guard = NodeData::Empty(box EmptyNode {
                                left: None,
                                right: page_ref,
                            });
                        }
                    } else {
                        break;
                    }
                }
            }
            if let Some((next_page, ptrs)) = next_non_empty(&page, &mut empty_nodes) {
                if !is_empty {
                    if page.right_bound() >= bound {
                        break;
                    } else {
                        prev_node_guard = Some(page);
                    }
                } else {
                    empty_nodes.push(page);
                }
                page = next_page;
                next_live_ptrs = Some(ptrs);
            } else {
                break;
            }
        } else {
            unreachable!();
        }
    }
    let last_page_ref = page.node_ref();
    for empty in &mut empty_nodes {
        // make nodes empty and set its next ptr to last non empty node
        **empty = NodeData::Empty(box EmptyNode {
            left: None,
            right: last_page_ref.clone(),
        });
    }
    // Reverse the list so only latest pushed will take effect
    altered_list.reverse();
    altered_list
}

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());
        debug_assert!(read_unchecked::<KS, PS>(&right_right_most).is_ext());
        debug_assert!(!Arc::ptr_eq(
            &right_right_most.inner,
            &left_most_leaf_guards.last().unwrap().node_ref().inner
        ));

        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;

        debug_assert!(&new_first_node_ext.right_bound > &prune_bound);

        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>
extern crate num;
use num::complex::Complex;

#[derive(Debug)]
struct ModelSingle {
    num_kernels: usize,
    kernels: Vec<Vec<Complex<f64>>>,
}

#[derive(Debug)]
struct ModelArray {
    num_models: usize,
    models: Vec<ModelSingle>,
}

#[derive(Debug)]
struct KernelIterator<'a> {
    num_kernels: usize,
    state: (usize, usize),
    ref_mod: &'a ModelArray,
    full_rank: Vec<usize>,
    reduced_rank: Vec<usize>,
}

impl<'b> KernelIterator<'b> {
    fn new(num_kernels: usize, ref_mod: &'b ModelArray) -> KernelIterator {
        let num_models = ref_mod.num_models;
        let mut full_rank     = vec![0; num_models];
        let mut reduced_rank  = vec![0; num_models];
        let mut total_kernels = 0;
        for i in 0..num_models {
            let nkernels = ref_mod.models[i].num_kernels;
            full_rank[i] = nkernels;
            total_kernels += nkernels;
        }
        let delta = total_kernels - num_kernels;

        println!("total_kernels: {} num_kernels: {} delta: {}",
                 total_kernels, num_kernels, delta);

        for i in 0..num_models {
            let nkernels = full_rank[i];
            let fraction = nkernels as f64 / total_kernels as f64;
            reduced_rank[i] = nkernels - ((delta as f64 * fraction) as usize);
        }

        KernelIterator { num_kernels: num_kernels, state : (0, 0), ref_mod: ref_mod,
                         full_rank: full_rank, reduced_rank: reduced_rank}
    }

    fn valid_state(&self) -> bool {
        let (model_index, kernel_index) = self.state;
        model_index  < self.ref_mod.num_models &&
        kernel_index < self.reduced_rank[model_index]
    }

    fn advance_state(&mut self) -> bool {
        let (model_index, kernel_index) = self.state;
        if kernel_index+1 < self.reduced_rank[model_index] {
            self.state = (model_index, kernel_index + 1);
            true
        } else if model_index < self.ref_mod.num_models {
            self.state = (model_index + 1, 0);
            true
        } else {
            self.state = (model_index + 1, kernel_index);
            false
        }
    }
}

impl<'c> Iterator for KernelIterator<'c> {
    type Item = Vec<Complex<f64>>;

    fn next(&mut self) -> Option<Vec<Complex<f64>>> {
        let ret = if self.valid_state() {
            let (model_index, kernel_index) = self.state;
            let m = self.ref_mod;
            let ref model = m.models[model_index];
            println!("model_index: {} kernel_index: {}", model_index, kernel_index);
            Some(model.kernels[kernel_index].clone())
        } else {
            None
        };
        self.advance_state();
        ret
    }
}

fn main() {
    println!("This is a kernel_iterator test");

    let kq1 = ModelSingle{num_kernels: 21,  kernels: vec![Vec::new(); 21]};
    let kq2 = ModelSingle{num_kernels: 75,  kernels: vec![Vec::new(); 75]};
    let kq3 = ModelSingle{num_kernels: 121, kernels: vec![Vec::new(); 121]};

    let mut mod_array = ModelArray {num_models: 3, models: Vec::new()};
    mod_array.models.push(kq1);
    mod_array.models.push(kq2);
    mod_array.models.push(kq3);

    println!("mod_array is constructed and properly initialized: {:?}", mod_array);

    let ki = KernelIterator::new(17, &mod_array);

    println!("KernelIterator is constructed and properly initialized: {:?}", ki);

    for _ in ki {
    }
}

<commit_msg>a nasty bug with index out of bounds fixed.<commit_after>
extern crate num;
use num::complex::Complex;

#[derive(Debug)]
struct ModelSingle {
    num_kernels: usize,
    kernels: Vec<Vec<Complex<f64>>>,
}

#[derive(Debug)]
struct ModelArray {
    num_models: usize,
    models: Vec<ModelSingle>,
}

#[derive(Debug)]
struct KernelIterator<'a> {
    num_kernels: usize,
    state: (usize, usize),
    ref_mod: &'a ModelArray,
    full_rank: Vec<usize>,
    reduced_rank: Vec<usize>,
}

impl<'b> KernelIterator<'b> {
    fn new(num_kernels: usize, ref_mod: &'b ModelArray) -> KernelIterator {
        let num_models = ref_mod.num_models;
        let mut full_rank     = vec![0; num_models];
        let mut reduced_rank  = vec![0; num_models];
        let mut total_kernels = 0;
        for i in 0..num_models {
            let nkernels = ref_mod.models[i].num_kernels;
            full_rank[i] = nkernels;
            total_kernels += nkernels;
        }
        let delta = total_kernels - num_kernels;

        println!("total_kernels: {} num_kernels: {} delta: {}",
                 total_kernels, num_kernels, delta);

        for i in 0..num_models {
            let nkernels = full_rank[i];
            let fraction = nkernels as f64 / total_kernels as f64;
            reduced_rank[i] = nkernels - ((delta as f64 * fraction) as usize);
        }

        KernelIterator { num_kernels: num_kernels, state : (0, 0), ref_mod: ref_mod,
                         full_rank: full_rank, reduced_rank: reduced_rank}
    }

    fn valid_state(&self) -> bool {
        let (model_index, kernel_index) = self.state;
        println!("inside `valid_state` model_index: {} kernel_index: {}", model_index, kernel_index);
        println!("model_index  < self.ref_mod.num_models : {}", model_index  < self.ref_mod.num_models);
        model_index  < self.ref_mod.num_models //&&
        //kernel_index < self.reduced_rank[model_index]
    }

    fn advance_state(&mut self) -> bool {
        let (model_index, kernel_index) = self.state;
        if kernel_index+1 < self.reduced_rank[model_index] {
            self.state = (model_index, kernel_index + 1);
            true
        } else if model_index < self.ref_mod.num_models {
            self.state = (model_index + 1, 0);
            true
        } else {
            self.state = (model_index + 1, kernel_index);
            false
        }
    }
}

impl<'c> Iterator for KernelIterator<'c> {
    type Item = Vec<Complex<f64>>;

    fn next(&mut self) -> Option<Vec<Complex<f64>>> {
        let ret = if self.valid_state() {
            let (model_index, kernel_index) = self.state;
            let m = self.ref_mod;
            let ref model = m.models[model_index];
            println!("model_index: {} kernel_index: {}", model_index, kernel_index);
            self.advance_state();
            Some(model.kernels[kernel_index].clone())
        } else {
            None
        };
        ret
    }
}

fn main() {
    println!("This is a kernel_iterator test");

    let kq1 = ModelSingle{num_kernels: 21,  kernels: vec![Vec::new(); 21]};
    let kq2 = ModelSingle{num_kernels: 75,  kernels: vec![Vec::new(); 75]};
    let kq3 = ModelSingle{num_kernels: 121, kernels: vec![Vec::new(); 121]};

    let mut mod_array = ModelArray {num_models: 3, models: Vec::new()};
    mod_array.models.push(kq1);
    mod_array.models.push(kq2);
    mod_array.models.push(kq3);

    println!("mod_array is constructed and properly initialized: {:?}", mod_array);

    let ki = KernelIterator::new(17, &mod_array);

    println!("KernelIterator is constructed and properly initialized: {:?}", ki);

    for z in ki {
        let kernel: &Vec<Complex<f64>> = &z;
    }
}

<|endoftext|>


<commit_before>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

mod address;

use super::{
    DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings,
    State, StateChangeImpl, Validation,
};
use slog_scope::{error, info, trace};
use std::path::PathBuf;

pub(crate) use address::{
    AbortDownloadResponse, Addr, Message, ProbeResponse, Response, StateResponse,
};

pub(super) struct StateMachine {
    state: State,
    context: Context,
}

pub struct Context {
    pub(super) communication: Channel<(Message, async_channel::Sender<Result<Response>>)>,
    pub(super) waker: Channel<()>,
    pub settings: Settings,
    pub runtime_settings: RuntimeSettings,
    pub firmware: Metadata,
}

pub(super) struct Channel<T> {
    pub(super) sender: async_channel::Sender<T>,
    pub(super) receiver: async_channel::Receiver<T>,
}

impl<T> Channel<T> {
    fn new(cap: usize) -> Self {
        let (sender, receiver) = async_channel::bounded(cap);
        Channel { sender, receiver }
    }
}

impl CommunicationState for State {}

#[async_trait::async_trait]
pub(super) trait CommunicationState: StateChangeImpl {
    async fn handle_communication(
        &self,
        msg: address::Message,
        responder: async_channel::Sender<Result<address::Response>>,
        context: &mut Context,
    ) -> Option<State> {
        trace!("received external request: {:?}", msg);

        let res = match msg {
            address::Message::Info => {
                let state = self.name().to_owned();
                Ok((
                    address::Response::Info(Box::new(sdk::api::info::Response {
                        state,
                        version: crate::version().to_string(),
                        config: context.settings.0.clone(),
                        firmware: context.firmware.0.clone(),
                        runtime_settings: context.runtime_settings.inner.clone(),
                    })),
                    None,
                ))
            }
            address::Message::Probe(custom_server) => self
                .handle_probe(context, custom_server)
                .await
                .map(|(res, st)| (address::Response::Probe(res), st)),
            address::Message::AbortDownload => self
                .handle_abort_download(context)
                .await
                .map(|(res, st)| (address::Response::AbortDownload(res), st)),
            address::Message::LocalInstall(update_file) => self
                .handle_local_install(context, update_file)
                .await
                .map(|(res, st)| (address::Response::LocalInstall(res), st)),
            address::Message::RemoteInstall(url) => self
                .handle_remote_install(context, url)
                .await
                .map(|(res, st)| (address::Response::RemoteInstall(res), st)),
        };

        match res {
            Ok((response, state)) => {
                responder.send(Ok(response)).await.ok()?;
                state
            }
            Err(e) => {
                error!("Request failed with: {}", e);
                responder.send(Err(e)).await.ok()?;
                None
            }
        }
    }

    async fn handle_probe(
        &self,
        context: &mut Context,
        custom_server: Option<String>,
    ) -> Result<(address::ProbeResponse, Option<State>)> {
        use chrono::Utc;
        use cloud::api::ProbeResponse;

        if !self.is_preemptive_state() {
            let name = self.name().to_owned();
            return Ok((address::ProbeResponse::Busy(name), None));
        }
        // Starting logging a new scope of operation since we are
        // starting to handle a user request
        crate::logger::start_memory_logging();
        info!("Probing the server as requested by the user");

        if let Some(server_address) = custom_server {
            context.runtime_settings.set_custom_server_address(&server_address);
        }

        match crate::CloudClient::new(context.server_address())
            .probe(context.runtime_settings.retries(), context.firmware.as_cloud_metadata())
            .await?
        {
            ProbeResponse::ExtraPoll(s) => {
                info!("server responded with extra poll of {} seconds", s);
                Ok((address::ProbeResponse::Delayed(s), None))
            }

            ProbeResponse::NoUpdate => {
                info!("no update is current available for this device");
                context.waker.sender.send(()).await?;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((address::ProbeResponse::Unavailable, Some(State::EntryPoint(EntryPoint {}))))
            }

            ProbeResponse::Update(package, sign) => {
                info!("update received: {} ({})", package.version(), package.package_uid());
                context.waker.sender.send(()).await?;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((
                    address::ProbeResponse::Available,
                    Some(State::Validation(Validation { package, sign, require_download: true })),
                ))
            }
        }
    }

    async fn handle_abort_download(
        &self,
        _: &Context,
    ) -> Result<(address::AbortDownloadResponse, Option<State>)> {
        if self.is_handling_download() {
            Ok((
                address::AbortDownloadResponse::RequestAccepted,
                Some(State::EntryPoint(EntryPoint {})),
            ))
        } else {
            Ok((address::AbortDownloadResponse::InvalidState, None))
        }
    }

    async fn handle_local_install(
        &self,
        context: &Context,
        update_file: PathBuf,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();
        if self.is_preemptive_state() {
            // Starting logging a new scope of operation since we are
            // starting to handle a user request
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await?;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::PrepareLocalInstall(PrepareLocalInstall { update_file })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }

    async fn handle_remote_install(
        &self,
        context: &Context,
        url: String,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();

        if self.is_preemptive_state() {
            // Starting logging a new scope of operation since we are
            // starting to handle a user request
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await?;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::DirectDownload(DirectDownload { url })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }
}

impl Context {
    pub(crate) fn new(
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        Context {
            communication: Channel::new(10),
            waker: Channel::new(1),
            settings,
            runtime_settings,
            firmware,
        }
    }

    pub(super) fn server_address(&self) -> &str {
        self.runtime_settings
            .custom_server_address()
            .unwrap_or(&self.settings.network.server_address)
    }
}

#[derive(Debug)]
pub(super) enum StepTransition {
    Delayed(chrono::Duration),
    Immediate,
    Never,
}

impl StateMachine {
    pub(super) fn new(
        state: State,
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        StateMachine { state, context: Context::new(settings, runtime_settings, firmware) }
    }

    pub(super) fn address(&self) -> Addr {
        Addr { message: self.context.communication.sender.clone() }
    }

    pub(super) async fn start(mut self) {
        loop {
            // Since the loop is already currently running, we can
            // discharges any wake message received.
            let _ = self.context.waker.receiver.try_recv();

            self.consume_pending_communication().await;

            let (state, transition) = self
                .state
                .handle(&mut self.context)
                .await
                .unwrap_or_else(|e| (State::from(e), StepTransition::Immediate));
            self.state = state;

            match transition {
                StepTransition::Immediate => {}
                StepTransition::Delayed(t) => {
                    trace!("delaying transition for: {} seconds", t.num_seconds());
                    let waker = self.context.waker.receiver.clone();

                    let sleep_fut = tokio::time::sleep(t.to_std().unwrap());
                    let waker_fut = async {
                        let _ = waker.recv().await;
                    };
                    let comm_fut = self.await_communication();

                    futures_util::pin_mut!(sleep_fut);
                    futures_util::pin_mut!(waker_fut);
                    futures_util::pin_mut!(comm_fut);

                    let _ = futures_util::future::select(
                        futures_util::future::select(sleep_fut, waker_fut),
                        comm_fut,
                    )
                    .await;
                }
                StepTransition::Never => {
                    trace!("stopping transition until awoken");
                    let waker_recv = self.context.waker.receiver.clone();
                    let recv_fut = waker_recv.recv();
                    let comm_fut = async {
                        self.await_communication().await;
                        std::result::Result::<_, async_channel::RecvError>::Ok(())
                    };

                    // recv_fut dones't need to be pinned as it doesn't capture any context
                    futures_util::pin_mut!(comm_fut);
                    let _ = futures_util::future::select(recv_fut, comm_fut).await;
                }
            }
        }
    }

    async fn consume_pending_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.try_recv() {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }

    async fn await_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.recv().await {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }
}
<commit_msg>updatehub: states: machine: fix typo<commit_after>// Copyright (C) 2020 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

mod address;

use super::{
    DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings,
    State, StateChangeImpl, Validation,
};
use slog_scope::{error, info, trace};
use std::path::PathBuf;

pub(crate) use address::{
    AbortDownloadResponse, Addr, Message, ProbeResponse, Response, StateResponse,
};

pub(super) struct StateMachine {
    state: State,
    context: Context,
}

pub struct Context {
    pub(super) communication: Channel<(Message, async_channel::Sender<Result<Response>>)>,
    pub(super) waker: Channel<()>,
    pub settings: Settings,
    pub runtime_settings: RuntimeSettings,
    pub firmware: Metadata,
}

pub(super) struct Channel<T> {
    pub(super) sender: async_channel::Sender<T>,
    pub(super) receiver: async_channel::Receiver<T>,
}

impl<T> Channel<T> {
    fn new(cap: usize) -> Self {
        let (sender, receiver) = async_channel::bounded(cap);
        Channel { sender, receiver }
    }
}

impl CommunicationState for State {}

#[async_trait::async_trait]
pub(super) trait CommunicationState: StateChangeImpl {
    async fn handle_communication(
        &self,
        msg: address::Message,
        responder: async_channel::Sender<Result<address::Response>>,
        context: &mut Context,
    ) -> Option<State> {
        trace!("received external request: {:?}", msg);

        let res = match msg {
            address::Message::Info => {
                let state = self.name().to_owned();
                Ok((
                    address::Response::Info(Box::new(sdk::api::info::Response {
                        state,
                        version: crate::version().to_string(),
                        config: context.settings.0.clone(),
                        firmware: context.firmware.0.clone(),
                        runtime_settings: context.runtime_settings.inner.clone(),
                    })),
                    None,
                ))
            }
            address::Message::Probe(custom_server) => self
                .handle_probe(context, custom_server)
                .await
                .map(|(res, st)| (address::Response::Probe(res), st)),
            address::Message::AbortDownload => self
                .handle_abort_download(context)
                .await
                .map(|(res, st)| (address::Response::AbortDownload(res), st)),
            address::Message::LocalInstall(update_file) => self
                .handle_local_install(context, update_file)
                .await
                .map(|(res, st)| (address::Response::LocalInstall(res), st)),
            address::Message::RemoteInstall(url) => self
                .handle_remote_install(context, url)
                .await
                .map(|(res, st)| (address::Response::RemoteInstall(res), st)),
        };

        match res {
            Ok((response, state)) => {
                responder.send(Ok(response)).await.ok()?;
                state
            }
            Err(e) => {
                error!("Request failed with: {}", e);
                responder.send(Err(e)).await.ok()?;
                None
            }
        }
    }

    async fn handle_probe(
        &self,
        context: &mut Context,
        custom_server: Option<String>,
    ) -> Result<(address::ProbeResponse, Option<State>)> {
        use chrono::Utc;
        use cloud::api::ProbeResponse;

        if !self.is_preemptive_state() {
            let name = self.name().to_owned();
            return Ok((address::ProbeResponse::Busy(name), None));
        }
        // Starting logging a new scope of operation since we are
        // starting to handle a user request
        crate::logger::start_memory_logging();
        info!("Probing the server as requested by the user");

        if let Some(server_address) = custom_server {
            context.runtime_settings.set_custom_server_address(&server_address);
        }

        match crate::CloudClient::new(context.server_address())
            .probe(context.runtime_settings.retries(), context.firmware.as_cloud_metadata())
            .await?
        {
            ProbeResponse::ExtraPoll(s) => {
                info!("server responded with extra poll of {} seconds", s);
                Ok((address::ProbeResponse::Delayed(s), None))
            }

            ProbeResponse::NoUpdate => {
                info!("no update is current available for this device");
                context.waker.sender.send(()).await?;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((address::ProbeResponse::Unavailable, Some(State::EntryPoint(EntryPoint {}))))
            }

            ProbeResponse::Update(package, sign) => {
                info!("update received: {} ({})", package.version(), package.package_uid());
                context.waker.sender.send(()).await?;

                // Store timestamp of last polling
                context.runtime_settings.set_last_polling(Utc::now())?;
                Ok((
                    address::ProbeResponse::Available,
                    Some(State::Validation(Validation { package, sign, require_download: true })),
                ))
            }
        }
    }

    async fn handle_abort_download(
        &self,
        _: &Context,
    ) -> Result<(address::AbortDownloadResponse, Option<State>)> {
        if self.is_handling_download() {
            Ok((
                address::AbortDownloadResponse::RequestAccepted,
                Some(State::EntryPoint(EntryPoint {})),
            ))
        } else {
            Ok((address::AbortDownloadResponse::InvalidState, None))
        }
    }

    async fn handle_local_install(
        &self,
        context: &Context,
        update_file: PathBuf,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();
        if self.is_preemptive_state() {
            // Starting logging a new scope of operation since we are
            // starting to handle a user request
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await?;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::PrepareLocalInstall(PrepareLocalInstall { update_file })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }

    async fn handle_remote_install(
        &self,
        context: &Context,
        url: String,
    ) -> Result<(address::StateResponse, Option<State>)> {
        let name = self.name().to_owned();

        if self.is_preemptive_state() {
            // Starting logging a new scope of operation since we are
            // starting to handle a user request
            crate::logger::start_memory_logging();
            context.waker.sender.send(()).await?;

            Ok((
                address::StateResponse::RequestAccepted(name),
                Some(State::DirectDownload(DirectDownload { url })),
            ))
        } else {
            Ok((address::StateResponse::InvalidState(name), None))
        }
    }
}

impl Context {
    pub(crate) fn new(
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        Context {
            communication: Channel::new(10),
            waker: Channel::new(1),
            settings,
            runtime_settings,
            firmware,
        }
    }

    pub(super) fn server_address(&self) -> &str {
        self.runtime_settings
            .custom_server_address()
            .unwrap_or(&self.settings.network.server_address)
    }
}

#[derive(Debug)]
pub(super) enum StepTransition {
    Delayed(chrono::Duration),
    Immediate,
    Never,
}

impl StateMachine {
    pub(super) fn new(
        state: State,
        settings: Settings,
        runtime_settings: RuntimeSettings,
        firmware: Metadata,
    ) -> Self {
        StateMachine { state, context: Context::new(settings, runtime_settings, firmware) }
    }

    pub(super) fn address(&self) -> Addr {
        Addr { message: self.context.communication.sender.clone() }
    }

    pub(super) async fn start(mut self) {
        loop {
            // Since the loop is already currently running, we can
            // discharges any wake message received.
            let _ = self.context.waker.receiver.try_recv();

            self.consume_pending_communication().await;

            let (state, transition) = self
                .state
                .handle(&mut self.context)
                .await
                .unwrap_or_else(|e| (State::from(e), StepTransition::Immediate));
            self.state = state;

            match transition {
                StepTransition::Immediate => {}
                StepTransition::Delayed(t) => {
                    trace!("delaying transition for: {} seconds", t.num_seconds());
                    let waker = self.context.waker.receiver.clone();

                    let sleep_fut = tokio::time::sleep(t.to_std().unwrap());
                    let waker_fut = async {
                        let _ = waker.recv().await;
                    };
                    let comm_fut = self.await_communication();

                    futures_util::pin_mut!(sleep_fut);
                    futures_util::pin_mut!(waker_fut);
                    futures_util::pin_mut!(comm_fut);

                    let _ = futures_util::future::select(
                        futures_util::future::select(sleep_fut, waker_fut),
                        comm_fut,
                    )
                    .await;
                }
                StepTransition::Never => {
                    trace!("stopping transition until awoken");
                    let waker_recv = self.context.waker.receiver.clone();
                    let recv_fut = waker_recv.recv();
                    let comm_fut = async {
                        self.await_communication().await;
                        std::result::Result::<_, async_channel::RecvError>::Ok(())
                    };

                    // recv_fut doesn't need to be pinned as it doesn't capture any context
                    futures_util::pin_mut!(comm_fut);
                    let _ = futures_util::future::select(recv_fut, comm_fut).await;
                }
            }
        }
    }

    async fn consume_pending_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.try_recv() {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }

    async fn await_communication(&mut self) {
        while let Ok((msg, responder)) = self.context.communication.receiver.recv().await {
            if let Some(new_state) =
                self.state.handle_communication(msg, responder, &mut self.context).await
            {
                self.state = new_state;
            }
        }
    }
}
<|endoftext|>


<commit_before>
/// Implemented by contexts that can color.
///
/// The context can color to type `T` by adding value `U`.
pub trait CanColor<'a, T, U> {
    /// Create a new context with color.
    fn color(&'a self, value: U) -> T;
}

/// Implemented by contexts that can set rectangle information.
pub trait CanRectangle<'a, T, U> {
    /// Create new context with recangle information.
    fn rectangle(&'a self, value: U) -> T;
}

/// Implemented by contexts that can transform.
///
/// The context can transform to type `T` by adding value `U`.
pub trait CanTransform<'a, T, U> {
    /// Create a new context with transformation.
    fn transform(&'a self, value: U) -> T;
}

/// Implemented by contexts that uses type `U` as current color.
pub trait HasColor<'a, U> {
    /// Returns the current color.
    fn get_color(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current rectangle.
pub trait HasRectangle<'a, U> {
    /// Returns the current rectangle.
    fn get_rectangle(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current transform.
///
/// This helps to remove redundant code.
pub trait HasTransform<'a, U> {
    /// Returns the current transform.
    fn get_transform(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current transform.
///
/// This helps to remove redundant code.
pub trait HasViewTransform<'a, U> {
    /// Returns the current view transform.
    fn get_view_transform(&'a self) -> &'a U;
}

<commit_msg>Added `CanViewTransform` to `internal`<commit_after>
/// Implemented by contexts that can color.
///
/// The context can color to type `T` by adding value `U`.
pub trait CanColor<'a, T, U> {
    /// Create a new context with color.
    fn color(&'a self, value: U) -> T;
}

/// Implemented by contexts that can set rectangle information.
pub trait CanRectangle<'a, T, U> {
    /// Create new context with recangle information.
    fn rectangle(&'a self, value: U) -> T;
}

/// Implemented by contexts that can transform.
///
/// The context can transform to type `T` by adding value `U`.
pub trait CanTransform<'a, T, U> {
    /// Create a new context with transformation.
    fn transform(&'a self, value: U) -> T;
}

/// Implemented by contexts that can view transform.
///
/// The context can view transform to type `T` by adding value `U`.
pub trait CanViewTransform<'a, T, U> {
    /// Create a new context with view transformation.
    fn view_transform(&'a self, value: U) -> T;
}

/// Implemented by contexts that uses type `U` as current color.
pub trait HasColor<'a, U> {
    /// Returns the current color.
    fn get_color(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current rectangle.
pub trait HasRectangle<'a, U> {
    /// Returns the current rectangle.
    fn get_rectangle(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current transform.
///
/// This helps to remove redundant code.
pub trait HasTransform<'a, U> {
    /// Returns the current transform.
    fn get_transform(&'a self) -> &'a U;
}

/// Implemented by contexts that uses type `U` as current view transform.
///
/// This helps to remove redundant code.
pub trait HasViewTransform<'a, U> {
    /// Returns the current view transform.
    fn get_view_transform(&'a self) -> &'a U;
}

<|endoftext|>