<commit_before>use std::cmp::Ordering;

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    // TODO: boundary_pm_final

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail,
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Some(self.lists[index - 1].lookahead1),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>boundary_pm_final [not complete]<commit_after>use std::cmp::Ordering;

use typed_arena::Arena;

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

struct Thing<'a> {
    node_arena: &'a Arena<Node<'a>>,
    leaves: Vec<Leaf>,
    lists: Vec<List<'a>>,
}

#[derive(Debug)]
struct Node<'a> {
    weight: usize,
    count: usize,
    tail: Option<&'a Node<'a>>,
}

#[derive(Debug)]
struct Leaf {
    weight: usize,
    count: usize,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug, Clone)]
struct List<'a> {
    lookahead0: &'a Node<'a>,
    lookahead1: &'a Node<'a>,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[usize], max_bits: usize) -> Vec<u32> {
    let num_freqs = frequencies.len();
    let mut bit_lengths = vec![0; num_freqs];

    // Count used symbols and place them in the leaves.
    let mut leaves: Vec<_> = frequencies.iter()
        .enumerate()
        .filter(|&(_, &freq)| freq != 0)
        .map(|(i, &freq)| Leaf { weight: freq, count: i })
        .collect();

    let num_symbols = leaves.len();

    // Short circuit some special cases

    // TODO:
    // if ((1 << maxbits) < numsymbols) {
    //   free(leaves);
    //   return 1;  /* Error, too few maxbits to represent symbols. */
    // }

    if num_symbols == 0 {
        // There are no non-zero frequencies.
        return bit_lengths;
    }
    if num_symbols == 1 {
        bit_lengths[leaves[0].count] = 1;
        return bit_lengths;
    }
    if num_symbols == 2 {
        bit_lengths[leaves[0].count] = 1;
        bit_lengths[leaves[1].count] = 1;
        return bit_lengths;
    }

    // Sort the leaves from least frequent to most frequent.
    leaves.sort();

    let max_bits = if num_symbols - 1 < max_bits {
        num_symbols - 1
    } else {
        max_bits
    };

    let arena_capacity = max_bits * 2 * num_symbols;
    let node_arena: Arena<Node> = Arena::with_capacity(arena_capacity);

    let node0 = node_arena.alloc(Node {
        weight: leaves[0].weight,
        count: 1,
        tail: None,
    });

    let node1 = node_arena.alloc(Node {
        weight: leaves[1].weight,
        count: 2,
        tail: None,
    });

    let lists: Vec<List> = vec![
        List {
            lookahead0: node0,
            lookahead1: node1,
        };
        max_bits
    ];

    let mut thing = Thing {
        node_arena: &node_arena,
        leaves,
        lists,
    };

    // In the last list, 2 * num_symbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = 2 * num_symbols - 4;
    for _ in 0..num_boundary_pm_runs - 1 {
        thing.boundary_pm(max_bits - 1);
    }

    thing.boundary_pm_final(max_bits - 1);

    // let mut a = lists.pop().unwrap().lookahead1.leaf_counts.into_iter().rev().peekable();
    //
    // let mut bitlength_value = 1;
    // while let Some(leaf_count) = a.next() {
    //     let next_count = *a.peek().unwrap_or(&0);
    //     for leaf in &leaves[next_count..leaf_count] {
    //         bit_lengths[leaf.count] = bitlength_value;
    //     }
    //     bitlength_value += 1;
    // }

    bit_lengths
}

impl<'a> Thing<'a> {
    fn boundary_pm(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        let last_count = self.lists[index].lookahead1.count; // Count of last chain of list.

        if index == 0 && last_count >= num_symbols {
            return
        }

        self.lists[index].lookahead0 = self.lists[index].lookahead1;

        if index == 0 {
            // New leaf node in list 0.
            let new_chain = self.node_arena.alloc(Node {
               weight: self.leaves[last_count].weight,
               count: last_count + 1,
               tail: self.lists[index].lookahead0.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            let weight_sum = {
                let previous_list = &self.lists[index - 1];
                previous_list.lookahead0.weight + previous_list.lookahead1.weight
            };

            if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
                // New leaf inserted in list, so count is incremented.
                let new_chain = self.node_arena.alloc(Node {
                   weight: self.leaves[last_count].weight,
                   count: last_count + 1,
                   tail: self.lists[index].lookahead0.tail,
                });
                self.lists[index].lookahead1 = new_chain;
            } else {
                let new_chain = self.node_arena.alloc(Node {
                   weight: weight_sum,
                   count: last_count,
                   tail: Some(self.lists[index - 1].lookahead1),
                });
                self.lists[index].lookahead1 = new_chain;

                // Two lookahead chains of previous list used up, create new ones.
                self.boundary_pm(index - 1);
                self.boundary_pm(index - 1);
            }
        }
    }

    fn boundary_pm_final(&mut self, index: usize) {
        let num_symbols = self.leaves.len();

        // Count of last chain of list.
        let last_count = self.lists[index].lookahead1.count;

        let weight_sum = {
            let previous_list = &self.lists[index - 1];
            previous_list.lookahead0.weight + previous_list.lookahead1.weight
        };

        if last_count < num_symbols && weight_sum > self.leaves[last_count].weight {
            let new_chain = self.node_arena.alloc(Node {
               weight: 0,
               count: last_count + 1,
               tail: self.lists[index].lookahead1.tail,
            });
            self.lists[index].lookahead1 = new_chain;
        } else {
            // ???
        }
    }
}

// fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let mut current_list = &mut lists[current_list_index];
//
//     // The next leaf goes next; counting itself makes the leaf_count increase by one.
//     current_list.lookahead1.weight = leaves[current_list.next_leaf_index].weight;
//     current_list.lookahead1.leaf_counts.clear();
//     current_list.lookahead1.leaf_counts.extend_from_slice(&current_list.lookahead0.leaf_counts);
//     let last_index = current_list.lookahead1.leaf_counts.len() - 1;
//     current_list.lookahead1.leaf_counts[last_index] += 1;
//     current_list.next_leaf_index += 1;
// }

// fn next_tree(weight_sum: usize, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     {
//         let (head, tail) = lists.split_at_mut(current_list_index);
//         let prev_list = head.last_mut().unwrap();
//         let current_list = tail.first_mut().unwrap();
//
//         let previous_list_leaf_counts = &prev_list.lookahead1.leaf_counts;
//
//         // Make a tree from the lookaheads from the previous list; that goes next.
//         // This is not a leaf node, so the leaf count stays the same.
//         current_list.lookahead1.weight = weight_sum;
//         current_list.lookahead1.leaf_counts.clear();
//
//         current_list.lookahead1.leaf_counts.extend_from_slice(previous_list_leaf_counts);
//         current_list.lookahead1.leaf_counts.push(*current_list.lookahead0.leaf_counts.last().unwrap());
//     }
//
//     // The previous list needs two new lookahead nodes.
//     boundary_pm(lists, leaves, current_list_index - 1);
//     boundary_pm(lists, leaves, current_list_index - 1);
// }

// fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
//     let last_index = lists.len() - 1;
//     boundary_pm(lists, leaves, last_index);
// }

// fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
//     let next_leaf_index = lists[current_list_index].next_leaf_index;
//     let num_symbols = leaves.len();
//
//     if current_list_index == 0 && next_leaf_index == num_symbols {
//         // We've added all the leaves to the lowest list, so we're done here
//         return;
//     }
//
//     mem::swap(&mut lists[current_list_index].lookahead0, &mut lists[current_list_index].lookahead1);
//
//     if current_list_index == 0 {
//         lowest_list(lists, leaves);
//     } else {
//         // We're at a list other than the lowest list.
//         let weight_sum = {
//             let previous_list = &lists[current_list_index - 1];
//             previous_list.lookahead0.weight + previous_list.lookahead1.weight
//         };
//
//         if next_leaf_index < num_symbols && weight_sum > leaves[next_leaf_index].weight {
//             next_leaf(lists, leaves, current_list_index);
//         } else {
//             next_tree(weight_sum, lists, leaves, current_list_index);
//         }
//     }
// }

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn max_bits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>


<commit_before>use af;
use af::{Dim4, Array};
use af::MatProp;

use activations;
use initializations;
use layer::{Layer, Input};

pub enum LSTMIndex {
  Input,
  Forget,
  Output,
  CellTilda,
}

pub enum ActivationIndex {
  Input,
  Inner,
  Outer,
}

pub enum DataIndex {
  Input,
  Recurrence,
}

pub struct LSTM {
  weights: Vec<Array>,
  recurrent_weights: Vec<Array>,
  bias: Vec<Array>,
  inner_activation: &'static str,
  outer_activation: &'static str,
  return_sequences: bool,
  delta: Array,
  inputs: Input,
}

impl LSTM {
  pub fn new(input_size: u64, output_size: u64
             , outer_activation: &'static str
             , inner_activation: &'static str
             , w_init: &'static str
             , w_inner_init: &'static str
             , bias_init: &'static str
             , forget_bias_init: &'static str
             , return_sequences: bool) -> LSTM
  {
    LSTM{
      weights: vec![initializations::get_initialization(w_init, Dim4::new(&[output_size, input_size, 1, 1])).unwrap()],
      recurrent_weights: vec![initializations::get_initialization(w_inner_init, Dim4::new(&[output_size, output_size, 1, 1])).unwrap(); 4],
      bias: vec![initializations::get_initialization(b_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(forget_bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()],
      inner_activation: inner_activation,
      outer_activation: outer_activation,
      return_sequences: return_sequences,
    }
  }
}

impl RTRL for LSTM {
  pub fn rtrl(&self, dW_tm1: &mut Array  // previous W derivatives for [I, F, Ct]
              , dU_tm1: &mut Array       // previous U derivatives for [I, F, Ct]
              , db_tm1: &mut Array       // previous b derivatives for [I, F, Ct]
              , z_t: &Array              // current time activation
              , inputs: &Input)          // x_t & h_{t-1}
  {
    let block_size = z_t.dims().unwrap()[0];
    assert!(block_size as f32 % 5.0f32 == 0); // there are 5 data pieces we need
    let chunk_size = block_size / 5;
    // chunk out i, f, ct, c_tm1
    let if_t = activations::get_activation(self.inner_activation,
                                            , &af::rows(z_t, 0, 2*chunk_size).unwrap).unwrap();
    let ct_t = activations::get_activation(self.outer_activation,
                                           , &af::rows(z_t, 3*chunk_size, 4*chunk_size).unwrap).unwrap();
    let c_tm1 = af::rows(z_t, 4*chunk_size, 5*chunk_size).unwrap();

    // diff(z_i), diff(z_f), diff(z_ct)
    let dz = vec![&activations::get_activation_derivative(self.inner_activation, &af::rows(z_t, 0, chunk_size).unwrap()).unwrap() //d(z)
                  , &activations::get_activation_derivative(self.inner_activation, &af::rows(z_t, chunk_size, 2*chunk_size).unwrap()).unwrap()
                  , &activations::get_activation_derivative(self.outer_activation, &af::rows(z_t, 3*chunk_size, 4*chunk_size).unwrap()).unwrap()];
    let ct_ctm1_i = vec![&ct_t, &c_tm1, &af::rows(if_t, 0, chunk_size).unwrap()];

    // [Ct_t; C_{t-1}; i_t] * dz
    let dzprod = af::mul(&af::join_many(0, ct_ctm1_i).unwrap()
                         , af::join_many(0, dzvec).unwrap(), false).unwrap();

    // dC_t/dWi  = (dC_{t-1}/dWi  * f_t) + ct_t  * inner_activation(z_i) * x_t
    // dC_t/dWf  = (dC_{t-1}/dWf  * f_t) + c_tm1 * inner_activation(z_f) * x_t
    // dC_t/dWct = (dC_{t-1}/dWct * f_t) + i_t   * outer_activation(Ct)  * x_t
    let w_lhs = af::mul(dW_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/dW * f_t
    let w_rhs = af::mul(&dzprod, &inputs.data[DataIndex::Input]);
    dW_tm1 = af::add(&w_lhs, &w_rhs).unwrap();

    // dC_t/dUi = (dC_{t-1}/dUi * f_t) + ct_t * inner_activation(z_i) * h_{t-1}
    // dC_t/dUf = (dC_{t-1}/dUf * f_t) + c_tm1 * inner_activation(z_f) * h_{t-1}
    // dC_t/dUct = (dC_{t-1}/dUct * f_t) + outer_activation(Ct) * x_t * h_{t-1}
    let u_lhs = af::mul(dU_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/dU * f_t
    let u_rhs = af::mul(&dzprod, &inputs.data[DataIndex::Recurrence]);
    dU_tm1 = af::add(&u_lhs, &u_rhs).unwrap();

    // dC_t/dbi = (dC_{t-1}/dbi * f_t) + ct_t * inner_activation(z_i)
    // dC_t/dbf = (dC_{t-1}/dbf * f_t) + c_tm1 * inner_activation(z_f)
    // dC_t/dbct = (dC_{t-1}/dbct * f_t) + outer_activation(Ct)
    let b_lhs = af::mul(db_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/db * f_t
    dW_tm1 = af::add(&b_lhs, &dzprod).unwrap();
  }
}

impl Layer for LSTM {

  pub fn forward(&mut self, inputs:& Input) {
    // keep previous layer's outputs
    self.inputs = inputs.clone();

    // apply the activation to the previous layer [Optimization: Memory saving]
    let activated_input = activations::get_activation(inputs.activation[ActivationIndex::Inner]
                                                      , &inputs.data[DataIndex::Input]).unwrap();

    // extract the sub-block of each gate [i_tm1, f_tm1, o_tm1, ct_tm1, c_tm2]
    let block_size = inputs.data[DataIndex::Recurrence].dims().unwrap()[0];
    assert!(block_size as f32 % 5.0f32 == 0); // there are 5 data pieces we need
    let chunk_size = block_size / 5;
    let ifo_tm1 = activations::get_activation(inputs.activation[ActivationIndex::Inner]
                                              , &af::rows(&inputs.data[DataIndex::Recurrence], 0, 3 * chunk_size).unwrap).unwrap();
    let ct_tm1 = activations::get_activation(inputs.activation[ActivationIndex::Outer]
                                             , &af::rows(&inputs.data[DataIndex::Recurrence], 3 * chunk_size, 4 * chunk_size).unwrap).unwrap();
    let c_tm2 = af::rows(&inputs.data[DataIndex::Recurrence], 4 * chunk_size, 5 * chunk_size).unwrap();

    // calculate c_tm1 & h_tm1
    let c_tm1 = af::add(&af::mul(&af::rows(&ifo_tm1, 0, chunk_size).unwrap(), &ct_tm1, false).unwrap()
                        , &af::mul(&af::rows(&ifo_tm1, chunk_size, 2 * chunk_size).unwrap(), &c_tm2, false).unwrap()
                        , false).unwrap();
    let h_tm1 = af::mul(&o_tm1, activations::get_activation(inputs.activation[ActivationIndex::Outer]
                                                            , &c_tm1), false).unwrap();

    // forward pass in a batch for performance
    let weights_ref    = vec![&self.weights[LSTMIndex::Input]
                              , &self.weights[LSTMIndex::Forget]
                              , &self.weights[LSTMIndex::Output]
                              , &self.weights[LSTMIndex::CellTilda]];
    let recurrents_ref = vec![&self.recurrent_weights[LSTMIndex::Input]
                              , &self.recurrent_weights[LSTMIndex::Forget]
                              , &self.recurrent_weights[LSTMIndex::Output]
                              , &self.recurrent_weights[LSTMIndex::CellTilda]];
    let bias_ref       = vec![&self.bias[LSTMIndex::Input]
                              , &self.bias[LSTMIndex::Forget]
                              , &self.bias[LSTMIndex::Output]
                              , &self.bias[LSTMIndex::CellTilda]];
    // [ifo_ct] = W*x + U*h_tm1 + b
    let z_t = af::add(&af::add(&af::matmul(&af::join_many(0, weights_ref).unwrap(), &activated_input).unwrap()
                               , &af::matmul(&af::join_many(0, recurrents_ref).unwrap(), &h_tm1).unwrap(), false).unwrap()
                      , &af::join_many(0, bias_ref).unwrap(), true).unwrap();
    self.rtrl(&self.dW, &self.dU, &self.db, &z_t, inputs);
    // since we are RTRL'ing I, F, Ct w.r.t. C we just need to shoot out O
    if self.return_sequences {
      Input { data: af::join_many(0, vec![&z_t, &c_tm1]).unwrap()
              , activation: vec![self.inner_activation, self.outer_activation] }
    }else { //TODO: Fix this
      Input { data: af::join_many(0, vec![&ifo_tm1, &ct_tm1, &c_tm1]).unwrap()
              , activation: vec![self.inner_activation, self.outer_activation] }
    }
  }

  fn backward(&mut self, delta: &Array) -> Array {
  self.delta = delta.clone();

  let activation_prev = activations::get_activation(self.inputs.activation[0], &self.inputs.data[DataIndex::Input]).unwrap();
  let d_activation_prev = activations::get_activation_derivative(self.inputs.activation[0], &activation_prev).unwrap();
  let delta_prev = af::mul(&af::matmul(&self.weights[0], delta, af::MatProp::TRANS, af::MatProp::NONE).unwrap()
                           , &d_activation_prev, false).unwrap();
  delta_prev
}

fn get_delta(&self) -> Array {
self.delta.clone()
  }

  fn get_weights(&self) -> Vec<Array> {
    self.recurrent_weights.extend(self.weights.iter().cloned()).clone()
  }

  fn set_weights(&mut self, weights: &Array, index: usize) {
    self.weights[index] = weights.clone();
  }

  fn get_bias(&self) -> Vec<Array> {
    self.bias.clone()
  }

  fn set_bias(&mut self, bias: &Array, index: usize) {
    self.bias[index] = bias.clone();
  }

  fn get_bias_dims(&self) -> Vec<Dim4> {
    let mut dims = Vec::new();
    for b in &self.bias {
      dims.push(b.dims().unwrap().clone())
    }
    dims
  }

  fn get_weight_dims(&self) -> Vec<Dim4> {
    let mut dims = Vec::new();
    for w in &self.weights {
      dims.push(w.dims().unwrap().clone())
    }
    for w in &self.recurrent_weights {
      dims.push(w.dims().unwrap().clone())
    }
    dims
  }

  fn get_input(&self) -> Input {
    self.inputs.clone()
  }

  fn output_size(&self) -> u64 {
    let weight_dims = self.get_weight_dims();
    weight_dims[weight_dims.len() - 1][1]
  }

  fn input_size(&self) -> u64 {
    let weight_dims = self.get_weight_dims();
    weight_dims[0][0]
  }

  fn get_activation_type(&self) -> &'static str {
    &self.activation
  }
}
<commit_msg>style cleanup<commit_after>use af;
use af::{Dim4, Array};
use af::MatProp;

use activations;
use initializations;
use layer::{Layer, Input};

pub enum LSTMIndex {
  Input,
  Forget,
  Output,
  CellTilda,
}

pub enum ActivationIndex {
  Input,
  Inner,
  Outer,
}

pub enum DataIndex {
  Input,
  Recurrence,
}

pub struct LSTM {
  weights: Vec<Array>,
  recurrent_weights: Vec<Array>,
  bias: Vec<Array>,
  inner_activation: &'static str,
  outer_activation: &'static str,
  return_sequences: bool,
  delta: Array,
  inputs: Input,
}

impl LSTM {
  pub fn new(input_size: u64, output_size: u64
             , outer_activation: &'static str
             , inner_activation: &'static str
             , w_init: &'static str
             , w_inner_init: &'static str
             , bias_init: &'static str
             , forget_bias_init: &'static str
             , return_sequences: bool) -> LSTM
  {
    LSTM{
      weights: vec![initializations::get_initialization(w_init, Dim4::new(&[output_size, input_size, 1, 1])).unwrap()],
      recurrent_weights: vec![initializations::get_initialization(w_inner_init, Dim4::new(&[output_size, output_size, 1, 1])).unwrap(); 4],
      bias: vec![initializations::get_initialization(b_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()
                 , initializations::get_initialization(forget_bias_init, Dim4::new(&[output_size, 1, 1, 1])).unwrap()],
      inner_activation: inner_activation,
      outer_activation: outer_activation,
      return_sequences: return_sequences,
    }
  }
}

impl RTRL for LSTM {
  pub fn rtrl(&self, dW_tm1: &mut Array  // previous W derivatives for [I, F, Ct]
              , dU_tm1: &mut Array       // previous U derivatives for [I, F, Ct]
              , db_tm1: &mut Array       // previous b derivatives for [I, F, Ct]
              , z_t: &Array              // current time activation
              , inputs: &Input)          // x_t & h_{t-1}
  {
    let block_size = z_t.dims().unwrap()[0];
    assert!(block_size as f32 % 5.0f32 == 0); // there are 5 data pieces we need
    let chunk_size = block_size / 5;
    // chunk out i, f, ct, c_tm1
    let if_t = activations::get_activation(self.inner_activation,
                                            , &af::rows(z_t, 0, 2*chunk_size).unwrap).unwrap();
    let ct_t = activations::get_activation(self.outer_activation,
                                           , &af::rows(z_t, 3*chunk_size, 4*chunk_size).unwrap).unwrap();
    let c_tm1 = af::rows(z_t, 4*chunk_size, 5*chunk_size).unwrap();

    // diff(z_i), diff(z_f), diff(z_ct)
    let dz = vec![&activations::get_activation_derivative(self.inner_activation, &af::rows(z_t, 0, chunk_size).unwrap()).unwrap() //d(z)
                  , &activations::get_activation_derivative(self.inner_activation, &af::rows(z_t, chunk_size, 2*chunk_size).unwrap()).unwrap()
                  , &activations::get_activation_derivative(self.outer_activation, &af::rows(z_t, 3*chunk_size, 4*chunk_size).unwrap()).unwrap()];
    let ct_ctm1_i = vec![&ct_t, &c_tm1, &af::rows(if_t, 0, chunk_size).unwrap()];

    // [Ct_t; C_{t-1}; i_t] * dz
    let dzprod = af::mul(&af::join_many(0, ct_ctm1_i).unwrap()
                         , af::join_many(0, dzvec).unwrap(), false).unwrap();

    // dC_t/dWi  = (dC_{t-1}/dWi  * f_t) + ct_t  * inner_activation(z_i) * x_t
    // dC_t/dWf  = (dC_{t-1}/dWf  * f_t) + c_tm1 * inner_activation(z_f) * x_t
    // dC_t/dWct = (dC_{t-1}/dWct * f_t) + i_t   * outer_activation(Ct)  * x_t
    let w_lhs = af::mul(dW_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/dW * f_t
    let w_rhs = af::mul(&dzprod, &inputs.data[DataIndex::Input]);
    dW_tm1 = af::add(&w_lhs, &w_rhs).unwrap();

    // dC_t/dUi = (dC_{t-1}/dUi * f_t) + ct_t * inner_activation(z_i) * h_{t-1}
    // dC_t/dUf = (dC_{t-1}/dUf * f_t) + c_tm1 * inner_activation(z_f) * h_{t-1}
    // dC_t/dUct = (dC_{t-1}/dUct * f_t) + outer_activation(Ct) * x_t * h_{t-1}
    let u_lhs = af::mul(dU_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/dU * f_t
    let u_rhs = af::mul(&dzprod, &inputs.data[DataIndex::Recurrence]);
    dU_tm1 = af::add(&u_lhs, &u_rhs).unwrap();

    // dC_t/dbi = (dC_{t-1}/dbi * f_t) + ct_t * inner_activation(z_i)
    // dC_t/dbf = (dC_{t-1}/dbf * f_t) + c_tm1 * inner_activation(z_f)
    // dC_t/dbct = (dC_{t-1}/dbct * f_t) + outer_activation(Ct)
    let b_lhs = af::mul(db_tm1, &af::rows(if_t, chunk_size, 2*chunk_size).unwrap(), true).unwrap(); // dC_{t-1}/db * f_t
    dW_tm1 = af::add(&b_lhs, &dzprod).unwrap();
  }
}

impl Layer for LSTM {

  pub fn forward(&mut self, inputs:& Input) {
    // keep previous layer's outputs
    self.inputs = inputs.clone();

    // apply the activation to the previous layer [Optimization: Memory saving]
    let activated_input = activations::get_activation(inputs.activation[ActivationIndex::Inner]
                                                      , &inputs.data[DataIndex::Input]).unwrap();

    // extract the sub-block of each gate [i_tm1, f_tm1, o_tm1, ct_tm1, c_tm2]
    let block_size = inputs.data[DataIndex::Recurrence].dims().unwrap()[0];
    assert!(block_size as f32 % 5.0f32 == 0); // there are 5 data pieces we need
    let chunk_size = block_size / 5;
    let ifo_tm1 = activations::get_activation(inputs.activation[ActivationIndex::Inner]
                                              , &af::rows(&inputs.data[DataIndex::Recurrence], 0, 3 * chunk_size).unwrap).unwrap();
    let ct_tm1 = activations::get_activation(inputs.activation[ActivationIndex::Outer]
                                             , &af::rows(&inputs.data[DataIndex::Recurrence], 3 * chunk_size, 4 * chunk_size).unwrap).unwrap();
    let c_tm2 = af::rows(&inputs.data[DataIndex::Recurrence], 4 * chunk_size, 5 * chunk_size).unwrap();

    // calculate c_tm1 & h_tm1
    let c_tm1 = af::add(&af::mul(&af::rows(&ifo_tm1, 0, chunk_size).unwrap(), &ct_tm1, false).unwrap()
                        , &af::mul(&af::rows(&ifo_tm1, chunk_size, 2 * chunk_size).unwrap(), &c_tm2, false).unwrap()
                        , false).unwrap();
    let h_tm1 = af::mul(&o_tm1, activations::get_activation(inputs.activation[ActivationIndex::Outer]
                                                            , &c_tm1), false).unwrap();

    // forward pass in a batch for performance
    let weights_ref    = vec![&self.weights[LSTMIndex::Input]
                              , &self.weights[LSTMIndex::Forget]
                              , &self.weights[LSTMIndex::Output]
                              , &self.weights[LSTMIndex::CellTilda]];
    let recurrents_ref = vec![&self.recurrent_weights[LSTMIndex::Input]
                              , &self.recurrent_weights[LSTMIndex::Forget]
                              , &self.recurrent_weights[LSTMIndex::Output]
                              , &self.recurrent_weights[LSTMIndex::CellTilda]];
    let bias_ref       = vec![&self.bias[LSTMIndex::Input]
                              , &self.bias[LSTMIndex::Forget]
                              , &self.bias[LSTMIndex::Output]
                              , &self.bias[LSTMIndex::CellTilda]];
    // [ifo_ct] = W*x + U*h_tm1 + b
    let z_t = af::add(&af::add(&af::matmul(&af::join_many(0, weights_ref).unwrap(), &activated_input).unwrap()
                               , &af::matmul(&af::join_many(0, recurrents_ref).unwrap(), &h_tm1).unwrap(), false).unwrap()
                      , &af::join_many(0, bias_ref).unwrap(), true).unwrap();
    self.rtrl(&self.dW, &self.dU, &self.db, &z_t, inputs);
    // since we are RTRL'ing I, F, Ct w.r.t. C we just need to shoot out O
    if self.return_sequences {
      Input { data: af::join_many(0, vec![&z_t, &c_tm1]).unwrap()
              , activation: vec![self.inner_activation, self.outer_activation] }
    }else { //TODO: Fix this
      Input { data: af::join_many(0, vec![&ifo_tm1, &ct_tm1, &c_tm1]).unwrap()
              , activation: vec![self.inner_activation, self.outer_activation] }
    }
  }

  fn backward(&mut self, delta: &Array) -> Array {
    self.delta = delta.clone();

    let activation_prev = activations::get_activation(self.inputs.activation[0], &self.inputs.data[DataIndex::Input]).unwrap();
    let d_activation_prev = activations::get_activation_derivative(self.inputs.activation[0], &activation_prev).unwrap();
    let delta_prev = af::mul(&af::matmul(&self.weights[0], delta, af::MatProp::TRANS, af::MatProp::NONE).unwrap()
                             , &d_activation_prev, false).unwrap();
    delta_prev
  }

  fn get_delta(&self) -> Array {
    self.delta.clone()
  }

  fn get_weights(&self) -> Vec<Array> {
    self.recurrent_weights.extend(self.weights.iter().cloned()).clone()
  }

  fn set_weights(&mut self, weights: &Array, index: usize) {
    self.weights[index] = weights.clone();
  }

  fn get_bias(&self) -> Vec<Array> {
    self.bias.clone()
  }

  fn set_bias(&mut self, bias: &Array, index: usize) {
    self.bias[index] = bias.clone();
  }

  fn get_bias_dims(&self) -> Vec<Dim4> {
    let mut dims = Vec::new();
    for b in &self.bias {
      dims.push(b.dims().unwrap().clone())
    }
    dims
  }

  fn get_weight_dims(&self) -> Vec<Dim4> {
    let mut dims = Vec::new();
    for w in &self.weights {
      dims.push(w.dims().unwrap().clone())
    }
    for w in &self.recurrent_weights {
      dims.push(w.dims().unwrap().clone())
    }
    dims
  }

  fn get_input(&self) -> Input {
    self.inputs.clone()
  }

  fn output_size(&self) -> u64 {
    let weight_dims = self.get_weight_dims();
    weight_dims[weight_dims.len() - 1][1]
  }

  fn input_size(&self) -> u64 {
    let weight_dims = self.get_weight_dims();
    weight_dims[0][0]
  }

  fn get_activation_type(&self) -> &'static str {
    &self.activation
  }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

pub mod AnyPointer {
    use std;
    use capability::{ClientHook, FromClientHook, PipelineHook, PipelineOp};
    use layout::{PointerReader, PointerBuilder, FromStructReader, FromStructBuilder,
                 HasStructSize, StructReader};
    use blob::{Text, Data};

    pub struct Reader<'a> {
        priv reader : PointerReader<'a>
    }

    impl <'a> Reader<'a> {
        #[inline]
        pub fn new<'b>(reader : PointerReader<'b>) -> Reader<'b> {
            Reader { reader : reader }
        }

        #[inline]
        pub fn is_null(&self) -> bool {
            self.reader.is_null()
        }

        #[inline]
        pub fn get_as_struct<T : FromStructReader<'a>>(&self) -> T {
            FromStructReader::new(self.reader.get_struct(std::ptr::null()))
        }

        pub fn get_as_text(&self) -> Text::Reader<'a> {
            self.reader.get_text(std::ptr::null(), 0)
        }

        pub fn get_as_data(&self) -> Data::Reader<'a> {
            self.reader.get_data(std::ptr::null(), 0)
        }

        pub fn get_as_capability<T : FromClientHook>(&self) -> T {
            FromClientHook::new(self.reader.get_capability())
        }
    }

    pub struct Builder<'a> {
        priv builder : PointerBuilder<'a>
    }

    impl <'a> Builder<'a> {
        #[inline]
        pub fn new<'b>(builder : PointerBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn get_as_struct<T : FromStructBuilder<'a> + HasStructSize>(&self) -> T {
            FromStructBuilder::new(
                self.builder.get_struct(HasStructSize::struct_size(None::<T>), std::ptr::null()))
        }

        pub fn init_as_struct<T : FromStructBuilder<'a> + HasStructSize>(&self) -> T {
            FromStructBuilder::new(
                self.builder.init_struct(
                    HasStructSize::struct_size(None::<T>)))
        }

        // XXX value should be a user struct.
        pub fn set_as_struct(&self, value : &StructReader) {
            self.builder.set_struct(value);
        }

        // XXX value should be a user client.
        pub fn set_as_capability(&self, value : ~ClientHook) {
            self.builder.set_capability(value);
        }

        pub fn set_as_text(&self, value : &str) {
            self.builder.set_text(value);
        }

        pub fn set_as_data(&self, value : &[u8]) {
            self.builder.set_data(value);
        }

        #[inline]
        pub fn clear(&self) {
            self.builder.clear()
        }
    }

    pub struct Pipeline {
        hook : ~PipelineHook,
        ops : ~[PipelineOp::Type],
    }

    impl Pipeline {
        pub fn new(hook : ~PipelineHook) -> Pipeline {
            Pipeline { hook : hook, ops : box[] }
        }

        pub fn noop(&self) -> Pipeline {
            Pipeline { hook : self.hook.copy(), ops : self.ops.clone() }
        }

        pub fn get_pointer_field(&self, pointer_index : u16) -> Pipeline {
            let mut new_ops = std::vec::with_capacity(self.ops.len() + 1);
            for &op in self.ops.iter() {
                new_ops.push(op)
            }
            new_ops.push(PipelineOp::GetPointerField(pointer_index));
            Pipeline { hook : self.hook.copy(), ops : new_ops }
        }

        pub fn as_cap(&self) -> ~ClientHook {
            self.hook.get_pipelined_cap(self.ops.clone())
        }
    }
}
<commit_msg>AnyPointer::Reader::get_pipelined_cap()<commit_after>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

pub mod AnyPointer {
    use std;
    use capability::{ClientHook, FromClientHook, PipelineHook, PipelineOp};
    use layout::{PointerReader, PointerBuilder, FromStructReader, FromStructBuilder,
                 HasStructSize, StructReader};
    use blob::{Text, Data};

    pub struct Reader<'a> {
        priv reader : PointerReader<'a>
    }

    impl <'a> Reader<'a> {
        #[inline]
        pub fn new<'b>(reader : PointerReader<'b>) -> Reader<'b> {
            Reader { reader : reader }
        }

        #[inline]
        pub fn is_null(&self) -> bool {
            self.reader.is_null()
        }

        #[inline]
        pub fn get_as_struct<T : FromStructReader<'a>>(&self) -> T {
            FromStructReader::new(self.reader.get_struct(std::ptr::null()))
        }

        pub fn get_as_text(&self) -> Text::Reader<'a> {
            self.reader.get_text(std::ptr::null(), 0)
        }

        pub fn get_as_data(&self) -> Data::Reader<'a> {
            self.reader.get_data(std::ptr::null(), 0)
        }

        pub fn get_as_capability<T : FromClientHook>(&self) -> T {
            FromClientHook::new(self.reader.get_capability())
        }

        //# Used by RPC system to implement pipelining. Applications
        //# generally shouldn't use this directly.
        pub fn get_pipelined_cap(&self, ops : &[PipelineOp::Type]) -> ~ClientHook {
            let mut pointer = self.reader;

            for op in ops.iter() {
                match op {
                    &PipelineOp::Noop =>  { }
                    &PipelineOp::GetPointerField(idx) => {
                        pointer = pointer.get_struct(std::ptr::null()).get_pointer_field(idx as uint)
                    }
                }
            }

            pointer.get_capability()
        }
    }

    pub struct Builder<'a> {
        priv builder : PointerBuilder<'a>
    }

    impl <'a> Builder<'a> {
        #[inline]
        pub fn new<'b>(builder : PointerBuilder<'a>) -> Builder<'a> {
            Builder { builder : builder }
        }

        pub fn get_as_struct<T : FromStructBuilder<'a> + HasStructSize>(&self) -> T {
            FromStructBuilder::new(
                self.builder.get_struct(HasStructSize::struct_size(None::<T>), std::ptr::null()))
        }

        pub fn init_as_struct<T : FromStructBuilder<'a> + HasStructSize>(&self) -> T {
            FromStructBuilder::new(
                self.builder.init_struct(
                    HasStructSize::struct_size(None::<T>)))
        }

        // XXX value should be a user struct.
        pub fn set_as_struct(&self, value : &StructReader) {
            self.builder.set_struct(value);
        }

        // XXX value should be a user client.
        pub fn set_as_capability(&self, value : ~ClientHook) {
            self.builder.set_capability(value);
        }

        pub fn set_as_text(&self, value : &str) {
            self.builder.set_text(value);
        }

        pub fn set_as_data(&self, value : &[u8]) {
            self.builder.set_data(value);
        }

        #[inline]
        pub fn clear(&self) {
            self.builder.clear()
        }
    }

    pub struct Pipeline {
        hook : ~PipelineHook,
        ops : ~[PipelineOp::Type],
    }

    impl Pipeline {
        pub fn new(hook : ~PipelineHook) -> Pipeline {
            Pipeline { hook : hook, ops : box[] }
        }

        pub fn noop(&self) -> Pipeline {
            Pipeline { hook : self.hook.copy(), ops : self.ops.clone() }
        }

        pub fn get_pointer_field(&self, pointer_index : u16) -> Pipeline {
            let mut new_ops = std::vec::with_capacity(self.ops.len() + 1);
            for &op in self.ops.iter() {
                new_ops.push(op)
            }
            new_ops.push(PipelineOp::GetPointerField(pointer_index));
            Pipeline { hook : self.hook.copy(), ops : new_ops }
        }

        pub fn as_cap(&self) -> ~ClientHook {
            self.hook.get_pipelined_cap(self.ops.clone())
        }
    }
}
<|endoftext|>


<commit_before>#[doc = "An atomically reference counted wrapper that can be used
hare immutable data between tasks."]

export arc, get, clone;

#[abi = "cdecl"]
native mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: *mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: *mut libc::intptr_t)
        -> libc::intptr_t;
}

type arc_data<T> = {
    mut count: libc::intptr_t,
    data: T
};

resource arc_destruct<T>(data: *arc_data<T>) {
    unsafe {
        let ptr = ptr::mut_addr_of((*data).count);

        let new_count = rustrt::rust_atomic_decrement(ptr);
        assert new_count >= 0;
        if new_count == 0 {
            let _ptr : ~arc_data<T> = unsafe::reinterpret_cast(data);
            // drop glue takes over.
        }
    }
}

type arc<T> = arc_destruct<T>;

fn arc<T>(-data: T) -> arc<T> {
    let data = ~{mut count: 1, data: data};
    unsafe {
        let ptr = unsafe::reinterpret_cast(data);
        unsafe::forget(data);
        arc_destruct(ptr)
    }
}

fn get<T>(rc: &a.arc<T>) -> &a.T {
    unsafe {
        &(***rc).data
    }
}

fn clone<T>(rc: &arc<T>) -> arc<T> {
    let data = **rc;
    unsafe {
        rustrt::rust_atomic_increment(
            ptr::mut_addr_of((*data).count));
    }
    arc_destruct(**rc)
}
<commit_msg>rewrite arc to use region & expressions (also making it pass borrowck)<commit_after>#[doc = "An atomically reference counted wrapper that can be used
hare immutable data between tasks."]

export arc, get, clone;

#[abi = "cdecl"]
native mod rustrt {
    #[rust_stack]
    fn rust_atomic_increment(p: &mut libc::intptr_t)
        -> libc::intptr_t;

    #[rust_stack]
    fn rust_atomic_decrement(p: &mut libc::intptr_t)
        -> libc::intptr_t;
}

type arc_data<T> = {
    mut count: libc::intptr_t,
    data: T
};

resource arc_destruct<T>(data: *arc_data<T>) {
    unsafe {
        let ptr = &mut (*data).count;

        let new_count = rustrt::rust_atomic_decrement(ptr);
        assert new_count >= 0;
        if new_count == 0 {
            let _ptr : ~arc_data<T> = unsafe::reinterpret_cast(data);
            // drop glue takes over.
        }
    }
}

type arc<T> = arc_destruct<T>;

fn arc<T>(-data: T) -> arc<T> {
    let data = ~{mut count: 1, data: data};
    unsafe {
        let ptr = unsafe::reinterpret_cast(data);
        unsafe::forget(data);
        arc_destruct(ptr)
    }
}

fn get<T>(rc: &a.arc<T>) -> &a.T {
    unsafe {
        &(***rc).data
    }
}

fn clone<T>(rc: &arc<T>) -> arc<T> {
    let data = **rc;
    unsafe {
        rustrt::rust_atomic_increment(&mut (*data).count);
    }
    arc_destruct(**rc)
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Traits for the built-in operators

#[allow(missing_doc)];

#[lang="drop"]
pub trait Drop {
    fn drop(&self);
}

#[lang="add"]
pub trait Add<RHS,Result> {
    fn add(&self, rhs: &RHS) -> Result;
}

#[lang="sub"]
pub trait Sub<RHS,Result> {
    fn sub(&self, rhs: &RHS) -> Result;
}

#[lang="mul"]
pub trait Mul<RHS,Result> {
    fn mul(&self, rhs: &RHS) -> Result;
}

#[lang="div"]
pub trait Div<RHS,Result> {
    fn div(&self, rhs: &RHS) -> Result;
}

#[lang="rem"]
pub trait Rem<RHS,Result> {
    fn rem(&self, rhs: &RHS) -> Result;
}

#[lang="neg"]
pub trait Neg<Result> {
    fn neg(&self) -> Result;
}

#[lang="not"]
pub trait Not<Result> {
    fn not(&self) -> Result;
}

#[lang="bitand"]
pub trait BitAnd<RHS,Result> {
    fn bitand(&self, rhs: &RHS) -> Result;
}

#[lang="bitor"]
pub trait BitOr<RHS,Result> {
    fn bitor(&self, rhs: &RHS) -> Result;
}

#[lang="bitxor"]
pub trait BitXor<RHS,Result> {
    fn bitxor(&self, rhs: &RHS) -> Result;
}

#[lang="shl"]
pub trait Shl<RHS,Result> {
    fn shl(&self, rhs: &RHS) -> Result;
}

#[lang="shr"]
pub trait Shr<RHS,Result> {
    fn shr(&self, rhs: &RHS) -> Result;
}

#[lang="index"]
pub trait Index<Index,Result> {
    fn index(&self, index: &Index) -> Result;
}
<commit_msg>std: add benchmark for allocating-and-dropping a struct with a dtor.<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Traits for the built-in operators

#[allow(missing_doc)];

#[lang="drop"]
pub trait Drop {
    fn drop(&self);
}

#[lang="add"]
pub trait Add<RHS,Result> {
    fn add(&self, rhs: &RHS) -> Result;
}

#[lang="sub"]
pub trait Sub<RHS,Result> {
    fn sub(&self, rhs: &RHS) -> Result;
}

#[lang="mul"]
pub trait Mul<RHS,Result> {
    fn mul(&self, rhs: &RHS) -> Result;
}

#[lang="div"]
pub trait Div<RHS,Result> {
    fn div(&self, rhs: &RHS) -> Result;
}

#[lang="rem"]
pub trait Rem<RHS,Result> {
    fn rem(&self, rhs: &RHS) -> Result;
}

#[lang="neg"]
pub trait Neg<Result> {
    fn neg(&self) -> Result;
}

#[lang="not"]
pub trait Not<Result> {
    fn not(&self) -> Result;
}

#[lang="bitand"]
pub trait BitAnd<RHS,Result> {
    fn bitand(&self, rhs: &RHS) -> Result;
}

#[lang="bitor"]
pub trait BitOr<RHS,Result> {
    fn bitor(&self, rhs: &RHS) -> Result;
}

#[lang="bitxor"]
pub trait BitXor<RHS,Result> {
    fn bitxor(&self, rhs: &RHS) -> Result;
}

#[lang="shl"]
pub trait Shl<RHS,Result> {
    fn shl(&self, rhs: &RHS) -> Result;
}

#[lang="shr"]
pub trait Shr<RHS,Result> {
    fn shr(&self, rhs: &RHS) -> Result;
}

#[lang="index"]
pub trait Index<Index,Result> {
    fn index(&self, index: &Index) -> Result;
}

#[cfg(test)]
mod bench {

    use extra::test::BenchHarness;
    use ops::Drop;

    // Overhead of dtors

    struct HasDtor {
        x: int
    }

    impl Drop for HasDtor {
        fn drop(&self) {
        }
    }

    #[bench]
    fn alloc_obj_with_dtor(bh: &mut BenchHarness) {
        do bh.iter {
            HasDtor { x : 10 };
        }
    }
}<|endoftext|>


<commit_before>//! Items Entity

mod path;
mod kind;
mod item;
mod content;

use std::collections::BTreeMap;
use hyper::method::Method;
use serde_json::{self, Value};
use ::connection::{Connection, Helper};
use ::odata::Parameters;
use ::api::MultiOption;
use ::{Result, Error};

pub use self::path::Path;
pub use self::kind::Kind;
pub use self::item::Item;
pub use self::content::Content;

/// Items Entity implementation.
///
/// Items struct implemets methods of [Items API Entity](http://api.sharefile.com/rest/docs/resource.aspx?name=Items)
pub struct Items {
	conn: Connection,
	meta: bool,
}

impl Items {
	/// Create a new instance of Item Entities API
	pub fn new(conn: Connection) -> Self {
		Items {
			conn: conn,
			meta: false,
		}
	}

	/// Set the flag to inform all future API requests to include item meta information
	/// or omit it.
	pub fn include_meta(&mut self, include: bool) {
		self.meta = include;
	}

	/// Resolve given `path` to the Item ID. On success returns `Some(Path::Id(id))`
	// and None otherwise.
	pub fn resolve_path(&self, path: Path) -> Option<Path> {
		if let &Path::Id(_) = &path {
			// The path is the ID already
			Some(path)
		}
		else {
			// Resolve the path to the ID
			match self.stat(path, None) {
				Ok(MultiOption::One(item)) => Some(item.path()),
				_ => None
			}
		}
	}

	/// Search for Item(s) at the `path` given and return found. `parameters` can be used
	/// to provide additional options to the API request, like `includeDeleted`.
	pub fn stat(&self, path: Path, parameters: Option<Parameters>) -> Result<MultiOption<Item>> {
		self.query_items(path.entity_and_parameters(None, parameters))
	}

	/// List all items at the `path` given. `parameters` can be used
	/// to provide additional options to the API request, like `includeDeleted`.
	pub fn list(&self, path: Path, parameters: Option<Parameters>) -> Result<MultiOption<Item>> {
		match self.stat(path, None) {
			Ok(MultiOption::One(item)) => match item.kind {
				Kind::Folder => {
					self.query_items(item.path().entity_and_parameters(Some("/Children"), parameters))
				},
				Kind::File => {
					Ok(MultiOption::One(item))
				}
			},
			Ok(other) => Ok(other),
			Err(err) => err.result()
		}
	}

	/// Create folder with `parent` item and `name` given. On success returns the path with the ID
	/// of the folder created.
	pub fn create_folder<T>(&self, parent: Path, name: T, description: Option<T>, overwite: bool) -> Result<Path>
	where T: Into<String> {
		if let Some(path) = self.resolve_path(parent) {
			// Prepare folder details
			let mut data = BTreeMap::new();
			data.insert(String::from("Name"), Value::String(name.into()));

			if let Some(desc) = description {
				data.insert(String::from("Description"), Value::String(desc.into()));
			}

			let body = Value::Object(data);

			// Create folder
			let parameters = Parameters::new()
				.custom(vec![("overwrite", super::bool_to_string(overwite)), ("passthrough", String::from("false"))]);

			let url = path.entity_and_parameters(Some("/Folder"), Some(parameters));

			self.query_create(url, Some(body))
				.and_then(|r| Error::from_json(r))
				.and_then(|v| Path::from_json(v))
		}
		else {
			Error::new("Cannot resolve parent ID").result()
		}
	}

	/// Download the item identified by `path`. The method returns reader which can be used
	/// to read data in any convenient manner.
	///
	/// The snippet of how the remote file can be downloaded to local.
	///
	/// ```ignore
	/// // Lets assume we have the opened connection already
	/// let path = Path::Absolute("/my_folder/remote_file.txt");
	///	let mut stream = conn.items().download(path).unwrap();
	/// let mut file = File::create("local_file.txt").unwrap();
	/// let mut buf = [0; 1024];
	/// loop {
	/// 	match stream.read(&mut buf) {
	///			Ok(0) => break, // Zero bytes mean the end of file
	///			Ok(n) => file.write_all(&buf[0..n]).unwrap(),
	///			Err(err) => panic!("{:?}", err)
	///		};
	///	}
	/// ```
	pub fn download(&self, path: Path) -> Result<Content> {
		if let Path::Id(id) = path {
			// We have the ID alredy so just start download
			Content::open_for_read(self.conn.clone(), id)
		}
		else {
			// We have a path which should be resolved to the id first
			match self.stat(path, None) {
				Ok(MultiOption::One(item)) => Content::open_for_read(self.conn.clone(), item.id),
				Ok(MultiOption::Many(_)) => Error::new("There are more than one Item on path").result(),
				Ok(MultiOption::None) => Error::new("The Item is not found").result(),
				Err(err) => err.result()
			}
		}
	}

	// Do API request which creates new Items (POST)
	fn query_create(&self, uri: String, data: Option<Value>) -> Result<Value> {
		let body = match data {
			Some(ref value) => Some(serde_json::to_string(value).unwrap()),
			None => None
		};

		match self.conn.query_string(Method::Post, uri, Some(Helper::json_headers()), body) {
			Ok(ref json) => match serde_json::from_str(json) {
				Ok(data) => Ok(data),
				Err(err) => Error::new("JSON parse failed").because(err).result()
			},
			Err(err) => err.result()
		}
	}

	// Do API request which returns Item Collection (GET)
	fn query_items(&self, uri: String) -> Result<MultiOption<Item>> {
		match self.conn.query_string(Method::Get, uri, None, None) {
			Ok(ref json) => match serde_json::from_str(json) {
				Ok(data) => Item::from_value(data, self.meta),
				Err(err) => Error::new("JSON parse failed").because(err).result()
			},
			Err(err) => err.result()
		}
	}
}
<commit_msg>Add Remove item method<commit_after>//! Items Entity

mod path;
mod kind;
mod item;
mod content;

use std::collections::BTreeMap;
use hyper::method::Method;
use serde_json::{self, Value};
use ::connection::{Connection, Helper};
use ::odata::Parameters;
use ::api::MultiOption;
use ::{Result, Error};

pub use self::path::Path;
pub use self::kind::Kind;
pub use self::item::Item;
pub use self::content::Content;

/// Items Entity implementation.
///
/// Items struct implemets methods of [Items API Entity](http://api.sharefile.com/rest/docs/resource.aspx?name=Items)
pub struct Items {
	conn: Connection,
	meta: bool,
}

impl Items {
	/// Create a new instance of Item Entities API
	pub fn new(conn: Connection) -> Self {
		Items {
			conn: conn,
			meta: false,
		}
	}

	/// Set the flag to inform all future API requests to include item meta information
	/// or omit it.
	pub fn include_meta(&mut self, include: bool) {
		self.meta = include;
	}

	/// Resolve given `path` to the Item ID. On success returns `Some(Path::Id(id))`
	// and None otherwise.
	pub fn resolve_path(&self, path: Path) -> Option<Path> {
		if let &Path::Id(_) = &path {
			// The path is the ID already
			Some(path)
		}
		else {
			// Resolve the path to the ID
			match self.stat(path, None) {
				Ok(MultiOption::One(item)) => Some(item.path()),
				_ => None
			}
		}
	}

	/// Search for Item(s) at the `path` given and return found. `parameters` can be used
	/// to provide additional options to the API request, like `includeDeleted`.
	pub fn stat(&self, path: Path, parameters: Option<Parameters>) -> Result<MultiOption<Item>> {
		self.query_items(path.entity_and_parameters(None, parameters))
	}

	/// List all items at the `path` given. `parameters` can be used
	/// to provide additional options to the API request, like `includeDeleted`.
	pub fn list(&self, path: Path, parameters: Option<Parameters>) -> Result<MultiOption<Item>> {
		match self.stat(path, None) {
			Ok(MultiOption::One(item)) => match item.kind {
				Kind::Folder => {
					self.query_items(item.path().entity_and_parameters(Some("/Children"), parameters))
				},
				Kind::File => {
					Ok(MultiOption::One(item))
				}
			},
			Ok(other) => Ok(other),
			Err(err) => err.result()
		}
	}

	/// Create folder with `parent` item and `name` given. On success returns the `Path` with the ID
	/// of the folder created.
	pub fn mkdir<T>(&self, parent: Path, name: T, description: Option<T>, overwite: bool) -> Result<Path>
	where T: Into<String> {
		if let Some(path) = self.resolve_path(parent) {
			// Prepare folder details
			let mut data = BTreeMap::new();
			data.insert(String::from("Name"), Value::String(name.into()));

			if let Some(desc) = description {
				data.insert(String::from("Description"), Value::String(desc.into()));
			}

			let body = Value::Object(data);

			// Create folder
			let parameters = Parameters::new()
				.custom(vec![
					("overwrite", super::bool_to_string(overwite)),
					("passthrough", String::from("false")) ]);

			let url = path.entity_and_parameters(Some("/Folder"), Some(parameters));

			self.query_create(url, Some(body))
				.and_then(|r| Error::from_json(r))
				.and_then(|v| Path::from_json(v))
		}
		else {
			Error::new("Cannot resolve parent ID").result()
		}
	}

	/// Removes the item with the ID given. `single_version` set to `true` will delete
	/// only the specified version rather than all sibling files with the same filename
	/// and `force_sync` set to `true` will block the operation from taking place
	/// asynchronously.
	pub fn remove(&self, path: Path, single_version: bool, force_sync: bool) -> Result<()> {
		self.resolve_path(path)
			.ok_or(Error::new("The Item is not found"))
			.and_then(|path| {
				let parameters = Parameters::new()
					.custom(vec![
						("singleversion", super::bool_to_string(single_version)),
						("forceSync", super::bool_to_string(force_sync)) ]);

				let url = path.entity_and_parameters(None, Some(parameters));

				self.query_delete(url)
					.map(|_| ())
			})
	}

	/// Download the item identified by `path`. The method returns reader which can be used
	/// to read data in any convenient manner.
	///
	/// The snippet of how the remote file can be downloaded to local.
	///
	/// ```ignore
	/// // Lets assume we have the opened connection already
	/// let path = Path::Absolute("/my_folder/remote_file.txt");
	///	let mut stream = conn.items().download(path).unwrap();
	/// let mut file = File::create("local_file.txt").unwrap();
	/// let mut buf = [0; 1024];
	/// loop {
	/// 	match stream.read(&mut buf) {
	///			Ok(0) => break, // Zero bytes mean the end of file
	///			Ok(n) => file.write_all(&buf[0..n]).unwrap(),
	///			Err(err) => panic!("{:?}", err)
	///		};
	///	}
	/// ```
	pub fn download(&self, path: Path) -> Result<Content> {
		if let Path::Id(id) = path {
			// We have the ID alredy so just start download
			Content::open_for_read(self.conn.clone(), id)
		}
		else {
			// We have a path which should be resolved to the id first
			match self.stat(path, None) {
				Ok(MultiOption::One(item)) => Content::open_for_read(self.conn.clone(), item.id),
				Ok(MultiOption::Many(_)) => Error::new("There are more than one Item on path").result(),
				Ok(MultiOption::None) => Error::new("The Item is not found").result(),
				Err(err) => err.result()
			}
		}
	}

	// Do API request which creates new Items (POST)
	fn query_create(&self, uri: String, data: Option<Value>) -> Result<Value> {
		let body = match data {
			Some(ref value) => Some(serde_json::to_string(value).unwrap()),
			None => None
		};

		match self.conn.query_string(Method::Post, uri, Some(Helper::json_headers()), body) {
			Ok(ref json) => match serde_json::from_str(json) {
				Ok(data) => Ok(data),
				Err(err) => Error::new("JSON parse failed").because(err).result()
			},
			Err(err) => err.result()
		}
	}

	// Do API DELETE request
	fn query_delete(&self, uri: String) -> Result<()> {
		self.conn.query(Method::Delete, uri, None, None)
			.and_then(|response| {
				if response.status.is_success() {
					Ok(())
				}
				else {
					Error::new(format!("API DELETE request failed with status {}", response.status)).result()
				}
			})
	}

	// Do API request which returns Item Collection (GET)
	fn query_items(&self, uri: String) -> Result<MultiOption<Item>> {
		match self.conn.query_string(Method::Get, uri, None, None) {
			Ok(ref json) => match serde_json::from_str(json) {
				Ok(data) => Item::from_value(data, self.meta),
				Err(err) => Error::new("JSON parse failed").because(err).result()
			},
			Err(err) => err.result()
		}
	}
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Macros for kernel-level logging

// #[cfg(debug_assertions)]
#[macro_use]
macro_rules! debug {
    ( $($args:tt)* ) => {
        { use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[debug][{}:{}] {}: {}\n"
                      , file!(), line!()
                      , module_path!()
                      , format_args!($($args)*)); }
    }
}

#[macro_use]
macro_rules! info {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {} {}"
                      , module_path!()
                      , $msg, $status);
        print!("{:<38}{:>40}", concat!($dots, $msg), $status );
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        { use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {}"
                      , module_path!()
                      , format_args!($($args)*));
        print!("{}{}", $dots, format_args!($($args)*)); }
    };
    ( $($args:tt)* ) => {
        { use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {}"
                      , module_path!()
                      , format_args!($($args)*));
        print!( $($args)* ); }
    };
}

#[macro_use]
macro_rules! infoln {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {} {}\n"
                      , module_path!()
                      , $msg, $status);
        println!("{:<38}{:>40}", concat!($dots, $msg), $status );
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        { use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {}\n"
                      , module_path!()
                      , format_args!($($args)*));
        println!("{}{}", $dots, format_args!($($args)*)); }
    };
    ( $($args:tt)* ) => {
        { use core::fmt::Write;
        // TODO: only do this if we're in debug mode?
        // suppress warnings because we don't care if there's no serial port
        let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                      , "[info] {}: {}\n"
                      , module_path!()
                      , format_args!($($args)*));
        println!( $($args)* ); }
    };
}
<commit_msg>[kernel] debug logging now only happens in debug mode! :D<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Macros for kernel-level logging

#[cfg(debug_assertions)]
#[macro_use]
macro_rules! debug {
    ( $($args:tt)* ) => {
        {
            use core::fmt::Write;
            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[debug][{}:{}] {}: {}\n"
                          , file!(), line!()
                          , module_path!()
                          , format_args!($($args)*));
        }
    }
}

#[cfg(not(debug_assertions))]
macro_rules! debug {
    ( $($args:tt)* ) => {
        // do nothing
    }
}

#[cfg(debug_assertions)]
#[macro_use]
macro_rules! info {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {} {}"
                          , module_path!()
                          , $msg, $status);
            print!("{:<38}{:>40}", concat!($dots, $msg), $status );
        }
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {}"
                          , module_path!()
                          , format_args!($($args)*));
            print!("{}{}", $dots, format_args!($($args)*));
        }
    };
    ( $($args:tt)* ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {}"
                          , module_path!()
                          , format_args!($($args)*));
            print!( $($args)* );
        }
    };
}

#[cfg(debug_assertions)]
#[macro_use]
macro_rules! infoln {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {} {}\n"
                          , module_path!()
                          , $msg, $status);
            println!("{:<38}{:>40}", concat!($dots, $msg), $status );
        }
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {}\n"
                          , module_path!()
                          , format_args!($($args)*));
            println!("{}{}", $dots, format_args!($($args)*));
        }
    };
    ( $($args:tt)* ) => {
        {
            use core::fmt::Write;

            // suppress warnings because we don't care if there's no serial port
            let _ = write!( $crate::arch::drivers::serial::COM1.lock()
                          , "[info] {}: {}\n"
                          , module_path!()
                          , format_args!($($args)*));
            println!( $($args)* );
        }
    };
}

#[cfg(not(debug_assertions))]
#[macro_use]
macro_rules! info {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        print!("{:<38}{:>40}", concat!($dots, $msg), $status );
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        print!("{}{}", $dots, format_args!($($args)*));
    };
    ( $($args:tt)* ) => {
        print!( $($args)* );
    };
}

#[cfg(not(debug_assertions))]
#[macro_use]
macro_rules! infoln {
    ( dots: $dots:expr, $msg:expr, status: $status:expr ) => {
        println!("{:<38}{:>40}", concat!($dots, $msg), $status );
    };
    ( dots: $dots:expr, $($args:tt)* ) => {
        println!("{}{}", $dots, format_args!($($args)*));
    };
    ( $($args:tt)* ) => {
        println!( $($args)* );
    };
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::collections::{VecDeque, vec_deque};
use std::str;
use std::cmp;
use std::mem;

use itertools::Itertools;
use rust_htslib::bcf;

pub fn match_variants(matchbcf: &str, max_dist: u32, max_len_diff: u32) -> Result<(), Box<Error>> {
    let inbcf = try!(bcf::Reader::new(&"-"));
    let mut header = bcf::Header::with_template(&inbcf.header);

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 is points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    header.push_record(
        format!("##rust-bio-tools={}", env!("CARGO_PKG_VERSION")).as_bytes()
    );
    header.push_record(
        b"##rust-bio-tools-subcommand=vcf-match"
    );
    let mut outbcf = try!(bcf::Writer::new(&"-", &header, false, false));
    let mut buffer = RecordBuffer::new(try!(bcf::Reader::new(&matchbcf)), max_dist);

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header.rid2name(rid);
            let pos = rec.pos();
            // move buffer to pos
            try!(buffer.fill(chrom, pos));

            let var = try!(Variant::new(&mut rec, &mut i));
            let matching = var.alleles.iter().map(|a| {
                for v in buffer.iter() {
                    if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                        return id as i32;
                    }
                }
                -1
            }).collect_vec();

            try!(rec.push_info_integer(b"MATCHING", &matching));
            try!(outbcf.write(&rec));
        } else {
            try!(outbcf.write(&rec));
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {

    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen[0] as u32)
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype[0].to_owned())
        } else { None };
        let end = if let Ok(Some(end)) = rec.info(b"END").integer() {
            Some(end[0] as u32)
        } else { None };
        let inslen = if let Ok(Some(inslen)) = rec.info(b"INSLEN").integer() {
            Some(inslen[0] as u32)
        } else { None };
        let alleles = rec.alleles();
        let refallele = alleles[0];

        let _alleles: Vec<VariantType> = if let Some(svtype) = svtype {
            vec![
                if svtype == b"INS" {
                    let svlen = match (svlen, inslen) {
                        (Some(svlen), _)     => svlen,
                        (None, Some(inslen)) => inslen,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or INSLEN".to_owned())));
                        }
                    };
                    VariantType::Insertion(svlen)
                } else if svtype == b"DEL" {
                    let svlen = match(svlen, end) {
                        (Some(svlen), _)  => svlen,
                        (None, Some(end)) => end - 1 - pos,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or END".to_owned())));
                        }
                    };
                    VariantType::Deletion(svlen)
                } else {
                    warn!("Unsupported variant {}", try!(str::from_utf8(&svtype)));
                    VariantType::Unsupported
                }
            ]
        } else {
            let mut _alleles = Vec::with_capacity(alleles.len() - 1);
            for a in &alleles[1..] {
                _alleles.push(
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        warn!("Unsupported variant {} -> {}", try!(str::from_utf8(refallele)), try!(str::from_utf8(a)));
                        VariantType::Unsupported
                    }
                );
            }
            _alleles
        };

        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32,
            &VariantType::Unsupported => panic!("Unsupported variant.")
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        for (j, b) in other.alleles.iter().enumerate() {
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                // TODO: for now, ignore complex variants
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32),
    Unsupported
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        MissingTag(tag: String) {
            description("missing tag")
            display("missing tag {}", tag)
        }
    }
}

pub struct RecordBuffer {
    reader: bcf::Reader,
    ringbuffer: VecDeque<Variant>,
    ringbuffer2: VecDeque<Variant>,
    window: u32,
    i: u32
}

impl RecordBuffer {
    pub fn new(reader: bcf::Reader, window: u32) -> Self {
        RecordBuffer {
            reader: reader,
            ringbuffer: VecDeque::new(),
            ringbuffer2: VecDeque::new(),
            window: window,
            i: 0
        }
    }

    pub fn last_rid(&self) -> Option<u32> {
        self.ringbuffer.back().map(|var| var.rid)
    }

    pub fn next_rid(&self) -> Option<u32> {
        self.ringbuffer2.back().map(|var| var.rid)
    }

    fn swap_buffers(&mut self) {
        // swap with buffer for next rid
        mem::swap(&mut self.ringbuffer2, &mut self.ringbuffer);
        // clear second buffer
        self.ringbuffer2.clear();
    }

    fn drain_left(&mut self, rid: u32, window_start: u32) {
        // remove records too far left or from wrong rid
        let to_remove = self.ringbuffer.iter().take_while(|var| var.pos < window_start || var.rid != rid).count();
        self.ringbuffer.drain(..to_remove);
    }

    pub fn fill(&mut self, chrom: &[u8], pos: u32) -> Result<(), Box<Error>> {
        let window_start = cmp::max(pos as i32 - self.window as i32 - 1, 0) as u32;
        let window_end = pos + self.window;
        let rid = try!(self.reader.header.name2rid(chrom));

        match (self.last_rid(), self.next_rid()) {
            (Some(last_rid), _) => {
                if last_rid != rid {
                    self.swap_buffers();
                } else {
                    self.drain_left(rid, window_start);
                }
            },
            (_, Some(_)) => {
                self.swap_buffers();
                self.drain_left(rid, window_start);
            },
            _ => ()
        }

        if !self.ringbuffer2.is_empty() {
            // We have already read beyond the current rid. Hence we can't extend to the right for
            // this rid.
            return Ok(())
        }

        // extend to the right
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = self.reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            let pos = rec.pos();
            if let Some(rec_rid) = rec.rid() {
                if rec_rid == rid {
                    if pos > window_end {
                        // Record is beyond our window. Store it anyways but stop.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                        break;
                    } else if pos >= window_start {
                        // Record is within out window.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    }
                } else if rec_rid > rid {
                    // record comes from next rid. Store it in second buffer but stop filling.
                    self.ringbuffer2.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    break;
                } else {
                    // Record comes from previous rid. Ignore it.
                    continue;
                }
            } else {
                // skip records without proper rid
                continue;
            }
        }

        Ok(())
    }

    pub fn iter(&self) -> vec_deque::Iter<Variant> {
        self.ringbuffer.iter()
    }
}
<commit_msg>Fix handling of unsupported variants.<commit_after>use std::error::Error;
use std::collections::{VecDeque, vec_deque};
use std::str;
use std::cmp;
use std::mem;

use itertools::Itertools;
use rust_htslib::bcf;

pub fn match_variants(matchbcf: &str, max_dist: u32, max_len_diff: u32) -> Result<(), Box<Error>> {
    let inbcf = try!(bcf::Reader::new(&"-"));
    let mut header = bcf::Header::with_template(&inbcf.header);

    header.push_record(
        format!("##INFO=<ID=MATCHING,Number=A,Type=Integer,\
        Description=\"For each alternative allele, -1 if it does not match a variant in another VCF/BCF. \
        If it matches a variant, an id i>=0 is points to the i-th variant in the VCF/BCF (counting each \
        alternative allele separately). For indels, matching is fuzzy: distance of centres <= {}, difference of \
        lengths <= {}\">", max_dist, max_len_diff).as_bytes()
    );
    header.push_record(
        format!("##rust-bio-tools={}", env!("CARGO_PKG_VERSION")).as_bytes()
    );
    header.push_record(
        b"##rust-bio-tools-subcommand=vcf-match"
    );
    let mut outbcf = try!(bcf::Writer::new(&"-", &header, false, false));
    let mut buffer = RecordBuffer::new(try!(bcf::Reader::new(&matchbcf)), max_dist);

    let mut rec = bcf::Record::new();
    let mut i = 0;
    loop {
        if let Err(e) = inbcf.read(&mut rec) {
            if e.is_eof() {
                break;
            }
            return Err(Box::new(e));
        }
        outbcf.translate(&mut rec);

        if let Some(rid) = rec.rid() {
            let chrom = inbcf.header.rid2name(rid);
            let pos = rec.pos();
            // move buffer to pos
            try!(buffer.fill(chrom, pos));

            let var = try!(Variant::new(&mut rec, &mut i));
            let matching = var.alleles.iter().map(|a| {
                for v in buffer.iter() {
                    if let Some(id) = var.matches(v, a, max_dist, max_len_diff) {
                        return id as i32;
                    }
                }
                -1
            }).collect_vec();

            try!(rec.push_info_integer(b"MATCHING", &matching));
            try!(outbcf.write(&rec));
        } else {
            try!(outbcf.write(&rec));
        }

        if (i) % 1000 == 0 {
            info!("{} variants written.", i);
        }
    }
    info!("{} variants written.", i);

    Ok(())
}

#[derive(Debug)]
pub struct Variant {
    id: u32,
    rid: u32,
    pos: u32,
    alleles: Vec<VariantType>
}

impl Variant {

    pub fn new(rec: &mut bcf::Record, id: &mut u32) -> Result<Self, Box<Error>> {
        let pos = rec.pos();

        let svlen = if let Ok(Some(svlen)) = rec.info(b"SVLEN").integer() {
            Some(svlen[0] as u32)
        } else { None };
        let svtype = if let Ok(Some(svtype)) = rec.info(b"SVTYPE").string() {
            Some(svtype[0].to_owned())
        } else { None };
        let end = if let Ok(Some(end)) = rec.info(b"END").integer() {
            Some(end[0] as u32)
        } else { None };
        let inslen = if let Ok(Some(inslen)) = rec.info(b"INSLEN").integer() {
            Some(inslen[0] as u32)
        } else { None };
        let alleles = rec.alleles();
        let refallele = alleles[0];

        let _alleles: Vec<VariantType> = if let Some(svtype) = svtype {
            vec![
                if svtype == b"INS" {
                    let svlen = match (svlen, inslen) {
                        (Some(svlen), _)     => svlen,
                        (None, Some(inslen)) => inslen,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or INSLEN".to_owned())));
                        }
                    };
                    VariantType::Insertion(svlen)
                } else if svtype == b"DEL" {
                    let svlen = match(svlen, end) {
                        (Some(svlen), _)  => svlen,
                        (None, Some(end)) => end - 1 - pos,
                        _ => {
                            return Err(Box::new(MatchError::MissingTag("SVLEN or END".to_owned())));
                        }
                    };
                    VariantType::Deletion(svlen)
                } else {
                    warn!("Unsupported variant {}", try!(str::from_utf8(&svtype)));
                    VariantType::Unsupported
                }
            ]
        } else {
            let mut _alleles = Vec::with_capacity(alleles.len() - 1);
            for a in &alleles[1..] {
                _alleles.push(
                    if a.len() < refallele.len() {
                        VariantType::Deletion((refallele.len() - a.len()) as u32)
                    } else if a.len() > refallele.len() {
                        VariantType::Insertion((a.len() - refallele.len()) as u32)
                    } else if a.len() == 1 {
                        VariantType::SNV(a[0])
                    } else {
                        warn!("Unsupported variant {} -> {}", try!(str::from_utf8(refallele)), try!(str::from_utf8(a)));
                        VariantType::Unsupported
                    }
                );
            }
            _alleles
        };

        let var = Variant {
            id: *id,
            rid: rec.rid().unwrap(),
            pos: pos,
            alleles: _alleles
        };
        *id += alleles.len() as u32 - 1;
        Ok(var)
    }

    pub fn centerpoint(&self, allele: &VariantType) -> u32 {
        match allele {
            &VariantType::SNV(_) => self.pos,
            &VariantType::Insertion(_) => self.pos,
            &VariantType::Deletion(len) => (self.pos as f64 + len as f64 / 2.0) as u32,
            &VariantType::Unsupported => panic!("Unsupported variant.")
        }
    }

    pub fn matches(&self, other: &Variant, allele: &VariantType, max_dist: u32, max_len_diff: u32) -> Option<u32> {
        if allele.is_unsupported() {
            return None;
        }
        for (j, b) in other.alleles.iter().enumerate() {
            if b.is_unsupported() {
                continue;
            }
            let dist = (self.centerpoint(allele) as i32 - other.centerpoint(b) as i32).abs() as u32;
            match (allele, b) {
                (&VariantType::SNV(a), &VariantType::SNV(b)) => {
                    if a == b && dist == 0 {
                        return Some(other.id(j));
                    }
                },
                (&VariantType::Insertion(l1), &VariantType::Insertion(l2)) |
                (&VariantType::Deletion(l1), &VariantType::Deletion(l2)) => {
                    if (l1 as i32 - l2 as i32).abs() as u32 <= max_len_diff && dist <= max_dist {
                        return Some(other.id(j));
                    }
                },
                // TODO: for now, ignore complex variants
                _ => continue
            }
        }
        None
    }

    pub fn id(&self, allele: usize) -> u32 {
        self.id + allele as u32
    }
}

#[derive(Debug)]
pub enum VariantType {
    SNV(u8),
    Insertion(u32),
    Deletion(u32),
    Unsupported
}

impl VariantType {
    pub fn is_unsupported(&self) -> bool {
        match self {
            &VariantType::Unsupported => true,
            _ => false
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum MatchError {
        MissingTag(tag: String) {
            description("missing tag")
            display("missing tag {}", tag)
        }
    }
}

pub struct RecordBuffer {
    reader: bcf::Reader,
    ringbuffer: VecDeque<Variant>,
    ringbuffer2: VecDeque<Variant>,
    window: u32,
    i: u32
}

impl RecordBuffer {
    pub fn new(reader: bcf::Reader, window: u32) -> Self {
        RecordBuffer {
            reader: reader,
            ringbuffer: VecDeque::new(),
            ringbuffer2: VecDeque::new(),
            window: window,
            i: 0
        }
    }

    pub fn last_rid(&self) -> Option<u32> {
        self.ringbuffer.back().map(|var| var.rid)
    }

    pub fn next_rid(&self) -> Option<u32> {
        self.ringbuffer2.back().map(|var| var.rid)
    }

    fn swap_buffers(&mut self) {
        // swap with buffer for next rid
        mem::swap(&mut self.ringbuffer2, &mut self.ringbuffer);
        // clear second buffer
        self.ringbuffer2.clear();
    }

    fn drain_left(&mut self, rid: u32, window_start: u32) {
        // remove records too far left or from wrong rid
        let to_remove = self.ringbuffer.iter().take_while(|var| var.pos < window_start || var.rid != rid).count();
        self.ringbuffer.drain(..to_remove);
    }

    pub fn fill(&mut self, chrom: &[u8], pos: u32) -> Result<(), Box<Error>> {
        let window_start = cmp::max(pos as i32 - self.window as i32 - 1, 0) as u32;
        let window_end = pos + self.window;
        let rid = try!(self.reader.header.name2rid(chrom));

        match (self.last_rid(), self.next_rid()) {
            (Some(last_rid), _) => {
                if last_rid != rid {
                    self.swap_buffers();
                } else {
                    self.drain_left(rid, window_start);
                }
            },
            (_, Some(_)) => {
                self.swap_buffers();
                self.drain_left(rid, window_start);
            },
            _ => ()
        }

        if !self.ringbuffer2.is_empty() {
            // We have already read beyond the current rid. Hence we can't extend to the right for
            // this rid.
            return Ok(())
        }

        // extend to the right
        let mut rec = bcf::Record::new();
        loop {
            if let Err(e) = self.reader.read(&mut rec) {
                if e.is_eof() {
                    break;
                }
                return Err(Box::new(e));
            }
            let pos = rec.pos();
            if let Some(rec_rid) = rec.rid() {
                if rec_rid == rid {
                    if pos > window_end {
                        // Record is beyond our window. Store it anyways but stop.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                        break;
                    } else if pos >= window_start {
                        // Record is within out window.
                        self.ringbuffer.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    }
                } else if rec_rid > rid {
                    // record comes from next rid. Store it in second buffer but stop filling.
                    self.ringbuffer2.push_back(try!(Variant::new(&mut rec, &mut self.i)));
                    break;
                } else {
                    // Record comes from previous rid. Ignore it.
                    continue;
                }
            } else {
                // skip records without proper rid
                continue;
            }
        }

        Ok(())
    }

    pub fn iter(&self) -> vec_deque::Iter<Variant> {
        self.ringbuffer.iter()
    }
}
<|endoftext|>