<commit_before>//!
//!	This module contains the code to sum up n-mer aligned read counts based on a
//! SeqTable instance.
//!
use htslib::bam;
use htslib::bam::Read;
use htslib::bam::Reader;
use htslib::bam::Records;
use std::fs::File;
use std::io::Read as ioRead;
use std::io::Seek;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Result;
use std::io::ErrorKind;
use std::io::Error as ioError;
use std::error::Error;
use std::process::exit;
use std::ops::Range;
use std::iter::Peekable;
use seqtable::{SeqTableParams,SeqTable,SequenceInfo};
use std::cmp::Ordering;
use std::str;
use filter::{RecordCheck, PairedChecker, SingleChecker};

struct KeyIter<'a> {
    kmer: Vec<u8>,
    alph: [char; 4],
    mask: Option<&'a Vec<bool>>,
}

impl<'a> KeyIter<'a> {
    fn new(k: u8, mask: Option<&'a Vec<bool>>) -> KeyIter {
        let mut kmer = vec![1; k as usize];
        kmer[(k - 1) as usize] = 0;
        KeyIter { kmer: kmer, alph: ['A','C','G','T'], mask: mask }
    }

    fn key(&self) -> String {
        let mut key = String::new();

        match self.mask {
            Some(flags) => {
                let kiter = &mut self.kmer.iter();
                for &flag in flags {
                    if flag {
                        let b = kiter.next().unwrap();
                        key.push(self.alph[(b - 1) as usize]);
                    } else {
                        key.push('N');
                    }
                }
            },
            None => {
                for b in &self.kmer {
                    key.push(self.alph[(b - 1) as usize]);
                }
            }
        }
        key
    }
}

impl<'a> Iterator for KeyIter<'a> {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        let mut idx = self.kmer.len() - 1;

        loop {
            self.kmer[idx] += 1;
            if self.kmer[idx] > 4 {
                self.kmer[idx] = 1;

                if idx > 0 {
                  idx -= 1;
                } else {
                  return None;
                }
            } else {
                break;
            }
        }

        Some(self.key())
    }
}

// Data for iterator over BED regions
struct BedRanges {
  sets: Vec<Vec<(u32, u32)>>,
  chrom_idx: usize,
  row_idx: usize,
}

impl BedRanges {

  pub fn parse(filename: &str, chroms: Vec<String>) -> Result<BedRanges> {
    let file = try!(File::open(filename));
    let reader = BufReader::new(file);
    let mut n_lines = 0;
    let mut n_no_chrom = 0;
    let mut sets = Vec::new();

    for _ in 0..chroms.len() {
      sets.push(Vec::new());
    }

    for res in reader.lines() {
      let line = try!(res);
      let parts: Vec<&str> = line.split('\t').collect();

      //
      if parts.len() < 3 {
        return Err(ioError::new(ErrorKind::Other, format!("not enough fields in line '{}'", line)));
      }
      n_lines += 1;

      // locate chromosome in set
      let idx = match chroms.iter().position(|chr| chr == parts[0]) {
        Some(idx) => idx,
        None => {
          n_no_chrom += 1;
          continue;
        }
      };

      // add region to set
      let start: u32 = match parts[1].parse() {
        Ok(value) => value,
        Err(_) => return Err(ioError::new(ErrorKind::Other, format!("invalid field value in line '{}'", line))),
      };
      let end: u32 = match parts[2].parse() {
        Ok(value) => value,
        Err(_) => return Err(ioError::new(ErrorKind::Other, format!("invalid field value in line '{}'", line))),
      };
      sets[idx].push((start, end));
    }

    // sort sets by starting position
    for i in 0..sets.len() {
      sets[i].sort_by(|a, b| a.0.cmp(&b.0)); // TODO: consider converting this to sort_by_key - requires rust 1.7
    }

    // check that no stored range overlaps
    for i in 0..sets.len() {
      for j in 1..sets[i].len() {
        if sets[i][j].0 < sets[i][j - 1].1 {
          return Err(ioError::new(ErrorKind::Other, format!("overlapping ranges in {}: {}-{} & {}-{}", chroms[i], sets[i][j - 1].0, sets[i][j - 1].1, sets[i][j].0, sets[i][j].1)));
        }
      }
    }

    if n_no_chrom > 0 {
      println!("WARN:{}: {} rows ({:.2} %) from unknown chromosomes.", filename, n_no_chrom, n_no_chrom as f32 / n_lines as f32 * 100.0);
    }

    Ok(BedRanges{ sets: sets, chrom_idx: 0, row_idx: 0})
  }

  fn contains(&self, chrom_idx: usize, position: i32) -> bool {
      if position < 0 { return false; }
      let position = position as u32;
      self.sets[chrom_idx].binary_search_by(|probe| {
          if probe.0 <= position && probe.1 > position {
              Ordering::Equal
          } else if probe.0 > position {
              Ordering::Greater
          } else {
              Ordering::Less
          }
      }).is_ok()
  }
}

// Iterator over BED regions returns chromosome index (as per supplied map) and the region coordinates in the form of a iteratable Range
impl Iterator for BedRanges {
  type Item = (usize, Range<u32>);

  fn next(&mut self) -> Option<(usize, Range<u32>)> {
    if self.row_idx >= self.sets[self.chrom_idx].len() {
      self.row_idx = 0;
      self.chrom_idx += 1;
    }
    if self.chrom_idx >= self.sets.len() {
      return None;
    }

    let result = Some((self.chrom_idx, Range { start: self.sets[self.chrom_idx][self.row_idx].0, end: self.sets[self.chrom_idx][self.row_idx].1}));
    self.row_idx += 1;
    result
  }
}

fn process_bam_seq<R: ioRead+Seek, C: RecordCheck>(counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, bamrecs: &mut Peekable<Records<Reader>>, tid: &mut i32, map: &Vec<usize>, checker: &C, regions: Option<&BedRanges>) -> bool {
    // skip unmapped sequences (tid = -1)
    if *tid < 0 {
        match bamrecs.next() {
            Some(Ok(ref rec)) => { *tid = rec.tid(); return true; },
            Some(Err(_)) => return false,
            None => return false,
        }
    }

    let chrom_idx = map[*tid as usize];
    let len_idx = table.len_by_idx(chrom_idx).ok().expect("read sequence length") as i32;
    let mut rdr = table.get_sequence_by_idx(chrom_idx).ok().expect("read sequence");

    loop {
        // check if we changed sequence
        match bamrecs.peek() {
            Some(&Ok(ref rec)) => if rec.tid() != *tid { *tid = rec.tid(); return true; },
            Some(&Err(_)) => return false,
            None => return false,
        }

        // if not count position
        if let Some(Ok(record)) = bamrecs.next() {
            // point in regions
            let good = match regions.as_ref() {
                Some(ref ranges) => ranges.contains(chrom_idx, checker.vir_pos(&record)),
                None => true,
            };

            //
            if good && checker.valid(&record) {
                let pos = checker.vir_pos(&record);
                if pos >= len_idx {
                    println!("WARN:{}: read extends beyond reference sequence length ({} > {}). Read ignored.", str::from_utf8(record.qname()).unwrap(), pos + 1, len_idx);
                } else {
                    let pair = rdr.vir_get(pos).unwrap();

                    if record.is_reverse() {
                        counts[pair.1 as usize].3 += 1;
                    } else {
                        counts[pair.0 as usize].2 += 1;
                    }
                }
            }
        }
    }
}

fn bed_regions<R: ioRead + Seek>(filename: &str, table: &mut SeqTable<R>) -> BedRanges {
    let chroms: Vec<String> = table.sequences().iter().map(|sinfo| sinfo.name.clone()).collect();

    match BedRanges::parse(filename, chroms) {
        Ok(iter) => iter,
        Err(e) => {
            println!("Error: Failed to read BED file {}: {}", filename, e.description());
            exit(1);
        },
    }
}

fn region_counts<R: ioRead + Seek>(table: &mut SeqTable<R>, bedregions: &str) -> Vec<(u64, u64, u64, u64)> {
    let bediter = bed_regions(bedregions, table);

    // allocate counts table
    let mut counts: Vec<(u64, u64, u64, u64)> = Vec::new();
    let nmer_count = table.params.nmer_count();

    for _ in 0..nmer_count {
        counts.push((0,0,0,0));
    }

    let n_seqs = table.len();

    for idx in 0..n_seqs {
        let mut rdr = table.get_sequence_by_idx(idx).ok().expect("read sequence");

        for &(start,end) in &bediter.sets[idx] {
            for position in start..end {
                let pair = rdr.get(position).unwrap();
                counts[pair.0 as usize].0 += 1;
                counts[pair.1 as usize].1 += 1;
            }
        }

    }

    counts
}

fn tabulate_bam<R: ioRead + Seek>(bamfilename: String, seqinfos: &Vec<SequenceInfo>, pair_range: &Option<(i32, i32)>, paired: bool, rlen: usize, minqual: u8, counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, regions: Option<&BedRanges>, exact_length: bool, tail_edge: bool) {
    println!("# tabulate {}", bamfilename);

    let bam = match bam::Reader::from_path(&bamfilename) {
        Ok(value) => value,
        Err(err) => {
            println!("Error: Failed to open BAM '{}': {}", &bamfilename, err.description());
            exit(1);
        },
    };
    let names = bam.header().target_names();
    let mut cur_tid = 0;

    // map BAM tid's to SeqTable idx's
    let map : Vec<usize> = names.iter().map(|&id| {
        let chrom = String::from_utf8_lossy(id);
        for i in 0..seqinfos.len() {
            if seqinfos[i].name.eq(&chrom) {
                return i;
            }
        }
        println!("Error: Unknown sequence name in BAM: {}", chrom);
        exit(1);
    }
    ).collect();

    // reads
    let mut iter = bam.records().peekable();

    if pair_range.is_some() || paired {
        let checker = match *pair_range {
            Some((min, max)) => PairedChecker {
                tail_edge: tail_edge,
                exact_length: exact_length,
                read_length: rlen,
                min_quality: minqual,
                min_dist: min,
                max_dist: max,
                force_paired: paired,
                max_distance: true,
            },
            None => PairedChecker {
                tail_edge: tail_edge,
                exact_length: exact_length,
                read_length: rlen,
                min_quality: minqual,
                min_dist: 0,
                max_dist: 0,
                force_paired: paired,
                max_distance: false,
            },
        };
        while process_bam_seq(counts, table, &mut iter, &mut cur_tid, &map, &checker, regions) {}
    } else {
        let checker = SingleChecker { tail_edge: tail_edge, exact_length: exact_length, read_length: rlen, min_quality: minqual };
        while process_bam_seq(counts, table, &mut iter, &mut cur_tid, &map, &checker, regions) {}
    }
}

pub fn tabulate(seqfile: &str, bamfile: Option<&Vec<String>>, minqual: u8, regions: Option<String>, pair_range: Option<(i32, i32)>, paired: bool, exact_length: bool, tail_edge: bool) -> Vec<(u64, u64, u64, u64)> {
    // read
    let file = match File::open(seqfile) {
        Ok(value) => value,
        Err(err) => {
            println!("Error: Failed to open sequence table file '{}': {}", seqfile, err.description());
            exit(1);
        },
    };
    let mut table = match SeqTable::open(file) {
        Ok(value) => value,
        Err(e) => {
            println!("Error:tabulate: {}", e.to_string());
            exit(1);
        },
    };

    // get counts table from file
    let mut counts = match regions.as_ref() {
        Some(regfile) => region_counts(&mut table, regfile),
        None => table.counts().unwrap(),
    };

    //
    let rlen = table.params.read_length as usize;
    let seqinfos = table.sequences();

    // if we received a BAM file, parse it
    if let Some(bamfilenames) = bamfile {
        // bed regions
        let ranges = if let Some(regfile) = regions {
            Some(bed_regions(&regfile, &mut table))
        } else {
            None
        };

        for bamfilename in bamfilenames {
            tabulate_bam(bamfilename.clone(), &seqinfos, &pair_range, paired, rlen, minqual, &mut counts, &mut table, ranges.as_ref(), exact_length, tail_edge);
        }
    }

    counts
}

pub fn print_counts(counts: &Vec<(u64, u64, u64, u64)>, with_bam: bool, params: &SeqTableParams) {
    let mut keys = KeyIter::new(params.unmasked_count, params.mask.as_ref());

    if with_bam {
        for i in 1..counts.len() {
            let (plus, minus, bam_plus, bam_minus) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}\t{}\t{}", i, key, plus, minus, bam_plus, bam_minus);
        }
    } else {
        for i in 1..counts.len() {
            let (plus, minus, _, _) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}", i, key, plus, minus);
        }
    }
}<commit_msg>reverse_complement for minus counts<commit_after>//!
//!	This module contains the code to sum up n-mer aligned read counts based on a
//! SeqTable instance.
//!
use htslib::bam;
use htslib::bam::Read;
use htslib::bam::Reader;
use htslib::bam::Records;
use std::fs::File;
use std::io::Read as ioRead;
use std::io::Seek;
use std::io::BufRead;
use std::io::BufReader;
use std::io::Result;
use std::io::ErrorKind;
use std::io::Error as ioError;
use std::error::Error;
use std::process::exit;
use std::ops::Range;
use std::iter::Peekable;
use seqtable::{SeqTableParams,SeqTable,SequenceInfo};
use std::cmp::Ordering;
use std::str;
use filter::{RecordCheck, PairedChecker, SingleChecker};

struct KeyIter<'a> {
    kmer: Vec<u8>,
    alph: [char; 4],
    mask: Option<&'a Vec<bool>>,
}

impl<'a> KeyIter<'a> {
    fn new(k: u8, mask: Option<&'a Vec<bool>>) -> KeyIter {
        let mut kmer = vec![1; k as usize];
        kmer[(k - 1) as usize] = 0;
        KeyIter { kmer: kmer, alph: ['A','C','G','T'], mask: mask }
    }

    fn key(&self) -> String {
        let mut key = String::new();

        match self.mask {
            Some(flags) => {
                let kiter = &mut self.kmer.iter();
                for &flag in flags {
                    if flag {
                        let b = kiter.next().unwrap();
                        key.push(self.alph[(b - 1) as usize]);
                    } else {
                        key.push('N');
                    }
                }
            },
            None => {
                for b in &self.kmer {
                    key.push(self.alph[(b - 1) as usize]);
                }
            }
        }
        key
    }
}

impl<'a> Iterator for KeyIter<'a> {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        let mut idx = self.kmer.len() - 1;

        loop {
            self.kmer[idx] += 1;
            if self.kmer[idx] > 4 {
                self.kmer[idx] = 1;

                if idx > 0 {
                  idx -= 1;
                } else {
                  return None;
                }
            } else {
                break;
            }
        }

        Some(self.key())
    }
}

// Data for iterator over BED regions
struct BedRanges {
  sets: Vec<Vec<(u32, u32)>>,
  chrom_idx: usize,
  row_idx: usize,
}

impl BedRanges {

  pub fn parse(filename: &str, chroms: Vec<String>) -> Result<BedRanges> {
    let file = try!(File::open(filename));
    let reader = BufReader::new(file);
    let mut n_lines = 0;
    let mut n_no_chrom = 0;
    let mut sets = Vec::new();

    for _ in 0..chroms.len() {
      sets.push(Vec::new());
    }

    for res in reader.lines() {
      let line = try!(res);
      let parts: Vec<&str> = line.split('\t').collect();

      //
      if parts.len() < 3 {
        return Err(ioError::new(ErrorKind::Other, format!("not enough fields in line '{}'", line)));
      }
      n_lines += 1;

      // locate chromosome in set
      let idx = match chroms.iter().position(|chr| chr == parts[0]) {
        Some(idx) => idx,
        None => {
          n_no_chrom += 1;
          continue;
        }
      };

      // add region to set
      let start: u32 = match parts[1].parse() {
        Ok(value) => value,
        Err(_) => return Err(ioError::new(ErrorKind::Other, format!("invalid field value in line '{}'", line))),
      };
      let end: u32 = match parts[2].parse() {
        Ok(value) => value,
        Err(_) => return Err(ioError::new(ErrorKind::Other, format!("invalid field value in line '{}'", line))),
      };
      sets[idx].push((start, end));
    }

    // sort sets by starting position
    for i in 0..sets.len() {
      sets[i].sort_by(|a, b| a.0.cmp(&b.0)); // TODO: consider converting this to sort_by_key - requires rust 1.7
    }

    // check that no stored range overlaps
    for i in 0..sets.len() {
      for j in 1..sets[i].len() {
        if sets[i][j].0 < sets[i][j - 1].1 {
          return Err(ioError::new(ErrorKind::Other, format!("overlapping ranges in {}: {}-{} & {}-{}", chroms[i], sets[i][j - 1].0, sets[i][j - 1].1, sets[i][j].0, sets[i][j].1)));
        }
      }
    }

    if n_no_chrom > 0 {
      println!("WARN:{}: {} rows ({:.2} %) from unknown chromosomes.", filename, n_no_chrom, n_no_chrom as f32 / n_lines as f32 * 100.0);
    }

    Ok(BedRanges{ sets: sets, chrom_idx: 0, row_idx: 0})
  }

  fn contains(&self, chrom_idx: usize, position: i32) -> bool {
      if position < 0 { return false; }
      let position = position as u32;
      self.sets[chrom_idx].binary_search_by(|probe| {
          if probe.0 <= position && probe.1 > position {
              Ordering::Equal
          } else if probe.0 > position {
              Ordering::Greater
          } else {
              Ordering::Less
          }
      }).is_ok()
  }
}

// Iterator over BED regions returns chromosome index (as per supplied map) and the region coordinates in the form of a iteratable Range
impl Iterator for BedRanges {
  type Item = (usize, Range<u32>);

  fn next(&mut self) -> Option<(usize, Range<u32>)> {
    if self.row_idx >= self.sets[self.chrom_idx].len() {
      self.row_idx = 0;
      self.chrom_idx += 1;
    }
    if self.chrom_idx >= self.sets.len() {
      return None;
    }

    let result = Some((self.chrom_idx, Range { start: self.sets[self.chrom_idx][self.row_idx].0, end: self.sets[self.chrom_idx][self.row_idx].1}));
    self.row_idx += 1;
    result
  }
}

// based on https://www.biostars.org/p/113640/#424278
pub fn reverse_complement(mer: u64, kmersize: u8) -> u64 {
    let mut res = (!mer).wrapping_add(1);
//    println!("mer '{}'", mer);
//    println!("kmersize '{}'", kmersize);
    res = (res >> 2 & 0x3333333333333333) | (res & 0x3333333333333333) << 2;
    res = res >> 4 & 0xf0f0f0f0f0f0f0f | (res & 0xf0f0f0f0f0f0f0f) << 4;
    res = res >> 8 & 0xff00ff00ff00ff | (res & 0xff00ff00ff00ff) << 8;
    res = res >> 16 & 0xffff0000ffff | (res & 0xffff0000ffff) << 16;
    res = res >> 32 & 0xffffffff | (res & 0xffffffff) << 32;
    res = (res >> 2 * (32 - kmersize)).wrapping_add(1);
//    println!("RESULT '{}'", res);
    return res;
}

fn process_bam_seq<R: ioRead+Seek, C: RecordCheck>(counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, bamrecs: &mut Peekable<Records<Reader>>, tid: &mut i32, map: &Vec<usize>, checker: &C, regions: Option<&BedRanges>) -> bool {
    // skip unmapped sequences (tid = -1)
    if *tid < 0 {
        match bamrecs.next() {
            Some(Ok(ref rec)) => { *tid = rec.tid(); return true; },
            Some(Err(_)) => return false,
            None => return false,
        }
    }

    let chrom_idx = map[*tid as usize];
    let len_idx = table.len_by_idx(chrom_idx).ok().expect("read sequence length") as i32;

    let mut rdr = table.get_sequence_by_idx(chrom_idx).ok().expect("read sequence");

    loop {
        // check if we changed sequence
        match bamrecs.peek() {
            Some(&Ok(ref rec)) => if rec.tid() != *tid { *tid = rec.tid(); return true; },
            Some(&Err(_)) => return false,
            None => return false,
        }

        // if not count position
        if let Some(Ok(record)) = bamrecs.next() {
            // point in regions
            let good = match regions.as_ref() {
                Some(ref ranges) => ranges.contains(chrom_idx, checker.vir_pos(&record)),
                None => true,
            };

            //
            if good && checker.valid(&record) {
                let pos = checker.vir_pos(&record);
                if pos >= len_idx {
                    println!("WARN:{}: read extends beyond reference sequence length ({} > {}). Read ignored.", str::from_utf8(record.qname()).unwrap(), pos + 1, len_idx);
                } else {
                    let pair = rdr.vir_get(pos).unwrap();

                    if record.is_reverse() {
		    // reverse complement here
                        counts[reverse_complement(pair.1 as u64, ksze) as usize].3 += 1;
                    } else {
                        counts[pair.0 as usize].2 += 1;
                    }
                }
            }
        }
    }
}

fn bed_regions<R: ioRead + Seek>(filename: &str, table: &mut SeqTable<R>) -> BedRanges {
    let chroms: Vec<String> = table.sequences().iter().map(|sinfo| sinfo.name.clone()).collect();

    match BedRanges::parse(filename, chroms) {
        Ok(iter) => iter,
        Err(e) => {
            println!("Error: Failed to read BED file {}: {}", filename, e.description());
            exit(1);
        },
    }
}

fn region_counts<R: ioRead + Seek>(table: &mut SeqTable<R>, bedregions: &str) -> Vec<(u64, u64, u64, u64)> {
    let bediter = bed_regions(bedregions, table);

    // allocate counts table
    let mut counts: Vec<(u64, u64, u64, u64)> = Vec::new();
    let nmer_count = table.params.nmer_count();

    for _ in 0..nmer_count {
        counts.push((0,0,0,0));
    }

    let n_seqs = table.len();

    for idx in 0..n_seqs {
        let mut rdr = table.get_sequence_by_idx(idx).ok().expect("read sequence");

        for &(start,end) in &bediter.sets[idx] {
            for position in start..end {
                let pair = rdr.get(position).unwrap();
                counts[pair.0 as usize].0 += 1;
                counts[pair.1 as usize].1 += 1;
            }
        }

    }

    counts
}

fn tabulate_bam<R: ioRead + Seek>(bamfilename: String, seqinfos: &Vec<SequenceInfo>, pair_range: &Option<(i32, i32)>, paired: bool, rlen: usize, minqual: u8, counts: &mut Vec<(u64, u64, u64, u64)>, table: &mut SeqTable<R>, regions: Option<&BedRanges>, exact_length: bool, tail_edge: bool) {
    println!("# tabulate {}", bamfilename);

    let bam = match bam::Reader::from_path(&bamfilename) {
        Ok(value) => value,
        Err(err) => {
            println!("Error: Failed to open BAM '{}': {}", &bamfilename, err.description());
            exit(1);
        },
    };
    let names = bam.header().target_names();
    let mut cur_tid = 0;

    // map BAM tid's to SeqTable idx's
    let map : Vec<usize> = names.iter().map(|&id| {
        let chrom = String::from_utf8_lossy(id);
        for i in 0..seqinfos.len() {
            if seqinfos[i].name.eq(&chrom) {
                return i;
            }
        }
        println!("Error: Unknown sequence name in BAM: {}", chrom);
        exit(1);
    }
    ).collect();

    // reads
    let mut iter = bam.records().peekable();

    if pair_range.is_some() || paired {
        let checker = match *pair_range {
            Some((min, max)) => PairedChecker {
                tail_edge: tail_edge,
                exact_length: exact_length,
                read_length: rlen,
                min_quality: minqual,
                min_dist: min,
                max_dist: max,
                force_paired: paired,
                max_distance: true,
            },
            None => PairedChecker {
                tail_edge: tail_edge,
                exact_length: exact_length,
                read_length: rlen,
                min_quality: minqual,
                min_dist: 0,
                max_dist: 0,
                force_paired: paired,
                max_distance: false,
            },
        };
        while process_bam_seq(counts, table, &mut iter, &mut cur_tid, &map, &checker, regions) {}
    } else {
        let checker = SingleChecker { tail_edge: tail_edge, exact_length: exact_length, read_length: rlen, min_quality: minqual };
        while process_bam_seq(counts, table, &mut iter, &mut cur_tid, &map, &checker, regions) {}
    }
}

pub fn tabulate(seqfile: &str, bamfile: Option<&Vec<String>>, minqual: u8, regions: Option<String>, pair_range: Option<(i32, i32)>, paired: bool, exact_length: bool, tail_edge: bool) -> Vec<(u64, u64, u64, u64)> {
    // read
    let file = match File::open(seqfile) {
        Ok(value) => value,
        Err(err) => {
            println!("Error: Failed to open sequence table file '{}': {}", seqfile, err.description());
            exit(1);
        },
    };
    let mut table = match SeqTable::open(file) {
        Ok(value) => value,
        Err(e) => {
            println!("Error:tabulate: {}", e.to_string());
            exit(1);
        },
    };

    // get counts table from file
    let mut counts = match regions.as_ref() {
        Some(regfile) => region_counts(&mut table, regfile),
        None => table.counts().unwrap(),
    };

    //
    let rlen = table.params.read_length as usize;
    let seqinfos = table.sequences();

    // if we received a BAM file, parse it
    if let Some(bamfilenames) = bamfile {
        // bed regions
        let ranges = if let Some(regfile) = regions {
            Some(bed_regions(&regfile, &mut table))
        } else {
            None
        };

        for bamfilename in bamfilenames {
            tabulate_bam(bamfilename.clone(), &seqinfos, &pair_range, paired, rlen, minqual, &mut counts, &mut table, ranges.as_ref(), exact_length, tail_edge);
        }
    }

    counts
}

pub fn print_counts(counts: &Vec<(u64, u64, u64, u64)>, with_bam: bool, params: &SeqTableParams) {
    let mut keys = KeyIter::new(params.unmasked_count, params.mask.as_ref());

    if with_bam {
        for i in 1..counts.len() {
            let (plus, minus, bam_plus, bam_minus) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}\t{}\t{}", i, key, plus, minus, bam_plus, bam_minus);
        }
    } else {
        for i in 1..counts.len() {
            let (plus, minus, _, _) = counts[i];
            let key = keys.next().unwrap();
            println!("{}\t{}\t{}\t{}", i, key, plus, minus);
        }
    }
}
<|endoftext|>


<commit_before>use bson::{self, Bson};
use {Error, Client, ThreadedClient, Result};
use wire_protocol::flags::OpQueryFlags;
use wire_protocol::operations::Message;
use std::collections::vec_deque::VecDeque;
use std::io::{Read, Write};

pub const DEFAULT_BATCH_SIZE: i32 = 20;

/// Maintains a connection to the server and lazily returns documents from a
/// query.
///
/// # Fields
///
/// `client` - The client to read from.
/// `namespace` - The namespace to read and write from.
/// `batch_size` - How many documents to fetch at a given time from the server.
/// `cursor_id` - Uniquely identifies the cursor being returned by the reply.
/// `limit` - An upper bound on the total number of documents this cursor
///           should return.
/// `count` - How many documents have been returned so far.
/// `buffer` - A cache for documents received from the query that have not
///            yet been returned.
pub struct Cursor {
    client: Client,
    namespace: String,
    batch_size: i32,
    cursor_id: i64,
    limit: i32,
    count: i32,
    buffer: VecDeque<bson::Document>,
}

impl Cursor {
    /// Construcs a new Cursor for a database command.
    ///
    /// # Arguments
    ///
    /// `client` - Client making the request.
    /// `db` - Which database the command is being sent to.
    /// `doc` - Specifies the command that is being run.
    ///
    /// # Return value
    ///
    /// Returns the newly created Cursor on success, or an Error on failure.
    pub fn command_cursor(client: Client, db: &str,
                          doc: bson::Document) -> Result<Cursor> {
        Cursor::query_with_batch_size(client.clone(), format!("{}.$cmd", db),
                                      1, OpQueryFlags::no_flags(), 0, 0,
                                      doc, None, true)
    }

    fn get_bson_and_cid_from_message(message: Message) -> Result<(VecDeque<bson::Document>, i64)> {
        match message {
            Message::OpReply { header: _, flags: _, cursor_id: cid,
                               starting_from: _, number_returned: _,
                               documents: docs } => {
                let mut v = VecDeque::new();

                for doc in docs {
                    v.push_back(doc.clone());
                }

                Ok((v, cid))
            },
            _ => Err(Error::CursorNotFoundError)
        }
    }

    fn get_bson_and_cursor_info_from_command_message(message: Message) -> Result<(VecDeque<bson::Document>, i64, String)> {
        let (v, _) = try!(Cursor::get_bson_and_cid_from_message(message));
        if v.len() != 1 {
            return Err(Error::CursorNotFoundError);
        }

        let ref doc = v[0];

        // Extract cursor information
        if let Some(&Bson::Document(ref cursor)) = doc.get("cursor") {
            if let Some(&Bson::I64(ref id)) = cursor.get("id") {
                if let Some(&Bson::String(ref ns)) = cursor.get("ns") {
                    if let Some(&Bson::Array(ref batch)) = cursor.get("firstBatch") {

                        // Extract first batch documents
                        let map = batch.iter().filter_map(|bdoc| {
                            if let &Bson::Document(ref doc) = bdoc {
                                Some(doc.clone())
                            } else {
                                None
                            }
                        }).collect();

                        return Ok((map, *id, ns.to_owned()))
                    }
                }
            }
        }

        Err(Error::CursorNotFoundError)
    }

    /// Executes a query where the batch size of the returned cursor is
    /// specified.
    ///
    /// # Arguments
    ///
    /// `client` - The client to read from.
    /// `namespace` - The namespace to read and write from.
    /// `batch_size` - How many documents the cursor should return at a time.
    /// `flags` - Bit vector of query options.
    /// `number_to_skip` - The number of initial documents to skip over in the
    ///                    query results.
    /// `number_to_return - The total number of documents that should be
    ///                     returned by the query.
    /// `return_field_selector - An optional projection of which fields should
    ///                          be present in the documents to be returned by
    ///                          the query.
    /// `is_cmd_cursor` - Whether or not the Cursor is for a database command.
    ///
    /// # Return value
    ///
    /// Returns the cursor for the query results on success, or an Error on
    /// failure.
    pub fn query_with_batch_size(client: Client, namespace: String,
                                 batch_size: i32, flags: OpQueryFlags,
                                 number_to_skip: i32, number_to_return: i32,
                                 query: bson::Document,
                                 return_field_selector: Option<bson::Document>,
                                 is_cmd_cursor: bool) -> Result<Cursor> {
        let result = Message::new_query(client.get_req_id(), flags,
                                        namespace.to_owned(),
                                        number_to_skip, batch_size,
                                        query, return_field_selector);

        let stream = try!(client.acquire_stream());
        let mut socket = stream.get_socket();

        let message = try!(result);
        try!(message.write(&mut socket));
        let reply = try!(Message::read(&mut socket));

        let (buf, cursor_id, namespace) = if is_cmd_cursor {
            try!(Cursor::get_bson_and_cursor_info_from_command_message(reply))
        } else {
            let (buf, id) = try!(Cursor::get_bson_and_cid_from_message(reply));
            (buf, id, namespace)
        };

        Ok(Cursor { client: client, namespace: namespace,
                    batch_size: batch_size, cursor_id: cursor_id,
                    limit: number_to_return, count: 0, buffer: buf, })
    }

    /// Executes a query with the default batch size.
    ///
    /// # Arguments
    ///
    /// `client` - The client to read from.
    /// `namespace` - The namespace to read and write from.
    /// `flags` - Bit vector of query options.
    /// `number_to_skip` - The number of initial documents to skip over in the
    ///                    query results.
    /// `number_to_return - The total number of documents that should be
    ///                     returned by the query.
    /// `query` - Specifies which documents to return.
    /// `return_field_selector - An optional projection of which fields should
    ///                          be present in the documents to be returned by
    ///                          the query.
    /// `is_cmd_cursor` - Whether or not the Cursor is for a database command.
    ///
    /// # Return value
    ///
    /// Returns the cursor for the query results on success, or an error string
    /// on failure.
    pub fn query(client: Client, namespace: String,
                 flags: OpQueryFlags, number_to_skip: i32,
                 number_to_return: i32, query: bson::Document,
                 return_field_selector: Option<bson::Document>,
                 is_cmd_cursor: bool) -> Result<Cursor> {
        Cursor::query_with_batch_size(client.clone(), namespace, DEFAULT_BATCH_SIZE, flags,
                                      number_to_skip,
                                      number_to_return, query,
                                      return_field_selector, is_cmd_cursor)
    }

    /// Helper method to create a "get more" request.
    ///
    /// # Return value
    ///
    /// Returns the newly-created method.
    fn new_get_more_request(&mut self) -> Message {
        Message::new_get_more(self.client.get_req_id(),
                              self.namespace.to_owned(),
                              self.batch_size, self.cursor_id)
    }

    fn get_from_stream(&mut self) -> Result<()> {
        let stream = try!(self.client.acquire_stream());
        let mut socket = stream.get_socket();

        let get_more = self.new_get_more_request();
        try!(get_more.write(&mut socket));
        let reply = try!(Message::read(&mut socket));

        let (v, _) = try!(Cursor::get_bson_and_cid_from_message(reply));
        self.buffer.extend(v);
        Ok(())
    }

    /// Attempts to read a specified number of BSON documents from the cursor.
    ///
    /// # Arguments
    ///
    /// `n` - The number of documents to read.
    ///
    /// # Return value
    ///
    /// Returns a vector containing the BSON documents that were read.
    pub fn next_n(&mut self, n: i32) -> Result<Vec<bson::Document>> {
        let mut vec = vec![];

        for _ in 0..n {
            let bson_option = self.next();

            match bson_option {
                Some(Ok(bson)) => vec.push(bson),
                Some(Err(err)) => return Err(err),
                None => break,
            };
        }

        Ok(vec)
    }

    /// Attempts to read a batch of BSON documents from the cursor.
    ///
    /// # Return value
    ///
    /// Returns a vector containing the BSON documents that were read.
    pub fn next_batch(&mut self) -> Result<Vec<bson::Document>> {
        let n = self.batch_size;
        self.next_n(n)
    }

    /// Checks whether there are any more documents for the cursor to return.
    ///
    /// # Return value
    ///
    /// Returns `true` if the cursor is not yet exhausted, or `false` if it is.
    pub fn has_next(&mut self) -> Result<bool> {
        if self.limit > 0 && self.count >= self.limit {
            Ok(false)
        } else {
            if self.buffer.is_empty() && self.limit != 1 && self.cursor_id != 0 {
                try!(self.get_from_stream());
            }
            Ok(!self.buffer.is_empty())
        }
    }
}

impl Iterator for Cursor {
    type Item = Result<bson::Document>;

    /// Attempts to read a BSON document from the cursor.
    ///
    /// # Return value
    ///
    /// Returns a BSON document if there is another one to return; `None` if
    /// there are no more documents to return; or an Error if the request for
    /// another document fails.
    fn next(&mut self) -> Option<Result<bson::Document>> {
        match self.has_next() {
            Ok(true) => {
                self.count += 1;
                match self.buffer.pop_front() {
                    Some(bson) => Some(Ok(bson)),
                    None => None,
                }
            },
            Ok(false) => None,
            Err(err) => Some(Err(err)),
        }
    }
}
<commit_msg>Return specific cursor errors<commit_after>use bson::{self, Bson};
use {Error, Client, ThreadedClient, Result};
use wire_protocol::flags::OpQueryFlags;
use wire_protocol::operations::Message;
use std::collections::vec_deque::VecDeque;
use std::io::{Read, Write};

pub const DEFAULT_BATCH_SIZE: i32 = 20;

/// Maintains a connection to the server and lazily returns documents from a
/// query.
///
/// # Fields
///
/// `client` - The client to read from.
/// `namespace` - The namespace to read and write from.
/// `batch_size` - How many documents to fetch at a given time from the server.
/// `cursor_id` - Uniquely identifies the cursor being returned by the reply.
/// `limit` - An upper bound on the total number of documents this cursor
///           should return.
/// `count` - How many documents have been returned so far.
/// `buffer` - A cache for documents received from the query that have not
///            yet been returned.
pub struct Cursor {
    client: Client,
    namespace: String,
    batch_size: i32,
    cursor_id: i64,
    limit: i32,
    count: i32,
    buffer: VecDeque<bson::Document>,
}

impl Cursor {
    /// Construcs a new Cursor for a database command.
    ///
    /// # Arguments
    ///
    /// `client` - Client making the request.
    /// `db` - Which database the command is being sent to.
    /// `doc` - Specifies the command that is being run.
    ///
    /// # Return value
    ///
    /// Returns the newly created Cursor on success, or an Error on failure.
    pub fn command_cursor(client: Client, db: &str,
                          doc: bson::Document) -> Result<Cursor> {
        Cursor::query_with_batch_size(client.clone(), format!("{}.$cmd", db),
                                      1, OpQueryFlags::no_flags(), 0, 0,
                                      doc, None, true)
    }

    fn get_bson_and_cid_from_message(message: Message) -> Result<(VecDeque<bson::Document>, i64)> {
        match message {
            Message::OpReply { header: _, flags: _, cursor_id: cid,
                               starting_from: _, number_returned: _,
                               documents: docs } => {
                let mut v = VecDeque::new();

                if !docs.is_empty() {
                    if let Some(&Bson::I32(ref code)) = docs[0].get("code") {
                        // If command doesn't exist or namespace not found, return
                        // an empty array instead of throwing an error.
                        if *code == 59 || *code == 26 {
                            return Ok((v, cid));
                        } else if let Some(&Bson::String(ref msg)) = docs[0].get("errmsg") {
                            return Err(Error::OperationError(msg.to_owned()));
                        }
                    }
                }

                for doc in docs {
                    v.push_back(doc.clone());
                }

                Ok((v, cid))
            },
            _ => Err(Error::CursorNotFoundError)
        }
    }

    fn get_bson_and_cursor_info_from_command_message(message: Message) -> Result<(VecDeque<bson::Document>, i64, String)> {
        let (v, _) = try!(Cursor::get_bson_and_cid_from_message(message));
        if v.len() != 1 {
            return Err(Error::CursorNotFoundError);
        }

        let ref doc = v[0];

        // Extract cursor information
        if let Some(&Bson::Document(ref cursor)) = doc.get("cursor") {
            if let Some(&Bson::I64(ref id)) = cursor.get("id") {
                if let Some(&Bson::String(ref ns)) = cursor.get("ns") {
                    if let Some(&Bson::Array(ref batch)) = cursor.get("firstBatch") {

                        // Extract first batch documents
                        let map = batch.iter().filter_map(|bdoc| {
                            if let &Bson::Document(ref doc) = bdoc {
                                Some(doc.clone())
                            } else {
                                None
                            }
                        }).collect();

                        return Ok((map, *id, ns.to_owned()))
                    }
                }
            }
        }

        Err(Error::CursorNotFoundError)
    }

    /// Executes a query where the batch size of the returned cursor is
    /// specified.
    ///
    /// # Arguments
    ///
    /// `client` - The client to read from.
    /// `namespace` - The namespace to read and write from.
    /// `batch_size` - How many documents the cursor should return at a time.
    /// `flags` - Bit vector of query options.
    /// `number_to_skip` - The number of initial documents to skip over in the
    ///                    query results.
    /// `number_to_return - The total number of documents that should be
    ///                     returned by the query.
    /// `return_field_selector - An optional projection of which fields should
    ///                          be present in the documents to be returned by
    ///                          the query.
    /// `is_cmd_cursor` - Whether or not the Cursor is for a database command.
    ///
    /// # Return value
    ///
    /// Returns the cursor for the query results on success, or an Error on
    /// failure.
    pub fn query_with_batch_size(client: Client, namespace: String,
                                 batch_size: i32, flags: OpQueryFlags,
                                 number_to_skip: i32, number_to_return: i32,
                                 query: bson::Document,
                                 return_field_selector: Option<bson::Document>,
                                 is_cmd_cursor: bool) -> Result<Cursor> {
        let result = Message::new_query(client.get_req_id(), flags,
                                        namespace.to_owned(),
                                        number_to_skip, batch_size,
                                        query, return_field_selector);

        let stream = try!(client.acquire_stream());
        let mut socket = stream.get_socket();

        let message = try!(result);
        try!(message.write(&mut socket));
        let reply = try!(Message::read(&mut socket));

        let (buf, cursor_id, namespace) = if is_cmd_cursor {
            try!(Cursor::get_bson_and_cursor_info_from_command_message(reply))
        } else {
            let (buf, id) = try!(Cursor::get_bson_and_cid_from_message(reply));
            (buf, id, namespace)
        };

        Ok(Cursor { client: client, namespace: namespace,
                    batch_size: batch_size, cursor_id: cursor_id,
                    limit: number_to_return, count: 0, buffer: buf, })
    }

    /// Executes a query with the default batch size.
    ///
    /// # Arguments
    ///
    /// `client` - The client to read from.
    /// `namespace` - The namespace to read and write from.
    /// `flags` - Bit vector of query options.
    /// `number_to_skip` - The number of initial documents to skip over in the
    ///                    query results.
    /// `number_to_return - The total number of documents that should be
    ///                     returned by the query.
    /// `query` - Specifies which documents to return.
    /// `return_field_selector - An optional projection of which fields should
    ///                          be present in the documents to be returned by
    ///                          the query.
    /// `is_cmd_cursor` - Whether or not the Cursor is for a database command.
    ///
    /// # Return value
    ///
    /// Returns the cursor for the query results on success, or an error string
    /// on failure.
    pub fn query(client: Client, namespace: String,
                 flags: OpQueryFlags, number_to_skip: i32,
                 number_to_return: i32, query: bson::Document,
                 return_field_selector: Option<bson::Document>,
                 is_cmd_cursor: bool) -> Result<Cursor> {
        Cursor::query_with_batch_size(client.clone(), namespace, DEFAULT_BATCH_SIZE, flags,
                                      number_to_skip,
                                      number_to_return, query,
                                      return_field_selector, is_cmd_cursor)
    }

    /// Helper method to create a "get more" request.
    ///
    /// # Return value
    ///
    /// Returns the newly-created method.
    fn new_get_more_request(&mut self) -> Message {
        Message::new_get_more(self.client.get_req_id(),
                              self.namespace.to_owned(),
                              self.batch_size, self.cursor_id)
    }

    fn get_from_stream(&mut self) -> Result<()> {
        let stream = try!(self.client.acquire_stream());
        let mut socket = stream.get_socket();

        let get_more = self.new_get_more_request();
        try!(get_more.write(&mut socket));
        let reply = try!(Message::read(&mut socket));

        let (v, _) = try!(Cursor::get_bson_and_cid_from_message(reply));
        self.buffer.extend(v);
        Ok(())
    }

    /// Attempts to read a specified number of BSON documents from the cursor.
    ///
    /// # Arguments
    ///
    /// `n` - The number of documents to read.
    ///
    /// # Return value
    ///
    /// Returns a vector containing the BSON documents that were read.
    pub fn next_n(&mut self, n: i32) -> Result<Vec<bson::Document>> {
        let mut vec = vec![];

        for _ in 0..n {
            let bson_option = self.next();

            match bson_option {
                Some(Ok(bson)) => vec.push(bson),
                Some(Err(err)) => return Err(err),
                None => break,
            };
        }

        Ok(vec)
    }

    /// Attempts to read a batch of BSON documents from the cursor.
    ///
    /// # Return value
    ///
    /// Returns a vector containing the BSON documents that were read.
    pub fn next_batch(&mut self) -> Result<Vec<bson::Document>> {
        let n = self.batch_size;
        self.next_n(n)
    }

    /// Checks whether there are any more documents for the cursor to return.
    ///
    /// # Return value
    ///
    /// Returns `true` if the cursor is not yet exhausted, or `false` if it is.
    pub fn has_next(&mut self) -> Result<bool> {
        if self.limit > 0 && self.count >= self.limit {
            Ok(false)
        } else {
            if self.buffer.is_empty() && self.limit != 1 && self.cursor_id != 0 {
                try!(self.get_from_stream());
            }
            Ok(!self.buffer.is_empty())
        }
    }
}

impl Iterator for Cursor {
    type Item = Result<bson::Document>;

    /// Attempts to read a BSON document from the cursor.
    ///
    /// # Return value
    ///
    /// Returns a BSON document if there is another one to return; `None` if
    /// there are no more documents to return; or an Error if the request for
    /// another document fails.
    fn next(&mut self) -> Option<Result<bson::Document>> {
        match self.has_next() {
            Ok(true) => {
                self.count += 1;
                match self.buffer.pop_front() {
                    Some(bson) => Some(Ok(bson)),
                    None => None,
                }
            },
            Ok(false) => None,
            Err(err) => Some(Err(err)),
        }
    }
}
<|endoftext|>


<commit_before>// Unfortunately, compiletest_rs does not support 'only-linux',
// so we need to ignore Windows and OS X instead
// ignore-macos: Uses Linux-only APIs
// ignore-windows: Uses Linux-only APIs
#![feature(rustc_private)]
extern crate libc;

fn main() {
    let mut buf = [0u8; 5];
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_getrandom, 0 as *mut libc::c_void, 0 as libc::size_t, 0 as libc::c_uint), 0);
        assert_eq!(libc::syscall(libc::SYS_getrandom, buf.as_mut_ptr() as *mut libc::c_void, 5 as libc::size_t, 0 as libc::c_uint), 5);

        assert_eq!(libc::getrandom(0 as *mut libc::c_void, 0 as libc::size_t, 0 as libc::c_uint), 0);
        assert_eq!(libc::getrandom(buf.as_mut_ptr() as *mut libc::c_void, 5 as libc::size_t, 0 as libc::c_uint), 5);
    }
}
<commit_msg>Improve formatting<commit_after>// Unfortunately, compiletest_rs does not support 'only-linux',
// so we need to ignore Windows and macOS instead.
// ignore-macos: Uses Linux-only APIs
// ignore-windows: Uses Linux-only APIs
#![feature(rustc_private)]
extern crate libc;

fn main() {
    let mut buf = [0u8; 5];
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_getrandom, 0 as *mut libc::c_void, 0 as libc::size_t, 0 as libc::c_uint), 0);
        assert_eq!(libc::syscall(libc::SYS_getrandom, buf.as_mut_ptr() as *mut libc::c_void, 5 as libc::size_t, 0 as libc::c_uint), 5);

        assert_eq!(libc::getrandom(0 as *mut libc::c_void, 0 as libc::size_t, 0 as libc::c_uint), 0);
        assert_eq!(libc::getrandom(buf.as_mut_ptr() as *mut libc::c_void, 5 as libc::size_t, 0 as libc::c_uint), 5);
    }
}
<|endoftext|>


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

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import Method %>

// Non-standard properties that Gecko uses for XUL elements.
<% data.new_style_struct("XUL", inherited=False) %>

${helpers.single_keyword("-moz-box-align", "stretch start center baseline end",
                         products="gecko", gecko_ffi_name="mBoxAlign",
                         gecko_enum_prefix="StyleBoxAlign",
                         gecko_inexhaustive=True,
                         animatable=False,
                         alias="-webkit-box-align",
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-align)")}

${helpers.predefined_type("-moz-box-flex", "Number", "0.0", "parse_non_negative",
                          products="gecko", gecko_ffi_name="mBoxFlex",
                          needs_context=False,
                          animatable=False,
                          alias="-webkit-box-flex",
                          spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-flex)")}

${helpers.single_keyword("-moz-stack-sizing", "stretch-to-fit ignore",
                         products="gecko", gecko_ffi_name="mStretchStack",
                         gecko_constant_prefix="NS_STYLE_STACK_SIZING",
                         animatable=False,
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-stack-sizing)")}
<commit_msg>Auto merge of #15817 - Manishearth:xul, r=bholley<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />
<% from data import Method %>

// Non-standard properties that Gecko uses for XUL elements.
<% data.new_style_struct("XUL", inherited=False) %>

${helpers.single_keyword("-moz-box-align", "stretch start center baseline end",
                         products="gecko", gecko_ffi_name="mBoxAlign",
                         gecko_enum_prefix="StyleBoxAlign",
                         gecko_inexhaustive=True,
                         animatable=False,
                         alias="-webkit-box-align",
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-align)")}

${helpers.single_keyword("-moz-box-direction", "normal reverse",
                         products="gecko", gecko_ffi_name="mBoxDirection",
                         gecko_enum_prefix="StyleBoxDirection",
                         animatable=False,
                         alias="-webkit-box-direction",
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-direction)")}

${helpers.predefined_type("-moz-box-flex", "Number", "0.0", "parse_non_negative",
                          products="gecko", gecko_ffi_name="mBoxFlex",
                          needs_context=False,
                          animatable=False,
                          alias="-webkit-box-flex",
                          spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-flex)")}

${helpers.single_keyword("-moz-box-orient", "horizontal vertical",
                         products="gecko", gecko_ffi_name="mBoxOrient",
                         gecko_enum_prefix="StyleBoxOrient",
                         animatable=False,
                         alias="-webkit-box-orient",
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-orient)")}

${helpers.single_keyword("-moz-box-pack", "start center end justify",
                         products="gecko", gecko_ffi_name="mBoxPack",
                         gecko_enum_prefix="StyleBoxPack",
                         animatable=False,
                         alias="-webkit-box-pack",
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/box-pack)")}

${helpers.single_keyword("-moz-stack-sizing", "stretch-to-fit ignore",
                         products="gecko", gecko_ffi_name="mStretchStack",
                         gecko_constant_prefix="NS_STYLE_STACK_SIZING",
                         animatable=False,
                         spec="Nonstandard (https://developer.mozilla.org/en-US/docs/Web/CSS/-moz-stack-sizing)")}
<|endoftext|>


<commit_before>use ffmpeg;

use errors::*;

/// Initialize the encoding stuff. Should be called once.
pub fn initialize() -> Result<()> {
    ffmpeg::init().chain_err(|| "error initializing ffmpeg")?;

    Ok(())
}

command!(cap_set_video_encoder, |engine| {
    let mut args = engine.args();

    if args.len() != 2 {
        let mut buf = String::new();

        buf.push_str("Usage:\n");
        buf.push_str("    cap_set_video_encoder <encoder>\n");
        buf.push_str("     - Set the video encoder by name.\n");
        buf.push_str("Example:\n");
        buf.push_str("    cap_set_video_encoder libx264\n");

        engine.con_print(&buf);
        return;
    }

    let encoder_name = args.nth(1).unwrap();

    if let Some(encoder) = ffmpeg::encoder::find_by_name(&encoder_name) {
        if let Ok(video) = encoder.video() {
            let mut buf = String::new();

            buf.push_str(&format!("Selected encoder: {}\n", encoder_name));
            buf.push_str(&format!("Description: {}\n", encoder.description()));
            buf.push_str("Pixel formats: ");

            if let Some(formats) = video.formats() {
                buf.push_str(&format!("{:?}\n", formats.collect::<Vec<_>>()));
            } else {
                buf.push_str("any\n");
            }

            engine.con_print(&buf);
        } else {
            engine.con_print(&format!("Invalid encoder type '{}'\n", encoder_name));
        }
    } else {
        engine.con_print(&format!("Unknown encoder '{}'\n", encoder_name));
    }
});
<commit_msg>Store the selected video encoder<commit_after>use ffmpeg;
use std::sync::Mutex;

use errors::*;

lazy_static! {
    static ref VIDEO_ENCODER: Mutex<Option<ffmpeg::codec::Video>> = Mutex::new(None);
}

/// Initialize the encoding stuff. Should be called once.
pub fn initialize() -> Result<()> {
    ffmpeg::init().chain_err(|| "error initializing ffmpeg")?;

    Ok(())
}

command!(cap_set_video_encoder, |engine| {
    let mut args = engine.args();

    if args.len() != 2 {
        let mut buf = String::new();

        buf.push_str("Usage:\n");
        buf.push_str("    cap_set_video_encoder <encoder>\n");
        buf.push_str("     - Set the video encoder by name.\n");
        buf.push_str("Example:\n");
        buf.push_str("    cap_set_video_encoder libx264\n");

        engine.con_print(&buf);
        return;
    }

    let encoder_name = args.nth(1).unwrap();

    if let Some(encoder) = ffmpeg::encoder::find_by_name(&encoder_name) {
        if let Ok(video) = encoder.video() {
            let mut buf = String::new();

            buf.push_str(&format!("Selected encoder: {}\n", encoder_name));
            buf.push_str(&format!("Description: {}\n", encoder.description()));
            buf.push_str("Pixel formats: ");

            if let Some(formats) = video.formats() {
                buf.push_str(&format!("{:?}\n", formats.collect::<Vec<_>>()));
            } else {
                buf.push_str("any\n");
            }

            engine.con_print(&buf);

            *VIDEO_ENCODER.lock().unwrap() = Some(video);
        } else {
            engine.con_print(&format!("Invalid encoder type '{}'\n", encoder_name));
        }
    } else {
        engine.con_print(&format!("Unknown encoder '{}'\n", encoder_name));
    }
});
<|endoftext|>


<commit_before>//! An adventure story modeled as actions on bits.
//!
//! Each state in the story is a sequence of bits, each representing something that happened.
//! There are two numbers stored for each state, one for the state and one for the options.
//! This type of modeling can be used to experiment with a story in an explorative way.
//!
//! By applying filtering to the data, one could build stories as abstractions
//! and reuse them in contexts where one needs control of the events.

type Data = Vec<(u64, u64)>;
type Actions = Vec<&'static str>;

fn is_set(state: u64, i: u64) -> bool {
    ((state >> i) & 0b1) == 0b1
}

fn set(state: u64, i: u64) -> u64 {
    state ^ (0b1 << i)
}

/// Computes the changed bit and whether it is inverse or not.
///
/// Returns index and whether it is inverted.
fn change(state: u64, prev_state: u64) -> uint {
    let diff = state ^ prev_state;
    for i in range(0u64, 64) {
        if is_set(diff, i) {
            if ((prev_state >> i) & 0b1) == 0b1 {
                fail!("Can not invert action.");
            } else {
                return i as uint;
            }
        }
    }

    fail!("State differs more or less than one bit");
}

fn with_options(state: u64, data: &Data, f: |u64|) {
    for &(id, options) in data.iter() {
        if id != state { continue; }
        for i in range(0u64, 64) {
            if is_set(options, i) {
                f(set(state, i))
            }
        }
    }
}

fn main() {
    let actions: Actions = vec!(
        "I used to live each summer in a house by the ocean.",
        "One day I went on a fishing trip but came out for a storm.",
        "The storm crushed my boat against a lonely island.",
        "Luckily a big boat picked me up before I sank.",

        // _0000
        "The boat was owned by am old fisherman.",
        "I search the whole island for water, but could not find any.",
        "Finally I settle down for the night in front of a fire.",
        "Suddenly I heard a roar in the dark.",

        // _0000_0000
        "A huge troll charged against me with a liften axe.",
        "Needless to say, the troll was soon dead.",
        "I climbed up in a tree faster than a squirrel in flight.",
        "Unfortunately the tree was no match for the troll's axe.",

        // _0000_0000_0000
        "I thanked him for his bravery.",
        "'You're welcome' he said and smiled with pointed teeth.",
        "That sight made me jump overboard and sim as fast as I could away from the boat.",
        "Fortunately the closest island was not far away.",

        // _0000_0000_0000_0000
        "Before I had time to reflect uppon the situation, I heard a squeeky sound.",
        "I looked down and saw a small troll grabbing my boots.",
        "A rope had tangled my feet and dragged me underwater.",
        "For a moment I could hold my breath, but then everything went black.",

        // _0000_0000_0000_0000_0000
        "I woke up in a strange place.",
        "There were people around me, but they were not quite human.",
        "Instead of arms and legs, they had fins.",
        "They told me every human that drown ends up here.",

        // _0000_0000_0000_0000_0000_0000
        "The tree hit the ground slowly enough for me to run.",
        "In the dark I did not saw where I was headed.",
        "My body fell down a cliff and everything went black.",
        "Instead of arms and legs, they had wings.",

        // _0000_0000_0000_0000_0000_0000_0000
        "They told me every human that falls down a cliff ends up here.",
        "Next day I looked in every direction, but there was no signs of land.",
        "I built a boat of the troll's skeleton.",
    );
	let data: Data = vec!(
        (0b0, 0b1),
        (0b1, 0b10),
        (0b11, 0b100_0000_0000_0000_1100),
        (0b_1011, 0b1_0000),
        (0b111, 0b10_0000),
        (0b10_0111, 0b100_0000),
        (0b110_0111, 0b1000_0000),
        (0b1110_0111, 0b1_0000_0000),
        (0b1_1110_0111, 0b110_0000_0000),
        (0b101_1110_0111, 0b1000_0000_0000),
        (0b11011, 0b1_0000_0000_0000),
        (0b1_0000_0001_1011, 0b10_0000_0000_0000),
        (0b11_0000_0001_1011, 0b100_0000_0000_0000),
        (0b111_0000_0001_1011, 0b1000_0000_0000_0000),
        (0b1111_0000_0001_1011, 0b1_0000_0000_0000_0000),
        (0b1_1111_0000_0001_1011, 0b10_0000_0000_0000_0000),
        (0b11_1111_0000_0001_1011, 0b1000_0000),
        (0b11_1111_0000_1001_1011, 0b1_0000_0000),
        (0b100_0000_0000_0000_0011, 0b1000_0000_0000_0000_0000),
        (0b1100_0000_0000_0000_0011, 0b1_0000_0000_0000_0000_0000),
        (0b1_1100_0000_0000_0000_0011, 0b10_0000_0000_0000_0000_0000),
        (0b11_1100_0000_0000_0000_0011, 0b100_0000_0000_0000_0000_0000),
        (0b111_1100_0000_0000_0000_0011, 0b1000_0000_0000_0000_0000_0000),
        (0b1101_1110_0111, 0b1_0000_0000_0000_0000_0000_0000),
        (0b1_0000_0000_0000_1101_1110_0111, 0b10_0000_0000_0000_0000_0000_0000),
        (0b11_0000_0000_0000_1101_1110_0111, 0b100_0000_0000_0000_0000_0000_0000),
        (0b111_0000_0000_0000_1101_1110_0111, 0b1_0000_0000_0000_0000_0000),
        (0b111_0001_0000_0000_1101_1110_0111, 0b10_0000_0000_0000_0000_0000),
        (0b111_0011_0000_0000_1101_1110_0111, 0b1000_0000_0000_0000_0000_0000_0000),
        (0b1111_0011_0000_0000_1101_1110_0111, 0b1_0000_0000_0000_0000_0000_0000_0000),
        (0b11_1110_0111, 0b10_0000_0000_0000_0000_0000_0000_0000),
        (0b10_0000_0000_0000_0000_0011_1110_0111, 0b100_0000_0000_0000_0000_0000_0000_0000),
    );
    let states: Vec<u64> = vec!(
        0b0,
        0b1,
        0b11,

            0b111,
            0b10_0111,
            0b110_0111,
            0b1110_0111,
            0b1_1110_0111,

                0b11_1110_0111,
                0b10_0000_0000_0000_0000_0011_1110_0111, // building a boat of the troll's skeleton

                // 0b101_1110_0111,
                // 0b1101_1110_0111,
                // 0b1_0000_0000_0000_1101_1110_0111,
                // 0b11_0000_0000_0000_1101_1110_0111,
                // 0b111_0000_0000_0000_1101_1110_0111,
                // 0b111_0001_0000_0000_1101_1110_0111,
                // 0b111_0011_0000_0000_1101_1110_0111, // waking up with bird people
                // 0b1111_0011_0000_0000_1101_1110_0111,

            // 0b_1011,
            // 0b1_1011, // the old fisherman picked me up.
            // 0b1_0000_0001_1011,
            // 0b11_0000_0001_1011,
            // 0b111_0000_0001_1011,
            // 0b1111_0000_0001_1011, // squeeky sound
            // 0b1_1111_0000_0001_1011,
            // 0b11_1111_0000_0001_1011,
            // 0b11_1111_0000_1001_1011,

            // 0b100_0000_0000_0000_0011, // dragged down under water
            // 0b1100_0000_0000_0000_0011,
            // 0b1_1100_0000_0000_0000_0011,
            // 0b11_1100_0000_0000_0000_0011,
            // 0b111_1100_0000_0000_0000_0011,
    );
    let mut i = 0;
    for state in states.as_slice().windows(2) {
        let ch = change(state[1], state[0]);
        println!("{}", actions.get(ch));
        if i == states.len() - 2 {
            with_options(state[1], &data, |opt: u64| {
                let opt_ch = change(opt, state[1]);
                println!("0b{:t}: {}", opt, actions.get(opt_ch));
            });
        }
        i += 1;
    }
    println!("0b{:t}", *states.get(states.len() - 1));
}

<commit_msg>Removed adventure example<commit_after><|endoftext|>


<commit_before>use std::cell::RefCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{AtomicUsize, Ordering};

use object::{Obj, Ref};
use threads::THREAD;

pub const HANDLE_SIZE: usize = 256;

pub struct HandleMemory {
    /// all buffers, Box is important since HandleBuffer
    /// is a big struct that needs to get moved/copied on resizes
    buffers: RefCell<Vec<Box<HandleBuffer>>>,

    // store addresses of inserted borders
    borders: RefCell<Vec<BorderData>>,

    // index of next free position in buffer
    free: AtomicUsize,
}

impl HandleMemory {
    pub fn new() -> HandleMemory {
        let buffer = box HandleBuffer::new();

        HandleMemory {
            buffers: RefCell::new(vec![buffer]),
            borders: RefCell::new(Vec::new()),
            free: AtomicUsize::new(0),
        }
    }

    pub fn root<T>(&self, obj: Ref<T>) -> Handle<T> {
        if self.free.load(Ordering::Relaxed) >= HANDLE_SIZE {
            self.push_buffer();
            self.free.store(0, Ordering::Relaxed);
        }

        let mut buffers = self.buffers.borrow_mut();
        let buffer = buffers.last_mut().unwrap();

        let idx = self.free.load(Ordering::Relaxed);
        let elem = &mut buffer.elements[idx];
        self.free.store(idx + 1, Ordering::Relaxed);

        *elem = obj.cast::<Obj>();

        Handle(elem as *mut Ref<Obj> as *mut Ref<T>)
    }

    fn push_buffer(&self) {
        self.buffers.borrow_mut().push(box HandleBuffer::new());
    }

    pub fn push_border(&self) {
        let buffer = self.buffers.borrow().len();
        let element = self.free.load(Ordering::Relaxed);

        self.borders.borrow_mut().push(BorderData {
            buffer: buffer,
            element: element,
        });
    }

    pub fn pop_border(&self) {
        let border = self.borders.borrow_mut().pop().expect("no border left");

        self.buffers.borrow_mut().truncate(border.buffer);
        self.free.store(border.element, Ordering::Relaxed);
    }

    pub fn iter(&self) -> HandleMemoryIter {
        let len = self.buffers.borrow().len();

        HandleMemoryIter {
            mem: self,
            buffer_idx: 0,
            element_idx: 0,
            full_buffer_len: if len == 0 { 0 } else { len - 1 },
            last_buffer_len: self.free.load(Ordering::Relaxed),
        }
    }
}

pub fn root<T>(obj: Ref<T>) -> Handle<T> {
    THREAD.with(|thread| thread.borrow().handles.root(obj))
}

struct HandleBuffer {
    elements: [Ref<Obj>; HANDLE_SIZE],
}

impl HandleBuffer {
    fn new() -> HandleBuffer {
        HandleBuffer {
            elements: [Ref::null(); HANDLE_SIZE],
        }
    }
}

struct BorderData {
    buffer: usize,
    element: usize,
}

pub struct Handle<T>(*mut Ref<T>);

impl<T> Handle<T> {
    pub fn direct(self) -> Ref<T> {
        unsafe { *self.0 }
    }

    pub fn raw(self) -> *mut Ref<T> {
        self.0
    }

    pub fn cast<R>(self) -> Handle<R> {
        Handle(self.0 as *mut Ref<R>)
    }
}

impl<T> Deref for Handle<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.0 }.deref()
    }
}

impl<T> DerefMut for Handle<T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.0 }.deref_mut()
    }
}

// known limitation of #[derive(Copy, Clone)]
// traits need to be implemented manually
impl<T> Copy for Handle<T> {}
impl<T> Clone for Handle<T> {
    fn clone(&self) -> Handle<T> {
        *self
    }
}

pub struct HandleMemoryIter<'a> {
    mem: &'a HandleMemory,
    buffer_idx: usize,
    element_idx: usize,
    full_buffer_len: usize,
    last_buffer_len: usize,
}

impl<'a> Iterator for HandleMemoryIter<'a> {
    type Item = Handle<Obj>;

    fn next(&mut self) -> Option<Handle<Obj>> {
        if self.buffer_idx < self.full_buffer_len {
            if self.element_idx < HANDLE_SIZE {
                let idx = self.element_idx;
                self.element_idx += 1;

                let mut buffers = self.mem.buffers.borrow_mut();
                let buffer = &mut buffers[self.buffer_idx];
                return Some(Handle(&mut buffer.elements[idx] as *mut Ref<Obj>));
            } else {
                self.buffer_idx += 1;
                self.element_idx = 0;
            }
        }

        if self.buffer_idx == self.full_buffer_len {
            if self.element_idx < self.last_buffer_len {
                let idx = self.element_idx;
                self.element_idx += 1;

                let mut buffers = self.mem.buffers.borrow_mut();
                let buffer = &mut buffers[self.buffer_idx];
                return Some(Handle(&mut buffer.elements[idx] as *mut Ref<Obj>));
            } else {
                return None;
            }
        }

        None
    }
}
<commit_msg>use Mutex for HandleMemory for now<commit_after>use parking_lot::{Mutex, MutexGuard};
use std::ops::{Deref, DerefMut};

use object::{Obj, Ref};
use threads::THREAD;

pub const HANDLE_SIZE: usize = 256;

pub struct HandleMemory {
    inner: Mutex<HandleMemoryInner>,
}

impl HandleMemory {
    pub fn new() -> HandleMemory {
        HandleMemory {
            inner: Mutex::new(HandleMemoryInner::new()),
        }
    }

    pub fn root<T>(&self, obj: Ref<T>) -> Handle<T> {
        self.inner.lock().root(obj)
    }

    pub fn push_border(&self) {
        self.inner.lock().push_border();
    }

    pub fn pop_border(&self) {
        self.inner.lock().pop_border();
    }

    pub fn iter(&self) -> HandleMemoryIter {
        let inner = self.inner.lock();
        let len = inner.buffers.len();
        let free = inner.free;

        HandleMemoryIter {
            mem: inner,
            buffer_idx: 0,
            element_idx: 0,
            full_buffer_len: if len == 0 { 0 } else { len - 1 },
            last_buffer_len: free,
        }
    }
}

pub struct HandleMemoryInner {
    /// all buffers, Box is important since HandleBuffer
    /// is a big struct that needs to get moved/copied on resizes
    buffers: Vec<Box<HandleBuffer>>,

    // store addresses of inserted borders
    borders: Vec<BorderData>,

    // index of next free position in buffer
    free: usize,
}

impl HandleMemoryInner {
    pub fn new() -> HandleMemoryInner {
        let buffer = box HandleBuffer::new();

        HandleMemoryInner {
            buffers: vec![buffer],
            borders: Vec::new(),
            free: 0,
        }
    }

    pub fn root<T>(&mut self, obj: Ref<T>) -> Handle<T> {
        if self.free >= HANDLE_SIZE {
            self.push_buffer();
            self.free = 0;
        }

        let buffer = self.buffers.last_mut().unwrap();

        let idx = self.free;
        let elem = &mut buffer.elements[idx];
        self.free = idx + 1;

        *elem = obj.cast::<Obj>();

        Handle(elem as *mut Ref<Obj> as *mut Ref<T>)
    }

    pub fn push_buffer(&mut self) {
        self.buffers.push(box HandleBuffer::new());
    }

    pub fn push_border(&mut self) {
        let buffer = self.buffers.len();
        let element = self.free;

        self.borders.push(BorderData {
            buffer: buffer,
            element: element,
        });
    }

    pub fn pop_border(&mut self) {
        let border = self.borders.pop().expect("no border left");

        self.buffers.truncate(border.buffer);
        self.free = border.element;
    }
}

pub fn root<T>(obj: Ref<T>) -> Handle<T> {
    THREAD.with(|thread| thread.borrow().handles.root(obj))
}

struct HandleBuffer {
    elements: [Ref<Obj>; HANDLE_SIZE],
}

impl HandleBuffer {
    fn new() -> HandleBuffer {
        HandleBuffer {
            elements: [Ref::null(); HANDLE_SIZE],
        }
    }
}

struct BorderData {
    buffer: usize,
    element: usize,
}

pub struct Handle<T>(*mut Ref<T>);

impl<T> Handle<T> {
    pub fn direct(self) -> Ref<T> {
        unsafe { *self.0 }
    }

    pub fn raw(self) -> *mut Ref<T> {
        self.0
    }

    pub fn cast<R>(self) -> Handle<R> {
        Handle(self.0 as *mut Ref<R>)
    }
}

impl<T> Deref for Handle<T> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { &*self.0 }.deref()
    }
}

impl<T> DerefMut for Handle<T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.0 }.deref_mut()
    }
}

// known limitation of #[derive(Copy, Clone)]
// traits need to be implemented manually
impl<T> Copy for Handle<T> {}
impl<T> Clone for Handle<T> {
    fn clone(&self) -> Handle<T> {
        *self
    }
}

pub struct HandleMemoryIter<'a> {
    mem: MutexGuard<'a, HandleMemoryInner>,
    buffer_idx: usize,
    element_idx: usize,
    full_buffer_len: usize,
    last_buffer_len: usize,
}

impl<'a> Iterator for HandleMemoryIter<'a> {
    type Item = Handle<Obj>;

    fn next(&mut self) -> Option<Handle<Obj>> {
        if self.buffer_idx < self.full_buffer_len {
            if self.element_idx < HANDLE_SIZE {
                let idx = self.element_idx;
                self.element_idx += 1;

                let buffer = &mut self.mem.buffers[self.buffer_idx];
                return Some(Handle(&mut buffer.elements[idx] as *mut Ref<Obj>));
            } else {
                self.buffer_idx += 1;
                self.element_idx = 0;
            }
        }

        if self.buffer_idx == self.full_buffer_len {
            if self.element_idx < self.last_buffer_len {
                let idx = self.element_idx;
                self.element_idx += 1;

                let buffer = &mut self.mem.buffers[self.buffer_idx];
                return Some(Handle(&mut buffer.elements[idx] as *mut Ref<Obj>));
            } else {
                return None;
            }
        }

        None
    }
}
<|endoftext|>


<commit_before>use base64;
use crate::buffer;
use crate::image;
use std::{fs, io};

use image_crate::ImageFormat::{Jpeg, Png};
use std::path::Path;
use crate::{Document, Error, Gltf, Result};

/// Return type of `import`.
type Import = (Document, Vec<buffer::Data>, Vec<image::Data>);

/// Represents the set of URI schemes the importer supports.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum Scheme<'a> {
    /// `data:[<media type>];base64,<data>`.
    Data(Option<&'a str>, &'a str),

    /// `file:[//]<absolute file path>`.
    ///
    /// Note: The file scheme does not implement authority.
    File(&'a str),

    /// `../foo`, etc.
    Relative,

    /// Placeholder for an unsupported URI scheme identifier.
    Unsupported,
}

impl<'a> Scheme<'a> {
    fn parse<'s>(uri: &'s str) -> Scheme<'s> {
        if uri.contains(":") {
            if uri.starts_with("data:") {
                let match0 = &uri["data:".len()..].split(";base64,").nth(0);
                let match1 = &uri["data:".len()..].split(";base64,").nth(1);
                if match1.is_some() {
                    Scheme::Data(Some(match0.unwrap()), match1.unwrap())
                } else if match0.is_some() {
                    Scheme::Data(None, match0.unwrap())
                } else {
                    Scheme::Unsupported
                }
            } else if uri.starts_with("file://") {
                Scheme::File(&uri["file://".len()..])
            } else if uri.starts_with("file:") {
                Scheme::File(&uri["file:".len()..])
            } else {
                Scheme::Unsupported
            }
        } else {
            Scheme::Relative
        }
    }

    fn read(base: &Path, uri: &str) -> Result<Vec<u8>> {
        match Scheme::parse(uri) {
            Scheme::Data(_, base64) => base64::decode(&base64).map_err(Error::Base64),
            Scheme::File(path) => read_to_end(path),
            Scheme::Relative => read_to_end(base.join(uri)),
            Scheme::Unsupported => Err(Error::UnsupportedScheme),
        }
    }
}

fn read_to_end<P>(path: P) -> Result<Vec<u8>>
where P: AsRef<Path>
{
    use io::Read;
    let file = fs::File::open(path.as_ref()).map_err(Error::Io)?;
    let length = file.metadata().map(|x| x.len()).unwrap_or(0);
    let mut reader = io::BufReader::new(file);
    let mut data = Vec::with_capacity(length as usize);
    reader.read_to_end(&mut data).map_err(Error::Io)?;
    Ok(data)
}

/// Import the buffer data referenced by a glTF document.
pub fn import_buffer_data(
    document: &Document,
    base: Option<&Path>,
    mut blob: Option<Vec<u8>>,
) -> Result<Vec<buffer::Data>> {
    let mut buffers = Vec::new();
    for buffer in document.buffers() {
        let mut data = match buffer.source() {
            buffer::Source::Uri(uri) if base.is_some() => Scheme::read(base.unwrap(), uri),
            buffer::Source::Bin => blob.take().ok_or(Error::MissingBlob),
            _ => Err(Error::ExternalReferenceInSliceImport)
        }?;
        if data.len() < buffer.length() {
            return Err(
                Error::BufferLength {
                    buffer: buffer.index(),
                    expected: buffer.length(),
                    actual: data.len(),
                }
            );
        }
        while data.len() % 4 != 0 {
            data.push(0);
        }
        buffers.push(buffer::Data(data));
    }
    Ok(buffers)
}

/// Import the image data referenced by a glTF document.
pub fn import_image_data(
    document: &Document,
    base: Option<&Path>,
    buffer_data: &[buffer::Data],
) -> Result<Vec<image::Data>> {
    let mut images = Vec::new();
    #[cfg(feature = "guess_mime_type")]
    let guess_format = |encoded_image: &[u8]| {
        match image_crate::guess_format(encoded_image) {
            Ok(image_crate::ImageFormat::Png) => Some(Png),
            Ok(image_crate::ImageFormat::Jpeg) => Some(Jpeg),
            _ => None,
        }
    };
    #[cfg(not(feature = "guess_mime_type"))]
    let guess_format = |_encoded_image: &[u8]| {
        None
    };
    for image in document.images() {
        match image.source() {
            image::Source::Uri { uri, mime_type } if base.is_some() => {
                match Scheme::parse(uri) {
                    Scheme::Data(Some(annoying_case), base64) => {
                        let encoded_image = base64::decode(&base64).map_err(Error::Base64)?;
                        let encoded_format = match annoying_case.as_ref() {
                            "image/png" => Png,
                            "image/jpeg" => Jpeg,
                            _ => match guess_format(&encoded_image) {
                                Some(format) => format,
                                None => return Err(Error::UnsupportedImageEncoding),
                            },
                        };
                        let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                        images.push(image::Data::new(decoded_image));
                        continue;
                    },
                    Scheme::Unsupported => return Err(Error::UnsupportedScheme),
                    _ => {},
                }
                let encoded_image = Scheme::read(base.unwrap(), uri)?;
                let encoded_format =  match mime_type {
                    Some("image/png") => Png,
                    Some("image/jpeg") => Jpeg,
                    Some(_) => match guess_format(&encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                    None => match uri.rsplit(".").next() {
                        Some("png") => Png,
                        Some("jpg") | Some("jpeg") => Jpeg,
                        _ => match guess_format(&encoded_image) {
                            Some(format) => format,
                            None => return Err(Error::UnsupportedImageEncoding),
                        },
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            image::Source::View { view, mime_type } => {
                let parent_buffer_data = &buffer_data[view.buffer().index()].0;
                let begin = view.offset();
                let end = begin + view.length();
                let encoded_image = &parent_buffer_data[begin..end];
                let encoded_format = match mime_type {
                    "image/png" => Png,
                    "image/jpeg" => Jpeg,
                    _ => match guess_format(encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            _ => return Err(Error::ExternalReferenceInSliceImport)
        }
    }

    Ok(images)
}

fn import_impl(Gltf { document, blob }: Gltf, base: Option<&Path>) -> Result<Import> {
    let buffer_data = import_buffer_data(&document, base, blob)?;
    let image_data = import_image_data(&document, base, &buffer_data)?;
    let import = (document, buffer_data, image_data);
    Ok(import)
}

fn import_path(path: &Path) -> Result<Import> {
    let base = path.parent().unwrap_or(Path::new("./"));
    let file = fs::File::open(path).map_err(Error::Io)?;
    let reader = io::BufReader::new(file);
    import_impl(Gltf::from_reader(reader)?, Some(base))
}

/// Import some glTF 2.0 from the file system.
///
/// ```
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.gltf";
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import(path)?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
///
/// ### Note
///
/// This function is provided as a convenience for loading glTF and associated
/// resources from the file system. It is suitable for real world use but may
/// not be suitable for all real world use cases. More complex import scenarios
/// such downloading from web URLs are not handled by this function. These
/// scenarios are delegated to the user.
///
/// You can read glTF without loading resources by constructing the [`Gltf`]
/// (standard glTF) or [`Glb`] (binary glTF) data structures explicitly.
///
/// [`Gltf`]: struct.Gltf.html
/// [`Glb`]: struct.Glb.html
pub fn import<P>(path: P) -> Result<Import>
    where P: AsRef<Path>
{
    import_path(path.as_ref())
}

pub fn import_slice_impl(slice: &[u8]) -> Result<Import> {
    import_impl(Gltf::from_slice(slice)?, None)
}

/// Import some glTF 2.0 from a slice
///
/// ```
/// # extern crate gltf;
/// # use std::fs;
/// # use std::io::Read;
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.glb";
/// # let mut file = fs::File::open(path).map_err(gltf::Error::Io)?;
/// # let mut bytes = Vec::new();
/// # file.read_to_end(&mut bytes).map_err(gltf::Error::Io)?;
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import_slice(bytes.as_slice())?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
pub fn import_slice<S>(slice: S) -> Result<Import>
    where S: AsRef<[u8]>
{
    import_slice_impl(slice.as_ref())
}
<commit_msg>Prevent an extra allocation and copy per imported file<commit_after>use base64;
use crate::buffer;
use crate::image;
use std::{fs, io};

use image_crate::ImageFormat::{Jpeg, Png};
use std::path::Path;
use crate::{Document, Error, Gltf, Result};

/// Return type of `import`.
type Import = (Document, Vec<buffer::Data>, Vec<image::Data>);

/// Represents the set of URI schemes the importer supports.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum Scheme<'a> {
    /// `data:[<media type>];base64,<data>`.
    Data(Option<&'a str>, &'a str),

    /// `file:[//]<absolute file path>`.
    ///
    /// Note: The file scheme does not implement authority.
    File(&'a str),

    /// `../foo`, etc.
    Relative,

    /// Placeholder for an unsupported URI scheme identifier.
    Unsupported,
}

impl<'a> Scheme<'a> {
    fn parse<'s>(uri: &'s str) -> Scheme<'s> {
        if uri.contains(":") {
            if uri.starts_with("data:") {
                let match0 = &uri["data:".len()..].split(";base64,").nth(0);
                let match1 = &uri["data:".len()..].split(";base64,").nth(1);
                if match1.is_some() {
                    Scheme::Data(Some(match0.unwrap()), match1.unwrap())
                } else if match0.is_some() {
                    Scheme::Data(None, match0.unwrap())
                } else {
                    Scheme::Unsupported
                }
            } else if uri.starts_with("file://") {
                Scheme::File(&uri["file://".len()..])
            } else if uri.starts_with("file:") {
                Scheme::File(&uri["file:".len()..])
            } else {
                Scheme::Unsupported
            }
        } else {
            Scheme::Relative
        }
    }

    fn read(base: &Path, uri: &str) -> Result<Vec<u8>> {
        match Scheme::parse(uri) {
            Scheme::Data(_, base64) => base64::decode(&base64).map_err(Error::Base64),
            Scheme::File(path) => read_to_end(path),
            Scheme::Relative => read_to_end(base.join(uri)),
            Scheme::Unsupported => Err(Error::UnsupportedScheme),
        }
    }
}

fn read_to_end<P>(path: P) -> Result<Vec<u8>>
where P: AsRef<Path>
{
    use io::Read;
    let file = fs::File::open(path.as_ref()).map_err(Error::Io)?;
    // Allocate one extra byte so the buffer doesn't need to grow before the
    // final `read` call at the end of the file.  Don't worry about `usize`
    // overflow because reading will fail regardless in that case.
    let length = file.metadata().map(|x| x.len() + 1).unwrap_or(0);
    let mut reader = io::BufReader::new(file);
    let mut data = Vec::with_capacity(length as usize);
    reader.read_to_end(&mut data).map_err(Error::Io)?;
    Ok(data)
}

/// Import the buffer data referenced by a glTF document.
pub fn import_buffer_data(
    document: &Document,
    base: Option<&Path>,
    mut blob: Option<Vec<u8>>,
) -> Result<Vec<buffer::Data>> {
    let mut buffers = Vec::new();
    for buffer in document.buffers() {
        let mut data = match buffer.source() {
            buffer::Source::Uri(uri) if base.is_some() => Scheme::read(base.unwrap(), uri),
            buffer::Source::Bin => blob.take().ok_or(Error::MissingBlob),
            _ => Err(Error::ExternalReferenceInSliceImport)
        }?;
        if data.len() < buffer.length() {
            return Err(
                Error::BufferLength {
                    buffer: buffer.index(),
                    expected: buffer.length(),
                    actual: data.len(),
                }
            );
        }
        while data.len() % 4 != 0 {
            data.push(0);
        }
        buffers.push(buffer::Data(data));
    }
    Ok(buffers)
}

/// Import the image data referenced by a glTF document.
pub fn import_image_data(
    document: &Document,
    base: Option<&Path>,
    buffer_data: &[buffer::Data],
) -> Result<Vec<image::Data>> {
    let mut images = Vec::new();
    #[cfg(feature = "guess_mime_type")]
    let guess_format = |encoded_image: &[u8]| {
        match image_crate::guess_format(encoded_image) {
            Ok(image_crate::ImageFormat::Png) => Some(Png),
            Ok(image_crate::ImageFormat::Jpeg) => Some(Jpeg),
            _ => None,
        }
    };
    #[cfg(not(feature = "guess_mime_type"))]
    let guess_format = |_encoded_image: &[u8]| {
        None
    };
    for image in document.images() {
        match image.source() {
            image::Source::Uri { uri, mime_type } if base.is_some() => {
                match Scheme::parse(uri) {
                    Scheme::Data(Some(annoying_case), base64) => {
                        let encoded_image = base64::decode(&base64).map_err(Error::Base64)?;
                        let encoded_format = match annoying_case.as_ref() {
                            "image/png" => Png,
                            "image/jpeg" => Jpeg,
                            _ => match guess_format(&encoded_image) {
                                Some(format) => format,
                                None => return Err(Error::UnsupportedImageEncoding),
                            },
                        };
                        let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                        images.push(image::Data::new(decoded_image));
                        continue;
                    },
                    Scheme::Unsupported => return Err(Error::UnsupportedScheme),
                    _ => {},
                }
                let encoded_image = Scheme::read(base.unwrap(), uri)?;
                let encoded_format =  match mime_type {
                    Some("image/png") => Png,
                    Some("image/jpeg") => Jpeg,
                    Some(_) => match guess_format(&encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                    None => match uri.rsplit(".").next() {
                        Some("png") => Png,
                        Some("jpg") | Some("jpeg") => Jpeg,
                        _ => match guess_format(&encoded_image) {
                            Some(format) => format,
                            None => return Err(Error::UnsupportedImageEncoding),
                        },
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(&encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            image::Source::View { view, mime_type } => {
                let parent_buffer_data = &buffer_data[view.buffer().index()].0;
                let begin = view.offset();
                let end = begin + view.length();
                let encoded_image = &parent_buffer_data[begin..end];
                let encoded_format = match mime_type {
                    "image/png" => Png,
                    "image/jpeg" => Jpeg,
                    _ => match guess_format(encoded_image) {
                        Some(format) => format,
                        None => return Err(Error::UnsupportedImageEncoding),
                    },
                };
                let decoded_image = image_crate::load_from_memory_with_format(encoded_image, encoded_format)?;
                images.push(image::Data::new(decoded_image));
            },
            _ => return Err(Error::ExternalReferenceInSliceImport)
        }
    }

    Ok(images)
}

fn import_impl(Gltf { document, blob }: Gltf, base: Option<&Path>) -> Result<Import> {
    let buffer_data = import_buffer_data(&document, base, blob)?;
    let image_data = import_image_data(&document, base, &buffer_data)?;
    let import = (document, buffer_data, image_data);
    Ok(import)
}

fn import_path(path: &Path) -> Result<Import> {
    let base = path.parent().unwrap_or(Path::new("./"));
    let file = fs::File::open(path).map_err(Error::Io)?;
    let reader = io::BufReader::new(file);
    import_impl(Gltf::from_reader(reader)?, Some(base))
}

/// Import some glTF 2.0 from the file system.
///
/// ```
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.gltf";
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import(path)?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
///
/// ### Note
///
/// This function is provided as a convenience for loading glTF and associated
/// resources from the file system. It is suitable for real world use but may
/// not be suitable for all real world use cases. More complex import scenarios
/// such downloading from web URLs are not handled by this function. These
/// scenarios are delegated to the user.
///
/// You can read glTF without loading resources by constructing the [`Gltf`]
/// (standard glTF) or [`Glb`] (binary glTF) data structures explicitly.
///
/// [`Gltf`]: struct.Gltf.html
/// [`Glb`]: struct.Glb.html
pub fn import<P>(path: P) -> Result<Import>
    where P: AsRef<Path>
{
    import_path(path.as_ref())
}

pub fn import_slice_impl(slice: &[u8]) -> Result<Import> {
    import_impl(Gltf::from_slice(slice)?, None)
}

/// Import some glTF 2.0 from a slice
///
/// ```
/// # extern crate gltf;
/// # use std::fs;
/// # use std::io::Read;
/// # fn run() -> Result<(), gltf::Error> {
/// # let path = "examples/Box.glb";
/// # let mut file = fs::File::open(path).map_err(gltf::Error::Io)?;
/// # let mut bytes = Vec::new();
/// # file.read_to_end(&mut bytes).map_err(gltf::Error::Io)?;
/// # #[allow(unused)]
/// let (document, buffers, images) = gltf::import_slice(bytes.as_slice())?;
/// # Ok(())
/// # }
/// # fn main() {
/// #     run().expect("test failure");
/// # }
/// ```
pub fn import_slice<S>(slice: S) -> Result<Import>
    where S: AsRef<[u8]>
{
    import_slice_impl(slice.as_ref())
}
<|endoftext|>