<commit_before>#![allow(dead_code)]

use std;
use std::sync::Arc;

use arrayvec::ArrayVec;

use smallvec::Array;
use small_string::SmallString;
use small_string_utils::{insert_at_char, split_string_near_byte};

// Internal node min/max values.
const MAX_CHILDREN: usize = 32;
const MIN_CHILDREN: usize = MAX_CHILDREN / 2;

// Leaf node min/max values.
const MAX_BYTES: usize = 384;
const MIN_BYTES: usize = MAX_BYTES / 2;

#[derive(Copy, Clone)]
struct BackingArray([u8; MAX_BYTES]);
unsafe impl Array for BackingArray {
    type Item = u8;
    fn size() -> usize {
        MAX_BYTES
    }
    fn ptr(&self) -> *const u8 {
        &self.0[0]
    }
    fn ptr_mut(&mut self) -> *mut u8 {
        &mut self.0[0]
    }
}

// Type alias used for char count, grapheme count, line count, etc. storage
// in nodes.
type Count = u32;

//-------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct Rope {
    root: Arc<Node>,
}

impl Rope {
    pub fn new() -> Rope {
        Rope { root: Arc::new(Node::new()) }
    }

    pub fn char_count(&self) -> Count {
        self.root.char_count()
    }

    pub fn insert(&mut self, char_pos: Count, text: &str) {
        let root = Arc::make_mut(&mut self.root);
        if let (char_count, Some(r_node)) = root.insert(char_pos, text) {
            let mut char_counts = ArrayVec::new();
            char_counts.push(char_count);
            char_counts.push(r_node.char_count());

            let mut children = ArrayVec::new();
            let mut l_node = Node::Empty;
            std::mem::swap(&mut l_node, root);
            children.push(Arc::new(l_node));
            children.push(Arc::new(r_node));

            *root = Node::Internal {
                char_counts: char_counts,
                children: children,
            };
        }
    }
}

//-------------------------------------------------------------

#[derive(Debug, Clone)]
enum Node {
    Empty,
    Leaf(SmallString<BackingArray>),
    Internal {
        char_counts: ArrayVec<[Count; MAX_CHILDREN]>,
        children: ArrayVec<[Arc<Node>; MAX_CHILDREN]>,
    },
}

impl Node {
    fn new() -> Node {
        Node::Empty
    }

    fn char_count(&self) -> Count {
        match self {
            &Node::Empty => 0,
            &Node::Leaf(ref text) => text.chars().count() as Count,
            &Node::Internal { ref char_counts, .. } => {
                char_counts.iter().fold(0, |a, b| a + b) as Count
            }
        }
    }

    /// Inserts the text at the given char index.
    ///
    /// Returns an updated char count for the node, and optionally a right-side
    /// residual node that overflowed the max char or node count.
    ///
    /// TODO: handle the situation where what's being inserted is larger
    /// than MAX_BYTES.
    fn insert(&mut self, char_pos: Count, text: &str) -> (Count, Option<Node>) {
        match self {
            // If it's empty, turn it into a leaf
            &mut Node::Empty => {
                *self = Node::Leaf(text.into());
                return (text.chars().count() as Count, None);
            }

            // If it's a leaf
            &mut Node::Leaf(ref mut cur_text) => {
                insert_at_char(cur_text, text, char_pos as usize);

                if cur_text.len() <= MAX_BYTES {
                    return (cur_text.chars().count() as Count, None);
                } else {
                    let split_pos = cur_text.len() / 2;
                    let right_text = split_string_near_byte(cur_text, split_pos);
                    cur_text.shrink_to_fit();

                    return (
                        cur_text.chars().count() as Count,
                        Some(Node::Leaf(right_text)),
                    );
                }
            }

            // If it's internal, things get a little more complicated
            &mut Node::Internal {
                ref mut char_counts,
                ref mut children,
            } => {
                // Find the child to traverse into along with its starting char
                // offset.
                let (child_i, start_char) = {
                    let mut child_i = 0;
                    let mut start_char = 0;
                    for c in char_counts.iter() {
                        let tmp = start_char + *c;
                        if char_pos <= tmp {
                            break;
                        }

                        start_char = tmp;
                        child_i += 1;
                    }
                    (child_i.min(children.len() - 1), start_char)
                };

                // Navigate into the appropriate child
                let (updated_char_count, residual) =
                    Arc::make_mut(&mut children[child_i]).insert(char_pos - start_char, text);
                char_counts[child_i] = updated_char_count;

                // Handle the new node, if any.
                if let Some(r_node) = residual {
                    // The new node will fit as a child of this node
                    if children.len() < MAX_CHILDREN {
                        char_counts.insert(child_i + 1, r_node.char_count());
                        children.insert(child_i + 1, Arc::new(r_node));
                        return (char_counts.iter().sum(), None);
                    }
                    // The new node won't fit!  Must split.
                    else {
                        let extra_count = char_counts
                            .try_insert(child_i + 1, r_node.char_count())
                            .err()
                            .unwrap()
                            .element();
                        let extra_child = children
                            .try_insert(child_i + 1, Arc::new(r_node))
                            .err()
                            .unwrap()
                            .element();

                        let mut r_char_counts = ArrayVec::new();
                        let mut r_children = ArrayVec::new();

                        let r_count = (children.len() + 1) / 2;
                        let l_count = (children.len() + 1) - r_count;

                        for _ in l_count..children.len() {
                            r_char_counts.push(char_counts.remove(l_count));
                            r_children.push(children.remove(l_count));
                        }
                        r_char_counts.push(extra_count);
                        r_children.push(extra_child);

                        return (
                            char_counts.iter().sum(),
                            Some(Node::Internal {
                                char_counts: r_char_counts,
                                children: r_children,
                            }),
                        );
                    }
                } else {
                    // No new node.  Easy.
                    return (char_counts.iter().sum(), None);
                }
            }
        }
    }
}
<commit_msg>Laying the groundwork for grapheme and line tracking.<commit_after>#![allow(dead_code)]

use std;
use std::sync::Arc;

use arrayvec::ArrayVec;

use smallvec::Array;
use small_string::SmallString;
use small_string_utils::{insert_at_char, split_string_near_byte};

// Internal node min/max values.
const MAX_CHILDREN: usize = 16;
const MIN_CHILDREN: usize = MAX_CHILDREN / 2;

// Leaf node min/max values.
const MAX_BYTES: usize = 320;
const MIN_BYTES: usize = MAX_BYTES / 2;

#[derive(Copy, Clone)]
struct BackingArray([u8; MAX_BYTES]);
unsafe impl Array for BackingArray {
    type Item = u8;
    fn size() -> usize {
        MAX_BYTES
    }
    fn ptr(&self) -> *const u8 {
        &self.0[0]
    }
    fn ptr_mut(&mut self) -> *mut u8 {
        &mut self.0[0]
    }
}

// Type alias used for char count, grapheme count, line count, etc. storage
// in nodes.
type Count = u32;

//-------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct Rope {
    root: Arc<Node>,
}

impl Rope {
    pub fn new() -> Rope {
        use std::mem::size_of;
        println!("Node size: {:?}", size_of::<Node>());
        Rope { root: Arc::new(Node::new()) }
    }

    pub fn char_count(&self) -> Count {
        self.root.text_info().chars
    }

    pub fn insert(&mut self, char_pos: Count, text: &str) {
        let root = Arc::make_mut(&mut self.root);
        if let Some(r_node) = root.insert(char_pos, text) {
            let mut info = ArrayVec::new();
            info.push(root.text_info());
            info.push(r_node.text_info());

            let mut children = ArrayVec::new();
            let mut l_node = Node::Empty;
            std::mem::swap(&mut l_node, root);
            children.push(Arc::new(l_node));
            children.push(Arc::new(r_node));

            *root = Node::Internal {
                info: info,
                children: children,
            };
        }
    }
}

//-------------------------------------------------------------

#[derive(Debug, Copy, Clone)]
struct TextInfo {
    chars: Count,
    graphemes: Count,
    newlines: Count,
}

impl TextInfo {
    fn new() -> TextInfo {
        TextInfo {
            chars: 0,
            graphemes: 0,
            newlines: 0,
        }
    }

    fn from_str(text: &str) -> TextInfo {
        TextInfo {
            chars: text.chars().count() as Count,
            graphemes: 0, // TODO
            newlines: 0, // TODO
        }
    }

    fn combine(&self, other: &TextInfo) -> TextInfo {
        TextInfo {
            chars: self.chars + other.chars,
            graphemes: self.graphemes + other.graphemes,
            newlines: self.newlines + other.newlines,
        }
    }
}

//-------------------------------------------------------------

#[derive(Debug, Clone)]
enum Node {
    Empty,
    Leaf(SmallString<BackingArray>),
    Internal {
        info: ArrayVec<[TextInfo; MAX_CHILDREN]>,
        children: ArrayVec<[Arc<Node>; MAX_CHILDREN]>,
    },
}

impl Node {
    fn new() -> Node {
        Node::Empty
    }

    fn text_info(&self) -> TextInfo {
        match self {
            &Node::Empty => TextInfo::new(),
            &Node::Leaf(ref text) => TextInfo::from_str(text),
            &Node::Internal { ref info, .. } => {
                info.iter().fold(TextInfo::new(), |a, b| a.combine(b))
            }
        }
    }

    /// Inserts the text at the given char index.
    ///
    /// Returns a right-side residual node if the insertion wouldn't fit
    /// within this node.
    ///
    /// TODO: handle the situation where what's being inserted is larger
    /// than MAX_BYTES.
    fn insert(&mut self, char_pos: Count, text: &str) -> Option<Node> {
        match self {
            // If it's empty, turn it into a leaf
            &mut Node::Empty => {
                *self = Node::Leaf(text.into());
                return None;
            }

            // If it's a leaf
            &mut Node::Leaf(ref mut cur_text) => {
                insert_at_char(cur_text, text, char_pos as usize);

                if cur_text.len() <= MAX_BYTES {
                    return None;
                } else {
                    let split_pos = cur_text.len() / 2;
                    let right_text = split_string_near_byte(cur_text, split_pos);
                    cur_text.shrink_to_fit();

                    return Some(Node::Leaf(right_text));
                }
            }

            // If it's internal, things get a little more complicated
            &mut Node::Internal {
                ref mut info,
                ref mut children,
            } => {
                // Find the child to traverse into along with its starting char
                // offset.
                let (child_i, start_char) = {
                    let mut child_i = 0;
                    let mut start_char = 0;
                    for &TextInfo { chars: c, .. } in info.iter() {
                        let tmp = start_char + c;
                        if char_pos <= tmp {
                            break;
                        }

                        start_char = tmp;
                        child_i += 1;
                    }
                    (child_i.min(children.len() - 1), start_char)
                };

                // Navigate into the appropriate child
                let residual =
                    Arc::make_mut(&mut children[child_i]).insert(char_pos - start_char, text);
                info[child_i] = children[child_i].text_info();

                // Handle the new node, if any.
                if let Some(r_node) = residual {
                    // The new node will fit as a child of this node
                    if children.len() < MAX_CHILDREN {
                        info.insert(child_i + 1, r_node.text_info());
                        children.insert(child_i + 1, Arc::new(r_node));
                        return None;
                    }
                    // The new node won't fit!  Must split.
                    else {
                        let extra_info = info.try_insert(child_i + 1, r_node.text_info())
                            .err()
                            .unwrap()
                            .element();
                        let extra_child = children
                            .try_insert(child_i + 1, Arc::new(r_node))
                            .err()
                            .unwrap()
                            .element();

                        let mut r_info = ArrayVec::new();
                        let mut r_children = ArrayVec::new();

                        let r_count = (children.len() + 1) / 2;
                        let l_count = (children.len() + 1) - r_count;

                        for _ in l_count..children.len() {
                            r_info.push(info.remove(l_count));
                            r_children.push(children.remove(l_count));
                        }
                        r_info.push(extra_info);
                        r_children.push(extra_child);

                        return Some(Node::Internal {
                            info: r_info,
                            children: r_children,
                        });
                    }
                } else {
                    // No new node.  Easy.
                    return None;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();

    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push(exon.start-1..exon.end);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    if exon.start < last.end && last.start < exon.end {
                        last.end = exon.end;
                        extend = true;
                    }
                }
                if !extend {
                    merged.push(exon.clone());
                }
            }
        }
    }

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in &keys {
        let exons = merged_exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.start as u32, exon.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.end && exon.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.end) -
                                        std::cmp::max(e.start, exon.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.end-exon.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.start,
                    end: exon.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if options.outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&options.outfile)?) });

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<commit_msg>produce merged and unmerged files<commit_after>#![recursion_limit="128"]
#![cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity, trivial_regex))]
use std::str;
use std::vec::Vec;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::io::stdout;
use std::ops::Range;

#[macro_use]
extern crate failure;

extern crate bam2bedgraph;
use bam2bedgraph::*;
use bam2bedgraph::error::*;
use bam2bedgraph::indexed_annotation::*;

extern crate rust_htslib;
use rust_htslib::bam::Read;
use rust_htslib::bam::IndexedReader;

extern crate structopt;
#[macro_use]
extern crate structopt_derive;
use structopt::StructOpt;

extern crate futures;
use futures::Future;
extern crate futures_cpupool;
use futures_cpupool::CpuPool;

extern crate num_cpus;

extern crate ordered_float;
use ordered_float::OrderedFloat;

#[derive(StructOpt, Debug)]
#[structopt(name = "intronrpkm", about = "Analyze RPKM values in intronic space")]
struct Options {
    // input files
    #[structopt(long="gff", help = "A genome annotation file in gff3 format", name="ANNOT_GFF_FILE")]
    annotfile_gff: Option<String>,
    #[structopt(long="gtf", help = "A genome annotation file in gtf format", name="ANNOT_GTF_FILE")]
    annotfile_gtf: Option<String>,
    #[structopt(long="bam1", short="1", help = "The set of stranded .bam files to analyze where read1 indicates strand", name="BAMFILE1")]
    bam1: Vec<String>,
    #[structopt(long="bam2", short="2", help = "The set of stranded .bam files to analyze where read2 indicates strand", name="BAMFILE2")]
    bam2: Vec<String>,
    #[structopt(long="bam", short="u", help = "The set of unstranded .bam files to analyze", name="BAMFILE")]
    bam: Vec<String>,
    #[structopt(long="chrmap", help = "Optional tab-delimited chr name mapping file", name="CHRMAP_FILE")]
    chrmap_file: Option<String>,
    #[structopt(long="vizchrmap", help = "Optional tab-delimited chr name mapping file for bigwig/bigbed exports", name="VIZCHRMAP_FILE")]
    vizchrmap_file: Option<String>,
    #[structopt(long="sizes", help = "Optional chr sizes file", name="SIZES_FILE")]
    sizes_file: Option<String>,
    // output file
    #[structopt(long="out", short="o", help = "Output file", name="OUT_FILE", default_value="-")]
    outfile: String,
    #[structopt(long="merged", short="m", help = "Merged Output file", name="MERGED_OUT_FILE", default_value="-")]
    merged_outfile: Option<String>,
    // feature types filter
    #[structopt(long="exon_type", help = "The exon type(s) to search for", name="EXON_TYPE")]
    exon_type: Vec<String>,
    #[structopt(long="transcript_type", help = "The transcript type(s) to search for", name="TRANSCRIPT_TYPE")]
    transcript_type: Vec<String>,
    #[structopt(long="gene_type", help = "The gene type(s) to search for", name="GENE_TYPE")]
    gene_type: Vec<String>,
    // flags
    #[structopt(long="cpu_threads", short="t", help = "How many threads to use for processing", default_value="0")]
    cpu_threads: usize,
}

#[derive(Ord, Eq, PartialOrd, PartialEq)]
struct Row {
    seqname: String,
    strand: String,
    start: u64,
    end: u64,
    cov: OrderedFloat<f64>,
    rpkm: OrderedFloat<f64>,
}

fn write_exon_cov(
    options: &Options,
    annot: &Arc<IndexedAnnotation>,
    total_reads: u64,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>)
    -> Result<()>
{
    let mut tidmaps = HashMap::<String,HashMap<String,u32>>::new();
    for bamfile in bamfiles {
        let bam = IndexedReader::from_path(bamfile)?;
        // build the tid map for this bam file
        let mut tidmap = HashMap::<String,u32>::new();
        {   let header = bam.header();
            for target_name in header.target_names() {
                let tid = header.tid(target_name).r()?;
                let target_name = String::from(std::str::from_utf8(target_name)?);
                let chr = annot.chrmap.get(&target_name).unwrap_or(&target_name);
                tidmap.insert(chr.clone(), tid);
            }
        }
        tidmaps.insert(bamfile.clone(), tidmap);
    }
    let tidmaps = Arc::new(tidmaps);

    let mut unmerged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    for (gene_row, gene) in annot.rows.iter().enumerate() {
        if options.gene_type.is_empty() || options.gene_type.contains(&gene.feature_type) {
            if let Some(feature_rows) = annot.row2children.get(&gene_row) {
                // get the transcript rows for this gene
                'transcript_row:
                for transcript_row in feature_rows {
                    let transcript = &annot.rows[*transcript_row];
                    // make sure this is a transcript type
                    // make sure transcript coordinates are consistent with the gene coordinates
                    if (options.transcript_type.is_empty() ||
                            options.transcript_type.contains(&transcript.feature_type)) &&
                        transcript.seqname == gene.seqname &&
                        transcript.strand == gene.strand
                    {
                        if let Some(exon_rows) = annot.row2children.get(transcript_row) {
                            // first get exon start/stop -> transcript associations
                            // and splice start/stop -> transcript associations
                            for exon_row in exon_rows {
                                let exon = &annot.rows[*exon_row];
                                if options.exon_type.is_empty() || options.exon_type.contains(&exon.feature_type) {
                                    let exon = &annot.rows[*exon_row];
                                    unmerged_exons.entry((exon.seqname.clone(),exon.strand.clone())).
                                        or_insert_with(Vec::new).push(exon.start-1..exon.end);
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    let mut merged_exons = HashMap::<(String,String),Vec<Range<u64>>>::new();
    let keys = unmerged_exons.keys().map(|k| k.clone()).collect::<Vec<_>>();
    for key in &keys {
        if let Some(ref mut unmerged) = unmerged_exons.get_mut(&key) {
            unmerged.sort_by_key(|u| u.start);
            let merged = merged_exons.entry((*key).clone()).or_insert_with(Vec::new);
            for exon in unmerged.iter() {
                let mut extend = false;
                if let Some(ref mut last) = merged.last_mut() {
                    if exon.start < last.end && last.start < exon.end {
                        last.end = exon.end;
                        extend = true;
                    }
                }
                if !extend {
                    merged.push(exon.clone());
                }
            }
        }
    }

    write_exon_cov_to_file(options,
        total_reads,
        &unmerged_exons,
        &options.outfile,
        bamfiles,
        bamstrand,
        &tidmaps,
    )?;
    if let Some(ref merged_outfile) = options.merged_outfile {
        write_exon_cov_to_file(options,
            total_reads,
            &merged_exons,
            &merged_outfile,
            bamfiles,
            bamstrand,
            &tidmaps,
        )?;
    }
    Ok(())
}

fn write_exon_cov_to_file(
    options: &Options,
    total_reads: u64,
    exons: &HashMap<(String,String),Vec<Range<u64>>>,
    outfile: &str,
    bamfiles: &Vec<String>,
    bamstrand: &Vec<Option<bool>>,
    tidmaps: &Arc<HashMap<String,HashMap<String,u32>>>,
    )
    -> Result<()>
{
    let mut output: BufWriter<Box<Write>> = BufWriter::new(
        if outfile == "-" { Box::new(stdout()) }
            else { Box::new(File::create(&outfile)?) });

    let num_cpus = num_cpus::get();
    let mut pair_futures = Vec::new();
    let pool = Arc::new(CpuPool::new(if options.cpu_threads==0 {num_cpus} else {options.cpu_threads}));
    for key in exons.keys() {
        let exons = exons[&key.clone()].clone();
        for exon in exons {
            let seqname = key.0.clone();
            let strand = key.1.clone();
            let chr = seqname.clone();
            let strand_is_plus = strand == "+";
            let bamfiles = bamfiles.clone();
            let bamstrand = bamstrand.clone();
            let tidmaps = tidmaps.clone();

            let pair_future = pool.spawn_fn(move ||->Result<Row> {
                let mut exon_cov = 0f64;
                let mut exon_reads = HashSet::<String>::new();
                for (i,bamfile) in bamfiles.iter().enumerate() {
                    let read1strand = bamstrand[i];
                    let tidmap = &tidmaps[bamfile];
                    if let Some(tid) = tidmap.get(&chr) {
                        let mut bam = IndexedReader::from_path(bamfile)?;
                        bam.fetch(*tid, exon.start as u32, exon.end as u32)?;
                        for read in bam.records() {
                            let read = read?;
                            // make sure the read's strand matches
                            if let Some(is_read1strand) = read1strand {
                                if !(((is_read1strand == read.is_first_in_template()) == !read.is_reverse()) == strand_is_plus) {
                                    continue;
                                }
                            }
                            let read_name = String::from(str::from_utf8(read.qname())?);
                            exon_reads.insert(read_name);
                            //eprintln!("Looking at read: {}", read_name);
                            let exons = cigar2exons(&read.cigar(), read.pos() as u64)?;
                            for e in exons {
                                if e.start < exon.end && exon.start < e.end {
                                    exon_cov +=
                                        (std::cmp::min(e.end, exon.end) -
                                        std::cmp::max(e.start, exon.start)) as f64;
                                }
                            }
                        }
                    }
                }
                let exon_length = exon.end-exon.start;
                let rpkm = (1e10f64 * exon_reads.len() as f64) / (total_reads as f64 * exon_length as f64);
                Ok(Row{
                    seqname: seqname,
                    strand: strand,
                    start: exon.start,
                    end: exon.end,
                    cov: OrderedFloat((exon_cov / exon_length as f64)),
                    rpkm: OrderedFloat(rpkm),
                })
            });
            pair_futures.push(pair_future);
        }
    }

    let mut rows = Vec::new();
    for future in pair_futures {
        match future.wait() {
            Ok(row) => {
                rows.push(row)
            }
            Err(ref e) => {
                eprintln!("Got Err in write_exon_cov: {:?}", e);
            }
        }
    }
    rows.sort();

    output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
                                  "seqname",
                                  "strand",
                                  "start",
                                  "end",
                                  "cov",
                                  "rpkm",
    ))?;
    for row in rows {
        output.write_fmt(format_args!("{}\t{}\t{}\t{}\t{}\t{}\n",
            row.seqname,
            row.strand,
            row.start,
            row.end,
            row.cov,
            row.rpkm,
        ))?;
    }
    Ok(())
}

fn run() -> Result<()> {
    let mut options = Options::from_args();
    // set defaults for feature types
    options.gene_type =
        (if options.gene_type.is_empty() { vec!["gene".to_string()] }
         else { options.gene_type.clone() }).into_iter().collect();
    //options.transcript_type =
    //    (if options.transcript_type.is_empty() { vec!["transcript".to_string()] }
    //     else { options.transcript_type.clone() }).into_iter().collect();
    options.exon_type =
        (if options.exon_type.is_empty() { vec!["exon".to_string()] }
         else { options.exon_type.clone() }).into_iter().collect();
    // set debug options if --debug flag is set

    // organize bamfiles by read strand type
    let mut bamfiles = Vec::<String>::new();
    bamfiles.append(&mut options.bam1.clone());
    bamfiles.append(&mut options.bam2.clone());
    bamfiles.append(&mut options.bam.clone());
    let mut bamstrand = Vec::<Option<bool>>::new();
    bamstrand.append(&mut options.bam1.iter().map(|_| Some(true)).collect());
    bamstrand.append(&mut options.bam2.iter().map(|_| Some(false)).collect());
    bamstrand.append(&mut options.bam.iter().map(|_| None).collect());
    // get the chromosome names and sizes from the first bam file
    if bamfiles.is_empty() {
        Options::clap().print_help()?;
        eprintln!("\n\nNo bam files were passed in!");
        std::process::exit(1)
    }
    let transcript_type = String::from("transcript");
    let mut annot = if let Some(annotfile_gff) = options.annotfile_gff.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gff);
        IndexedAnnotation::from_gff(
            &annotfile_gff,
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else if let Some(annotfile_gtf) = options.annotfile_gtf.clone() {
        eprintln!("Reading annotation file {:?}", &annotfile_gtf);
        IndexedAnnotation::from_gtf(&annotfile_gtf,
            options.gene_type.get(0).r()?,
            options.transcript_type.get(0).unwrap_or(&transcript_type),
            &options.chrmap_file,
            &options.vizchrmap_file)?
    } else {
        Options::clap().print_help()?;
        eprintln!("\n\nNo annotation file was given!");
        std::process::exit(1);
    };
    eprintln!("Getting refseq lengths from bam file {:?}", &bamfiles[0]);
    let refs = match options.sizes_file.clone() {
        Some(sizes_file) => read_sizes_file(&sizes_file, &annot.chrmap)?,
        None => get_bam_refs(&bamfiles[0], &annot.chrmap)?,
    };
    annot.refs = refs;

    // get the total bam reads
    eprintln!("Running samtools idxstats to get total bam read counts");
    let total_reads = get_bam_total_reads(&bamfiles)?;
    eprintln!("Found {} total reads", total_reads);
    write_exon_cov(&options, &Arc::new(annot), total_reads, &bamfiles, &bamstrand)?;
    Ok(())
}

fn main() {
    // enable stack traces
    std::env::set_var("RUST_BACKTRACE", "full");

    if let Err(ref e) = run() {
        eprintln!("error: {}", e);
        eprintln!("backtrace: {:?}", e.backtrace());
        ::std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>use std::io::Write;
use std::fmt::Write as FmtWrite;
use std::collections::HashMap;
use std::default::Default;
use std::borrow::Cow;

use codegen::{Codegen, Logic, Token, Scope, ContentType, Content, write_indent};

use Error;

///Visibility of modules and structures in Rust.
pub enum Visibility {
    Public,
    Private
}

///Code generator for Rust templates.
pub struct Rust<'a> {
    ///Generate a submodule. Defaults to none and should only be used when a
    ///whole module is generated.
    pub named_module: Option<(&'a str, Visibility)>,

    ///The template visibility. Defaults to public.
    pub visibility: Visibility
}

impl<'a> Rust<'a> {
    fn eval_logic<W: Write>(&self, w: &mut W, first: bool, cond: &Logic, params: &HashMap<String, ContentType>) -> Result<(), Error> {
        match cond {
            &Logic::And(ref conds) => {
                if !first {
                    try!(write!(w, "("));
                }
                for (i, cond) in conds.iter().enumerate() {
                    if i > 0 {
                        try!(write!(w, " && "));
                    }
                    try!(self.eval_logic(w, false, cond, params));
                }
                if !first {
                    try!(write!(w, ")"));
                }
            },
            &Logic::Or(ref conds) => {
                if !first {
                    try!(write!(w, "("));
                }
                for (i, cond) in conds.iter().enumerate() {
                    if i > 0 {
                        try!(write!(w, " || "));
                    }
                    try!(self.eval_logic(w, false, cond, params));
                }
                if !first {
                    try!(write!(w, ")"));
                }
            },
            &Logic::Not(ref cond) => {
                try!(write!(w, "!("));
                try!(self.eval_logic(w, false, cond, params));
                try!(write!(w, ")"));
            },
            &Logic::Value(ref val) => {
                match params.get(val) {
                    Some(&ContentType::String) | Some(&ContentType::Template) => try!(write!(w, "true")),
                    Some(&ContentType::OptionalString) | Some(&ContentType::OptionalTemplate) => try!(write!(w, "self.{}.is_some()", val)),
                    Some(&ContentType::Bool) => try!(write!(w, "self.{}", val)),
                    None => {}
                }
            }
        }

        Ok(())
    }
}

impl<'a> Default for Rust<'a> {
    fn default() -> Rust<'a> {
        Rust {
            named_module: None,
            visibility: Visibility::Public
        }
    }
}

impl<'a> Codegen for Rust<'a> {
    fn build_template<W: Write>(&self, w: &mut W, name: &str, mut indent: u8, params: &HashMap<String, ContentType>, tokens: &[Token]) -> Result<(), Error> {
        let public = match (&self.visibility, &self.named_module) {
            (&Visibility::Public, _) => true,
            (_, &Some(_)) => true,
            _ => false
        };

        let mut string_buf = String::new();
        let mut fmt_args = vec![];

        if public {
            if params.len() > 0 {
                line!(w, indent, "pub struct {}<'a> {{", name);
            } else {
                line!(w, indent, "pub struct {};", name);
            }
        } else {
            if params.len() > 0 {
                line!(w, indent, "struct {}<'a> {{", name);
            } else {
                line!(w, indent, "struct {};", name);
            }
        }

        for (parameter, ty) in params {
            match ty {
                &ContentType::String => line!(w, indent + 1, "pub {}: &'a str,", parameter),
                &ContentType::OptionalString => line!(w, indent + 1, "pub {}: Option<&'a str>,", parameter),
                &ContentType::Bool => line!(w, indent + 1, "pub {}: bool,", parameter),
                &ContentType::Template => line!(w, indent + 1, "pub {}: &'a ::symbiosis_rust::Template,", parameter),
                &ContentType::OptionalTemplate => line!(w, indent + 1, "pub {}: Option<&'a ::symbiosis_rust::Template>,", parameter),
            }
        }

        line!(w, indent, "}}");

        if params.len() > 0 {
            line!(w, indent, "impl<'a> ::symbiosis_rust::Template for {}<'a> {{", name);
        } else {
            line!(w, indent, "impl ::symbiosis_rust::Template for {} {{", name);
        }
        indent += 1;

        line!(w, indent, "fn render_to(&self, writer: &mut ::std::io::Write) -> ::std::io::Result<()> {{");
        indent += 1;

        for token in tokens {
            match token {
                &Token::BeginTag(ref name) => try!(write!(&mut string_buf, "<{}", name.as_slice())),
                &Token::EndTag(_self_close) => try!(write!(&mut string_buf, ">")),
                &Token::CloseTag(ref name) => try!(write!(&mut string_buf, "</{}>", name.as_slice())),
                &Token::BeginAttribute(ref name, ref content) => match content {
                    &Content::String(ref content) => try!(write!(&mut string_buf, " {}=\\\"{}", name.as_slice(), content)),
                    &Content::Placeholder(ref placeholder) => {

                        match params.get(placeholder) {
                            Some(&ContentType::String) => {
                                try!(write!(&mut string_buf, " {}=\\\"{{}}", name.as_slice()));
                                fmt_args.push(format!("self.{}", placeholder));
                            },
                            Some(&ContentType::OptionalString) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(val) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(write!(writer, \"{{}}\", val));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Template) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "try!(self.{}.render_to(writer));", placeholder);
                            },
                            Some(&ContentType::OptionalTemplate) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(template) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(template.render_to(writer));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Bool) => return Err(Error::Parse(vec![Cow::Borrowed("boolean values cannot be rendered as text")])),
                            None => {}
                        }
                    }
                },
                &Token::AppendToAttribute(ref text) | &Token::BeginText(ref text) | &Token::AppendToText(ref text) => match text {
                    &Content::String(ref content) => try!(write!(&mut string_buf, "{}", content)),
                    &Content::Placeholder(ref placeholder) => {
                        match params.get(placeholder) {
                            Some(&ContentType::String) => {
                                try!(write!(&mut string_buf, "{{}}"));
                                fmt_args.push(format!("self.{}", placeholder));
                            },
                            Some(&ContentType::OptionalString) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(val) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(write!(writer, \"{{}}\", val));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Template) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "try!(self.{}.render_to(writer));", placeholder);
                            },
                            Some(&ContentType::OptionalTemplate) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(template) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(template.render_to(writer));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Bool) => return Err(Error::Parse(vec![Cow::Borrowed("boolean values cannot be rendered as text")])),
                            None => {}
                        }
                    }
                },
                &Token::EndAttribute => try!(write!(&mut string_buf, "\\\"")),
                &Token::EndText => {},
                &Token::Scope(Scope::If(ref cond)) => {
                    try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                    try!(write_indent(w, indent));
                    try!(write!(w, "if "));
                    try!(self.eval_logic(w, true, &cond.flattened(), params));
                    try!(write!(w, " {{\n"));
                    indent += 1;
                },
                &Token::End => {
                    try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                    indent -= 1;
                    line!(w, indent, "}}");
                }
            }
        }

        try!(write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));
        indent -= 1;

        line!(w, indent, "}}");
        indent -= 1;

        line!(w, indent, "}}");
        Ok(())
    }

    fn build_module<W, F>(&self, w: &mut W, build_templates: F) -> Result<(), Error> where
        W: Write,
        F: FnOnce(&mut W, u8) -> Result<(), Error>
    {
        let indent = if let Some((ref module, ref visibility)) = self.named_module {
            match visibility {
                &Visibility::Public => line!(w, 0, "pub mod {} {{", module),
                &Visibility::Private => line!(w, 0, "mod {} {{", module)
            }
            1
        } else {
            0
        };

        try!(build_templates(w, indent));

        if self.named_module.is_some() {
            line!(w, 0, "}}");
        }

        Ok(())
    }
}

fn try_write_and_clear_fmt<W: Write>(w: &mut W, indent: u8, buf: &mut String, args: &mut Vec<String>) -> Result<(), Error> {
    if buf.len() > 0 {
        if args.len() == 0 {
            line!(w, indent, "try!(write!(writer, \"{}\"));", buf);
        } else {
            line!(w, indent, "try!(write!(writer, \"{}\", {}));", buf, args.connect(", "));
        }
        buf.clear();
        args.clear();
    }

    Ok(())
}

fn write_and_clear_fmt<W: Write>(w: &mut W, indent: u8, buf: &mut String, args: &mut Vec<String>) -> Result<(), Error> {
    if buf.len() > 0 {
        if args.len() == 0 {
            line!(w, indent, "write!(writer, \"{}\")", buf);
        } else {
            line!(w, indent, "write!(writer, \"{}\", {})", buf, args.connect(", "));
        }
        buf.clear();
        args.clear();
    }

    Ok(())
}
<commit_msg>Only close struct block if there are any fields.<commit_after>use std::io::Write;
use std::fmt::Write as FmtWrite;
use std::collections::HashMap;
use std::default::Default;
use std::borrow::Cow;

use codegen::{Codegen, Logic, Token, Scope, ContentType, Content, write_indent};

use Error;

///Visibility of modules and structures in Rust.
pub enum Visibility {
    Public,
    Private
}

///Code generator for Rust templates.
pub struct Rust<'a> {
    ///Generate a submodule. Defaults to none and should only be used when a
    ///whole module is generated.
    pub named_module: Option<(&'a str, Visibility)>,

    ///The template visibility. Defaults to public.
    pub visibility: Visibility
}

impl<'a> Rust<'a> {
    fn eval_logic<W: Write>(&self, w: &mut W, first: bool, cond: &Logic, params: &HashMap<String, ContentType>) -> Result<(), Error> {
        match cond {
            &Logic::And(ref conds) => {
                if !first {
                    try!(write!(w, "("));
                }
                for (i, cond) in conds.iter().enumerate() {
                    if i > 0 {
                        try!(write!(w, " && "));
                    }
                    try!(self.eval_logic(w, false, cond, params));
                }
                if !first {
                    try!(write!(w, ")"));
                }
            },
            &Logic::Or(ref conds) => {
                if !first {
                    try!(write!(w, "("));
                }
                for (i, cond) in conds.iter().enumerate() {
                    if i > 0 {
                        try!(write!(w, " || "));
                    }
                    try!(self.eval_logic(w, false, cond, params));
                }
                if !first {
                    try!(write!(w, ")"));
                }
            },
            &Logic::Not(ref cond) => {
                try!(write!(w, "!("));
                try!(self.eval_logic(w, false, cond, params));
                try!(write!(w, ")"));
            },
            &Logic::Value(ref val) => {
                match params.get(val) {
                    Some(&ContentType::String) | Some(&ContentType::Template) => try!(write!(w, "true")),
                    Some(&ContentType::OptionalString) | Some(&ContentType::OptionalTemplate) => try!(write!(w, "self.{}.is_some()", val)),
                    Some(&ContentType::Bool) => try!(write!(w, "self.{}", val)),
                    None => {}
                }
            }
        }

        Ok(())
    }
}

impl<'a> Default for Rust<'a> {
    fn default() -> Rust<'a> {
        Rust {
            named_module: None,
            visibility: Visibility::Public
        }
    }
}

impl<'a> Codegen for Rust<'a> {
    fn build_template<W: Write>(&self, w: &mut W, name: &str, mut indent: u8, params: &HashMap<String, ContentType>, tokens: &[Token]) -> Result<(), Error> {
        let public = match (&self.visibility, &self.named_module) {
            (&Visibility::Public, _) => true,
            (_, &Some(_)) => true,
            _ => false
        };

        let mut string_buf = String::new();
        let mut fmt_args = vec![];

        if public {
            if params.len() > 0 {
                line!(w, indent, "pub struct {}<'a> {{", name);
            } else {
                line!(w, indent, "pub struct {};", name);
            }
        } else {
            if params.len() > 0 {
                line!(w, indent, "struct {}<'a> {{", name);
            } else {
                line!(w, indent, "struct {};", name);
            }
        }

        for (parameter, ty) in params {
            match ty {
                &ContentType::String => line!(w, indent + 1, "pub {}: &'a str,", parameter),
                &ContentType::OptionalString => line!(w, indent + 1, "pub {}: Option<&'a str>,", parameter),
                &ContentType::Bool => line!(w, indent + 1, "pub {}: bool,", parameter),
                &ContentType::Template => line!(w, indent + 1, "pub {}: &'a ::symbiosis_rust::Template,", parameter),
                &ContentType::OptionalTemplate => line!(w, indent + 1, "pub {}: Option<&'a ::symbiosis_rust::Template>,", parameter),
            }
        }

        if params.len() > 0 {
            line!(w, indent, "}}");
        }

        if params.len() > 0 {
            line!(w, indent, "impl<'a> ::symbiosis_rust::Template for {}<'a> {{", name);
        } else {
            line!(w, indent, "impl ::symbiosis_rust::Template for {} {{", name);
        }
        indent += 1;

        line!(w, indent, "fn render_to(&self, writer: &mut ::std::io::Write) -> ::std::io::Result<()> {{");
        indent += 1;

        for token in tokens {
            match token {
                &Token::BeginTag(ref name) => try!(write!(&mut string_buf, "<{}", name.as_slice())),
                &Token::EndTag(_self_close) => try!(write!(&mut string_buf, ">")),
                &Token::CloseTag(ref name) => try!(write!(&mut string_buf, "</{}>", name.as_slice())),
                &Token::BeginAttribute(ref name, ref content) => match content {
                    &Content::String(ref content) => try!(write!(&mut string_buf, " {}=\\\"{}", name.as_slice(), content)),
                    &Content::Placeholder(ref placeholder) => {

                        match params.get(placeholder) {
                            Some(&ContentType::String) => {
                                try!(write!(&mut string_buf, " {}=\\\"{{}}", name.as_slice()));
                                fmt_args.push(format!("self.{}", placeholder));
                            },
                            Some(&ContentType::OptionalString) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(val) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(write!(writer, \"{{}}\", val));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Template) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "try!(self.{}.render_to(writer));", placeholder);
                            },
                            Some(&ContentType::OptionalTemplate) => {
                                try!(write!(&mut string_buf, " {}=\\\"", name.as_slice()));
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(template) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(template.render_to(writer));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Bool) => return Err(Error::Parse(vec![Cow::Borrowed("boolean values cannot be rendered as text")])),
                            None => {}
                        }
                    }
                },
                &Token::AppendToAttribute(ref text) | &Token::BeginText(ref text) | &Token::AppendToText(ref text) => match text {
                    &Content::String(ref content) => try!(write!(&mut string_buf, "{}", content)),
                    &Content::Placeholder(ref placeholder) => {
                        match params.get(placeholder) {
                            Some(&ContentType::String) => {
                                try!(write!(&mut string_buf, "{{}}"));
                                fmt_args.push(format!("self.{}", placeholder));
                            },
                            Some(&ContentType::OptionalString) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(val) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(write!(writer, \"{{}}\", val));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Template) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "try!(self.{}.render_to(writer));", placeholder);
                            },
                            Some(&ContentType::OptionalTemplate) => {
                                try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                                line!(w, indent, "if let Some(template) = self.{} {{", placeholder);
                                line!(w, indent + 1, "try!(template.render_to(writer));");
                                line!(w, indent, "}}");
                            },
                            Some(&ContentType::Bool) => return Err(Error::Parse(vec![Cow::Borrowed("boolean values cannot be rendered as text")])),
                            None => {}
                        }
                    }
                },
                &Token::EndAttribute => try!(write!(&mut string_buf, "\\\"")),
                &Token::EndText => {},
                &Token::Scope(Scope::If(ref cond)) => {
                    try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                    try!(write_indent(w, indent));
                    try!(write!(w, "if "));
                    try!(self.eval_logic(w, true, &cond.flattened(), params));
                    try!(write!(w, " {{\n"));
                    indent += 1;
                },
                &Token::End => {
                    try!(try_write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));

                    indent -= 1;
                    line!(w, indent, "}}");
                }
            }
        }

        try!(write_and_clear_fmt(w, indent, &mut string_buf, &mut fmt_args));
        indent -= 1;

        line!(w, indent, "}}");
        indent -= 1;

        line!(w, indent, "}}");
        Ok(())
    }

    fn build_module<W, F>(&self, w: &mut W, build_templates: F) -> Result<(), Error> where
        W: Write,
        F: FnOnce(&mut W, u8) -> Result<(), Error>
    {
        let indent = if let Some((ref module, ref visibility)) = self.named_module {
            match visibility {
                &Visibility::Public => line!(w, 0, "pub mod {} {{", module),
                &Visibility::Private => line!(w, 0, "mod {} {{", module)
            }
            1
        } else {
            0
        };

        try!(build_templates(w, indent));

        if self.named_module.is_some() {
            line!(w, 0, "}}");
        }

        Ok(())
    }
}

fn try_write_and_clear_fmt<W: Write>(w: &mut W, indent: u8, buf: &mut String, args: &mut Vec<String>) -> Result<(), Error> {
    if buf.len() > 0 {
        if args.len() == 0 {
            line!(w, indent, "try!(write!(writer, \"{}\"));", buf);
        } else {
            line!(w, indent, "try!(write!(writer, \"{}\", {}));", buf, args.connect(", "));
        }
        buf.clear();
        args.clear();
    }

    Ok(())
}

fn write_and_clear_fmt<W: Write>(w: &mut W, indent: u8, buf: &mut String, args: &mut Vec<String>) -> Result<(), Error> {
    if buf.len() > 0 {
        if args.len() == 0 {
            line!(w, indent, "write!(writer, \"{}\")", buf);
        } else {
            line!(w, indent, "write!(writer, \"{}\", {})", buf, args.connect(", "));
        }
        buf.clear();
        args.clear();
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::fmt;

use groups::{GroupIter, GroupWindow};
use window::Window;

pub trait LayoutClone {
    fn clone_box(&self) -> Box<Layout>;
}

impl<T> LayoutClone for T
    where T: 'static + Layout + Clone
{
    fn clone_box(&self) -> Box<Layout> {
        Box::new(self.clone())
    }
}

pub trait Layout: LayoutClone {
    fn name(&self) -> &str;
    fn layout(&self, width: i32, height: i32, focused: Option<GroupWindow>, stack: GroupIter);
}

impl Clone for Box<Layout> {
    fn clone(&self) -> Box<Layout> {
        self.clone_box()
    }
}

impl fmt::Debug for Layout {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Layout {{ \"{}\" }}", self.name())
    }
}

#[derive(Clone)]
pub struct TiledLayout {
    name: String,
}

impl TiledLayout {
    pub fn new(name: String) -> Box<Layout> {
        Box::new(TiledLayout { name })
    }
}

impl Layout for TiledLayout {
    fn name(&self) -> &str {
        &self.name
    }

    fn layout(&self, width: i32, height: i32, _: Option<GroupWindow>, stack: GroupIter) {
        if stack.len() == 0 {
            return;
        }

        let tile_height = height / stack.len() as i32;

        for (i, window) in stack.enumerate() {
            window.without_focus_tracking(|window| {
                                              window.map();
                                              window.configure(0,
                                                               i as i32 * tile_height,
                                                               width,
                                                               tile_height);
                                          });
        }
    }
}

#[derive(Clone)]
pub struct StackLayout {
    name: String,
}

impl StackLayout {
    pub fn new(name: String) -> Box<Layout> {
        Box::new(StackLayout { name })
    }
}

impl Layout for StackLayout {
    fn name(&self) -> &str {
        &self.name
    }

    fn layout(&self, width: i32, height: i32, focused: Option<GroupWindow>, stack: GroupIter) {
        if stack.len() == 0 {
            return;
        }

        for window in stack {
            window.without_focus_tracking(|window| window.unmap());
        }
        focused.map(|window| {
                        window.without_focus_tracking(|window| {
                                                          window.map();
                                                          window.configure(0, 0, width, height);
                                                      })
                    });
    }
}
<commit_msg>StackLayout: don't unmap the focused window.<commit_after>use std::fmt;

use groups::{GroupIter, GroupWindow};
use window::Window;

pub trait LayoutClone {
    fn clone_box(&self) -> Box<Layout>;
}

impl<T> LayoutClone for T
    where T: 'static + Layout + Clone
{
    fn clone_box(&self) -> Box<Layout> {
        Box::new(self.clone())
    }
}

pub trait Layout: LayoutClone {
    fn name(&self) -> &str;
    fn layout(&self, width: i32, height: i32, focused: Option<GroupWindow>, stack: GroupIter);
}

impl Clone for Box<Layout> {
    fn clone(&self) -> Box<Layout> {
        self.clone_box()
    }
}

impl fmt::Debug for Layout {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Layout {{ \"{}\" }}", self.name())
    }
}

#[derive(Clone)]
pub struct TiledLayout {
    name: String,
}

impl TiledLayout {
    pub fn new(name: String) -> Box<Layout> {
        Box::new(TiledLayout { name })
    }
}

impl Layout for TiledLayout {
    fn name(&self) -> &str {
        &self.name
    }

    fn layout(&self, width: i32, height: i32, _: Option<GroupWindow>, stack: GroupIter) {
        if stack.len() == 0 {
            return;
        }

        let tile_height = height / stack.len() as i32;

        for (i, window) in stack.enumerate() {
            window.without_focus_tracking(|window| {
                                              window.map();
                                              window.configure(0,
                                                               i as i32 * tile_height,
                                                               width,
                                                               tile_height);
                                          });
        }
    }
}

#[derive(Clone)]
pub struct StackLayout {
    name: String,
}

impl StackLayout {
    pub fn new(name: String) -> Box<Layout> {
        Box::new(StackLayout { name })
    }
}

impl Layout for StackLayout {
    fn name(&self) -> &str {
        &self.name
    }

    fn layout(&self, width: i32, height: i32, focused: Option<GroupWindow>, stack: GroupIter) {
        if stack.len() == 0 {
            return;
        }

        {
            let unfocused = stack.filter(|window| {
                                             focused
                                                 .as_ref()
                                                 .map_or(true, |focused_window| {
                    window.id() != focused_window.id()
                })
                                         });
            for window in unfocused {
                window.without_focus_tracking(|window| window.unmap());
            }
        }
        focused.map(|window| {
                        window.without_focus_tracking(|window| {
                                                          window.map();
                                                          window.configure(0, 0, width, height);
                                                      })
                    });
    }
}
<|endoftext|>


<commit_before>extern crate walkdir;

use constant::*;
use self::walkdir::{DirEntry as WalkDirEntry, WalkDir, WalkDirIterator};
use std::collections::BTreeSet;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::io::Error as IoError;
use std::path::{MAIN_SEPARATOR, Path, PathBuf};
use std::result::Result;

macro_rules! path_buf {
    ($($x: expr),*) => {
        PathBuf::new()
            $(.join($x))*
    };
}

pub fn append_prefix_if_need<S: AsRef<str>>(path: S) -> String {
    let path         = path.as_ref();
    let prefix       = format!(".{}", MAIN_SEPARATOR);
    let prefix: &str = prefix.as_ref();
    format!("{}{}", if path.starts_with(prefix) { "" } else { prefix }, path)
}

pub fn extract_file_name(full_path: &PathBuf) -> Option<&str> {
    full_path
        .file_name()
        .and_then(OsStr::to_str)
}

pub fn ls<P: AsRef<Path>>(path: &P) -> Result<BTreeSet<String>, IoError> {
    let dirs = try!(fs::read_dir(path));
    let dirs = dirs
        .filter_map(Result::ok)
        .map(|d| d.file_name())
        .map(OsString::into_string)
        .filter_map(Result::ok)
        .collect::<BTreeSet<String>>();
    Ok(dirs)
}

fn is_hidden(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, |s| s.starts_with(".") && s.len() > 1 && s != "..")
}

fn is_pinned(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, |s| s == WORKING_DIR_NAME || s == KEEPED_FILE_NAME)
}

fn is_target(entry: &WalkDirEntry) -> bool {
    !is_hidden(entry) && !is_pinned(entry)
}

fn to_string(entry: WalkDirEntry) -> Result<String, OsString> {
    entry
        .path()
        .as_os_str()
        .to_os_string()
        .into_string()
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> BTreeSet<String> {
    let walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(is_target)
        .filter_map(Result::ok)
        .filter(|e| !e.file_type().is_dir())
        .collect::<Vec<WalkDirEntry>>();

    walker
        .into_iter()
        .map(to_string)
        .filter_map(Result::ok)
        .collect::<BTreeSet<String>>()
}

pub fn dirs_ordered_by_descending_depth<P: AsRef<Path>>(root: P) -> Vec<String> {
    let mut walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(is_target)
        .filter_entry(|e| e.file_type().is_dir())
        .filter_map(Result::ok)
        .collect::<Vec<WalkDirEntry>>();

    walker.sort_by(|l, r| {
        let l_depth = &l.depth();
        let r_depth = &r.depth();
        r_depth.cmp(l_depth)
    });

    walker
        .into_iter()
        .map(to_string)
        .filter_map(Result::ok)
        .collect::<Vec<String>>()
}
<commit_msg>[Refactoring] Do not AsRef() because of useless<commit_after>extern crate walkdir;

use constant::*;
use self::walkdir::{DirEntry as WalkDirEntry, WalkDir, WalkDirIterator};
use std::collections::BTreeSet;
use std::ffi::{OsStr, OsString};
use std::fs;
use std::io::Error as IoError;
use std::path::{MAIN_SEPARATOR, Path, PathBuf};
use std::result::Result;

macro_rules! path_buf {
    ($($x: expr),*) => {
        PathBuf::new()
            $(.join($x))*
    };
}

pub fn append_prefix_if_need(path: &String) -> String {
    let prefix       = format!(".{}", MAIN_SEPARATOR);
    let prefix: &str = prefix.as_ref();
    format!("{}{}", if path.starts_with(prefix) { "" } else { prefix }, path)
}

pub fn extract_file_name(full_path: &PathBuf) -> Option<&str> {
    full_path
        .file_name()
        .and_then(OsStr::to_str)
}

pub fn ls<P: AsRef<Path>>(path: &P) -> Result<BTreeSet<String>, IoError> {
    let dirs = try!(fs::read_dir(path));
    let dirs = dirs
        .filter_map(Result::ok)
        .map(|d| d.file_name())
        .map(OsString::into_string)
        .filter_map(Result::ok)
        .collect::<BTreeSet<String>>();
    Ok(dirs)
}

fn is_hidden(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, |s| s.starts_with(".") && s.len() > 1 && s != "..")
}

fn is_pinned(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, |s| s == WORKING_DIR_NAME || s == KEEPED_FILE_NAME)
}

fn is_target(entry: &WalkDirEntry) -> bool {
    !is_hidden(entry) && !is_pinned(entry)
}

fn to_string(entry: WalkDirEntry) -> Result<String, OsString> {
    entry
        .path()
        .as_os_str()
        .to_os_string()
        .into_string()
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> BTreeSet<String> {
    let walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(is_target)
        .filter_map(Result::ok)
        .filter(|e| !e.file_type().is_dir())
        .collect::<Vec<WalkDirEntry>>();

    walker
        .into_iter()
        .map(to_string)
        .filter_map(Result::ok)
        .collect::<BTreeSet<String>>()
}

pub fn dirs_ordered_by_descending_depth<P: AsRef<Path>>(root: P) -> Vec<String> {
    let mut walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(is_target)
        .filter_entry(|e| e.file_type().is_dir())
        .filter_map(Result::ok)
        .collect::<Vec<WalkDirEntry>>();

    walker.sort_by(|l, r| {
        let l_depth = &l.depth();
        let r_depth = &r.depth();
        r_depth.cmp(l_depth)
    });

    walker
        .into_iter()
        .map(to_string)
        .filter_map(Result::ok)
        .collect::<Vec<String>>()
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::mem;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::Duration;
use std::time::{UNIX_EPOCH, Instant, SystemTime};

use chrono;
use futures::{Async, Future, Poll, Stream};
use ifaces;
use timer;

use DataType;
use Row;
use backoff::Backoff;

pub fn duration_to_ms(duration: &Duration) -> u64 {
    duration.as_secs() * 1000 + u64::from(duration.subsec_nanos()) / 1_000_000
}

pub fn fmt_hex<T>(f: &mut fmt::Formatter, bytes: &[T]) -> fmt::Result where T: fmt::LowerHex {
    if bytes.is_empty() {
        return write!(f, "0x")
    }
    try!(write!(f, "{:#x}", bytes[0]));
    for b in &bytes[1..] {
        try!(write!(f, "{:x}", b));
    }
    Ok(())
}

pub fn time_to_us(time: &SystemTime) -> i64 {
    // TODO: do overflow checking
    match time.duration_since(UNIX_EPOCH) {
        Ok(duration) => {
            (duration.as_secs() * 1_000_000 + u64::from(duration.subsec_nanos()) / 1000) as i64
        },
        Err(error) => {
            let duration = error.duration();
            (- ((duration.as_secs() * 1_000_000 + u64::from(duration.subsec_nanos()) / 1000) as i64))
        }
    }
}

pub fn us_to_time(us: i64) -> SystemTime {
    let abs = us.abs() as u64;

    let s = abs / 1_000_000;
    let ns = (abs % 1_000_000) as u32 * 1000;

    if us.is_negative() {
        UNIX_EPOCH - Duration::new(s, ns)
    } else {
        UNIX_EPOCH + Duration::new(s, ns)
    }
}

pub fn fmt_timestamp(f: &mut fmt::Formatter, timestamp: SystemTime) -> fmt::Result {
    let datetime = if timestamp < UNIX_EPOCH {
        chrono::NaiveDateTime::from_timestamp(0, 0) -
            chrono::Duration::from_std(UNIX_EPOCH.duration_since(timestamp).unwrap()).unwrap()
    } else {
        chrono::NaiveDateTime::from_timestamp(0, 0) +
            chrono::Duration::from_std(timestamp.duration_since(UNIX_EPOCH).unwrap()).unwrap()
    };

    write!(f, "{}", datetime.format("%Y-%m-%dT%H:%M:%S%.6fZ"))
}

pub fn fmt_cell(f: &mut fmt::Formatter, row: &Row, idx: usize) -> fmt::Result {
    match row.schema().columns()[idx].data_type() {
        DataType::Bool => write!(f, "{}", row.get::<bool>(idx).unwrap()),
        DataType::Int8 => write!(f, "{}", row.get::<i8>(idx).unwrap()),
        DataType::Int16 => write!(f, "{}", row.get::<i16>(idx).unwrap()),
        DataType::Int32 => write!(f, "{}", row.get::<i32>(idx).unwrap()),
        DataType::Int64 => write!(f, "{}", row.get::<i64>(idx).unwrap()),
        DataType::Timestamp => fmt_timestamp(f, row.get::<SystemTime>(idx).unwrap()),
        DataType::Float => write!(f, "{}", row.get::<f32>(idx).unwrap()),
        DataType::Double => write!(f, "{}", row.get::<f64>(idx).unwrap()),
        DataType::Binary => fmt_hex(f, row.get::<&[u8]>(idx).unwrap()),
        DataType::String => write!(f, "{:?}", row.get::<&str>(idx).unwrap()),
    }
}

pub fn dummy_addr() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
}

lazy_static! {
    static ref LOCAL_ADDRS: HashSet<IpAddr> = {
        let mut addrs = HashSet::new();
        match ifaces::Interface::get_all() {
            Ok(ifaces) => {
                for iface in ifaces {
                    if let Some(addr) = iface.addr {
                        addrs.insert(addr.ip());
                    }
                }
            },
            Err(error) => warn!("failed to resolve local interface addresses: {}", error),
        }
        addrs
    };
}

/// Returns `true` if socket addr is for a local interface.
pub fn is_local_addr(addr: &IpAddr) -> bool {
    LOCAL_ADDRS.contains(addr) || addr.is_loopback()
}

pub fn cmp_socket_addrs(a: &SocketAddr, b: &SocketAddr) -> Ordering {
    match (a, b) {
        (&SocketAddr::V4(ref a), &SocketAddr::V4(ref b)) => (a.ip(), a.port()).cmp(&(b.ip(), b.port())),
        (&SocketAddr::V6(ref a), &SocketAddr::V6(ref b)) => (a.ip(), a.port()).cmp(&(b.ip(), b.port())),
        (&SocketAddr::V4(_), &SocketAddr::V6(_)) => Ordering::Less,
        (&SocketAddr::V6(_), &SocketAddr::V4(_)) => Ordering::Greater,
    }
}

/// Returns a stream which yields elements according to the backoff policy.
pub fn backoff_stream(mut backoff: Backoff, timer: timer::Timer) -> BackoffStream {
    let sleep = timer.sleep(backoff.next_backoff());
    BackoffStream {
        backoff: backoff,
        timer: timer,
        sleep: sleep,
    }
}
/// Stream which yields elements according to a backoff policy.
#[must_use = "streams do nothing unless polled"]
pub struct BackoffStream {
    backoff: Backoff,
    timer: timer::Timer,
    sleep: timer::Sleep
}
impl Stream for BackoffStream {
    type Item = ();
    type Error = ();
    fn poll(&mut self) -> Poll<Option<()>, ()> {
        try_ready!(self.sleep.poll());
        let backoff = self.backoff.next_backoff();
        self.sleep = self.timer.sleep(backoff);
        Ok(Async::Ready(Some(())))
    }
}

pub fn retry_with_backoff<R, F>(timer: timer::Timer,
                                mut backoff: Backoff,
                                mut retry: R)
                                -> RetryWithBackoff<R, F>
where R: FnMut(Instant, RetryCause<F::Error>) -> F,
      F: Future,
{
    let duration = backoff.next_backoff();
    let future = retry(Instant::now() + duration, RetryCause::Initial);
    let sleep = timer.sleep(duration);
    RetryWithBackoff {
        backoff: backoff,
        timer: timer,
        sleep: sleep,
        retry: retry,
        try: Try::Future(future),
    }
}

pub enum RetryCause<E> {
    Initial,
    TimedOut,
    Err(E),
}

enum Try<F> where F: Future {
    Future(F),
    Err(F::Error),
    None,
}
impl <F> Try<F> where F: Future {
    fn take(&mut self) -> Result<F, F::Error> {
        match mem::replace(self, Try::None) {
            Try::Future(f) => Ok(f),
            Try::Err(error) => Err(error),
            Try::None => unreachable!(),
        }
    }
}

#[must_use = "futures do nothing unless polled"]
pub struct RetryWithBackoff<R, F>
where R: FnMut(Instant, RetryCause<F::Error>) -> F,
      F: Future,
{
    backoff: Backoff,
    timer: timer::Timer,
    sleep: timer::Sleep,
    retry: R,
    try: Try<F>,
}

impl <R, F> Future for RetryWithBackoff<R, F>
where R: FnMut(Instant, RetryCause<F::Error>) -> F,
      F: Future,
{
    type Item = F::Item;
    type Error = ();

    fn poll(&mut self) -> Poll<F::Item, ()> {
        loop {
            {
                let poll = if let Try::Future(ref mut f) = self.try {
                    f.poll()
                } else {
                    Ok(Async::NotReady)
                };
                match poll {
                    Ok(Async::Ready(item)) => return Ok(Async::Ready(item)),
                    Ok(Async::NotReady) => (),
                    Err(error) => self.try = Try::Err(error),
                }
            }

            // Unwrap here is unfortunate, but we really have no way to handle
            // the timer being out of capacity.
            match self.sleep.poll().expect("timer sleep failed") {
                Async::Ready(_) => {
                    let duration = self.backoff.next_backoff();

                    let cause = match self.try.take() {
                        Ok(_) => RetryCause::TimedOut,
                        Err(error) => RetryCause::Err(error),
                    };

                    self.try = Try::Future((self.retry)(Instant::now() + duration, cause));
                    self.sleep = self.timer.sleep(duration);
                },
                Async::NotReady => return Ok(Async::NotReady),
            }
        }
    }
}

#[cfg(test)]
mod tests {

    use std::time::{Duration, UNIX_EPOCH};
    use std::net::ToSocketAddrs;

    use env_logger;
    use quickcheck::{quickcheck, TestResult};

    use schema;
    use super::*;

    #[test]
    fn timestamp_conversion() {
        let _ = env_logger::init();

        fn roundtrip(us: i64) -> TestResult {
            TestResult::from_bool(us == time_to_us(&us_to_time(us)))
        }

        quickcheck(roundtrip as fn(i64) -> TestResult);
    }

    #[test]
    fn test_format_timestamp() {
        let _ = env_logger::init();
        let schema = schema::tests::all_types_schema();
        let mut row = schema.new_row();

        row.set_by_name("timestamp", UNIX_EPOCH - Duration::from_millis(1234)).unwrap();
        assert_eq!("Timestamp \"timestamp\"=1969-12-31T23:59:58.766000Z",
                   &format!("{:?}", row));

        row.set_by_name("timestamp", UNIX_EPOCH + Duration::from_millis(1234)).unwrap();
        assert_eq!("Timestamp \"timestamp\"=1970-01-01T00:00:01.234000Z",
                   &format!("{:?}", row));
    }

    #[test]
    fn test_is_local_addr() {
        let _ = env_logger::init();
        let addr = "127.0.1.1:0".to_socket_addrs().unwrap().next().unwrap().ip();
        assert!(is_local_addr(&addr));
        let addr = "127.0.0.1:0".to_socket_addrs().unwrap().next().unwrap().ip();
        assert!(is_local_addr(&addr));
    }
}
<commit_msg>remove unused util combinators<commit_after>use std::cmp::Ordering;
use std::collections::HashSet;
use std::fmt;
use std::mem;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::Duration;
use std::time::{UNIX_EPOCH, Instant, SystemTime};

use chrono;
use futures::{Async, Future, Poll, Stream};
use ifaces;
use timer;

use DataType;
use Row;
use backoff::Backoff;

pub fn duration_to_ms(duration: &Duration) -> u64 {
    duration.as_secs() * 1000 + u64::from(duration.subsec_nanos()) / 1_000_000
}

pub fn fmt_hex<T>(f: &mut fmt::Formatter, bytes: &[T]) -> fmt::Result where T: fmt::LowerHex {
    if bytes.is_empty() {
        return write!(f, "0x")
    }
    try!(write!(f, "{:#x}", bytes[0]));
    for b in &bytes[1..] {
        try!(write!(f, "{:x}", b));
    }
    Ok(())
}

pub fn time_to_us(time: &SystemTime) -> i64 {
    // TODO: do overflow checking
    match time.duration_since(UNIX_EPOCH) {
        Ok(duration) => {
            (duration.as_secs() * 1_000_000 + u64::from(duration.subsec_nanos()) / 1000) as i64
        },
        Err(error) => {
            let duration = error.duration();
            (- ((duration.as_secs() * 1_000_000 + u64::from(duration.subsec_nanos()) / 1000) as i64))
        }
    }
}

pub fn us_to_time(us: i64) -> SystemTime {
    let abs = us.abs() as u64;

    let s = abs / 1_000_000;
    let ns = (abs % 1_000_000) as u32 * 1000;

    if us.is_negative() {
        UNIX_EPOCH - Duration::new(s, ns)
    } else {
        UNIX_EPOCH + Duration::new(s, ns)
    }
}

pub fn fmt_timestamp(f: &mut fmt::Formatter, timestamp: SystemTime) -> fmt::Result {
    let datetime = if timestamp < UNIX_EPOCH {
        chrono::NaiveDateTime::from_timestamp(0, 0) -
            chrono::Duration::from_std(UNIX_EPOCH.duration_since(timestamp).unwrap()).unwrap()
    } else {
        chrono::NaiveDateTime::from_timestamp(0, 0) +
            chrono::Duration::from_std(timestamp.duration_since(UNIX_EPOCH).unwrap()).unwrap()
    };

    write!(f, "{}", datetime.format("%Y-%m-%dT%H:%M:%S%.6fZ"))
}

pub fn fmt_cell(f: &mut fmt::Formatter, row: &Row, idx: usize) -> fmt::Result {
    match row.schema().columns()[idx].data_type() {
        DataType::Bool => write!(f, "{}", row.get::<bool>(idx).unwrap()),
        DataType::Int8 => write!(f, "{}", row.get::<i8>(idx).unwrap()),
        DataType::Int16 => write!(f, "{}", row.get::<i16>(idx).unwrap()),
        DataType::Int32 => write!(f, "{}", row.get::<i32>(idx).unwrap()),
        DataType::Int64 => write!(f, "{}", row.get::<i64>(idx).unwrap()),
        DataType::Timestamp => fmt_timestamp(f, row.get::<SystemTime>(idx).unwrap()),
        DataType::Float => write!(f, "{}", row.get::<f32>(idx).unwrap()),
        DataType::Double => write!(f, "{}", row.get::<f64>(idx).unwrap()),
        DataType::Binary => fmt_hex(f, row.get::<&[u8]>(idx).unwrap()),
        DataType::String => write!(f, "{:?}", row.get::<&str>(idx).unwrap()),
    }
}

pub fn dummy_addr() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
}

lazy_static! {
    static ref LOCAL_ADDRS: HashSet<IpAddr> = {
        let mut addrs = HashSet::new();
        match ifaces::Interface::get_all() {
            Ok(ifaces) => {
                for iface in ifaces {
                    if let Some(addr) = iface.addr {
                        addrs.insert(addr.ip());
                    }
                }
            },
            Err(error) => warn!("failed to resolve local interface addresses: {}", error),
        }
        addrs
    };
}

/// Returns `true` if socket addr is for a local interface.
pub fn is_local_addr(addr: &IpAddr) -> bool {
    LOCAL_ADDRS.contains(addr) || addr.is_loopback()
}

pub fn cmp_socket_addrs(a: &SocketAddr, b: &SocketAddr) -> Ordering {
    match (a, b) {
        (&SocketAddr::V4(ref a), &SocketAddr::V4(ref b)) => (a.ip(), a.port()).cmp(&(b.ip(), b.port())),
        (&SocketAddr::V6(ref a), &SocketAddr::V6(ref b)) => (a.ip(), a.port()).cmp(&(b.ip(), b.port())),
        (&SocketAddr::V4(_), &SocketAddr::V6(_)) => Ordering::Less,
        (&SocketAddr::V6(_), &SocketAddr::V4(_)) => Ordering::Greater,
    }
}

#[cfg(test)]
mod tests {

    use std::time::{Duration, UNIX_EPOCH};
    use std::net::ToSocketAddrs;

    use env_logger;
    use quickcheck::{quickcheck, TestResult};

    use schema;
    use super::*;

    #[test]
    fn timestamp_conversion() {
        let _ = env_logger::init();

        fn roundtrip(us: i64) -> TestResult {
            TestResult::from_bool(us == time_to_us(&us_to_time(us)))
        }

        quickcheck(roundtrip as fn(i64) -> TestResult);
    }

    #[test]
    fn test_format_timestamp() {
        let _ = env_logger::init();
        let schema = schema::tests::all_types_schema();
        let mut row = schema.new_row();

        row.set_by_name("timestamp", UNIX_EPOCH - Duration::from_millis(1234)).unwrap();
        assert_eq!("Timestamp \"timestamp\"=1969-12-31T23:59:58.766000Z",
                   &format!("{:?}", row));

        row.set_by_name("timestamp", UNIX_EPOCH + Duration::from_millis(1234)).unwrap();
        assert_eq!("Timestamp \"timestamp\"=1970-01-01T00:00:01.234000Z",
                   &format!("{:?}", row));
    }

    #[test]
    fn test_is_local_addr() {
        let _ = env_logger::init();
        let addr = "127.0.1.1:0".to_socket_addrs().unwrap().next().unwrap().ip();
        assert!(is_local_addr(&addr));
        let addr = "127.0.0.1:0".to_socket_addrs().unwrap().next().unwrap().ip();
        assert!(is_local_addr(&addr));
    }
}
<|endoftext|>