<commit_before>//! Wrapper for `std::io::Read`.

use std::fmt;
use std::mem;
use std::io;

macro_rules! impl_read_little_endian_integer {
    ($ty:ident, $name:ident, $size:expr) => {
/// Reads a little-endian value and returns it.
        fn $name(&mut self) -> io::Result<$ty> {
            assert_eq!($size, mem::size_of::<$ty>());

            let mut data: $ty = 0;
            let mut slice = unsafe {
                ::std::slice::from_raw_parts_mut(&mut data as *mut $ty as *mut u8, $size)
            };
            try!(self.read_exact(slice));
            Ok($ty::from_le(data))
        }
    }
}

/// The `ReadLittleEndian` trait allows for reading little-endian primitive type values from a
/// source.
pub trait ReadLittleEndian: io::Read {
    impl_read_little_endian_integer!(u8, read_u8, 1);
    impl_read_little_endian_integer!(u32, read_u32, 4);
    impl_read_little_endian_integer!(u64, read_u64, 8);
    impl_read_little_endian_integer!(i16, read_i16, 2);
    impl_read_little_endian_integer!(i32, read_i32, 4);
    impl_read_little_endian_integer!(i64, read_i64, 8);

    /// Reads a little-endian value and returns it.
    fn read_f32(&mut self) -> io::Result<f32> {
        let integer = try!(self.read_u32());
        let val = unsafe { mem::transmute(integer) };
        Ok(val)
    }

    /// Reads a little-endian value and returns it.
    fn read_f64(&mut self) -> io::Result<f64> {
        let integer = try!(self.read_u64());
        let val = unsafe { mem::transmute(integer) };
        Ok(val)
    }
}

impl<R: io::Read> ReadLittleEndian for R {}

/// Source stream for `BinaryParser`.
pub trait ParserSource: fmt::Debug + io::Read {
    /// Returns the current position from the start of the stream.
    fn position(&self) -> u64;

    /// Skips to the given position.
    ///
    /// # Panics
    /// Panics if a byte at the given position has been already read.
    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()>;
}

/// Reader with position info.
pub struct BasicSource<R> {
    /// Source reader.
    source: R,
    /// Current position from the start of the stream..
    position: u64,
}

impl<R: io::Read> BasicSource<R> {
    /// Creates a new `BasicSource`.
    pub fn new(source: R) -> Self {
        BasicSource {
            source: source,
            position: 0,
        }
    }
}

impl<R: io::Read> io::Read for BasicSource<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let read_len = try!(self.source.read(buf));
        self.position += read_len as u64;
        Ok(read_len)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        try!(self.source.read_exact(buf));
        self.position += buf.len() as u64;
        Ok(())
    }
}

impl<R> fmt::Debug for BasicSource<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("BasicSource")
            .field("position", &self.position)
            .finish()
    }
}

impl<R: io::Read> ParserSource for BasicSource<R> {
    fn position(&self) -> u64 {
        self.position
    }

    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()> {
        use std::io::Read;

        assert!(dest_pos >= self.position(),
                "Destination position should be after current position: dest_pos={}, position={}",
                dest_pos,
                self.position());
        const TEMP_BUF_LEN: usize = 256;
        let mut temp_buf = [0u8; TEMP_BUF_LEN];
        let mut rest_len = dest_pos - self.position();
        while rest_len > TEMP_BUF_LEN as u64 {
            try!(self.read_exact(&mut temp_buf));
            rest_len -= TEMP_BUF_LEN as u64;
        }
        try!(self.read_exact(&mut temp_buf[0..rest_len as usize]));

        assert_eq!(self.position(), dest_pos);
        Ok(())
    }
}

/// Reader with position info and seek feature.
///
/// This wrapper doesn't manage offset, i.e. the start of the source stream should be the start of
/// the FBX data.
pub struct SeekableSource<R> {
    /// Source reader.
    source: R,
    /// Current position from the start of the stream..
    position: u64,
}

impl<R: io::Read + io::Seek> SeekableSource<R> {
    /// Creates a new `SeekableSource`.
    pub fn new(source: R) -> Self {
        SeekableSource {
            source: source,
            position: 0,
        }
    }
}

impl<R: io::Read> io::Read for SeekableSource<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let read_len = try!(self.source.read(buf));
        self.position += read_len as u64;
        Ok(read_len)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        try!(self.source.read_exact(buf));
        self.position += buf.len() as u64;
        Ok(())
    }
}

impl<R: io::Read + io::Seek> io::Seek for SeekableSource<R> {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.position = try!(self.source.seek(pos));
        Ok(self.position)
    }
}

impl<R> fmt::Debug for SeekableSource<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SeekableSource")
            .field("position", &self.position)
            .finish()
    }
}

impl<R: io::Read + io::Seek> ParserSource for SeekableSource<R> {
    fn position(&self) -> u64 {
        self.position
    }

    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()> {
        use std::io::{Seek, SeekFrom};

        assert!(dest_pos >= self.position(),
                "Destination position should be after current position: dest_pos={}, position={}",
                dest_pos,
                self.position());
        try!(self.seek(SeekFrom::Start(dest_pos)));

        assert_eq!(self.position(), dest_pos);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor, Seek, SeekFrom};
    use super::BasicSource;

    fn do_test_skip_to(buf_size: usize, skip_dest: u64) {
        let mut short_buf = Cursor::new(vec![0; buf_size]);
        let short_count = {
            let mut reader = BasicSource::new(&mut short_buf);
            reader.skip_to(skip_dest).expect("Failed to skip");
            reader.position()
        };
        assert_eq!(short_count, skip_dest);
        assert_eq!(short_count,
                   short_buf.seek(SeekFrom::Current(0)).expect("Failed to seek"));
    }

    #[test]
    fn test_skip_to() {
        do_test_skip_to(0, 0);
        do_test_skip_to(15, 0);
        do_test_skip_to(30, 23);
        do_test_skip_to(512, 401);
        do_test_skip_to(64, 64);
    }
}
<commit_msg>Add `LimitedSeekReader`<commit_after>//! Wrapper for `std::io::Read`.

use std::fmt;
use std::mem;
use std::io;

macro_rules! impl_read_little_endian_integer {
    ($ty:ident, $name:ident, $size:expr) => {
/// Reads a little-endian value and returns it.
        fn $name(&mut self) -> io::Result<$ty> {
            assert_eq!($size, mem::size_of::<$ty>());

            let mut data: $ty = 0;
            let mut slice = unsafe {
                ::std::slice::from_raw_parts_mut(&mut data as *mut $ty as *mut u8, $size)
            };
            try!(self.read_exact(slice));
            Ok($ty::from_le(data))
        }
    }
}

/// The `ReadLittleEndian` trait allows for reading little-endian primitive type values from a
/// source.
pub trait ReadLittleEndian: io::Read {
    impl_read_little_endian_integer!(u8, read_u8, 1);
    impl_read_little_endian_integer!(u32, read_u32, 4);
    impl_read_little_endian_integer!(u64, read_u64, 8);
    impl_read_little_endian_integer!(i16, read_i16, 2);
    impl_read_little_endian_integer!(i32, read_i32, 4);
    impl_read_little_endian_integer!(i64, read_i64, 8);

    /// Reads a little-endian value and returns it.
    fn read_f32(&mut self) -> io::Result<f32> {
        let integer = try!(self.read_u32());
        let val = unsafe { mem::transmute(integer) };
        Ok(val)
    }

    /// Reads a little-endian value and returns it.
    fn read_f64(&mut self) -> io::Result<f64> {
        let integer = try!(self.read_u64());
        let val = unsafe { mem::transmute(integer) };
        Ok(val)
    }
}

impl<R: io::Read> ReadLittleEndian for R {}

/// Source stream for `BinaryParser`.
pub trait ParserSource: fmt::Debug + io::Read {
    /// Returns the current position from the start of the stream.
    fn position(&self) -> u64;

    /// Skips to the given position.
    ///
    /// # Panics
    /// Panics if a byte at the given position has been already read.
    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()>;
}

/// Reader with position info.
pub struct BasicSource<R> {
    /// Source reader.
    source: R,
    /// Current position from the start of the stream..
    position: u64,
}

impl<R: io::Read> BasicSource<R> {
    /// Creates a new `BasicSource`.
    pub fn new(source: R) -> Self {
        BasicSource {
            source: source,
            position: 0,
        }
    }
}

impl<R: io::Read> io::Read for BasicSource<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let read_len = try!(self.source.read(buf));
        self.position += read_len as u64;
        Ok(read_len)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        try!(self.source.read_exact(buf));
        self.position += buf.len() as u64;
        Ok(())
    }
}

impl<R> fmt::Debug for BasicSource<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("BasicSource")
            .field("position", &self.position)
            .finish()
    }
}

impl<R: io::Read> ParserSource for BasicSource<R> {
    fn position(&self) -> u64 {
        self.position
    }

    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()> {
        use std::io::Read;

        assert!(dest_pos >= self.position(),
                "Destination position should be after current position: dest_pos={}, position={}",
                dest_pos,
                self.position());
        const TEMP_BUF_LEN: usize = 256;
        let mut temp_buf = [0u8; TEMP_BUF_LEN];
        let mut rest_len = dest_pos - self.position();
        while rest_len > TEMP_BUF_LEN as u64 {
            try!(self.read_exact(&mut temp_buf));
            rest_len -= TEMP_BUF_LEN as u64;
        }
        try!(self.read_exact(&mut temp_buf[0..rest_len as usize]));

        assert_eq!(self.position(), dest_pos);
        Ok(())
    }
}

/// Reader with position info and seek feature.
///
/// This wrapper doesn't manage offset, i.e. the start of the source stream should be the start of
/// the FBX data.
pub struct SeekableSource<R> {
    /// Source reader.
    source: R,
    /// Current position from the start of the stream..
    position: u64,
}

impl<R: io::Read + io::Seek> SeekableSource<R> {
    /// Creates a new `SeekableSource`.
    pub fn new(source: R) -> Self {
        SeekableSource {
            source: source,
            position: 0,
        }
    }
}

impl<R: io::Read> io::Read for SeekableSource<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let read_len = try!(self.source.read(buf));
        self.position += read_len as u64;
        Ok(read_len)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        try!(self.source.read_exact(buf));
        self.position += buf.len() as u64;
        Ok(())
    }
}

impl<R: io::Read + io::Seek> io::Seek for SeekableSource<R> {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        self.position = try!(self.source.seek(pos));
        Ok(self.position)
    }
}

impl<R> fmt::Debug for SeekableSource<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SeekableSource")
            .field("position", &self.position)
            .finish()
    }
}

impl<R: io::Read + io::Seek> ParserSource for SeekableSource<R> {
    fn position(&self) -> u64 {
        self.position
    }

    fn skip_to(&mut self, dest_pos: u64) -> io::Result<()> {
        use std::io::{Seek, SeekFrom};

        assert!(dest_pos >= self.position(),
                "Destination position should be after current position: dest_pos={}, position={}",
                dest_pos,
                self.position());
        try!(self.seek(SeekFrom::Start(dest_pos)));

        assert_eq!(self.position(), dest_pos);
        Ok(())
    }
}

/// Reader which can read and seek limited area of a stream.
pub struct LimitedSeekReader<R> {
    /// Source stream.
    source: R,
    /// Current position
    current: u64,
    /// Start position.
    begin: u64,
    /// End position.
    end: u64,
}

impl<R: io::Read> LimitedSeekReader<R> {
    /// Creates a new `LimitedSeekReader`.
    pub fn new(source: R, current: u64, begin: u64, end: u64) -> Self {
        assert!(begin <= end);
        assert!(begin <= current);
        assert!(current <= end);

        LimitedSeekReader {
            source: source,
            current: current,
            begin: begin,
            end: end,
        }
    }
}

impl<R> LimitedSeekReader<R> {
    /// Returns length of region allowed to be read in bytes.
    pub fn len(&self) -> u64 {
        assert!(self.begin <= self.end);
        self.end - self.begin
    }

    /// Returns rest length of region allowed to be read in bytes.
    fn rest_len(&self) -> u64 {
        assert!(self.current <= self.end);
        self.end - self.current
    }

    /// Returns the current position relative to the start of the limited region.
    fn rel_pos(&self) -> u64 {
        assert!(self.begin <= self.current);
        self.current - self.begin
    }
}

impl<R: io::Read> io::Read for LimitedSeekReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let limit = self.rest_len();
        let size = try!(self.source.by_ref().take(limit).read(buf));
        self.current += size as u64;
        Ok(size)
    }
}

impl<R: io::Seek> io::Seek for LimitedSeekReader<R> {
    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
        match pos {
            io::SeekFrom::Start(val) => {
                let offset = ::std::cmp::min(self.len(), val);
                let target = self.begin + offset;
                assert!(target >= self.begin);
                assert!(target <= self.end);
                self.current = try!(self.source.seek(io::SeekFrom::Start(target)));
            },
            io::SeekFrom::Current(val) => {
                let offset = if val < 0 {
                    ::std::cmp::max(-(self.rel_pos() as i64), val)
                } else {
                    ::std::cmp::min(self.rest_len(), val as u64) as i64
                };
                let target = self.current as i64 + offset;
                assert!(target as u64 >= self.begin);
                assert!(target as u64 <= self.end);
                self.current = try!(self.source.seek(io::SeekFrom::Current(offset)));
            },
            io::SeekFrom::End(val) => {
                let offset = ::std::cmp::max(-(self.len() as i64), ::std::cmp::min(0, val));
                let target = (self.end as i64 + offset) as u64;
                assert!(target >= self.begin);
                assert!(target <= self.end);
                self.current = try!(self.source.seek(io::SeekFrom::Start(target)));
            },
        }
        Ok(self.rel_pos())
    }
}

impl<R> fmt::Debug for LimitedSeekReader<R> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("LimitedSeekReader")
            .field("current", &self.current)
            .field("begin", &self.begin)
            .field("end", &self.end)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor, Seek, SeekFrom};
    use super::BasicSource;

    fn do_test_skip_to(buf_size: usize, skip_dest: u64) {
        let mut short_buf = Cursor::new(vec![0; buf_size]);
        let short_count = {
            let mut reader = BasicSource::new(&mut short_buf);
            reader.skip_to(skip_dest).expect("Failed to skip");
            reader.position()
        };
        assert_eq!(short_count, skip_dest);
        assert_eq!(short_count,
                   short_buf.seek(SeekFrom::Current(0)).expect("Failed to seek"));
    }

    #[test]
    fn test_skip_to() {
        do_test_skip_to(0, 0);
        do_test_skip_to(15, 0);
        do_test_skip_to(30, 23);
        do_test_skip_to(512, 401);
        do_test_skip_to(64, 64);
    }
}
<|endoftext|>


<commit_before>use std::collections::VecMap;
use std::iter::repeat;

pub struct BOM {
    m: usize,
    table: Vec<VecMap<usize>>
}

impl BOM {
    pub fn new(pattern: &[u8]) -> Self {
        let m = pattern.len();
        let maxsym = *pattern.iter().max().expect("Expecting non-empty pattern.") as usize;
        let mut table: Vec<VecMap<usize>> = Vec::with_capacity(m);
        // init suffix table, initially all values unknown
        // suff[i] is the state in which the longest suffix of
        // pattern[..i+1] ends that does not end in i
        let mut suff: Vec<Option<usize>> = repeat(None).take(m + 1).collect();

        for i in 1..m+1 {
            let a = pattern[i - 1] as usize;
            let mut delta = VecMap::with_capacity(maxsym);
            // reading symbol a leads into state i (this is an inner edge)
            delta.insert(a, i);
            // now, add edges for substrings ending with a
            let mut k = suff[i - 1];

            // for this iterator over the known suffixes until
            // reaching an edge labelled with a or the start
            loop {
                match k {
                    Some(k_) => {
                        if table[k_].contains_key(&a) {
                            break;
                        }
                        table[k_].insert(a, i);
                        k = suff[k_];
                    },
                    None => break
                }
            }

            // the longest suffix is either 0 or the state
            // reached by the edge labelled with a
            suff[i] = Some(match k {
                Some(k) => *table[k].get(&a).unwrap(),
                None => 0
            });

            table.push(delta);
        }

        BOM { m: m, table: table }
    }

    fn delta(&self, q: usize, a: u8) -> Option<usize> {
        match self.table[q].get(&(a as usize)) {
            Some(&q) => Some(q),
            None => None
        }
    }

    pub fn find_all<'a>(&'a self, text: &'a [u8]) -> FindAll {
        FindAll { bom: self, text: text, window: self.m }
    }
}

pub struct FindAll<'a> {
    bom: &'a BOM,
    text: &'a [u8],
    window: usize
}

impl<'a> Iterator for FindAll<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        while self.window <= self.text.len() {
            let (mut q, mut j) = (Some(0), 1);
            while j <= self.bom.m {
                match q {
                    Some(q_) => {
                        q = self.bom.delta(q_, self.text[self.window - j]);
                        j += 1;
                    },
                    None => break
                }
            }
            // putative start position
            let i = self.window - self.bom.m;
            self.window += self.bom.m - j + 2;
            if q.is_some() {
                // return match
                return Some(i);
            }
        }
        None
    }
}

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

    #[test]
    fn test_delta() {
        let pattern = b"nannannnq";
        let bom = BOM::new(pattern);
        assert_eq!(bom.delta(0, b'n'), Some(1));
        assert_eq!(bom.delta(1, b'a'), Some(2));
        assert_eq!(bom.delta(2, b'n'), Some(3));
        assert_eq!(bom.delta(3, b'n'), Some(4));
        assert_eq!(bom.delta(4, b'a'), Some(5));
        assert_eq!(bom.delta(5, b'n'), Some(6));
        assert_eq!(bom.delta(6, b'n'), Some(7));
        assert_eq!(bom.delta(7, b'n'), Some(8));
        assert_eq!(bom.delta(8, b'q'), Some(9));

        assert_eq!(bom.delta(0, b'a'), Some(2));
        assert_eq!(bom.delta(0, b'q'), Some(9));
        assert_eq!(bom.delta(1, b'n'), Some(4));
        assert_eq!(bom.delta(1, b'q'), Some(9));
        assert_eq!(bom.delta(4, b'n'), Some(8));
        assert_eq!(bom.delta(4, b'q'), Some(9));
    }
}
<commit_msg>Docs.<commit_after>use std::collections::VecMap;
use std::iter::repeat;

/// Backward oracle matching algorithm.
///
/// # Example
///
/// ```
/// use bio::pattern_matching::bom::BOM;
/// let text = b"ACGGCTAGGAAAAAGACTGAGGACTGAAAA";
/// let pattern = b"AAAA";
/// let bom = BOM::new(pattern);
/// let occ: Vec<usize> = bom.find_all(text).collect();
/// assert_eq!(occ, [9, 10, 26]);
/// ```
pub struct BOM {
    m: usize,
    table: Vec<VecMap<usize>>
}

impl BOM {
    pub fn new(pattern: &[u8]) -> Self {
        let m = pattern.len();
        let maxsym = *pattern.iter().max().expect("Expecting non-empty pattern.") as usize;
        let mut table: Vec<VecMap<usize>> = Vec::with_capacity(m);
        // init suffix table, initially all values unknown
        // suff[i] is the state in which the longest suffix of
        // pattern[..i+1] ends that does not end in i
        let mut suff: Vec<Option<usize>> = repeat(None).take(m + 1).collect();

        for i in 1..m+1 {
            let a = pattern[i - 1] as usize;
            let mut delta = VecMap::with_capacity(maxsym);
            // reading symbol a leads into state i (this is an inner edge)
            delta.insert(a, i);
            // now, add edges for substrings ending with a
            let mut k = suff[i - 1];

            // for this iterator over the known suffixes until
            // reaching an edge labelled with a or the start
            loop {
                match k {
                    Some(k_) => {
                        if table[k_].contains_key(&a) {
                            break;
                        }
                        table[k_].insert(a, i);
                        k = suff[k_];
                    },
                    None => break
                }
            }

            // the longest suffix is either 0 or the state
            // reached by the edge labelled with a
            suff[i] = Some(match k {
                Some(k) => *table[k].get(&a).unwrap(),
                None => 0
            });

            table.push(delta);
        }

        BOM { m: m, table: table }
    }

    fn delta(&self, q: usize, a: u8) -> Option<usize> {
        match self.table[q].get(&(a as usize)) {
            Some(&q) => Some(q),
            None => None
        }
    }

    /// Find all exact occurrences of the pattern in the given text.
    ///
    /// # Arguments
    ///
    /// * `text` - the given text
    pub fn find_all<'a>(&'a self, text: &'a [u8]) -> FindAll {
        FindAll { bom: self, text: text, window: self.m }
    }
}

pub struct FindAll<'a> {
    bom: &'a BOM,
    text: &'a [u8],
    window: usize
}

impl<'a> Iterator for FindAll<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        while self.window <= self.text.len() {
            let (mut q, mut j) = (Some(0), 1);
            while j <= self.bom.m {
                match q {
                    Some(q_) => {
                        q = self.bom.delta(q_, self.text[self.window - j]);
                        j += 1;
                    },
                    None => break
                }
            }
            // putative start position
            let i = self.window - self.bom.m;
            self.window += self.bom.m - j + 2;
            if q.is_some() {
                // return match
                return Some(i);
            }
        }
        None
    }
}

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

    #[test]
    fn test_delta() {
        let pattern = b"nannannnq";
        let bom = BOM::new(pattern);
        assert_eq!(bom.delta(0, b'n'), Some(1));
        assert_eq!(bom.delta(1, b'a'), Some(2));
        assert_eq!(bom.delta(2, b'n'), Some(3));
        assert_eq!(bom.delta(3, b'n'), Some(4));
        assert_eq!(bom.delta(4, b'a'), Some(5));
        assert_eq!(bom.delta(5, b'n'), Some(6));
        assert_eq!(bom.delta(6, b'n'), Some(7));
        assert_eq!(bom.delta(7, b'n'), Some(8));
        assert_eq!(bom.delta(8, b'q'), Some(9));

        assert_eq!(bom.delta(0, b'a'), Some(2));
        assert_eq!(bom.delta(0, b'q'), Some(9));
        assert_eq!(bom.delta(1, b'n'), Some(4));
        assert_eq!(bom.delta(1, b'q'), Some(9));
        assert_eq!(bom.delta(4, b'n'), Some(8));
        assert_eq!(bom.delta(4, b'q'), Some(9));
    }
}
<|endoftext|>


<commit_before>use crate::{
    arch::{gdt, interrupt::InterruptStack},
    context,
    ptrace,
    syscall,
    syscall::flag::{PTRACE_FLAG_IGNORE, PTRACE_STOP_PRE_SYSCALL, PTRACE_STOP_POST_SYSCALL},
};
use x86::msr;

pub unsafe fn init() {
    // IA32_STAR[31:0] are reserved.

    // The base selector of the two consecutive segments for kernel code and the immediately
    // suceeding stack (data).
    let syscall_cs_ss_base = (gdt::GDT_KERNEL_CODE as u16) << 3;
    // The base selector of the three consecutive segments (of which two are used) for user code
    // and user data. It points to a 32-bit code segment, which must be followed by a data segment
    // (stack), and a 64-bit code segment.
    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | 3;
    let star_high = u32::from(syscall_cs_ss_base) | (u32::from(sysret_cs_ss_base) << 16);

    msr::wrmsr(msr::IA32_STAR, u64::from(star_high) << 32);
    msr::wrmsr(msr::IA32_LSTAR, syscall_instruction as u64);
    msr::wrmsr(msr::IA32_FMASK, 0x0300); // Clear trap flag and interrupt enable

    let efer = msr::rdmsr(msr::IA32_EFER);
    msr::wrmsr(msr::IA32_EFER, efer | 1);
}

macro_rules! with_interrupt_stack {
    (|$stack:ident| $code:block) => {{
        let allowed = ptrace::breakpoint_callback(PTRACE_STOP_PRE_SYSCALL, None)
            .and_then(|_| ptrace::next_breakpoint().map(|f| !f.contains(PTRACE_FLAG_IGNORE)));

        if allowed.unwrap_or(true) {
            // If the syscall is `clone`, the clone won't return here. Instead,
            // it'll return early and leave any undropped values. This is
            // actually GOOD, because any references are at that point UB
            // anyway, because they are based on the wrong stack.
            let $stack = &mut *$stack;
            (*$stack).scratch.rax = $code;
        }

        ptrace::breakpoint_callback(PTRACE_STOP_POST_SYSCALL, None);
    }}
}

#[no_mangle]
pub unsafe extern "C" fn __inner_syscall_instruction(stack: *mut InterruptStack) {
    let _guard = ptrace::set_process_regs(stack);
    with_interrupt_stack!(|stack| {
        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, scratch.rdi, scratch.rsi, scratch.rdx, scratch.r10, scratch.r8, rbp, stack)
    });
}

function!(syscall_instruction => {
    // Yes, this is magic. No, you don't need to understand
    "
        swapgs                    // Set gs segment to TSS
        mov gs:[0x70], rsp        // Save userspace stack pointer
        mov rsp, gs:[4]           // Load kernel stack pointer
        push QWORD PTR 5 * 8 + 3  // Push fake userspace SS (resembling iret frame)
        push QWORD PTR gs:[0x70]  // Push userspace rsp
        push r11                  // Push rflags
        push QWORD PTR 6 * 8 + 3  // Push fake userspace CS (resembling iret frame)
        push rcx                  // Push userspace return pointer
    ",

    // Push context registers
    "push rax\n",
    push_scratch!(),
    push_preserved!(),
    push_fs!(),

    // TODO: Map PTI
    // $crate::arch::x86_64::pti::map();

    // Call inner funtion
    "mov rdi, rsp\n",
    "call __inner_syscall_instruction\n",

    // TODO: Unmap PTI
    // $crate::arch::x86_64::pti::unmap();

    // Pop context registers
    pop_fs!(),
    pop_preserved!(),
    pop_scratch!(),

    // Return
    "
        pop rcx                 // Pop userspace return pointer
        add rsp, 8              // Pop CS
        pop r11                 // Pop rflags
        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
        add rsp, 8              // Pop SS
        mov rsp, gs:[0x70]      // Restore userspace stack pointer
        swapgs                  // Restore gs from TSS to user data
        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
    ",
});

interrupt_stack!(syscall, |stack| {
    with_interrupt_stack!(|stack| {
        {
            let contexts = context::contexts();
            let context = contexts.current();
            if let Some(current) = context {
                let current = current.read();
                println!("Warning: Context {} used deprecated `int 0x80` construct", *current.name.read());
            } else {
                println!("Warning: Unknown context used deprecated `int 0x80` construct");
            }
        }

        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, stack.preserved.rbx, scratch.rcx, scratch.rdx, scratch.rsi, scratch.rdi, rbp, stack)
    })
});

function!(clone_ret => {
    // The address of this instruction is injected by `clone` in process.rs, on
    // top of the stack syscall->inner in this file, which is done using the rbp
    // register we save there.
    //
    // The top of our stack here is the address pointed to by rbp, which is:
    //
    // - the previous rbp
    // - the return location
    //
    // Our goal is to return from the parent function, inner, so we restore
    // rbp...
    "pop rbp\n",
    // ...and we return to the address at the top of the stack
    "ret\n",
});
<commit_msg>Check whether RCX is canonical in sysretq.<commit_after>use crate::{
    arch::{gdt, interrupt::InterruptStack},
    context,
    ptrace,
    syscall,
    syscall::flag::{PTRACE_FLAG_IGNORE, PTRACE_STOP_PRE_SYSCALL, PTRACE_STOP_POST_SYSCALL},
};
use x86::msr;

pub unsafe fn init() {
    // IA32_STAR[31:0] are reserved.

    // The base selector of the two consecutive segments for kernel code and the immediately
    // suceeding stack (data).
    let syscall_cs_ss_base = (gdt::GDT_KERNEL_CODE as u16) << 3;
    // The base selector of the three consecutive segments (of which two are used) for user code
    // and user data. It points to a 32-bit code segment, which must be followed by a data segment
    // (stack), and a 64-bit code segment.
    let sysret_cs_ss_base = ((gdt::GDT_USER_CODE32_UNUSED as u16) << 3) | 3;
    let star_high = u32::from(syscall_cs_ss_base) | (u32::from(sysret_cs_ss_base) << 16);

    msr::wrmsr(msr::IA32_STAR, u64::from(star_high) << 32);
    msr::wrmsr(msr::IA32_LSTAR, syscall_instruction as u64);
    msr::wrmsr(msr::IA32_FMASK, 0x0300); // Clear trap flag and interrupt enable

    let efer = msr::rdmsr(msr::IA32_EFER);
    msr::wrmsr(msr::IA32_EFER, efer | 1);
}

macro_rules! with_interrupt_stack {
    (|$stack:ident| $code:block) => {{
        let allowed = ptrace::breakpoint_callback(PTRACE_STOP_PRE_SYSCALL, None)
            .and_then(|_| ptrace::next_breakpoint().map(|f| !f.contains(PTRACE_FLAG_IGNORE)));

        if allowed.unwrap_or(true) {
            // If the syscall is `clone`, the clone won't return here. Instead,
            // it'll return early and leave any undropped values. This is
            // actually GOOD, because any references are at that point UB
            // anyway, because they are based on the wrong stack.
            let $stack = &mut *$stack;
            (*$stack).scratch.rax = $code;
        }

        ptrace::breakpoint_callback(PTRACE_STOP_POST_SYSCALL, None);
    }}
}

#[no_mangle]
pub unsafe extern "C" fn __inner_syscall_instruction(stack: *mut InterruptStack) {
    let _guard = ptrace::set_process_regs(stack);
    with_interrupt_stack!(|stack| {
        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, scratch.rdi, scratch.rsi, scratch.rdx, scratch.r10, scratch.r8, rbp, stack)
    });
}

macro_rules! fast_sysretq {
    () => { "
        pop rcx                 // Pop userspace return pointer
        sub rsp, 8              // Pop fake userspace CS
        pop r11                 // Pop rflags
        pop QWORD PTR gs:[0x70] // Pop userspace stack pointer
        mov rsp, gs:[0x70]      // Restore userspace stack pointer
        swapgs                  // Restore gs from TSS to user data
        ud2
        sysretq                 // Return into userspace; RCX=>RIP,R11=>RFLAGS
    " }
}

function!(syscall_instruction => {
    // Yes, this is magic. No, you don't need to understand
    "
        swapgs                    // Set gs segment to TSS
        mov gs:[0x70], rsp        // Save userspace stack pointer
        mov rsp, gs:[4]           // Load kernel stack pointer
        push QWORD PTR 5 * 8 + 3  // Push fake SS (resembling iret stack frame)
        push QWORD PTR gs:[0x70]  // Push userspace rsp
        push r11                  // Push rflags
        push QWORD PTR 6 * 8 + 3  // Push fake CS (resembling iret stack frame)
        push rcx                  // Push userspace return pointer
    ",

    // Push context registers
    "push rax\n",
    push_scratch!(),
    push_preserved!(),
    push_fs!(),

    // TODO: Map PTI
    // $crate::arch::x86_64::pti::map();

    // Call inner funtion
    "mov rdi, rsp\n",
    "call __inner_syscall_instruction\n",

    // TODO: Unmap PTI
    // $crate::arch::x86_64::pti::unmap();

    // Pop context registers
    pop_fs!(),
    pop_preserved!(),
    pop_scratch!(),

    // Return
    "
        // Test whether RCX is canonical. This is not strictly necessary, but could be fatal if
        // some kernel bug would allow RCX to be modified by user code.

        // Peek at the preserved RCX, to double-check that it is canonical.
        // Set ZF iff bit 47 (i.e. the bit that must be sign extended) is set.

        // TODO: Any hacky one-instruction canonicalness check?

        bt QWORD PTR [rsp], 47

        // Jump to the canonicalness check.
        jnz 2f

        // Otherwise, continue with the fast return.
        1:
        // Fast sysretq
    ",

    fast_sysretq!(),

    "
        // Jumps to aligned addresses are faster, and it is not always the case the userspace
        // doesn't call from higher half (though unlikely).
        .align 8, 0x90

        2:

        // Load userspace return RIP.
        mov rcx, [rsp]

        // Shift it to the right by 47, only getting the bits which must all be zero or one.
        shr rcx, 47
        // Since we have already checked that RCX[47] == 1, RCX[63:47] must all be one.
        cmp rcx, 0x1FFFF
        // If the address is canonical, go to the slower iretq.
        jne 3f
    ",

        fast_sysretq!(),

    "
        3:
        // Slow iretq

        xor rcx, rcx
        xor r11, r11
        swapgs
        iretq
    ",
});

interrupt_stack!(syscall, |stack| {
    with_interrupt_stack!(|stack| {
        {
            let contexts = context::contexts();
            let context = contexts.current();
            if let Some(current) = context {
                let current = current.read();
                println!("Warning: Context {} used deprecated `int 0x80` construct", *current.name.read());
            } else {
                println!("Warning: Unknown context used deprecated `int 0x80` construct");
            }
        }

        // Set a restore point for clone
        let rbp;
        asm!("mov {}, rbp", out(reg) rbp);

        let scratch = &stack.scratch;
        syscall::syscall(scratch.rax, stack.preserved.rbx, scratch.rcx, scratch.rdx, scratch.rsi, scratch.rdi, rbp, stack)
    })
});

function!(clone_ret => {
    // The address of this instruction is injected by `clone` in process.rs, on
    // top of the stack syscall->inner in this file, which is done using the rbp
    // register we save there.
    //
    // The top of our stack here is the address pointed to by rbp, which is:
    //
    // - the previous rbp
    // - the return location
    //
    // Our goal is to return from the parent function, inner, so we restore
    // rbp...
    "pop rbp\n",
    // ...and we return to the address at the top of the stack
    "ret\n",
});
<|endoftext|>


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

//! Abstraction of a task pool for basic parallelism.

use core::prelude::*;

use task;
use task::spawn;
use vec::Vec;
use comm::{channel, Sender};

enum Msg<T> {
    Execute(proc(&T):Send),
    Quit
}

/// A task pool used to execute functions in parallel.
pub struct TaskPool<T> {
    channels: Vec<Sender<Msg<T>>>,
    next_index: uint,
}

#[unsafe_destructor]
impl<T> Drop for TaskPool<T> {
    fn drop(&mut self) {
        for channel in self.channels.iter_mut() {
            channel.send(Quit);
        }
    }
}

impl<T> TaskPool<T> {
    /// Spawns a new task pool with `n_tasks` tasks. The provided
    /// `init_fn_factory` returns a function which, given the index of the
    /// task, should return local data to be kept around in that task.
    ///
    /// # Panics
    ///
    /// This function will panic if `n_tasks` is less than 1.
    pub fn new(n_tasks: uint,
               init_fn_factory: || -> proc(uint):Send -> T)
               -> TaskPool<T> {
        assert!(n_tasks >= 1);

        let channels = Vec::from_fn(n_tasks, |i| {
            let (tx, rx) = channel::<Msg<T>>();
            let init_fn = init_fn_factory();

            let task_body = proc() {
                let local_data = init_fn(i);
                loop {
                    match rx.recv() {
                        Execute(f) => f(&local_data),
                        Quit => break
                    }
                }
            };

            // Run on this scheduler.
            task::spawn(task_body);

            tx
        });

        return TaskPool {
            channels: channels,
            next_index: 0,
        };
    }

    /// Executes the function `f` on a task in the pool. The function
    /// receives a reference to the local data returned by the `init_fn`.
    pub fn execute(&mut self, f: proc(&T):Send) {
        self.channels[self.next_index].send(Execute(f));
        self.next_index += 1;
        if self.next_index == self.channels.len() { self.next_index = 0; }
    }
}

#[test]
fn test_task_pool() {
    let f: || -> proc(uint):Send -> uint = || { proc(i) i };
    let mut pool = TaskPool::new(4, f);
    for _ in range(0u, 8) {
        pool.execute(proc(i) println!("Hello from thread {}!", *i));
    }
}

#[test]
#[should_fail]
fn test_zero_tasks_panic() {
    let f: || -> proc(uint):Send -> uint = || { proc(i) i };
    TaskPool::new(0, f);
}
<commit_msg>Rewrite std::sync::TaskPool to be load balancing and panic-resistant<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.

//! Abstraction of a task pool for basic parallelism.

use core::prelude::*;

use task::spawn;
use comm::{channel, Sender, Receiver};
use sync::{Arc, Mutex};

struct Sentinel<'a> {
    jobs: &'a Arc<Mutex<Receiver<proc(): Send>>>,
    active: bool
}

impl<'a> Sentinel<'a> {
    fn new(jobs: &Arc<Mutex<Receiver<proc(): Send>>>) -> Sentinel {
        Sentinel {
            jobs: jobs,
            active: true
        }
    }

    // Cancel and destroy this sentinel.
    fn cancel(mut self) {
        self.active = false;
    }
}

#[unsafe_destructor]
impl<'a> Drop for Sentinel<'a> {
    fn drop(&mut self) {
        if self.active {
            spawn_in_pool(self.jobs.clone())
        }
    }
}

/// A task pool used to execute functions in parallel.
///
/// Spawns `n` worker tasks and replenishes the pool if any worker tasks
/// panic.
///
/// # Example
///
/// ```rust
/// # use sync::TaskPool;
/// # use iter::AdditiveIterator;
///
/// let pool = TaskPool::new(4u);
///
/// let (tx, rx) = channel();
/// for _ in range(0, 8u) {
///     let tx = tx.clone();
///     pool.execute(proc() {
///         tx.send(1u);
///     });
/// }
///
/// assert_eq!(rx.iter().take(8u).sum(), 8u);
/// ```
pub struct TaskPool {
    // How the taskpool communicates with subtasks.
    //
    // This is the only such Sender, so when it is dropped all subtasks will
    // quit.
    jobs: Sender<proc(): Send>
}

impl TaskPool {
    /// Spawns a new task pool with `tasks` tasks.
    ///
    /// # Panics
    ///
    /// This function will panic if `tasks` is 0.
    pub fn new(tasks: uint) -> TaskPool {
        assert!(tasks >= 1);

        let (tx, rx) = channel::<proc(): Send>();
        let rx = Arc::new(Mutex::new(rx));

        // Taskpool tasks.
        for _ in range(0, tasks) {
            spawn_in_pool(rx.clone());
        }

        TaskPool { jobs: tx }
    }

    /// Executes the function `job` on a task in the pool.
    pub fn execute(&self, job: proc():Send) {
        self.jobs.send(job);
    }
}

fn spawn_in_pool(jobs: Arc<Mutex<Receiver<proc(): Send>>>) {
    spawn(proc() {
        // Will spawn a new task on panic unless it is cancelled.
        let sentinel = Sentinel::new(&jobs);

        loop {
            let message = {
                // Only lock jobs for the time it takes
                // to get a job, not run it.
                let lock = jobs.lock();
                lock.recv_opt()
            };

            match message {
                Ok(job) => job(),

                // The Taskpool was dropped.
                Err(..) => break
            }
        }

        sentinel.cancel();
    })
}

#[cfg(test)]
mod test {
    use core::prelude::*;
    use super::*;
    use comm::channel;
    use iter::range;

    const TEST_TASKS: uint = 4u;

    #[test]
    fn test_works() {
        use iter::AdditiveIterator;

        let pool = TaskPool::new(TEST_TASKS);

        let (tx, rx) = channel();
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            pool.execute(proc() {
                tx.send(1u);
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    #[should_fail]
    fn test_zero_tasks_panic() {
        TaskPool::new(0);
    }

    #[test]
    fn test_recovery_from_subtask_panic() {
        use iter::AdditiveIterator;

        let pool = TaskPool::new(TEST_TASKS);

        // Panic all the existing tasks.
        for _ in range(0, TEST_TASKS) {
            pool.execute(proc() { panic!() });
        }

        // Ensure new tasks were spawned to compensate.
        let (tx, rx) = channel();
        for _ in range(0, TEST_TASKS) {
            let tx = tx.clone();
            pool.execute(proc() {
                tx.send(1u);
            });
        }

        assert_eq!(rx.iter().take(TEST_TASKS).sum(), TEST_TASKS);
    }

    #[test]
    fn test_should_not_panic_on_drop_if_subtasks_panic_after_drop() {
        use sync::{Arc, Barrier};

        let pool = TaskPool::new(TEST_TASKS);
        let waiter = Arc::new(Barrier::new(TEST_TASKS + 1));

        // Panic all the existing tasks in a bit.
        for _ in range(0, TEST_TASKS) {
            let waiter = waiter.clone();
            pool.execute(proc() {
                waiter.wait();
                panic!();
            });
        }

        drop(pool);

        // Kick off the failure.
        waiter.wait();
    }
}

<|endoftext|>


<commit_before>use rustc_index::bit_set::BitSet;
use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor};
use rustc_middle::mir::{self, Local, Location};

use crate::dataflow::{AnalysisDomain, Backward, BottomValue, GenKill, GenKillAnalysis};

/// A [live-variable dataflow analysis][liveness].
///
/// This analysis considers references as being used only at the point of the
/// borrow. In other words, this analysis does not track uses because of references that already
/// exist. See [this `mir-datalow` test][flow-test] for an example. You almost never want to use
/// this analysis without also looking at the results of [`MaybeBorrowedLocals`].
///
/// [`MaybeBorrowedLocals`]: ../struct.MaybeBorrowedLocals.html
/// [flow-test]: https://github.com/rust-lang/rust/blob/a08c47310c7d49cbdc5d7afb38408ba519967ecd/src/test/ui/mir-dataflow/liveness-ptr.rs
/// [liveness]: https://en.wikipedia.org/wiki/Live_variable_analysis
pub struct MaybeLiveLocals;

impl MaybeLiveLocals {
    fn transfer_function<T>(&self, trans: &'a mut T) -> TransferFunction<'a, T> {
        TransferFunction(trans)
    }
}

impl BottomValue for MaybeLiveLocals {
    // bottom = not live
    const BOTTOM_VALUE: bool = false;
}

impl AnalysisDomain<'tcx> for MaybeLiveLocals {
    type Idx = Local;
    type Direction = Backward;

    const NAME: &'static str = "liveness";

    fn bits_per_block(&self, body: &mir::Body<'tcx>) -> usize {
        body.local_decls.len()
    }

    fn initialize_start_block(&self, _: &mir::Body<'tcx>, _: &mut BitSet<Self::Idx>) {
        // No variables are live until we observe a use
    }
}

impl GenKillAnalysis<'tcx> for MaybeLiveLocals {
    fn statement_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        self.transfer_function(trans).visit_statement(statement, location);
    }

    fn terminator_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        terminator: &mir::Terminator<'tcx>,
        location: Location,
    ) {
        self.transfer_function(trans).visit_terminator(terminator, location);
    }

    fn call_return_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        _block: mir::BasicBlock,
        _func: &mir::Operand<'tcx>,
        _args: &[mir::Operand<'tcx>],
        dest_place: mir::Place<'tcx>,
    ) {
        if let Some(local) = dest_place.as_local() {
            trans.kill(local);
        }
    }

    fn yield_resume_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        _resume_block: mir::BasicBlock,
        resume_place: mir::Place<'tcx>,
    ) {
        if let Some(local) = resume_place.as_local() {
            trans.kill(local);
        }
    }
}

struct TransferFunction<'a, T>(&'a mut T);

impl<'tcx, T> Visitor<'tcx> for TransferFunction<'_, T>
where
    T: GenKill<Local>,
{
    fn visit_local(&mut self, &local: &Local, context: PlaceContext, _: Location) {
        match DefUse::for_place(context) {
            Some(DefUse::Def) => self.0.kill(local),
            Some(DefUse::Use) => self.0.gen(local),
            _ => {}
        }
    }
}

#[derive(Eq, PartialEq, Clone)]
enum DefUse {
    Def,
    Use,
}

impl DefUse {
    fn for_place(context: PlaceContext) -> Option<DefUse> {
        match context {
            PlaceContext::NonUse(_) => None,

            PlaceContext::MutatingUse(MutatingUseContext::Store) => Some(DefUse::Def),

            // `MutatingUseContext::Call` and `MutatingUseContext::Yield` indicate that this is the
            // destination place for a `Call` return or `Yield` resume respectively. Since this is
            // only a `Def` when the function returns succesfully, we handle this case separately
            // in `call_return_effect` above.
            PlaceContext::MutatingUse(MutatingUseContext::Call | MutatingUseContext::Yield) => None,

            // All other contexts are uses...
            PlaceContext::MutatingUse(
                MutatingUseContext::AddressOf
                | MutatingUseContext::AsmOutput
                | MutatingUseContext::Borrow
                | MutatingUseContext::Drop
                | MutatingUseContext::Projection
                | MutatingUseContext::Retag,
            )
            | PlaceContext::NonMutatingUse(
                NonMutatingUseContext::AddressOf
                | NonMutatingUseContext::Copy
                | NonMutatingUseContext::Inspect
                | NonMutatingUseContext::Move
                | NonMutatingUseContext::Projection
                | NonMutatingUseContext::ShallowBorrow
                | NonMutatingUseContext::SharedBorrow
                | NonMutatingUseContext::UniqueBorrow,
            ) => Some(DefUse::Use),
        }
    }
}
<commit_msg>Handle stores to projections correctly in liveness analysis<commit_after>use rustc_index::bit_set::BitSet;
use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor};
use rustc_middle::mir::{self, Local, Location};

use crate::dataflow::{AnalysisDomain, Backward, BottomValue, GenKill, GenKillAnalysis};

/// A [live-variable dataflow analysis][liveness].
///
/// This analysis considers references as being used only at the point of the
/// borrow. In other words, this analysis does not track uses because of references that already
/// exist. See [this `mir-datalow` test][flow-test] for an example. You almost never want to use
/// this analysis without also looking at the results of [`MaybeBorrowedLocals`].
///
/// [`MaybeBorrowedLocals`]: ../struct.MaybeBorrowedLocals.html
/// [flow-test]: https://github.com/rust-lang/rust/blob/a08c47310c7d49cbdc5d7afb38408ba519967ecd/src/test/ui/mir-dataflow/liveness-ptr.rs
/// [liveness]: https://en.wikipedia.org/wiki/Live_variable_analysis
pub struct MaybeLiveLocals;

impl MaybeLiveLocals {
    fn transfer_function<T>(&self, trans: &'a mut T) -> TransferFunction<'a, T> {
        TransferFunction(trans)
    }
}

impl BottomValue for MaybeLiveLocals {
    // bottom = not live
    const BOTTOM_VALUE: bool = false;
}

impl AnalysisDomain<'tcx> for MaybeLiveLocals {
    type Idx = Local;
    type Direction = Backward;

    const NAME: &'static str = "liveness";

    fn bits_per_block(&self, body: &mir::Body<'tcx>) -> usize {
        body.local_decls.len()
    }

    fn initialize_start_block(&self, _: &mir::Body<'tcx>, _: &mut BitSet<Self::Idx>) {
        // No variables are live until we observe a use
    }
}

impl GenKillAnalysis<'tcx> for MaybeLiveLocals {
    fn statement_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        self.transfer_function(trans).visit_statement(statement, location);
    }

    fn terminator_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        terminator: &mir::Terminator<'tcx>,
        location: Location,
    ) {
        self.transfer_function(trans).visit_terminator(terminator, location);
    }

    fn call_return_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        _block: mir::BasicBlock,
        _func: &mir::Operand<'tcx>,
        _args: &[mir::Operand<'tcx>],
        dest_place: mir::Place<'tcx>,
    ) {
        if let Some(local) = dest_place.as_local() {
            trans.kill(local);
        }
    }

    fn yield_resume_effect(
        &self,
        trans: &mut impl GenKill<Self::Idx>,
        _resume_block: mir::BasicBlock,
        resume_place: mir::Place<'tcx>,
    ) {
        if let Some(local) = resume_place.as_local() {
            trans.kill(local);
        }
    }
}

struct TransferFunction<'a, T>(&'a mut T);

impl<'tcx, T> Visitor<'tcx> for TransferFunction<'_, T>
where
    T: GenKill<Local>,
{
    fn visit_place(&mut self, place: &mir::Place<'tcx>, context: PlaceContext, location: Location) {
        let mir::Place { projection, local } = *place;

        // We purposefully do not call `super_place` here to avoid calling `visit_local` for this
        // place with one of the `Projection` variants of `PlaceContext`.
        self.visit_projection(local, projection, context, location);

        match DefUse::for_place(context) {
            // Treat derefs as a use of the base local. `*p = 4` is not a def of `p` but a use.
            Some(_) if place.is_indirect() => self.0.gen(local),

            Some(DefUse::Def) if projection.is_empty() => self.0.kill(local),
            Some(DefUse::Use) => self.0.gen(local),
            _ => {}
        }
    }

    fn visit_local(&mut self, &local: &Local, context: PlaceContext, _: Location) {
        // Because we do not call `super_place` above, `visit_local` is only called for locals that
        // do not appear as part of  a `Place` in the MIR. This handles cases like the implicit use
        // of the return place in a `Return` terminator or the index in an `Index` projection.
        match DefUse::for_place(context) {
            Some(DefUse::Def) => self.0.kill(local),
            Some(DefUse::Use) => self.0.gen(local),
            _ => {}
        }
    }
}

#[derive(Eq, PartialEq, Clone)]
enum DefUse {
    Def,
    Use,
}

impl DefUse {
    fn for_place(context: PlaceContext) -> Option<DefUse> {
        match context {
            PlaceContext::NonUse(_) => None,

            PlaceContext::MutatingUse(MutatingUseContext::Store) => Some(DefUse::Def),

            // `MutatingUseContext::Call` and `MutatingUseContext::Yield` indicate that this is the
            // destination place for a `Call` return or `Yield` resume respectively. Since this is
            // only a `Def` when the function returns succesfully, we handle this case separately
            // in `call_return_effect` above.
            PlaceContext::MutatingUse(MutatingUseContext::Call | MutatingUseContext::Yield) => None,

            // All other contexts are uses...
            PlaceContext::MutatingUse(
                MutatingUseContext::AddressOf
                | MutatingUseContext::AsmOutput
                | MutatingUseContext::Borrow
                | MutatingUseContext::Drop
                | MutatingUseContext::Retag,
            )
            | PlaceContext::NonMutatingUse(
                NonMutatingUseContext::AddressOf
                | NonMutatingUseContext::Copy
                | NonMutatingUseContext::Inspect
                | NonMutatingUseContext::Move
                | NonMutatingUseContext::ShallowBorrow
                | NonMutatingUseContext::SharedBorrow
                | NonMutatingUseContext::UniqueBorrow,
            ) => Some(DefUse::Use),

            PlaceContext::MutatingUse(MutatingUseContext::Projection)
            | PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection) => {
                unreachable!("A projection could be a def or a use and must be handled separately")
            }
        }
    }
}
<|endoftext|>


<commit_before>use rustls::ClientConfig;
use tokio_rustls::ClientConfigExt;
use rustls::internal::pemfile;
use rustls::ClientSession;
use hyper::client::{HttpConnector, Connect};
use hyper::client::connect::{Destination, Connected};
use std::io;
use std::fmt;
use std::sync::Arc;
use webpki::{DNSName, DNSNameRef};
use webpki_roots;
use futures::{Future, Poll};
use tokio_rustls::TlsStream;
use tokio::net::TcpStream;

/// Connector for Application-Layer Protocol Negotiation to form a TLS
/// connection for Hyper.
pub struct AlpnConnector {
    tls: Arc<ClientConfig>,
    http: HttpConnector,
}

type AlpnStream = TlsStream<TcpStream, ClientSession>;

impl AlpnConnector {
    /// Construct a new `AlpnConnector`.
    pub fn new() -> Self {
        Self::with_client_config(ClientConfig::new())
    }

    /// Construct a new `AlpnConnector` with a custom certificate and private
    /// key, which should be in PEM format.
    pub fn with_client_cert(
        cert_pem: &[u8],
        key_pem: &[u8],
    ) -> Result<Self, io::Error> {
        let parsed_keys = pemfile::pkcs8_private_keys(&mut io::BufReader::new(key_pem)).or({
            trace!("AlpnConnector::with_client_cert error reading private key");
            Err(io::Error::new(io::ErrorKind::InvalidData, "private key"))
        })?;

        if let Some(key) = parsed_keys.first() {
            let mut config = ClientConfig::new();
            let parsed_cert = pemfile::certs(&mut io::BufReader::new(cert_pem)).or({
                trace!("AlpnConnector::with_client_cert error reading certificate");
                Err(io::Error::new(io::ErrorKind::InvalidData, "certificate"))
            })?;

            config.set_single_client_cert(parsed_cert, key.clone());

            Ok(Self::with_client_config(config))
        } else {
            trace!("AlpnConnector::with_client_cert no private keys found from the given PEM");
            Err(io::Error::new(io::ErrorKind::InvalidData, "private key"))
        }
    }

    fn with_client_config(mut config: ClientConfig) -> Self {
        config.alpn_protocols.push("h2".to_owned());
        config
            .root_store
            .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);

        let mut http = HttpConnector::new(4);
        http.enforce_http(false);

        AlpnConnector {
            tls: Arc::new(config),
            http: http,
        }
    }
}

impl fmt::Debug for AlpnConnector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("AlpnConnector").finish()
    }
}

impl Connect for AlpnConnector {
    type Transport = AlpnStream;
    type Error = io::Error;
    type Future = AlpnConnecting;

    fn connect(&self, dst: Destination) -> Self::Future {
        trace!("AlpnConnector::call ({:?})", dst);

        let host: DNSName = match DNSNameRef::try_from_ascii_str(dst.host()) {
            Ok(host) => host.into(),
            Err(err) => {
                return AlpnConnecting(Box::new(::futures::future::err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("invalid url: {:?}", err),
                ))))
            }
        };

        let tls = self.tls.clone();
        let connecting = self.http
            .connect(dst)
            .and_then(move |(tcp, connected)| {
                trace!("AlpnConnector::call got TCP, trying TLS");
                tls.connect_async(host.as_ref(), tcp)
                    .map_err(|e| {
                        trace!("AlpnConnector::call got error forming a TLS connection.");
                        io::Error::new(io::ErrorKind::Other, e)
                    })
                    .map(move |tls| {
                        (tls, connected)
                    })
            })
            .map_err(|e| {
                trace!("AlpnConnector::call got error reading a TLS stream (#{}).", e);
                io::Error::new(io::ErrorKind::Other, e)
            });

        AlpnConnecting(Box::new(connecting))
    }
}

pub struct AlpnConnecting(Box<Future<Item = (AlpnStream, Connected), Error = io::Error> + Send + 'static>);

impl Future for AlpnConnecting {
    type Item = (AlpnStream, Connected);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

impl fmt::Debug for AlpnConnecting {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("AlpnConnecting")
    }
}
<commit_msg>Import Connect from a different location<commit_after>use rustls::ClientConfig;
use tokio_rustls::ClientConfigExt;
use rustls::internal::pemfile;
use rustls::ClientSession;
use hyper::client::{HttpConnector};
use hyper::client::connect::{Destination, Connected, Connect};
use std::io;
use std::fmt;
use std::sync::Arc;
use webpki::{DNSName, DNSNameRef};
use webpki_roots;
use futures::{Future, Poll};
use tokio_rustls::TlsStream;
use tokio::net::TcpStream;

/// Connector for Application-Layer Protocol Negotiation to form a TLS
/// connection for Hyper.
pub struct AlpnConnector {
    tls: Arc<ClientConfig>,
    http: HttpConnector,
}

type AlpnStream = TlsStream<TcpStream, ClientSession>;

impl AlpnConnector {
    /// Construct a new `AlpnConnector`.
    pub fn new() -> Self {
        Self::with_client_config(ClientConfig::new())
    }

    /// Construct a new `AlpnConnector` with a custom certificate and private
    /// key, which should be in PEM format.
    pub fn with_client_cert(
        cert_pem: &[u8],
        key_pem: &[u8],
    ) -> Result<Self, io::Error> {
        let parsed_keys = pemfile::pkcs8_private_keys(&mut io::BufReader::new(key_pem)).or({
            trace!("AlpnConnector::with_client_cert error reading private key");
            Err(io::Error::new(io::ErrorKind::InvalidData, "private key"))
        })?;

        if let Some(key) = parsed_keys.first() {
            let mut config = ClientConfig::new();
            let parsed_cert = pemfile::certs(&mut io::BufReader::new(cert_pem)).or({
                trace!("AlpnConnector::with_client_cert error reading certificate");
                Err(io::Error::new(io::ErrorKind::InvalidData, "certificate"))
            })?;

            config.set_single_client_cert(parsed_cert, key.clone());

            Ok(Self::with_client_config(config))
        } else {
            trace!("AlpnConnector::with_client_cert no private keys found from the given PEM");
            Err(io::Error::new(io::ErrorKind::InvalidData, "private key"))
        }
    }

    fn with_client_config(mut config: ClientConfig) -> Self {
        config.alpn_protocols.push("h2".to_owned());
        config
            .root_store
            .add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);

        let mut http = HttpConnector::new(4);
        http.enforce_http(false);

        AlpnConnector {
            tls: Arc::new(config),
            http: http,
        }
    }
}

impl fmt::Debug for AlpnConnector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("AlpnConnector").finish()
    }
}

impl Connect for AlpnConnector {
    type Transport = AlpnStream;
    type Error = io::Error;
    type Future = AlpnConnecting;

    fn connect(&self, dst: Destination) -> Self::Future {
        trace!("AlpnConnector::call ({:?})", dst);

        let host: DNSName = match DNSNameRef::try_from_ascii_str(dst.host()) {
            Ok(host) => host.into(),
            Err(err) => {
                return AlpnConnecting(Box::new(::futures::future::err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("invalid url: {:?}", err),
                ))))
            }
        };

        let tls = self.tls.clone();
        let connecting = self.http
            .connect(dst)
            .and_then(move |(tcp, connected)| {
                trace!("AlpnConnector::call got TCP, trying TLS");
                tls.connect_async(host.as_ref(), tcp)
                    .map_err(|e| {
                        trace!("AlpnConnector::call got error forming a TLS connection.");
                        io::Error::new(io::ErrorKind::Other, e)
                    })
                    .map(move |tls| {
                        (tls, connected)
                    })
            })
            .map_err(|e| {
                trace!("AlpnConnector::call got error reading a TLS stream (#{}).", e);
                io::Error::new(io::ErrorKind::Other, e)
            });

        AlpnConnecting(Box::new(connecting))
    }
}

pub struct AlpnConnecting(Box<Future<Item = (AlpnStream, Connected), Error = io::Error> + Send + 'static>);

impl Future for AlpnConnecting {
    type Item = (AlpnStream, Connected);
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

impl fmt::Debug for AlpnConnecting {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("AlpnConnecting")
    }
}
<|endoftext|>


<commit_before>// Parse ANSI attr code
use std::default::Default;
use std::mem;

use beef::lean::Cow;
use tuikit::prelude::*;
use vte::Perform;

/// An ANSI Parser, will parse one line at a time.
///
/// It will cache the latest attribute used, that means if an attribute affect multiple
/// lines, the parser will recognize it.
pub struct ANSIParser {
    partial_str: String,
    last_attr: Attr,

    stripped: String,
    fragments: Vec<(Attr, (u32, u32))>,
}

impl Default for ANSIParser {
    fn default() -> Self {
        ANSIParser {
            partial_str: String::new(),
            last_attr: Attr::default(),

            stripped: String::new(),
            fragments: Vec::new(),
        }
    }
}

impl Perform for ANSIParser {
    fn print(&mut self, ch: char) {
        self.partial_str.push(ch);
    }

    fn execute(&mut self, byte: u8) {
        match byte {
            // \b to delete character back
            0x08 => {
                self.partial_str.pop();
            }
            // put back \0 \r \n \t
            0x00 | 0x0d | 0x0A | 0x09 => self.partial_str.push(byte as char),
            // ignore all others
            _ => trace!("AnsiParser:execute ignored {:?}", byte),
        }
    }

    fn hook(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool) {
        trace!("AnsiParser:hook ignored {:?}", params);
    }

    fn put(&mut self, byte: u8) {
        trace!("AnsiParser:put ignored {:?}", byte);
    }

    fn unhook(&mut self) {
        trace!("AnsiParser:unhook ignored");
    }

    fn osc_dispatch(&mut self, params: &[&[u8]]) {
        trace!("AnsiParser:osc ignored {:?}", params);
    }

    fn csi_dispatch(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool, mode: char) {
        // https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters
        // Only care about graphic modes, ignore all others

        if mode != 'm' {
            trace!("ignore: params: {:?}, mode: {:?}", params, mode);
            return;
        }

        // \[[m => means reset
        let mut attr = if params.is_empty() {
            Attr::default()
        } else {
            self.last_attr
        };

        let mut iter = params.iter();
        while let Some(&code) = iter.next() {
            match code {
                0 => attr = Attr::default(),
                1 => attr.effect |= Effect::BOLD,
                2 => attr.effect |= !Effect::BOLD,
                4 => attr.effect |= Effect::UNDERLINE,
                5 => attr.effect |= Effect::BLINK,
                7 => attr.effect |= Effect::REVERSE,
                num if num >= 30 && num <= 37 => {
                    attr.fg = Color::AnsiValue((num - 30) as u8);
                }
                38 => match iter.next() {
                    Some(2) => {
                        // ESC[ 38;2;<r>;<g>;<b> m Select RGB foreground color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.fg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 38;5;<n> m Select foreground color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.fg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("error on parsing CSI {:?} m", params);
                    }
                },
                39 => attr.fg = Color::Default,
                num if num >= 40 && num <= 47 => {
                    attr.bg = Color::AnsiValue((num - 40) as u8);
                }
                48 => match iter.next() {
                    Some(2) => {
                        // ESC[ 48;2;<r>;<g>;<b> m Select RGB background color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.bg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 48;5;<n> m Select background color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.bg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("ignore CSI {:?} m", params);
                    }
                },
                49 => attr.bg = Color::Default,
                _ => {
                    trace!("ignore CSI {:?} m", params);
                }
            }
        }

        self.attr_change(attr);
    }

    fn esc_dispatch(&mut self, _params: &[i64], _intermediates: &[u8], _ignore: bool, _byte: u8) {
        // ESC characters are replaced with \[
        self.partial_str.push('"');
        self.partial_str.push('[');
    }
}

impl ANSIParser {
    /// save the partial_str into fragments with current attr
    fn save_str(&mut self) {
        if self.partial_str.is_empty() {
            return;
        }

        let string = mem::replace(&mut self.partial_str, String::new());
        self.fragments.push((
            self.last_attr,
            (self.stripped.len() as u32, (self.stripped.len() + string.len()) as u32),
        ));
        self.stripped.push_str(&string);
    }

    // accept a new attr
    fn attr_change(&mut self, new_attr: Attr) {
        if new_attr == self.last_attr {
            return;
        }

        self.save_str();
        self.last_attr = new_attr;
    }

    pub fn parse_ansi(&mut self, text: &str) -> AnsiString<'static> {
        let mut statemachine = vte::Parser::new();

        for byte in text.as_bytes() {
            statemachine.advance(self, *byte);
        }
        self.save_str();

        let stripped = mem::replace(&mut self.stripped, String::new());
        let fragments = mem::replace(&mut self.fragments, Vec::new());
        AnsiString::new(stripped, fragments)
    }
}

#[derive(Clone, Debug)]
/// A String that contains ANSI state (e.g. colors)
///
/// It is internally represented as Vec<(attr, string)>
pub struct AnsiString<'a> {
    stripped: Cow<'a, str>,
    // attr: start, end
    fragments: Option<Vec<(Attr, (u32, u32))>>,
}

impl<'a> AnsiString<'a> {
    pub fn new_empty() -> Self {
        Self {
            stripped: Cow::borrowed(""),
            fragments: None,
        }
    }

    fn new_string(string: String) -> Self {
        Self {
            stripped: Cow::owned(string),
            fragments: None,
        }
    }

    fn new_str(str_ref: &'a str) -> Self {
        Self {
            stripped: Cow::borrowed(str_ref),
            fragments: None,
        }
    }

    fn new(stripped: String, fragments: Vec<(Attr, (u32, u32))>) -> Self {
        let fragments_empty = fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Attr::default());
        Self {
            stripped: Cow::owned(stripped),
            fragments: if fragments_empty { None } else { Some(fragments) },
        }
    }

    pub fn parse(raw: &'a str) -> AnsiString<'static> {
        ANSIParser::default().parse_ansi(raw)
    }

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

    #[inline]
    pub fn into_inner(self) -> std::borrow::Cow<'a, str> {
        std::borrow::Cow::Owned(self.stripped.into_owned())
    }

    pub fn iter(&'a self) -> Box<dyn Iterator<Item = (char, Attr)> + 'a> {
        if self.fragments.is_none() {
            return Box::new(self.stripped.chars().map(|c| (c, Attr::default())));
        }

        Box::new(AnsiStringIterator::new(
            &self.stripped,
            self.fragments.as_ref().unwrap(),
        ))
    }

    pub fn has_attrs(&self) -> bool {
        self.fragments.is_some()
    }

    #[inline]
    pub fn stripped(&self) -> &str {
        &self.stripped
    }
}

impl<'a> From<&'a str> for AnsiString<'a> {
    fn from(s: &'a str) -> AnsiString<'a> {
        AnsiString::new_str(s)
    }
}

impl From<String> for AnsiString<'static> {
    fn from(s: String) -> Self {
        AnsiString::new_string(s)
    }
}

/// An iterator over all the (char, attr) characters.
pub struct AnsiStringIterator<'a> {
    fragments: &'a [(Attr, (u32, u32))],
    fragment_idx: usize,
    chars_iter: std::str::CharIndices<'a>,
}

impl<'a> AnsiStringIterator<'a> {
    pub fn new(stripped: &'a str, fragments: &'a [(Attr, (u32, u32))]) -> Self {
        Self {
            fragments,
            fragment_idx: 0,
            chars_iter: stripped.char_indices(),
        }
    }
}

impl<'a> Iterator for AnsiStringIterator<'a> {
    type Item = (char, Attr);

    fn next(&mut self) -> Option<Self::Item> {
        match self.chars_iter.next() {
            Some((byte_idx, char)) => {
                // update fragment_idx
                loop {
                    if self.fragment_idx >= self.fragments.len() {
                        return None;
                    }

                    let (_attr, (_start, end)) = self.fragments[self.fragment_idx];
                    if byte_idx < (end as usize) {
                        break;
                    } else {
                        self.fragment_idx += 1;
                    }
                }

                let (attr, _) = self.fragments[self.fragment_idx];
                Some((char, attr))
            }
            None => None,
        }
    }
}

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

    #[test]
    fn test_ansi_iterator() {
        let input = "\x1B[48;2;5;10;15m\x1B[38;2;70;130;180mhi\x1B[0m";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::Rgb(70, 130, 180),
            bg: Color::Rgb(5, 10, 15),
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_normal_string() {
        let input = "ab";
        let ansistring = ANSIParser::default().parse_ansi(input);

        assert_eq!(false, ansistring.has_attrs());

        let mut it = ansistring.iter();
        assert_eq!(Some(('a', Attr::default())), it.next());
        assert_eq!(Some(('b', Attr::default())), it.next());
        assert_eq!(None, it.next());

        assert_eq!(ansistring.stripped(), "ab");
    }

    #[test]
    fn test_multiple_attributes() {
        let input = "\x1B[1;31mhi";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::AnsiValue(1),
            effect: Effect::BOLD,
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_reset() {
        let input = "\x1B[35mA\x1B[mB";
        let ansistring = ANSIParser::default().parse_ansi(input);
        assert_eq!(ansistring.fragments.as_ref().map(|x| x.len()).unwrap(), 2);
        assert_eq!(ansistring.stripped(), "AB");
    }
}
<commit_msg>[ansi] support highlight-only indices<commit_after>// Parse ANSI attr code
use std::default::Default;
use std::mem;

use beef::lean::Cow;
use std::str::Chars;
use tuikit::prelude::*;
use vte::Perform;

/// An ANSI Parser, will parse one line at a time.
///
/// It will cache the latest attribute used, that means if an attribute affect multiple
/// lines, the parser will recognize it.
pub struct ANSIParser {
    partial_str: String,
    last_attr: Attr,

    stripped: String,
    fragments: Vec<(Attr, (u32, u32))>,
}

impl Default for ANSIParser {
    fn default() -> Self {
        ANSIParser {
            partial_str: String::new(),
            last_attr: Attr::default(),

            stripped: String::new(),
            fragments: Vec::new(),
        }
    }
}

impl Perform for ANSIParser {
    fn print(&mut self, ch: char) {
        self.partial_str.push(ch);
    }

    fn execute(&mut self, byte: u8) {
        match byte {
            // \b to delete character back
            0x08 => {
                self.partial_str.pop();
            }
            // put back \0 \r \n \t
            0x00 | 0x0d | 0x0A | 0x09 => self.partial_str.push(byte as char),
            // ignore all others
            _ => trace!("AnsiParser:execute ignored {:?}", byte),
        }
    }

    fn hook(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool) {
        trace!("AnsiParser:hook ignored {:?}", params);
    }

    fn put(&mut self, byte: u8) {
        trace!("AnsiParser:put ignored {:?}", byte);
    }

    fn unhook(&mut self) {
        trace!("AnsiParser:unhook ignored");
    }

    fn osc_dispatch(&mut self, params: &[&[u8]]) {
        trace!("AnsiParser:osc ignored {:?}", params);
    }

    fn csi_dispatch(&mut self, params: &[i64], _intermediates: &[u8], _ignore: bool, mode: char) {
        // https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters
        // Only care about graphic modes, ignore all others

        if mode != 'm' {
            trace!("ignore: params: {:?}, mode: {:?}", params, mode);
            return;
        }

        // \[[m => means reset
        let mut attr = if params.is_empty() {
            Attr::default()
        } else {
            self.last_attr
        };

        let mut iter = params.iter();
        while let Some(&code) = iter.next() {
            match code {
                0 => attr = Attr::default(),
                1 => attr.effect |= Effect::BOLD,
                2 => attr.effect |= !Effect::BOLD,
                4 => attr.effect |= Effect::UNDERLINE,
                5 => attr.effect |= Effect::BLINK,
                7 => attr.effect |= Effect::REVERSE,
                num if num >= 30 && num <= 37 => {
                    attr.fg = Color::AnsiValue((num - 30) as u8);
                }
                38 => match iter.next() {
                    Some(2) => {
                        // ESC[ 38;2;<r>;<g>;<b> m Select RGB foreground color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.fg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 38;5;<n> m Select foreground color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.fg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("error on parsing CSI {:?} m", params);
                    }
                },
                39 => attr.fg = Color::Default,
                num if num >= 40 && num <= 47 => {
                    attr.bg = Color::AnsiValue((num - 40) as u8);
                }
                48 => match iter.next() {
                    Some(2) => {
                        // ESC[ 48;2;<r>;<g>;<b> m Select RGB background color
                        let or = iter.next();
                        let og = iter.next();
                        let ob = iter.next();
                        if ob.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }

                        let r = *or.unwrap() as u8;
                        let g = *og.unwrap() as u8;
                        let b = *ob.unwrap() as u8;

                        attr.bg = Color::Rgb(r, g, b);
                    }
                    Some(5) => {
                        // ESC[ 48;5;<n> m Select background color
                        let color = iter.next();
                        if color.is_none() {
                            trace!("ignore CSI {:?} m", params);
                            continue;
                        }
                        attr.bg = Color::AnsiValue(*color.unwrap() as u8);
                    }
                    _ => {
                        trace!("ignore CSI {:?} m", params);
                    }
                },
                49 => attr.bg = Color::Default,
                _ => {
                    trace!("ignore CSI {:?} m", params);
                }
            }
        }

        self.attr_change(attr);
    }

    fn esc_dispatch(&mut self, _params: &[i64], _intermediates: &[u8], _ignore: bool, _byte: u8) {
        // ESC characters are replaced with \[
        self.partial_str.push('"');
        self.partial_str.push('[');
    }
}

impl ANSIParser {
    /// save the partial_str into fragments with current attr
    fn save_str(&mut self) {
        if self.partial_str.is_empty() {
            return;
        }

        let string = mem::replace(&mut self.partial_str, String::new());
        self.fragments.push((
            self.last_attr,
            (self.stripped.len() as u32, (self.stripped.len() + string.len()) as u32),
        ));
        self.stripped.push_str(&string);
    }

    // accept a new attr
    fn attr_change(&mut self, new_attr: Attr) {
        if new_attr == self.last_attr {
            return;
        }

        self.save_str();
        self.last_attr = new_attr;
    }

    pub fn parse_ansi(&mut self, text: &str) -> AnsiString<'static> {
        let mut statemachine = vte::Parser::new();

        for byte in text.as_bytes() {
            statemachine.advance(self, *byte);
        }
        self.save_str();

        let stripped = mem::replace(&mut self.stripped, String::new());
        let fragments = mem::replace(&mut self.fragments, Vec::new());
        AnsiString::new_string(stripped, fragments)
    }
}

/// A String that contains ANSI state (e.g. colors)
///
/// It is internally represented as Vec<(attr, string)>
#[derive(Clone, Debug)]
pub struct AnsiString<'a> {
    stripped: Cow<'a, str>,
    // attr: start, end
    fragments: Option<Vec<(Attr, (u32, u32))>>,
}

impl<'a> AnsiString<'a> {
    pub fn new_empty() -> Self {
        Self {
            stripped: Cow::borrowed(""),
            fragments: None,
        }
    }

    fn new_raw_string(string: String) -> Self {
        Self {
            stripped: Cow::owned(string),
            fragments: None,
        }
    }

    fn new_raw_str(str_ref: &'a str) -> Self {
        Self {
            stripped: Cow::borrowed(str_ref),
            fragments: None,
        }
    }

    pub fn new_str(stripped: &'a str, fragments: Vec<(Attr, (u32, u32))>) -> Self {
        let fragments_empty = fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Attr::default());
        Self {
            stripped: Cow::borrowed(stripped),
            fragments: if fragments_empty { None } else { Some(fragments) },
        }
    }

    pub fn new_string(stripped: String, fragments: Vec<(Attr, (u32, u32))>) -> Self {
        let fragments_empty = fragments.is_empty() || (fragments.len() == 1 && fragments[0].0 == Attr::default());
        Self {
            stripped: Cow::owned(stripped),
            fragments: if fragments_empty { None } else { Some(fragments) },
        }
    }

    pub fn parse(raw: &'a str) -> AnsiString<'static> {
        ANSIParser::default().parse_ansi(raw)
    }

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

    #[inline]
    pub fn into_inner(self) -> std::borrow::Cow<'a, str> {
        std::borrow::Cow::Owned(self.stripped.into_owned())
    }

    pub fn iter(&'a self) -> Box<dyn Iterator<Item = (char, Attr)> + 'a> {
        if self.fragments.is_none() {
            return Box::new(self.stripped.chars().map(|c| (c, Attr::default())));
        }

        Box::new(AnsiStringIterator::new(
            &self.stripped,
            self.fragments.as_ref().unwrap(),
        ))
    }

    pub fn has_attrs(&self) -> bool {
        self.fragments.is_some()
    }

    #[inline]
    pub fn stripped(&self) -> &str {
        &self.stripped
    }
}

impl<'a> From<&'a str> for AnsiString<'a> {
    fn from(s: &'a str) -> AnsiString<'a> {
        AnsiString::new_raw_str(s)
    }
}

impl From<String> for AnsiString<'static> {
    fn from(s: String) -> Self {
        AnsiString::new_raw_string(s)
    }
}

// (text, indices, highlight attribute) -> AnsiString
impl<'a> From<(&'a str, &'a [usize], Attr)> for AnsiString<'a> {
    fn from((text, indices, attr): (&'a str, &'a [usize], Attr)) -> Self {
        let fragments = indices
            .iter()
            .map(|&idx| (attr, (idx as u32, 1 + idx as u32)))
            .collect();
        AnsiString::new_str(text, fragments)
    }
}

/// An iterator over all the (char, attr) characters.
pub struct AnsiStringIterator<'a> {
    fragments: &'a [(Attr, (u32, u32))],
    fragment_idx: usize,
    chars_iter: std::str::CharIndices<'a>,
}

impl<'a> AnsiStringIterator<'a> {
    pub fn new(stripped: &'a str, fragments: &'a [(Attr, (u32, u32))]) -> Self {
        Self {
            fragments,
            fragment_idx: 0,
            chars_iter: stripped.char_indices(),
        }
    }
}

impl<'a> Iterator for AnsiStringIterator<'a> {
    type Item = (char, Attr);

    fn next(&mut self) -> Option<Self::Item> {
        match self.chars_iter.next() {
            Some((char_idx, char)) => {
                // update fragment_idx
                loop {
                    if self.fragment_idx >= self.fragments.len() {
                        break;
                    }

                    let (_attr, (_start, end)) = self.fragments[self.fragment_idx];
                    if char_idx < (end as usize) {
                        break;
                    } else {
                        self.fragment_idx += 1;
                    }
                }

                let (attr, (start, end)) = if self.fragment_idx >= self.fragments.len() {
                    (Attr::default(), (char_idx as u32, 1 + char_idx as u32))
                } else {
                    self.fragments[self.fragment_idx]
                };

                if (start as usize) <= char_idx && char_idx < (end as usize) {
                    Some((char, attr))
                } else {
                    Some((char, Attr::default()))
                }
            }
            None => None,
        }
    }
}

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

    #[test]
    fn test_ansi_iterator() {
        let input = "\x1B[48;2;5;10;15m\x1B[38;2;70;130;180mhi\x1B[0m";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::Rgb(70, 130, 180),
            bg: Color::Rgb(5, 10, 15),
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_highlight_indices() {
        let text = "abc";
        let indices: Vec<usize> = vec![1];
        let attr = Attr {
            fg: Color::Rgb(70, 130, 180),
            bg: Color::Rgb(5, 10, 15),
            ..Attr::default()
        };

        let ansistring = AnsiString::from((text, &indices as &[usize], attr));
        let mut it = ansistring.iter();

        assert_eq!(Some(('a', Attr::default())), it.next());
        assert_eq!(Some(('b', attr)), it.next());
        assert_eq!(Some(('c', Attr::default())), it.next());
        assert_eq!(None, it.next());
    }

    #[test]
    fn test_normal_string() {
        let input = "ab";
        let ansistring = ANSIParser::default().parse_ansi(input);

        assert_eq!(false, ansistring.has_attrs());

        let mut it = ansistring.iter();
        assert_eq!(Some(('a', Attr::default())), it.next());
        assert_eq!(Some(('b', Attr::default())), it.next());
        assert_eq!(None, it.next());

        assert_eq!(ansistring.stripped(), "ab");
    }

    #[test]
    fn test_multiple_attributes() {
        let input = "\x1B[1;31mhi";
        let ansistring = ANSIParser::default().parse_ansi(input);
        let mut it = ansistring.iter();
        let attr = Attr {
            fg: Color::AnsiValue(1),
            effect: Effect::BOLD,
            ..Attr::default()
        };

        assert_eq!(Some(('h', attr)), it.next());
        assert_eq!(Some(('i', attr)), it.next());
        assert_eq!(None, it.next());
        assert_eq!(ansistring.stripped(), "hi");
    }

    #[test]
    fn test_reset() {
        let input = "\x1B[35mA\x1B[mB";
        let ansistring = ANSIParser::default().parse_ansi(input);
        assert_eq!(ansistring.fragments.as_ref().map(|x| x.len()).unwrap(), 2);
        assert_eq!(ansistring.stripped(), "AB");
    }
}
<|endoftext|>


<commit_before>use std::io::stdout;
use std::path::PathBuf;

use structopt::clap::{AppSettings, Shell};
use structopt::StructOpt;

use anyhow::Result;
use lazy_static::lazy_static;
use strum::VariantNames;
use strum_macros::{EnumString, EnumVariantNames, ToString};

#[derive(Debug, StructOpt)]
#[structopt(about="A utility like pkg-audit for Arch Linux.", global_settings = &[AppSettings::ColoredHelp, AppSettings::DeriveDisplayOrder])]
pub struct Args {
    /// Show only vulnerable package names and their versions. Set twice to hide the versions as well.
    #[structopt(long, short = "q", parse(from_occurrences))]
    pub quiet: u8,
    /// Prints packages that depend on vulnerable packages and are thus potentially vulnerable as well. Set twice to show ALL the packages that requires them.
    #[structopt(long, short = "r", parse(from_occurrences))]
    pub recursive: u8,
    /// Show packages which are in the [testing] repos. See https://wiki.archlinux.org/index.php/Official_repositories#Testing_repositories
    #[structopt(long = "show-testing", short = "t")]
    pub testing: bool,
    /// Show only packages that have already been fixed
    #[structopt(long, short = "u", long)]
    pub upgradable: bool,
    /// Bypass tty detection for colors
    #[structopt(long, short = "C", default_value = "auto", possible_values=&Color::VARIANTS)]
    pub color: Color,
    /// Set an alternate database location
    #[structopt(
        long,
        short = "b",
        parse(from_os_str),
        default_value = "/var/lib/pacman"
    )]
    pub dbpath: PathBuf,
    /// Specify a format to control the output. Placeholders are %n (pkgname), %c (CVEs), %v (fixed version), %t (type), %s (severity), and %r (required by, only when -r is also used).
    #[structopt(long, short = "f")]
    pub format: Option<String>,
    /// Print json output
    #[structopt(long)]
    pub json: bool,
    /// Specify the URL or file path to the security tracker json data
    #[structopt(long)]
    pub source: Option<String>,
    /// Send requests through a proxy
    #[structopt(long)]
    pub proxy: Option<String>,
    /// Do not use a proxy even if one is configured
    #[structopt(long)]
    pub no_proxy: bool,
    /// Specify how to sort the output
    #[structopt(long, use_delimiter = true, possible_values = &SortBy::VARIANTS, default_value = &SORT_BY_DEFAULT_VALUE)]
    pub sort: Vec<SortBy>,
    /// Print the CVE numbers.
    #[structopt(long, short = "c")]
    pub show_cve: bool,
    #[structopt(subcommand)]
    pub subcommand: Option<SubCommand>,
}

#[derive(Debug, StructOpt)]
pub enum SubCommand {
    /// Generate shell completions
    #[structopt(name = "completions")]
    Completions(Completions),
}

#[derive(Debug, StructOpt, ToString, EnumString, EnumVariantNames)]
#[strum(serialize_all = "lowercase")]
pub enum Color {
    Auto,
    Always,
    Never,
}

impl Default for Color {
    fn default() -> Self {
        Self::Auto
    }
}

#[derive(Debug, StructOpt, ToString, EnumString, EnumVariantNames)]
#[strum(serialize_all = "snake_case")]
pub enum SortBy {
    Severity,
    Pkgname,
    Upgradable,
    Reverse,
}

lazy_static! {
    static ref SORT_BY_DEFAULT_VALUE: String = vec![SortBy::Severity, SortBy::Pkgname,]
        .iter()
        .map(|e| e.to_string())
        .collect::<Vec<String>>()
        .join(",");
}

#[derive(Debug, StructOpt)]
pub struct Completions {
    #[structopt(possible_values=&Shell::variants())]
    pub shell: Shell,
}

pub fn gen_completions(args: &Completions) -> Result<()> {
    Args::clap().gen_completions_to("dfrs", args.shell, &mut stdout());
    Ok(())
}
<commit_msg>fix(completion): fix gen completions to target<commit_after>use std::io::stdout;
use std::path::PathBuf;

use structopt::clap::{AppSettings, Shell};
use structopt::StructOpt;

use anyhow::Result;
use lazy_static::lazy_static;
use strum::VariantNames;
use strum_macros::{EnumString, EnumVariantNames, ToString};

#[derive(Debug, StructOpt)]
#[structopt(about="A utility like pkg-audit for Arch Linux.", global_settings = &[AppSettings::ColoredHelp, AppSettings::DeriveDisplayOrder])]
pub struct Args {
    /// Show only vulnerable package names and their versions. Set twice to hide the versions as well.
    #[structopt(long, short = "q", parse(from_occurrences))]
    pub quiet: u8,
    /// Prints packages that depend on vulnerable packages and are thus potentially vulnerable as well. Set twice to show ALL the packages that requires them.
    #[structopt(long, short = "r", parse(from_occurrences))]
    pub recursive: u8,
    /// Show packages which are in the [testing] repos. See https://wiki.archlinux.org/index.php/Official_repositories#Testing_repositories
    #[structopt(long = "show-testing", short = "t")]
    pub testing: bool,
    /// Show only packages that have already been fixed
    #[structopt(long, short = "u", long)]
    pub upgradable: bool,
    /// Bypass tty detection for colors
    #[structopt(long, short = "C", default_value = "auto", possible_values=&Color::VARIANTS)]
    pub color: Color,
    /// Set an alternate database location
    #[structopt(
        long,
        short = "b",
        parse(from_os_str),
        default_value = "/var/lib/pacman"
    )]
    pub dbpath: PathBuf,
    /// Specify a format to control the output. Placeholders are %n (pkgname), %c (CVEs), %v (fixed version), %t (type), %s (severity), and %r (required by, only when -r is also used).
    #[structopt(long, short = "f")]
    pub format: Option<String>,
    /// Print json output
    #[structopt(long)]
    pub json: bool,
    /// Specify the URL or file path to the security tracker json data
    #[structopt(long)]
    pub source: Option<String>,
    /// Send requests through a proxy
    #[structopt(long)]
    pub proxy: Option<String>,
    /// Do not use a proxy even if one is configured
    #[structopt(long)]
    pub no_proxy: bool,
    /// Specify how to sort the output
    #[structopt(long, use_delimiter = true, possible_values = &SortBy::VARIANTS, default_value = &SORT_BY_DEFAULT_VALUE)]
    pub sort: Vec<SortBy>,
    /// Print the CVE numbers.
    #[structopt(long, short = "c")]
    pub show_cve: bool,
    #[structopt(subcommand)]
    pub subcommand: Option<SubCommand>,
}

#[derive(Debug, StructOpt)]
pub enum SubCommand {
    /// Generate shell completions
    #[structopt(name = "completions")]
    Completions(Completions),
}

#[derive(Debug, StructOpt, ToString, EnumString, EnumVariantNames)]
#[strum(serialize_all = "lowercase")]
pub enum Color {
    Auto,
    Always,
    Never,
}

impl Default for Color {
    fn default() -> Self {
        Self::Auto
    }
}

#[derive(Debug, StructOpt, ToString, EnumString, EnumVariantNames)]
#[strum(serialize_all = "snake_case")]
pub enum SortBy {
    Severity,
    Pkgname,
    Upgradable,
    Reverse,
}

lazy_static! {
    static ref SORT_BY_DEFAULT_VALUE: String = vec![SortBy::Severity, SortBy::Pkgname,]
        .iter()
        .map(|e| e.to_string())
        .collect::<Vec<String>>()
        .join(",");
}

#[derive(Debug, StructOpt)]
pub struct Completions {
    #[structopt(possible_values=&Shell::variants())]
    pub shell: Shell,
}

pub fn gen_completions(args: &Completions) -> Result<()> {
    Args::clap().gen_completions_to("arch-audit", args.shell, &mut stdout());
    Ok(())
}
<|endoftext|>