<commit_before>mod Worksheet;
// Basic excel

struct Basic_Excel {
    pub file_: Compound_File,                  // Compound file handler.
    pub workbook_: Workbook,                   // Raw Workbook.
    pub worksheets_: Vec<Worksheet::Worksheet>,// Raw Worksheets.
    pub yesheets_: Vec<Basic_Excel_Worksheet>, // Parsed Worksheets.
}

impl Basic_Excel {

    pub fn new() -> Basic_Excel {
        return Basic_Excel{file_: Compound_File, workbook_: Workbook, worksheets_: vec![Worksheet::Worksheet], yesheets_: vec![Basic_Excel_Worksheet::new()]};
    }

    // File functions
    pub fn new_with_sheets(sheets: i32) {

    }

    pub fn load<T>(filename: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn save() -> bool {
        let mut result = false;

        return result;
    }

    pub fn save_as<T>(filename: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn close() {

    }

    // Worksheet functions
    pub fn get_total_worksheets() -> i32 {
        let mut result = 0;

        return result;
    }

    pub fn get_worksheet<T>(sheet_index: T) -> Basic_Excel_Worksheet {
        let mut result = Basic_Excel_Worksheet::new();

        return result;
    }

    pub fn add_worksheet<T>(sheet_index: T) -> Basic_Excel_Worksheet {
        let mut result = Basic_Excel_Worksheet::new();

        return result;
    }

    pub fn delete_worksheet<T>(sheet_index: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn get_ansi_sheet_name(sheet_index: i32) {
        panic!("Rust operates completely in UTF-8 Unicode strings and as such there's no reason return an ANSI version.");
    }

    pub fn get_unicode_sheet_name(sheet_index: i32) -> String {
        let mut result = String::new();

        return result;
    }

    pub fn get_sheet_name(sheet_index: i32) -> (bool, String) {
        let mut result = false;
        let mut name = String::new();

        return (result, name);
    }

    pub fn rename_worksheet_from_index(sheet_index: i32, to: &str) -> bool {
        let mut result = false;

        return result;
    }

    pub fn rename_worksheet_from_str(from: &str, to: &str) -> bool {
        let mut result = false;

        return result;
    }

}

struct Basic_Excel_Worksheet {
    pub excel_: Basic_Excel,
    pub sheet_index_: i32,
    pub max_rows_: i32,
    pub max_cols_: i32,
    pub cells_: Vec<Basic_Excel_Cell>,
    pub col_infos_: Worksheet::Col_Infos,
}

impl Basic_Excel_Worksheet {
    pub fn new() -> Basic_Excel_Worksheet {
        return Basic_Excel_Worksheet{excel_: Basic_Excel::new(), sheet_index_: 0, max_rows_: 0, max_cols_: 0, cells_: vec![Basic_Excel_Cell], col_infos_: Worksheet::Col_Infos};
    }
}

struct Compound_File;

struct Workbook;

struct Basic_Excel_Cell;

<commit_msg>Skeleton for Basic_Excel_Worksheet<commit_after>mod Worksheet;
// Basic excel

struct Basic_Excel {
    pub file_: Compound_File,                  // Compound file handler.
    pub workbook_: Workbook,                   // Raw Workbook.
    pub worksheets_: Vec<Worksheet::Worksheet>,// Raw Worksheets.
    pub yesheets_: Vec<Basic_Excel_Worksheet>, // Parsed Worksheets.
}

impl Basic_Excel {

    pub fn new() -> Basic_Excel {
        return Basic_Excel{file_: Compound_File, workbook_: Workbook, worksheets_: vec![Worksheet::Worksheet], yesheets_: vec![Basic_Excel_Worksheet::new()]};
    }

    // File functions
    pub fn new_with_sheets(sheets: i32) {
        unimplemented!();
    }

    pub fn load<T>(filename: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn save() -> bool {
        let mut result = false;

        return result;
    }

    pub fn save_as<T>(filename: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn close() {
        unimplemented!();
    }

    // Worksheet functions
    pub fn get_total_worksheets() -> i32 {
        let mut result = 0;

        return result;
    }

    pub fn get_worksheet<T>(sheet_index: T) -> Basic_Excel_Worksheet {
        let mut result = Basic_Excel_Worksheet::new();

        return result;
    }

    pub fn add_worksheet<T>(sheet_index: T) -> Basic_Excel_Worksheet {
        let mut result = Basic_Excel_Worksheet::new();

        return result;
    }

    pub fn delete_worksheet<T>(sheet_index: T) -> bool {
        let mut result = false;

        return result;
    }

    pub fn get_ansi_sheet_name(sheet_index: i32) {
        panic!("Rust operates completely in UTF-8 Unicode strings and as such there's no reason return an ANSI version.");
    }

    pub fn get_unicode_sheet_name(sheet_index: i32) -> String {
        let mut result = String::new();

        return result;
    }

    pub fn get_sheet_name(sheet_index: i32) -> (bool, String) {
        let mut result = false;
        let mut name = String::new();

        return (result, name);
    }

    pub fn rename_worksheet_from_index(sheet_index: i32, to: &str) -> bool {
        let mut result = false;

        return result;
    }

    pub fn rename_worksheet_from_str(from: &str, to: &str) -> bool {
        let mut result = false;

        return result;
    }

}

struct Basic_Excel_Worksheet {
    pub excel_: Basic_Excel,
    pub sheet_index_: i32,
    pub max_rows_: i32,
    pub max_cols_: i32,
    pub cells_: Vec<Basic_Excel_Cell>,
    pub col_infos_: Worksheet::Col_Infos,
}

impl Basic_Excel_Worksheet {
    // Worksheet functions
    pub fn new() -> Basic_Excel_Worksheet {
        return Basic_Excel_Worksheet{excel_: Basic_Excel::new(), sheet_index_: 0, max_rows_: 0, max_cols_: 0, cells_: vec![Basic_Excel_Cell], col_infos_: Worksheet::Col_Infos};
    }

    pub fn get_sheet_name() -> (bool, String) {
        unimplemented!();
    }

    pub fn rename(to: &str) -> bool {
        unimplemented!();
    }

    pub fn print() {
        unimplemented!();
    }

    // Cell functions
    pub fn get_total_rows() -> i32 {
        unimplemented!();
    }

    pub fn get_total_cols() -> i32 {
        unimplemented!();
    }

    pub fn cell(row: i32, col: i32) -> *mut Basic_Excel_Cell {
        unimplemented!();
    }

    pub fn erase_cell(row: i32, col: i32) -> bool {
        unimplemented!();
    }

    pub fn set_col_width(col_index: i32, col_width: u16) {
        unimplemented!();
    }

    pub fn get_col_width(col_index: i32) -> u16 {
        unimplemented!();
    }

    pub fn merge_cells(row: i32, col: i32, row_range: u16, col_range: u16) {
        unimplemented!();
    }

    // Private
    fn update_cells() {
        unimplemented!();
    }
}

struct Compound_File;

struct Workbook;

struct Basic_Excel_Cell;

<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::path::PathBuf;
use std::str::Split;

use clap::ArgMatches;
use toml::Value;

use libimagstore::store::EntryHeader;
use libimagrt::runtime::Runtime;
use libimagutil::key_value_split::IntoKeyValue;

pub fn build_entry_path(rt: &Runtime, path_elem: &str) -> PathBuf {
    debug!("Building path...");
    let mut path = rt.store().path().clone();

    if path_elem.chars().next() == Some('/') {
        path.push(&path_elem[1..path_elem.len()]);
    } else {
        path.push(path_elem);
    }

    path
}

pub fn build_toml_header(matches: &ArgMatches, header: EntryHeader) -> EntryHeader {
    debug!("Building header from cli spec");
    if let Some(headerspecs) = matches.values_of("header") {
        let mut main = BTreeMap::new();
        for tpl in headerspecs.into_iter().filter_map(|hs| String::from(hs).into_kv()) {
            let (key, value) = tpl.into();
            debug!("Splitting: {:?}", key);
            let mut split = key.split(".");
            let current = split.next();
            if current.is_some() {
                insert_key_into(String::from(current.unwrap()), &mut split, value, &mut main);
            }
        }
    }
    header
}

fn insert_key_into(current: String,
                   rest_path: &mut Split<&str>,
                   value: String,
                   map: &mut BTreeMap<String, Value>) {
    let next = rest_path.next();

    if next.is_none() {
        map.insert(current, parse_value(value));
    } else {
        if map.contains_key(&current) {
            match map.get_mut(&current).unwrap() {
                &mut Value::Table(ref mut t) => {
                    insert_key_into(String::from(next.unwrap()), rest_path, value, t);
                },
                _ => unreachable!(),
            }
        } else {
            let mut submap = BTreeMap::new();
            insert_key_into(String::from(next.unwrap()), rest_path, value, &mut submap);
            map.insert(current, Value::Table(submap));
        }
    }
}

fn parse_value(value: String) -> Value {
    fn is_ary(v: &String) -> bool {
        v.chars().next() == Some('[') && v.chars().last() == Some(']') && v.len() >= 3
    }

    debug!("Building value out of: {:?}", value);
    if value == "true" {
        Value::Boolean(true)
    } else if value == "false" {
        Value::Boolean(false)
    } else if is_ary(&value) {
        let sub = &value[1..(value.len()-1)];
        Value::Array(sub.split(",").map(|v| parse_value(String::from(v))).collect())
    } else {
        Value::String(value)
    }
}

<commit_msg>Fix: build_toml_header() should get the header object as mutable<commit_after>use std::collections::BTreeMap;
use std::path::PathBuf;
use std::str::Split;

use clap::ArgMatches;
use toml::Value;

use libimagstore::store::EntryHeader;
use libimagrt::runtime::Runtime;
use libimagutil::key_value_split::IntoKeyValue;

pub fn build_entry_path(rt: &Runtime, path_elem: &str) -> PathBuf {
    debug!("Building path...");
    let mut path = rt.store().path().clone();

    if path_elem.chars().next() == Some('/') {
        path.push(&path_elem[1..path_elem.len()]);
    } else {
        path.push(path_elem);
    }

    path
}

pub fn build_toml_header(matches: &ArgMatches, mut header: EntryHeader) -> EntryHeader {
    debug!("Building header from cli spec");
    if let Some(headerspecs) = matches.values_of("header") {
        let mut main = header.toml_mut();
        for tpl in headerspecs.into_iter().filter_map(|hs| String::from(hs).into_kv()) {
            let (key, value) = tpl.into();
            debug!("Splitting: {:?}", key);
            let mut split = key.split(".");
            let current = split.next();
            if current.is_some() {
                insert_key_into(String::from(current.unwrap()), &mut split, value, &mut main);
            }
        }
    }
    header
}

fn insert_key_into(current: String,
                   rest_path: &mut Split<&str>,
                   value: String,
                   map: &mut BTreeMap<String, Value>) {
    let next = rest_path.next();

    if next.is_none() {
        map.insert(current, parse_value(value));
    } else {
        if map.contains_key(&current) {
            match map.get_mut(&current).unwrap() {
                &mut Value::Table(ref mut t) => {
                    insert_key_into(String::from(next.unwrap()), rest_path, value, t);
                },
                _ => unreachable!(),
            }
        } else {
            let mut submap = BTreeMap::new();
            insert_key_into(String::from(next.unwrap()), rest_path, value, &mut submap);
            map.insert(current, Value::Table(submap));
        }
    }
}

fn parse_value(value: String) -> Value {
    fn is_ary(v: &String) -> bool {
        v.chars().next() == Some('[') && v.chars().last() == Some(']') && v.len() >= 3
    }

    debug!("Building value out of: {:?}", value);
    if value == "true" {
        Value::Boolean(true)
    } else if value == "false" {
        Value::Boolean(false)
    } else if is_ary(&value) {
        let sub = &value[1..(value.len()-1)];
        Value::Array(sub.split(",").map(|v| parse_value(String::from(v))).collect())
    } else {
        Value::String(value)
    }
}

<|endoftext|>


<commit_before>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::BPlusTree;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
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;

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_non_empty(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<KSA, PSA>(src_tree: &BPlusTree<KSA, PSA>, dest_tree: &LevelTree) -> usize
where
    KSA: Slice<EntryKey> + Debug + 'static,
    PSA: Slice<NodeCellRef> + 'static,
{
    let left_most_leaf_guards = select::<KSA, PSA>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    debug!("Merge selected {} pages", left_most_leaf_guards.len());

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

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

    merge_page_len
}
<commit_msg>review len sync after level merged<commit_after>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::BPlusTree;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
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_non_empty(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<KSA, PSA>(src_tree: &BPlusTree<KSA, PSA>, dest_tree: &LevelTree) -> usize
where
    KSA: Slice<EntryKey> + Debug + 'static,
    PSA: Slice<NodeCellRef> + 'static,
{
    let left_most_leaf_guards = select::<KSA, PSA>(&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()
            .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::<KSA, PSA>(&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();

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

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::codegen::InheritTypes::EventCast;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::error::Error::{InvalidAccess, Syntax};
use dom::bindings::global::{GlobalField, GlobalRef};
use dom::bindings::js::Root;
use dom::bindings::refcounted::Trusted;
use dom::bindings::str::USVString;
use dom::bindings::trace::JSTraceable;
use dom::bindings::utils::reflect_dom_object;
use dom::closeevent::CloseEvent;
use dom::event::{Event, EventBubbles, EventCancelable, EventHelpers};
use dom::eventtarget::{EventTarget, EventTargetHelpers, EventTargetTypeId};
use script_task::Runnable;
use script_task::ScriptMsg;
use std::cell::{Cell, RefCell};
use std::borrow::ToOwned;
use util::str::DOMString;

use websocket::Message;
use websocket::ws::sender::Sender as Sender_Object;
use websocket::client::sender::Sender;
use websocket::client::receiver::Receiver;
use websocket::stream::WebSocketStream;
use websocket::client::request::Url;
use websocket::Client;

#[derive(JSTraceable, PartialEq, Copy, Clone)]
enum WebSocketRequestState {
    Connecting = 0,
    Open = 1,
    Closing = 2,
    Closed = 3,
}

no_jsmanaged_fields!(Sender<WebSocketStream>);
no_jsmanaged_fields!(Receiver<WebSocketStream>);

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: Url,
    global: GlobalField,
    ready_state: Cell<WebSocketRequestState>,
    sender: RefCell<Option<Sender<WebSocketStream>>>,
    receiver: RefCell<Option<Receiver<WebSocketStream>>>,
    failed: Cell<bool>, //Flag to tell if websocket was closed due to failure
    full: Cell<bool>, //Flag to tell if websocket queue is full
    clean_close: Cell<bool>, //Flag to tell if the websocket closed cleanly (not due to full or fail)
    code: Cell<u16>, //Closing code
    reason: DOMRefCell<DOMString>, //Closing reason
    data: DOMRefCell<DOMString>, //Data from send - TODO: Remove after buffer is added.
    sendCloseFrame: Cell<bool>
}

fn parse_web_socket_url(url_str: &str) -> Fallible<(Url, String, u16, String, bool)> {
    // https://html.spec.whatwg.org/multipage/#parse-a-websocket-url's-components
    // Steps 1 and 2
    let parsed_url = Url::parse(url_str);
    let parsed_url = match parsed_url {
        Ok(parsed_url) => parsed_url,
        Err(_) => return Err(Error::Syntax),
    };

    // Step 4
    if parsed_url.fragment != None {
        return Err(Error::Syntax);
    }

    // Steps 3 and 5
    let secure = match parsed_url.scheme.as_ref() {
        "ws" => false,
        "wss" => true,
        _ => return Err(Error::Syntax), // step 3
    };

    let host = parsed_url.host().unwrap().serialize(); // Step 6
    let port = parsed_url.port_or_default().unwrap(); // Steps 7 and 8
    let mut resource = parsed_url.path().unwrap().connect("/"); // Step 9
    if resource.is_empty() {
        resource = "/".to_owned(); // Step 10
    }

    // Step 11
    if let Some(ref query) = parsed_url.query {
        resource.push('?');
        resource.push_str(query);
    }

    // Step 12
    // FIXME remove parsed_url once it's no longer used in WebSocket::new
    Ok((parsed_url, host, port, resource, secure))
}

impl WebSocket {
    pub fn new_inherited(global: GlobalRef, url: Url) -> WebSocket {
        WebSocket {
            eventtarget: EventTarget::new_inherited(EventTargetTypeId::WebSocket),
            url: url,
            global: GlobalField::from_rooted(&global),
            ready_state: Cell::new(WebSocketRequestState::Connecting),
            failed: Cell::new(false),
            sender: RefCell::new(None),
            receiver: RefCell::new(None),
            full: Cell::new(false),
            clean_close: Cell::new(true),
            code: Cell::new(0),
            reason: DOMRefCell::new("".to_owned()),
            data: DOMRefCell::new("".to_owned()),
            sendCloseFrame: Cell::new(false)
        }

    }

    pub fn new(global: GlobalRef, url: DOMString) -> Fallible<Root<WebSocket>> {
        // Step 1.
        // FIXME extract the right variables once Client::connect
        // implementation is fixed to follow the RFC 6455 properly.
        let (url, _, _, _, _) = try!(parse_web_socket_url(&url));

        /*TODO: This constructor is only a prototype, it does not accomplish the specs
          defined here:
          http://html.spec.whatwg.org
          The remaining 8 items must be satisfied.
          TODO: This constructor should be responsible for spawning a thread for the
          receive loop after ws.r().Open() - See comment
        */
        let ws = reflect_dom_object(box WebSocket::new_inherited(global, url.clone()),
                                    global,
                                    WebSocketBinding::Wrap);

        // TODO Client::connect does not conform to RFC 6455
        // see https://github.com/cyderize/rust-websocket/issues/38
        let request = match Client::connect(url) {
            Ok(request) => request,
            Err(_) => {
                let global_root = ws.r().global.root();
                let address = Trusted::new(global_root.r().get_cx(), ws.r(), global_root.r().script_chan().clone());
                let task = box WebSocketTaskHandler::new(address, WebSocketTask::Close);
                global_root.r().script_chan().send(ScriptMsg::RunnableMsg(task)).unwrap();
                return Ok(ws);
            }
        };
        let response = request.send().unwrap();
        response.validate().unwrap();
        ws.r().ready_state.set(WebSocketRequestState::Open);
        //Check to see if ready_state is Closing or Closed and failed = true - means we failed the websocket
        //if so return without setting any states
        let ready_state = ws.r().ready_state.get();
        let failed = ws.r().failed.get();
        if failed && (ready_state == WebSocketRequestState::Closed || ready_state == WebSocketRequestState::Closing) {
            //Do nothing else. Let the close finish.
            return Ok(ws);
        }

        let (temp_sender, temp_receiver) = response.begin().split();
        *ws.r().sender.borrow_mut() = Some(temp_sender);
        *ws.r().receiver.borrow_mut() = Some(temp_receiver);

        //Create everything necessary for starting the open asynchronous task, then begin the task.
        let global_root = ws.r().global.root();
        let addr: Trusted<WebSocket> =
            Trusted::new(global_root.r().get_cx(), ws.r(), global_root.r().script_chan().clone());
        let open_task = box WebSocketTaskHandler::new(addr.clone(), WebSocketTask::Open);
        global_root.r().script_chan().send(ScriptMsg::RunnableMsg(open_task)).unwrap();
        //TODO: Spawn thread here for receive loop
        /*TODO: Add receive loop here and make new thread run this
          Receive is an infinite loop "similiar" the one shown here:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: The receive loop however does need to follow the spec. These are outlined here
          under "WebSocket message has been received" items 1-5:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: The receive loop also needs to dispatch an asynchronous event as stated here:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: When the receive loop receives a close message from the server,
          it confirms the websocket is now closed. This requires the close event
          to be fired (dispatch_close fires the close event - see implementation below)
        */
        Ok(ws)
    }

    pub fn Constructor(global: GlobalRef, url: DOMString) -> Fallible<Root<WebSocket>> {
        WebSocket::new(global, url)
    }
}

impl<'a> WebSocketMethods for &'a WebSocket {
    event_handler!(open, GetOnopen, SetOnopen);
    event_handler!(close, GetOnclose, SetOnclose);
    event_handler!(error, GetOnerror, SetOnerror);

    fn Url(self) -> DOMString {
        self.url.serialize()
    }

    fn ReadyState(self) -> u16 {
        self.ready_state.get() as u16
    }

    fn Send(self, data: Option<USVString>)-> Fallible<()>{
        /*TODO: This is not up to spec see http://html.spec.whatwg.org/multipage/comms.html search for
                "If argument is a string"
          TODO: Need to buffer data
          TODO: bufferedAmount attribute returns the size of the buffer in bytes -
                this is a required attribute defined in the websocket.webidl file
          TODO: The send function needs to flag when full by using the following
          self.full.set(true). This needs to be done when the buffer is full
        */
        let mut other_sender = self.sender.borrow_mut();
        let my_sender = other_sender.as_mut().unwrap();
        if self.sendCloseFrame.get() { //TODO: Also check if the buffer is full
            self.sendCloseFrame.set(false);
            let _ = my_sender.send_message(Message::Close(None));
            return Ok(());
        }
        let _ = my_sender.send_message(Message::Text(data.unwrap().0));
        return Ok(())
    }

    fn Close(self, code: Option<u16>, reason: Option<USVString>) -> Fallible<()>{
        if let Some(code) = code {
            //Check code is NOT 1000 NOR in the range of 3000-4999 (inclusive)
            if  code != 1000 && (code < 3000 || code > 4999) {
                return Err(Error::InvalidAccess);
            }
        }
        if let Some(ref reason) = reason {
            if reason.0.as_bytes().len() > 123 { //reason cannot be larger than 123 bytes
                return Err(Error::Syntax);
            }
        }

        match self.ready_state.get() {
            WebSocketRequestState::Closing | WebSocketRequestState::Closed  => {} //Do nothing
            WebSocketRequestState::Connecting => { //Connection is not yet established
                /*By setting the state to closing, the open function
                  will abort connecting the websocket*/
                self.ready_state.set(WebSocketRequestState::Closing);
                self.failed.set(true);
                self.sendCloseFrame.set(true);
                //Dispatch send task to send close frame
                //TODO: Sending here is just empty string, though no string is really needed. Another send, empty
                //      send, could be used.
                let _ = self.Send(None);
                //Note: After sending the close message, the receive loop confirms a close message from the server and
                //      must fire a close event
            }
            WebSocketRequestState::Open => {
                //Closing handshake not started - still in open
                //Start the closing by setting the code and reason if they exist
                if let Some(code) = code {
                    self.code.set(code);
                }
                if let Some(reason) = reason {
                    *self.reason.borrow_mut() = reason.0;
                }
                self.ready_state.set(WebSocketRequestState::Closing);
                self.sendCloseFrame.set(true);
                //Dispatch send task to send close frame
                let _ = self.Send(None);
                //Note: After sending the close message, the receive loop confirms a close message from the server and
                //      must fire a close event
            }
        }
        Ok(()) //Return Ok
    }
}

pub enum WebSocketTask {
    Open,
    Close,
}

pub struct WebSocketTaskHandler {
    addr: Trusted<WebSocket>,
    task: WebSocketTask,
}

impl WebSocketTaskHandler {
    pub fn new(addr: Trusted<WebSocket>, task: WebSocketTask) -> WebSocketTaskHandler {
        WebSocketTaskHandler {
            addr: addr,
            task: task,
        }
    }

    fn dispatch_open(&self) {
        /*TODO: Items 1, 3, 4, & 5 under "WebSocket connection is established" as specified here:
          https://html.spec.whatwg.org/multipage/#feedback-from-the-protocol
        */
        let ws = self.addr.root(); //Get root
        let ws = ws.r(); //Get websocket reference
        let global = ws.global.root();
        let event = Event::new(global.r(),
            "open".to_owned(),
            EventBubbles::DoesNotBubble,
            EventCancelable::NotCancelable);
        let target = EventTargetCast::from_ref(ws);
        event.r().fire(target);
    }

    fn dispatch_close(&self) {
        let ws = self.addr.root();
        let ws = ws.r();
        let global = ws.global.root();
        ws.ready_state.set(WebSocketRequestState::Closed);
        //If failed or full, fire error event
        if ws.failed.get() || ws.full.get() {
            ws.failed.set(false);
            ws.full.set(false);
            //A Bad close
            ws.clean_close.set(false);
            let event = Event::new(global.r(),
                                   "error".to_owned(),
                                   EventBubbles::DoesNotBubble,
                                   EventCancelable::Cancelable);
            let target = EventTargetCast::from_ref(ws);
            event.r().fire(target);
        }
        let rsn = ws.reason.borrow();
        let rsn_clone = rsn.clone();
        /*In addition, we also have to fire a close even if error event fired
         https://html.spec.whatwg.org/multipage/#closeWebSocket
        */
        let close_event = CloseEvent::new(global.r(),
                                          "close".to_owned(),
                                          EventBubbles::DoesNotBubble,
                                          EventCancelable::NotCancelable,
                                          ws.clean_close.get(),
                                          ws.code.get(),
                                          rsn_clone);
        let target = EventTargetCast::from_ref(ws);
        let event = EventCast::from_ref(close_event.r());
        event.fire(target);
    }
}

impl Runnable for WebSocketTaskHandler {
    fn handler(self: Box<WebSocketTaskHandler>) {
        match self.task {
            WebSocketTask::Open => {
                self.dispatch_open();
            }
            WebSocketTask::Close => {
                self.dispatch_close();
            }
        }
    }
}

<commit_msg>Auto merge of #6610 - Ms2ger:task-readystate, r=jdm<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::cell::DOMRefCell;
use dom::bindings::codegen::Bindings::WebSocketBinding;
use dom::bindings::codegen::Bindings::WebSocketBinding::WebSocketMethods;
use dom::bindings::codegen::Bindings::EventHandlerBinding::EventHandlerNonNull;
use dom::bindings::codegen::InheritTypes::EventTargetCast;
use dom::bindings::codegen::InheritTypes::EventCast;
use dom::bindings::error::{Error, Fallible};
use dom::bindings::error::Error::{InvalidAccess, Syntax};
use dom::bindings::global::{GlobalField, GlobalRef};
use dom::bindings::js::Root;
use dom::bindings::refcounted::Trusted;
use dom::bindings::str::USVString;
use dom::bindings::trace::JSTraceable;
use dom::bindings::utils::reflect_dom_object;
use dom::closeevent::CloseEvent;
use dom::event::{Event, EventBubbles, EventCancelable, EventHelpers};
use dom::eventtarget::{EventTarget, EventTargetHelpers, EventTargetTypeId};
use script_task::Runnable;
use script_task::ScriptMsg;
use std::cell::{Cell, RefCell};
use std::borrow::ToOwned;
use util::str::DOMString;

use websocket::Message;
use websocket::ws::sender::Sender as Sender_Object;
use websocket::client::sender::Sender;
use websocket::client::receiver::Receiver;
use websocket::stream::WebSocketStream;
use websocket::client::request::Url;
use websocket::Client;

#[derive(JSTraceable, PartialEq, Copy, Clone)]
enum WebSocketRequestState {
    Connecting = 0,
    Open = 1,
    Closing = 2,
    Closed = 3,
}

no_jsmanaged_fields!(Sender<WebSocketStream>);
no_jsmanaged_fields!(Receiver<WebSocketStream>);

#[dom_struct]
pub struct WebSocket {
    eventtarget: EventTarget,
    url: Url,
    global: GlobalField,
    ready_state: Cell<WebSocketRequestState>,
    sender: RefCell<Option<Sender<WebSocketStream>>>,
    receiver: RefCell<Option<Receiver<WebSocketStream>>>,
    failed: Cell<bool>, //Flag to tell if websocket was closed due to failure
    full: Cell<bool>, //Flag to tell if websocket queue is full
    clean_close: Cell<bool>, //Flag to tell if the websocket closed cleanly (not due to full or fail)
    code: Cell<u16>, //Closing code
    reason: DOMRefCell<DOMString>, //Closing reason
    data: DOMRefCell<DOMString>, //Data from send - TODO: Remove after buffer is added.
    sendCloseFrame: Cell<bool>
}

fn parse_web_socket_url(url_str: &str) -> Fallible<(Url, String, u16, String, bool)> {
    // https://html.spec.whatwg.org/multipage/#parse-a-websocket-url's-components
    // Steps 1 and 2
    let parsed_url = Url::parse(url_str);
    let parsed_url = match parsed_url {
        Ok(parsed_url) => parsed_url,
        Err(_) => return Err(Error::Syntax),
    };

    // Step 4
    if parsed_url.fragment != None {
        return Err(Error::Syntax);
    }

    // Steps 3 and 5
    let secure = match parsed_url.scheme.as_ref() {
        "ws" => false,
        "wss" => true,
        _ => return Err(Error::Syntax), // step 3
    };

    let host = parsed_url.host().unwrap().serialize(); // Step 6
    let port = parsed_url.port_or_default().unwrap(); // Steps 7 and 8
    let mut resource = parsed_url.path().unwrap().connect("/"); // Step 9
    if resource.is_empty() {
        resource = "/".to_owned(); // Step 10
    }

    // Step 11
    if let Some(ref query) = parsed_url.query {
        resource.push('?');
        resource.push_str(query);
    }

    // Step 12
    // FIXME remove parsed_url once it's no longer used in WebSocket::new
    Ok((parsed_url, host, port, resource, secure))
}

impl WebSocket {
    pub fn new_inherited(global: GlobalRef, url: Url) -> WebSocket {
        WebSocket {
            eventtarget: EventTarget::new_inherited(EventTargetTypeId::WebSocket),
            url: url,
            global: GlobalField::from_rooted(&global),
            ready_state: Cell::new(WebSocketRequestState::Connecting),
            failed: Cell::new(false),
            sender: RefCell::new(None),
            receiver: RefCell::new(None),
            full: Cell::new(false),
            clean_close: Cell::new(true),
            code: Cell::new(0),
            reason: DOMRefCell::new("".to_owned()),
            data: DOMRefCell::new("".to_owned()),
            sendCloseFrame: Cell::new(false)
        }

    }

    pub fn new(global: GlobalRef, url: DOMString) -> Fallible<Root<WebSocket>> {
        // Step 1.
        // FIXME extract the right variables once Client::connect
        // implementation is fixed to follow the RFC 6455 properly.
        let (url, _, _, _, _) = try!(parse_web_socket_url(&url));

        /*TODO: This constructor is only a prototype, it does not accomplish the specs
          defined here:
          http://html.spec.whatwg.org
          The remaining 8 items must be satisfied.
          TODO: This constructor should be responsible for spawning a thread for the
          receive loop after ws.r().Open() - See comment
        */
        let ws = reflect_dom_object(box WebSocket::new_inherited(global, url.clone()),
                                    global,
                                    WebSocketBinding::Wrap);

        // TODO Client::connect does not conform to RFC 6455
        // see https://github.com/cyderize/rust-websocket/issues/38
        let request = match Client::connect(url) {
            Ok(request) => request,
            Err(_) => {
                let global_root = ws.r().global.root();
                let address = Trusted::new(global_root.r().get_cx(), ws.r(), global_root.r().script_chan().clone());
                let task = box WebSocketTaskHandler::new(address, WebSocketTask::Close);
                global_root.r().script_chan().send(ScriptMsg::RunnableMsg(task)).unwrap();
                return Ok(ws);
            }
        };
        let response = request.send().unwrap();
        response.validate().unwrap();

        let (temp_sender, temp_receiver) = response.begin().split();
        *ws.r().sender.borrow_mut() = Some(temp_sender);
        *ws.r().receiver.borrow_mut() = Some(temp_receiver);

        //Create everything necessary for starting the open asynchronous task, then begin the task.
        let global_root = ws.r().global.root();
        let addr: Trusted<WebSocket> =
            Trusted::new(global_root.r().get_cx(), ws.r(), global_root.r().script_chan().clone());
        let open_task = box WebSocketTaskHandler::new(addr, WebSocketTask::ConnectionEstablished);
        global_root.r().script_chan().send(ScriptMsg::RunnableMsg(open_task)).unwrap();
        //TODO: Spawn thread here for receive loop
        /*TODO: Add receive loop here and make new thread run this
          Receive is an infinite loop "similiar" the one shown here:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: The receive loop however does need to follow the spec. These are outlined here
          under "WebSocket message has been received" items 1-5:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: The receive loop also needs to dispatch an asynchronous event as stated here:
          https://github.com/cyderize/rust-websocket/blob/master/examples/client.rs#L64
          TODO: When the receive loop receives a close message from the server,
          it confirms the websocket is now closed. This requires the close event
          to be fired (dispatch_close fires the close event - see implementation below)
        */
        Ok(ws)
    }

    pub fn Constructor(global: GlobalRef, url: DOMString) -> Fallible<Root<WebSocket>> {
        WebSocket::new(global, url)
    }
}

impl<'a> WebSocketMethods for &'a WebSocket {
    event_handler!(open, GetOnopen, SetOnopen);
    event_handler!(close, GetOnclose, SetOnclose);
    event_handler!(error, GetOnerror, SetOnerror);

    fn Url(self) -> DOMString {
        self.url.serialize()
    }

    fn ReadyState(self) -> u16 {
        self.ready_state.get() as u16
    }

    fn Send(self, data: Option<USVString>)-> Fallible<()>{
        /*TODO: This is not up to spec see http://html.spec.whatwg.org/multipage/comms.html search for
                "If argument is a string"
          TODO: Need to buffer data
          TODO: bufferedAmount attribute returns the size of the buffer in bytes -
                this is a required attribute defined in the websocket.webidl file
          TODO: The send function needs to flag when full by using the following
          self.full.set(true). This needs to be done when the buffer is full
        */
        let mut other_sender = self.sender.borrow_mut();
        let my_sender = other_sender.as_mut().unwrap();
        if self.sendCloseFrame.get() { //TODO: Also check if the buffer is full
            self.sendCloseFrame.set(false);
            let _ = my_sender.send_message(Message::Close(None));
            return Ok(());
        }
        let _ = my_sender.send_message(Message::Text(data.unwrap().0));
        return Ok(())
    }

    fn Close(self, code: Option<u16>, reason: Option<USVString>) -> Fallible<()>{
        if let Some(code) = code {
            //Check code is NOT 1000 NOR in the range of 3000-4999 (inclusive)
            if  code != 1000 && (code < 3000 || code > 4999) {
                return Err(Error::InvalidAccess);
            }
        }
        if let Some(ref reason) = reason {
            if reason.0.as_bytes().len() > 123 { //reason cannot be larger than 123 bytes
                return Err(Error::Syntax);
            }
        }

        match self.ready_state.get() {
            WebSocketRequestState::Closing | WebSocketRequestState::Closed  => {} //Do nothing
            WebSocketRequestState::Connecting => { //Connection is not yet established
                /*By setting the state to closing, the open function
                  will abort connecting the websocket*/
                self.ready_state.set(WebSocketRequestState::Closing);
                self.failed.set(true);
                self.sendCloseFrame.set(true);
                //Dispatch send task to send close frame
                //TODO: Sending here is just empty string, though no string is really needed. Another send, empty
                //      send, could be used.
                let _ = self.Send(None);
                //Note: After sending the close message, the receive loop confirms a close message from the server and
                //      must fire a close event
            }
            WebSocketRequestState::Open => {
                //Closing handshake not started - still in open
                //Start the closing by setting the code and reason if they exist
                if let Some(code) = code {
                    self.code.set(code);
                }
                if let Some(reason) = reason {
                    *self.reason.borrow_mut() = reason.0;
                }
                self.ready_state.set(WebSocketRequestState::Closing);
                self.sendCloseFrame.set(true);
                //Dispatch send task to send close frame
                let _ = self.Send(None);
                //Note: After sending the close message, the receive loop confirms a close message from the server and
                //      must fire a close event
            }
        }
        Ok(()) //Return Ok
    }
}

pub enum WebSocketTask {
    /// Task queued when *the WebSocket connection is established*.
    ConnectionEstablished,
    Close,
}

pub struct WebSocketTaskHandler {
    addr: Trusted<WebSocket>,
    task: WebSocketTask,
}

impl WebSocketTaskHandler {
    pub fn new(addr: Trusted<WebSocket>, task: WebSocketTask) -> WebSocketTaskHandler {
        WebSocketTaskHandler {
            addr: addr,
            task: task,
        }
    }

    fn connection_established(&self) {
        /*TODO: Items 1, 3, 4, & 5 under "WebSocket connection is established" as specified here:
          https://html.spec.whatwg.org/multipage/#feedback-from-the-protocol
        */
        let ws = self.addr.root();

        // Step 1: Protocols.

        // Step 2.
        ws.ready_state.set(WebSocketRequestState::Open);

        // Step 3: Extensions.
        // Step 4: Protocols.
        // Step 5: Cookies.

        // Step 6.
        let global = ws.global.root();
        let event = Event::new(global.r(), "open".to_owned(),
                               EventBubbles::DoesNotBubble,
                               EventCancelable::NotCancelable);
        event.fire(EventTargetCast::from_ref(ws.r()));
    }

    fn dispatch_close(&self) {
        let ws = self.addr.root();
        let ws = ws.r();
        let global = ws.global.root();
        ws.ready_state.set(WebSocketRequestState::Closed);
        //If failed or full, fire error event
        if ws.failed.get() || ws.full.get() {
            ws.failed.set(false);
            ws.full.set(false);
            //A Bad close
            ws.clean_close.set(false);
            let event = Event::new(global.r(),
                                   "error".to_owned(),
                                   EventBubbles::DoesNotBubble,
                                   EventCancelable::Cancelable);
            let target = EventTargetCast::from_ref(ws);
            event.r().fire(target);
        }
        let rsn = ws.reason.borrow();
        let rsn_clone = rsn.clone();
        /*In addition, we also have to fire a close even if error event fired
         https://html.spec.whatwg.org/multipage/#closeWebSocket
        */
        let close_event = CloseEvent::new(global.r(),
                                          "close".to_owned(),
                                          EventBubbles::DoesNotBubble,
                                          EventCancelable::NotCancelable,
                                          ws.clean_close.get(),
                                          ws.code.get(),
                                          rsn_clone);
        let target = EventTargetCast::from_ref(ws);
        let event = EventCast::from_ref(close_event.r());
        event.fire(target);
    }
}

impl Runnable for WebSocketTaskHandler {
    fn handler(self: Box<WebSocketTaskHandler>) {
        match self.task {
            WebSocketTask::ConnectionEstablished => {
                self.connection_established();
            }
            WebSocketTask::Close => {
                self.dispatch_close();
            }
        }
    }
}

<|endoftext|>


<commit_before>/*!
  This module provides the syntax extension `iterator!()` that enables mapping, filtering and
  nesting of iterators using a comprehension syntax as well as the associated types `Prepend1`…
  `Prependn` which return the elements of an iterator in the last position of a 1…n element tuple,
  the preceding positions being constant.
*/

#![feature(macro_rules,phase)]

#![crate_id = "iteratorcomprehensions#0.1.0"]
#![crate_type = "lib"]

pub use prepend::{prepend1,Prepend1,prepend2,Prepend2};

/**
  Contains the macros that implement the comprehension syntax.
*/
pub mod macros {
  #![macro_escape]
  /**
    The `iterator!()` macro implements the following comprehension syntax:

    ```notrust
    iterator!(
      map_expr for var_1 in gen_expr_1 [if filter_expr_1]
      [… for var_n in gen_expr_n [if filter_expr]]
    )
    ```

    * `var_1`… `var_n` identify the iteration variables associated with each of the nested
      iterators.
    * `gen_expr_1`… `gen_expr_n` are expressions that evaluate to an `Iterator`. `gen_expr_i` can
      refer to all "outer" iteration variables `var_1`… `var_(i-1)`.
    * `map_expr` is an expression that constructs the elements of the iterator comprehension from
      the iteration variables.
    * `filter_expr_1`… `filter_expr_n` are expressions that evaluate to a boolean which filters
      the iterator elements based on the "outer" iteration variables.

    `iterator!()` evaluates to an expression which itself implements the `Iterator` trait.

    The expression

    ```notrust
    iterator!(
      (i, j) for i in range(0,3) for j in range(0, i + 1) if (i + j) % 2 == 0
    )
    ```
    evaluates to an iterator that contains the elements

    ```notrust
    (0,0), (1,1), (2,0), (2,2)
    ```
  */
  #[macro_export]
  macro_rules! iterator(
    (
      $map:expr
      for $var:ident in $gen:expr
      $(if $filter:expr)*
    ) => (
      $gen
      $(.filter(|&$var| { $filter }))*
      .map(|$var| { $map })
    );
    (
      $map:expr
      for $var1:ident in $gen1:expr
      $(if $filter1:expr)*
      for $var2:ident in $gen2:expr
      $(if $filter2:expr)*
    ) => (
      $gen1
      $(.filter(|&$var1| { $filter1 }))*
      .flat_map(|$var1| { iteratorcomprehensions::prepend1($var1, $gen2) })
      $(.filter(|&($var1, $var2)| { $filter2 }))*
      .map(|($var1, $var2)| { $map })
    );
    (
      $map:expr
      for $var1:ident in $gen1:expr
      $(if $filter1:expr)*
      for $var2:ident in $gen2:expr
      $(if $filter2:expr)*
      for $var3:ident in $gen3:expr
      $(if $filter3:expr)*
    ) => (
      $gen1
      $(.filter(|&$var1| { $filter1 }))*
      .flat_map(|$var1| { iteratorcomprehensions::prepend1($var1, $gen2) })
      $(.filter(|&($var1, $var2)| { $filter2 }))*
      .flat_map(|($var1, $var2)| { iteratorcomprehensions::prepend2($var1, $var2, $gen3) })
      $(.filter(|&($var1, $var2, $var3)| { $filter3 }))*
      .map(|($var1, $var2, $var3)| { $map })
    );
  )
}

/**
  Contains the types `Prepend1`… `Prependn` as well as the convenience functions `prepend1`… `prependn`.
*/
pub mod prepend {
  /**
    Maps each element `i` in `iter` to the tuple `(val1, i)`.
  */
  pub struct Prepend1<T1, T, It> {
    iter: It,
    val1: T1
  }

  impl<T1: Clone, T, It: Iterator<T>> Prepend1<T1, T, It> {
    #[inline]
    fn new(v1: T1, it: It) -> Prepend1<T1, T, It> {
      Prepend1 { iter: it, val1: v1 }
    }
  }

  impl<T1: Clone, T, It: Iterator<T>> Iterator<(T1, T)> for Prepend1<T1, T, It> {
    #[inline]
    fn next(&mut self) -> Option<(T1, T)> {
      self.iter.next().map(|val| (self.val1.clone(), val))
    }
  }

  /**
    Constructs a new `Prepend1`.
  */
  #[inline]
  pub fn prepend1<T1: Clone, T, It: Iterator<T>>(v1: T1, it: It) -> Prepend1<T1, T, It> {
    Prepend1::new(v1, it)
  }

  /**
    Maps each element `i` in `iter` to the tuple `(val1, val2, i)`.
  */
  pub struct Prepend2<T1, T2, T, It> {
    iter: It,
    val1: T1,
    val2: T2
  }

  impl<T1: Clone, T2: Clone, T, It: Iterator<T>> Prepend2<T1, T2, T, It> {
    #[inline]
    fn new(v1: T1, v2: T2, it: It) -> Prepend2<T1, T2, T, It> {
      Prepend2 { iter: it, val1: v1, val2: v2 }
    }
  }

  impl<T1: Clone, T2: Clone, T, It: Iterator<T>> Iterator<(T1, T2, T)> for Prepend2<T1, T2, T, It> {
    #[inline]
    fn next(&mut self) -> Option<(T1, T2, T)> {
      self.iter.next().map(|val| (self.val1.clone(), self.val2.clone(), val))
    }
  }

  /**
    Constructs a new `Prepend2`.
  */
  #[inline]
  pub fn prepend2<T1: Clone, T2: Clone, T, It: Iterator<T>>(v1: T1, v2: T2, it: It)
  -> Prepend2<T1, T2, T, It> {
    Prepend2::new(v1, v2, it)
  }
}

#[cfg(test)]
mod tests {
  use iteratorcomprehensions = prepend; // Yikes!

  #[test]
  fn iterator1_test() {
    let xs: Vec<int> = iterator!( i for i in range(0i, 3i) ).collect();
    assert_eq!(xs, vec!(0, 1, 2));
  }

  #[test]
  fn iterator1_map_test() {
    let xs: Vec<int> = iterator!( i + 1 for i in range(0i, 3i) ).collect();
    assert_eq!(xs, vec!(1, 2, 3));
  }

  #[test]
  fn iterator1_filter_test() {
    let xs: Vec<int> = iterator!( i for i in range(0i, 3i) if i % 2 == 1 ).collect();
    assert_eq!(xs, vec!(1));
  }

  #[test]
  fn iterator1_filter_map_test() {
    let xs: Vec<int> = iterator!( i * 2 for i in range(0i, 3i) if i % 2 == 1 ).collect();
    assert_eq!(xs, vec!(2));
  }

  #[test]
  fn iterator2_map_test() {
    let xs: Vec<int> = iterator!( i * j for i in range(1i, 3i) for j in range(2i, 4i) ).collect();
    assert_eq!(xs, vec!(2, 3, 4, 6));
  }

  #[test]
  fn iterator2_filter_map_test() {
    let xs: Vec<int> = iterator!(
      i / j for i in range(6i, 9i) for j in range(1i, 4i) if i % j == 0
    ).collect();
    assert_eq!(xs, vec!(6, 3, 2, 7, 8, 4));
  }

  #[test]
  fn iterator2_example_test() {
    let xs: Vec<(int,int)> = iterator!(
      (i,j) for i in range(0i, 3i) for j in range(0i, i + 1i) if (i + j) % 2 == 0
    ).collect();
    assert_eq!(xs, vec!((0, 0), (1, 1), (2, 0), (2, 2)));
  }

  #[test]
  fn iterator3_map_test() {
    let xs: Vec<int> = iterator!(
      i * j * k for i in range(1i, 3i) for j in range(2i, 4i) for k in range(3i, 5i)
    ).collect();
    assert_eq!(xs, vec!(6, 8, 9, 12, 12, 16, 18, 24));
  }

  #[test]
  fn iterator3_filter_map_test() {
    let xs: Vec<int> = iterator!(
      i + j + k for i in range(0i, 10i) for j in range(0i, 10i) for k in range(0i, 10i)
      if i == 1 && j == 1 && k == 1
    ).collect();
    assert_eq!(xs, vec!(3));
  }

  #[test]
  fn iterator3_empty_test() {
    let xs: Vec<int> = iterator!(
      i * j * k for i in range(0i, 2i) for j in range(0i, i) for k in range(0i, 1i)
    ).collect();
    assert_eq!(xs, vec!(0));
  }
}
<commit_msg>fix example in docs.<commit_after>/*!
  This module provides the syntax extension `iterator!()` that enables mapping, filtering and
  nesting of iterators using a comprehension syntax as well as the associated types `Prepend1`…
  `Prependn` which return the elements of an iterator in the last position of a 1…n element tuple,
  the preceding positions being constant.
*/

#![feature(macro_rules,phase)]

#![crate_id = "iteratorcomprehensions#0.1.0"]
#![crate_type = "lib"]

pub use prepend::{prepend1,Prepend1,prepend2,Prepend2};

/**
  Contains the macros that implement the comprehension syntax.
*/
pub mod macros {
  #![macro_escape]
  /**
    The `iterator!()` macro implements the following comprehension syntax:

    ```notrust
    iterator!(
      map_expr for var_1 in gen_expr_1 [if filter_expr_1]
      [… for var_n in gen_expr_n [if filter_expr]]
    )
    ```

    * `var_1`… `var_n` identify the iteration variables associated with each of the nested
      iterators.
    * `gen_expr_1`… `gen_expr_n` are expressions that evaluate to an `Iterator`. `gen_expr_i` can
      refer to all "outer" iteration variables `var_1`… `var_(i-1)`.
    * `map_expr` is an expression that constructs the elements of the iterator comprehension from
      the iteration variables.
    * `filter_expr_1`… `filter_expr_n` are expressions that evaluate to a boolean which filters
      the iterator elements based on the "outer" iteration variables.

    `iterator!()` evaluates to an expression which itself implements the `Iterator` trait.

    The expression

    ```notrust
    iterator!(
      (i, j) for i in range(0i, 3i) for j in range(0i, i + 1i) if (i + j) % 2 == 0
    )
    ```
    evaluates to an iterator that contains the elements

    ```notrust
    (0,0), (1,1), (2,0), (2,2)
    ```
  */
  #[macro_export]
  macro_rules! iterator(
    (
      $map:expr
      for $var:ident in $gen:expr
      $(if $filter:expr)*
    ) => (
      $gen
      $(.filter(|&$var| { $filter }))*
      .map(|$var| { $map })
    );
    (
      $map:expr
      for $var1:ident in $gen1:expr
      $(if $filter1:expr)*
      for $var2:ident in $gen2:expr
      $(if $filter2:expr)*
    ) => (
      $gen1
      $(.filter(|&$var1| { $filter1 }))*
      .flat_map(|$var1| { iteratorcomprehensions::prepend1($var1, $gen2) })
      $(.filter(|&($var1, $var2)| { $filter2 }))*
      .map(|($var1, $var2)| { $map })
    );
    (
      $map:expr
      for $var1:ident in $gen1:expr
      $(if $filter1:expr)*
      for $var2:ident in $gen2:expr
      $(if $filter2:expr)*
      for $var3:ident in $gen3:expr
      $(if $filter3:expr)*
    ) => (
      $gen1
      $(.filter(|&$var1| { $filter1 }))*
      .flat_map(|$var1| { iteratorcomprehensions::prepend1($var1, $gen2) })
      $(.filter(|&($var1, $var2)| { $filter2 }))*
      .flat_map(|($var1, $var2)| { iteratorcomprehensions::prepend2($var1, $var2, $gen3) })
      $(.filter(|&($var1, $var2, $var3)| { $filter3 }))*
      .map(|($var1, $var2, $var3)| { $map })
    );
  )
}

/**
  Contains the types `Prepend1`… `Prependn` as well as the convenience functions `prepend1`… `prependn`.
*/
pub mod prepend {
  /**
    Maps each element `i` in `iter` to the tuple `(val1, i)`.
  */
  pub struct Prepend1<T1, T, It> {
    iter: It,
    val1: T1
  }

  impl<T1: Clone, T, It: Iterator<T>> Prepend1<T1, T, It> {
    #[inline]
    fn new(v1: T1, it: It) -> Prepend1<T1, T, It> {
      Prepend1 { iter: it, val1: v1 }
    }
  }

  impl<T1: Clone, T, It: Iterator<T>> Iterator<(T1, T)> for Prepend1<T1, T, It> {
    #[inline]
    fn next(&mut self) -> Option<(T1, T)> {
      self.iter.next().map(|val| (self.val1.clone(), val))
    }
  }

  /**
    Constructs a new `Prepend1`.
  */
  #[inline]
  pub fn prepend1<T1: Clone, T, It: Iterator<T>>(v1: T1, it: It) -> Prepend1<T1, T, It> {
    Prepend1::new(v1, it)
  }

  /**
    Maps each element `i` in `iter` to the tuple `(val1, val2, i)`.
  */
  pub struct Prepend2<T1, T2, T, It> {
    iter: It,
    val1: T1,
    val2: T2
  }

  impl<T1: Clone, T2: Clone, T, It: Iterator<T>> Prepend2<T1, T2, T, It> {
    #[inline]
    fn new(v1: T1, v2: T2, it: It) -> Prepend2<T1, T2, T, It> {
      Prepend2 { iter: it, val1: v1, val2: v2 }
    }
  }

  impl<T1: Clone, T2: Clone, T, It: Iterator<T>> Iterator<(T1, T2, T)> for Prepend2<T1, T2, T, It> {
    #[inline]
    fn next(&mut self) -> Option<(T1, T2, T)> {
      self.iter.next().map(|val| (self.val1.clone(), self.val2.clone(), val))
    }
  }

  /**
    Constructs a new `Prepend2`.
  */
  #[inline]
  pub fn prepend2<T1: Clone, T2: Clone, T, It: Iterator<T>>(v1: T1, v2: T2, it: It)
  -> Prepend2<T1, T2, T, It> {
    Prepend2::new(v1, v2, it)
  }
}

#[cfg(test)]
mod tests {
  use iteratorcomprehensions = prepend; // Yikes!

  #[test]
  fn iterator1_test() {
    let xs: Vec<int> = iterator!( i for i in range(0i, 3i) ).collect();
    assert_eq!(xs, vec!(0, 1, 2));
  }

  #[test]
  fn iterator1_map_test() {
    let xs: Vec<int> = iterator!( i + 1 for i in range(0i, 3i) ).collect();
    assert_eq!(xs, vec!(1, 2, 3));
  }

  #[test]
  fn iterator1_filter_test() {
    let xs: Vec<int> = iterator!( i for i in range(0i, 3i) if i % 2 == 1 ).collect();
    assert_eq!(xs, vec!(1));
  }

  #[test]
  fn iterator1_filter_map_test() {
    let xs: Vec<int> = iterator!( i * 2 for i in range(0i, 3i) if i % 2 == 1 ).collect();
    assert_eq!(xs, vec!(2));
  }

  #[test]
  fn iterator2_map_test() {
    let xs: Vec<int> = iterator!( i * j for i in range(1i, 3i) for j in range(2i, 4i) ).collect();
    assert_eq!(xs, vec!(2, 3, 4, 6));
  }

  #[test]
  fn iterator2_filter_map_test() {
    let xs: Vec<int> = iterator!(
      i / j for i in range(6i, 9i) for j in range(1i, 4i) if i % j == 0
    ).collect();
    assert_eq!(xs, vec!(6, 3, 2, 7, 8, 4));
  }

  #[test]
  fn iterator2_example_test() {
    let xs: Vec<(int,int)> = iterator!(
      (i,j) for i in range(0i, 3i) for j in range(0i, i + 1i) if (i + j) % 2 == 0
    ).collect();
    assert_eq!(xs, vec!((0, 0), (1, 1), (2, 0), (2, 2)));
  }

  #[test]
  fn iterator3_map_test() {
    let xs: Vec<int> = iterator!(
      i * j * k for i in range(1i, 3i) for j in range(2i, 4i) for k in range(3i, 5i)
    ).collect();
    assert_eq!(xs, vec!(6, 8, 9, 12, 12, 16, 18, 24));
  }

  #[test]
  fn iterator3_filter_map_test() {
    let xs: Vec<int> = iterator!(
      i + j + k for i in range(0i, 10i) for j in range(0i, 10i) for k in range(0i, 10i)
      if i == 1 && j == 1 && k == 1
    ).collect();
    assert_eq!(xs, vec!(3));
  }

  #[test]
  fn iterator3_empty_test() {
    let xs: Vec<int> = iterator!(
      i * j * k for i in range(0i, 2i) for j in range(0i, i) for k in range(0i, 1i)
    ).collect();
    assert_eq!(xs, vec!(0));
  }
}
<|endoftext|>