<commit_before>use neb::ram::schema::Field;
use neb::ram::cell::{MAX_CELL_SIZE, Cell};
use neb::ram::types::{TypeId, Id, Map, Value, id_io, u32_io, key_hash};
use neb::client::transaction::{Transaction, TxnError};

use std::collections::BTreeSet;

use utils::transaction::set_map_by_key_id;

pub const NEXT_KEY: &'static str = "_next";
pub const LIST_KEY: &'static str = "_list";

pub enum IdListError {
    ContainerCellNotFound,
    FormatError,
    Unexpected,
    TxnError(TxnError)
}

pub static ID_LIST_SCHEMA_ID: u32 = 100;

lazy_static! {
    pub static ref ID_LINKED_LIST: Vec<Field> = vec![
            Field::new(&String::from(NEXT_KEY), TypeId::Id as u32, false, false, None),
            Field::new(&String::from(LIST_KEY), TypeId::Id as u32, false, true, None)
          ];
    pub static ref LIST_CAPACITY: usize =
        ((MAX_CELL_SIZE - u32_io::size(0) - id_io::size(0)) / id_io::size(0));
    pub static ref NEXT_KEY_ID: u64 = key_hash(&String::from(NEXT_KEY));
    pub static ref LIST_KEY_ID: u64 = key_hash(&String::from(LIST_KEY));
    pub static ref NEXT_KEY_ID_VEC: Vec<u64> = vec![*NEXT_KEY_ID];
}

pub struct IdList<'a> {
    txn: &'a mut Transaction,
    container_id: Id,
    field_id: u64
}

fn empty_list_segment(container_id: &Id, level: usize) -> (Id, Value) {
    let str_id = format!("IDLIST-{},{}-{}", container_id.higher, container_id.lower, level);
    let list_id = Id::new(container_id.higher, key_hash(&str_id));
    let mut list_map = Map::new();
    list_map.insert_key_id(*NEXT_KEY_ID, Value::Id(Id::unit_id()));
    list_map.insert_key_id(*LIST_KEY_ID, Value::Array(Vec::<Value>::new()));
    return (list_id, Value::Map(list_map));
}

fn count_cell_list(seg: &Cell) -> Result<usize, IdListError> {
    if let &Value::Map(ref map) = &seg.data {
        if let &Value::Array(ref array) = map.get_by_key_id(*LIST_KEY_ID) {
            Ok(array.len())
        } else {
            Err(IdListError::FormatError)
        }
    } else {
        Err(IdListError::FormatError)
    }
}

fn val_is_id(val: &Value, id: &Id) -> bool {
    if let &Value::Id(ref val_id) = val {
        return val_id != id;
    } else {
        return true;
    }
}

fn seg_cell_by_id(txn: &mut Transaction, id: Option<Id>) -> Result<Option<Cell>, IdListError> {
    match id {
        Some(id) => match txn.read(&id) {
            Ok(cell) => Ok(cell),
            Err(e) => Err(IdListError::TxnError(e))
        },
        None => Ok(None)
    }
}

impl<'a> IdList <'a> {
    pub fn from_txn_and_container(txn: &'a mut Transaction, container_id: &Id, field_id: u64) -> IdList<'a> {
        IdList {
            txn: txn,
            container_id: *container_id,
            field_id: field_id,
        }
    }
    fn get_root_list_id(&mut self, ensure_container: bool) -> Result<Id, IdListError> {
        match self.txn.read_selected(&self.container_id, &vec![self.field_id]) {
            Err(e) => Err(IdListError::TxnError(e)),
            Ok(Some(fields)) => {
                if let Some(&Value::Id(id)) = fields.get(0) {
                    if id == Id::unit_id() && ensure_container {
                        let (list_id, list_value) = empty_list_segment(&self.container_id, 0);
                        let list_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &list_id, list_value);
                        match self.txn.write(&list_cell) {
                            Ok(()) => {},
                            Err(e) => return Err(IdListError::TxnError(e))
                        }
                        set_map_by_key_id(self.txn, &self.container_id, self.field_id, Value::Id(list_id)).map_err(IdListError::TxnError)?;
                        Ok(list_id)
                    } else {
                        Ok(id)
                    }
                } else {
                    Err(IdListError::FormatError)
                }
            },
            Ok(None) => {
                Err(IdListError::ContainerCellNotFound)
            }
        }
    }
    pub fn iter(&mut self) -> Result<IdListIterator, IdListError> {
        let list_root_id = self.get_root_list_id(false)?;
        let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
        let current_seg = segments.next();
        Ok(IdListIterator {
            segments: segments,
            current_seg: current_seg,
            current_pos: 0,
        })
    }
    pub fn all(&mut self) -> Result<Vec<Id>, IdListError> {
        Ok(self.iter()?.collect())
    }
    pub fn count(&mut self) -> Result<usize, IdListError> {
        Ok(self.iter()?.count())
    }
    pub fn add(&mut self, id: &Id) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let mut list_level = 0;
        let mut last_seg = {
            let last_seg_id = {
                let mut segments = IdListSegmentIdIterator::new(&mut self.txn, list_root_id);
                let mut last_seg_id = None;
                for seg in segments {
                    list_level += 1;
                    last_seg_id = Some(seg);
                }
                last_seg_id
            };
            let last_seg = seg_cell_by_id(&mut self.txn, last_seg_id)?;
            if let Some(seg) = last_seg { seg } else { return Err(IdListError::Unexpected); }
        };
        if count_cell_list(&mut last_seg)? >= *LIST_CAPACITY { // create new segment to prevent cell overflow
            list_level += 1;
            let (next_seg_id, next_seg_value) = empty_list_segment(&self.container_id, list_level);
            let next_seg_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &next_seg_id, next_seg_value);
            match self.txn.write(&next_seg_cell) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
            set_map_by_key_id(&mut self.txn, &last_seg.id(), *NEXT_KEY_ID, Value::Id(next_seg_id)).map_err(IdListError::TxnError)?;
            last_seg = next_seg_cell;
        }
        if let &mut Value::Map(ref mut map) = &mut last_seg.data {
            if let &mut Value::Array(ref mut array) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                array.push(Value::Id(*id));
            } else {
                return Err(IdListError::FormatError);
            }
        } else {
            return Err(IdListError::FormatError);
        }
        match self.txn.update(&last_seg) {
            Ok(()) => Ok(()),
            Err(e) => Err(IdListError::TxnError(e))
        }
    }
    pub fn remove(&mut self, id: &Id, all: bool) -> Result<(), IdListError> {
        let id_value = Value::Id(*id);
        let mut contained_segs = { // collect affected segment cell ids
            let mut iter = self.iter()?;
            let mut seg_ids = BTreeSet::new();
            while true {
                if let Some(iter_id) = iter.next() {
                    if iter_id == *id {
                        if let Some(ref seg) = iter.current_seg {
                            seg_ids.insert(seg.id());
                        } else {
                            return Err(IdListError::Unexpected);
                        }
                    }
                } else {
                    break;
                }
            }
            seg_ids
        };
        for seg_id in &contained_segs { // mutate cell array
            match self.txn.read(seg_id) {
                Ok(Some(mut seg)) => {
                    if let &mut Value::Map(ref mut map) = &mut seg.data {
                        if let &mut Value::Array(ref mut array) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                            if all {
                                array.retain(|v| { !val_is_id(v, id) });
                            } else {
                                let index = match array.iter().position(|v| { val_is_id(v, id) }) {
                                    Some(pos) => pos, None => return Err(IdListError::Unexpected)
                                };
                                array.remove(index);
                            }
                        } else {
                            return Err(IdListError::FormatError);
                        }
                    } else {
                        return Err(IdListError::FormatError);
                    }
                    self.txn.update(&seg);
                    if !all { break; }
                },
                Ok(None) => return Err(IdListError::Unexpected),
                Err(e) => return Err(IdListError::TxnError(e)),
            }
        }
        return Ok(());
    }
    pub fn clear(&mut self) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let segments: Vec<_> = IdListSegmentIdIterator::new(&mut self.txn, list_root_id).collect();
        for seg_id in segments {
            match self.txn.remove(&seg_id) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
        }
        set_map_by_key_id(self.txn, &self.container_id, self.field_id, Value::Id(Id::unit_id()));
        return Ok(())
    }
}

pub struct IdListSegmentIdIterator<'a> {
    txn: &'a mut Transaction,
    next: Id,
    level: u32
}

impl <'a> IdListSegmentIdIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIdIterator<'a> {
        IdListSegmentIdIterator {
            txn: txn,
            next: head_id,
            level: 1
        }
    }
}

impl <'a> Iterator for IdListSegmentIdIterator<'a> {

    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.next.is_unit_id() {
            match self.txn.read_selected(&self.next, &*NEXT_KEY_ID_VEC) {
                Ok(Some(fields)) => {
                    let current_id = self.next;
                    if let Some(&Value::Id(ref id)) = fields.get(0) {
                        self.next = *id;
                        self.level += 1;
                        return Some(current_id);
                    }
                },
                _ => {}
            }
        }
        None
    }
}

pub struct IdListSegmentIterator<'a> {
    id_iter: IdListSegmentIdIterator<'a>
}

impl <'a>IdListSegmentIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIterator<'a> {
        IdListSegmentIterator {
            id_iter: IdListSegmentIdIterator::new(txn, head_id)
        }
    }
}

impl <'a> Iterator for IdListSegmentIterator <'a> {
    type Item = Cell;

    fn next(&mut self) -> Option<Self::Item> {
        let next_id = self.id_iter.next();
        if let Ok(Some(cell)) = seg_cell_by_id(self.id_iter.txn, next_id) {
            Some(cell)
        } else {
            None
        }
    }
}

pub struct IdListIterator<'a> {
    segments: IdListSegmentIterator<'a>,
    current_seg: Option<Cell>,
    current_pos: u32
}

impl <'a> IdListIterator <'a> {
    pub fn next_seg(&mut self) {
        self.current_seg = self.segments.next();
        self.current_pos = 0;
    }
}

impl <'a> Iterator for IdListIterator<'a> {
    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        let mut need_next_seg = false;
        if let Some(ref cell) = self.current_seg {
            let pos = self.current_pos;
            self.current_pos += 1;
            if let &Value::Map(ref map) = &cell.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.get(pos as usize) {
                        return Some(id);
                    } else {
                        need_next_seg = true
                    }
                }
            }
        };
        if need_next_seg {
            self.next_seg();
            self.next()
        } else {
            None
        }
    }

    fn last(self) -> Option<Self::Item> where Self: Sized {
        if let Some(last_seg) = self.segments.last() {
            if let &Value::Map(ref map) = &last_seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.last() {
                        return Some(id);
                    }
                }
            }
        }
        return None;
    }
    fn count(self) -> usize where Self: Sized {
        let mut count = 0;
        for seg in self.segments {
            if let &Value::Map(ref map) = &seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    count += list.len();
                }
            }
        }
        return count;
    }
}
<commit_msg>defining id type list<commit_after>use neb::ram::schema::Field;
use neb::ram::cell::{MAX_CELL_SIZE, Cell};
use neb::ram::types::{TypeId, Id, Map, Value, id_io, u32_io, key_hash};
use neb::client::transaction::{Transaction, TxnError};

use std::collections::BTreeSet;

use utils::transaction::set_map_by_key_id;

pub const NEXT_KEY: &'static str = "_next";
pub const LIST_KEY: &'static str = "_list";

pub const ID_TYPES_MAP_KEY: &'static str = "_edges";
pub const ID_TYPE_SCHEMA_ID_KEY: &'static str = "_type";
pub const ID_TYPE_ID_LIST_KEY: &'static str = "_list";

pub enum IdListError {
    ContainerCellNotFound,
    FormatError,
    Unexpected,
    TxnError(TxnError)
}

pub static ID_LIST_SCHEMA_ID: u32 = 100;

lazy_static! {
    pub static ref ID_TYPE_LIST: Vec<Field> = vec![
        Field::new(&String::from(ID_TYPES_MAP_KEY), TypeId::Map as u32, false, true,
            Some(vec![
                Field::new(&String::from(ID_TYPE_SCHEMA_ID_KEY), TypeId::U32 as u32, false, false, None),
                Field::new(&String::from(ID_TYPE_ID_LIST_KEY), TypeId::Id as u32, false, true, None)
            ]))
    ];
    pub static ref ID_LINKED_LIST: Vec<Field> = vec![
        Field::new(&String::from(NEXT_KEY), TypeId::Id as u32, false, false, None),
        Field::new(&String::from(LIST_KEY), TypeId::Id as u32, false, true, None)
    ];
    pub static ref LIST_CAPACITY: usize =
        ((MAX_CELL_SIZE - u32_io::size(0) - id_io::size(0)) / id_io::size(0));
    pub static ref NEXT_KEY_ID: u64 = key_hash(&String::from(NEXT_KEY));
    pub static ref LIST_KEY_ID: u64 = key_hash(&String::from(LIST_KEY));
    pub static ref NEXT_KEY_ID_VEC: Vec<u64> = vec![*NEXT_KEY_ID];
}

pub struct IdList<'a> {
    txn: &'a mut Transaction,
    container_id: Id,
    field_id: u64
}

fn empty_list_segment(container_id: &Id, level: usize) -> (Id, Value) {
    let str_id = format!("IDLIST-{},{}-{}", container_id.higher, container_id.lower, level);
    let list_id = Id::new(container_id.higher, key_hash(&str_id));
    let mut list_map = Map::new();
    list_map.insert_key_id(*NEXT_KEY_ID, Value::Id(Id::unit_id()));
    list_map.insert_key_id(*LIST_KEY_ID, Value::Array(Vec::<Value>::new()));
    return (list_id, Value::Map(list_map));
}

fn count_cell_list(seg: &Cell) -> Result<usize, IdListError> {
    if let &Value::Map(ref map) = &seg.data {
        if let &Value::Array(ref array) = map.get_by_key_id(*LIST_KEY_ID) {
            Ok(array.len())
        } else {
            Err(IdListError::FormatError)
        }
    } else {
        Err(IdListError::FormatError)
    }
}

fn val_is_id(val: &Value, id: &Id) -> bool {
    if let &Value::Id(ref val_id) = val {
        return val_id != id;
    } else {
        return true;
    }
}

fn seg_cell_by_id(txn: &mut Transaction, id: Option<Id>) -> Result<Option<Cell>, IdListError> {
    match id {
        Some(id) => match txn.read(&id) {
            Ok(cell) => Ok(cell),
            Err(e) => Err(IdListError::TxnError(e))
        },
        None => Ok(None)
    }
}

impl<'a> IdList <'a> {
    pub fn from_txn_and_container(txn: &'a mut Transaction, container_id: &Id, field_id: u64) -> IdList<'a> {
        IdList {
            txn: txn,
            container_id: *container_id,
            field_id: field_id,
        }
    }
    fn get_root_list_id(&mut self, ensure_container: bool) -> Result<Id, IdListError> {
        match self.txn.read_selected(&self.container_id, &vec![self.field_id]) {
            Err(e) => Err(IdListError::TxnError(e)),
            Ok(Some(fields)) => {
                if let Some(&Value::Id(id)) = fields.get(0) {
                    if id == Id::unit_id() && ensure_container {
                        let (list_id, list_value) = empty_list_segment(&self.container_id, 0);
                        let list_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &list_id, list_value);
                        match self.txn.write(&list_cell) {
                            Ok(()) => {},
                            Err(e) => return Err(IdListError::TxnError(e))
                        }
                        set_map_by_key_id(self.txn, &self.container_id, self.field_id, Value::Id(list_id)).map_err(IdListError::TxnError)?;
                        Ok(list_id)
                    } else {
                        Ok(id)
                    }
                } else {
                    Err(IdListError::FormatError)
                }
            },
            Ok(None) => {
                Err(IdListError::ContainerCellNotFound)
            }
        }
    }
    pub fn iter(&mut self) -> Result<IdListIterator, IdListError> {
        let list_root_id = self.get_root_list_id(false)?;
        let mut segments = IdListSegmentIterator::new(&mut self.txn, list_root_id);
        let current_seg = segments.next();
        Ok(IdListIterator {
            segments: segments,
            current_seg: current_seg,
            current_pos: 0,
        })
    }
    pub fn all(&mut self) -> Result<Vec<Id>, IdListError> {
        Ok(self.iter()?.collect())
    }
    pub fn count(&mut self) -> Result<usize, IdListError> {
        Ok(self.iter()?.count())
    }
    pub fn add(&mut self, id: &Id) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let mut list_level = 0;
        let mut last_seg = {
            let last_seg_id = {
                let mut segments = IdListSegmentIdIterator::new(&mut self.txn, list_root_id);
                let mut last_seg_id = None;
                for seg in segments {
                    list_level += 1;
                    last_seg_id = Some(seg);
                }
                last_seg_id
            };
            let last_seg = seg_cell_by_id(&mut self.txn, last_seg_id)?;
            if let Some(seg) = last_seg { seg } else { return Err(IdListError::Unexpected); }
        };
        if count_cell_list(&mut last_seg)? >= *LIST_CAPACITY { // create new segment to prevent cell overflow
            list_level += 1;
            let (next_seg_id, next_seg_value) = empty_list_segment(&self.container_id, list_level);
            let next_seg_cell = Cell::new_with_id(ID_LIST_SCHEMA_ID, &next_seg_id, next_seg_value);
            match self.txn.write(&next_seg_cell) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
            set_map_by_key_id(&mut self.txn, &last_seg.id(), *NEXT_KEY_ID, Value::Id(next_seg_id)).map_err(IdListError::TxnError)?;
            last_seg = next_seg_cell;
        }
        if let &mut Value::Map(ref mut map) = &mut last_seg.data {
            if let &mut Value::Array(ref mut array) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                array.push(Value::Id(*id));
            } else {
                return Err(IdListError::FormatError);
            }
        } else {
            return Err(IdListError::FormatError);
        }
        match self.txn.update(&last_seg) {
            Ok(()) => Ok(()),
            Err(e) => Err(IdListError::TxnError(e))
        }
    }
    pub fn remove(&mut self, id: &Id, all: bool) -> Result<(), IdListError> {
        let id_value = Value::Id(*id);
        let mut contained_segs = { // collect affected segment cell ids
            let mut iter = self.iter()?;
            let mut seg_ids = BTreeSet::new();
            while true {
                if let Some(iter_id) = iter.next() {
                    if iter_id == *id {
                        if let Some(ref seg) = iter.current_seg {
                            seg_ids.insert(seg.id());
                        } else {
                            return Err(IdListError::Unexpected);
                        }
                    }
                } else {
                    break;
                }
            }
            seg_ids
        };
        for seg_id in &contained_segs { // mutate cell array
            match self.txn.read(seg_id) {
                Ok(Some(mut seg)) => {
                    if let &mut Value::Map(ref mut map) = &mut seg.data {
                        if let &mut Value::Array(ref mut array) = map.get_mut_by_key_id(*LIST_KEY_ID) {
                            if all {
                                array.retain(|v| { !val_is_id(v, id) });
                            } else {
                                let index = match array.iter().position(|v| { val_is_id(v, id) }) {
                                    Some(pos) => pos, None => return Err(IdListError::Unexpected)
                                };
                                array.remove(index);
                            }
                        } else {
                            return Err(IdListError::FormatError);
                        }
                    } else {
                        return Err(IdListError::FormatError);
                    }
                    self.txn.update(&seg);
                    if !all { break; }
                },
                Ok(None) => return Err(IdListError::Unexpected),
                Err(e) => return Err(IdListError::TxnError(e)),
            }
        }
        return Ok(());
    }
    pub fn clear(&mut self) -> Result<(), IdListError> {
        let list_root_id = self.get_root_list_id(true)?;
        let segments: Vec<_> = IdListSegmentIdIterator::new(&mut self.txn, list_root_id).collect();
        for seg_id in segments {
            match self.txn.remove(&seg_id) {
                Ok(()) => {},
                Err(e) => return Err(IdListError::TxnError(e))
            }
        }
        set_map_by_key_id(self.txn, &self.container_id, self.field_id, Value::Id(Id::unit_id()));
        return Ok(())
    }
}

pub struct IdListSegmentIdIterator<'a> {
    txn: &'a mut Transaction,
    next: Id,
    level: u32
}

impl <'a> IdListSegmentIdIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIdIterator<'a> {
        IdListSegmentIdIterator {
            txn: txn,
            next: head_id,
            level: 1
        }
    }
}

impl <'a> Iterator for IdListSegmentIdIterator<'a> {

    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.next.is_unit_id() {
            match self.txn.read_selected(&self.next, &*NEXT_KEY_ID_VEC) {
                Ok(Some(fields)) => {
                    let current_id = self.next;
                    if let Some(&Value::Id(ref id)) = fields.get(0) {
                        self.next = *id;
                        self.level += 1;
                        return Some(current_id);
                    }
                },
                _ => {}
            }
        }
        None
    }
}

pub struct IdListSegmentIterator<'a> {
    id_iter: IdListSegmentIdIterator<'a>
}

impl <'a>IdListSegmentIterator<'a> {
    pub fn new(txn: &'a mut Transaction, head_id: Id) -> IdListSegmentIterator<'a> {
        IdListSegmentIterator {
            id_iter: IdListSegmentIdIterator::new(txn, head_id)
        }
    }
}

impl <'a> Iterator for IdListSegmentIterator <'a> {
    type Item = Cell;

    fn next(&mut self) -> Option<Self::Item> {
        let next_id = self.id_iter.next();
        if let Ok(Some(cell)) = seg_cell_by_id(self.id_iter.txn, next_id) {
            Some(cell)
        } else {
            None
        }
    }
}

pub struct IdListIterator<'a> {
    segments: IdListSegmentIterator<'a>,
    current_seg: Option<Cell>,
    current_pos: u32
}

impl <'a> IdListIterator <'a> {
    pub fn next_seg(&mut self) {
        self.current_seg = self.segments.next();
        self.current_pos = 0;
    }
}

impl <'a> Iterator for IdListIterator<'a> {
    type Item = Id;

    fn next(&mut self) -> Option<Self::Item> {
        let mut need_next_seg = false;
        if let Some(ref cell) = self.current_seg {
            let pos = self.current_pos;
            self.current_pos += 1;
            if let &Value::Map(ref map) = &cell.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.get(pos as usize) {
                        return Some(id);
                    } else {
                        need_next_seg = true
                    }
                }
            }
        };
        if need_next_seg {
            self.next_seg();
            self.next()
        } else {
            None
        }
    }

    fn last(self) -> Option<Self::Item> where Self: Sized {
        if let Some(last_seg) = self.segments.last() {
            if let &Value::Map(ref map) = &last_seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    if let Some(&Value::Id(id)) = list.last() {
                        return Some(id);
                    }
                }
            }
        }
        return None;
    }
    fn count(self) -> usize where Self: Sized {
        let mut count = 0;
        for seg in self.segments {
            if let &Value::Map(ref map) = &seg.data {
                if let &Value::Array(ref list) = map.get_by_key_id(*LIST_KEY_ID) {
                    count += list.len();
                }
            }
        }
        return count;
    }
}
<|endoftext|>


<commit_before>use super::internal::*;
use super::*;
use std::cmp::min;
use std::iter::Chain;
use std::option::IntoIter as OptionIter;
use std::str::{Chars, Split};

/// Test if a byte is the start of a UTF-8 character.
/// (extracted from `str::is_char_boundary`)
fn is_char_boundary(b: u8) -> bool {
    // This is bit magic equivalent to: b < 128 || b >= 192
    (b as i8) >= -0x40
}

/// Find the index of a character boundary near the midpoint.
fn find_char_midpoint(chars: &str) -> usize {
    let mid = chars.len() / 2;

    // We want to split near the midpoint, but we need to find an actual
    // character boundary.  So we look at the raw bytes, first scanning
    // forward from the midpoint for a boundary, then trying backward.
    let (left, right) = chars.as_bytes().split_at(mid);
    right.iter()
        .cloned()
        .position(is_char_boundary)
        .map(|i| mid + i)
        .or_else(|| left.iter().cloned().rposition(is_char_boundary))
        .unwrap_or(0)
}

/// Parallel extensions for strings.
pub trait ParallelString {
    /// Returns a parallel iterator over the characters of a string.
    fn par_chars(&self) -> ParChars;

    /// Returns a parallel iterator over substrings separated by a
    /// given character, similar to `str::split`.
    fn par_split(&self, char) -> ParSplit;

    /// Returns a parallel iterator over substrings terminated by a
    /// given character, similar to `str::split_terminator`.  It's
    /// equivalent to `par_split`, except it doesn't produce an empty
    /// substring after a trailing terminator.
    fn par_split_terminator(&self, char) -> ParSplitTerminator;

    /// Returns a parallel iterator over the lines of a string, ending with an
    /// optional carriage return and with a newline (`\r\n` or just `\n`).
    /// The final line ending is optional, and line endings are not included in
    /// the output strings.
    fn par_lines(&self) -> ParLines;
}

impl ParallelString for str {
    fn par_chars(&self) -> ParChars {
        ParChars { chars: self }
    }

    fn par_split(&self, separator: char) -> ParSplit {
        ParSplit::new(self, separator)
    }

    fn par_split_terminator(&self, terminator: char) -> ParSplitTerminator {
        ParSplitTerminator::new(self, terminator)
    }

    fn par_lines(&self) -> ParLines {
        ParLines(self)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParChars<'a> {
    chars: &'a str,
}

impl<'a> ParallelIterator for ParChars<'a> {
    type Item = char;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParChars<'a> {
    fn split(&mut self) -> Option<Self> {
        let index = find_char_midpoint(self.chars);
        if index > 0 {
            let (left, right) = self.chars.split_at(index);
            self.chars = left;
            Some(ParChars { chars: right })
        } else {
            None
        }
    }
}

impl<'a> IntoIterator for ParChars<'a> {
    type Item = char;
    type IntoIter = Chars<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.chars.chars()
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplit<'a> {
    chars: &'a str,
    separator: char,

    /// Marks the endpoint beyond which we've already found no separators.
    tail: usize,
}

impl<'a> ParSplit<'a> {
    fn new(chars: &'a str, separator: char) -> Self {
        ParSplit {
            chars: chars,
            separator: separator,
            tail: chars.len(),
        }
    }
}

impl<'a> ParallelIterator for ParSplit<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParSplit<'a> {
    fn split(&mut self) -> Option<Self> {
        let ParSplit { chars, separator, tail } = *self;

        // First find a suitable UTF-8 boundary in the unsearched region.
        let char_index = find_char_midpoint(&chars[..tail]);

        // Look forward for the separator, and failing that look backward.
        let index = chars[char_index..tail]
            .find(separator)
            .map(|i| char_index + i)
            .or_else(|| chars[..char_index].rfind(separator));

        if let Some(index) = index {
            // Update `self` as the region before the separator.
            self.chars = &chars[..index];
            self.tail = min(char_index, index);

            // Create the right split following the separator.
            let right_index = index + separator.len_utf8();
            let mut right = ParSplit {
                chars: &chars[right_index..],
                separator: separator,
                tail: tail - right_index,
            };

            // If we scanned backwards to find the separator, everything in
            // the right side is exhausted, with no separators left to find.
            if index < char_index {
                right.tail = 0;
            }

            Some(right)

        } else {
            self.tail = 0;
            None
        }
    }
}

impl<'a> IntoIterator for ParSplit<'a> {
    type Item = &'a str;
    type IntoIter = Chain<Split<'a, char>, OptionIter<&'a str>>;

    fn into_iter(self) -> Self::IntoIter {
        let ParSplit { chars, separator, tail } = self;

        if tail == chars.len() {
            // No tail section, so just let `str::split` handle it.
            chars.split(separator).chain(None)

        } else if let Some(index) = chars[..tail].rfind(separator) {
            // We found the last separator to complete the tail, so
            // end with that slice after `str::split` finds the rest.
            let head = &chars[..index];
            let last = &chars[index + separator.len_utf8()..];
            head.split(separator).chain(Some(last))

        } else {
            // We know there are no separators at all.  Return our whole string,
            // but for type correctness we need to chain an emptied `Split` too.
            let mut empty = "".split('\0');
            Iterator::last(&mut empty);
            empty.chain(Some(chars))
        }
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitTerminator<'a> {
    splitter: ParSplit<'a>,
    endpoint: bool,
}

impl<'a> ParSplitTerminator<'a> {
    fn new(chars: &'a str, terminator: char) -> Self {
        ParSplitTerminator {
            splitter: ParSplit::new(chars, terminator),
            endpoint: true,
        }
    }
}

impl<'a> ParallelIterator for ParSplitTerminator<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParSplitTerminator<'a> {
    fn split(&mut self) -> Option<Self> {
        self.splitter.split().map(|right| {
            let endpoint = self.endpoint;
            self.endpoint = false;
            ParSplitTerminator { splitter: right, endpoint: endpoint }
        })
    }
}

impl<'a> IntoIterator for ParSplitTerminator<'a> {
    type Item = &'a str;
    type IntoIter = Chain<Split<'a, char>, OptionIter<&'a str>>;

    fn into_iter(self) -> Self::IntoIter {
        let skip_last = if self.endpoint {
            let chars = self.splitter.chars;
            let terminator = self.splitter.separator;
            chars.is_empty() || chars.ends_with(terminator)
        } else {
            false
        };

        let mut iter = self.splitter.into_iter();

        if skip_last {
            // eat the empty trailing substring
            iter.next_back();
        }

        iter
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParLines<'a>(&'a str);

impl<'a> ParallelIterator for ParLines<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0.par_split_terminator('\n').map(|line| {
            if line.ends_with('\r') {
                &line[..line.len() - 1]
            } else {
                line
            }
        }).drive_unindexed(consumer)
    }
}
<commit_msg>Run cargo fmt on string.rs<commit_after>use super::internal::*;
use super::*;
use std::cmp::min;
use std::iter::Chain;
use std::option::IntoIter as OptionIter;
use std::str::{Chars, Split};

/// Test if a byte is the start of a UTF-8 character.
/// (extracted from `str::is_char_boundary`)
fn is_char_boundary(b: u8) -> bool {
    // This is bit magic equivalent to: b < 128 || b >= 192
    (b as i8) >= -0x40
}

/// Find the index of a character boundary near the midpoint.
fn find_char_midpoint(chars: &str) -> usize {
    let mid = chars.len() / 2;

    // We want to split near the midpoint, but we need to find an actual
    // character boundary.  So we look at the raw bytes, first scanning
    // forward from the midpoint for a boundary, then trying backward.
    let (left, right) = chars.as_bytes().split_at(mid);
    right.iter()
        .cloned()
        .position(is_char_boundary)
        .map(|i| mid + i)
        .or_else(|| left.iter().cloned().rposition(is_char_boundary))
        .unwrap_or(0)
}

/// Parallel extensions for strings.
pub trait ParallelString {
    /// Returns a parallel iterator over the characters of a string.
    fn par_chars(&self) -> ParChars;

    /// Returns a parallel iterator over substrings separated by a
    /// given character, similar to `str::split`.
    fn par_split(&self, char) -> ParSplit;

    /// Returns a parallel iterator over substrings terminated by a
    /// given character, similar to `str::split_terminator`.  It's
    /// equivalent to `par_split`, except it doesn't produce an empty
    /// substring after a trailing terminator.
    fn par_split_terminator(&self, char) -> ParSplitTerminator;

    /// Returns a parallel iterator over the lines of a string, ending with an
    /// optional carriage return and with a newline (`\r\n` or just `\n`).
    /// The final line ending is optional, and line endings are not included in
    /// the output strings.
    fn par_lines(&self) -> ParLines;
}

impl ParallelString for str {
    fn par_chars(&self) -> ParChars {
        ParChars { chars: self }
    }

    fn par_split(&self, separator: char) -> ParSplit {
        ParSplit::new(self, separator)
    }

    fn par_split_terminator(&self, terminator: char) -> ParSplitTerminator {
        ParSplitTerminator::new(self, terminator)
    }

    fn par_lines(&self) -> ParLines {
        ParLines(self)
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParChars<'a> {
    chars: &'a str,
}

impl<'a> ParallelIterator for ParChars<'a> {
    type Item = char;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParChars<'a> {
    fn split(&mut self) -> Option<Self> {
        let index = find_char_midpoint(self.chars);
        if index > 0 {
            let (left, right) = self.chars.split_at(index);
            self.chars = left;
            Some(ParChars { chars: right })
        } else {
            None
        }
    }
}

impl<'a> IntoIterator for ParChars<'a> {
    type Item = char;
    type IntoIter = Chars<'a>;

    fn into_iter(self) -> Self::IntoIter {
        self.chars.chars()
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplit<'a> {
    chars: &'a str,
    separator: char,

    /// Marks the endpoint beyond which we've already found no separators.
    tail: usize,
}

impl<'a> ParSplit<'a> {
    fn new(chars: &'a str, separator: char) -> Self {
        ParSplit {
            chars: chars,
            separator: separator,
            tail: chars.len(),
        }
    }
}

impl<'a> ParallelIterator for ParSplit<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParSplit<'a> {
    fn split(&mut self) -> Option<Self> {
        let ParSplit { chars, separator, tail } = *self;

        // First find a suitable UTF-8 boundary in the unsearched region.
        let char_index = find_char_midpoint(&chars[..tail]);

        // Look forward for the separator, and failing that look backward.
        let index = chars[char_index..tail]
            .find(separator)
            .map(|i| char_index + i)
            .or_else(|| chars[..char_index].rfind(separator));

        if let Some(index) = index {
            // Update `self` as the region before the separator.
            self.chars = &chars[..index];
            self.tail = min(char_index, index);

            // Create the right split following the separator.
            let right_index = index + separator.len_utf8();
            let mut right = ParSplit {
                chars: &chars[right_index..],
                separator: separator,
                tail: tail - right_index,
            };

            // If we scanned backwards to find the separator, everything in
            // the right side is exhausted, with no separators left to find.
            if index < char_index {
                right.tail = 0;
            }

            Some(right)

        } else {
            self.tail = 0;
            None
        }
    }
}

impl<'a> IntoIterator for ParSplit<'a> {
    type Item = &'a str;
    type IntoIter = Chain<Split<'a, char>, OptionIter<&'a str>>;

    fn into_iter(self) -> Self::IntoIter {
        let ParSplit { chars, separator, tail } = self;

        if tail == chars.len() {
            // No tail section, so just let `str::split` handle it.
            chars.split(separator).chain(None)

        } else if let Some(index) = chars[..tail].rfind(separator) {
            // We found the last separator to complete the tail, so
            // end with that slice after `str::split` finds the rest.
            let head = &chars[..index];
            let last = &chars[index + separator.len_utf8()..];
            head.split(separator).chain(Some(last))

        } else {
            // We know there are no separators at all.  Return our whole string,
            // but for type correctness we need to chain an emptied `Split` too.
            let mut empty = "".split('\0');
            Iterator::last(&mut empty);
            empty.chain(Some(chars))
        }
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParSplitTerminator<'a> {
    splitter: ParSplit<'a>,
    endpoint: bool,
}

impl<'a> ParSplitTerminator<'a> {
    fn new(chars: &'a str, terminator: char) -> Self {
        ParSplitTerminator {
            splitter: ParSplit::new(chars, terminator),
            endpoint: true,
        }
    }
}

impl<'a> ParallelIterator for ParSplitTerminator<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        bridge_unindexed(self, consumer)
    }
}

impl<'a> UnindexedProducer for ParSplitTerminator<'a> {
    fn split(&mut self) -> Option<Self> {
        self.splitter.split().map(|right| {
            let endpoint = self.endpoint;
            self.endpoint = false;
            ParSplitTerminator {
                splitter: right,
                endpoint: endpoint,
            }
        })
    }
}

impl<'a> IntoIterator for ParSplitTerminator<'a> {
    type Item = &'a str;
    type IntoIter = Chain<Split<'a, char>, OptionIter<&'a str>>;

    fn into_iter(self) -> Self::IntoIter {
        let skip_last = if self.endpoint {
            let chars = self.splitter.chars;
            let terminator = self.splitter.separator;
            chars.is_empty() || chars.ends_with(terminator)
        } else {
            false
        };

        let mut iter = self.splitter.into_iter();

        if skip_last {
            // eat the empty trailing substring
            iter.next_back();
        }

        iter
    }
}

// /////////////////////////////////////////////////////////////////////////

pub struct ParLines<'a>(&'a str);

impl<'a> ParallelIterator for ParLines<'a> {
    type Item = &'a str;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
        where C: UnindexedConsumer<Self::Item>
    {
        self.0
            .par_split_terminator('\n')
            .map(|line| {
                if line.ends_with('\r') {
                    &line[..line.len() - 1]
                } else {
                    line
                }
            })
            .drive_unindexed(consumer)
    }
}
<|endoftext|>


<commit_before>//! Tests for various types of file (video, image, compressed, etc).
//!
//! Currently this is dependent on the file’s name and extension, because
//! those are the only metadata that we have access to without reading the
//! file’s contents.

use ansi_term::Style;

use crate::fs::File;
use crate::output::file_name::FileColours;
use crate::output::icons::FileIcon;

#[derive(Debug, Default, PartialEq)]
pub struct FileExtensions;

impl FileExtensions {

    /// An “immediate” file is something that can be run or activated somehow
    /// in order to kick off the build of a project. It’s usually only present
    /// in directories full of source code.
    fn is_immediate(&self, file: &File) -> bool {
        file.name.to_lowercase().starts_with("readme") ||
        file.name.ends_with(".ninja") ||
        file.name_is_one_of( &[
            "Makefile", "Cargo.toml", "SConstruct", "CMakeLists.txt",
            "build.gradle", "Rakefile", "package.json", "Gruntfile.js",
            "Gruntfile.coffee", "BUILD", "BUILD.bazel", "WORKSPACE", "build.xml",
            "webpack.config.js", "meson.build",
        ])
    }

    fn is_image(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "png", "jpeg", "jpg", "gif", "bmp", "tiff", "tif",
            "ppm", "pgm", "pbm", "pnm", "webp", "raw", "arw",
            "svg", "stl", "eps", "dvi", "ps", "cbr", "jpf",
            "cbz", "xpm", "ico", "cr2", "orf", "nef",
        ])
    }

    fn is_video(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "avi", "flv", "m2v", "m4v", "mkv", "mov", "mp4", "mpeg",
            "mpg", "ogm", "ogv", "vob", "wmv", "webm", "m2ts",
        ])
    }

    fn is_music(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "aac", "m4a", "mp3", "ogg", "wma", "mka", "opus",
        ])
    }

    // Lossless music, rather than any other kind of data...
    fn is_lossless(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "alac", "ape", "flac", "wav",
        ])
    }

    fn is_crypto(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "asc", "enc", "gpg", "pgp", "sig", "signature", "pfx", "p12",
        ])
    }

    fn is_document(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "djvu", "doc", "docx", "dvi", "eml", "eps", "fotd",
            "odp", "odt", "pdf", "ppt", "pptx", "rtf",
            "xls", "xlsx",
        ])
    }

    fn is_compressed(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "zip", "tar", "Z", "z", "gz", "bz2", "a", "ar", "7z",
            "iso", "dmg", "tc", "rar", "par", "tgz", "xz", "txz",
            "lz", "tlz", "lzma", "deb", "rpm", "zst",
        ])
    }

    fn is_temp(&self, file: &File) -> bool {
        file.name.ends_with('~')
            || (file.name.starts_with('#') && file.name.ends_with('#'))
            || file.extension_is_one_of( &[ "tmp", "swp", "swo", "swn", "bak", "bk" ])
    }

    fn is_compiled(&self, file: &File) -> bool {
        if file.extension_is_one_of( &[ "class", "elc", "hi", "o", "pyc" ]) {
            true
        }
        else if let Some(dir) = file.parent_dir {
            file.get_source_files().iter().any(|path| dir.contains(path))
        }
        else {
            false
        }
    }
}

impl FileColours for FileExtensions {
    fn colour_file(&self, file: &File) -> Option<Style> {
        use ansi_term::Colour::*;

        Some(match file {
            f if self.is_temp(f)        => Fixed(244).normal(),
            f if self.is_immediate(f)   => Yellow.bold().underline(),
            f if self.is_image(f)       => Fixed(133).normal(),
            f if self.is_video(f)       => Fixed(135).normal(),
            f if self.is_music(f)       => Fixed(92).normal(),
            f if self.is_lossless(f)    => Fixed(93).normal(),
            f if self.is_crypto(f)      => Fixed(109).normal(),
            f if self.is_document(f)    => Fixed(105).normal(),
            f if self.is_compressed(f)  => Red.normal(),
            f if self.is_compiled(f)    => Fixed(137).normal(),
            _                           => return None,
        })
    }
}

impl FileIcon for FileExtensions {
    fn icon_file(&self, file: &File) -> Option<char> {
        use crate::output::icons::Icons;

        Some(match file {
            f if self.is_music(f) || self.is_lossless(f) => Icons::Audio.value(),
            f if self.is_image(f) => Icons::Image.value(),
            f if self.is_video(f) => Icons::Video.value(),
            _ => return None,
        })
    }
}
<commit_msg>Add pom.xml and zwc filetypes<commit_after>//! Tests for various types of file (video, image, compressed, etc).
//!
//! Currently this is dependent on the file’s name and extension, because
//! those are the only metadata that we have access to without reading the
//! file’s contents.

use ansi_term::Style;

use crate::fs::File;
use crate::output::file_name::FileColours;
use crate::output::icons::FileIcon;

#[derive(Debug, Default, PartialEq)]
pub struct FileExtensions;

impl FileExtensions {

    /// An “immediate” file is something that can be run or activated somehow
    /// in order to kick off the build of a project. It’s usually only present
    /// in directories full of source code.
    fn is_immediate(&self, file: &File) -> bool {
        file.name.to_lowercase().starts_with("readme") ||
        file.name.ends_with(".ninja") ||
        file.name_is_one_of( &[
            "Makefile", "Cargo.toml", "SConstruct", "CMakeLists.txt",
            "build.gradle", "pom.xml", "Rakefile", "package.json", "Gruntfile.js",
            "Gruntfile.coffee", "BUILD", "BUILD.bazel", "WORKSPACE", "build.xml",
            "webpack.config.js", "meson.build",
        ])
    }

    fn is_image(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "png", "jpeg", "jpg", "gif", "bmp", "tiff", "tif",
            "ppm", "pgm", "pbm", "pnm", "webp", "raw", "arw",
            "svg", "stl", "eps", "dvi", "ps", "cbr", "jpf",
            "cbz", "xpm", "ico", "cr2", "orf", "nef",
        ])
    }

    fn is_video(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "avi", "flv", "m2v", "m4v", "mkv", "mov", "mp4", "mpeg",
            "mpg", "ogm", "ogv", "vob", "wmv", "webm", "m2ts",
        ])
    }

    fn is_music(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "aac", "m4a", "mp3", "ogg", "wma", "mka", "opus",
        ])
    }

    // Lossless music, rather than any other kind of data...
    fn is_lossless(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "alac", "ape", "flac", "wav",
        ])
    }

    fn is_crypto(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "asc", "enc", "gpg", "pgp", "sig", "signature", "pfx", "p12",
        ])
    }

    fn is_document(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "djvu", "doc", "docx", "dvi", "eml", "eps", "fotd",
            "odp", "odt", "pdf", "ppt", "pptx", "rtf",
            "xls", "xlsx",
        ])
    }

    fn is_compressed(&self, file: &File) -> bool {
        file.extension_is_one_of( &[
            "zip", "tar", "Z", "z", "gz", "bz2", "a", "ar", "7z",
            "iso", "dmg", "tc", "rar", "par", "tgz", "xz", "txz",
            "lz", "tlz", "lzma", "deb", "rpm", "zst",
        ])
    }

    fn is_temp(&self, file: &File) -> bool {
        file.name.ends_with('~')
            || (file.name.starts_with('#') && file.name.ends_with('#'))
            || file.extension_is_one_of( &[ "tmp", "swp", "swo", "swn", "bak", "bk" ])
    }

    fn is_compiled(&self, file: &File) -> bool {
        if file.extension_is_one_of( &[ "class", "elc", "hi", "o", "pyc", "zwc" ]) {
            true
        }
        else if let Some(dir) = file.parent_dir {
            file.get_source_files().iter().any(|path| dir.contains(path))
        }
        else {
            false
        }
    }
}

impl FileColours for FileExtensions {
    fn colour_file(&self, file: &File) -> Option<Style> {
        use ansi_term::Colour::*;

        Some(match file {
            f if self.is_temp(f)        => Fixed(244).normal(),
            f if self.is_immediate(f)   => Yellow.bold().underline(),
            f if self.is_image(f)       => Fixed(133).normal(),
            f if self.is_video(f)       => Fixed(135).normal(),
            f if self.is_music(f)       => Fixed(92).normal(),
            f if self.is_lossless(f)    => Fixed(93).normal(),
            f if self.is_crypto(f)      => Fixed(109).normal(),
            f if self.is_document(f)    => Fixed(105).normal(),
            f if self.is_compressed(f)  => Red.normal(),
            f if self.is_compiled(f)    => Fixed(137).normal(),
            _                           => return None,
        })
    }
}

impl FileIcon for FileExtensions {
    fn icon_file(&self, file: &File) -> Option<char> {
        use crate::output::icons::Icons;

        Some(match file {
            f if self.is_music(f) || self.is_lossless(f) => Icons::Audio.value(),
            f if self.is_image(f) => Icons::Image.value(),
            f if self.is_video(f) => Icons::Video.value(),
            _ => return None,
        })
    }
}
<|endoftext|>


<commit_before>/// This file contains commands to generally control and administer the bot.
extern crate irc;

use irc::Client;
use irc::interface::CommandEvent;

fn say(event: &mut CommandEvent) {
    if event.args[0].starts_with("#") {
        if !event.client.is_admin(event) {
            return;
        }
        event.client.send_message(event.args[0], &*event.args.slice_from(1).connect(" "));
    } else {
        event.client.send_message(event.channel, &*event.args.connect(" "));
    }
}

fn quit(event: &mut CommandEvent) {
    if !event.client.is_admin(event) {
        return
    }
    event.client.send_command("QUIT".to_string(), &[":See you!"]);
}

fn raw(event: &mut CommandEvent) {
    event.client.send_raw(event.args.connect(" "));
}

pub fn register(client: &mut Client) {
    client.add_command("say", say);
    client.add_command("quit", quit);
    client.add_command("raw", raw);
}
<commit_msg>Remove unrequired extern crate irc in plugin<commit_after>/// This file contains commands to generally control and administer the bot.
use irc::Client;
use irc::interface::CommandEvent;

fn say(event: &mut CommandEvent) {
    if event.args[0].starts_with("#") {
        if !event.client.is_admin(event) {
            return;
        }
        event.client.send_message(event.args[0], &*event.args.slice_from(1).connect(" "));
    } else {
        event.client.send_message(event.channel, &*event.args.connect(" "));
    }
}

fn quit(event: &mut CommandEvent) {
    if !event.client.is_admin(event) {
        return
    }
    event.client.send_command("QUIT".to_string(), &[":See you!"]);
}

fn raw(event: &mut CommandEvent) {
    event.client.send_raw(event.args.connect(" "));
}

pub fn register(client: &mut Client) {
    client.add_command("say", say);
    client.add_command("quit", quit);
    client.add_command("raw", raw);
}
<|endoftext|>


<commit_before>// Copyright 2014 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
use std::mem;
use libc::c_int;
use num::complex::{
    Complex32,
    Complex64,
};
use error::Error;
use ll::*;
use matrix::{
    Matrix,
};
use scalar::Scalar;
use types::{Order, Transpose};
use util::ColMem;

pub trait Gels: Sized {
    fn gels(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<(), Error> {
        //TODO: nancheck

        let work_len = try!(Gels::gels_work_len(a_trans, a, b));
        let mut work: Vec<_> = Vec::with_capacity(work_len as usize);
        unsafe {
            work.set_len(work_len as usize);
        }

        Gels::gels_work(layout, a_trans, a, b, &mut work[..])
    }

    fn gels_work(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>, work: &mut [Self]) -> Result<(), Error>;
    fn gels_work_len(a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<usize, Error>;
}

macro_rules! least_sq_impl(($($t: ident), +) => ($(
    impl Gels for $t {
        fn gels_work(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>, work: &mut [Self]) -> Result<(), Error> {
            let mut info: c_int = 0;

            let m = a.rows();
            let n = a.cols();
            let nrhs = b.cols();

            let mut a_mem = ColMem::new(layout, a);
            let mut b_mem = ColMem::new(layout, b);

            let lda = a_mem.lead();
            let ldb = b_mem.lead();

            unsafe {
                prefix!($t, gels_)(
                    a_trans.as_i8().as_mut(),
                    m.as_mut(), n.as_mut(),
                    nrhs.as_mut(),
                    a_mem.as_mut_ptr(), lda.as_mut(),
                    b_mem.as_mut_ptr(), ldb.as_mut(),
                    work.as_mut_ptr(), (work.len() as c_int).as_mut(),
                    &mut info as *mut c_int);
            }

            match info {
                0 => Ok(()),
                x if x < 0 => Err(Error::IllegalParameter(-x as usize)),
                x => Err(Error::DiagonalElementZero(x as usize)),
            }
        }

        fn gels_work_len(a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<usize, Error> {
            let mut info: c_int = 0;
            let mut len_info: $t = unsafe { mem::zeroed() };
            let len_ptr = (&mut len_info) as *mut $t;

            let m = a.rows();
            let n = a.cols();
            let nrhs = b.cols();
            let lda_t = ::std::cmp::max(1, m);
            let ldb_t = ::std::cmp::max(1, ::std::cmp::max(m, n));

            unsafe {
                prefix!($t, gels_)(
                    a_trans.as_i8().as_mut(),
                    m.as_mut(), n.as_mut(),
                    nrhs.as_mut(),
                    a.as_mut_ptr(), lda_t.as_mut(),
                    b.as_mut_ptr(), ldb_t.as_mut(),
                    len_ptr, (-1 as c_int).as_mut(),
                    &mut info as *mut c_int);
            };

            match info {
                0 => Ok(len_info.as_work()),
                x if x < 0 => Err(Error::IllegalParameter(-x as usize)),
                x => Err(Error::DiagonalElementZero(x as usize)),
            }
        }
    }
)+));

least_sq_impl!(f32, f64, Complex32, Complex64);

#[cfg(test)]
mod gels_tests {
    use types::Order::*;
    use types::Transpose;
    use matrix::tests::M;
    use least_squares::Gels;

    #[test]
    fn col_major() {
        let mut a = M(3i32, 2i32, vec![2.0f32,4.0,7.0,3.0,9.0,4.0]);
        let mut b = M(3i32, 2i32, vec![2.0f32,4.0,7.0,6.0,18.0,8.0]);

        Gels::gels(ColMajor, &Transpose::None, &mut a, &mut b).unwrap();

        let M(_, _, x) = b;
        assert_eq!(x, vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0]);
    }

    #[test]
    fn row_major() {
        let mut a = M(3i32, 2i32, vec![2.0f32,3.0,4.0,9.0,7.0,4.0]);
        let mut b = M(3i32, 2i32, vec![2.0f32,3.0,4.0,9.0,7.0,4.0]);

        Gels::gels(RowMajor, &Transpose::None, &mut a, &mut b).unwrap();

        let M(_, _, x) = b;
        assert_eq!(x, vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]);
    }
}
<commit_msg>Clean up<commit_after>// Copyright 2014 Michael Yang. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
use std::mem;
use libc::c_int;
use num::complex::{
    Complex32,
    Complex64,
};
use error::Error;
use ll::*;
use matrix::{
    Matrix,
};
use scalar::Scalar;
use types::{Order, Transpose};
use util::ColMem;

pub trait Gels: Sized {
    fn gels(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<(), Error> {
        //TODO: nancheck

        let work_len = try!(Gels::gels_work_len(a_trans, a, b));
        let mut work: Vec<_> = Vec::with_capacity(work_len as usize);
        unsafe {
            work.set_len(work_len as usize);
        }

        Gels::gels_work(layout, a_trans, a, b, &mut work[..])
    }

    fn gels_work(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>, work: &mut [Self]) -> Result<(), Error>;
    fn gels_work_len(a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<usize, Error>;
}

macro_rules! least_sq_impl(($($t: ident), +) => ($(
    impl Gels for $t {
        fn gels_work(layout: Order, a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>, work: &mut [Self]) -> Result<(), Error> {
            let mut info: c_int = 0;

            let m = a.rows();
            let n = a.cols();
            let nrhs = b.cols();
            let mut a_mem = ColMem::new(layout, a);
            let mut b_mem = ColMem::new(layout, b);

            unsafe {
                prefix!($t, gels_)(
                    a_trans.as_i8().as_mut(),
                    m.as_mut(), n.as_mut(),
                    nrhs.as_mut(),
                    a_mem.as_mut_ptr(), a_mem.lead().as_mut(),
                    b_mem.as_mut_ptr(), b_mem.lead().as_mut(),
                    work.as_mut_ptr(), (work.len() as c_int).as_mut(),
                    &mut info as *mut c_int);
            }

            match info {
                0 => Ok(()),
                x if x < 0 => Err(Error::IllegalParameter(-x as usize)),
                x => Err(Error::DiagonalElementZero(x as usize)),
            }
        }

        fn gels_work_len(a_trans: &Transpose, a: &mut Matrix<Self>, b: &mut Matrix<Self>) -> Result<usize, Error> {
            let mut info: c_int = 0;
            let mut len_info: $t = unsafe { mem::zeroed() };
            let len_ptr = (&mut len_info) as *mut $t;

            let m = a.rows();
            let n = a.cols();
            let nrhs = b.cols();
            let lda_t = ::std::cmp::max(1, m);
            let ldb_t = ::std::cmp::max(1, ::std::cmp::max(m, n));

            unsafe {
                prefix!($t, gels_)(
                    a_trans.as_i8().as_mut(),
                    m.as_mut(), n.as_mut(),
                    nrhs.as_mut(),
                    a.as_mut_ptr(), lda_t.as_mut(),
                    b.as_mut_ptr(), ldb_t.as_mut(),
                    len_ptr, (-1 as c_int).as_mut(),
                    &mut info as *mut c_int);
            };

            match info {
                0 => Ok(len_info.as_work()),
                x if x < 0 => Err(Error::IllegalParameter(-x as usize)),
                x => Err(Error::DiagonalElementZero(x as usize)),
            }
        }
    }
)+));

least_sq_impl!(f32, f64, Complex32, Complex64);

#[cfg(test)]
mod gels_tests {
    use types::Order::*;
    use types::Transpose;
    use matrix::tests::M;
    use least_squares::Gels;

    #[test]
    fn col_major() {
        let mut a = M(3i32, 2i32, vec![2.0f32,4.0,7.0,3.0,9.0,4.0]);
        let mut b = M(3i32, 2i32, vec![2.0f32,4.0,7.0,6.0,18.0,8.0]);

        Gels::gels(ColMajor, &Transpose::None, &mut a, &mut b).unwrap();

        let M(_, _, x) = b;
        assert_eq!(x, vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0]);
    }

    #[test]
    fn row_major() {
        let mut a = M(3i32, 2i32, vec![2.0f32,3.0,4.0,9.0,7.0,4.0]);
        let mut b = M(3i32, 2i32, vec![2.0f32,3.0,4.0,9.0,7.0,4.0]);

        Gels::gels(RowMajor, &Transpose::None, &mut a, &mut b).unwrap();

        let M(_, _, x) = b;
        assert_eq!(x, vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]);
    }
}
<|endoftext|>


<commit_before>
extern crate regex;

use std::process::{Command, Stdio, Child};
use std::io::prelude::*;
use std::io::BufReader;
use std::convert::Into;
use std::path::PathBuf;
use std::str;

use lib::DownloadError;

use mysql::conn::pool::{MyPool,MyPooledConn};
use mysql::conn::Stmt;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

/// Converter containing all necessary methods to merge, extract and convert audio with/from video files
pub struct Converter<'a> {
    ffmpeg_dir: PathBuf,
    mp3_quality: &'a i16,
    pool: MyPool,
}

/// Struct containing file information needed for progress calculation
struct FileInfo {
	duration: f32, // seconds
	frames: f32,
}

impl<'a> Converter<'a> {
    pub fn new(ffmpeg_dir: &'a str,mp3_quality: &'a i16, pool: MyPool) -> Converter<'a> {
    	debug!("ffmpeg dir: {}",ffmpeg_dir);
        Converter{ffmpeg_dir: PathBuf::from(ffmpeg_dir),mp3_quality: mp3_quality, pool: pool}
    }

    /// Merge audo & video files to one
    /// As ffmpeg uses \r for progress updates, we'll have to read untill this delimiter
    /// ffmpeg prints only to stderr
    pub fn merge_files(&self, qid: &i64, video_file: &'a str,audio_file: &'a str, output_file: &'a str) -> Result<(), DownloadError>{
        let file_info = try!(self.get_file_info(video_file));
        trace!("Total frames: {}",file_info.frames);

        let mut child = try!(self.run_merge_cmd(audio_file, video_file, output_file));
        trace!("started merge process");
        let mut stdout = BufReader::new(child.stderr.take().unwrap());

        let mut conn = self.pool.get_conn().unwrap();
        let mut statement = self.prepare_progress_updater(&mut conn);
        let re = regex!(r"frame=\s*(\d+)");

		let mut buf = vec![];
		buf.reserve(128);

		let mut cur_frame: f32;

		loop {
			match stdout.read_until(b'\r', &mut buf) {
				Ok(0) => break,
				Ok(_) => {},
				Err(why) => {error!("couldn't read cmd stdout: {}", why); panic!(); }
			}

			{
				let line = str::from_utf8(&buf).unwrap();
				debug!("ffmpeg: {}",line);
				if let Some(cap) = re.captures(&line) {
	                debug!("frame: {}", cap.at(1).unwrap());
	                cur_frame = cap.at(1).unwrap().parse::<f32>().unwrap();
	                try!(self.update_progress(&mut statement, format!("{:.2}",(cur_frame / file_info.frames ) * 100.0), qid)
	                );
				}
			}

			buf.clear();

		}

        try!(child.wait());

        Ok(())
    }

    /// Extract audio from video files
    /// If set audio will be converted to mp3 on the fly
    pub fn extract_audio(&self, qid: &i64, video_file: &'a str, output_file: &'a str, convert_mp3: bool) -> Result<(), DownloadError> {
		let file_info = try!(self.get_file_info(video_file));
		debug!("duration: {}",file_info.duration);

        let mut child;
        if convert_mp3 {
            child = try!(self.run_audio_extract_to_mp3(video_file, output_file));
        }else {
            child = try!(self.run_audio_extract(video_file, output_file));
        }

        let mut stdout = BufReader::new(child.stderr.take().unwrap());

        let mut conn = self.pool.get_conn().unwrap();
        let mut statement = self.prepare_progress_updater(&mut conn);
        let re = regex!(r"time=(\d+):(\d+):(\d+.?\d*)");

		let mut buf = vec![];
		buf.reserve(128);

		loop {
			match stdout.read_until(b'\r', &mut buf) {
				Ok(0) => break,
				Ok(_) => {},
				Err(why) => {error!("couldn't read cmd stdout: {}", why); panic!(); }
			}

			{
				let line = str::from_utf8(&buf).unwrap();
				debug!("ffmpeg: {}",line);
				if let Some(cap) = re.captures(&line) {
					let mut seconds: i32 = cap.at(2).unwrap().parse::<i32>().unwrap() * 60; // minutes
		            seconds += cap.at(1).unwrap().parse::<i32>().unwrap() * 60 * 60 ; // hours

		            let seconds: f32 = seconds as f32 + cap.at(3).unwrap().parse::<f32>().unwrap();
		            try!(self.update_progress(&mut statement, format!("{:.2}",(seconds / file_info.duration ) * 100.0), qid)
	                );
				}

			}

			buf.clear();

		}

        try!(child.wait());

        Ok(())
    }

    /// Retrive file information
    fn get_file_info(&self, video_file: &str) -> Result<FileInfo,DownloadError> {
        let stdout = try!(self.run_file_probe(video_file));

		let regex_duration = regex!(r"duration=(\d+\.?\d*)");
		let regex_fps = regex!(r"r_frame_rate=(\d+)/(\d+)");

		if let Some(cap) = regex_duration.captures(&stdout) {
			let mut file_info = FileInfo { duration: -1.0, frames: -1.0 };
			file_info.duration = cap.at(1).unwrap().parse::<f32>().unwrap();

			if let Some(cap) = regex_fps.captures(&stdout) {
				let x: f32 = cap.at(1).unwrap().parse::<f32>().unwrap();
				let y: f32 = cap.at(2).unwrap().parse::<f32>().unwrap();
				file_info.frames =  ( x / y ) * file_info.duration;
			}

			Ok(file_info)
		}else{
			Err(DownloadError::FFMPEGError(format!("Couldn't get max frames: {}",stdout)))
		}
    }

    /// Runs a file probe and returns its output, used in progress calculation
    fn run_file_probe(&self, video_file: &str) -> Result<String, DownloadError> {
    	let mut command = self.create_ffmpeg_base("ffprobe");
    	command.args(&["-select_streams","0"]);
    	command.args(&["-show_entries","stream=duration,r_frame_rate"]);
    	command.args(&["-of","default=noprint_wrappers=1"]);
    	command.arg(video_file);

    	match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(mut process) => {
            	let mut stdout_buffer = BufReader::new(process.stdout.take().unwrap());
	        	let mut stdout: String = String::new();
	      		try!(stdout_buffer.read_to_string(&mut stdout));
	      		try!(process.wait());
	      		debug!("ffprobe: {}",stdout);
	      		Ok(stdout)
            },
    	}
    }

    /// Merges an audio & an video file together.
    /// ffmpeg uses \r as \n
    fn run_merge_cmd(&self, audio_file: &str, video_file: &str, output_file: &str) -> Result<Child,DownloadError> {
    	let mut command = self.create_ffmpeg_base("ffmpeg");
    	command.args(&["-threads","0"]);
    	command.args(&["-i",video_file]);
    	command.args(&["-i",audio_file]);
    	command.args(&["-map","0"]);
    	command.args(&["-map","1"]);
    	command.args(&["-codec","copy"]);
    	command.arg("-shortest");
    	command.arg(output_file);
    	//-stats -threads 0 -i "{}" -i "{}" -map 0 -map 1 -codec copy -shortest "{}"
    	match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    ///Create a ffmpeg instance with the audio extract cmd
    fn run_audio_extract(&self, video_file: &str, output_file: &str) -> Result<Child, DownloadError> {
        let mut command = self.create_ffmpeg_base("ffmpeg");
        command.args(&["-threads", "0"]);
        command.args(&["-i",video_file]);
        command.args(&["-vn", "-acodec","copy"]);
        command.arg(output_file);

        match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    fn run_audio_extract_to_mp3(&self, video_file: &str, output_file: &str) -> Result<Child, DownloadError> {
        let mut command = self.create_ffmpeg_base("ffmpeg");
        command.args(&["-threads", "0"]);
        command.args(&["-i",video_file]);
        command.args(&["-codec:a", "libmp3lame"]);
        command.args(&["-qscale:a",&self.mp3_quality.to_string()]);
        command.arg(output_file);

        match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    /// Create FFMPEG basic command
    /// executable is the called ffmpeg binary
    fn create_ffmpeg_base(&self, executable: &'static str) -> Command {
    	let mut cmd = Command::new(self.ffmpeg_dir.join(executable));
    	cmd.args(&["-v","error"]);
    	cmd.stdin(Stdio::null());
    	cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
    	cmd
    }

    // MyPooledConn does only live when MyOpts is alive -> lifetime needs to be declared
    //TODO: think about generalizing this, while using the local pool
    fn prepare_progress_updater(&'a self,conn: &'a mut MyPooledConn) -> Stmt<'a> { // no livetime needed: struct livetime used
        conn.prepare("UPDATE querydetails SET progress = ? WHERE qid = ?").unwrap()
    }

    ///updater called from the stdout progress
    fn update_progress(&self,stmt: &mut Stmt, progress: String, qid: &i64) -> Result<(),DownloadError>{
    	trace!("updating progress {}",progress);
        try!(stmt.execute((&progress,qid)).map(|_| Ok(())))
    }

}<commit_msg>fixes #4 duration of webm containers seem to be only in the format section, not in the vp9 stream itself<commit_after>
extern crate regex;

use std::process::{Command, Stdio, Child};
use std::io::prelude::*;
use std::io::BufReader;
use std::convert::Into;
use std::path::PathBuf;
use std::str;

use lib::DownloadError;

use mysql::conn::pool::{MyPool,MyPooledConn};
use mysql::conn::Stmt;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

/// Converter containing all necessary methods to merge, extract and convert audio with/from video files
pub struct Converter<'a> {
    ffmpeg_dir: PathBuf,
    mp3_quality: &'a i16,
    pool: MyPool,
}

/// Struct containing file information needed for progress calculation
struct FileInfo {
	duration: f32, // seconds
	frames: f32,
}

impl<'a> Converter<'a> {
    pub fn new(ffmpeg_dir: &'a str,mp3_quality: &'a i16, pool: MyPool) -> Converter<'a> {
    	debug!("ffmpeg dir: {}",ffmpeg_dir);
        Converter{ffmpeg_dir: PathBuf::from(ffmpeg_dir),mp3_quality: mp3_quality, pool: pool}
    }

    /// Merge audo & video files to one
    /// As ffmpeg uses \r for progress updates, we'll have to read untill this delimiter
    /// ffmpeg prints only to stderr
    pub fn merge_files(&self, qid: &i64, video_file: &'a str,audio_file: &'a str, output_file: &'a str) -> Result<(), DownloadError>{
        let file_info = try!(self.get_file_info(video_file));
        trace!("Total frames: {}",file_info.frames);

        let mut child = try!(self.run_merge_cmd(audio_file, video_file, output_file));
        trace!("started merge process");
        let mut stdout = BufReader::new(child.stderr.take().unwrap());

        let mut conn = self.pool.get_conn().unwrap();
        let mut statement = self.prepare_progress_updater(&mut conn);
        let re = regex!(r"frame=\s*(\d+)");

		let mut buf = vec![];
		buf.reserve(128);

		let mut cur_frame: f32;

		loop {
			match stdout.read_until(b'\r', &mut buf) {
				Ok(0) => break,
				Ok(_) => {},
				Err(why) => {error!("couldn't read cmd stdout: {}", why); panic!(); }
			}

			{
				let line = str::from_utf8(&buf).unwrap();
				debug!("ffmpeg: {}",line);
				if let Some(cap) = re.captures(&line) {
	                debug!("frame: {}", cap.at(1).unwrap());
	                cur_frame = cap.at(1).unwrap().parse::<f32>().unwrap();
	                try!(self.update_progress(&mut statement, format!("{:.2}",(cur_frame / file_info.frames ) * 100.0), qid)
	                );
				}
			}

			buf.clear();

		}

        try!(child.wait());

        Ok(())
    }

    /// Extract audio from video files
    /// If set audio will be converted to mp3 on the fly
    pub fn extract_audio(&self, qid: &i64, video_file: &'a str, output_file: &'a str, convert_mp3: bool) -> Result<(), DownloadError> {
		let file_info = try!(self.get_file_info(video_file));
		debug!("duration: {}",file_info.duration);

        let mut child;
        if convert_mp3 {
            child = try!(self.run_audio_extract_to_mp3(video_file, output_file));
        }else {
            child = try!(self.run_audio_extract(video_file, output_file));
        }

        let mut stdout = BufReader::new(child.stderr.take().unwrap());

        let mut conn = self.pool.get_conn().unwrap();
        let mut statement = self.prepare_progress_updater(&mut conn);
        let re = regex!(r"time=(\d+):(\d+):(\d+.?\d*)");

		let mut buf = vec![];
		buf.reserve(128);

		loop {
			match stdout.read_until(b'\r', &mut buf) {
				Ok(0) => break,
				Ok(_) => {},
				Err(why) => {error!("couldn't read cmd stdout: {}", why); panic!(); }
			}

			{
				let line = str::from_utf8(&buf).unwrap();
				debug!("ffmpeg: {}",line);
				if let Some(cap) = re.captures(&line) {
					let mut seconds: i32 = cap.at(2).unwrap().parse::<i32>().unwrap() * 60; // minutes
		            seconds += cap.at(1).unwrap().parse::<i32>().unwrap() * 60 * 60 ; // hours

		            let seconds: f32 = seconds as f32 + cap.at(3).unwrap().parse::<f32>().unwrap();
		            try!(self.update_progress(&mut statement, format!("{:.2}",(seconds / file_info.duration ) * 100.0), qid)
	                );
				}

			}

			buf.clear();

		}

        try!(child.wait());

        Ok(())
    }

    /// Retrive file information
    fn get_file_info(&self, video_file: &str) -> Result<FileInfo,DownloadError> {
        let stdout = try!(self.run_file_probe(video_file));

		let regex_duration = regex!(r"duration=(\d+\.?\d*)");
		let regex_fps = regex!(r"r_frame_rate=(\d+)/(\d+)");

		if let Some(cap) = regex_duration.captures(&stdout) {
			let mut file_info = FileInfo { duration: -1.0, frames: -1.0 };
			file_info.duration = cap.at(1).unwrap().parse::<f32>().unwrap();

			if let Some(cap) = regex_fps.captures(&stdout) {
				let x: f32 = cap.at(1).unwrap().parse::<f32>().unwrap();
				let y: f32 = cap.at(2).unwrap().parse::<f32>().unwrap();
				file_info.frames =  ( x / y ) * file_info.duration;
			}

			Ok(file_info)
		}else{
			Err(DownloadError::FFMPEGError(format!("Couldn't get max frames: {}",stdout)))
		}
    }

    /// Runs a file probe and returns its output, used in progress calculation
    fn run_file_probe(&self, video_file: &str) -> Result<String, DownloadError> {
    	let mut command = self.create_ffmpeg_base("ffprobe");
    	command.args(&["-select_streams","0"]);
    	command.args(&["-show_entries","format=duration:stream=r_frame_rate"]);
    	command.args(&["-of","default=noprint_wrappers=1"]);
    	command.arg(video_file);

    	match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(mut process) => {
            	let mut stdout_buffer = BufReader::new(process.stdout.take().unwrap());
	        	let mut stdout: String = String::new();
	      		try!(stdout_buffer.read_to_string(&mut stdout));
	      		try!(process.wait());
	      		debug!("ffprobe: {}",stdout);
	      		Ok(stdout)
            },
    	}
    }

    /// Merges an audio & an video file together.
    /// ffmpeg uses \r as \n
    fn run_merge_cmd(&self, audio_file: &str, video_file: &str, output_file: &str) -> Result<Child,DownloadError> {
    	let mut command = self.create_ffmpeg_base("ffmpeg");
    	command.args(&["-threads","0"]);
    	command.args(&["-i",video_file]);
    	command.args(&["-i",audio_file]);
    	command.args(&["-map","0"]);
    	command.args(&["-map","1"]);
    	command.args(&["-codec","copy"]);
    	command.arg("-shortest");
    	command.arg(output_file);
    	//-stats -threads 0 -i "{}" -i "{}" -map 0 -map 1 -codec copy -shortest "{}"
    	match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    ///Create a ffmpeg instance with the audio extract cmd
    fn run_audio_extract(&self, video_file: &str, output_file: &str) -> Result<Child, DownloadError> {
        let mut command = self.create_ffmpeg_base("ffmpeg");
        command.args(&["-threads", "0"]);
        command.args(&["-i",video_file]);
        command.args(&["-vn", "-acodec","copy"]);
        command.arg(output_file);

        match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    fn run_audio_extract_to_mp3(&self, video_file: &str, output_file: &str) -> Result<Child, DownloadError> {
        let mut command = self.create_ffmpeg_base("ffmpeg");
        command.args(&["-threads", "0"]);
        command.args(&["-i",video_file]);
        command.args(&["-codec:a", "libmp3lame"]);
        command.args(&["-qscale:a",&self.mp3_quality.to_string()]);
        command.arg(output_file);

        match command.spawn() {
    		Err(why) => Err(why.into()),
            Ok(process) => Ok(process),
    	}
    }

    /// Create FFMPEG basic command
    /// executable is the called ffmpeg binary
    fn create_ffmpeg_base(&self, executable: &'static str) -> Command {
    	let mut cmd = Command::new(self.ffmpeg_dir.join(executable));
    	cmd.args(&["-v","error"]);
    	cmd.stdin(Stdio::null());
    	cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());
    	cmd
    }

    // MyPooledConn does only live when MyOpts is alive -> lifetime needs to be declared
    //TODO: think about generalizing this, while using the local pool
    fn prepare_progress_updater(&'a self,conn: &'a mut MyPooledConn) -> Stmt<'a> { // no livetime needed: struct livetime used
        conn.prepare("UPDATE querydetails SET progress = ? WHERE qid = ?").unwrap()
    }

    ///updater called from the stdout progress
    fn update_progress(&self,stmt: &mut Stmt, progress: String, qid: &i64) -> Result<(),DownloadError>{
    	trace!("updating progress {}",progress);
        try!(stmt.execute((&progress,qid)).map(|_| Ok(())))
    }

}<|endoftext|>