<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 std::borrow::Cow;
use std::cell::RefCell;
use std::fmt::Display;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Arc;

use dbus;
use dbus::Connection;
use dbus::Message;
use dbus::MessageItem;
use dbus::NameFlag;

use dbus::tree::Factory;
use dbus::tree::MethodErr;
use dbus::tree::MethodFn;
use dbus::tree::MethodResult;
use dbus::tree::Property;
use dbus::tree::Tree;

use dbus_consts::*;

use engine::Engine;
use pool::StratisPool;
use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

fn list_pools(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let result = engine.borrow().list_pools();

    let mut msg_vec = Vec::new();

    // TODO: deal with failure
    let pool_tree = result.unwrap();

    for (name, pool) in pool_tree {
        let entry = vec![MessageItem::Str(format!("{}", name)),
                         MessageItem::UInt16(0),
                         MessageItem::Str(String::from("Ok"))];
        msg_vec.push(MessageItem::Struct(entry));
    }

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));

    Ok(vec![m.method_return().append1(item_array)])
}

fn create_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let raid_level: u16 = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner()
                .map_err(|_| MethodErr::invalid_arg(&i))
        }));

    let devs = match try!(items.pop().ok_or_else(MethodErr::no_arg)) {
        MessageItem::Array(x, _) => x,
        x => return Err(MethodErr::invalid_arg(&x)),
    };

    // Get the name of the pool from the parameters
    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    let blockdevs = devs.into_iter()
        .map(|x| PathBuf::from(x.inner::<&str>().unwrap()))
        .collect::<Vec<_>>();

    let result = engine.borrow_mut().create_pool(&name, &blockdevs, raid_level);

    Ok(vec![m.method_return().append3("/dbus/newpool/path", 0, "Ok")])
}

fn destroy_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let item: MessageItem = try!(m.get1().ok_or_else(MethodErr::no_arg));

    let name: &String = try!(item.inner().map_err(|_| MethodErr::invalid_arg(&item)));

    let result = engine.borrow_mut().destroy_pool(&name);

    let message = m.method_return();

    let msg = match result {
        Ok(_) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
        Err(err) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
    };

    Ok(vec![msg])

}

fn get_pool_object_path(m: &Message) -> MethodResult {

    Ok(vec![m.method_return().append3("/dbus/pool/path", 0, "Ok")])
}

fn get_volume_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/volume/path", 0, "Ok")])
}

fn get_dev_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/dev/path", 0, "Ok")])
}

fn get_cache_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/cache/path", 0, "Ok")])
}

fn get_list_items<T, I>(m: &Message, iter: I) -> MethodResult
    where T: HasCodes + Display,
          I: Iterator<Item = T>
{
    let msg_vec = iter.map(|item| {
            MessageItem::Struct(vec![MessageItem::Str(format!("{}", item)),
                                     MessageItem::UInt16(item.get_error_int()),
                                     MessageItem::Str(format!("{}", item.get_error_string()))])
        })
        .collect::<Vec<MessageItem>>();

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));
    Ok(vec![m.method_return().append1(item_array)])
}

fn get_error_codes(m: &Message) -> MethodResult {
    get_list_items(m, StratisErrorEnum::iter_variants())
}

fn get_raid_levels(m: &Message) -> MethodResult {
    get_list_items(m, StratisRaidType::iter_variants())
}

fn get_dev_types(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let _name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

pub fn get_base_tree<'a>(c: &'a Connection,
                         engine: Rc<RefCell<Engine>>)
                         -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let engine_clone = engine.clone();

    let createpool_method = f.method(CREATE_POOL, move |m, _, _| create_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .in_arg(("dev_list", "as"))
        .in_arg(("raid_type", "q"))
        .out_arg(("object_path", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| destroy_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| list_pools(m, &engine_clone))
        .out_arg(("pool_names", "as"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getpoolobjectpath_method =
        f.method(GET_POOL_OBJECT_PATH, move |m, _, _| get_pool_object_path(m))
            .in_arg(("pool_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH,
                move |m, _, _| get_volume_object_path(m))
        .in_arg(("pool_name", "s"))
        .in_arg(("volume_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getdevobjectpath_method =
        f.method(GET_DEV_OBJECT_PATH, move |m, _, _| get_dev_object_path(m))
            .in_arg(("dev_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH,
                move |m, _, _| get_cache_object_path(m))
        .in_arg(("cache_dev_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| get_error_codes(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| get_raid_levels(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| get_dev_types(m));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_MANAGER_INTERFACE)
            .add_m(listpools_method)
            .add_m(createpool_method)
            .add_m(destroypool_method)
            .add_m(getpoolobjectpath_method)
            .add_m(getvolumeobjectpath_method)
            .add_m(getdevobjectpath_method)
            .add_m(getcacheobjectpath_method)
            .add_m(geterrorcodes_method)
            .add_m(getraidlevels_method)
            .add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<commit_msg>Remove unused import.<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 std::borrow::Cow;
use std::cell::RefCell;
use std::fmt::Display;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Arc;

use dbus;
use dbus::Connection;
use dbus::Message;
use dbus::MessageItem;
use dbus::NameFlag;

use dbus::tree::Factory;
use dbus::tree::MethodErr;
use dbus::tree::MethodFn;
use dbus::tree::MethodResult;
use dbus::tree::Property;
use dbus::tree::Tree;

use dbus_consts::*;

use engine::Engine;
use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

fn list_pools(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let result = engine.borrow().list_pools();

    let mut msg_vec = Vec::new();

    // TODO: deal with failure
    let pool_tree = result.unwrap();

    for (name, pool) in pool_tree {
        let entry = vec![MessageItem::Str(format!("{}", name)),
                         MessageItem::UInt16(0),
                         MessageItem::Str(String::from("Ok"))];
        msg_vec.push(MessageItem::Struct(entry));
    }

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));

    Ok(vec![m.method_return().append1(item_array)])
}

fn create_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let raid_level: u16 = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner()
                .map_err(|_| MethodErr::invalid_arg(&i))
        }));

    let devs = match try!(items.pop().ok_or_else(MethodErr::no_arg)) {
        MessageItem::Array(x, _) => x,
        x => return Err(MethodErr::invalid_arg(&x)),
    };

    // Get the name of the pool from the parameters
    let name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    let blockdevs = devs.into_iter()
        .map(|x| PathBuf::from(x.inner::<&str>().unwrap()))
        .collect::<Vec<_>>();

    let result = engine.borrow_mut().create_pool(&name, &blockdevs, raid_level);

    Ok(vec![m.method_return().append3("/dbus/newpool/path", 0, "Ok")])
}

fn destroy_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let item: MessageItem = try!(m.get1().ok_or_else(MethodErr::no_arg));

    let name: &String = try!(item.inner().map_err(|_| MethodErr::invalid_arg(&item)));

    let result = engine.borrow_mut().destroy_pool(&name);

    let message = m.method_return();

    let msg = match result {
        Ok(_) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
        Err(err) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
    };

    Ok(vec![msg])

}

fn get_pool_object_path(m: &Message) -> MethodResult {

    Ok(vec![m.method_return().append3("/dbus/pool/path", 0, "Ok")])
}

fn get_volume_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/volume/path", 0, "Ok")])
}

fn get_dev_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/dev/path", 0, "Ok")])
}

fn get_cache_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/cache/path", 0, "Ok")])
}

fn get_list_items<T, I>(m: &Message, iter: I) -> MethodResult
    where T: HasCodes + Display,
          I: Iterator<Item = T>
{
    let msg_vec = iter.map(|item| {
            MessageItem::Struct(vec![MessageItem::Str(format!("{}", item)),
                                     MessageItem::UInt16(item.get_error_int()),
                                     MessageItem::Str(format!("{}", item.get_error_string()))])
        })
        .collect::<Vec<MessageItem>>();

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));
    Ok(vec![m.method_return().append1(item_array)])
}

fn get_error_codes(m: &Message) -> MethodResult {
    get_list_items(m, StratisErrorEnum::iter_variants())
}

fn get_raid_levels(m: &Message) -> MethodResult {
    get_list_items(m, StratisRaidType::iter_variants())
}

fn get_dev_types(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let _name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

pub fn get_base_tree<'a>(c: &'a Connection,
                         engine: Rc<RefCell<Engine>>)
                         -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let engine_clone = engine.clone();

    let createpool_method = f.method(CREATE_POOL, move |m, _, _| create_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .in_arg(("dev_list", "as"))
        .in_arg(("raid_type", "q"))
        .out_arg(("object_path", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| destroy_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| list_pools(m, &engine_clone))
        .out_arg(("pool_names", "as"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getpoolobjectpath_method =
        f.method(GET_POOL_OBJECT_PATH, move |m, _, _| get_pool_object_path(m))
            .in_arg(("pool_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH,
                move |m, _, _| get_volume_object_path(m))
        .in_arg(("pool_name", "s"))
        .in_arg(("volume_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getdevobjectpath_method =
        f.method(GET_DEV_OBJECT_PATH, move |m, _, _| get_dev_object_path(m))
            .in_arg(("dev_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH,
                move |m, _, _| get_cache_object_path(m))
        .in_arg(("cache_dev_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| get_error_codes(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| get_raid_levels(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| get_dev_types(m));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_MANAGER_INTERFACE)
            .add_m(listpools_method)
            .add_m(createpool_method)
            .add_m(destroypool_method)
            .add_m(getpoolobjectpath_method)
            .add_m(getvolumeobjectpath_method)
            .add_m(getdevobjectpath_method)
            .add_m(getcacheobjectpath_method)
            .add_m(geterrorcodes_method)
            .add_m(getraidlevels_method)
            .add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<|endoftext|>


<commit_before>use std::kinds::marker::{ContravariantLifetime, NoCopy};

use objc::runtime::Object;
use objc::{Id, IdVector, IntoIdVector, Owned, Ownership, Shared, ShareId};

use {class, INSCopying, INSMutableCopying, INSObject};

#[repr(C)]
pub enum NSComparisonResult {
    Ascending  = -1i,
    Same       = 0i,
    Descending = 1i,
}

impl NSComparisonResult {
    pub fn from_ordering(order: Ordering) -> NSComparisonResult {
        match order {
            Less => NSComparisonResult::Ascending,
            Equal => NSComparisonResult::Same,
            Greater => NSComparisonResult::Descending,
        }
    }

    pub fn as_ordering(&self) -> Ordering {
        match *self {
            NSComparisonResult::Ascending => Less,
            NSComparisonResult::Same => Equal,
            NSComparisonResult::Descending => Greater,
        }
    }
}

pub struct NSRange {
    pub location: uint,
    pub length: uint,
}

pub struct NSEnumerator<'a, T> {
    id: Id<Object>,
    marker: ContravariantLifetime<'a>,
}

impl<'a, T> NSEnumerator<'a, T> {
    pub unsafe fn from_ptr(ptr: *mut Object) -> NSEnumerator<'a, T> {
        NSEnumerator { id: Id::from_ptr(ptr), marker: ContravariantLifetime }
    }
}

impl<'a, T> Iterator<&'a T> for NSEnumerator<'a, T> {
    fn next(&mut self) -> Option<&'a T> {
        unsafe {
            let obj = msg_send![self.id nextObject] as *mut T;
            obj.as_ref()
        }
    }
}

pub trait INSArray<T: INSObject, O: Ownership> : INSObject {
    fn count(&self) -> uint {
        let result = unsafe {
            msg_send![self count]
        };
        result as uint
    }

    fn object_at(&self, index: uint) -> &T {
        unsafe {
            let obj = msg_send![self objectAtIndex:index] as *const T;
            &*obj
        }
    }

    fn first_object(&self) -> Option<&T> {
        unsafe {
            let obj = msg_send![self firstObject] as *const T;
            obj.as_ref()
        }
    }

    fn last_object(&self) -> Option<&T> {
        unsafe {
            let obj = msg_send![self lastObject] as *const T;
            obj.as_ref()
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<T> {
        unsafe {
            let result = msg_send![self objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    unsafe fn from_refs(refs: &[&T]) -> Id<Self> {
        let cls = class::<Self>();
        let obj = msg_send![cls alloc];
        let obj = msg_send![obj initWithObjects:refs.as_ptr() count:refs.len()];
        Id::from_retained_ptr(obj as *mut Self)
    }

    fn from_vec(vec: Vec<Id<T, O>>) -> Id<Self> {
        let refs = vec.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn objects_in_range(&self, start: uint, len: uint) -> Vec<&T> {
        let vec: Vec<*mut T> = Vec::from_elem(len, RawPtr::null());
        let range = NSRange { location: start, length: len };
        unsafe {
            msg_send![self getObjects:vec.as_ptr() range:range];
            vec.map_in_place(|ptr| &*ptr)
        }
    }

    fn to_vec(&self) -> Vec<&T> {
        self.objects_in_range(0, self.count())
    }

    fn into_vec(array: Id<Self>) -> Vec<Id<T, O>> {
        let vec = array.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub trait INSOwnedArray<T: INSObject> : INSArray<T, Owned> {
    fn mut_object_at(&mut self, index: uint) -> &mut T {
        unsafe {
            let result = msg_send![self objectAtIndex:index] as *mut T;
            &mut *result
        }
    }
}

pub trait INSSharedArray<T: INSObject> : INSArray<T, Shared> {
    fn shared_object_at(&self, index: uint) -> ShareId<T> {
        let obj = self.object_at(index);
        unsafe {
            Id::from_ptr(obj as *const _ as *mut T)
        }
    }

    fn from_slice(slice: &[ShareId<T>]) -> Id<Self> {
        let refs = slice.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn to_shared_vec(&self) -> Vec<ShareId<T>> {
        let vec = self.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub struct NSArray<T, O = Owned> {
    nocopy: NoCopy,
}

object_impl!(NSArray<T, O>)

impl<T: INSObject, O: Ownership> INSArray<T, O> for NSArray<T, O> { }

impl<T: INSObject> INSOwnedArray<T> for NSArray<T, Owned> { }

impl<T: INSObject> INSSharedArray<T> for NSArray<T, Shared> { }

impl<T> INSCopying<NSSharedArray<T>> for NSArray<T, Shared> { }

impl<T> INSMutableCopying<NSMutableSharedArray<T>> for NSArray<T, Shared> { }

impl<T: INSObject, O: Ownership> Index<uint, T> for NSArray<T, O> {
    fn index(&self, index: &uint) -> &T {
        self.object_at(*index)
    }
}

pub type NSSharedArray<T> = NSArray<T, Shared>;

pub trait INSMutableArray<T: INSObject, O: Ownership> : INSArray<T, O> {
    fn add_object(&mut self, obj: Id<T, O>) {
        unsafe {
            msg_send![self addObject:obj];
        }
    }

    fn insert_object_at(&mut self, index: uint, obj: Id<T, O>) {
        unsafe {
            msg_send![self insertObject:obj atIndex:index];
        }
    }

    fn replace_object_at(&mut self, index: uint, obj: Id<T, O>) -> Id<T, O> {
        let old_obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut T)
        };
        unsafe {
            msg_send![self replaceObjectAtIndex:index withObject:obj];
        }
        old_obj
    }

    fn remove_object_at(&mut self, index: uint) -> Id<T, O> {
        let obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut T)
        };
        unsafe {
            msg_send![self removeObjectAtIndex:index];
        }
        obj
    }

    fn remove_last_object(&mut self) -> Id<T, O> {
        let obj = self.last_object().map(|obj| unsafe {
            Id::from_ptr(obj as *const _ as *mut T)
        });
        unsafe {
            msg_send![self removeLastObject];
        }
        // removeLastObject would have failed if the array is empty,
        // so we know this won't be None
        obj.unwrap()
    }

    fn remove_all_objects(&mut self) {
        unsafe {
            msg_send![self removeAllObjects];
        }
    }

    fn sort_with_context<C>(&mut self,
            compare: extern fn(&T, &T, Option<&C>) -> NSComparisonResult,
            context: Option<&C>) {
        unsafe {
            msg_send![self sortUsingFunction:compare context:context];
        }
    }
}

pub struct NSMutableArray<T, O = Owned> {
    nocopy: NoCopy,
}

object_impl!(NSMutableArray<T, O>)

impl<T: INSObject, O: Ownership> INSArray<T, O> for NSMutableArray<T, O> { }

impl<T: INSObject> INSOwnedArray<T> for NSMutableArray<T, Owned> { }

impl<T: INSObject> INSSharedArray<T> for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> INSMutableArray<T, O> for NSMutableArray<T, O> { }

impl<T> INSCopying<NSSharedArray<T>> for NSMutableArray<T, Shared> { }

impl<T> INSMutableCopying<NSMutableSharedArray<T>> for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> Index<uint, T> for NSMutableArray<T, O> {
    fn index(&self, index: &uint) -> &T {
        self.object_at(*index)
    }
}

pub type NSMutableSharedArray<T> = NSMutableArray<T, Shared>;

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSObject, INSString, NSObject, NSString};
    use super::{
        INSArray, INSMutableArray,
        NSArray, NSMutableArray, NSComparisonResult
    };

    fn sample_array(len: uint) -> Id<NSArray<NSObject>> {
        let vec: Vec<Id<NSObject>> = Vec::from_fn(len, |_| INSObject::new());
        INSArray::from_vec(vec)
    }

    #[test]
    fn test_count() {
        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.count() == 0);

        let array = sample_array(4);
        assert!(array.count() == 4);
    }

    #[test]
    fn test_object_at() {
        let array = sample_array(4);
        assert!(array.object_at(0) != array.object_at(3));
        assert!(array.first_object().unwrap() == array.object_at(0));
        assert!(array.last_object().unwrap() == array.object_at(3));

        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.first_object().is_none());
        assert!(empty_array.last_object().is_none());
    }

    #[test]
    fn test_object_enumerator() {
        let array = sample_array(4);

        assert!(array.object_enumerator().count() == 4);
        assert!(array.object_enumerator()
                     .enumerate()
                     .all(|(i, obj)| obj == array.object_at(i)));
    }

    #[test]
    fn test_objects_in_range() {
        let array = sample_array(4);

        let middle_objs = array.objects_in_range(1, 2);
        assert!(middle_objs.len() == 2);
        assert!(middle_objs[0] == array.object_at(1));
        assert!(middle_objs[1] == array.object_at(2));

        let empty_objs = array.objects_in_range(1, 0);
        assert!(empty_objs.len() == 0);

        let all_objs = array.objects_in_range(0, 4);
        assert!(all_objs.len() == 4);
    }

    #[test]
    fn test_into_vec() {
        let array = sample_array(4);

        let vec = INSArray::into_vec(array);
        assert!(vec.len() == 4);
    }

    #[test]
    fn test_add_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        assert!(array.count() == 1);
        assert!(array.object_at(0) == array.object_at(0));

        let obj: Id<NSObject> = INSObject::new();
        array.insert_object_at(0, obj);
        assert!(array.count() == 2);
    }

    #[test]
    fn test_replace_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        let obj: Id<NSObject> = INSObject::new();
        let old_obj = array.replace_object_at(0, obj);
        assert!(&*old_obj != array.object_at(0));
    }

    #[test]
    fn test_remove_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        for _ in range(0u, 4) {
            let obj: Id<NSObject> = INSObject::new();
            array.add_object(obj);
        }

        array.remove_object_at(1);
        assert!(array.count() == 3);

        array.remove_last_object();
        assert!(array.count() == 2);

        array.remove_all_objects();
        assert!(array.count() == 0);
    }

    #[test]
    fn test_sort() {
        let strings: Vec<Id<NSString>> = vec![
            INSString::from_str("hello"),
            INSString::from_str("hi"),
        ];
        let mut strings: Id<NSMutableArray<_>> = INSArray::from_vec(strings);

        extern fn compare_strings_by_len(s1: &NSString, s2: &NSString,
                _: Option<&()>) -> NSComparisonResult {
            let order = s1.as_str().len().cmp(&s2.as_str().len());
            NSComparisonResult::from_ordering(order)
        }

        strings.sort_with_context(compare_strings_by_len, None);
        assert!(strings[0].as_str() == "hi");
        assert!(strings[1].as_str() == "hello");
    }
}
<commit_msg>Replace sort_with_context with sort_by.<commit_after>use std::kinds::marker::{ContravariantLifetime, NoCopy};

use objc::runtime::Object;
use objc::{Id, IdVector, IntoIdVector, Owned, Ownership, Shared, ShareId};

use {class, INSCopying, INSMutableCopying, INSObject};

#[repr(C)]
pub enum NSComparisonResult {
    Ascending  = -1i,
    Same       = 0i,
    Descending = 1i,
}

impl NSComparisonResult {
    pub fn from_ordering(order: Ordering) -> NSComparisonResult {
        match order {
            Less => NSComparisonResult::Ascending,
            Equal => NSComparisonResult::Same,
            Greater => NSComparisonResult::Descending,
        }
    }

    pub fn as_ordering(&self) -> Ordering {
        match *self {
            NSComparisonResult::Ascending => Less,
            NSComparisonResult::Same => Equal,
            NSComparisonResult::Descending => Greater,
        }
    }
}

pub struct NSRange {
    pub location: uint,
    pub length: uint,
}

pub struct NSEnumerator<'a, T> {
    id: Id<Object>,
    marker: ContravariantLifetime<'a>,
}

impl<'a, T> NSEnumerator<'a, T> {
    pub unsafe fn from_ptr(ptr: *mut Object) -> NSEnumerator<'a, T> {
        NSEnumerator { id: Id::from_ptr(ptr), marker: ContravariantLifetime }
    }
}

impl<'a, T> Iterator<&'a T> for NSEnumerator<'a, T> {
    fn next(&mut self) -> Option<&'a T> {
        unsafe {
            let obj = msg_send![self.id nextObject] as *mut T;
            obj.as_ref()
        }
    }
}

pub trait INSArray<T: INSObject, O: Ownership> : INSObject {
    fn count(&self) -> uint {
        let result = unsafe {
            msg_send![self count]
        };
        result as uint
    }

    fn object_at(&self, index: uint) -> &T {
        unsafe {
            let obj = msg_send![self objectAtIndex:index] as *const T;
            &*obj
        }
    }

    fn first_object(&self) -> Option<&T> {
        unsafe {
            let obj = msg_send![self firstObject] as *const T;
            obj.as_ref()
        }
    }

    fn last_object(&self) -> Option<&T> {
        unsafe {
            let obj = msg_send![self lastObject] as *const T;
            obj.as_ref()
        }
    }

    fn object_enumerator(&self) -> NSEnumerator<T> {
        unsafe {
            let result = msg_send![self objectEnumerator];
            NSEnumerator::from_ptr(result)
        }
    }

    unsafe fn from_refs(refs: &[&T]) -> Id<Self> {
        let cls = class::<Self>();
        let obj = msg_send![cls alloc];
        let obj = msg_send![obj initWithObjects:refs.as_ptr() count:refs.len()];
        Id::from_retained_ptr(obj as *mut Self)
    }

    fn from_vec(vec: Vec<Id<T, O>>) -> Id<Self> {
        let refs = vec.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn objects_in_range(&self, start: uint, len: uint) -> Vec<&T> {
        let vec: Vec<*mut T> = Vec::from_elem(len, RawPtr::null());
        let range = NSRange { location: start, length: len };
        unsafe {
            msg_send![self getObjects:vec.as_ptr() range:range];
            vec.map_in_place(|ptr| &*ptr)
        }
    }

    fn to_vec(&self) -> Vec<&T> {
        self.objects_in_range(0, self.count())
    }

    fn into_vec(array: Id<Self>) -> Vec<Id<T, O>> {
        let vec = array.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub trait INSOwnedArray<T: INSObject> : INSArray<T, Owned> {
    fn mut_object_at(&mut self, index: uint) -> &mut T {
        unsafe {
            let result = msg_send![self objectAtIndex:index] as *mut T;
            &mut *result
        }
    }
}

pub trait INSSharedArray<T: INSObject> : INSArray<T, Shared> {
    fn shared_object_at(&self, index: uint) -> ShareId<T> {
        let obj = self.object_at(index);
        unsafe {
            Id::from_ptr(obj as *const _ as *mut T)
        }
    }

    fn from_slice(slice: &[ShareId<T>]) -> Id<Self> {
        let refs = slice.as_refs_slice();
        unsafe {
            INSArray::from_refs(refs)
        }
    }

    fn to_shared_vec(&self) -> Vec<ShareId<T>> {
        let vec = self.to_vec();
        unsafe {
            vec.into_id_vec()
        }
    }
}

pub struct NSArray<T, O = Owned> {
    nocopy: NoCopy,
}

object_impl!(NSArray<T, O>)

impl<T: INSObject, O: Ownership> INSArray<T, O> for NSArray<T, O> { }

impl<T: INSObject> INSOwnedArray<T> for NSArray<T, Owned> { }

impl<T: INSObject> INSSharedArray<T> for NSArray<T, Shared> { }

impl<T> INSCopying<NSSharedArray<T>> for NSArray<T, Shared> { }

impl<T> INSMutableCopying<NSMutableSharedArray<T>> for NSArray<T, Shared> { }

impl<T: INSObject, O: Ownership> Index<uint, T> for NSArray<T, O> {
    fn index(&self, index: &uint) -> &T {
        self.object_at(*index)
    }
}

pub type NSSharedArray<T> = NSArray<T, Shared>;

pub trait INSMutableArray<T: INSObject, O: Ownership> : INSArray<T, O> {
    fn add_object(&mut self, obj: Id<T, O>) {
        unsafe {
            msg_send![self addObject:obj];
        }
    }

    fn insert_object_at(&mut self, index: uint, obj: Id<T, O>) {
        unsafe {
            msg_send![self insertObject:obj atIndex:index];
        }
    }

    fn replace_object_at(&mut self, index: uint, obj: Id<T, O>) -> Id<T, O> {
        let old_obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut T)
        };
        unsafe {
            msg_send![self replaceObjectAtIndex:index withObject:obj];
        }
        old_obj
    }

    fn remove_object_at(&mut self, index: uint) -> Id<T, O> {
        let obj = unsafe {
            let obj = self.object_at(index);
            Id::from_ptr(obj as *const _ as *mut T)
        };
        unsafe {
            msg_send![self removeObjectAtIndex:index];
        }
        obj
    }

    fn remove_last_object(&mut self) -> Id<T, O> {
        let obj = self.last_object().map(|obj| unsafe {
            Id::from_ptr(obj as *const _ as *mut T)
        });
        unsafe {
            msg_send![self removeLastObject];
        }
        // removeLastObject would have failed if the array is empty,
        // so we know this won't be None
        obj.unwrap()
    }

    fn remove_all_objects(&mut self) {
        unsafe {
            msg_send![self removeAllObjects];
        }
    }

    fn sort_by(&mut self, compare: |&T, &T| -> Ordering) {
        extern fn compare_with_closure<T>(obj1: &T, obj2: &T,
                compare: &mut |&T, &T| -> Ordering) -> NSComparisonResult {
            NSComparisonResult::from_ordering((*compare)(obj1, obj2))
        }

        let mut closure = compare;
        unsafe {
            msg_send![self sortUsingFunction:compare_with_closure::<T>
                                     context:&mut closure];
        }
    }
}

pub struct NSMutableArray<T, O = Owned> {
    nocopy: NoCopy,
}

object_impl!(NSMutableArray<T, O>)

impl<T: INSObject, O: Ownership> INSArray<T, O> for NSMutableArray<T, O> { }

impl<T: INSObject> INSOwnedArray<T> for NSMutableArray<T, Owned> { }

impl<T: INSObject> INSSharedArray<T> for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> INSMutableArray<T, O> for NSMutableArray<T, O> { }

impl<T> INSCopying<NSSharedArray<T>> for NSMutableArray<T, Shared> { }

impl<T> INSMutableCopying<NSMutableSharedArray<T>> for NSMutableArray<T, Shared> { }

impl<T: INSObject, O: Ownership> Index<uint, T> for NSMutableArray<T, O> {
    fn index(&self, index: &uint) -> &T {
        self.object_at(*index)
    }
}

pub type NSMutableSharedArray<T> = NSMutableArray<T, Shared>;

#[cfg(test)]
mod tests {
    use objc::{Id};
    use {INSObject, INSString, NSObject, NSString};
    use super::{
        INSArray, INSMutableArray,
        NSArray, NSMutableArray, NSComparisonResult
    };

    fn sample_array(len: uint) -> Id<NSArray<NSObject>> {
        let vec: Vec<Id<NSObject>> = Vec::from_fn(len, |_| INSObject::new());
        INSArray::from_vec(vec)
    }

    #[test]
    fn test_count() {
        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.count() == 0);

        let array = sample_array(4);
        assert!(array.count() == 4);
    }

    #[test]
    fn test_object_at() {
        let array = sample_array(4);
        assert!(array.object_at(0) != array.object_at(3));
        assert!(array.first_object().unwrap() == array.object_at(0));
        assert!(array.last_object().unwrap() == array.object_at(3));

        let empty_array: Id<NSArray<NSObject>> = INSObject::new();
        assert!(empty_array.first_object().is_none());
        assert!(empty_array.last_object().is_none());
    }

    #[test]
    fn test_object_enumerator() {
        let array = sample_array(4);

        assert!(array.object_enumerator().count() == 4);
        assert!(array.object_enumerator()
                     .enumerate()
                     .all(|(i, obj)| obj == array.object_at(i)));
    }

    #[test]
    fn test_objects_in_range() {
        let array = sample_array(4);

        let middle_objs = array.objects_in_range(1, 2);
        assert!(middle_objs.len() == 2);
        assert!(middle_objs[0] == array.object_at(1));
        assert!(middle_objs[1] == array.object_at(2));

        let empty_objs = array.objects_in_range(1, 0);
        assert!(empty_objs.len() == 0);

        let all_objs = array.objects_in_range(0, 4);
        assert!(all_objs.len() == 4);
    }

    #[test]
    fn test_into_vec() {
        let array = sample_array(4);

        let vec = INSArray::into_vec(array);
        assert!(vec.len() == 4);
    }

    #[test]
    fn test_add_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        assert!(array.count() == 1);
        assert!(array.object_at(0) == array.object_at(0));

        let obj: Id<NSObject> = INSObject::new();
        array.insert_object_at(0, obj);
        assert!(array.count() == 2);
    }

    #[test]
    fn test_replace_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        let obj: Id<NSObject> = INSObject::new();
        array.add_object(obj);

        let obj: Id<NSObject> = INSObject::new();
        let old_obj = array.replace_object_at(0, obj);
        assert!(&*old_obj != array.object_at(0));
    }

    #[test]
    fn test_remove_object() {
        let mut array: Id<NSMutableArray<NSObject>> = INSObject::new();
        for _ in range(0u, 4) {
            let obj: Id<NSObject> = INSObject::new();
            array.add_object(obj);
        }

        array.remove_object_at(1);
        assert!(array.count() == 3);

        array.remove_last_object();
        assert!(array.count() == 2);

        array.remove_all_objects();
        assert!(array.count() == 0);
    }

    #[test]
    fn test_sort() {
        let strings: Vec<Id<NSString>> = vec![
            INSString::from_str("hello"),
            INSString::from_str("hi"),
        ];
        let mut strings: Id<NSMutableArray<_>> = INSArray::from_vec(strings);

        strings.sort_by(|s1, s2| s1.as_str().len().cmp(&s2.as_str().len()));
        assert!(strings[0].as_str() == "hi");
        assert!(strings[1].as_str() == "hello");
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;
use std::fs::File;

const INLINE_TAGS: [&'static str; 16] = [
    "area", "base", "br", "col", "command", "embed", "hr", "img",
    "input", "keygen", "link", "meta", "param", "source", "track", "wbr"];

const DOCTYPE_HTML: &'static str = "html";
const DOCTYPE_XML: &'static str = "xml";
const DOCTYPE_JSON: &'static str = "json";

pub struct Document {
    path: &'static str,
    contents: String,
    doctype: &'static str
}

impl Document {
    pub fn new(path: &'static str) -> Result<Document, io::Error> {
        let mut f = try!(File::open(path));

        let mut contents = String::new();
        try!(f.read_to_string(&mut contents));

        Ok(Document {
            path: path,
            contents: contents,
            doctype: "unknown"
        })
    }

    pub fn validate(&mut self) -> Option<&'static str> {
        if self.contents.len() == 0 {
            return Some("The file was empty");
        }

        if !self.contents.starts_with("doctype") && !self.contents.starts_with("extends") {
            return Some("The document must start with a 'doctype' or 'extends'");
        }

        // TODO: this isn't future proof, will need fixing
        match self.contents.lines().collect::<Vec<_>>()[0] {
            "doctype html" => self.doctype = DOCTYPE_HTML,
            "doctype xml" => self.doctype = DOCTYPE_XML,
            "doctype json" => self.doctype = DOCTYPE_JSON,
            _ => return Some("Unknown 'doctype' supplied")
        }

        None
    }

    pub fn compile(self, pretty: bool) -> String {
        let nodes = parse(self);

        let mut parent_stack: Vec<&DocumentNode> = Vec::new();
        let mut output = String::new();
        let len = nodes.len();
        let mut i = 0;

        for n in &nodes {
            if i == 0 {
                if n.tokens[0] == "doctype" {
                    if let Some(doctype) = generate_doctype(&n.content) {
                        output.push_str(&doctype.to_string());

                        if pretty {
                            output.push('\n');
                        }
                    } else {
                        // TODO: proper error handling
                        panic!("Unknown 'doctype' suppied");
                    }
                } /*else if n.tokens[0] == "extends" {
                    // TODO: when extends is added
                }*/

                i += 1;
                continue;
            }

            // TODO: sort this code nesting out, it's too deep
            let has_next = i + 1 < len;
            let has_sub = has_next && nodes[i + 1].depth > n.depth;

            if has_sub {
                output.push_str(&n.render_open(pretty).to_string());
                parent_stack.push(n);
            } else {
                output.push_str(&n.render(pretty).to_string());

                if has_next && nodes[i + 1].depth < n.depth {
                    loop {
                        match parent_stack.pop() {
                            None =>  { break; },
                            Some(p) => {
                                output.push_str(&p.render_end(pretty).to_string());

                                if p.depth == nodes[i + 1].depth {
                                    break;
                                }
                            }
                        }
                    }
                }
            }

            if !has_next && parent_stack.len() > 0 {
                loop {
                    match parent_stack.pop() {
                        None =>  { break; },
                        Some(p) => {
                            output.push_str(&p.render_end(pretty).to_string());
                        }
                    }
                }
            }

            if !pretty && has_next {
                let next = &nodes[i + 1];

                if n.tokens[0] == "|" && next.tokens[0] == "|" {
                    output.push(' ');
                }
            }

            i += 1;
        }

        output
    }
}

pub struct DocumentNode {
    depth: usize,
    tokens: Vec<String>,
    content: String,
    is_self_closing: bool,
    ignore_sub_content: bool
}

impl DocumentNode {
    fn from_content(indent: usize, content: String) -> DocumentNode {
        DocumentNode {
            depth: indent,
            tokens: vec!["|".to_string()],
            content: content,
            is_self_closing: false,
            ignore_sub_content: true
        }
    }

    fn from(line: &str, indent: usize, doctype: &str) -> Option<DocumentNode> {
        let (tokens, content) = match split_tokens(String::from(line.trim())) {
            Some(result) => result,
            None => return None
        };

        let mut is_self_closing = false;
        let mut ignore_sub_content = false;
        let last_token = &*tokens.last().unwrap().to_string();

        if tokens[0] == "//" || tokens[0] == "|" || tokens[0] == "style" || tokens[0] == "script" {
            ignore_sub_content = true;
        } else if last_token == "/" || (doctype == DOCTYPE_HTML && INLINE_TAGS.contains(&&*tokens[0].to_string())) {
            is_self_closing = true;
        }

        Some(DocumentNode {
            depth: indent,
            tokens: tokens,
            content: content,
            is_self_closing: is_self_closing,
            ignore_sub_content: ignore_sub_content
        })
    }

    fn render(&self, pretty: bool) -> String {
        let mut output = String::new();

        if self.tokens[0] == "|" {
            output.push_str(&self.content.to_string());

            if pretty {
                return pretty_print(&output, self.depth);
            } else {
                return output;
            }
        }

        output.push_str(&self.render_open(false).to_string());

        if self.is_self_closing {
            if pretty {
                return pretty_print(&output, self.depth);
            }

            return output;
        }

        if !self.content.is_empty() {
            output.push_str(&self.content.to_string());
        }

        output.push_str(&self.render_end(false).to_string());

        if pretty {
            return pretty_print(&output, self.depth);
        }

        output
    }

    fn render_open(&self, pretty: bool) -> String {
        let tag_tokens = &self.tokens;
        let mut output = String::new();

        if tag_tokens[0] == "//" {
            output.push_str("<!--");
        } else if tag_tokens[0] == "|" {
            output.push_str(&self.content.to_string());
        } else {
            output.push_str(&format!("<{}", tag_tokens[0]).to_string());

            if tag_tokens.len() > 1 {
                let mut classes = String::new();
                let mut attributes = String::new();

                for t in &tag_tokens[1..] {
                    if t.starts_with('#') {
                        output.push_str(&format!(" id=\"{}\"", t[1..].to_string()))
                    } else if t.starts_with('.') {
                        classes.push_str(&format!(" {}", t[1..].to_string()));
                    } else if t.starts_with('(') {
                        let attrs: Vec<_> = t[1..].split(',').collect();

                        attributes.push_str(&attrs.concat().to_string());
                    }
                }

                if !classes.is_empty() {
                    output.push_str(&format!(" class=\"{}\"", classes.trim()));
                }

                if !attributes.is_empty() {
                    output.push_str(&format!(" {}", attributes).to_string());
                }
            }

            if self.is_self_closing {
                output.push('/');
            }

            output.push('>');
        }

        if pretty {
            return pretty_print(&output, self.depth);
        }

        output
    }

    fn render_end(&self, pretty: bool) -> String {
        let output = match self.tokens[0].as_ref() {
            "//" => "-->".to_string(),
            "|" => String::new(),
            _ => format!("</{}>", &self.tokens[0])
        };

        if pretty && !output.is_empty() {
            return pretty_print(&output, self.depth);
        }

        output
    }
}

fn parse(doc: Document) -> Vec<DocumentNode> {
    let mut nodes: Vec<DocumentNode> = Vec::new();
    let mut parsable_indent = 0;
    let mut mode = 0;

    for line in doc.contents.lines() {
        if line.is_empty() {
            continue;
        }

        let mut indent: usize = 0;

        for c in line.chars() {
            if c.is_whitespace() {
                indent += 1;
            } else {
                break;
            }
        }

        if mode == 2 && indent > parsable_indent {
            continue;
        } else if mode == 1 && indent > parsable_indent {
            nodes.push(DocumentNode::from_content(indent, line.trim().to_string()));
            continue;
        }

        if let Some(node) = DocumentNode::from(line, indent, doc.doctype) {
            if node.ignore_sub_content {
                mode = 1;
            } else {
                mode = 0;
            }

            nodes.push(node);
        } else {
            mode = 2;
        }

        parsable_indent = indent;
    }

    nodes
}

fn split_tokens(s: String) -> Option<(Vec<String>, String)> {
    let mut tokens: Vec<String> = Vec::new();

    let mut start = 0;
    let len = s.len();
    let mut mode = 0;
    let mut content = String::new();

    if s.starts_with("|") {
        tokens.push("|".to_string());
        content.push_str(&s[1..].trim().to_string());
    } else if s.starts_with("//-") {
        return None;
    } else if s.starts_with("//") {
        tokens.push("//".to_string());
        content.push_str(&s[2..].trim().to_string());
    }

    if tokens.len() > 0 {
        return Some((tokens, content));
    }

    for (i, c) in s.chars().enumerate() {
        if c.is_whitespace() && mode == 0 {
            tokens.push(s[start..i].to_string());
            content.push_str(&s[i..len].trim().to_string());
            break;
        }

        if c == ')' {
            if mode != 1 {
                // TODO: proper error handling
                panic!("Invalid attribute closing brace");
            }

            tokens.push(s[start..i].to_string());

            mode = 0;
            start = i;
        } else {
            if c == '/' && mode == 0 {
                tokens.push(s[start..i].to_string());
                tokens.push("/".to_string());
                start = i + 1;
            } else if ((c == '#' || c == '.' || c == '(') && mode == 0) || i == len - 1 {
                if i == len - 1 {
                    tokens.push(s[start..].to_string());
                } else if i > start {
                    tokens.push(s[start..i].to_string());
                }

                if c == '(' {
                    mode = 1;
                }

                start = i;
            }
        }
    }

    // TODO: pretty sure this can be done using a pattern in the starts_with
    if tokens.len() > 0 && (tokens[0].starts_with(".") || tokens[0].starts_with("#")) {
        tokens.insert(0, "div".to_string());
    }

    Some((tokens, content))
}

// TODO: need to abstract this stuff to support multiple doctypes
fn generate_doctype(content: &String) -> Option<String> {
    return match content.as_ref() {
        "html" => Some(format!("<!DOCTYPE {}>", content)),
        "xml" => Some("<?xml version=\"1.0\" encoding=\"utf-8\" ?>".to_string()),
        _ => None
    };
}

// TODO: look at making this a macro
// TODO: reduce if statements when using this by applying it always with a flag to indicate which format to use
fn pretty_print(content: &String, depth: usize) -> String {
    let mut output = String::new();

    for _ in 0..depth {
        output.push(' ');
    }

    output.push_str(&content.to_string());
    output.push('\n');
    output
}
<commit_msg>fixed typo<commit_after>use std::io;
use std::io::prelude::*;
use std::fs::File;

const INLINE_TAGS: [&'static str; 16] = [
    "area", "base", "br", "col", "command", "embed", "hr", "img",
    "input", "keygen", "link", "meta", "param", "source", "track", "wbr"];

const DOCTYPE_HTML: &'static str = "html";
const DOCTYPE_XML: &'static str = "xml";
const DOCTYPE_JSON: &'static str = "json";

pub struct Document {
    path: &'static str,
    contents: String,
    doctype: &'static str
}

impl Document {
    pub fn new(path: &'static str) -> Result<Document, io::Error> {
        let mut f = try!(File::open(path));

        let mut contents = String::new();
        try!(f.read_to_string(&mut contents));

        Ok(Document {
            path: path,
            contents: contents,
            doctype: "unknown"
        })
    }

    pub fn validate(&mut self) -> Option<&'static str> {
        if self.contents.len() == 0 {
            return Some("The file was empty");
        }

        if !self.contents.starts_with("doctype") && !self.contents.starts_with("extends") {
            return Some("The document must start with a 'doctype' or 'extends'");
        }

        // TODO: this isn't future proof, will need fixing
        match self.contents.lines().collect::<Vec<_>>()[0] {
            "doctype html" => self.doctype = DOCTYPE_HTML,
            "doctype xml" => self.doctype = DOCTYPE_XML,
            "doctype json" => self.doctype = DOCTYPE_JSON,
            _ => return Some("Unknown 'doctype' supplied")
        }

        None
    }

    // TODO: move this into the compiler mod
    pub fn compile(self, pretty: bool) -> String {
        let nodes = parse(self);

        let mut parent_stack: Vec<&DocumentNode> = Vec::new();
        let mut output = String::new();
        let len = nodes.len();
        let mut i = 0;

        for n in &nodes {
            if i == 0 {
                if n.tokens[0] == "doctype" {
                    if let Some(doctype) = generate_doctype(&n.content) {
                        output.push_str(&doctype.to_string());

                        if pretty {
                            output.push('\n');
                        }
                    } else {
                        // TODO: proper error handling
                        panic!("Unknown 'doctype' suppied");
                    }
                } /*else if n.tokens[0] == "extends" {
                    // TODO: when extends is added
                }*/

                i += 1;
                continue;
            }

            // TODO: sort this code nesting out, it's too deep
            let has_next = i + 1 < len;
            let has_sub = has_next && nodes[i + 1].depth > n.depth;

            if has_sub {
                output.push_str(&n.render_open(pretty).to_string());
                parent_stack.push(n);
            } else {
                output.push_str(&n.render(pretty).to_string());

                if has_next && nodes[i + 1].depth < n.depth {
                    loop {
                        match parent_stack.pop() {
                            None =>  { break; },
                            Some(p) => {
                                output.push_str(&p.render_end(pretty).to_string());

                                if p.depth == nodes[i + 1].depth {
                                    break;
                                }
                            }
                        }
                    }
                }
            }

            if !has_next && parent_stack.len() > 0 {
                loop {
                    match parent_stack.pop() {
                        None =>  { break; },
                        Some(p) => {
                            output.push_str(&p.render_end(pretty).to_string());
                        }
                    }
                }
            }

            if !pretty && has_next {
                let next = &nodes[i + 1];

                if n.tokens[0] == "|" && next.tokens[0] == "|" {
                    output.push(' ');
                }
            }

            i += 1;
        }

        output
    }
}

pub struct DocumentNode {
    depth: usize,
    tokens: Vec<String>,
    content: String,
    is_self_closing: bool,
    ignore_sub_content: bool
}

impl DocumentNode {
    fn from_content(indent: usize, content: String) -> DocumentNode {
        DocumentNode {
            depth: indent,
            tokens: vec!["|".to_string()],
            content: content,
            is_self_closing: false,
            ignore_sub_content: true
        }
    }

    fn from(line: &str, indent: usize, doctype: &str) -> Option<DocumentNode> {
        let (tokens, content) = match split_tokens(String::from(line.trim())) {
            Some(result) => result,
            None => return None
        };

        let mut is_self_closing = false;
        let mut ignore_sub_content = false;
        let last_token = &*tokens.last().unwrap().to_string();

        if tokens[0] == "//" || tokens[0] == "|" || tokens[0] == "style" || tokens[0] == "script" {
            ignore_sub_content = true;
        } else if last_token == "/" || (doctype == DOCTYPE_HTML && INLINE_TAGS.contains(&&*tokens[0].to_string())) {
            is_self_closing = true;
        }

        Some(DocumentNode {
            depth: indent,
            tokens: tokens,
            content: content,
            is_self_closing: is_self_closing,
            ignore_sub_content: ignore_sub_content
        })
    }

    fn render(&self, pretty: bool) -> String {
        let mut output = String::new();

        if self.tokens[0] == "|" {
            output.push_str(&self.content.to_string());

            if pretty {
                return pretty_print(&output, self.depth);
            } else {
                return output;
            }
        }

        output.push_str(&self.render_open(false).to_string());

        if self.is_self_closing {
            if pretty {
                return pretty_print(&output, self.depth);
            }

            return output;
        }

        if !self.content.is_empty() {
            output.push_str(&self.content.to_string());
        }

        output.push_str(&self.render_end(false).to_string());

        if pretty {
            return pretty_print(&output, self.depth);
        }

        output
    }

    fn render_open(&self, pretty: bool) -> String {
        let tag_tokens = &self.tokens;
        let mut output = String::new();

        if tag_tokens[0] == "//" {
            output.push_str("<!--");
        } else if tag_tokens[0] == "|" {
            output.push_str(&self.content.to_string());
        } else {
            output.push_str(&format!("<{}", tag_tokens[0]).to_string());

            if tag_tokens.len() > 1 {
                let mut classes = String::new();
                let mut attributes = String::new();

                for t in &tag_tokens[1..] {
                    if t.starts_with('#') {
                        output.push_str(&format!(" id=\"{}\"", t[1..].to_string()))
                    } else if t.starts_with('.') {
                        classes.push_str(&format!(" {}", t[1..].to_string()));
                    } else if t.starts_with('(') {
                        let attrs: Vec<_> = t[1..].split(',').collect();

                        attributes.push_str(&attrs.concat().to_string());
                    }
                }

                if !classes.is_empty() {
                    output.push_str(&format!(" class=\"{}\"", classes.trim()));
                }

                if !attributes.is_empty() {
                    output.push_str(&format!(" {}", attributes).to_string());
                }
            }

            if self.is_self_closing {
                output.push('/');
            }

            output.push('>');
        }

        if pretty {
            return pretty_print(&output, self.depth);
        }

        output
    }

    fn render_end(&self, pretty: bool) -> String {
        let output = match self.tokens[0].as_ref() {
            "//" => "-->".to_string(),
            "|" => String::new(),
            _ => format!("</{}>", &self.tokens[0])
        };

        if pretty && !output.is_empty() {
            return pretty_print(&output, self.depth);
        }

        output
    }
}

fn parse(doc: Document) -> Vec<DocumentNode> {
    let mut nodes: Vec<DocumentNode> = Vec::new();
    let mut parsable_indent = 0;
    let mut mode = 0;

    for line in doc.contents.lines() {
        if line.is_empty() {
            continue;
        }

        let mut indent: usize = 0;

        for c in line.chars() {
            if c.is_whitespace() {
                indent += 1;
            } else {
                break;
            }
        }

        if mode == 2 && indent > parsable_indent {
            continue;
        } else if mode == 1 && indent > parsable_indent {
            nodes.push(DocumentNode::from_content(indent, line.trim().to_string()));
            continue;
        }

        if let Some(node) = DocumentNode::from(line, indent, doc.doctype) {
            if node.ignore_sub_content {
                mode = 1;
            } else {
                mode = 0;
            }

            nodes.push(node);
        } else {
            mode = 2;
        }

        parsable_indent = indent;
    }

    nodes
}

fn split_tokens(s: String) -> Option<(Vec<String>, String)> {
    let mut tokens: Vec<String> = Vec::new();

    let mut start = 0;
    let len = s.len();
    let mut mode = 0;
    let mut content = String::new();

    if s.starts_with("|") {
        tokens.push("|".to_string());
        content.push_str(&s[1..].trim().to_string());
    } else if s.starts_with("//-") {
        return None;
    } else if s.starts_with("//") {
        tokens.push("//".to_string());
        content.push_str(&s[2..].trim().to_string());
    }

    if tokens.len() > 0 {
        return Some((tokens, content));
    }

    for (i, c) in s.chars().enumerate() {
        if c.is_whitespace() && mode == 0 {
            tokens.push(s[start..i].to_string());
            content.push_str(&s[i..len].trim().to_string());
            break;
        }

        if c == ')' {
            if mode != 1 {
                // TODO: proper error handling
                panic!("Invalid attribute closing brace");
            }

            tokens.push(s[start..i].to_string());

            mode = 0;
            start = i;
        } else {
            if c == '/' && mode == 0 {
                tokens.push(s[start..i].to_string());
                tokens.push("/".to_string());
                start = i + 1;
            } else if ((c == '#' || c == '.' || c == '(') && mode == 0) || i == len - 1 {
                if i == len - 1 {
                    tokens.push(s[start..].to_string());
                } else if i > start {
                    tokens.push(s[start..i].to_string());
                }

                if c == '(' {
                    mode = 1;
                }

                start = i;
            }
        }
    }

    // TODO: pretty sure this can be done using a pattern in the starts_with
    if tokens.len() > 0 && (tokens[0].starts_with(".") || tokens[0].starts_with("#")) {
        tokens.insert(0, "div".to_string());
    }

    Some((tokens, content))
}

// TODO: need to abstract this stuff to support multiple doctypes
fn generate_doctype(content: &String) -> Option<String> {
    return match content.as_ref() {
        "html" => Some(format!("<!DOCTYPE {}>", content)),
        "xml" => Some("<?xml version=\"1.0\" encoding=\"utf-8\" ?>".to_string()),
        _ => None
    };
}

// TODO: look at making this a macro
// TODO: reduce if statements when using this by applying it always with a flag to indicate which format to use
fn pretty_print(content: &String, depth: usize) -> String {
    let mut output = String::new();

    for _ in 0..depth {
        output.push(' ');
    }

    output.push_str(&content.to_string());
    output.push('\n');
    output
}
<|endoftext|>


<commit_before>use libc::{c_int, c_double, c_char, c_void};
use std::mem;

/// Typedef for initialization function.
pub type EnergyMonInitFn = extern fn(*mut EnergyMon) -> c_int;
/// Typedef for read function.
pub type EnergyMonReadTotalFn = extern fn(*mut EnergyMon) -> c_double;
/// Typedef for cleanup function.
pub type EnergyMonFinishFn = extern fn(*mut EnergyMon) -> c_int;
/// Typedef for function to get human-readable name.
pub type EnergyMonGetSourceFn = extern fn(*mut c_char) -> *mut c_char;

#[repr(C)]
/// Representation of native C struct `em_impl`.
pub struct EnergyMon {
    /// Native C function type signature that initializes the `EnergyMon`.
    pub finit: EnergyMonInitFn,
    /// Native C function type signature that reads energy data from the `EnergyMon`.
    pub fread: EnergyMonReadTotalFn,
    /// Native C function type signature that cleans up the `EnergyMon`.
    pub ffinish: EnergyMonFinishFn,
    /// Native C function type signature that gets the `EnergyMon`'s human-readable name.
    pub fsource: EnergyMonGetSourceFn,
    /// Native C pointer used by the underlying `EnergyMon` implementation for storing state.
    pub state: *mut c_void,
}

extern {
    /// Native C function that fills in the EnergyMon struct values and may allocate other
    /// resources.
    fn em_impl_get(em: *mut EnergyMon) -> c_int;
}

impl EnergyMon {
    /// Create an `EnergyMon` which represents a C struct.
    pub fn new() -> Result<EnergyMon, &'static str> {
        unsafe {
            let mut em: EnergyMon = mem::uninitialized();
            let result: i32 = em_impl_get(&mut em);
            if result != 0 {
                return Err("Failed to create energymon");
            }
            Ok(em)
        }
    }

    /// Initialize the `EnergyMon` - may allocate more resources and/or start threads.
    pub fn init(&mut self) -> Result<(), &'static str> {
        let result: i32 = (self.finit)(self);
        if result != 0 {
            return Err("Failed to initialize energymon");
        }
        Ok(())
    }

    /// Read the total energy from the `EnergyMon`.
    pub fn read(&mut self) -> f64 {
        (self.fread)(self)
    }

    /// Cleanup the `EnergyMon` - may deallocate associated resources and/or stop threads.
    pub fn finish(&mut self) -> Result<(), &'static str> {
        let result: i32 = (self.ffinish)(self);
        if result != 0 {
            return Err("Failed to cleanup energymon");
        }
        Ok(())
    }

    /// Utility function to get a human-readable name of the `EnergyMon`'s source.
    pub fn source(&mut self) -> Result<String, &'static str> {
        const BUFSIZE: usize = 100;
        let mut buf: [c_char; BUFSIZE] = [0; BUFSIZE];
        let ret: *mut c_char = (self.fsource)(buf.as_mut_ptr());
        if ret.is_null() {
            return Err("Failed to get energymon source");
        }
        let buf: &[u8] = unsafe {
            mem::transmute::<&[c_char], &[u8]>(&buf)
        };
        Ok(String::from_utf8_lossy(buf).into_owned())
    }
}

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

    #[test]
    fn test_interface() {
        let mut em: EnergyMon = EnergyMon::new().unwrap();
        em.init().unwrap();
        let val = em.read();
        assert!(val >= 0.0);
        let source = em.source().unwrap();
        em.finish().unwrap();
        println!("Read {} from {}", val, source);
    }

}
<commit_msg>Update to handle long long return value for uJ instead of double for Joules<commit_after>use libc::{c_int, c_longlong, c_char, c_void};
use std::mem;

/// Typedef for initialization function.
pub type EnergyMonInitFn = extern fn(*mut EnergyMon) -> c_int;
/// Typedef for read function.
pub type EnergyMonReadTotalFn = extern fn(*mut EnergyMon) -> c_longlong;
/// Typedef for cleanup function.
pub type EnergyMonFinishFn = extern fn(*mut EnergyMon) -> c_int;
/// Typedef for function to get human-readable name.
pub type EnergyMonGetSourceFn = extern fn(*mut c_char) -> *mut c_char;

#[repr(C)]
/// Representation of native C struct `em_impl`.
pub struct EnergyMon {
    /// Native C function type signature that initializes the `EnergyMon`.
    pub finit: EnergyMonInitFn,
    /// Native C function type signature that reads energy data from the `EnergyMon`.
    pub fread: EnergyMonReadTotalFn,
    /// Native C function type signature that cleans up the `EnergyMon`.
    pub ffinish: EnergyMonFinishFn,
    /// Native C function type signature that gets the `EnergyMon`'s human-readable name.
    pub fsource: EnergyMonGetSourceFn,
    /// Native C pointer used by the underlying `EnergyMon` implementation for storing state.
    pub state: *mut c_void,
}

extern {
    /// Native C function that fills in the EnergyMon struct values and may allocate other
    /// resources.
    fn em_impl_get(em: *mut EnergyMon) -> c_int;
}

impl EnergyMon {
    /// Create an `EnergyMon` which represents a C struct.
    pub fn new() -> Result<EnergyMon, &'static str> {
        unsafe {
            let mut em: EnergyMon = mem::uninitialized();
            let result: i32 = em_impl_get(&mut em);
            if result != 0 {
                return Err("Failed to create energymon");
            }
            Ok(em)
        }
    }

    /// Initialize the `EnergyMon` - may allocate more resources and/or start threads.
    pub fn init(&mut self) -> Result<(), &'static str> {
        let result: i32 = (self.finit)(self);
        if result != 0 {
            return Err("Failed to initialize energymon");
        }
        Ok(())
    }

    /// Read the total energy from the `EnergyMon`.
    pub fn read(&mut self) -> i64 {
        (self.fread)(self)
    }

    /// Cleanup the `EnergyMon` - may deallocate associated resources and/or stop threads.
    pub fn finish(&mut self) -> Result<(), &'static str> {
        let result: i32 = (self.ffinish)(self);
        if result != 0 {
            return Err("Failed to cleanup energymon");
        }
        Ok(())
    }

    /// Utility function to get a human-readable name of the `EnergyMon`'s source.
    pub fn source(&mut self) -> Result<String, &'static str> {
        const BUFSIZE: usize = 100;
        let mut buf: [c_char; BUFSIZE] = [0; BUFSIZE];
        let ret: *mut c_char = (self.fsource)(buf.as_mut_ptr());
        if ret.is_null() {
            return Err("Failed to get energymon source");
        }
        let buf: &[u8] = unsafe {
            mem::transmute::<&[c_char], &[u8]>(&buf)
        };
        Ok(String::from_utf8_lossy(buf).into_owned())
    }
}

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

    #[test]
    fn test_interface() {
        let mut em: EnergyMon = EnergyMon::new().unwrap();
        em.init().unwrap();
        let val = em.read();
        assert!(val >= 0);
        let source = em.source().unwrap();
        em.finish().unwrap();
        println!("Read {} from {}", val, source);
    }

}
<|endoftext|>


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

use minhashes::KmerCount;
use statistics::hist;

/// Used to pass around filter options for sketching
#[derive(Debug)]
pub struct FilterParams {
    pub filter_on: Option<bool>,
    pub abun_filter: (Option<u16>, Option<u16>),
    pub err_filter: f32,
    pub strand_filter: f32,
}

/// Applies filter options to a sketch
pub fn filter_sketch(hashes: &[KmerCount], filters: &FilterParams) -> (Vec<KmerCount>, HashMap<String, String>) {
    let filter_on = filters.filter_on.expect("Sorry! Filter should have either been passed or set during detection");
    let mut filter_stats: HashMap<String, String> = HashMap::new();

    let mut low_abun_filter = filters.abun_filter.0;

    if filter_on && filters.err_filter > 0f32 {
        let cutoff = guess_filter_threshold(&hashes, filters.err_filter);
        if let None = low_abun_filter {
            low_abun_filter = Some(cutoff);
            filter_stats.insert(String::from("errFilter"), filters.err_filter.to_string());
        }
    }

    let mut filtered_hashes = hashes.to_vec();
    if filter_on && (low_abun_filter != None || filters.abun_filter.1 != None) {
        filtered_hashes = filter_abundance(&filtered_hashes, low_abun_filter, filters.abun_filter.1);
        if let Some(v) = low_abun_filter {
            filter_stats.insert(String::from("minCopies"), v.to_string());
        }
        if let Some(v) = filters.abun_filter.1 {
            filter_stats.insert(String::from("maxCopies"), v.to_string());
        }
    }

    if filter_on && filters.strand_filter > 0f32 {
        filtered_hashes = filter_strands(&filtered_hashes, filters.strand_filter);
        filter_stats.insert(String::from("strandFilter"), filters.strand_filter.to_string());
    }

    (filtered_hashes, filter_stats)
}

/// Determines a dynamic filtering threshold for low abundance kmers
///
/// Useful for removing, e.g. kmers containing sequencing errors
///
pub fn guess_filter_threshold(sketch: &[KmerCount], filter_level: f32) -> u16 {
    let hist_data = hist(sketch);
    let total_counts = hist_data.iter().enumerate().map(|t| (t.0 as u64 + 1) * t.1).sum::<u64>() as f32;
    let cutoff_amt = filter_level * total_counts;

    // calculate the coverage that N% of the weighted data is above
    let mut wgt_cutoff: usize = 1;
    let mut cum_count: u64 = 0;
    for count in &hist_data {
        cum_count += wgt_cutoff as u64 * *count as u64;
        if cum_count as f32 > cutoff_amt {
            break;
        }
        wgt_cutoff += 1;
    }

    if wgt_cutoff <= 2 {
        return wgt_cutoff as u16;
    }

    // now find the right-most global maxima
    let win_size = cmp::max(1, wgt_cutoff / 20);
    let mut sum: u64 = hist_data[..win_size].iter().sum();
    let mut lowest_val = sum;
    let mut lowest_idx = win_size;
    for (i, j) in (0..wgt_cutoff - win_size).zip(win_size..wgt_cutoff) {
        if sum <= lowest_val {
            lowest_val = sum;
            lowest_idx = j;
        }
        sum -= hist_data[i];
        sum += hist_data[j];
    }

    lowest_idx as u16
}

#[test]
fn test_guess_filter_threshold() {
    let sketch = vec![];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 9, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 8);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 9, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.1);
    assert_eq!(cutoff, 8);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 2, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 4, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.1);
    assert_eq!(cutoff, 1);
}

pub fn filter_abundance(sketch: &[KmerCount], low: Option<u16>, high: Option<u16>) -> Vec<KmerCount> {
    let mut filtered = Vec::new();
    let lo_threshold = low.unwrap_or(0u16);
    let hi_threshold = high.unwrap_or(u16::max_value());
    for kmer in sketch {
        if lo_threshold <= kmer.count && kmer.count <= hi_threshold {
            filtered.push(kmer.clone());
        }
    }
    filtered
}

#[test]
fn test_filter_abundance() {
    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
    ];
    let filtered = filter_abundance(&sketch, Some(1), None);
    assert_eq!(filtered.len(), 2);
    assert_eq!(filtered[0].hash, 1);
    assert_eq!(filtered[1].hash, 2);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 9, extra_count: 0},
    ];
    let filtered = filter_abundance(&sketch, Some(9), None);
    assert_eq!(filtered.len(), 3);
    assert_eq!(filtered[0].hash, 2);
    assert_eq!(filtered[1].hash, 3);
    assert_eq!(filtered[2].hash, 4);

    let filtered = filter_abundance(&sketch, Some(2), Some(9));
    assert_eq!(filtered.len(), 1);
    assert_eq!(filtered[0].hash, 4);
}

/// Filter out kmers that have a large abundance difference between being seen in the
/// "forward" and "reverse" orientations (picked arbitrarily which is which).
///
/// These tend to be sequencing adapters.
pub fn filter_strands(sketch: &[KmerCount], ratio_cutoff: f32) -> Vec<KmerCount> {
    let mut filtered = Vec::new();
    for kmer in sketch {
        // "special-case" anything with fewer than 16 kmers -> these are too stochastic to accurately
        // determine if something is an adapter or not. The odds of randomly picking less than 10%
        // (0 or 1 reversed kmers) in 16 should be ~ 17 / 2 ** 16 or 1/4000 so we're avoiding
        // removing "good" kmers
        if kmer.count < 16u16 {
            filtered.push(kmer.clone());
            continue;
        }

        // check the forward/reverse ratio and only add if it's within bounds
        let lowest_strand_count: u16 = cmp::min(kmer.extra_count, kmer.count - kmer.extra_count);
        if (lowest_strand_count as f32 / kmer.count as f32) >= ratio_cutoff {
            filtered.push(kmer.clone());
        }
    }
    filtered
}

#[test]
fn test_filter_strands() {
    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 10, extra_count: 1},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 2},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 8},
        KmerCount {hash: 4, kmer: vec![], count: 10, extra_count: 9},
    ];
    let filtered = filter_strands(&sketch, 0.15);
    assert_eq!(filtered.len(), 4);
    assert_eq!(filtered[0].hash, 1);
    assert_eq!(filtered[3].hash, 4);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 16, extra_count: 1},
        KmerCount {hash: 2, kmer: vec![], count: 16, extra_count: 2},
        KmerCount {hash: 3, kmer: vec![], count: 16, extra_count: 8},
        KmerCount {hash: 4, kmer: vec![], count: 16, extra_count: 9},
    ];
    let filtered = filter_strands(&sketch, 0.15);
    assert_eq!(filtered.len(), 2);
    assert_eq!(filtered[0].hash, 3);
    assert_eq!(filtered[1].hash, 4);
}
<commit_msg>Fix bug where all reads would be filtered out when files are heavily adapter contaminated<commit_after>use std::cmp;
use std::collections::HashMap;

use minhashes::KmerCount;
use statistics::hist;

/// Used to pass around filter options for sketching
#[derive(Debug)]
pub struct FilterParams {
    pub filter_on: Option<bool>,
    pub abun_filter: (Option<u16>, Option<u16>),
    pub err_filter: f32,
    pub strand_filter: f32,
}

/// Applies filter options to a sketch
pub fn filter_sketch(hashes: &[KmerCount], filters: &FilterParams) -> (Vec<KmerCount>, HashMap<String, String>) {
    let filter_on = filters.filter_on.expect("Sorry! Filter should have either been passed or set during detection");
    let mut filter_stats: HashMap<String, String> = HashMap::new();

    let mut low_abun_filter = filters.abun_filter.0;
    let mut filtered_hashes = hashes.to_vec();

    if filter_on && filters.strand_filter > 0f32 {
        filtered_hashes = filter_strands(&filtered_hashes, filters.strand_filter);
        filter_stats.insert(String::from("strandFilter"), filters.strand_filter.to_string());
    }

    if filter_on && filters.err_filter > 0f32 {
        let cutoff = guess_filter_threshold(&filtered_hashes, filters.err_filter);
        if let None = low_abun_filter {
            low_abun_filter = Some(cutoff);
            filter_stats.insert(String::from("errFilter"), filters.err_filter.to_string());
        }
    }

    if filter_on && (low_abun_filter != None || filters.abun_filter.1 != None) {
        filtered_hashes = filter_abundance(&filtered_hashes, low_abun_filter, filters.abun_filter.1);
        if let Some(v) = low_abun_filter {
            filter_stats.insert(String::from("minCopies"), v.to_string());
        }
        if let Some(v) = filters.abun_filter.1 {
            filter_stats.insert(String::from("maxCopies"), v.to_string());
        }
    }

    (filtered_hashes, filter_stats)
}

/// Determines a dynamic filtering threshold for low abundance kmers
///
/// Useful for removing, e.g. kmers containing sequencing errors
///
pub fn guess_filter_threshold(sketch: &[KmerCount], filter_level: f32) -> u16 {
    let hist_data = hist(sketch);
    let total_counts = hist_data.iter().enumerate().map(|t| (t.0 as u64 + 1) * t.1).sum::<u64>() as f32;
    let cutoff_amt = filter_level * total_counts;

    // calculate the coverage that N% of the weighted data is above
    let mut wgt_cutoff: usize = 1;
    let mut cum_count: u64 = 0;
    for count in &hist_data {
        cum_count += wgt_cutoff as u64 * *count as u64;
        if cum_count as f32 > cutoff_amt {
            break;
        }
        wgt_cutoff += 1;
    }

    if wgt_cutoff <= 2 {
        return wgt_cutoff as u16;
    }

    // now find the right-most global maxima
    let win_size = cmp::max(1, wgt_cutoff / 20);
    let mut sum: u64 = hist_data[..win_size].iter().sum();
    let mut lowest_val = sum;
    let mut lowest_idx = win_size;
    for (i, j) in (0..wgt_cutoff - win_size).zip(win_size..wgt_cutoff) {
        if sum <= lowest_val {
            lowest_val = sum;
            lowest_idx = j;
        }
        sum -= hist_data[i];
        sum += hist_data[j];
    }

    lowest_idx as u16
}

#[test]
fn test_guess_filter_threshold() {
    let sketch = vec![];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 1);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 9, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.2);
    assert_eq!(cutoff, 8);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 9, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.1);
    assert_eq!(cutoff, 8);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 2, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 4, extra_count: 0},
    ];
    let cutoff = guess_filter_threshold(&sketch, 0.1);
    assert_eq!(cutoff, 1);
}

pub fn filter_abundance(sketch: &[KmerCount], low: Option<u16>, high: Option<u16>) -> Vec<KmerCount> {
    let mut filtered = Vec::new();
    let lo_threshold = low.unwrap_or(0u16);
    let hi_threshold = high.unwrap_or(u16::max_value());
    for kmer in sketch {
        if lo_threshold <= kmer.count && kmer.count <= hi_threshold {
            filtered.push(kmer.clone());
        }
    }
    filtered
}

#[test]
fn test_filter_abundance() {
    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 1, extra_count: 0},
    ];
    let filtered = filter_abundance(&sketch, Some(1), None);
    assert_eq!(filtered.len(), 2);
    assert_eq!(filtered[0].hash, 1);
    assert_eq!(filtered[1].hash, 2);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 1, extra_count: 0},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 0},
        KmerCount {hash: 4, kmer: vec![], count: 9, extra_count: 0},
    ];
    let filtered = filter_abundance(&sketch, Some(9), None);
    assert_eq!(filtered.len(), 3);
    assert_eq!(filtered[0].hash, 2);
    assert_eq!(filtered[1].hash, 3);
    assert_eq!(filtered[2].hash, 4);

    let filtered = filter_abundance(&sketch, Some(2), Some(9));
    assert_eq!(filtered.len(), 1);
    assert_eq!(filtered[0].hash, 4);
}

/// Filter out kmers that have a large abundance difference between being seen in the
/// "forward" and "reverse" orientations (picked arbitrarily which is which).
///
/// These tend to be sequencing adapters.
pub fn filter_strands(sketch: &[KmerCount], ratio_cutoff: f32) -> Vec<KmerCount> {
    let mut filtered = Vec::new();
    for kmer in sketch {
        // "special-case" anything with fewer than 16 kmers -> these are too stochastic to accurately
        // determine if something is an adapter or not. The odds of randomly picking less than 10%
        // (0 or 1 reversed kmers) in 16 should be ~ 17 / 2 ** 16 or 1/4000 so we're avoiding
        // removing "good" kmers
        if kmer.count < 16u16 {
            filtered.push(kmer.clone());
            continue;
        }

        // check the forward/reverse ratio and only add if it's within bounds
        let lowest_strand_count: u16 = cmp::min(kmer.extra_count, kmer.count - kmer.extra_count);
        if (lowest_strand_count as f32 / kmer.count as f32) >= ratio_cutoff {
            filtered.push(kmer.clone());
        }
    }
    filtered
}

#[test]
fn test_filter_strands() {
    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 10, extra_count: 1},
        KmerCount {hash: 2, kmer: vec![], count: 10, extra_count: 2},
        KmerCount {hash: 3, kmer: vec![], count: 10, extra_count: 8},
        KmerCount {hash: 4, kmer: vec![], count: 10, extra_count: 9},
    ];
    let filtered = filter_strands(&sketch, 0.15);
    assert_eq!(filtered.len(), 4);
    assert_eq!(filtered[0].hash, 1);
    assert_eq!(filtered[3].hash, 4);

    let sketch = vec![
        KmerCount {hash: 1, kmer: vec![], count: 16, extra_count: 1},
        KmerCount {hash: 2, kmer: vec![], count: 16, extra_count: 2},
        KmerCount {hash: 3, kmer: vec![], count: 16, extra_count: 8},
        KmerCount {hash: 4, kmer: vec![], count: 16, extra_count: 9},
    ];
    let filtered = filter_strands(&sketch, 0.15);
    assert_eq!(filtered.len(), 2);
    assert_eq!(filtered[0].hash, 3);
    assert_eq!(filtered[1].hash, 4);
}
<|endoftext|>