<commit_before>
use std::ffi;
use std::convert::AsRef;
use std::path::Path;

pub mod record;
pub mod header;

use htslib;
use bcf::header::{HeaderView, SampleSubset};
use utils;

pub use bcf::header::Header;
pub use bcf::record::Record;

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
}

unsafe impl Send for Reader {}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Result<Self, BCFError> {
        let htsfile = try!(bcf_open(path, b"r"));
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Ok(Reader { inner: htsfile, header: HeaderView::new(header) })
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Writer {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
    subset: Option<SampleSubset>,
}

unsafe impl Send for Writer {}

impl Writer {
    pub fn new<P: AsRef<Path>>(path: &P, header: &Header, uncompressed: bool, vcf: bool) -> Result<Self, BCFError> {
        let mode: &[u8] = match (uncompressed, vcf) {
            (true, true)   => b"w",
            (false, true)  => b"wz",
            (true, false)  => b"wu",
            (false, false) => b"wb",
        };

        let htsfile = try!(bcf_open(path, mode));
        unsafe { htslib::vcf::bcf_hdr_write(htsfile, header.inner) };
        Ok(Writer {
            inner: htsfile,
            header: HeaderView::new(unsafe { htslib::vcf::bcf_hdr_dup(header.inner) }),
            subset: header.subset.clone()
        })
    }

    /// Translate record to header of this writer.
    pub fn translate(&mut self, record: &mut record::Record) {
        unsafe {
            htslib::vcf::bcf_translate(self.header.inner, record.header, record.inner);
        }
        record.header = self.header.inner;
    }

    /// Subset samples of record to match header of this writer.
    pub fn subset(&mut self, record: &mut record::Record) {
        match self.subset {
            Some(ref mut subset) => unsafe {
                htslib::vcf::bcf_subset(self.header.inner, record.inner, subset.len() as i32, subset.as_mut_ptr());
            },
            None         => ()
        }
    }

    pub fn write(&mut self, record: &record::Record) -> Result<(), WriteError> {
        if unsafe { htslib::vcf::bcf_write(self.inner, self.header.inner, record.inner) } == -1 {
            Err(WriteError::WriteError)
        }
        else {
            Ok(())
        }
    }
}

impl Drop for Writer {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum BCFError {
        InvalidPath {
            description("invalid path")
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> Result<*mut htslib::vcf::htsFile, BCFError> {
    if let Some(p) = utils::path_to_cstring(path) {
        Ok(unsafe {
            htslib::vcf::hts_open(
                p.as_ptr(),
                ffi::CString::new(mode).unwrap().as_ptr()
            )
        })
    }
    else {
        Err(BCFError::InvalidPath)
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum ReadError {
        Invalid {
            description("invalid record")
        }
        NoMoreRecord {
            description("no more record")
        }
    }
}

impl ReadError {
    /// Returns true if no record has been read because the end of the file was reached.
    pub fn is_eof(&self) -> bool {
        match self {
            &ReadError::NoMoreRecord => true,
            _ => false
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum WriteError {
        WriteError {
            description("failed to write record")
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use std::path::Path;

    fn _test_read<P: AsRef<Path>>(path: &P) {
        let bcf = Reader::new(path).ok().expect("Error opening file.");
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.sample_count(), 1);

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }

    #[test]
    fn test_read() {
        _test_read(&"test/test.bcf");
    }

    #[test]
    fn test_write() {
        let bcf = Reader::new(&"test/test_multi.bcf").ok().expect("Error opening file.");
        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bcfpath = tmp.path().join("test.bcf");
        println!("{:?}", bcfpath);
        {
            let header = Header::subset_template(&bcf.header, &[b"NA12878.subsample-0.25-0"]).ok().expect("Error subsetting samples.");
            let mut writer = Writer::new(&bcfpath, &header, false, false).ok().expect("Error opening file.");
            for rec in bcf.records() {
                let mut record = rec.ok().expect("Error reading record.");
                writer.translate(&mut record);
                writer.subset(&mut record);
                record.trim_alleles().ok().expect("Error trimming alleles.");
                writer.write(&record).ok().expect("Error writing record");
            }
        }
        {
            _test_read(&bcfpath);
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }

    #[test]
    fn test_strings() {
        let vcf = Reader::new(&"test/test_string.vcf").ok().expect("Error opening file.");
        let fs1 = [&b"LongString1"[..], &b"LongString2"[..], &b"."[..], &b"LongString4"[..], &b"evenlength"[..], &b"ss6"[..]];
        for (i, rec) in vcf.records().enumerate() {
            println!("record {}", i);
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.info(b"S1").string().ok().expect("Error reading string.")[0], format!("string{}", i + 1).as_bytes());
            println!("{}", String::from_utf8_lossy(record.format(b"FS1").string().ok().expect("Error reading string.")[0]));
            assert_eq!(record.format(b"FS1").string().ok().expect("Error reading string.")[0], fs1[i]);
        }
    }

    #[test]
    fn test_missing() {
        let vcf = Reader::new(&"test/test_missing.vcf").ok().expect("Error opening file.");
        let fn4 = [&[record::MISSING_INTEGER, record::MISSING_INTEGER, record::MISSING_INTEGER, record::MISSING_INTEGER][..], &[record::MISSING_INTEGER][..]];
        let f1 = [false, true];
        for (i, rec) in vcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.info(b"F1").float().ok().expect("Error reading float.")[0].is_nan(), f1[i]);
            assert_eq!(record.format(b"FN4").integer().ok().expect("Error reading integer.")[1], fn4[i]);
            println!("{:?}", record.format(b"FF4").float().ok().expect("Error reading float.")[1]);
            assert!(record.format(b"FF4").float().ok().expect("Error reading float.")[1].iter().all(|v| v.is_nan()));
        }
    }

    #[test]
    fn test_genotypes() {
        let vcf = Reader::new(&"test/test_string.vcf").ok().expect("Error opening file.");
        let expected = ["./1", "1|1", "0/1", "0|1", "1/1", "1/1"];
        for (rec, exp_gt) in vcf.records().zip(expected.into_iter()) {
            let mut rec = rec.ok().expect("Error reading record.");
            let genotypes = rec.genotypes().expect("Error reading genotypes");
            assert_eq!(&format!("{}", genotypes.get(0)), exp_gt);
        }
    }
}
<commit_msg>Define missing values.<commit_after>
use std::ffi;
use std::convert::AsRef;
use std::path::Path;

pub mod record;
pub mod header;

use htslib;
use bcf::header::{HeaderView, SampleSubset};
use utils;

pub use bcf::header::Header;
pub use bcf::record::Record;

pub const MISSING_FLOAT: f32 = 0x7F800001 as f32;
#[allow(overflowing_literals)]
pub const MISSING_INTEGER: i32 = 0x80000000;

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
}

unsafe impl Send for Reader {}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Result<Self, BCFError> {
        let htsfile = try!(bcf_open(path, b"r"));
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Ok(Reader { inner: htsfile, header: HeaderView::new(header) })
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Writer {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
    subset: Option<SampleSubset>,
}

unsafe impl Send for Writer {}

impl Writer {
    pub fn new<P: AsRef<Path>>(path: &P, header: &Header, uncompressed: bool, vcf: bool) -> Result<Self, BCFError> {
        let mode: &[u8] = match (uncompressed, vcf) {
            (true, true)   => b"w",
            (false, true)  => b"wz",
            (true, false)  => b"wu",
            (false, false) => b"wb",
        };

        let htsfile = try!(bcf_open(path, mode));
        unsafe { htslib::vcf::bcf_hdr_write(htsfile, header.inner) };
        Ok(Writer {
            inner: htsfile,
            header: HeaderView::new(unsafe { htslib::vcf::bcf_hdr_dup(header.inner) }),
            subset: header.subset.clone()
        })
    }

    /// Translate record to header of this writer.
    pub fn translate(&mut self, record: &mut record::Record) {
        unsafe {
            htslib::vcf::bcf_translate(self.header.inner, record.header, record.inner);
        }
        record.header = self.header.inner;
    }

    /// Subset samples of record to match header of this writer.
    pub fn subset(&mut self, record: &mut record::Record) {
        match self.subset {
            Some(ref mut subset) => unsafe {
                htslib::vcf::bcf_subset(self.header.inner, record.inner, subset.len() as i32, subset.as_mut_ptr());
            },
            None         => ()
        }
    }

    pub fn write(&mut self, record: &record::Record) -> Result<(), WriteError> {
        if unsafe { htslib::vcf::bcf_write(self.inner, self.header.inner, record.inner) } == -1 {
            Err(WriteError::WriteError)
        }
        else {
            Ok(())
        }
    }
}

impl Drop for Writer {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum BCFError {
        InvalidPath {
            description("invalid path")
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> Result<*mut htslib::vcf::htsFile, BCFError> {
    if let Some(p) = utils::path_to_cstring(path) {
        Ok(unsafe {
            htslib::vcf::hts_open(
                p.as_ptr(),
                ffi::CString::new(mode).unwrap().as_ptr()
            )
        })
    }
    else {
        Err(BCFError::InvalidPath)
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum ReadError {
        Invalid {
            description("invalid record")
        }
        NoMoreRecord {
            description("no more record")
        }
    }
}

impl ReadError {
    /// Returns true if no record has been read because the end of the file was reached.
    pub fn is_eof(&self) -> bool {
        match self {
            &ReadError::NoMoreRecord => true,
            _ => false
        }
    }
}

quick_error! {
    #[derive(Debug)]
    pub enum WriteError {
        WriteError {
            description("failed to write record")
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use std::path::Path;

    fn _test_read<P: AsRef<Path>>(path: &P) {
        let bcf = Reader::new(path).ok().expect("Error opening file.");
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.sample_count(), 1);

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }

    #[test]
    fn test_read() {
        _test_read(&"test/test.bcf");
    }

    #[test]
    fn test_write() {
        let bcf = Reader::new(&"test/test_multi.bcf").ok().expect("Error opening file.");
        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bcfpath = tmp.path().join("test.bcf");
        println!("{:?}", bcfpath);
        {
            let header = Header::subset_template(&bcf.header, &[b"NA12878.subsample-0.25-0"]).ok().expect("Error subsetting samples.");
            let mut writer = Writer::new(&bcfpath, &header, false, false).ok().expect("Error opening file.");
            for rec in bcf.records() {
                let mut record = rec.ok().expect("Error reading record.");
                writer.translate(&mut record);
                writer.subset(&mut record);
                record.trim_alleles().ok().expect("Error trimming alleles.");
                writer.write(&record).ok().expect("Error writing record");
            }
        }
        {
            _test_read(&bcfpath);
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }

    #[test]
    fn test_strings() {
        let vcf = Reader::new(&"test/test_string.vcf").ok().expect("Error opening file.");
        let fs1 = [&b"LongString1"[..], &b"LongString2"[..], &b"."[..], &b"LongString4"[..], &b"evenlength"[..], &b"ss6"[..]];
        for (i, rec) in vcf.records().enumerate() {
            println!("record {}", i);
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.info(b"S1").string().ok().expect("Error reading string.")[0], format!("string{}", i + 1).as_bytes());
            println!("{}", String::from_utf8_lossy(record.format(b"FS1").string().ok().expect("Error reading string.")[0]));
            assert_eq!(record.format(b"FS1").string().ok().expect("Error reading string.")[0], fs1[i]);
        }
    }

    #[test]
    fn test_missing() {
        let vcf = Reader::new(&"test/test_missing.vcf").ok().expect("Error opening file.");
        let fn4 = [&[record::MISSING_INTEGER, record::MISSING_INTEGER, record::MISSING_INTEGER, record::MISSING_INTEGER][..], &[record::MISSING_INTEGER][..]];
        let f1 = [false, true];
        for (i, rec) in vcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.info(b"F1").float().ok().expect("Error reading float.")[0].is_nan(), f1[i]);
            assert_eq!(record.format(b"FN4").integer().ok().expect("Error reading integer.")[1], fn4[i]);
            println!("{:?}", record.format(b"FF4").float().ok().expect("Error reading float.")[1]);
            assert!(record.format(b"FF4").float().ok().expect("Error reading float.")[1].iter().all(|v| v.is_nan()));
        }
    }

    #[test]
    fn test_genotypes() {
        let vcf = Reader::new(&"test/test_string.vcf").ok().expect("Error opening file.");
        let expected = ["./1", "1|1", "0/1", "0|1", "1/1", "1/1"];
        for (rec, exp_gt) in vcf.records().zip(expected.into_iter()) {
            let mut rec = rec.ok().expect("Error reading record.");
            let genotypes = rec.genotypes().expect("Error reading genotypes");
            assert_eq!(&format!("{}", genotypes.get(0)), exp_gt);
        }
    }
}
<|endoftext|>


<commit_before>pub type BitKmer = u64;

pub struct BitKmerIter<'a> {
    k: u8,
    pos: usize,
    cur_kmer: BitKmer,
    buffer: &'a [u8],
}

pub fn fast_extend_kmer(kmer: &BitKmer, k: &u8, new_char: &u8) -> BitKmer {
    // new_char must be in ACGTE (E gets counted as T unfortunately)
    assert!(new_char & 223 == 84 || new_char & 217 == 65);
    let new_kmer = (kmer << 2) + (((new_char >> 2) ^ (new_char >> 1)) & 3u8) as BitKmer;

    // mask out any overflowed bits
    new_kmer & (BitKmer::pow(2, (2 * k) as u32) - 1) as BitKmer
}

pub fn extend_kmer(kmer: &BitKmer, k: &u8, new_char: &u8) -> BitKmer {
    let new_kmer = (kmer << 2) + match new_char {
        &b'A' | &b'a' => 0 as BitKmer,
        &b'C' | &b'c' => 1 as BitKmer,
        &b'G' | &b'g' => 2 as BitKmer,
        &b'T' | &b't' => 3 as BitKmer,
        _ => panic!("Received a non ACGT nucleotide"),
    };

    // mask out any overflowed bits
    new_kmer & (BitKmer::pow(2, (2 * k) as u32) - 1) as BitKmer
}

impl<'a> BitKmerIter<'a> {
    pub fn new(slice: &'a [u8], k: u8) -> BitKmerIter<'a> {
        // TODO: raise an error if the slice is shorter than k
        let mut kmer = 0u64;
        if slice.len() >= k as usize {
            for i in 0..k - 1 {
                kmer = fast_extend_kmer(&kmer, &k, &slice[i as usize]);
            }
        }

        BitKmerIter {
            k: k,
            pos: (k - 2) as usize,
            cur_kmer: kmer,
            buffer: slice,
        }
    }
}

impl<'a> Iterator for BitKmerIter<'a> {
    type Item = BitKmer;

    fn next(&mut self) -> Option<BitKmer> {
        self.pos += 1;
        if self.pos > self.buffer.len() {
            return None;
        }
        self.cur_kmer = fast_extend_kmer(&self.cur_kmer, &self.k, &self.buffer[self.pos]);

        Some(self.cur_kmer)
    }
}

#[test]
fn test_iterator() {
    let seq = "ACGTA".as_bytes();
    let mut kmer_iter = BitKmerIter::new(seq, 3);
    assert_eq!(kmer_iter.next(), Some(6));
    assert_eq!(kmer_iter.next(), Some(27));
    assert_eq!(kmer_iter.next(), Some(44));
}

fn reverse_complement(kmer: BitKmer, k: u8) -> BitKmer {
    // FIXME: this is not going to work with BitKmers of u128 or u32
    // inspired from https://www.biostars.org/p/113640/
    let mut new_kmer = kmer;
    // reverse it
    new_kmer = (new_kmer >> 2 & 0x3333333333333333) | (new_kmer & 0x3333333333333333) << 2;
    new_kmer = (new_kmer >> 4 & 0x0F0F0F0F0F0F0F0F) | (new_kmer & 0x0F0F0F0F0F0F0F0F) << 4;
    new_kmer = (new_kmer >> 8 & 0x00FF00FF00FF00FF) | (new_kmer & 0x00FF00FF00FF00FF) << 8;
    new_kmer = (new_kmer >> 16 & 0x0000FFFF0000FFFF) | (new_kmer & 0x0000FFFF0000FFFF) << 16;
    new_kmer = (new_kmer >> 32 & 0x00000000FFFFFFFF) | (new_kmer & 0x00000000FFFFFFFF) << 32;
    // complement it
    new_kmer ^= 0xFFFFFFFFFFFFFFFF;
    // shift it to the right size
    new_kmer = new_kmer >> (2 * (32 - k));
    new_kmer
}

#[test]
fn test_reverse_complement() {
  assert_eq!(reverse_complement(0b000000, 3), 0b111111);
  assert_eq!(reverse_complement(0b111111, 3), 0b000000);
  assert_eq!(reverse_complement(0b00000000, 4), 0b11111111);
  assert_eq!(reverse_complement(0b00011011, 4), 0b00011011);
}

fn canonical(kmer: BitKmer, k: u8) -> BitKmer {
    let rc = reverse_complement(kmer, k);
    if kmer > rc {
        rc
    } else {
        kmer
    }
}

fn print_ikmer(kmer: BitKmer, k: u8) -> String {
    let mut new_kmer = kmer;
    let mut new_kmer_str = String::new();
    let offset = (k - 1) * 2;
    let bitmask = BitKmer::pow(2, (2 * k - 1) as u32) + BitKmer::pow(2, (2 * k - 2) as u32);

    for i in 0..k {
        let new_char = (new_kmer & bitmask) >> offset;
        new_kmer <<= 2;
        new_kmer_str.push(match new_char {
            0 => 'A',
            1 => 'C',
            2 => 'G',
            3 => 'T',
            _ => panic!("Mathematical impossibility"),
        });
    }
    new_kmer_str
}

#[test]
fn test_print_ikmer() {
    assert_eq!(print_ikmer(1 as BitKmer, 1), String::from("C"));
    assert_eq!(print_ikmer(60 as BitKmer, 3), String::from("TTA"));
    assert_eq!(print_ikmer(0 as BitKmer, 3), String::from("AAA"));
}
<commit_msg>Fix off-by-one error in bitkmers.rs<commit_after>pub type BitKmer = u64;

pub struct BitKmerIter<'a> {
    k: u8,
    pos: usize,
    cur_kmer: BitKmer,
    buffer: &'a [u8],
}

pub fn fast_extend_kmer(kmer: &BitKmer, k: &u8, new_char: &u8) -> BitKmer {
    // new_char must be in ACGTE (E gets counted as T unfortunately)
    assert!(new_char & 223 == 84 || new_char & 217 == 65);
    let new_kmer = (kmer << 2) + (((new_char >> 2) ^ (new_char >> 1)) & 3u8) as BitKmer;

    // mask out any overflowed bits
    new_kmer & (BitKmer::pow(2, (2 * k) as u32) - 1) as BitKmer
}

pub fn extend_kmer(kmer: &BitKmer, k: &u8, new_char: &u8) -> BitKmer {
    let new_kmer = (kmer << 2) + match new_char {
        &b'A' | &b'a' => 0 as BitKmer,
        &b'C' | &b'c' => 1 as BitKmer,
        &b'G' | &b'g' => 2 as BitKmer,
        &b'T' | &b't' => 3 as BitKmer,
        _ => panic!("Received a non ACGT nucleotide"),
    };

    // mask out any overflowed bits
    new_kmer & (BitKmer::pow(2, (2 * k) as u32) - 1) as BitKmer
}

impl<'a> BitKmerIter<'a> {
    pub fn new(slice: &'a [u8], k: u8) -> BitKmerIter<'a> {
        // TODO: raise an error if the slice is shorter than k
        let mut kmer = 0u64;
        if slice.len() >= k as usize {
            for i in 0..k - 1 {
                kmer = fast_extend_kmer(&kmer, &k, &slice[i as usize]);
            }
        }

        BitKmerIter {
            k: k,
            pos: (k - 2) as usize,
            cur_kmer: kmer,
            buffer: slice,
        }
    }
}

impl<'a> Iterator for BitKmerIter<'a> {
    type Item = BitKmer;

    fn next(&mut self) -> Option<BitKmer> {
        self.pos += 1;
        if self.pos >= self.buffer.len() {
            return None;
        }
        self.cur_kmer = fast_extend_kmer(&self.cur_kmer, &self.k, &self.buffer[self.pos]);

        Some(self.cur_kmer)
    }
}

#[test]
fn test_iterator() {
    let seq = "ACGTA".as_bytes();
    let mut kmer_iter = BitKmerIter::new(seq, 3);
    assert_eq!(kmer_iter.next(), Some(6));
    assert_eq!(kmer_iter.next(), Some(27));
    assert_eq!(kmer_iter.next(), Some(44));
    assert_eq!(kmer_iter.next(), None);
}

fn reverse_complement(kmer: BitKmer, k: u8) -> BitKmer {
    // FIXME: this is not going to work with BitKmers of u128 or u32
    // inspired from https://www.biostars.org/p/113640/
    let mut new_kmer = kmer;
    // reverse it
    new_kmer = (new_kmer >> 2 & 0x3333333333333333) | (new_kmer & 0x3333333333333333) << 2;
    new_kmer = (new_kmer >> 4 & 0x0F0F0F0F0F0F0F0F) | (new_kmer & 0x0F0F0F0F0F0F0F0F) << 4;
    new_kmer = (new_kmer >> 8 & 0x00FF00FF00FF00FF) | (new_kmer & 0x00FF00FF00FF00FF) << 8;
    new_kmer = (new_kmer >> 16 & 0x0000FFFF0000FFFF) | (new_kmer & 0x0000FFFF0000FFFF) << 16;
    new_kmer = (new_kmer >> 32 & 0x00000000FFFFFFFF) | (new_kmer & 0x00000000FFFFFFFF) << 32;
    // complement it
    new_kmer ^= 0xFFFFFFFFFFFFFFFF;
    // shift it to the right size
    new_kmer = new_kmer >> (2 * (32 - k));
    new_kmer
}

#[test]
fn test_reverse_complement() {
  assert_eq!(reverse_complement(0b000000, 3), 0b111111);
  assert_eq!(reverse_complement(0b111111, 3), 0b000000);
  assert_eq!(reverse_complement(0b00000000, 4), 0b11111111);
  assert_eq!(reverse_complement(0b00011011, 4), 0b00011011);
}

fn canonical(kmer: BitKmer, k: u8) -> BitKmer {
    let rc = reverse_complement(kmer, k);
    if kmer > rc {
        rc
    } else {
        kmer
    }
}

fn print_ikmer(kmer: BitKmer, k: u8) -> String {
    let mut new_kmer = kmer;
    let mut new_kmer_str = String::new();
    let offset = (k - 1) * 2;
    let bitmask = BitKmer::pow(2, (2 * k - 1) as u32) + BitKmer::pow(2, (2 * k - 2) as u32);

    for i in 0..k {
        let new_char = (new_kmer & bitmask) >> offset;
        new_kmer <<= 2;
        new_kmer_str.push(match new_char {
            0 => 'A',
            1 => 'C',
            2 => 'G',
            3 => 'T',
            _ => panic!("Mathematical impossibility"),
        });
    }
    new_kmer_str
}

#[test]
fn test_print_ikmer() {
    assert_eq!(print_ikmer(1 as BitKmer, 1), String::from("C"));
    assert_eq!(print_ikmer(60 as BitKmer, 3), String::from("TTA"));
    assert_eq!(print_ikmer(0 as BitKmer, 3), String::from("AAA"));
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! azure metadata fetcher

mod crypto;

use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};

use openssh_keys::PublicKey;
use reqwest::header::{HeaderName, HeaderValue};
use update_ssh_keys::AuthorizedKeyEntry;

use self::crypto::x509;
use errors::*;
use network;
use providers::MetadataProvider;
use retry;
use util;

static HDR_AGENT_NAME: &str = "x-ms-agent-name";
static HDR_VERSION: &str = "x-ms-version";
static HDR_CIPHER_NAME: &str = "x-ms-cipher-name";
static HDR_CERT: &str = "x-ms-guest-agent-public-x509-cert";

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

#[derive(Debug, Deserialize, Clone, Default)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize, Clone, Default)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize, Clone, Default)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance", default)]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize, Clone)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize, Clone)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String,
    #[serde(rename = "SharedConfig", default)]
    pub shared_config: String,
}

#[derive(Debug, Deserialize, Clone)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize, Clone)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize, Clone)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

#[derive(Debug, Deserialize, Clone)]
struct SharedConfig {
    #[serde(rename = "Incarnation")]
    pub incarnation: Incarnation,
    #[serde(rename = "Instances")]
    pub instances: Instances,
}

#[derive(Debug, Deserialize, Clone)]
struct Incarnation {
    pub instance: String,
}

#[derive(Debug, Deserialize, Clone)]
struct Instances {
    #[serde(rename = "Instance", default)]
    pub instances: Vec<Instance>,
}

#[derive(Debug, Deserialize, Clone)]
struct Instance {
    pub id: String,
    pub address: String,
    #[serde(rename = "InputEndpoints")]
    pub input_endpoints: InputEndpoints,
}

#[derive(Debug, Deserialize, Clone)]
struct InputEndpoints {
    #[serde(rename = "Endpoint", default)]
    pub endpoints: Vec<Endpoint>,
}

#[derive(Debug, Deserialize, Clone)]
struct Endpoint {
    #[serde(rename = "loadBalancedPublicAddress", default)]
    pub load_balanced_public_address: String,
}

#[derive(Debug, Copy, Clone, Default)]
struct Attributes {
    pub virtual_ipv4: Option<IpAddr>,
    pub dynamic_ipv4: Option<IpAddr>,
}

#[derive(Debug, Clone)]
pub struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
    goal_state: GoalState,
}

impl Azure {
    pub fn try_new() -> Result<Azure> {
        let addr = Azure::get_fabric_address()
            .chain_err(|| "failed to get fabric address")?;
        let client = retry::Client::try_new()?
            .header(HeaderName::from_static(HDR_AGENT_NAME),
                    HeaderValue::from_static(MS_AGENT_NAME))
            .header(HeaderName::from_static(HDR_VERSION),
                    HeaderValue::from_static(MS_VERSION));

        let mut azure = Azure {
            client,
            endpoint: addr,
            goal_state: GoalState::default(),
        };

        // make sure the metadata service is compatible with our version
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;

        // populate goalstate
        azure.goal_state = azure.get_goal_state()?;
        Ok(azure)
    }

    fn get_goal_state(&self) -> Result<GoalState> {
        self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
        .ok_or_else(|| "failed to get goal state: not found response".into())
    }

    fn get_fabric_address() -> Result<IpAddr> {
        let v = util::dns_lease_key_lookup(OPTION_245)?;
        // value is an 8 digit hex value. convert it to u32 and
        // then parse that into an ip. Ipv4Addr::from(u32)
        // performs conversion from big-endian
        trace!("found fabric address in hex - {:?}", v);
        let dec = u32::from_str_radix(&v, 16)
            .chain_err(|| format!("failed to convert '{}' from hex", v))?;
        Ok(IpAddr::V4(dec.into()))
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &self.goal_state.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs<S: AsRef<str>>(&self, mangled_pem: S) -> Result<String> {
        // get the certificates
        let endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(HeaderName::from_static(HDR_CIPHER_NAME),
                    HeaderValue::from_static("DES_EDE3_CBC"))
            .header(HeaderName::from_static(HDR_CERT),
                    HeaderValue::from_str(mangled_pem.as_ref())?)
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }

    fn get_attributes(&self) -> Result<Attributes> {
        let endpoint = &self.goal_state.container.role_instance_list.role_instances[0].configuration.shared_config;

        let shared_config: SharedConfig = self.client.get(retry::Xml, endpoint.to_string()).send()
            .chain_err(|| "failed to get shared configuration")?
            .ok_or_else(|| "failed to get shared configuration: not found")?;

        let mut attributes = Attributes::default();

        for instance in shared_config.instances.instances {
            if instance.id == shared_config.incarnation.instance {
                attributes.dynamic_ipv4 = Some(instance.address.parse()
                    .chain_err(|| format!("failed to parse instance ip address: {}", instance.address))?);
                for endpoint in instance.input_endpoints.endpoints {
                    attributes.virtual_ipv4 = match endpoint.load_balanced_public_address.parse::<SocketAddr>() {
                        Ok(lbpa) => Some(lbpa.ip()),
                        Err(_) => continue,
                    };
                }
            }
        }

        Ok(attributes)
    }
}

impl MetadataProvider for Azure {
    fn attributes(&self) -> Result<HashMap<String, String>> {
        let attributes = self.get_attributes()?;
        let mut out = HashMap::with_capacity(2);

        if let Some(virtual_ipv4) = attributes.virtual_ipv4 {
            out.insert("AZURE_IPV4_VIRTUAL".to_string(), virtual_ipv4.to_string());
        }

        if let Some(dynamic_ipv4) = attributes.dynamic_ipv4 {
            out.insert("AZURE_IPV4_DYNAMIC".to_string(), dynamic_ipv4.to_string());
        }

        Ok(out)
    }

    fn hostname(&self) -> Result<Option<String>> {
        Ok(None)
    }

    fn ssh_keys(&self) -> Result<Vec<AuthorizedKeyEntry>> {
        let key = self.get_ssh_pubkey()?;
        Ok(vec![AuthorizedKeyEntry::Valid{key}])
    }

    fn networks(&self) -> Result<Vec<network::Interface>> {
        Ok(vec![])
    }

    fn network_devices(&self) -> Result<Vec<network::Device>> {
        Ok(vec![])
    }
}
<commit_msg>azure: hardcode fallback for wireserver endpoint<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! azure metadata fetcher

mod crypto;

use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};

use openssh_keys::PublicKey;
use reqwest::header::{HeaderName, HeaderValue};
use update_ssh_keys::AuthorizedKeyEntry;

use self::crypto::x509;
use errors::*;
use network;
use providers::MetadataProvider;
use retry;
use util;

static HDR_AGENT_NAME: &str = "x-ms-agent-name";
static HDR_VERSION: &str = "x-ms-version";
static HDR_CIPHER_NAME: &str = "x-ms-cipher-name";
static HDR_CERT: &str = "x-ms-guest-agent-public-x509-cert";

const OPTION_245: &str = "OPTION_245";
const MS_AGENT_NAME: &str = "com.coreos.metadata";
const MS_VERSION: &str = "2012-11-30";
const SMIME_HEADER: &str = "\
MIME-Version:1.0
Content-Disposition: attachment; filename=/home/core/encrypted-ssh-cert.pem
Content-Type: application/x-pkcs7-mime; name=/home/core/encrypted-ssh-cert.pem
Content-Transfer-Encoding: base64

";

/// This is a known working wireserver endpoint within Azure.
/// See: https://blogs.msdn.microsoft.com/mast/2015/05/18/what-is-the-ip-address-168-63-129-16/
const FALLBACK_WIRESERVER_ADDR: [u8; 4] = [168, 63, 129, 16]; // for grep: 168.63.129.16

#[derive(Debug, Deserialize, Clone, Default)]
struct GoalState {
    #[serde(rename = "Container")]
    pub container: Container
}

#[derive(Debug, Deserialize, Clone, Default)]
struct Container {
    #[serde(rename = "RoleInstanceList")]
    pub role_instance_list: RoleInstanceList
}

#[derive(Debug, Deserialize, Clone, Default)]
struct RoleInstanceList {
    #[serde(rename = "RoleInstance", default)]
    pub role_instances: Vec<RoleInstance>
}

#[derive(Debug, Deserialize, Clone)]
struct RoleInstance {
    #[serde(rename = "Configuration")]
    pub configuration: Configuration
}

#[derive(Debug, Deserialize, Clone)]
struct Configuration {
    #[serde(rename = "Certificates", default)]
    pub certificates: String,
    #[serde(rename = "SharedConfig", default)]
    pub shared_config: String,
}

#[derive(Debug, Deserialize, Clone)]
struct CertificatesFile {
    #[serde(rename = "Data", default)]
    pub data: String
}

#[derive(Debug, Deserialize, Clone)]
struct Versions {
    #[serde(rename = "Supported")]
    pub supported: Supported
}

#[derive(Debug, Deserialize, Clone)]
struct Supported {
    #[serde(rename = "Version", default)]
    pub versions: Vec<String>
}

#[derive(Debug, Deserialize, Clone)]
struct SharedConfig {
    #[serde(rename = "Incarnation")]
    pub incarnation: Incarnation,
    #[serde(rename = "Instances")]
    pub instances: Instances,
}

#[derive(Debug, Deserialize, Clone)]
struct Incarnation {
    pub instance: String,
}

#[derive(Debug, Deserialize, Clone)]
struct Instances {
    #[serde(rename = "Instance", default)]
    pub instances: Vec<Instance>,
}

#[derive(Debug, Deserialize, Clone)]
struct Instance {
    pub id: String,
    pub address: String,
    #[serde(rename = "InputEndpoints")]
    pub input_endpoints: InputEndpoints,
}

#[derive(Debug, Deserialize, Clone)]
struct InputEndpoints {
    #[serde(rename = "Endpoint", default)]
    pub endpoints: Vec<Endpoint>,
}

#[derive(Debug, Deserialize, Clone)]
struct Endpoint {
    #[serde(rename = "loadBalancedPublicAddress", default)]
    pub load_balanced_public_address: String,
}

#[derive(Debug, Copy, Clone, Default)]
struct Attributes {
    pub virtual_ipv4: Option<IpAddr>,
    pub dynamic_ipv4: Option<IpAddr>,
}

#[derive(Debug, Clone)]
pub struct Azure {
    client: retry::Client,
    endpoint: IpAddr,
    goal_state: GoalState,
}

impl Azure {
    pub fn try_new() -> Result<Azure> {
        let addr = Azure::get_fabric_address();
        let client = retry::Client::try_new()?
            .header(HeaderName::from_static(HDR_AGENT_NAME),
                    HeaderValue::from_static(MS_AGENT_NAME))
            .header(HeaderName::from_static(HDR_VERSION),
                    HeaderValue::from_static(MS_VERSION));

        let mut azure = Azure {
            client,
            endpoint: addr,
            goal_state: GoalState::default(),
        };

        // make sure the metadata service is compatible with our version
        azure.is_fabric_compatible(MS_VERSION)
            .chain_err(|| "failed version compatibility check")?;

        // populate goalstate
        azure.goal_state = azure.get_goal_state()?;
        Ok(azure)
    }

    fn get_goal_state(&self) -> Result<GoalState> {
        self.client.get(retry::Xml, format!("http://{}/machine/?comp=goalstate", self.endpoint)).send()
            .chain_err(|| "failed to get goal state")?
        .ok_or_else(|| "failed to get goal state: not found response".into())
    }

    fn get_fabric_address() -> IpAddr {
        // try to fetch from dhcp, else use fallback; this is similar to what WALinuxAgent does
        Azure::get_fabric_address_from_dhcp().unwrap_or_else(|e| {
            warn!("Failed to get fabric address from DHCP: {}", e);
            info!("Using fallback address");
            IpAddr::from(FALLBACK_WIRESERVER_ADDR)
        })
    }

    fn get_fabric_address_from_dhcp() -> Result<IpAddr> {
        let v = util::dns_lease_key_lookup(OPTION_245)?;
        // value is an 8 digit hex value. convert it to u32 and
        // then parse that into an ip. Ipv4Addr::from(u32)
        // performs conversion from big-endian
        trace!("found fabric address in hex - {:?}", v);
        let dec = u32::from_str_radix(&v, 16)
            .chain_err(|| format!("failed to convert '{}' from hex", v))?;
        Ok(IpAddr::V4(dec.into()))
    }

    fn is_fabric_compatible(&self, version: &str) -> Result<()> {
        let versions: Versions = self.client.get(retry::Xml, format!("http://{}/?comp=versions", self.endpoint)).send()
            .chain_err(|| "failed to get versions")?
            .ok_or_else(|| "failed to get versions: not found")?;

        if versions.supported.versions.iter().any(|v| v == version) {
            Ok(())
        } else {
            Err(format!("fabric version {} not compatible with fabric address {}", MS_VERSION, self.endpoint).into())
        }
    }

    fn get_certs_endpoint(&self) -> Result<String> {
        // grab the certificates endpoint from the xml and return it
        let cert_endpoint: &str = &self.goal_state.container.role_instance_list.role_instances[0].configuration.certificates;
        Ok(String::from(cert_endpoint))
    }

    fn get_certs<S: AsRef<str>>(&self, mangled_pem: S) -> Result<String> {
        // get the certificates
        let endpoint = self.get_certs_endpoint()
            .chain_err(|| "failed to get certs endpoint")?;

        let certs: CertificatesFile = self.client.get(retry::Xml, endpoint)
            .header(HeaderName::from_static(HDR_CIPHER_NAME),
                    HeaderValue::from_static("DES_EDE3_CBC"))
            .header(HeaderName::from_static(HDR_CERT),
                    HeaderValue::from_str(mangled_pem.as_ref())?)
            .send()
            .chain_err(|| "failed to get certificates")?
            .ok_or_else(|| "failed to get certificates: not found")?;

        // the cms decryption expects it to have MIME information on the top
        // since cms is really for email attachments...don't tell the cops.
        let mut smime = String::from(SMIME_HEADER);
        smime.push_str(&certs.data);

        Ok(smime)
    }

    // put it all together
    fn get_ssh_pubkey(&self) -> Result<PublicKey> {
        // first we have to get the certificates endoint.
        // we have to generate the rsa public/private keypair and the x509 cert
        // that we use to make the request. this is equivalent to
        // `openssl req -x509 -nodes -subj /CN=LinuxTransport -days 365 -newkey rsa:2048 -keyout private.pem -out cert.pem`
        let (x509, pkey) = x509::generate_cert(&x509::Config::new(2048, 365))
            .chain_err(|| "failed to generate keys")?;

        // mangle the pem file for the request
        let mangled_pem = crypto::mangle_pem(&x509)
            .chain_err(|| "failed to mangle pem")?;

        // fetch the encrypted cms blob from the certs endpoint
        let smime = self.get_certs(mangled_pem)
            .chain_err(|| "failed to get certs")?;

        // decrypt the cms blob
        let p12 = crypto::decrypt_cms(smime.as_bytes(), &pkey, &x509)
            .chain_err(|| "failed to decrypt cms blob")?;

        // convert that to the OpenSSH public key format
        let ssh_pubkey = crypto::p12_to_ssh_pubkey(&p12)
            .chain_err(|| "failed to convert pkcs12 blob to ssh pubkey")?;

        Ok(ssh_pubkey)
    }

    fn get_attributes(&self) -> Result<Attributes> {
        let endpoint = &self.goal_state.container.role_instance_list.role_instances[0].configuration.shared_config;

        let shared_config: SharedConfig = self.client.get(retry::Xml, endpoint.to_string()).send()
            .chain_err(|| "failed to get shared configuration")?
            .ok_or_else(|| "failed to get shared configuration: not found")?;

        let mut attributes = Attributes::default();

        for instance in shared_config.instances.instances {
            if instance.id == shared_config.incarnation.instance {
                attributes.dynamic_ipv4 = Some(instance.address.parse()
                    .chain_err(|| format!("failed to parse instance ip address: {}", instance.address))?);
                for endpoint in instance.input_endpoints.endpoints {
                    attributes.virtual_ipv4 = match endpoint.load_balanced_public_address.parse::<SocketAddr>() {
                        Ok(lbpa) => Some(lbpa.ip()),
                        Err(_) => continue,
                    };
                }
            }
        }

        Ok(attributes)
    }
}

impl MetadataProvider for Azure {
    fn attributes(&self) -> Result<HashMap<String, String>> {
        let attributes = self.get_attributes()?;
        let mut out = HashMap::with_capacity(2);

        if let Some(virtual_ipv4) = attributes.virtual_ipv4 {
            out.insert("AZURE_IPV4_VIRTUAL".to_string(), virtual_ipv4.to_string());
        }

        if let Some(dynamic_ipv4) = attributes.dynamic_ipv4 {
            out.insert("AZURE_IPV4_DYNAMIC".to_string(), dynamic_ipv4.to_string());
        }

        Ok(out)
    }

    fn hostname(&self) -> Result<Option<String>> {
        Ok(None)
    }

    fn ssh_keys(&self) -> Result<Vec<AuthorizedKeyEntry>> {
        let key = self.get_ssh_pubkey()?;
        Ok(vec![AuthorizedKeyEntry::Valid{key}])
    }

    fn networks(&self) -> Result<Vec<network::Interface>> {
        Ok(vec![])
    }

    fn network_devices(&self) -> Result<Vec<network::Device>> {
        Ok(vec![])
    }
}
<|endoftext|>


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

//! OS-based thread local storage
//!
//! This module provides an implementation of OS-based thread local storage,
//! using the native OS-provided facilities (think `TlsAlloc` or
//! `pthread_setspecific`). The interface of this differs from the other types
//! of thread-local-storage provided in this crate in that OS-based TLS can only
//! get/set pointers,
//!
//! This module also provides two flavors of TLS. One is intended for static
//! initialization, and does not contain a `Drop` implementation to deallocate
//! the OS-TLS key. The other is a type which does implement `Drop` and hence
//! has a safe interface.
//!
//! # Usage
//!
//! This module should likely not be used directly unless other primitives are
//! being built on. types such as `thread_local::scoped::Key` are likely much
//! more useful in practice than this OS-based version which likely requires
//! unsafe code to interoperate with.
//!
//! # Example
//!
//! Using a dynamically allocated TLS key. Note that this key can be shared
//! among many threads via an `Arc`.
//!
//! ```rust,ignore
//! let key = Key::new(None);
//! assert!(key.get().is_null());
//! key.set(1 as *mut u8);
//! assert!(!key.get().is_null());
//!
//! drop(key); // deallocate this TLS slot.
//! ```
//!
//! Sometimes a statically allocated key is either required or easier to work
//! with, however.
//!
//! ```rust,ignore
//! static KEY: StaticKey = INIT;
//!
//! unsafe {
//!     assert!(KEY.get().is_null());
//!     KEY.set(1 as *mut u8);
//! }
//! ```

#![allow(non_camel_case_types)]

use prelude::*;

use kinds::marker;
use mem;
use rustrt::exclusive::Exclusive;
use rustrt;
use sync::atomic::{mod, AtomicUint};
use sync::{Once, ONCE_INIT};

use sys::thread_local as imp;

/// A type for TLS keys that are statically allocated.
///
/// This type is entirely `unsafe` to use as it does not protect against
/// use-after-deallocation or use-during-deallocation.
///
/// The actual OS-TLS key is lazily allocated when this is used for the first
/// time. The key is also deallocated when the Rust runtime exits or `destroy`
/// is called, whichever comes first.
///
/// # Example
///
/// ```ignore
/// use tls::os::{StaticKey, INIT};
///
/// static KEY: StaticKey = INIT;
///
/// unsafe {
///     assert!(KEY.get().is_null());
///     KEY.set(1 as *mut u8);
/// }
/// ```
pub struct StaticKey {
    /// Inner static TLS key (internals), created with by `INIT_INNER` in this
    /// module.
    pub inner: StaticKeyInner,
    /// Destructor for the TLS value.
    ///
    /// See `Key::new` for information about when the destructor runs and how
    /// it runs.
    pub dtor: Option<unsafe extern fn(*mut u8)>,
}

/// Inner contents of `StaticKey`, created by the `INIT_INNER` constant.
pub struct StaticKeyInner {
    key: AtomicUint,
    nc: marker::NoCopy,
}

/// A type for a safely managed OS-based TLS slot.
///
/// This type allocates an OS TLS key when it is initialized and will deallocate
/// the key when it falls out of scope. When compared with `StaticKey`, this
/// type is entirely safe to use.
///
/// Implementations will likely, however, contain unsafe code as this type only
/// operates on `*mut u8`, an unsafe pointer.
///
/// # Example
///
/// ```rust,ignore
/// use tls::os::Key;
///
/// let key = Key::new(None);
/// assert!(key.get().is_null());
/// key.set(1 as *mut u8);
/// assert!(!key.get().is_null());
///
/// drop(key); // deallocate this TLS slot.
/// ```
pub struct Key {
    key: imp::Key,
}

/// Constant initialization value for static TLS keys.
///
/// This value specifies no destructor by default.
pub const INIT: StaticKey = StaticKey {
    inner: INIT_INNER,
    dtor: None,
};

/// Constant initialization value for the inner part of static TLS keys.
///
/// This value allows specific configuration of the destructor for a TLS key.
pub const INIT_INNER: StaticKeyInner = StaticKeyInner {
    key: atomic::INIT_ATOMIC_UINT,
    nc: marker::NoCopy,
};

static INIT_KEYS: Once = ONCE_INIT;
static mut KEYS: *mut Exclusive<Vec<imp::Key>> = 0 as *mut _;

impl StaticKey {
    /// Gets the value associated with this TLS key
    ///
    /// This will lazily allocate a TLS key from the OS if one has not already
    /// been allocated.
    #[inline]
    pub unsafe fn get(&self) -> *mut u8 { imp::get(self.key()) }

    /// Sets this TLS key to a new value.
    ///
    /// This will lazily allocate a TLS key from the OS if one has not already
    /// been allocated.
    #[inline]
    pub unsafe fn set(&self, val: *mut u8) { imp::set(self.key(), val) }

    /// Deallocates this OS TLS key.
    ///
    /// This function is unsafe as there is no guarantee that the key is not
    /// currently in use by other threads or will not ever be used again.
    ///
    /// Note that this does *not* run the user-provided destructor if one was
    /// specified at definition time. Doing so must be done manually.
    pub unsafe fn destroy(&self) {
        match self.inner.key.swap(0, atomic::SeqCst) {
            0 => {}
            n => { unregister_key(n as imp::Key); imp::destroy(n as imp::Key) }
        }
    }

    #[inline]
    unsafe fn key(&self) -> imp::Key {
        match self.inner.key.load(atomic::Relaxed) {
            0 => self.lazy_init() as imp::Key,
            n => n as imp::Key
        }
    }

    unsafe fn lazy_init(&self) -> uint {
        let key = imp::create(self.dtor);
        assert!(key != 0);
        match self.inner.key.compare_and_swap(0, key as uint, atomic::SeqCst) {
            // The CAS succeeded, so we've created the actual key
            0 => {
                register_key(key);
                key as uint
            }
            // If someone beat us to the punch, use their key instead
            n => { imp::destroy(key); n }
        }
    }
}

impl Key {
    /// Create a new managed OS TLS key.
    ///
    /// This key will be deallocated when the key falls out of scope.
    ///
    /// The argument provided is an optionally-specified destructor for the
    /// value of this TLS key. When a thread exits and the value for this key
    /// is non-null the destructor will be invoked. The TLS value will be reset
    /// to null before the destructor is invoked.
    ///
    /// Note that the destructor will not be run when the `Key` goes out of
    /// scope.
    #[inline]
    pub fn new(dtor: Option<unsafe extern fn(*mut u8)>) -> Key {
        Key { key: unsafe { imp::create(dtor) } }
    }

    /// See StaticKey::get
    #[inline]
    pub fn get(&self) -> *mut u8 {
        unsafe { imp::get(self.key) }
    }

    /// See StaticKey::set
    #[inline]
    pub fn set(&self, val: *mut u8) {
        unsafe { imp::set(self.key, val) }
    }
}

impl Drop for Key {
    fn drop(&mut self) {
        unsafe { imp::destroy(self.key) }
    }
}

fn init_keys() {
    let keys = box Exclusive::new(Vec::<imp::Key>::new());
    unsafe {
        KEYS = mem::transmute(keys);
    }

    rustrt::at_exit(proc() unsafe {
        let keys: Box<Exclusive<Vec<imp::Key>>> = mem::transmute(KEYS);
        KEYS = 0 as *mut _;
        let keys = keys.lock();
        for key in keys.iter() {
            imp::destroy(*key);
        }
    });
}

fn register_key(key: imp::Key) {
    INIT_KEYS.doit(init_keys);
    let mut keys = unsafe { (*KEYS).lock() };
    keys.push(key);
}

fn unregister_key(key: imp::Key) {
    INIT_KEYS.doit(init_keys);
    let mut keys = unsafe { (*KEYS).lock() };
    keys.retain(|k| *k != key);
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use super::{Key, StaticKey, INIT_INNER};

    fn assert_sync<T: Sync>() {}
    fn assert_send<T: Send>() {}

    #[test]
    fn smoke() {
        assert_sync::<Key>();
        assert_send::<Key>();

        let k1 = Key::new(None);
        let k2 = Key::new(None);
        assert!(k1.get().is_null());
        assert!(k2.get().is_null());
        k1.set(1 as *mut _);
        k2.set(2 as *mut _);
        assert_eq!(k1.get() as uint, 1);
        assert_eq!(k2.get() as uint, 2);
    }

    #[test]
    fn statik() {
        static K1: StaticKey = StaticKey { inner: INIT_INNER, dtor: None };
        static K2: StaticKey = StaticKey { inner: INIT_INNER, dtor: None };

        unsafe {
            assert!(K1.get().is_null());
            assert!(K2.get().is_null());
            K1.set(1 as *mut _);
            K2.set(2 as *mut _);
            assert_eq!(K1.get() as uint, 1);
            assert_eq!(K2.get() as uint, 2);
        }
    }
}

<commit_msg>auto merge of #19285 : alexcrichton/rust/issue-19280, r=aturon<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! OS-based thread local storage
//!
//! This module provides an implementation of OS-based thread local storage,
//! using the native OS-provided facilities (think `TlsAlloc` or
//! `pthread_setspecific`). The interface of this differs from the other types
//! of thread-local-storage provided in this crate in that OS-based TLS can only
//! get/set pointers,
//!
//! This module also provides two flavors of TLS. One is intended for static
//! initialization, and does not contain a `Drop` implementation to deallocate
//! the OS-TLS key. The other is a type which does implement `Drop` and hence
//! has a safe interface.
//!
//! # Usage
//!
//! This module should likely not be used directly unless other primitives are
//! being built on. types such as `thread_local::scoped::Key` are likely much
//! more useful in practice than this OS-based version which likely requires
//! unsafe code to interoperate with.
//!
//! # Example
//!
//! Using a dynamically allocated TLS key. Note that this key can be shared
//! among many threads via an `Arc`.
//!
//! ```rust,ignore
//! let key = Key::new(None);
//! assert!(key.get().is_null());
//! key.set(1 as *mut u8);
//! assert!(!key.get().is_null());
//!
//! drop(key); // deallocate this TLS slot.
//! ```
//!
//! Sometimes a statically allocated key is either required or easier to work
//! with, however.
//!
//! ```rust,ignore
//! static KEY: StaticKey = INIT;
//!
//! unsafe {
//!     assert!(KEY.get().is_null());
//!     KEY.set(1 as *mut u8);
//! }
//! ```

#![allow(non_camel_case_types)]

use prelude::*;

use kinds::marker;
use rustrt::exclusive::Exclusive;
use sync::atomic::{mod, AtomicUint};
use sync::{Once, ONCE_INIT};

use sys::thread_local as imp;

/// A type for TLS keys that are statically allocated.
///
/// This type is entirely `unsafe` to use as it does not protect against
/// use-after-deallocation or use-during-deallocation.
///
/// The actual OS-TLS key is lazily allocated when this is used for the first
/// time. The key is also deallocated when the Rust runtime exits or `destroy`
/// is called, whichever comes first.
///
/// # Example
///
/// ```ignore
/// use tls::os::{StaticKey, INIT};
///
/// static KEY: StaticKey = INIT;
///
/// unsafe {
///     assert!(KEY.get().is_null());
///     KEY.set(1 as *mut u8);
/// }
/// ```
pub struct StaticKey {
    /// Inner static TLS key (internals), created with by `INIT_INNER` in this
    /// module.
    pub inner: StaticKeyInner,
    /// Destructor for the TLS value.
    ///
    /// See `Key::new` for information about when the destructor runs and how
    /// it runs.
    pub dtor: Option<unsafe extern fn(*mut u8)>,
}

/// Inner contents of `StaticKey`, created by the `INIT_INNER` constant.
pub struct StaticKeyInner {
    key: AtomicUint,
    nc: marker::NoCopy,
}

/// A type for a safely managed OS-based TLS slot.
///
/// This type allocates an OS TLS key when it is initialized and will deallocate
/// the key when it falls out of scope. When compared with `StaticKey`, this
/// type is entirely safe to use.
///
/// Implementations will likely, however, contain unsafe code as this type only
/// operates on `*mut u8`, an unsafe pointer.
///
/// # Example
///
/// ```rust,ignore
/// use tls::os::Key;
///
/// let key = Key::new(None);
/// assert!(key.get().is_null());
/// key.set(1 as *mut u8);
/// assert!(!key.get().is_null());
///
/// drop(key); // deallocate this TLS slot.
/// ```
pub struct Key {
    key: imp::Key,
}

/// Constant initialization value for static TLS keys.
///
/// This value specifies no destructor by default.
pub const INIT: StaticKey = StaticKey {
    inner: INIT_INNER,
    dtor: None,
};

/// Constant initialization value for the inner part of static TLS keys.
///
/// This value allows specific configuration of the destructor for a TLS key.
pub const INIT_INNER: StaticKeyInner = StaticKeyInner {
    key: atomic::INIT_ATOMIC_UINT,
    nc: marker::NoCopy,
};

static INIT_KEYS: Once = ONCE_INIT;
static mut KEYS: *mut Exclusive<Vec<imp::Key>> = 0 as *mut _;

impl StaticKey {
    /// Gets the value associated with this TLS key
    ///
    /// This will lazily allocate a TLS key from the OS if one has not already
    /// been allocated.
    #[inline]
    pub unsafe fn get(&self) -> *mut u8 { imp::get(self.key()) }

    /// Sets this TLS key to a new value.
    ///
    /// This will lazily allocate a TLS key from the OS if one has not already
    /// been allocated.
    #[inline]
    pub unsafe fn set(&self, val: *mut u8) { imp::set(self.key(), val) }

    /// Deallocates this OS TLS key.
    ///
    /// This function is unsafe as there is no guarantee that the key is not
    /// currently in use by other threads or will not ever be used again.
    ///
    /// Note that this does *not* run the user-provided destructor if one was
    /// specified at definition time. Doing so must be done manually.
    pub unsafe fn destroy(&self) {
        match self.inner.key.swap(0, atomic::SeqCst) {
            0 => {}
            n => { imp::destroy(n as imp::Key) }
        }
    }

    #[inline]
    unsafe fn key(&self) -> imp::Key {
        match self.inner.key.load(atomic::Relaxed) {
            0 => self.lazy_init() as imp::Key,
            n => n as imp::Key
        }
    }

    unsafe fn lazy_init(&self) -> uint {
        let key = imp::create(self.dtor);
        assert!(key != 0);
        match self.inner.key.compare_and_swap(0, key as uint, atomic::SeqCst) {
            // The CAS succeeded, so we've created the actual key
            0 => key as uint,
            // If someone beat us to the punch, use their key instead
            n => { imp::destroy(key); n }
        }
    }
}

impl Key {
    /// Create a new managed OS TLS key.
    ///
    /// This key will be deallocated when the key falls out of scope.
    ///
    /// The argument provided is an optionally-specified destructor for the
    /// value of this TLS key. When a thread exits and the value for this key
    /// is non-null the destructor will be invoked. The TLS value will be reset
    /// to null before the destructor is invoked.
    ///
    /// Note that the destructor will not be run when the `Key` goes out of
    /// scope.
    #[inline]
    pub fn new(dtor: Option<unsafe extern fn(*mut u8)>) -> Key {
        Key { key: unsafe { imp::create(dtor) } }
    }

    /// See StaticKey::get
    #[inline]
    pub fn get(&self) -> *mut u8 {
        unsafe { imp::get(self.key) }
    }

    /// See StaticKey::set
    #[inline]
    pub fn set(&self, val: *mut u8) {
        unsafe { imp::set(self.key, val) }
    }
}

impl Drop for Key {
    fn drop(&mut self) {
        unsafe { imp::destroy(self.key) }
    }
}

#[cfg(test)]
mod tests {
    use prelude::*;
    use super::{Key, StaticKey, INIT_INNER};

    fn assert_sync<T: Sync>() {}
    fn assert_send<T: Send>() {}

    #[test]
    fn smoke() {
        assert_sync::<Key>();
        assert_send::<Key>();

        let k1 = Key::new(None);
        let k2 = Key::new(None);
        assert!(k1.get().is_null());
        assert!(k2.get().is_null());
        k1.set(1 as *mut _);
        k2.set(2 as *mut _);
        assert_eq!(k1.get() as uint, 1);
        assert_eq!(k2.get() as uint, 2);
    }

    #[test]
    fn statik() {
        static K1: StaticKey = StaticKey { inner: INIT_INNER, dtor: None };
        static K2: StaticKey = StaticKey { inner: INIT_INNER, dtor: None };

        unsafe {
            assert!(K1.get().is_null());
            assert!(K2.get().is_null());
            K1.set(1 as *mut _);
            K2.set(2 as *mut _);
            assert_eq!(K1.get() as uint, 1);
            assert_eq!(K2.get() as uint, 2);
        }
    }
}

<|endoftext|>


<commit_before>use {Action, InputTranslator, InputRebind, MouseTranslationData};
use input::Button;
use window::Size;
use std::convert::Into;
use std::default::Default;

/// Convenience object for constructing an InputMap.
pub struct RebindBuilder<A: Action> {
    input_remappings: Vec<(Button, A)>,
    mouse_data: MouseTranslationData
}

impl<A: Action> RebindBuilder<A> {
    /// Creates a new `RebindBuilder` with the specified viewport size.
    pub fn new<S: Into<Size> + Sized>(size: S) -> Self {
        RebindBuilder {
            input_remappings: vec![],
            mouse_data: MouseTranslationData::new(size)
        }
    }

    /// Set whether the x scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn x_scroll_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.x_axis_scroll_inverted = invert;
        self
    }

    /// Returns true if the x scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_scroll_inverted
    }

    /// Set whether the y scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn y_scroll_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.y_axis_scroll_inverted = invert;
        self
    }

    /// Returns true if the y scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_scroll_inverted
    }

    /// Set whether the x axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn x_motion_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.x_axis_motion_inverted = invert;
        self
    }

    /// Returns true if the x axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_motion_inverted
    }

    /// Set whether the y axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn y_motion_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.y_axis_motion_inverted = invert;
        self
    }

    /// Returns true if the y axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_motion_inverted
    }

    /// Set the mouse sensitivity
    pub fn mouse_sensitivity(mut self, sensitivity: f64) -> Self {
        self.mouse_data.sensitivity = sensitivity;
        self
    }

    /// Returns the mouse sensitivity
    pub fn get_mouse_sensitivity(&self) -> &f64 {
        &self.mouse_data.sensitivity
    }

    /// Sets the viewport size used for mouse position calculations.
    pub fn viewport_size(mut self, size: Size) -> Self {
        self.mouse_data.viewport_size = size;
        self
    }

    /// Returns the currently set viewport size.
    pub fn get_viewport_size(&self) -> &Size {
        &self.mouse_data.viewport_size
    }

    /// Add an association between the Button and Action in the built object.
    pub fn with_action_mapping(mut self, button: Button, action: A) -> Self {
        self.input_remappings.push((button, action));
        self
    }

    /// Creates an `InputTranslator` from this builder object. Consumes self.
    pub fn build_translator(self) -> InputTranslator<A> { self.into() }

    /// Creates an `InputRebind` from this builder object. Consumes self.
    pub fn build_rebind(self) -> InputRebind<A> { self.into() }
}

/// Creates a new `RebindBuilder`. The viewport size is set to (800, 600).
impl<A: Action> Default for RebindBuilder<A> {
    fn default() -> Self {
        Self::new((800, 600))
    }
}

impl<A: Action> Into<InputTranslator<A>> for RebindBuilder<A> {
    fn into(self) -> InputTranslator<A> {
        let mut input_map = InputTranslator::new(self.mouse_data.viewport_size);

        input_map.mouse_translator.data = self.mouse_data;
        input_map.keymap = self.input_remappings.iter().cloned().collect();

        input_map
    }
}

impl<A: Action> Into<InputRebind<A>> for RebindBuilder<A> {
    #[allow(dead_code, unused_variables, unreachable_code)]
    fn into(self) -> InputRebind<A> {
        unimplemented!();

        let mut input_rebind = InputRebind::new(self.mouse_data.viewport_size);

        input_rebind.mouse_data = self.mouse_data;
        //input_rebind.keymap.btn_map = self.input_remappings.iter().map(|x| x.clone()).collect();

        input_rebind
    }
}

<commit_msg>Complete implementation of `Into<InputRebind<A>> for RebindBuilder<A>`<commit_after>use {Action, ButtonTuple, InputTranslator, InputRebind, MouseTranslationData};
use input::Button;
use itertools::Itertools;
use window::Size;
use std::convert::Into;
use std::default::Default;

/// Convenience object for constructing an InputMap.
pub struct RebindBuilder<A: Action> {
    input_remappings: Vec<(Button, A)>,
    mouse_data: MouseTranslationData
}

impl<A: Action> RebindBuilder<A> {
    /// Creates a new `RebindBuilder` with the specified viewport size.
    pub fn new<S: Into<Size> + Sized>(size: S) -> Self {
        RebindBuilder {
            input_remappings: vec![],
            mouse_data: MouseTranslationData::new(size)
        }
    }

    /// Set whether the x scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn x_scroll_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.x_axis_scroll_inverted = invert;
        self
    }

    /// Returns true if the x scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn get_x_scroll_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_scroll_inverted
    }

    /// Set whether the y scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn y_scroll_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.y_axis_scroll_inverted = invert;
        self
    }

    /// Returns true if the y scroll is inverted on the builder (and thus
    /// on the built object).
    pub fn get_y_scroll_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_scroll_inverted
    }

    /// Set whether the x axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn x_motion_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.x_axis_motion_inverted = invert;
        self
    }

    /// Returns true if the x axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn get_x_motion_inverted(&self) -> &bool {
        &self.mouse_data.x_axis_motion_inverted
    }

    /// Set whether the y axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn y_motion_inverted(mut self, invert: bool) -> Self {
        self.mouse_data.y_axis_motion_inverted = invert;
        self
    }

    /// Returns true if the y axis motion is inverted on the builder (and thus
    /// on the built object).
    pub fn get_y_motion_inverted(&self) -> &bool {
        &self.mouse_data.y_axis_motion_inverted
    }

    /// Set the mouse sensitivity
    pub fn mouse_sensitivity(mut self, sensitivity: f64) -> Self {
        self.mouse_data.sensitivity = sensitivity;
        self
    }

    /// Returns the mouse sensitivity
    pub fn get_mouse_sensitivity(&self) -> &f64 {
        &self.mouse_data.sensitivity
    }

    /// Sets the viewport size used for mouse position calculations.
    pub fn viewport_size(mut self, size: Size) -> Self {
        self.mouse_data.viewport_size = size;
        self
    }

    /// Returns the currently set viewport size.
    pub fn get_viewport_size(&self) -> &Size {
        &self.mouse_data.viewport_size
    }

    /// Add an association between the Button and Action in the built object.
    pub fn with_action_mapping(mut self, button: Button, action: A) -> Self {
        self.input_remappings.push((button, action));
        self
    }

    /// Creates an `InputTranslator` from this builder object. Consumes self.
    pub fn build_translator(self) -> InputTranslator<A> { self.into() }

    /// Creates an `InputRebind` from this builder object. Consumes self.
    pub fn build_rebind(self) -> InputRebind<A> { self.into() }
}

/// Creates a new `RebindBuilder`. The viewport size is set to (800, 600).
impl<A: Action> Default for RebindBuilder<A> {
    fn default() -> Self {
        Self::new((800, 600))
    }
}

impl<A: Action> Into<InputTranslator<A>> for RebindBuilder<A> {
    fn into(self) -> InputTranslator<A> {
        let mut input_map = InputTranslator::new(self.mouse_data.viewport_size);

        input_map.mouse_translator.data = self.mouse_data;
        input_map.keymap = self.input_remappings.iter().cloned().collect();

        input_map
    }
}

impl<A: Action> Into<InputRebind<A>> for RebindBuilder<A> {
    fn into(self) -> InputRebind<A> {
        let mut input_rebind = InputRebind::new(self.mouse_data.viewport_size);

        input_rebind.mouse_data = self.mouse_data;
        input_rebind.keymap = self.input_remappings.iter()
                                                   .map(|&(b, a)| (a, vec![Some(b)]))
                                                   .sorted_by(|&(a0, _), &(a1, _)| Ord::cmp(&a0, &a1))
                                                   .into_iter()
                                                   .coalesce(|(a0, b0), (a1, b1)| if a0 == a1 {
                                                       Ok((a0, b0.into_iter().chain(b1).collect()))
                                                   } else {
                                                       Err(((a0, b0), (a1, b1)))
                                                   })
                                                   .map(|(a, bs)| {
                                                      let buttons = &bs.iter()
                                                                       .cloned()
                                                                       .pad_using(3, |_| None)
                                                                       .take(3)
                                                                       .collect_vec();

                                                       if buttons.len() >= 3 {
                                                            (a, ButtonTuple(buttons[0],
                                                                            buttons[1],
                                                                            buttons[2]))
                                                       } else {
                                                           unreachable!();
                                                       }
                                                   })
                                                   .collect();

        input_rebind
    }
}

<|endoftext|>


<commit_before>use std::{fs, io, path};

use syntax::{ast, visit, codemap};

use tree::{Tree, Visibility};

pub struct Config {
    pub include_orphans: bool,
    pub ignored_files: Vec<path::PathBuf>,
}

pub struct Builder<'a> {
    tree: Tree,
    path: Vec<String>,
    config: Config,
    codemap: &'a codemap::CodeMap,
}

impl<'a> Builder<'a> {
    pub fn new(config: Config, crate_name: String, codemap: &'a codemap::CodeMap) -> Self {
        let tree = Tree::new_crate(crate_name);
        Builder {
            tree: tree,
            path: vec![],
            config: config,
            codemap: codemap,
        }
    }

    pub fn tree(&self) -> &Tree {
        &self.tree
    }

    fn find_orphan_candidates(path: &[String],
                              ignored_paths: &[path::PathBuf])
                              -> Result<Vec<String>, io::Error> {
        let mut dir_path = "./src/".to_string();
        for name in path.iter() {
            dir_path.push_str(name);
            dir_path.push_str("/");
        }
        fn is_mod(entry: &fs::DirEntry) -> bool {
            let path = entry.path();
            if path.is_file() {
                if path.extension().map_or(false, |ext| ext == "rs") {
                    return file_name(entry)
                        .map_or(false, |s| (s != "mod") && (s != "lib") && (s != "main"));
                }
            } else if path.join("mod.rs").exists() {
                return true;
            }
            false
        }
        fn file_name(entry: &fs::DirEntry) -> Option<String> {
            entry.path().file_stem().map_or(None, |s| s.to_str().map(|s| s.to_string()))
        }
        let mut orphans = vec![];
        if try!(fs::metadata(&dir_path)).is_dir() {
            for entry in try!(fs::read_dir(&dir_path)).filter_map(|e| e.ok()).filter(is_mod) {
                if !ignored_paths.contains(&entry.path()) {
                    if let Some(name) = file_name(&entry) {
                        orphans.push(name);
                    }
                }
            }
        }
        Ok(orphans)
    }

    fn sanitize_condition(condition: String) -> String {
        let words: Vec<&str> = condition.split_whitespace().collect();
        words.join(" ")
    }
}

impl<'v, 'a> visit::Visitor<'v> for Builder<'a> {
    fn visit_item(&mut self, item: &ast::Item) {
        let condition = item.attrs
            .iter()
            .find(|attr| {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if (n == "cfg") => true,
                    _ => false,
                }
            })
            .map(|attr| {
                self.codemap
                    .span_to_snippet(attr.span)
                    .map(|string| Builder::sanitize_condition(string))
                    .unwrap_or("".to_string())
            });
        if let ast::ItemKind::Mod(_) = item.node {
            let name = item.ident.to_string();
            {
                let mut tree = self.tree.subtree_at_path(&self.path).unwrap();
                let visibility = if item.vis == ast::Visibility::Public {
                    Visibility::Public
                } else {
                    Visibility::Private
                };
                tree.insert(Tree::new_module(name.clone(), visibility, condition));
            }
            self.path.push(name);
            visit::walk_item(self, item);
            let _ = self.path.pop();
        } else {
            visit::walk_item(self, item);
        }
    }

    fn visit_mod(&mut self, m: &'v ast::Mod, _s: codemap::Span, _n: ast::NodeId) {
        visit::walk_mod(self, m);
        if !self.config.include_orphans {
            return;
        }
        let path = &self.path;
        if let Ok(candidates) = Builder::find_orphan_candidates(path, &self.config.ignored_files) {
            let mut tree = self.tree.subtree_at_path(path).unwrap();
            let names: Vec<_> = tree.subtree_names();
            for candidate in candidates {
                if !names.contains(&candidate) {
                    tree.insert(Tree::new_orphan(candidate.clone()));
                }
            }
        }
    }

    fn visit_mac(&mut self, mac: &'v ast::Mac) {
        visit::walk_mac(self, mac);
    }
}
<commit_msg>fixed one clippy warning & small refactoring (#5)<commit_after>use std::{fs, io, path};

use syntax::{ast, visit, codemap};

use tree::{Tree, Visibility};

pub struct Config {
    pub include_orphans: bool,
    pub ignored_files: Vec<path::PathBuf>,
}

pub struct Builder<'a> {
    tree: Tree,
    path: Vec<String>,
    config: Config,
    codemap: &'a codemap::CodeMap,
}

impl<'a> Builder<'a> {
    pub fn new(config: Config, crate_name: String, codemap: &'a codemap::CodeMap) -> Self {
        let tree = Tree::new_crate(crate_name);
        Builder {
            tree: tree,
            path: vec![],
            config: config,
            codemap: codemap,
        }
    }

    pub fn tree(&self) -> &Tree {
        &self.tree
    }

    fn find_orphan_candidates(path: &[String],
                              ignored_paths: &[path::PathBuf])
                              -> Result<Vec<String>, io::Error> {
        let mut dir_path = "./src/".to_string();
        for name in path.iter() {
            dir_path.push_str(name);
            dir_path.push_str("/");
        }
        fn is_mod(entry: &fs::DirEntry) -> bool {
            let path = entry.path();
            if path.is_file() {
                if path.extension().map_or(false, |ext| ext == "rs") {
                    return file_name(entry)
                        .map_or(false, |s| (s != "mod") && (s != "lib") && (s != "main"));
                }
            } else if path.join("mod.rs").exists() {
                return true;
            }
            false
        }
        fn file_name(entry: &fs::DirEntry) -> Option<String> {
            entry.path().file_stem().map_or(None, |s| s.to_str().map(|s| s.to_string()))
        }
        Ok(try!(fs::read_dir(&dir_path)).filter_map(|e| e.ok())
                                        .filter(is_mod)
                                        .filter(|e| !ignored_paths.contains(&e.path()))
                                        .filter_map(|e| file_name(&e))
                                        .collect::<Vec<_>>())
    }

    fn sanitize_condition(condition: String) -> String {
        let words: Vec<&str> = condition.split_whitespace().collect();
        words.join(" ")
    }
}

impl<'v, 'a> visit::Visitor<'v> for Builder<'a> {
    fn visit_item(&mut self, item: &ast::Item) {
        let condition = item.attrs
            .iter()
            .find(|attr| {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if (n == "cfg") => true,
                    _ => false,
                }
            })
            .map(|attr| {
                self.codemap
                    .span_to_snippet(attr.span)
                    .map(Builder::sanitize_condition)
                    .unwrap_or_else(|_| "".to_string())
            });
        if let ast::ItemKind::Mod(_) = item.node {
            let name = item.ident.to_string();
            {
                let mut tree = self.tree.subtree_at_path(&self.path).unwrap();
                let visibility = if item.vis == ast::Visibility::Public {
                    Visibility::Public
                } else {
                    Visibility::Private
                };
                tree.insert(Tree::new_module(name.clone(), visibility, condition));
            }
            self.path.push(name);
            visit::walk_item(self, item);
            let _ = self.path.pop();
        } else {
            visit::walk_item(self, item);
        }
    }

    fn visit_mod(&mut self, m: &'v ast::Mod, _s: codemap::Span, _n: ast::NodeId) {
        visit::walk_mod(self, m);
        if !self.config.include_orphans {
            return;
        }
        let path = &self.path;
        if let Ok(candidates) = Builder::find_orphan_candidates(path, &self.config.ignored_files) {
            let mut tree = self.tree.subtree_at_path(path).unwrap();
            let names: Vec<_> = tree.subtree_names();
            for candidate in candidates {
                if !names.contains(&candidate) {
                    tree.insert(Tree::new_orphan(candidate.clone()));
                }
            }
        }
    }

    fn visit_mac(&mut self, mac: &'v ast::Mac) {
        visit::walk_mac(self, mac);
    }
}
<|endoftext|>