<commit_before>use std::collections::HashMap;
use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::iter::Filter;
use std::fs;
use std::path::Path;

use bio::io::gff::{self, GffType};
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use {Coord, Exon, ExonFeatureKind as EFK, Gene, GBuilder, Strand, Transcript, Error};

// Various commonly-used feature column values
const GENE_STR: &'static str = "gene";
const TRANSCRIPT_STR: &'static str = "transcript";
const EXON_STR: &'static str = "exon";
const UTR_STR: &'static str = "UTR";
const UTR5_STR: &'static str = "UTR5";
const UTR3_STR: &'static str = "UTR3";
const CDS_STR: &'static str = "CDS";
const START_CODON_STR: &'static str = "start_codon";
const STOP_CODON_STR: &'static str = "stop_codon";

// Value for unknown columns.
const UNK_STR: &'static str = ".";

// Commonly-used attribute keys.
const GENE_ID_STR: &'static str = "gene_id";
const TRANSCRIPT_ID_STR: &'static str = "transcript_id";

// Initial coordinate for features.
const INIT_COORD: (u64, u64) = (::std::u64::MAX, ::std::u64::MIN);

pub struct Reader<R: io::Read> {
    inner: gff::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R, gff_type: GffType) -> Reader<R> {
        Reader {
            inner: gff::Reader::new(in_reader, gff_type),
        }
    }

    pub fn genes(&mut self) -> GffGenes<R> {
        GffGenes {
            inner: self.records()
                .filter(GffGenes::<R>::gene_filter_func as GxFilterFunc)
                .group_by(GffGenes::<R>::gene_group_func),
        }
    }

    fn records(&mut self) -> GffRecords<R> {
        GffRecords {
            inner: self.inner.records()
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P, gff_type: GffType) -> io::Result<Self> {
        fs::File::open(path).map(|file| Reader::from_reader(file, gff_type))
    }
}

struct GffRecords<'a, R: 'a> where R: io::Read {
    inner: gff::Records<'a, R>,
}

impl<'a, R> Iterator for GffRecords<'a, R> where R: io::Read {

    type Item = Result<gff::Record, Error>;

    fn next(&mut self) -> Option<Result<gff::Record, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
            })
    }
}

pub struct GffGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GxGroupKey, GxRecords<'a, R>, GxGroupFunc>,
}

type GxFilterFunc = fn(&Result<gff::Record, Error>) -> bool;

type GxRecords<'a, R> = Filter<GffRecords<'a, R>, GxFilterFunc>;

type GxGroupKey = Option<(String, String, Strand)>;

type GxGroupFunc = fn(&Result<gff::Record, Error>) -> GxGroupKey;

type GxGroupedRecs<'a, 'b, R> = Group<'b, GxGroupKey, GxRecords<'a, R>, GxGroupFunc>;

type RawTrxCoord = (Coord<u64>, Vec<Coord<u64>>, Option<Coord<u64>>);

impl<'a, R> GffGenes<'a, R> where R: io::Read {

    fn gene_filter_func(result: &Result<gff::Record, Error>) -> bool {
        match result {
            &Ok(ref rec) => rec.attributes().contains_key(GENE_ID_STR),
            &Err(_) => true,  // Err(_) is supposed to be handled elsewhere
        }
    }

    fn gene_group_func(result: &Result<gff::Record, Error>) -> GxGroupKey {
        result.as_ref().ok()
            .map(|res| {
                let gene_id = res.attributes().get(GENE_ID_STR)
                    .expect("'gene_id' attribute not found")
                    .clone();
                let seq_name = String::from(res.seqname());
                let strand = res.strand().unwrap_or(Strand::Unknown);

                (gene_id, seq_name, strand)
            })
    }

    fn group_to_gene<'b>(group: (GxGroupKey, GxGroupedRecs<'a, 'b, R>)) -> Result<Gene, Error> {
        let (group_key, records) = group;
        match group_key {

            None => Err(records.filter_map(|x| x.err()).next().unwrap()),

            Some((gid, seq_name, strand)) => {

                let mut gene_coord = INIT_COORD;
                let mut trx_coords: LinkedHashMap<String, RawTrxCoord> = LinkedHashMap::new();

                for record in records {

                    let mut rec = record?;
                    rec.attributes_mut().remove(GENE_ID_STR);

                    let rtrx_entry = rec.attributes_mut()
                        .remove(TRANSCRIPT_ID_STR)
                        .ok_or(Error::Gff("required 'transcript_id' attribute not found"))
                        .map(|id| trx_coords.entry(id).or_insert((INIT_COORD, vec![], None)));

                    match rec.feature_type() {
                        GENE_STR => {
                            gene_coord = adjust_coord(gene_coord, &rec);
                        },
                        TRANSCRIPT_STR => {
                            let trx_entry = rtrx_entry?;
                            trx_entry.0 = adjust_coord(trx_entry.0, &rec);
                        },
                        EXON_STR => {
                            let trx_entry = rtrx_entry?;
                            (trx_entry.1).push((*rec.start() - 1, *rec.end()));
                        },
                        CDS_STR => {
                            let trx_entry = rtrx_entry?;
                            trx_entry.2 = (trx_entry.2).or(Some(INIT_COORD))
                                .map(|coord| adjust_coord(coord, &rec));
                        },
                        // TODO: How to best handle nested features? They may be out of order ...
                        _ => {},
                    }
                }

                GBuilder::new(seq_name, gene_coord.0, gene_coord.1)
                    .id(gid)
                    .strand(strand)
                    .transcript_coords(trx_coords)
                    .build()
            },
        }
    }
}

impl<'a, R> Iterator for GffGenes<'a, R> where R: io::Read {

    type Item = Result<Gene, Error>;

    fn next(&mut self) -> Option<Result<Gene, Error>> {
        self.inner.into_iter().map(Self::group_to_gene).next()
    }
}

struct GffRow(String, String, String, u64, u64, String, String, String, HashMap<String, String>);

impl From<GffRow> for gff::Record {
    fn from(row: GffRow) -> Self {
        let mut rec = gff::Record::default();
        *rec.seqname_mut() = row.0;
        *rec.source_mut() = row.1;
        *rec.feature_type_mut() = row.2;
        *rec.start_mut() = row.3 + 1;
        *rec.end_mut() = row.4;
        *rec.score_mut() = row.5;
        *rec.strand_mut() = row.6;
        *rec.frame_mut() = row.7;
        *rec.attributes_mut() = row.8;
        rec
    }
}

impl Gene {

    #[inline(always)]
    fn num_records(&self) -> usize {
        1 + self.transcripts().values()
            .map(|ref trx| trx.num_records())
            .fold(0, |acc, x| acc + x)
    }

    // TODO: also handle gene-level features
    fn to_gff_records(mut self) -> Result<Vec<gff::Record>, Error> {

        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());
        let gene_id = match self.id.take() {
            Some(gid) => {
                self.attributes.insert(GENE_ID_STR.to_owned(), gid.clone());
                gid
            },
            None =>  {
                self.attributes.get(GENE_ID_STR)
                    .ok_or(Error::Gff("required 'gene_id' attribute not found"))
                    .map(|gid| gid.clone())?
            },
        };

        let mut recs = Vec::with_capacity(self.num_records());

        let row = GffRow(
            self.seq_name().to_owned(), source, GENE_STR.to_owned(),
            self.interval().start, self.interval().end, score, strand.to_owned(),
            UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(row));

        for (_, mut transcript) in self.into_transcripts() {
            let mut trx_recs = transcript.to_gff_records(gene_id.as_str())?;
            recs.append(&mut trx_recs);
        }

        Ok(recs)
    }
}

impl Transcript {

    #[inline(always)]
    fn num_records(&self) -> usize {
        1 + self.exons().iter()
            .map(|ref exn| 1 + exn.features.len())
            .fold(0, |acc, x| acc + x)
    }

    // TODO: also handle transcript-level features
    fn to_gff_records(mut self, gene_id: &str) -> Result<Vec<gff::Record>, Error> {

        self.attributes.insert(GENE_ID_STR.to_owned(), gene_id.to_owned());
        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());
        let transcript_id = match self.id.take() {
            Some(tid) => {
                self.attributes.insert(TRANSCRIPT_ID_STR.to_owned(), tid.clone());
                tid
            },
            None =>  {
                self.attributes.get(TRANSCRIPT_ID_STR)
                    .ok_or(Error::Gff("required 'transcript_id' attribute not found"))
                    .map(|tid| tid.clone())?
            },
        };

        let mut recs = Vec::with_capacity(self.num_records());

        let trx_row =
            GffRow(self.seq_name().to_owned(), source, TRANSCRIPT_STR.to_owned(),
                   self.interval().start, self.interval().end, score,
                   strand.to_owned(), UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(trx_row));

        for exon in self.into_exons() {
            let mut exon_recs = exon.to_gff_records(gene_id, transcript_id.as_str())?;
            recs.append(&mut exon_recs);
        }

        Ok(recs)
    }
}

impl EFK {

    #[inline(always)]
    fn get_feature_frame(&self) -> (String, String) {
        let (feature, frame) = match self {
            &EFK::UTR => (UTR_STR, UNK_STR),
            &EFK::UTR5 => (UTR5_STR, UNK_STR),
            &EFK::UTR3 => (UTR3_STR, UNK_STR),
            &EFK::CDS { frame: ref f } => (CDS_STR, frame_to_str(f)),
            &EFK::StopCodon { frame: ref f } => (STOP_CODON_STR, frame_to_str(f)),
            &EFK::StartCodon { frame: ref f } => (START_CODON_STR, frame_to_str(f)),
            &EFK::Any(ref s) => (s.as_str(), UNK_STR),
        };
        (feature.to_owned(), frame.to_owned())
    }
}

impl Exon {

    fn to_gff_records(mut self, gene_id: &str, transcript_id: &str
    ) -> Result<Vec<gff::Record>, Error> {

        self.attributes.insert(GENE_ID_STR.to_owned(), gene_id.to_owned());
        self.attributes.insert(TRANSCRIPT_ID_STR.to_owned(), transcript_id.to_owned());
        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());

        let mut recs = Vec::with_capacity(1 + self.features.len());

        let exon_row =
            GffRow(self.seq_name().to_owned(), source.clone(), EXON_STR.to_owned(),
                   self.interval().start, self.interval().end, score, strand.to_owned(),
                   UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(exon_row));

        for fx in self.features.iter() {
            let (feature, frame) = fx.kind.get_feature_frame();
            let fx_row =
                GffRow(self.seq_name().to_owned(), source.clone(), feature,
                       self.interval().start, self.interval().end, UNK_STR.to_owned(),
                       strand.to_owned(), UNK_STR.to_owned(), self.attributes.clone());
            recs.push(gff::Record::from(fx_row));
        }

        Ok(recs)
    }
}

pub struct Writer<W: io::Write> {
    inner: gff::Writer<W>,
}

impl<W: io::Write> Writer<W> {

    pub fn from_writer(in_writer: W, gff_type: GffType) -> Writer<W> {
        Writer {
            inner: gff::Writer::new(in_writer, gff_type),
        }
    }

    pub fn write_gene(&mut self, mut gene: Gene) -> Result<(), Error> {
        for record in gene.to_gff_records()?.iter() {
            self.inner.write(&record)?;
        }
        Ok(())
    }
}

#[inline(always)]
fn adjust_coord(cur_coord: Coord<u64>, record: &gff::Record) -> Coord<u64> {
    (min(cur_coord.0, *record.start() - 1),
     max(cur_coord.1, *record.end()))
}

#[inline(always)]
fn extract_source_score(attributes: &mut HashMap<String, String>) -> (String, String) {
    let source = attributes.remove("source").unwrap_or(UNK_STR.to_owned());
    let score = attributes.remove("score").unwrap_or(UNK_STR.to_owned());
    (source, score)
}

#[inline(always)]
fn strand_to_string(strand: &Strand) -> String {
    let value = match strand {
        &Strand::Forward => "+",
        &Strand::Reverse => "-",
        &Strand::Unknown => UNK_STR,
    };
    value.to_owned()
}

#[inline(always)]
fn frame_to_str(frame: &Option<u8>) -> &str {
    match frame {
        &Some(0) => "0",
        &Some(1) => "1",
        &Some(2) => "2",
        _ => UNK_STR,
    }
}
<commit_msg>Update gene group key order<commit_after>use std::collections::HashMap;
use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::iter::Filter;
use std::fs;
use std::path::Path;

use bio::io::gff::{self, GffType};
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use {Coord, Exon, ExonFeatureKind as EFK, Gene, GBuilder, Strand, Transcript, Error};

// Various commonly-used feature column values
const GENE_STR: &'static str = "gene";
const TRANSCRIPT_STR: &'static str = "transcript";
const EXON_STR: &'static str = "exon";
const UTR_STR: &'static str = "UTR";
const UTR5_STR: &'static str = "UTR5";
const UTR3_STR: &'static str = "UTR3";
const CDS_STR: &'static str = "CDS";
const START_CODON_STR: &'static str = "start_codon";
const STOP_CODON_STR: &'static str = "stop_codon";

// Value for unknown columns.
const UNK_STR: &'static str = ".";

// Commonly-used attribute keys.
const GENE_ID_STR: &'static str = "gene_id";
const TRANSCRIPT_ID_STR: &'static str = "transcript_id";

// Initial coordinate for features.
const INIT_COORD: (u64, u64) = (::std::u64::MAX, ::std::u64::MIN);

pub struct Reader<R: io::Read> {
    inner: gff::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R, gff_type: GffType) -> Reader<R> {
        Reader {
            inner: gff::Reader::new(in_reader, gff_type),
        }
    }

    pub fn genes(&mut self) -> GffGenes<R> {
        GffGenes {
            inner: self.records()
                .filter(GffGenes::<R>::gene_filter_func as GxFilterFunc)
                .group_by(GffGenes::<R>::gene_group_func),
        }
    }

    fn records(&mut self) -> GffRecords<R> {
        GffRecords {
            inner: self.inner.records()
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P, gff_type: GffType) -> io::Result<Self> {
        fs::File::open(path).map(|file| Reader::from_reader(file, gff_type))
    }
}

struct GffRecords<'a, R: 'a> where R: io::Read {
    inner: gff::Records<'a, R>,
}

impl<'a, R> Iterator for GffRecords<'a, R> where R: io::Read {

    type Item = Result<gff::Record, Error>;

    fn next(&mut self) -> Option<Result<gff::Record, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
            })
    }
}

pub struct GffGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GxGroupKey, GxRecords<'a, R>, GxGroupFunc>,
}

type GxFilterFunc = fn(&Result<gff::Record, Error>) -> bool;

type GxRecords<'a, R> = Filter<GffRecords<'a, R>, GxFilterFunc>;

type GxGroupKey = Option<(String, String, Strand)>;

type GxGroupFunc = fn(&Result<gff::Record, Error>) -> GxGroupKey;

type GxGroupedRecs<'a, 'b, R> = Group<'b, GxGroupKey, GxRecords<'a, R>, GxGroupFunc>;

type RawTrxCoord = (Coord<u64>, Vec<Coord<u64>>, Option<Coord<u64>>);

impl<'a, R> GffGenes<'a, R> where R: io::Read {

    fn gene_filter_func(result: &Result<gff::Record, Error>) -> bool {
        match result {
            &Ok(ref rec) => rec.attributes().contains_key(GENE_ID_STR),
            &Err(_) => true,  // Err(_) is supposed to be handled elsewhere
        }
    }

    fn gene_group_func(result: &Result<gff::Record, Error>) -> GxGroupKey {
        result.as_ref().ok()
            .map(|res| {
                let gene_id = res.attributes().get(GENE_ID_STR)
                    .expect("'gene_id' attribute not found")
                    .clone();
                let seq_name = res.seqname().to_owned();
                let strand = res.strand().unwrap_or(Strand::Unknown);

                (seq_name, gene_id, strand)
            })
    }

    fn group_to_gene<'b>(group: (GxGroupKey, GxGroupedRecs<'a, 'b, R>)) -> Result<Gene, Error> {
        let (group_key, records) = group;
        match group_key {

            None => Err(records.filter_map(|x| x.err()).next().unwrap()),

            Some((seq_name, gid, strand)) => {

                let mut gene_coord = INIT_COORD;
                let mut trx_coords: LinkedHashMap<String, RawTrxCoord> = LinkedHashMap::new();

                for record in records {

                    let mut rec = record?;
                    rec.attributes_mut().remove(GENE_ID_STR);

                    let rtrx_entry = rec.attributes_mut()
                        .remove(TRANSCRIPT_ID_STR)
                        .ok_or(Error::Gff("required 'transcript_id' attribute not found"))
                        .map(|id| trx_coords.entry(id).or_insert((INIT_COORD, vec![], None)));

                    match rec.feature_type() {
                        GENE_STR => {
                            gene_coord = adjust_coord(gene_coord, &rec);
                        },
                        TRANSCRIPT_STR => {
                            let trx_entry = rtrx_entry?;
                            trx_entry.0 = adjust_coord(trx_entry.0, &rec);
                        },
                        EXON_STR => {
                            let trx_entry = rtrx_entry?;
                            (trx_entry.1).push((*rec.start() - 1, *rec.end()));
                        },
                        CDS_STR => {
                            let trx_entry = rtrx_entry?;
                            trx_entry.2 = (trx_entry.2).or(Some(INIT_COORD))
                                .map(|coord| adjust_coord(coord, &rec));
                        },
                        // TODO: How to best handle nested features? They may be out of order ...
                        _ => {},
                    }
                }

                GBuilder::new(seq_name, gene_coord.0, gene_coord.1)
                    .id(gid)
                    .strand(strand)
                    .transcript_coords(trx_coords)
                    .build()
            },
        }
    }
}

impl<'a, R> Iterator for GffGenes<'a, R> where R: io::Read {

    type Item = Result<Gene, Error>;

    fn next(&mut self) -> Option<Result<Gene, Error>> {
        self.inner.into_iter().map(Self::group_to_gene).next()
    }
}

struct GffRow(String, String, String, u64, u64, String, String, String, HashMap<String, String>);

impl From<GffRow> for gff::Record {
    fn from(row: GffRow) -> Self {
        let mut rec = gff::Record::default();
        *rec.seqname_mut() = row.0;
        *rec.source_mut() = row.1;
        *rec.feature_type_mut() = row.2;
        *rec.start_mut() = row.3 + 1;
        *rec.end_mut() = row.4;
        *rec.score_mut() = row.5;
        *rec.strand_mut() = row.6;
        *rec.frame_mut() = row.7;
        *rec.attributes_mut() = row.8;
        rec
    }
}

impl Gene {

    #[inline(always)]
    fn num_records(&self) -> usize {
        1 + self.transcripts().values()
            .map(|ref trx| trx.num_records())
            .fold(0, |acc, x| acc + x)
    }

    // TODO: also handle gene-level features
    fn to_gff_records(mut self) -> Result<Vec<gff::Record>, Error> {

        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());
        let gene_id = match self.id.take() {
            Some(gid) => {
                self.attributes.insert(GENE_ID_STR.to_owned(), gid.clone());
                gid
            },
            None =>  {
                self.attributes.get(GENE_ID_STR)
                    .ok_or(Error::Gff("required 'gene_id' attribute not found"))
                    .map(|gid| gid.clone())?
            },
        };

        let mut recs = Vec::with_capacity(self.num_records());

        let row = GffRow(
            self.seq_name().to_owned(), source, GENE_STR.to_owned(),
            self.interval().start, self.interval().end, score, strand.to_owned(),
            UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(row));

        for (_, mut transcript) in self.into_transcripts() {
            let mut trx_recs = transcript.to_gff_records(gene_id.as_str())?;
            recs.append(&mut trx_recs);
        }

        Ok(recs)
    }
}

impl Transcript {

    #[inline(always)]
    fn num_records(&self) -> usize {
        1 + self.exons().iter()
            .map(|ref exn| 1 + exn.features.len())
            .fold(0, |acc, x| acc + x)
    }

    // TODO: also handle transcript-level features
    fn to_gff_records(mut self, gene_id: &str) -> Result<Vec<gff::Record>, Error> {

        self.attributes.insert(GENE_ID_STR.to_owned(), gene_id.to_owned());
        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());
        let transcript_id = match self.id.take() {
            Some(tid) => {
                self.attributes.insert(TRANSCRIPT_ID_STR.to_owned(), tid.clone());
                tid
            },
            None =>  {
                self.attributes.get(TRANSCRIPT_ID_STR)
                    .ok_or(Error::Gff("required 'transcript_id' attribute not found"))
                    .map(|tid| tid.clone())?
            },
        };

        let mut recs = Vec::with_capacity(self.num_records());

        let trx_row =
            GffRow(self.seq_name().to_owned(), source, TRANSCRIPT_STR.to_owned(),
                   self.interval().start, self.interval().end, score,
                   strand.to_owned(), UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(trx_row));

        for exon in self.into_exons() {
            let mut exon_recs = exon.to_gff_records(gene_id, transcript_id.as_str())?;
            recs.append(&mut exon_recs);
        }

        Ok(recs)
    }
}

impl EFK {

    #[inline(always)]
    fn get_feature_frame(&self) -> (String, String) {
        let (feature, frame) = match self {
            &EFK::UTR => (UTR_STR, UNK_STR),
            &EFK::UTR5 => (UTR5_STR, UNK_STR),
            &EFK::UTR3 => (UTR3_STR, UNK_STR),
            &EFK::CDS { frame: ref f } => (CDS_STR, frame_to_str(f)),
            &EFK::StopCodon { frame: ref f } => (STOP_CODON_STR, frame_to_str(f)),
            &EFK::StartCodon { frame: ref f } => (START_CODON_STR, frame_to_str(f)),
            &EFK::Any(ref s) => (s.as_str(), UNK_STR),
        };
        (feature.to_owned(), frame.to_owned())
    }
}

impl Exon {

    fn to_gff_records(mut self, gene_id: &str, transcript_id: &str
    ) -> Result<Vec<gff::Record>, Error> {

        self.attributes.insert(GENE_ID_STR.to_owned(), gene_id.to_owned());
        self.attributes.insert(TRANSCRIPT_ID_STR.to_owned(), transcript_id.to_owned());
        let (source, score) = extract_source_score(&mut self.attributes);
        let strand = strand_to_string(self.strand());

        let mut recs = Vec::with_capacity(1 + self.features.len());

        let exon_row =
            GffRow(self.seq_name().to_owned(), source.clone(), EXON_STR.to_owned(),
                   self.interval().start, self.interval().end, score, strand.to_owned(),
                   UNK_STR.to_owned(), self.attributes.clone());
        recs.push(gff::Record::from(exon_row));

        for fx in self.features.iter() {
            let (feature, frame) = fx.kind.get_feature_frame();
            let fx_row =
                GffRow(self.seq_name().to_owned(), source.clone(), feature,
                       self.interval().start, self.interval().end, UNK_STR.to_owned(),
                       strand.to_owned(), UNK_STR.to_owned(), self.attributes.clone());
            recs.push(gff::Record::from(fx_row));
        }

        Ok(recs)
    }
}

pub struct Writer<W: io::Write> {
    inner: gff::Writer<W>,
}

impl<W: io::Write> Writer<W> {

    pub fn from_writer(in_writer: W, gff_type: GffType) -> Writer<W> {
        Writer {
            inner: gff::Writer::new(in_writer, gff_type),
        }
    }

    pub fn write_gene(&mut self, mut gene: Gene) -> Result<(), Error> {
        for record in gene.to_gff_records()?.iter() {
            self.inner.write(&record)?;
        }
        Ok(())
    }
}

#[inline(always)]
fn adjust_coord(cur_coord: Coord<u64>, record: &gff::Record) -> Coord<u64> {
    (min(cur_coord.0, *record.start() - 1),
     max(cur_coord.1, *record.end()))
}

#[inline(always)]
fn extract_source_score(attributes: &mut HashMap<String, String>) -> (String, String) {
    let source = attributes.remove("source").unwrap_or(UNK_STR.to_owned());
    let score = attributes.remove("score").unwrap_or(UNK_STR.to_owned());
    (source, score)
}

#[inline(always)]
fn strand_to_string(strand: &Strand) -> String {
    let value = match strand {
        &Strand::Forward => "+",
        &Strand::Reverse => "-",
        &Strand::Unknown => UNK_STR,
    };
    value.to_owned()
}

#[inline(always)]
fn frame_to_str(frame: &Option<u8>) -> &str {
    match frame {
        &Some(0) => "0",
        &Some(1) => "1",
        &Some(2) => "2",
        _ => UNK_STR,
    }
}
<|endoftext|>


<commit_before>use common::consts::get_version_string;
use platform::process::get_pid;
use platform::time::time_now;
use storage::Cache;
use storage::CacheError;
use storage::Key;
use storage::Value;
use tcp_transport::stats::TransportStats;

use super::cmd::Cmd;
use super::cmd::Delete;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Touch;
use super::cmd::Value as CmdValue;

// For use to get an early exit from a function. The first parameter is a bool
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Option<Resp>. In case
// of None no return is made ie. execution continues. In case of Some(resp) we
// perform return resp;
//
// maybe_reply_stmt!(!noreply, match rv {
//     Ok(_) => None,   // do not return
//     Err(_) => Some(Resp::Error),
// });
//
// ->
//
// let opt = expr_provided;
// if !opt.is_none() {
//     if !noreply {
//         return opt.unwrap();
//     } else {
//         return Resp::Empty;
//     }
// }
macro_rules! maybe_reply_stmt {
    ( $cond:expr, $opt:expr ) => {
        {
            if !$opt.is_none() {
                if $cond {
                    return $opt.unwrap();
                } else {
                    return Resp::Empty;
                }
            };
        }
    };
}

// For use to conditionally produce a return value. The first parameter is used
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Resp.
//
// maybe_reply_expr!(!noreply, match rv {
//     Ok(_) => Resp::Stored,
//     Err(_) => Resp::Error,
// });
//
// ->
//
// let resp = expr_provided;
// if !noreply {
//     resp
// } else {
//     Resp::Empty
// }
macro_rules! maybe_reply_expr {
    ( $cond:expr, $resp:expr ) => {
        {
            if $cond {
                $resp
            } else {
                Resp::Empty
            }
        }
    };
}

struct DriverStats {
    cmd_get: u64,
    cmd_set: u64,
    cmd_flush: u64,
    cmd_touch: u64,
}

impl DriverStats {
    pub fn new() -> DriverStats {
        DriverStats {
            cmd_get: 0,
            cmd_set: 0,
            cmd_flush: 0,
            cmd_touch: 0,
        }
    }
}

pub struct Driver {
    cache: Cache,
    time_start: f64,

    stats: DriverStats,
    transport_stats: TransportStats, // this is a global snapshot
}

impl Driver {
    pub fn new(cache: Cache) -> Driver {
        Driver {
            cache: cache,
            stats: DriverStats::new(),
            time_start: time_now(),
            transport_stats: TransportStats::new(),
        }
    }

    fn set_exptime(&self, value: &mut Value, exptime: u32) {
        // If exptime is greater than zero we need to set it on the value
        if exptime > 0 {
            let tm;

            // Is it an interval greater than 30 days? Then it's a timestamp
            if exptime > 60 * 60 * 24 * 30 {
                tm = exptime as f64;

            } else {
                // Otherwise it's relative from now
                tm = time_now() + exptime as f64;
            }

            value.set_exptime(tm);
        }
    }

    fn do_add(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If so it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => {
                Some(Resp::NotStored)
            },
            Ok(false) => None,
            Err(_) => {
                Some(Resp::Error)
            }
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_append(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
            _ => None,
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Append the data we just received to the blob that is there
        value.item.extend(set.data);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_delete(&mut self, delete: Delete) -> Resp {
        let key = Key::new(delete.key.clone().into_bytes());

        let rv = self.cache.remove(&key);

        maybe_reply_expr!(!delete.noreply, match rv {
            Ok(_) => Resp::Deleted,
            Err(_) => Resp::NotFound,
        })
    }

    fn do_get(&mut self, get: Get) -> Resp {
        // Update stats
        self.stats.cmd_get += 1;

        let mut values = vec![];

        for key in get.keys {
            let key_str = key.clone();

            let key_st = Key::new(key.into_bytes());
            let rv = self.cache.get(&key_st);

            match rv {
                Ok(value) => {
                    let val_st = CmdValue {
                        key: key_str,
                        flags: value.flags,
                        data: value.item.clone(),
                    };
                    values.push(val_st);
                }
                // Keys that were not found are skipped, no error given
                Err(_) => (),
            }
        }

        Resp::Values(values)
    }

    fn do_prepend(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(_) => None,
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Prepend the data we just received to the blob that is there
        let mut new_item = Vec::with_capacity(set.data.len() + value.item.len());
        new_item.extend(set.data);
        new_item.extend(value.item);
        value.item = new_item;

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_replace(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If not it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_set(&mut self, set: Set) -> Resp {
        // Update stats
        self.stats.cmd_set += 1;

        let key = Key::new(set.key.into_bytes());
        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_stats(&self) -> Resp {
        let storage = self.cache.get_stats();

        let pid = get_pid().to_string();
        let uptime = ((time_now() - self.time_start) as u64).to_string();
        let time = (time_now() as u64).to_string();
        let version = get_version_string();
        let cmd_get = self.stats.cmd_get.to_string();
        let cmd_set = self.stats.cmd_set.to_string();
        let cmd_flush = self.stats.cmd_flush.to_string();
        let cmd_touch = self.stats.cmd_touch.to_string();
        let get_hits = storage.get_hits.to_string();
        let get_misses = storage.get_misses.to_string();
        let delete_hits = storage.delete_hits.to_string();
        let delete_misses = storage.delete_misses.to_string();
        let bytes_read = self.transport_stats.bytes_read.to_string();
        let bytes_written = self.transport_stats.bytes_written.to_string();
        let limit_maxbytes = self.cache.capacity.to_string();
        let bytes = storage.bytes.to_string();
        let curr_items = self.cache.len().to_string();
        let total_items = storage.total_items.to_string();
        let evictions = storage.evictions.to_string();
        let reclaimed = storage.reclaimed.to_string();

        let st_pid = Stat::new("pid", pid);
        let st_uptime = Stat::new("uptime", uptime);
        let st_time = Stat::new("time", time);
        let st_version = Stat::new("version", version);
        let st_cmd_get = Stat::new("cmd_get", cmd_get);
        let st_cmd_set = Stat::new("cmd_set", cmd_set);
        let st_cmd_flush = Stat::new("cmd_flush", cmd_flush);
        let st_cmd_touch = Stat::new("cmd_touch", cmd_touch);
        let st_get_hits = Stat::new("get_hits", get_hits);
        let st_get_misses = Stat::new("get_misses", get_misses);
        let st_delete_hits = Stat::new("delete_hits", delete_hits);
        let st_delete_misses = Stat::new("delete_misses", delete_misses);
        let st_bytes_read = Stat::new("bytes_read", bytes_read);
        let st_bytes_written = Stat::new("bytes_written", bytes_written);
        let st_limit_maxbytes = Stat::new("limit_maxbytes", limit_maxbytes);
        let st_bytes = Stat::new("bytes", bytes);
        let st_curr_items = Stat::new("curr_items", curr_items);
        let st_total_items = Stat::new("total_items", total_items);
        let st_evictions = Stat::new("evictions", evictions);
        let st_reclaimed = Stat::new("reclaimed", reclaimed);

        Resp::Stats(vec![st_pid,
                         st_uptime,
                         st_time,
                         st_version,
                         st_cmd_get,
                         st_cmd_set,
                         st_cmd_flush,
                         st_cmd_touch,
                         st_get_hits,
                         st_get_misses,
                         st_delete_hits,
                         st_delete_misses,
                         st_bytes_read,
                         st_bytes_written,
                         st_limit_maxbytes,
                         st_bytes,
                         st_curr_items,
                         st_total_items,
                         st_evictions,
                         st_reclaimed])
    }

    pub fn do_touch(&mut self, touch: Touch) -> Resp {
        let key = Key::new(touch.key.into_bytes());

        // See if the key is set
        let rv = self.cache.contains_key(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotFound),
            Err(_) => Some(Resp::Error),
        });

        // Load the value
        let rv = self.cache.remove(&key);
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(_) => None,
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        self.set_exptime(&mut value, touch.exptime);

        // Set it
        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!touch.noreply, match rv {
            Ok(_) => Resp::Touched,
            Err(_) => Resp::Error,
        })
    }

    pub fn do_version(&self) -> Resp {
        Resp::Version(get_version_string())
    }

    pub fn run(&mut self, cmd: Cmd) -> Resp {
        match cmd {
            Cmd::Delete(del) => self.do_delete(del),
            Cmd::Get(get) => self.do_get(get),
            Cmd::Quit => Resp::Empty,  // handled at transport level
            Cmd::Set(set) => match set.instr {
                SetInstr::Add => self.do_add(set),
                SetInstr::Append => self.do_append(set),
                SetInstr::Prepend => self.do_prepend(set),
                SetInstr::Replace => self.do_replace(set),
                SetInstr::Set => self.do_set(set),
                _ => Resp::Error,  // TODO
            },
            Cmd::Stats => self.do_stats(),
            Cmd::Touch(touch) => self.do_touch(touch),
            Cmd::Version => self.do_version(),
        }
    }

    pub fn update_transport_stats(&mut self, stats: TransportStats) {
        self.transport_stats = stats;
    }
}
<commit_msg>update stats for touch<commit_after>use common::consts::get_version_string;
use platform::process::get_pid;
use platform::time::time_now;
use storage::Cache;
use storage::CacheError;
use storage::Key;
use storage::Value;
use tcp_transport::stats::TransportStats;

use super::cmd::Cmd;
use super::cmd::Delete;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Touch;
use super::cmd::Value as CmdValue;

// For use to get an early exit from a function. The first parameter is a bool
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Option<Resp>. In case
// of None no return is made ie. execution continues. In case of Some(resp) we
// perform return resp;
//
// maybe_reply_stmt!(!noreply, match rv {
//     Ok(_) => None,   // do not return
//     Err(_) => Some(Resp::Error),
// });
//
// ->
//
// let opt = expr_provided;
// if !opt.is_none() {
//     if !noreply {
//         return opt.unwrap();
//     } else {
//         return Resp::Empty;
//     }
// }
macro_rules! maybe_reply_stmt {
    ( $cond:expr, $opt:expr ) => {
        {
            if !$opt.is_none() {
                if $cond {
                    return $opt.unwrap();
                } else {
                    return Resp::Empty;
                }
            };
        }
    };
}

// For use to conditionally produce a return value. The first parameter is used
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Resp.
//
// maybe_reply_expr!(!noreply, match rv {
//     Ok(_) => Resp::Stored,
//     Err(_) => Resp::Error,
// });
//
// ->
//
// let resp = expr_provided;
// if !noreply {
//     resp
// } else {
//     Resp::Empty
// }
macro_rules! maybe_reply_expr {
    ( $cond:expr, $resp:expr ) => {
        {
            if $cond {
                $resp
            } else {
                Resp::Empty
            }
        }
    };
}

struct DriverStats {
    cmd_get: u64,
    cmd_set: u64,
    cmd_flush: u64,
    cmd_touch: u64,
}

impl DriverStats {
    pub fn new() -> DriverStats {
        DriverStats {
            cmd_get: 0,
            cmd_set: 0,
            cmd_flush: 0,
            cmd_touch: 0,
        }
    }
}

pub struct Driver {
    cache: Cache,
    time_start: f64,

    stats: DriverStats,
    transport_stats: TransportStats, // this is a global snapshot
}

impl Driver {
    pub fn new(cache: Cache) -> Driver {
        Driver {
            cache: cache,
            stats: DriverStats::new(),
            time_start: time_now(),
            transport_stats: TransportStats::new(),
        }
    }

    fn set_exptime(&self, value: &mut Value, exptime: u32) {
        // If exptime is greater than zero we need to set it on the value
        if exptime > 0 {
            let tm;

            // Is it an interval greater than 30 days? Then it's a timestamp
            if exptime > 60 * 60 * 24 * 30 {
                tm = exptime as f64;

            } else {
                // Otherwise it's relative from now
                tm = time_now() + exptime as f64;
            }

            value.set_exptime(tm);
        }
    }

    fn do_add(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If so it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => {
                Some(Resp::NotStored)
            },
            Ok(false) => None,
            Err(_) => {
                Some(Resp::Error)
            }
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_append(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
            _ => None,
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Append the data we just received to the blob that is there
        value.item.extend(set.data);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_delete(&mut self, delete: Delete) -> Resp {
        let key = Key::new(delete.key.clone().into_bytes());

        let rv = self.cache.remove(&key);

        maybe_reply_expr!(!delete.noreply, match rv {
            Ok(_) => Resp::Deleted,
            Err(_) => Resp::NotFound,
        })
    }

    fn do_get(&mut self, get: Get) -> Resp {
        // Update stats
        self.stats.cmd_get += 1;

        let mut values = vec![];

        for key in get.keys {
            let key_str = key.clone();

            let key_st = Key::new(key.into_bytes());
            let rv = self.cache.get(&key_st);

            match rv {
                Ok(value) => {
                    let val_st = CmdValue {
                        key: key_str,
                        flags: value.flags,
                        data: value.item.clone(),
                    };
                    values.push(val_st);
                }
                // Keys that were not found are skipped, no error given
                Err(_) => (),
            }
        }

        Resp::Values(values)
    }

    fn do_prepend(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(_) => None,
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Prepend the data we just received to the blob that is there
        let mut new_item = Vec::with_capacity(set.data.len() + value.item.len());
        new_item.extend(set.data);
        new_item.extend(value.item);
        value.item = new_item;

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_replace(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If not it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_set(&mut self, set: Set) -> Resp {
        // Update stats
        self.stats.cmd_set += 1;

        let key = Key::new(set.key.into_bytes());
        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_stats(&self) -> Resp {
        let storage = self.cache.get_stats();

        let pid = get_pid().to_string();
        let uptime = ((time_now() - self.time_start) as u64).to_string();
        let time = (time_now() as u64).to_string();
        let version = get_version_string();
        let cmd_get = self.stats.cmd_get.to_string();
        let cmd_set = self.stats.cmd_set.to_string();
        let cmd_flush = self.stats.cmd_flush.to_string();
        let cmd_touch = self.stats.cmd_touch.to_string();
        let get_hits = storage.get_hits.to_string();
        let get_misses = storage.get_misses.to_string();
        let delete_hits = storage.delete_hits.to_string();
        let delete_misses = storage.delete_misses.to_string();
        let bytes_read = self.transport_stats.bytes_read.to_string();
        let bytes_written = self.transport_stats.bytes_written.to_string();
        let limit_maxbytes = self.cache.capacity.to_string();
        let bytes = storage.bytes.to_string();
        let curr_items = self.cache.len().to_string();
        let total_items = storage.total_items.to_string();
        let evictions = storage.evictions.to_string();
        let reclaimed = storage.reclaimed.to_string();

        let st_pid = Stat::new("pid", pid);
        let st_uptime = Stat::new("uptime", uptime);
        let st_time = Stat::new("time", time);
        let st_version = Stat::new("version", version);
        let st_cmd_get = Stat::new("cmd_get", cmd_get);
        let st_cmd_set = Stat::new("cmd_set", cmd_set);
        let st_cmd_flush = Stat::new("cmd_flush", cmd_flush);
        let st_cmd_touch = Stat::new("cmd_touch", cmd_touch);
        let st_get_hits = Stat::new("get_hits", get_hits);
        let st_get_misses = Stat::new("get_misses", get_misses);
        let st_delete_hits = Stat::new("delete_hits", delete_hits);
        let st_delete_misses = Stat::new("delete_misses", delete_misses);
        let st_bytes_read = Stat::new("bytes_read", bytes_read);
        let st_bytes_written = Stat::new("bytes_written", bytes_written);
        let st_limit_maxbytes = Stat::new("limit_maxbytes", limit_maxbytes);
        let st_bytes = Stat::new("bytes", bytes);
        let st_curr_items = Stat::new("curr_items", curr_items);
        let st_total_items = Stat::new("total_items", total_items);
        let st_evictions = Stat::new("evictions", evictions);
        let st_reclaimed = Stat::new("reclaimed", reclaimed);

        Resp::Stats(vec![st_pid,
                         st_uptime,
                         st_time,
                         st_version,
                         st_cmd_get,
                         st_cmd_set,
                         st_cmd_flush,
                         st_cmd_touch,
                         st_get_hits,
                         st_get_misses,
                         st_delete_hits,
                         st_delete_misses,
                         st_bytes_read,
                         st_bytes_written,
                         st_limit_maxbytes,
                         st_bytes,
                         st_curr_items,
                         st_total_items,
                         st_evictions,
                         st_reclaimed])
    }

    pub fn do_touch(&mut self, touch: Touch) -> Resp {
        // Update stats
        self.stats.cmd_touch += 1;

        let key = Key::new(touch.key.into_bytes());

        // See if the key is set
        let rv = self.cache.contains_key(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotFound),
            Err(_) => Some(Resp::Error),
        });

        // Load the value
        let rv = self.cache.remove(&key);
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(_) => None,
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        self.set_exptime(&mut value, touch.exptime);

        // Set it
        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!touch.noreply, match rv {
            Ok(_) => Resp::Touched,
            Err(_) => Resp::Error,
        })
    }

    pub fn do_version(&self) -> Resp {
        Resp::Version(get_version_string())
    }

    pub fn run(&mut self, cmd: Cmd) -> Resp {
        match cmd {
            Cmd::Delete(del) => self.do_delete(del),
            Cmd::Get(get) => self.do_get(get),
            Cmd::Quit => Resp::Empty,  // handled at transport level
            Cmd::Set(set) => match set.instr {
                SetInstr::Add => self.do_add(set),
                SetInstr::Append => self.do_append(set),
                SetInstr::Prepend => self.do_prepend(set),
                SetInstr::Replace => self.do_replace(set),
                SetInstr::Set => self.do_set(set),
                _ => Resp::Error,  // TODO
            },
            Cmd::Stats => self.do_stats(),
            Cmd::Touch(touch) => self.do_touch(touch),
            Cmd::Version => self.do_version(),
        }
    }

    pub fn update_transport_stats(&mut self, stats: TransportStats) {
        self.transport_stats = stats;
    }
}
<|endoftext|>


<commit_before>//! Example of basic search-and-replace functionality implemented on top
//! of a Ropey rope.
//!
//! Usage:
//!     search_and_replace <search_pattern> <replacement_text> <input_filepath>
//!
//! The file contents with the search-and-replace performed on it is sent to
//! stdout.

extern crate ropey;

use std::fs::File;
use std::io;

use ropey::{iter::Chars, Rope, RopeSlice};

fn main() {
    // Get arguments from commandline
    let (search_pattern, replacement_text, filepath) = if std::env::args().count() > 3 {
        (
            std::env::args().nth(1).unwrap(),
            std::env::args().nth(2).unwrap(),
            std::env::args().nth(3).unwrap(),
        )
    } else {
        println!(
            "Usage:\n    search_and_replace <search_pattern> <replacement_text> <input_filepath>"
        );
        return;
    };

    // Load file contents into a rope.
    let mut text = Rope::from_reader(io::BufReader::new(File::open(&filepath).unwrap())).expect("Cannot read file: either it doesn't exist, file permissions don't allow reading, or is not utf8 text.");

    // Do the search-and-replace.
    search_and_replace(&mut text, &search_pattern, &replacement_text);

    // Print the new text to stdout.
    println!("{}", text);
}

/// Searches the rope for `search_pattern` and replaces all matches with
/// `replacement_text`.
///
/// There are several ways this could be done:
///
/// 1. Clone the rope and then do the search on the original while replacing
///    on the clone.  This isn't as awful as it sounds because the clone
///    operation is constant-time and the two ropes will share most of their
///    storage in typical cases.  However, this probably isn't the best
///    general solution because it will use a lot of additional space if a
///    large percentage of the text is bring replaced.
///
/// 2. A two-stage approach: first find and collect all the matches, then
///    do the replacements on the original rope.  This is a good solution
///    when a relatively small number of matches are expected.  However, if
///    there are a large number of matches then the space to store the
///    matches themselves can become large.
///
/// 3. A piece-meal approach: search for the first match, replace it, then
///    restart the match from the end of the replacement, repeat.  This is
///    a good solution for memory-constrained situations.  However,
///    computationally it is likely the most expensive when there are a large
///    number of matches and there are costs associated with repeatedly
///    restarting the search.
///
/// 4. Combine approaches #2 and #3: collect a fixed number of matches and
///    replace them, then collect another batch matches and replace them, and
///    so on.  This is probably the best general solution, because it
///    combines the best of both worlds: it allows you to collect the matches
///    in a bounded amount of space, and any costs associated with restarting
///    the search are amortized across multiple matches.
///
/// In this implementation we take approach #4 because it seems the
/// all-around best.
fn search_and_replace(rope: &mut Rope, search_pattern: &str, replacement_text: &str) {
    const BATCH_SIZE: usize = 256;

    let replacement_text_len = replacement_text.chars().count();

    let mut head = 0; // Char index of the search/replace head.
    let mut matches = Vec::with_capacity(BATCH_SIZE);
    loop {
        // Collect the next batch of matches.  Note that we don't use
        // `Iterator::collect()` to collect the batch because we want to
        // re-use the same Vec to avoid unnecessary allocations.
        matches.clear();
        for m in SearchIter::from_rope_slice(&rope.slice(head..), &search_pattern).take(BATCH_SIZE)
        {
            matches.push(m);
        }

        // If there are no matches, we're done!
        if matches.len() == 0 {
            break;
        }

        // Replace the collected matches.
        let mut index_diff: isize = 0;
        for &(start, end) in matches.iter() {
            // Get the properly offset indices.
            let start_d = (head as isize + start as isize + index_diff) as usize;
            let end_d = (head as isize + end as isize + index_diff) as usize;

            // Do the replacement.
            rope.remove(start_d..end_d);
            rope.insert(start_d, &replacement_text);

            // Update the index offset.
            let match_len = (end - start) as isize;
            index_diff = index_diff - match_len + replacement_text_len as isize;
        }

        // Update head for next iteration.
        head = (head as isize + index_diff + matches.last().unwrap().1 as isize) as usize;
    }
}

/// An iterator over simple textual matches in a RopeSlice.
///
/// This implementation is somewhat naive, and could be sped up by using a
/// more sophisticated text searching algorithm such as Boyer-Moore or
/// Knuth-Morris-Pratt.
///
/// The important thing, however, is the interface.  For example, a regex
/// implementation providing an equivalent interface could easily be dropped
/// in, and the search-and-replace function above would then work with it
/// quite happily.
struct SearchIter<'a> {
    char_iter: Chars<'a>,
    search_pattern: &'a str,
    search_pattern_char_len: usize,
    cur_index: usize, // The current char index of the search head.
    possible_matches: Vec<std::str::Chars<'a>>, // Tracks where we are in the search pattern for the current possible matches.
}

impl<'a> SearchIter<'a> {
    fn from_rope_slice<'b>(slice: &'b RopeSlice, search_pattern: &'b str) -> SearchIter<'b> {
        assert!(
            search_pattern.len() > 0,
            "Can't search using an empty search pattern."
        );
        SearchIter {
            char_iter: slice.chars(),
            search_pattern: search_pattern,
            search_pattern_char_len: search_pattern.chars().count(),
            cur_index: 0,
            possible_matches: Vec::new(),
        }
    }
}

impl<'a> Iterator for SearchIter<'a> {
    type Item = (usize, usize);

    // Return the start/end char indices of the next match.
    fn next(&mut self) -> Option<(usize, usize)> {
        while let Some(next_char) = self.char_iter.next() {
            // Push current position onto the possible matches list, to seed a
            // possible beginning of a match.
            self.possible_matches.push(self.search_pattern.chars());

            // Check it against the next character in each of the potential
            // matches, removing the potential matches that don't match.
            // We're using indexing instead of iteration here so that we can
            // remove the possible matches as we go.
            let mut i = 0;
            while i < self.possible_matches.len() {
                let pattern_char = self.possible_matches[i].next().unwrap();
                if next_char == pattern_char {
                    if self.possible_matches[i].clone().next() == None {
                        self.cur_index += 1;
                        let char_match_range = (
                            self.cur_index - self.search_pattern_char_len,
                            self.cur_index,
                        );
                        self.possible_matches.clear();
                        return Some(char_match_range);
                    }
                    i += 1;
                } else {
                    self.possible_matches.swap_remove(i);
                }
            }

            self.cur_index += 1;
        }

        return None;
    }
}
<commit_msg>Better comments and misc cleanup of search-and-replace example.<commit_after>//! Example of basic search-and-replace functionality implemented on top
//! of a Ropey rope.
//!
//! Usage:
//!     search_and_replace <search_pattern> <replacement_text> <input_filepath>
//!
//! The file contents with the search-and-replace performed on it is sent to
//! stdout.

extern crate ropey;

use std::fs::File;
use std::io;

use ropey::{iter::Chars, Rope, RopeSlice};

fn main() {
    // Get arguments from commandline
    let (search_pattern, replacement_text, filepath) = if std::env::args().count() > 3 {
        (
            std::env::args().nth(1).unwrap(),
            std::env::args().nth(2).unwrap(),
            std::env::args().nth(3).unwrap(),
        )
    } else {
        println!(
            "Usage:\n    search_and_replace <search_pattern> <replacement_text> <input_filepath>"
        );
        return;
    };

    // Load file contents into a rope.
    let mut text = Rope::from_reader(io::BufReader::new(File::open(&filepath).unwrap())).expect("Cannot read file: either it doesn't exist, file permissions don't allow reading, or is not utf8 text.");

    // Do the search-and-replace.
    search_and_replace(&mut text, &search_pattern, &replacement_text);

    // Print the new text to stdout.
    println!("{}", text);
}

/// Searches the rope for `search_pattern` and replaces all matches with
/// `replacement_text`.
///
/// There are several ways this could be done:
///
/// 1. Clone the rope and then do the search on the original while replacing
///    on the clone.  This isn't as awful as it sounds because the clone
///    operation is constant-time and the two ropes will share most of their
///    storage in typical cases.  However, this probably isn't the best
///    general solution because it will use a lot of additional space if a
///    large percentage of the text is bring replaced.
///
/// 2. A two-stage approach: first find and collect all the matches, then
///    do the replacements on the original rope.  This is a good solution
///    when a relatively small number of matches are expected.  However, if
///    there are a large number of matches then the space to store the
///    matches themselves can become large.
///
/// 3. A piece-meal approach: search for the first match, replace it, then
///    restart the match from the end of the replacement, repeat.  This is
///    a good solution for memory-constrained situations.  However,
///    computationally it is likely the most expensive when there are a large
///    number of matches and there are costs associated with repeatedly
///    restarting the search.
///
/// 4. Combine approaches #2 and #3: collect a fixed number of matches and
///    replace them, then collect another batch matches and replace them, and
///    so on.  This is probably the best general solution, because it
///    combines the best of both worlds: it allows you to collect the matches
///    in a bounded amount of space, and any costs associated with restarting
///    the search are amortized across multiple matches.
///
/// In this implementation we take approach #4 because it seems the
/// all-around best.
fn search_and_replace(rope: &mut Rope, search_pattern: &str, replacement_text: &str) {
    const BATCH_SIZE: usize = 256;

    let replacement_text_len = replacement_text.chars().count();

    let mut head = 0; // Char index of the search/replace head.
    let mut matches = Vec::with_capacity(BATCH_SIZE);
    loop {
        // Collect the next batch of matches.  Note that we don't use
        // `Iterator::collect()` to collect the batch because we want to
        // re-use the same Vec to avoid unnecessary allocations.
        matches.clear();
        for m in SearchIter::from_rope_slice(&rope.slice(head..), &search_pattern).take(BATCH_SIZE)
        {
            matches.push(m);
        }

        // If there are no matches, we're done!
        if matches.len() == 0 {
            break;
        }

        // Replace the collected matches.
        let mut index_diff: isize = 0;
        for &(start, end) in matches.iter() {
            // Get the properly offset indices.
            let start_d = (head as isize + start as isize + index_diff) as usize;
            let end_d = (head as isize + end as isize + index_diff) as usize;

            // Do the replacement.
            rope.remove(start_d..end_d);
            rope.insert(start_d, &replacement_text);

            // Update the index offset.
            let match_len = (end - start) as isize;
            index_diff = index_diff - match_len + replacement_text_len as isize;
        }

        // Update head for next iteration.
        head = (head as isize + index_diff + matches.last().unwrap().1 as isize) as usize;
    }
}

/// An iterator over simple textual matches in a RopeSlice.
///
/// This implementation is somewhat naive, and could be sped up by using a
/// more sophisticated text searching algorithm such as Boyer-Moore or
/// Knuth-Morris-Pratt.
///
/// The important thing, however, is the interface.  For example, a regex
/// implementation providing an equivalent interface could easily be dropped
/// in, and the search-and-replace function above would then work with it
/// quite happily.
struct SearchIter<'a> {
    char_iter: Chars<'a>,
    search_pattern: &'a str,
    search_pattern_char_len: usize,
    cur_index: usize, // The current char index of the search head.
    possible_matches: Vec<std::str::Chars<'a>>, // Tracks where we are in the search pattern for the current possible matches.
}

impl<'a> SearchIter<'a> {
    fn from_rope_slice<'b>(slice: &'b RopeSlice, search_pattern: &'b str) -> SearchIter<'b> {
        assert!(
            search_pattern.len() > 0,
            "Can't search using an empty search pattern."
        );
        SearchIter {
            char_iter: slice.chars(),
            search_pattern: search_pattern,
            search_pattern_char_len: search_pattern.chars().count(),
            cur_index: 0,
            possible_matches: Vec::new(),
        }
    }
}

impl<'a> Iterator for SearchIter<'a> {
    type Item = (usize, usize);

    // Return the start/end char indices of the next match.
    fn next(&mut self) -> Option<(usize, usize)> {
        while let Some(next_char) = self.char_iter.next() {
            self.cur_index += 1;

            // Push new potential match, for a possible match starting at the
            // current char.
            self.possible_matches.push(self.search_pattern.chars());

            // Check the rope's char against the next character in each of
            // the potential matches, removing the potential matches that
            // don't match.  We're using indexing instead of iteration here
            // so that we can remove the possible matches as we go.
            let mut i = 0;
            while i < self.possible_matches.len() {
                let pattern_char = self.possible_matches[i].next().unwrap();
                if next_char == pattern_char {
                    if self.possible_matches[i].clone().next() == None {
                        // We have a match!  Reset possible matches and
                        // return the successful match's char indices.
                        let char_match_range = (
                            self.cur_index - self.search_pattern_char_len,
                            self.cur_index,
                        );
                        self.possible_matches.clear();
                        return Some(char_match_range);
                    } else {
                        // Match isn't complete yet, move on to the next.
                        i += 1;
                    }
                } else {
                    // Doesn't match, remove it.
                    self.possible_matches.swap_remove(i);
                }
            }
        }

        return None;
    }
}
<|endoftext|>


<commit_before>pub struct Allergies(pub usize);

#[derive(PartialEq, Debug)]
pub enum Allergen { Eggs, Peanuts, Shellfish, Strawberries, Tomatoes, Chocolate, Pollen, Cats }

impl Allergies {
    pub fn new(score: usize) -> Allergies {
        Allergies(score)
    }

    pub fn is_allergic_to(&self, allergen: &Allergen) -> bool {
        let allergens = Allergies::allergens();
        let index = allergens.iter().position(|x: &Allergen| x == allergen).unwrap();
        match self.0 & 1 << index { 0 => false, _ => true }
    }

    pub fn allergies(&self) -> Vec<Allergen> {
        Allergies::allergens().into_iter().filter(|allergen| self.is_allergic_to(allergen)).collect()
    }

    fn allergens() -> Vec<Allergen> {
        vec![Allergen::Eggs, Allergen::Peanuts, Allergen::Shellfish,
            Allergen::Strawberries, Allergen::Tomatoes, Allergen::Chocolate,
            Allergen::Pollen, Allergen::Cats]
    }
}
<commit_msg>allergies: replace a match -> bool with a comparison<commit_after>pub struct Allergies(pub usize);

#[derive(PartialEq, Debug)]
pub enum Allergen { Eggs, Peanuts, Shellfish, Strawberries, Tomatoes, Chocolate, Pollen, Cats }

impl Allergies {
    pub fn new(score: usize) -> Allergies {
        Allergies(score)
    }

    pub fn is_allergic_to(&self, allergen: &Allergen) -> bool {
        let allergens = Allergies::allergens();
        let index = allergens.iter().position(|x: &Allergen| x == allergen).unwrap();
        (self.0 & (1 << index)) != 0
    }

    pub fn allergies(&self) -> Vec<Allergen> {
        Allergies::allergens().into_iter().filter(|allergen| self.is_allergic_to(allergen)).collect()
    }

    fn allergens() -> Vec<Allergen> {
        vec![Allergen::Eggs, Allergen::Peanuts, Allergen::Shellfish,
            Allergen::Strawberries, Allergen::Tomatoes, Allergen::Chocolate,
            Allergen::Pollen, Allergen::Cats]
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;

use self::byteorder::{ByteOrder, BigEndian};

pub fn checksum(bytes: &[u8]) -> Vec<u8> {
  let mut checksum = 0xFFFF;
  let generator_polynomial = 0x8408;

  for byte in bytes {
    checksum = checksum ^ byte;

    for i in (0..8) {
      if checksum & 1 == 1 {
        checksum = (checksum >> 1) ^ generator_polynomial
      } else {
        checksum = checksum >> 1
      }
    }
  }

  checksum = checksum ^ 0xFFFF;
  let mut buf = [0u8; 2];
  byteorder::BigEndian::write_u16(&mut buf, checksum as u16);

  return buf.to_vec();
}
<commit_msg>checksum now returns array instead of vector<commit_after>extern crate byteorder;

use self::byteorder::{ByteOrder, BigEndian};

pub fn checksum(bytes: &[u8]) -> [u8; 2] {
  let mut checksum = 0xFFFF;
  let generator_polynomial = 0x8408;

  for byte in bytes {
    checksum = checksum ^ byte;

    for i in (0..8) {
      if checksum & 1 == 1 {
        checksum = (checksum >> 1) ^ generator_polynomial
      } else {
        checksum = checksum >> 1
      }
    }
  }

  checksum = checksum ^ 0xFFFF;
  let mut buf = [0u8; 2];
  byteorder::BigEndian::write_u16(&mut buf, checksum as u16);

  return buf;
}
<|endoftext|>


<commit_before>use std::os::unix::net::UnixStream;
use std::sync::mpsc;
use std::thread;

use hyper::server::Server as HyperServer;
use hyperlocal::UnixSocketServer;
use ruru::{Proc, VM};

use config::{BindingType, Config};
use core::{Core, Handler, ResponseFuture};
use request::Request;

pub struct Server {
    config: Config,
}

impl Server {
    pub fn new(config: Config) -> Self {
        Server { config: config }
    }

    pub fn listen(self, ruby_handler: Proc) {
        println!("[rust] Start listening");

        let (handler_stream, core_stream) = UnixStream::pair().unwrap();

        let (request_sender, request_receiver) = mpsc::channel::<Request>();
        let (response_future_sender, response_future_receiver) = mpsc::channel::<ResponseFuture>();

        let handler = Handler::new(request_sender, response_future_receiver, handler_stream);

        let mut core = Core::new(response_future_sender,
                                 request_receiver,
                                 ruby_handler,
                                 core_stream,
                                 self.config.rack_thread_pool_size());

        thread::spawn(move || {
            let handler_function = || -> () {
                println!("[hyper] GVL released for server thread");

                let address = self.config.listen_on();
                let thread_pool_size = self.config.native_thread_pool_size();

                println!("[hyper] Spawning {} native thread(s)", thread_pool_size);
                println!("[hyper] Listening on {}", address);

                match *self.config.binding_type() {
                    BindingType::Unix => {
                        UnixSocketServer::new(address).unwrap()
                            .handle_threads(handler, thread_pool_size).unwrap();
                    },
                    BindingType::Tcp => {
                        HyperServer::http(address).unwrap()
                            .handle_threads(handler, thread_pool_size).unwrap();
                    }
                };
            };

            let unblock_function = || {
                ()
            };

            VM::thread_call_without_gvl(
                handler_function,
                Some(unblock_function)
            );
        });

        core.run();
    }
}
<commit_msg>Unlock GVL with Thread instead of VM<commit_after>use std::os::unix::net::UnixStream;
use std::sync::mpsc;
use std::thread;

use hyper::server::Server as HyperServer;
use hyperlocal::UnixSocketServer;
use ruru::{Proc, Thread};

use config::{BindingType, Config};
use core::{Core, Handler, ResponseFuture};
use request::Request;

pub struct Server {
    config: Config,
}

impl Server {
    pub fn new(config: Config) -> Self {
        Server { config: config }
    }

    pub fn listen(self, ruby_handler: Proc) {
        println!("[rust] Start listening");

        let (handler_stream, core_stream) = UnixStream::pair().unwrap();

        let (request_sender, request_receiver) = mpsc::channel::<Request>();
        let (response_future_sender, response_future_receiver) = mpsc::channel::<ResponseFuture>();

        let handler = Handler::new(request_sender, response_future_receiver, handler_stream);

        let mut core = Core::new(response_future_sender,
                                 request_receiver,
                                 ruby_handler,
                                 core_stream,
                                 self.config.rack_thread_pool_size());

        thread::spawn(move || {
            let handler_function = || -> () {
                println!("[hyper] GVL released for server thread");

                let address = self.config.listen_on();
                let thread_pool_size = self.config.native_thread_pool_size();

                println!("[hyper] Spawning {} native thread(s)", thread_pool_size);
                println!("[hyper] Listening on {}", address);

                match *self.config.binding_type() {
                    BindingType::Unix => {
                        UnixSocketServer::new(address).unwrap()
                            .handle_threads(handler, thread_pool_size).unwrap();
                    },
                    BindingType::Tcp => {
                        HyperServer::http(address).unwrap()
                            .handle_threads(handler, thread_pool_size).unwrap();
                    }
                };
            };

            let unblock_function = || {
                ()
            };

            Thread::call_without_gvl(
                handler_function,
                Some(unblock_function)
            );
        });

        core.run();
    }
}
<|endoftext|>


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

#![deny(warnings)]

#[macro_use]extern crate style;
extern crate app_units;
extern crate atomic_refcell;
extern crate cssparser;
extern crate env_logger;
extern crate euclid;
extern crate libc;
#[macro_use] extern crate log;
extern crate parking_lot;
extern crate selectors;
extern crate servo_url;
extern crate style_traits;

#[allow(non_snake_case)]
pub mod glue;
mod stylesheet_loader;

// FIXME(bholley): This should probably go away once we harmonize the allocators.
#[no_mangle]
pub extern "C" fn je_malloc_usable_size(_: *const ::libc::c_void) -> ::libc::size_t { 0 }
<commit_msg>Auto merge of #15348 - emilio:unused-macro-use, r=SimonSapin<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![deny(warnings)]

extern crate app_units;
extern crate atomic_refcell;
extern crate cssparser;
extern crate env_logger;
extern crate euclid;
extern crate libc;
#[macro_use] extern crate log;
extern crate parking_lot;
extern crate selectors;
extern crate servo_url;
extern crate style;
extern crate style_traits;

#[allow(non_snake_case)]
pub mod glue;
mod stylesheet_loader;

// FIXME(bholley): This should probably go away once we harmonize the allocators.
#[no_mangle]
pub extern "C" fn je_malloc_usable_size(_: *const ::libc::c_void) -> ::libc::size_t { 0 }
<|endoftext|>


<commit_before>#![feature(slice_concat_ext)]

extern crate compiletest_rs as compiletest;
extern crate colored;

use colored::*;

use std::slice::SliceConcatExt;
use std::path::{PathBuf, Path};
use std::io::Write;

macro_rules! eprintln {
    ($($arg:tt)*) => {
        let stderr = std::io::stderr();
        writeln!(stderr.lock(), $($arg)*).unwrap();
    }
}

fn miri_path() -> PathBuf {
    if rustc_test_suite().is_some() {
        PathBuf::from(option_env!("MIRI_PATH").unwrap())
    } else {
        PathBuf::from(concat!("target/", env!("PROFILE"), "/miri"))
    }
}

fn rustc_test_suite() -> Option<PathBuf> {
    option_env!("RUSTC_TEST_SUITE").map(PathBuf::from)
}

fn rustc_lib_path() -> PathBuf {
    option_env!("RUSTC_LIB_PATH").unwrap().into()
}

fn have_fullmir() -> bool {
    // We assume we have full MIR when MIRI_SYSROOT is set or when we are in rustc
    std::env::var("MIRI_SYSROOT").is_ok() || rustc_test_suite().is_some()
}

fn compile_fail(sysroot: &Path, path: &str, target: &str, host: &str, need_fullmir: bool) {
    if need_fullmir && !have_fullmir() {
        eprintln!("{}", format!(
            "## Skipping compile-fail tests in {} against miri for target {} due to missing mir",
            path,
            target
        ).yellow().bold());
        return;
    }

    eprintln!("{}", format!(
        "## Running compile-fail tests in {} against miri for target {}",
        path,
        target
    ).green().bold());
    let mut config = compiletest::Config::default().tempdir();
    config.mode = "compile-fail".parse().expect("Invalid mode");
    config.rustc_path = miri_path();
    let mut flags = Vec::new();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    flags.push(format!("--sysroot {}", sysroot.display()));
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    config.src_base = PathBuf::from(path.to_string());
    flags.push("-Zmir-opt-level=0".to_owned()); // optimization circumvents some stacked borrow checks
    flags.push("-Zmir-emit-validate=1".to_owned());
    config.target_rustcflags = Some(flags.join(" "));
    config.target = target.to_owned();
    config.host = host.to_owned();
    compiletest::run_tests(&config);
}

fn miri_pass(sysroot: &Path, path: &str, target: &str, host: &str, need_fullmir: bool, opt: bool) {
    if need_fullmir && !have_fullmir() {
        eprintln!("{}", format!(
            "## Skipping run-pass tests in {} against miri for target {} due to missing mir",
            path,
            target
        ).yellow().bold());
        return;
    }

    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running run-pass tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());
    let mut config = compiletest::Config::default().tempdir();
    config.mode = "ui".parse().expect("Invalid mode");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.host = host.to_owned();
    config.rustc_path = miri_path();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    let mut flags = Vec::new();
    flags.push(format!("--sysroot {}", sysroot.display()));
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    if opt {
        // FIXME: Using level 1 (instead of 3) for now, as the optimizer is pretty broken
        // and crashes...
        // Level 0 and 1 are not the same, so this still gives us *some* coverage.
        // See https://github.com/rust-lang/rust/issues/50411
        flags.push("-Zmir-opt-level=1".to_owned());
    } else {
        flags.push("-Zmir-opt-level=0".to_owned());
        // For now, only validate without optimizations.  Inlining breaks validation.
        flags.push("-Zmir-emit-validate=1".to_owned());
    }
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

fn is_target_dir<P: Into<PathBuf>>(path: P) -> bool {
    let mut path = path.into();
    path.push("lib");
    path.metadata().map(|m| m.is_dir()).unwrap_or(false)
}

fn for_all_targets<F: FnMut(String)>(sysroot: &Path, mut f: F) {
    let target_dir = sysroot.join("lib").join("rustlib");
    for entry in std::fs::read_dir(target_dir).expect("invalid sysroot") {
        let entry = entry.unwrap();
        if !is_target_dir(entry.path()) {
            continue;
        }
        let target = entry.file_name().into_string().unwrap();
        f(target);
    }
}

fn get_sysroot() -> PathBuf {
    let sysroot = std::env::var("MIRI_SYSROOT").unwrap_or_else(|_| {
        let sysroot = std::process::Command::new("rustc")
            .arg("--print")
            .arg("sysroot")
            .output()
            .expect("rustc not found")
            .stdout;
        String::from_utf8(sysroot).expect("sysroot is not utf8")
    });
    PathBuf::from(sysroot.trim())
}

fn get_host() -> String {
    let rustc = rustc_test_suite().unwrap_or(PathBuf::from("rustc"));
    println!("using rustc at {}", rustc.display());
    let host = std::process::Command::new(rustc)
        .arg("-vV")
        .output()
        .expect("rustc not found for -vV")
        .stdout;
    let host = std::str::from_utf8(&host).expect("sysroot is not utf8");
    let host = host.split("\nhost: ").nth(1).expect(
        "no host: part in rustc -vV",
    );
    let host = host.split('\n').next().expect("no \n after host");
    String::from(host)
}

fn run_pass_miri(opt: bool) {
    let sysroot = get_sysroot();
    let host = get_host();

    for_all_targets(&sysroot, |target| {
        miri_pass(&sysroot, "tests/run-pass", &target, &host, false, opt);
    });
    miri_pass(&sysroot, "tests/run-pass-fullmir", &host, &host, true, opt);
}

fn compile_fail_miri() {
    let sysroot = get_sysroot();
    let host = get_host();

    // FIXME: run tests for other targets, too
    compile_fail(&sysroot, "tests/compile-fail", &host, &host, false);
    compile_fail(&sysroot, "tests/compile-fail-fullmir", &host, &host, true);
}

#[test]
fn test() {
    // We put everything into a single test to avoid the parallelism `cargo test`
    // introduces.  We still get parallelism within our tests because `compiletest`
    // uses `libtest` which runs jobs in parallel.

    run_pass_miri(false);
    run_pass_miri(true);

    compile_fail_miri();
}
<commit_msg>opt level 3 works again :)<commit_after>#![feature(slice_concat_ext)]

extern crate compiletest_rs as compiletest;
extern crate colored;

use colored::*;

use std::slice::SliceConcatExt;
use std::path::{PathBuf, Path};
use std::io::Write;

macro_rules! eprintln {
    ($($arg:tt)*) => {
        let stderr = std::io::stderr();
        writeln!(stderr.lock(), $($arg)*).unwrap();
    }
}

fn miri_path() -> PathBuf {
    if rustc_test_suite().is_some() {
        PathBuf::from(option_env!("MIRI_PATH").unwrap())
    } else {
        PathBuf::from(concat!("target/", env!("PROFILE"), "/miri"))
    }
}

fn rustc_test_suite() -> Option<PathBuf> {
    option_env!("RUSTC_TEST_SUITE").map(PathBuf::from)
}

fn rustc_lib_path() -> PathBuf {
    option_env!("RUSTC_LIB_PATH").unwrap().into()
}

fn have_fullmir() -> bool {
    // We assume we have full MIR when MIRI_SYSROOT is set or when we are in rustc
    std::env::var("MIRI_SYSROOT").is_ok() || rustc_test_suite().is_some()
}

fn compile_fail(sysroot: &Path, path: &str, target: &str, host: &str, need_fullmir: bool) {
    if need_fullmir && !have_fullmir() {
        eprintln!("{}", format!(
            "## Skipping compile-fail tests in {} against miri for target {} due to missing mir",
            path,
            target
        ).yellow().bold());
        return;
    }

    eprintln!("{}", format!(
        "## Running compile-fail tests in {} against miri for target {}",
        path,
        target
    ).green().bold());
    let mut config = compiletest::Config::default().tempdir();
    config.mode = "compile-fail".parse().expect("Invalid mode");
    config.rustc_path = miri_path();
    let mut flags = Vec::new();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    flags.push(format!("--sysroot {}", sysroot.display()));
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    config.src_base = PathBuf::from(path.to_string());
    flags.push("-Zmir-opt-level=0".to_owned()); // optimization circumvents some stacked borrow checks
    flags.push("-Zmir-emit-validate=1".to_owned());
    config.target_rustcflags = Some(flags.join(" "));
    config.target = target.to_owned();
    config.host = host.to_owned();
    compiletest::run_tests(&config);
}

fn miri_pass(sysroot: &Path, path: &str, target: &str, host: &str, need_fullmir: bool, opt: bool) {
    if need_fullmir && !have_fullmir() {
        eprintln!("{}", format!(
            "## Skipping run-pass tests in {} against miri for target {} due to missing mir",
            path,
            target
        ).yellow().bold());
        return;
    }

    let opt_str = if opt { " with optimizations" } else { "" };
    eprintln!("{}", format!(
        "## Running run-pass tests in {} against miri for target {}{}",
        path,
        target,
        opt_str
    ).green().bold());
    let mut config = compiletest::Config::default().tempdir();
    config.mode = "ui".parse().expect("Invalid mode");
    config.src_base = PathBuf::from(path);
    config.target = target.to_owned();
    config.host = host.to_owned();
    config.rustc_path = miri_path();
    if rustc_test_suite().is_some() {
        config.run_lib_path = rustc_lib_path();
        config.compile_lib_path = rustc_lib_path();
    }
    let mut flags = Vec::new();
    flags.push(format!("--sysroot {}", sysroot.display()));
    flags.push("-Dwarnings -Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    if opt {
        flags.push("-Zmir-opt-level=3".to_owned());
    } else {
        flags.push("-Zmir-opt-level=0".to_owned());
        // For now, only validate without optimizations.  Inlining breaks validation.
        flags.push("-Zmir-emit-validate=1".to_owned());
    }
    config.target_rustcflags = Some(flags.join(" "));
    compiletest::run_tests(&config);
}

fn is_target_dir<P: Into<PathBuf>>(path: P) -> bool {
    let mut path = path.into();
    path.push("lib");
    path.metadata().map(|m| m.is_dir()).unwrap_or(false)
}

fn for_all_targets<F: FnMut(String)>(sysroot: &Path, mut f: F) {
    let target_dir = sysroot.join("lib").join("rustlib");
    for entry in std::fs::read_dir(target_dir).expect("invalid sysroot") {
        let entry = entry.unwrap();
        if !is_target_dir(entry.path()) {
            continue;
        }
        let target = entry.file_name().into_string().unwrap();
        f(target);
    }
}

fn get_sysroot() -> PathBuf {
    let sysroot = std::env::var("MIRI_SYSROOT").unwrap_or_else(|_| {
        let sysroot = std::process::Command::new("rustc")
            .arg("--print")
            .arg("sysroot")
            .output()
            .expect("rustc not found")
            .stdout;
        String::from_utf8(sysroot).expect("sysroot is not utf8")
    });
    PathBuf::from(sysroot.trim())
}

fn get_host() -> String {
    let rustc = rustc_test_suite().unwrap_or(PathBuf::from("rustc"));
    println!("using rustc at {}", rustc.display());
    let host = std::process::Command::new(rustc)
        .arg("-vV")
        .output()
        .expect("rustc not found for -vV")
        .stdout;
    let host = std::str::from_utf8(&host).expect("sysroot is not utf8");
    let host = host.split("\nhost: ").nth(1).expect(
        "no host: part in rustc -vV",
    );
    let host = host.split('\n').next().expect("no \n after host");
    String::from(host)
}

fn run_pass_miri(opt: bool) {
    let sysroot = get_sysroot();
    let host = get_host();

    for_all_targets(&sysroot, |target| {
        miri_pass(&sysroot, "tests/run-pass", &target, &host, false, opt);
    });
    miri_pass(&sysroot, "tests/run-pass-fullmir", &host, &host, true, opt);
}

fn compile_fail_miri() {
    let sysroot = get_sysroot();
    let host = get_host();

    // FIXME: run tests for other targets, too
    compile_fail(&sysroot, "tests/compile-fail", &host, &host, false);
    compile_fail(&sysroot, "tests/compile-fail-fullmir", &host, &host, true);
}

#[test]
fn test() {
    // We put everything into a single test to avoid the parallelism `cargo test`
    // introduces.  We still get parallelism within our tests because `compiletest`
    // uses `libtest` which runs jobs in parallel.

    run_pass_miri(false);
    run_pass_miri(true);

    compile_fail_miri();
}
<|endoftext|>