<commit_before>mod reader;
mod writer;

use tracker::{self, Announce, Response, TrackerResponse, Result, ResultExt, Error, ErrorKind, dns};
use std::time::{Instant, Duration};
use std::mem;
use std::sync::Arc;
use {PEER_ID, bencode, amy};
use self::writer::Writer;
use self::reader::Reader;
use std::collections::HashMap;
use std::net::SocketAddr;
use url::percent_encoding::{percent_encode_byte};
use url::Url;
use slog::Logger;
use socket::TSocket;

const TIMEOUT_MS: u64 = 2500;

pub struct Handler {
    reg: Arc<amy::Registrar>,
    connections: HashMap<usize, Tracker>,
    l: Logger
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    last_updated: Instant,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS { sock: TSocket, req: Vec<u8>, port: u16 },
    Writing { sock: TSocket, writer: Writer },
    Reading { sock: TSocket, reader: Reader },
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16 ) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<Option<TrackerResponse>> {
        let s = mem::replace(self, TrackerState::Error);
        let n = s.next(event)?;
        if let TrackerState::Complete(r) = n {
            Ok(Some(r))
        } else {
            mem::replace(self, n);
            Ok(None)
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (TrackerState::ResolvingDNS { sock, req, port }, Event::DNSResolved(r)) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr);
                Ok(TrackerState::Writing { sock, writer: Writer::new(req) }.next(Event::Writable)?)
            }
            (TrackerState::Writing { mut sock, mut writer }, Event::Writable) => {
                match writer.writable(&mut sock.conn)? {
                    Some(()) => {
                        let r = Reader::new();
                        Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                    }
                    None => {
                        Ok(TrackerState::Writing { sock, writer })
                    }
                }
            }
            (TrackerState::Reading { mut sock, mut reader }, Event::Readable) => {
                if reader.readable(&mut sock.conn)? {
                    let data = reader.consume();
                    let content = bencode::decode_buf(&data).chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                } else {
                    Ok(TrackerState::Reading { sock, reader })
                }
            }
            (s @ TrackerState::Writing { .. }, _) => Ok(s),
            (s @ TrackerState::Reading { .. }, _) => Ok(s),
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!")
        }
    }
}

impl Handler {
    pub fn new(reg: Arc<amy::Registrar>, l: Logger) -> Handler {
        Handler { reg, connections: HashMap::new(), l }
    }

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn readable(&mut self, id: usize) -> Option<Response> {
        debug!(self.l, "Announce reading: {:?}", id);
        if let Some(mut trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(Some(r)) => {
                    // TODO: deregister socket here
                    debug!(self.l, "Annoucne response received for {:?}, {:?}", id, r);
                    return Some(((trk.torrent, Ok(r))))
                }
                Ok(None) => { }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        debug!(self.l, "Announce writing: {:?}", id);
        if let Some(mut trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => {  }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        self.connections.retain(|_, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                resps.push((trk.torrent, Err(ErrorKind::Timeout.into())));
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        debug!(self.l, "Received a DNS resp for {:?}", resp.id);
        if let Some(mut trk) = self.connections.get_mut(&resp.id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => { }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn new_announce(&mut self, req: Announce, url: &Url, dns: &mut dns::Resolver) -> Result<()> {
        debug!(self.l, "Received a new announce req for {:?}", url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL:
        http_req.extend_from_slice(url.path().as_bytes());
        // The fact that I have to do this is genuinely depressing.
        // This will be rewritten as a proper http protocol
        // encoder in an event loop.
        http_req.extend_from_slice("?".as_bytes());
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "numwant", "50");
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "numwant", "20");
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {
                append_query_pair(&mut http_req, "numwant", "20");
            }
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str().ok_or::<Error>(
            ErrorKind::InvalidRequest(format!("Tracker announce url has no host!")).into()
        )?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        // Setup actual connection
        let (id, sock) = TSocket::new_v4(self.reg.clone()).chain_err(|| ErrorKind::IO)?;
        dns.new_query(id, host);
        self.connections.insert(id, Tracker {
            last_updated: Instant::now(),
            torrent: req.id,
            state: TrackerState::new(sock, http_req, port),
        });
        debug!(self.l, "Dispatching DNS req, id {:?}", id);

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice("=".as_bytes());
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice("&".as_bytes());
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        resp.push_str(percent_encode_byte(*byte));
    }
    resp
}
<commit_msg>Minor cleanups to code and comments<commit_after>mod reader;
mod writer;

use tracker::{self, Announce, Response, TrackerResponse, Result, ResultExt, Error, ErrorKind, dns};
use std::time::{Instant, Duration};
use std::mem;
use std::sync::Arc;
use {PEER_ID, bencode, amy};
use self::writer::Writer;
use self::reader::Reader;
use std::collections::HashMap;
use std::net::SocketAddr;
use url::percent_encoding::{percent_encode_byte};
use url::Url;
use slog::Logger;
use socket::TSocket;

const TIMEOUT_MS: u64 = 2500;

pub struct Handler {
    reg: Arc<amy::Registrar>,
    connections: HashMap<usize, Tracker>,
    l: Logger
}

enum Event {
    DNSResolved(dns::QueryResponse),
    Readable,
    Writable,
}

struct Tracker {
    torrent: usize,
    last_updated: Instant,
    state: TrackerState,
}

enum TrackerState {
    Error,
    ResolvingDNS { sock: TSocket, req: Vec<u8>, port: u16 },
    Writing { sock: TSocket, writer: Writer },
    Reading { sock: TSocket, reader: Reader },
    Complete(TrackerResponse),
}

impl TrackerState {
    fn new(sock: TSocket, req: Vec<u8>, port: u16 ) -> TrackerState {
        TrackerState::ResolvingDNS { sock, req, port }
    }

    fn handle(&mut self, event: Event) -> Result<Option<TrackerResponse>> {
        let s = mem::replace(self, TrackerState::Error);
        let n = s.next(event)?;
        if let TrackerState::Complete(r) = n {
            Ok(Some(r))
        } else {
            mem::replace(self, n);
            Ok(None)
        }
    }

    fn next(self, event: Event) -> Result<TrackerState> {
        match (self, event) {
            (TrackerState::ResolvingDNS { sock, req, port }, Event::DNSResolved(r)) => {
                let addr = SocketAddr::new(r.res?, port);
                sock.connect(addr);
                Ok(TrackerState::Writing { sock, writer: Writer::new(req) }.next(Event::Writable)?)
            }
            (TrackerState::Writing { mut sock, mut writer }, Event::Writable) => {
                match writer.writable(&mut sock.conn)? {
                    Some(()) => {
                        let r = Reader::new();
                        Ok(TrackerState::Reading { sock, reader: r }.next(Event::Readable)?)
                    }
                    None => {
                        Ok(TrackerState::Writing { sock, writer })
                    }
                }
            }
            (TrackerState::Reading { mut sock, mut reader }, Event::Readable) => {
                if reader.readable(&mut sock.conn)? {
                    let data = reader.consume();
                    let content = bencode::decode_buf(&data).chain_err(|| ErrorKind::InvalidResponse("Invalid BEncoded response!"))?;
                    let resp = TrackerResponse::from_bencode(content)?;
                    Ok(TrackerState::Complete(resp))
                } else {
                    Ok(TrackerState::Reading { sock, reader })
                }
            }
            (s @ TrackerState::Writing { .. }, _) => Ok(s),
            (s @ TrackerState::Reading { .. }, _) => Ok(s),
            (s @ TrackerState::ResolvingDNS { .. }, _) => Ok(s),
            _ => bail!("Unknown state transition encountered!")
        }
    }
}

impl Handler {
    pub fn new(reg: Arc<amy::Registrar>, l: Logger) -> Handler {
        Handler { reg, connections: HashMap::new(), l }
    }

    pub fn contains(&self, id: usize) -> bool {
        self.connections.contains_key(&id)
    }

    pub fn readable(&mut self, id: usize) -> Option<Response> {
        if let Some(mut trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Readable) {
                Ok(Some(r)) => {
                    debug!(self.l, "Announce response received for {:?} succesfully", id);
                    return Some(((trk.torrent, Ok(r))))
                }
                Ok(None) => { }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn writable(&mut self, id: usize) -> Option<Response> {
        if let Some(mut trk) = self.connections.get_mut(&id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::Writable) {
                Ok(_) => { }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn tick(&mut self) -> Vec<Response> {
        let mut resps = Vec::new();
        let ref l = self.l;
        self.connections.retain(|id, trk| {
            if trk.last_updated.elapsed() > Duration::from_millis(TIMEOUT_MS) {
                resps.push((trk.torrent, Err(ErrorKind::Timeout.into())));
                debug!(l, "Announce {:?} timed out", id);
                false
            } else {
                true
            }
        });
        resps
    }

    pub fn dns_resolved(&mut self, resp: dns::QueryResponse) -> Option<Response> {
        debug!(self.l, "Received a DNS resp for {:?}", resp.id);
        if let Some(mut trk) = self.connections.get_mut(&resp.id) {
            trk.last_updated = Instant::now();
            match trk.state.handle(Event::DNSResolved(resp)) {
                Ok(_) => { }
                Err(e) => {
                    return Some((trk.torrent, Err(e)));
                }
            }
        }
        None
    }

    pub fn new_announce(&mut self, req: Announce, url: &Url, dns: &mut dns::Resolver) -> Result<()> {
        debug!(self.l, "Received a new announce req for {:?}", url);
        let mut http_req = Vec::with_capacity(50);
        // Encode GET req
        http_req.extend_from_slice(b"GET ");

        // Encode the URL
        http_req.extend_from_slice(url.path().as_bytes());
        http_req.extend_from_slice("?".as_bytes());
        append_query_pair(&mut http_req, "info_hash", &encode_param(&req.hash));
        append_query_pair(&mut http_req, "peer_id", &encode_param(&PEER_ID[..]));
        append_query_pair(&mut http_req, "uploaded", &req.uploaded.to_string());
        append_query_pair(&mut http_req, "downloaded", &req.downloaded.to_string());
        append_query_pair(&mut http_req, "left", &req.left.to_string());
        append_query_pair(&mut http_req, "compact", "1");
        append_query_pair(&mut http_req, "port", &req.port.to_string());
        match req.event {
            Some(tracker::Event::Started) => {
                append_query_pair(&mut http_req, "numwant", "50");
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Stopped) => {
                append_query_pair(&mut http_req, "event", "started");
            }
            Some(tracker::Event::Completed) => {
                append_query_pair(&mut http_req, "numwant", "20");
                append_query_pair(&mut http_req, "event", "completed");
            }
            None => {
                append_query_pair(&mut http_req, "numwant", "20");
            }
        }

        // Encode HTTP protocol
        http_req.extend_from_slice(b" HTTP/1.1\r\n");
        // Encode host header
        http_req.extend_from_slice(b"Host: ");
        let host = url.host_str().ok_or::<Error>(
            ErrorKind::InvalidRequest(format!("Tracker announce url has no host!")).into()
        )?;
        let port = url.port().unwrap_or(80);
        http_req.extend_from_slice(host.as_bytes());
        http_req.extend_from_slice(b"\r\n");
        // Encode empty line to terminate request
        http_req.extend_from_slice(b"\r\n");

        // Setup actual connection and start DNS query
        let (id, sock) = TSocket::new_v4(self.reg.clone()).chain_err(|| ErrorKind::IO)?;
        self.connections.insert(id, Tracker {
            last_updated: Instant::now(),
            torrent: req.id,
            state: TrackerState::new(sock, http_req, port),
        });

        debug!(self.l, "Dispatching DNS req, id {:?}", id);
        dns.new_query(id, host);

        Ok(())
    }
}

fn append_query_pair(s: &mut Vec<u8>, k: &str, v: &str) {
    s.extend_from_slice(k.as_bytes());
    s.extend_from_slice("=".as_bytes());
    s.extend_from_slice(v.as_bytes());
    s.extend_from_slice("&".as_bytes());
}

fn encode_param(data: &[u8]) -> String {
    let mut resp = String::new();
    for byte in data {
        resp.push_str(percent_encode_byte(*byte));
    }
    resp
}
<|endoftext|>


<commit_before>extern crate jvmti;

#[cfg(test)]
mod tests {

    use jvmti::bytecode::class_stream::ClassStream;
    use jvmti::bytecode::constants::ConstantType;

    #[test]
    fn read_magic_bytes_accepts_only_valid_magic_bytes() {
        let valid_magic = vec![ 0xCA, 0xFE, 0xBA, 0xBE ];
        let invalid_magics = vec![
            vec![],
            vec![ 0x00, ],
            vec![ 0xCA, 0xFE, 0xBA ],
            vec![ 0xCA, 0xFE, 0xBE, 0xBA ]
        ];

        assert_eq!(true, ClassStream::new(&valid_magic).read_magic_bytes());

        for magic in invalid_magics {
            assert_eq!(false, ClassStream::new(&magic).read_magic_bytes());
        }
    }

    #[test]
    fn read_version_number_accepts_only_valid_version_numbers() {
        let valid_versions = vec![
            vec![ 0x00, 0x10, 0x00, 0x30 ],
            vec![ 0x00, 0x00, 0x00, 0x34 ], // Java 8 class file version
            vec![ 0xFF, 0xFF, 0x00, 0x27 ]
        ];

        let invalid_versions= vec![
            vec![ ],
            vec![ 0x00 ],
            vec![ 0x00, 0x01 ],
            vec![ 0x00, 0x02, 0x03 ]
        ];

        for version in valid_versions {
            let ovn = ClassStream::new(&version).read_version_number();
            assert!(ovn.is_some());
            let vn = ovn.unwrap();
            assert!(vn.1 >= 0x27 && vn.1 <= 0x34);
        }

        for version in invalid_versions {
            let ovn = ClassStream::new(&version).read_version_number();
            assert!(ovn.is_none());
        }
    }

    #[test]
    fn read_constant_pool_doesnt_accept_empty_bytes() {
        let empty_cp = vec![];
        let r = ClassStream::new(&empty_cp).read_constant_pool();
        assert!(r.is_none());
    }

    #[test]
    fn read_constant_pool_doesnt_invalid_pools() {
        let invalid_pools = vec![
            vec![ 0x00 ], // too short
            vec![ 0x00, 0x02 ], // missing first item
        ];

        for pool in invalid_pools {
            let r = ClassStream::new(&pool).read_constant_pool();
            assert!(r.is_none());
        }
    }

    #[test]
    fn read_constant_pool_reads_empty_pools() {
        let empty_pool = vec![ 0x00, 0x01 ];
        let or = ClassStream::new(&empty_pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(0, r.len());
    }

    #[test]
    fn read_constant_pool_handles_long_entries_correctly() {
        //                           INT
        let pool = vec![ 0x00, 0x02, 0x03, 0x01, 0x02, 0x03, 0x04 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x01020304, b),
            _ => assert!(false)
        }

        //                           LONG
        let pool = vec![ 0x00, 0x02, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x01020304, h);
                assert_eq!(0x05060708, l);
            },
            _ => assert!(false)
        }

        //                           INT                           LONG
        let pool = vec![ 0x00, 0x03, 0x03, 0x10, 0x20, 0x30, 0x40, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(2, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x10203040, b),
            _ => assert!(false)
        }

        match r[1] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x01020304, h);
                assert_eq!(0x05060708, l);
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn read_constant_pool_should_ignore_entries_off_limit_by_long_entries() {
        // The second INT entry should be ignored in this test because the LONG entry takes up both indexes
        //                           LONG                                                   INT
        let pool = vec![ 0x00, 0x03, 0x05, 0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80, 0x03, 0x04, 0x05, 0x06, 0x07 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x10203040, h);
                assert_eq!(0x50607080, l);
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn it_should_parse_multiple_entries() {
        let pool = vec![ 0, 5, 0x03, 1, 1, 1, 1, 0x07, 2, 2, 0x08, 3, 3, 0x03, 4, 4, 4, 4 ];

        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(4, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x01010101, b),
            _ => assert!(false)
        }
        match r[1] {
            ConstantType::Class { name_index: n } => assert_eq!(0x0202, n),
            _ => assert!(false)
        }
        match r[2] {
            ConstantType::String { string_index: s } => assert_eq!(0x0303, s),
            _ => assert!(false)
        }
        match r[3] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x04040404, b),
            _ => assert!(false)
        }
    }
}
<commit_msg>minor fixes<commit_after>extern crate jvmti;

#[cfg(test)]
mod tests {

    use jvmti::bytecode::class_stream::ClassStream;
    use jvmti::bytecode::constants::ConstantType;

    #[test]
    fn read_magic_bytes_accepts_only_valid_magic_bytes() {
        let valid_magic = vec![ 0xCA, 0xFE, 0xBA, 0xBE ];
        let invalid_magics = vec![
            vec![],
            vec![ 0x00, ],
            vec![ 0xCA, 0xFE, 0xBA ],
            vec![ 0xCA, 0xFE, 0xBE, 0xBA ]
        ];

        assert_eq!(true, ClassStream::new(&valid_magic).read_magic_bytes());

        for magic in invalid_magics {
            assert_eq!(false, ClassStream::new(&magic).read_magic_bytes());
        }
    }

    #[test]
    fn read_version_number_accepts_only_valid_version_numbers() {
        let valid_versions = vec![
            vec![ 0x00, 0x10, 0x00, 0x30 ],
            vec![ 0x00, 0x00, 0x00, 0x34 ], // Java 8 class file version
            vec![ 0xFF, 0xFF, 0x00, 0x27 ]
        ];

        let invalid_versions= vec![
            vec![ ],
            vec![ 0x00 ],
            vec![ 0x00, 0x01 ],
            vec![ 0x00, 0x02, 0x03 ]
        ];

        for version in valid_versions {
            let ovn = ClassStream::new(&version).read_version_number();
            assert!(ovn.is_some());
            let vn = ovn.unwrap();
            assert!(vn.1 >= 0x27 && vn.1 <= 0x34);
        }

        for version in invalid_versions {
            let ovn = ClassStream::new(&version).read_version_number();
            assert!(ovn.is_none());
        }
    }

    #[test]
    fn read_constant_pool_doesnt_accept_empty_bytes() {
        let empty_cp = vec![];
        let r = ClassStream::new(&empty_cp).read_constant_pool();
        assert!(r.is_none());
    }

    #[test]
    fn read_constant_pool_doesnt_invalid_pools() {
        let invalid_pools = vec![
            vec![ 0x00 ], // too short
            vec![ 0x00, 0x02 ], // missing first item
        ];

        for pool in invalid_pools {
            let r = ClassStream::new(&pool).read_constant_pool();
            assert!(r.is_none());
        }
    }

    #[test]
    fn read_constant_pool_reads_empty_pools() {
        let empty_pool = vec![ 0x00, 0x01 ];
        let or = ClassStream::new(&empty_pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(0, r.len());
    }

    #[test]
    fn read_constant_pool_handles_long_entries_correctly() {
        //                           INT
        let pool = vec![ 0x00, 0x02, 0x03, 0x01, 0x02, 0x03, 0x04 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x01020304, b),
            _ => assert!(false)
        }

        //                           LONG
        let pool = vec![ 0x00, 0x02, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x01020304, h);
                assert_eq!(0x05060708, l);
            },
            _ => assert!(false)
        }

        //                           INT                           LONG
        let pool = vec![ 0x00, 0x03, 0x03, 0x10, 0x20, 0x30, 0x40, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(2, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x10203040, b),
            _ => assert!(false)
        }

        match r[1] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x01020304, h);
                assert_eq!(0x05060708, l);
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn read_constant_pool_should_ignore_entries_off_limit_by_long_entries() {
        // The second INT entry should be ignored in this test because the LONG entry takes up both indexes
        //                           LONG                                                   INT
        let pool = vec![ 0x00, 0x03, 0x05, 0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80, 0x03, 0x04, 0x05, 0x06, 0x07 ];
        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(1, r.len());

        match r[0] {
            ConstantType::Long { high_bytes: h, low_bytes: l } => {
                assert_eq!(0x10203040, h);
                assert_eq!(0x50607080, l);
            },
            _ => assert!(false)
        }
    }

    #[test]
    fn it_should_parse_multiple_entries() {
        let pool = vec![ 0, 5, 0x03, 1, 1, 1, 1, 0x07, 2, 2, 0x08, 3, 3, 0x03, 4, 4, 4, 4 ];

        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(4, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x01010101, b),
            _ => assert!(false)
        }
        match r[1] {
            ConstantType::Class { name_index: n } => assert_eq!(0x0202, n),
            _ => assert!(false)
        }
        match r[2] {
            ConstantType::String { string_index: s } => assert_eq!(0x0303, s),
            _ => assert!(false)
        }
        match r[3] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x04040404, b),
            _ => assert!(false)
        }
    }

    #[test]
    fn it_should_parse_multiple_entries_with_long_items() {
        let pool = vec![ 0, 7, 0x03, 1, 1, 1, 1, 0x05, 0, 0, 0, 0, 0, 0, 0, 0, 0x07, 2, 2, 0x08, 3, 3, 0x03, 4, 4, 4, 4 ];

        let or = ClassStream::new(&pool).read_constant_pool();
        assert!(or.is_some());
        let r = or.unwrap();

        assert_eq!(5, r.len());

        match r[0] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x01010101, b),
            _ => assert!(false)
        }
        match r[1] {
            ConstantType::Long { high_bytes: h, low_bytes: _ } => assert_eq!(0, h),
            _ => assert!(false)
        }
        match r[2] {
            ConstantType::Class { name_index: n } => assert_eq!(0x0202, n),
            _ => assert!(false)
        }
        match r[3] {
            ConstantType::String { string_index: s } => assert_eq!(0x0303, s),
            _ => assert!(false)
        }
        match r[4] {
            ConstantType::Integer { bytes: b } => assert_eq!(0x04040404, b),
            _ => assert!(false)
        }
    }
}
<|endoftext|>


<commit_before>use crate::runtime::task::RawTask;

use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};

cfg_rt! {
    /// An owned permission to join on a task (await its termination).
    ///
    /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for
    /// a task rather than a thread.
    ///
    /// A `JoinHandle` *detaches* the associated task when it is dropped, which
    /// means that there is no longer any handle to the task, and no way to `join`
    /// on it.
    ///
    /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`]
    /// functions.
    ///
    /// # Examples
    ///
    /// Creation from [`task::spawn`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn(async {
    ///     // some work here
    /// });
    /// # }
    /// ```
    ///
    /// Creation from [`task::spawn_blocking`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| {
    ///     // some blocking work here
    /// });
    /// # }
    /// ```
    ///
    /// The generic parameter `T` in `JoinHandle<T>` is the return type of the spawned task.
    /// If the return value is an i32, the join handle has type `JoinHandle<i32>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<i32> = task::spawn(async {
    ///     5 + 3
    /// });
    /// # }
    ///
    /// ```
    ///
    /// If the task does not have a return value, the join handle has type `JoinHandle<()>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<()> = task::spawn(async {
    ///     println!("I return nothing.");
    /// });
    /// # }
    /// ```
    ///
    /// Note that `handle.await` doesn't give you the return type directly. It is wrapped in a
    /// `Result` because panics in the spawned task are caught by Tokio. The `?` operator has
    /// to be double chained to extract the returned value:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         Ok(5 + 3)
    ///     });
    ///
    ///     let result = join_handle.await??;
    ///     assert_eq!(result, 8);
    ///     Ok(())
    /// }
    /// ```
    ///
    /// If the task panics, the error is a [`JoinError`] that contains the panic:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    /// use std::panic;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         panic!("boom");
    ///     });
    ///
    ///     let err = join_handle.await.unwrap_err();
    ///     assert!(err.is_panic());
    ///     Ok(())
    /// }
    ///
    /// ```
    /// Child being detached and outliving its parent:
    ///
    /// ```no_run
    /// use tokio::task;
    /// use tokio::time;
    /// use std::time::Duration;
    ///
    /// # #[tokio::main] async fn main() {
    /// let original_task = task::spawn(async {
    ///     let _detached_task = task::spawn(async {
    ///         // Here we sleep to make sure that the first task returns before.
    ///         time::sleep(Duration::from_millis(10)).await;
    ///         // This will be called, even though the JoinHandle is dropped.
    ///         println!("♫ Still alive ♫");
    ///     });
    /// });
    ///
    /// original_task.await.expect("The task being joined has panicked");
    /// println!("Original task is joined.");
    ///
    /// // We make sure that the new task has time to run, before the main
    /// // task returns.
    ///
    /// time::sleep(Duration::from_millis(1000)).await;
    /// # }
    /// ```
    ///
    /// [`task::spawn`]: crate::task::spawn()
    /// [`task::spawn_blocking`]: crate::task::spawn_blocking
    /// [`std::thread::JoinHandle`]: std::thread::JoinHandle
    /// [`JoinError`]: crate::task::JoinError
    pub struct JoinHandle<T> {
        raw: Option<RawTask>,
        _p: PhantomData<T>,
    }
}

unsafe impl<T: Send> Send for JoinHandle<T> {}
unsafe impl<T: Send> Sync for JoinHandle<T> {}

impl<T> JoinHandle<T> {
    pub(super) fn new(raw: RawTask) -> JoinHandle<T> {
        JoinHandle {
            raw: Some(raw),
            _p: PhantomData,
        }
    }

    /// Abort the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will complete with a `Err(JoinError::Cancelled)`.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       true
    ///    }));
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       false
    ///    }));
    ///
    ///    for handle in &handles {
    ///        handle.abort();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    pub fn abort(&self) {
        if let Some(raw) = self.raw {
            raw.remote_abort();
        }
    }
}

impl<T> Unpin for JoinHandle<T> {}

impl<T> Future for JoinHandle<T> {
    type Output = super::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        let coop = ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task output. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // The function must go via the vtable, which requires erasing generic
        // types. To do this, the function "return" is placed on the stack
        // **before** calling the function and is passed into the function using
        // `*mut ()`.
        //
        // Safety:
        //
        // The type of `T` must match the task's output type.
        unsafe {
            raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
        }

        if ret.is_ready() {
            coop.made_progress();
        }

        ret
    }
}

impl<T> Drop for JoinHandle<T> {
    fn drop(&mut self) {
        if let Some(raw) = self.raw.take() {
            if raw.header().state.drop_join_handle_fast().is_ok() {
                return;
            }

            raw.drop_join_handle_slow();
        }
    }
}

impl<T> fmt::Debug for JoinHandle<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("JoinHandle").finish()
    }
}
<commit_msg>task: improve JoinHandle::abort cancellation doc (#4121)<commit_after>use crate::runtime::task::RawTask;

use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};

cfg_rt! {
    /// An owned permission to join on a task (await its termination).
    ///
    /// This can be thought of as the equivalent of [`std::thread::JoinHandle`] for
    /// a task rather than a thread.
    ///
    /// A `JoinHandle` *detaches* the associated task when it is dropped, which
    /// means that there is no longer any handle to the task, and no way to `join`
    /// on it.
    ///
    /// This `struct` is created by the [`task::spawn`] and [`task::spawn_blocking`]
    /// functions.
    ///
    /// # Examples
    ///
    /// Creation from [`task::spawn`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn(async {
    ///     // some work here
    /// });
    /// # }
    /// ```
    ///
    /// Creation from [`task::spawn_blocking`]:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<_> = task::spawn_blocking(|| {
    ///     // some blocking work here
    /// });
    /// # }
    /// ```
    ///
    /// The generic parameter `T` in `JoinHandle<T>` is the return type of the spawned task.
    /// If the return value is an i32, the join handle has type `JoinHandle<i32>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<i32> = task::spawn(async {
    ///     5 + 3
    /// });
    /// # }
    ///
    /// ```
    ///
    /// If the task does not have a return value, the join handle has type `JoinHandle<()>`:
    ///
    /// ```
    /// use tokio::task;
    ///
    /// # async fn doc() {
    /// let join_handle: task::JoinHandle<()> = task::spawn(async {
    ///     println!("I return nothing.");
    /// });
    /// # }
    /// ```
    ///
    /// Note that `handle.await` doesn't give you the return type directly. It is wrapped in a
    /// `Result` because panics in the spawned task are caught by Tokio. The `?` operator has
    /// to be double chained to extract the returned value:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         Ok(5 + 3)
    ///     });
    ///
    ///     let result = join_handle.await??;
    ///     assert_eq!(result, 8);
    ///     Ok(())
    /// }
    /// ```
    ///
    /// If the task panics, the error is a [`JoinError`] that contains the panic:
    ///
    /// ```
    /// use tokio::task;
    /// use std::io;
    /// use std::panic;
    ///
    /// #[tokio::main]
    /// async fn main() -> io::Result<()> {
    ///     let join_handle: task::JoinHandle<Result<i32, io::Error>> = tokio::spawn(async {
    ///         panic!("boom");
    ///     });
    ///
    ///     let err = join_handle.await.unwrap_err();
    ///     assert!(err.is_panic());
    ///     Ok(())
    /// }
    ///
    /// ```
    /// Child being detached and outliving its parent:
    ///
    /// ```no_run
    /// use tokio::task;
    /// use tokio::time;
    /// use std::time::Duration;
    ///
    /// # #[tokio::main] async fn main() {
    /// let original_task = task::spawn(async {
    ///     let _detached_task = task::spawn(async {
    ///         // Here we sleep to make sure that the first task returns before.
    ///         time::sleep(Duration::from_millis(10)).await;
    ///         // This will be called, even though the JoinHandle is dropped.
    ///         println!("♫ Still alive ♫");
    ///     });
    /// });
    ///
    /// original_task.await.expect("The task being joined has panicked");
    /// println!("Original task is joined.");
    ///
    /// // We make sure that the new task has time to run, before the main
    /// // task returns.
    ///
    /// time::sleep(Duration::from_millis(1000)).await;
    /// # }
    /// ```
    ///
    /// [`task::spawn`]: crate::task::spawn()
    /// [`task::spawn_blocking`]: crate::task::spawn_blocking
    /// [`std::thread::JoinHandle`]: std::thread::JoinHandle
    /// [`JoinError`]: crate::task::JoinError
    pub struct JoinHandle<T> {
        raw: Option<RawTask>,
        _p: PhantomData<T>,
    }
}

unsafe impl<T: Send> Send for JoinHandle<T> {}
unsafe impl<T: Send> Sync for JoinHandle<T> {}

impl<T> JoinHandle<T> {
    pub(super) fn new(raw: RawTask) -> JoinHandle<T> {
        JoinHandle {
            raw: Some(raw),
            _p: PhantomData,
        }
    }

    /// Abort the task associated with the handle.
    ///
    /// Awaiting a cancelled task might complete as usual if the task was
    /// already completed at the time it was cancelled, but most likely it
    /// will fail with a [cancelled] `JoinError`.
    ///
    /// ```rust
    /// use tokio::time;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///    let mut handles = Vec::new();
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       true
    ///    }));
    ///
    ///    handles.push(tokio::spawn(async {
    ///       time::sleep(time::Duration::from_secs(10)).await;
    ///       false
    ///    }));
    ///
    ///    for handle in &handles {
    ///        handle.abort();
    ///    }
    ///
    ///    for handle in handles {
    ///        assert!(handle.await.unwrap_err().is_cancelled());
    ///    }
    /// }
    /// ```
    /// [cancelled]: method@super::error::JoinError::is_cancelled
    pub fn abort(&self) {
        if let Some(raw) = self.raw {
            raw.remote_abort();
        }
    }
}

impl<T> Unpin for JoinHandle<T> {}

impl<T> Future for JoinHandle<T> {
    type Output = super::Result<T>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut ret = Poll::Pending;

        // Keep track of task budget
        let coop = ready!(crate::coop::poll_proceed(cx));

        // Raw should always be set. If it is not, this is due to polling after
        // completion
        let raw = self
            .raw
            .as_ref()
            .expect("polling after `JoinHandle` already completed");

        // Try to read the task output. If the task is not yet complete, the
        // waker is stored and is notified once the task does complete.
        //
        // The function must go via the vtable, which requires erasing generic
        // types. To do this, the function "return" is placed on the stack
        // **before** calling the function and is passed into the function using
        // `*mut ()`.
        //
        // Safety:
        //
        // The type of `T` must match the task's output type.
        unsafe {
            raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker());
        }

        if ret.is_ready() {
            coop.made_progress();
        }

        ret
    }
}

impl<T> Drop for JoinHandle<T> {
    fn drop(&mut self) {
        if let Some(raw) = self.raw.take() {
            if raw.header().state.drop_join_handle_fast().is_ok() {
                return;
            }

            raw.drop_join_handle_slow();
        }
    }
}

impl<T> fmt::Debug for JoinHandle<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.debug_struct("JoinHandle").finish()
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate nom;
extern crate proc_macro;
extern crate syn;

use proc_macro::TokenStream;

mod generator;
mod parser;
mod path;

// Holds metadata for the template, based on the `template()` attribute.
struct TemplateMeta {
    path: String,
    print: String,
}

// Returns a `TemplateMeta` based on the `template()` attribute data found
// in the parsed struct or enum. Will panic if it does not find the required
// template path, or if the `print` key has an unexpected value.
fn get_template_meta(ast: &syn::DeriveInput) -> TemplateMeta {
    let mut path = None;
    let mut print = "none".to_string();
    let attr = ast.attrs.iter().find(|a| a.name() == "template").unwrap();
    if let syn::MetaItem::List(_, ref inner) = attr.value {
        for nm_item in inner {
            if let syn::NestedMetaItem::MetaItem(ref item) = *nm_item {
                if let syn::MetaItem::NameValue(ref key, ref val) = *item {
                    match key.as_ref() {
                        "path" => if let syn::Lit::Str(ref s, _) = *val {
                            path = Some(s.clone());
                        } else {
                            panic!("template path must be string literal");
                        },
                        "print" => if let syn::Lit::Str(ref s, _) = *val {
                            print = s.clone();
                        } else {
                            panic!("print value must be string literal");
                        },
                        _ => { panic!("unsupported annotation key found") }
                    }
                }
            }
        }
    }
    if path.is_none() {
        panic!("template path not found in struct attributes");
    }
    TemplateMeta { path: path.unwrap(), print: print }
}

/// Takes a `syn::DeriveInput` and generates source code for it
///
/// Reads the metadata from the `template()` attribute to get the template
/// metadata, then fetches the source from the filesystem. The source is
/// parsed, and the parse tree is fed to the code generator. Will print
/// the parse tree and/or generated source according to the `print` key's
/// value as passed to the `template()` attribute.
fn build_template(ast: &syn::DeriveInput) -> String {
    let meta = get_template_meta(ast);
    let mut src = path::get_template_source(&meta.path);
    if src.ends_with('\n') {
        let _ = src.pop();
    }
    let nodes = parser::parse(&src);
    if meta.print == "ast" || meta.print == "all" {
        println!("{:?}", nodes);
    }
    let code = generator::generate(ast, &meta.path, nodes);
    if meta.print == "code" || meta.print == "all" {
        println!("{}", code);
    }
    code
}

#[proc_macro_derive(Template, attributes(template))]
pub fn derive_template(input: TokenStream) -> TokenStream {
    let ast = syn::parse_derive_input(&input.to_string()).unwrap();
    match ast.body {
        syn::Body::Struct(ref data) => data,
        _ => panic!("#[derive(Template)] can only be used with structs"),
    };
    build_template(&ast).parse().unwrap()
}
<commit_msg>Panic with better message if template attribute is not found (fixes #3)<commit_after>#[macro_use]
extern crate nom;
extern crate proc_macro;
extern crate syn;

use proc_macro::TokenStream;

mod generator;
mod parser;
mod path;

// Holds metadata for the template, based on the `template()` attribute.
struct TemplateMeta {
    path: String,
    print: String,
}

// Returns a `TemplateMeta` based on the `template()` attribute data found
// in the parsed struct or enum. Will panic if it does not find the required
// template path, or if the `print` key has an unexpected value.
fn get_template_meta(ast: &syn::DeriveInput) -> TemplateMeta {
    let attr = ast.attrs.iter().find(|a| a.name() == "template");
    if attr.is_none() {
        let msg = format!("'template' attribute not found on struct '{}'",
                          ast.ident.as_ref());
        panic!(msg);
    }

    let attr = attr.unwrap();
    let mut path = None;
    let mut print = "none".to_string();
    if let syn::MetaItem::List(_, ref inner) = attr.value {
        for nm_item in inner {
            if let syn::NestedMetaItem::MetaItem(ref item) = *nm_item {
                if let syn::MetaItem::NameValue(ref key, ref val) = *item {
                    match key.as_ref() {
                        "path" => if let syn::Lit::Str(ref s, _) = *val {
                            path = Some(s.clone());
                        } else {
                            panic!("template path must be string literal");
                        },
                        "print" => if let syn::Lit::Str(ref s, _) = *val {
                            print = s.clone();
                        } else {
                            panic!("print value must be string literal");
                        },
                        _ => { panic!("unsupported annotation key found") }
                    }
                }
            }
        }
    }
    if path.is_none() {
        panic!("template path not found in struct attributes");
    }
    TemplateMeta { path: path.unwrap(), print: print }
}

/// Takes a `syn::DeriveInput` and generates source code for it
///
/// Reads the metadata from the `template()` attribute to get the template
/// metadata, then fetches the source from the filesystem. The source is
/// parsed, and the parse tree is fed to the code generator. Will print
/// the parse tree and/or generated source according to the `print` key's
/// value as passed to the `template()` attribute.
fn build_template(ast: &syn::DeriveInput) -> String {
    let meta = get_template_meta(ast);
    let mut src = path::get_template_source(&meta.path);
    if src.ends_with('\n') {
        let _ = src.pop();
    }
    let nodes = parser::parse(&src);
    if meta.print == "ast" || meta.print == "all" {
        println!("{:?}", nodes);
    }
    let code = generator::generate(ast, &meta.path, nodes);
    if meta.print == "code" || meta.print == "all" {
        println!("{}", code);
    }
    code
}

#[proc_macro_derive(Template, attributes(template))]
pub fn derive_template(input: TokenStream) -> TokenStream {
    let ast = syn::parse_derive_input(&input.to_string()).unwrap();
    match ast.body {
        syn::Body::Struct(ref data) => data,
        _ => panic!("#[derive(Template)] can only be used with structs"),
    };
    build_template(&ast).parse().unwrap()
}
<|endoftext|>


<commit_before>use memory::MemSegment;
use super::Color;
use super::PPU;
use super::PaletteIndex;
use super::PaletteSet;

const TRANSPARENT: PaletteIndex = PaletteIndex{
    set: PaletteSet::Sprite,
    palette_id: 0,
    color_id: 0,
};

bitflags! {
    flags OAMAttr : u8 {
        const FLIP_VERT = 0b1000_0000,
        const FLIP_HORZ = 0b0100_0000,
        const BEHIND    = 0b0010_0000,
        const PALETTE1  = 0b0000_0010,
        const PALETTE2  = 0b0000_0001,
    }
}

impl OAMAttr {
    fn palette(&self) -> u8 {
        self.bits & 0x03
    }

    fn priority(&self) -> SpritePriority {
        if self.contains(BEHIND) {
            SpritePriority::Background
        }
        else {
            SpritePriority::Foreground
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct OAMEntry {
    y: u8,
    tile: u8,
    attr: OAMAttr,
    x: u8,
}

impl Default for OAMEntry {
    fn default() -> OAMEntry {
        OAMEntry {
            y: 0xFF,
            tile: 0,
            attr: OAMAttr::from_bits_truncate(0),
            x: 0xFF,
        }
    }
}

impl MemSegment for OAMEntry {
    fn read(&mut self, idx: u16) -> u8 {
        match idx {
            0 => self.y,
            1 => self.tile,
            2 => self.attr.bits(),
            3 => self.x,
            x => invalid_address!(x),
        }
    }

    fn write(&mut self, idx: u16, val: u8) {
        match idx {
            0 => self.y = val,
            1 => self.tile = val,
            2 => self.attr = OAMAttr::from_bits_truncate(val),
            3 => self.x = val,
            x => invalid_address!(x),
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct SpriteDetails {
    x: u8,
    attr: OAMAttr,
    tile: (u8, u8),
}

impl Default for SpriteDetails {
    fn default() -> SpriteDetails {
        SpriteDetails{
            x: 0xFF,
            attr: OAMAttr::empty(),
            tile: (0, 0),
        }
    }
}

pub struct SpriteRenderer {
    primary_oam: [OAMEntry; 64],
    secondary_oam: [SpriteDetails; 8],
}

impl Default for SpriteRenderer {
    fn default() -> SpriteRenderer {
        SpriteRenderer {
            primary_oam: [Default::default(); 64],
            secondary_oam: [Default::default(); 8],
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum SpritePriority {
    Foreground, Background,
}

///Reads the primary OAM table.
impl MemSegment for SpriteRenderer {
    fn read(&mut self, idx: u16) -> u8 {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].read( idx % 4 )
    }

    fn write(&mut self, idx: u16, val: u8) {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].write( idx % 4, val )
    }
}

impl PPU {
    pub fn visible_scanline_sprite(&mut self, pixel: u16, scanline: i16) {
        if pixel == 0 {
            self.sprite_eval(scanline);
        }
    }

    fn sprite_eval(&mut self, scanline: i16) {
        let mut n = 0;
        self.sprite_data.secondary_oam = [Default::default(); 8];
        for x in 0..64 {
            let oam = self.sprite_data.primary_oam[x];
            if self.is_on_scanline( oam, scanline ) {
                self.sprite_data.secondary_oam[n] = self.convert_oam_entry(oam, scanline);
                n += 1;
                if n == 8 {
                    return;
                }
            }
        }
    }

    fn is_on_scanline(&self, oam: OAMEntry, scanline: i16) -> bool {
        let y = oam.y as i16;
        y <= scanline && scanline < y + 8
    }

    fn convert_oam_entry(&mut self, oam: OAMEntry, sl: i16 ) -> SpriteDetails {
        let tile_id = oam.tile;
        let fine_y_scroll = sl as u16 - oam.y as u16;
        let tile_table = self.reg.ppuctrl.sprite_table();
        let tile = self.read_tile_pattern(tile_id, fine_y_scroll, tile_table);
        SpriteDetails {
            x: oam.x,
            attr: oam.attr,
            tile: tile,
        }
    }

    pub fn get_sprite_pixel(&mut self, x: u16, y: u16) -> (SpritePriority, PaletteIndex) {
        for n in 0..8 {
            let det_x = self.sprite_data.secondary_oam[n];
            if self.is_active( det_x, x ) {
                return self.do_get_pixel( det_x, x );
            }
        }
        return (SpritePriority::Background, TRANSPARENT);
    }

    fn is_active(&self, details: SpriteDetails, x: u16) -> bool {
        x.wrapping_sub(details.x as u16) <= 8
    }

    fn do_get_pixel(&mut self, details: SpriteDetails, x: u16) -> (SpritePriority, PaletteIndex) {
        let fine_x = x - details.x as u16;
        let attr = details.attr;
        let color_id = self.get_color_in_pattern(details.tile, fine_x as u32);
        let idx = PaletteIndex {
            set: PaletteSet::Sprite,
            palette_id: attr.palette(),
            color_id: color_id,
        };
        return (attr.priority(), idx);
    }
}

#[cfg(tests)]
mod tests {
    use memory::MemSegment;
    use ppu::PPU;
    use super::*;

    #[test]
    fn reading_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        for x in 0u8..63u8 {
            ppu.oam[x as usize] = OAMEntry::new(x, x, x, x);
        }
        ppu.reg.oamaddr = 0;
        assert_eq!(ppu.read(0x2004), 0);
        ppu.reg.oamaddr = 10;
        assert_eq!(ppu.read(0x2004), 2);
    }

    #[test]
    fn writing_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        ppu.reg.oamaddr = 0;
        ppu.write(0x2004, 12);
        assert_eq!(ppu.oam[0].y, 12);
        ppu.reg.oamaddr = 10;
        ppu.write(0x2004, 3);
        assert_eq!(ppu.oam[2].attr.bits(), 3);
    }

    #[test]
    fn test_sprite_on_scanline() {
    	let mut ppu = create_test_ppu();
    	let mut oam : OAMEntry = Default::default();
    	oam.y = 10;

    	assert!(!ppu.is_on_scanline(oam, 9));
    	for sl in 10..18 {
    	    assert!(ppu.is_on_scanline(oam, sl));
    	}
    	assert!(!ppu.is_on_scanline(oam, 18));
    }
}<commit_msg>Fixed bug where sprites rendered with a column of garbage on the right.<commit_after>use memory::MemSegment;
use super::Color;
use super::PPU;
use super::PaletteIndex;
use super::PaletteSet;

const TRANSPARENT: PaletteIndex = PaletteIndex{
    set: PaletteSet::Sprite,
    palette_id: 0,
    color_id: 0,
};

bitflags! {
    flags OAMAttr : u8 {
        const FLIP_VERT = 0b1000_0000,
        const FLIP_HORZ = 0b0100_0000,
        const BEHIND    = 0b0010_0000,
        const PALETTE1  = 0b0000_0010,
        const PALETTE2  = 0b0000_0001,
    }
}

impl OAMAttr {
    fn palette(&self) -> u8 {
        self.bits & 0x03
    }

    fn priority(&self) -> SpritePriority {
        if self.contains(BEHIND) {
            SpritePriority::Background
        }
        else {
            SpritePriority::Foreground
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct OAMEntry {
    y: u8,
    tile: u8,
    attr: OAMAttr,
    x: u8,
}

impl Default for OAMEntry {
    fn default() -> OAMEntry {
        OAMEntry {
            y: 0xFF,
            tile: 0,
            attr: OAMAttr::from_bits_truncate(0),
            x: 0xFF,
        }
    }
}

impl MemSegment for OAMEntry {
    fn read(&mut self, idx: u16) -> u8 {
        match idx {
            0 => self.y,
            1 => self.tile,
            2 => self.attr.bits(),
            3 => self.x,
            x => invalid_address!(x),
        }
    }

    fn write(&mut self, idx: u16, val: u8) {
        match idx {
            0 => self.y = val,
            1 => self.tile = val,
            2 => self.attr = OAMAttr::from_bits_truncate(val),
            3 => self.x = val,
            x => invalid_address!(x),
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct SpriteDetails {
    x: u8,
    attr: OAMAttr,
    tile: (u8, u8),
}

impl Default for SpriteDetails {
    fn default() -> SpriteDetails {
        SpriteDetails{
            x: 0xFF,
            attr: OAMAttr::empty(),
            tile: (0, 0),
        }
    }
}

pub struct SpriteRenderer {
    primary_oam: [OAMEntry; 64],
    secondary_oam: [SpriteDetails; 8],
}

impl Default for SpriteRenderer {
    fn default() -> SpriteRenderer {
        SpriteRenderer {
            primary_oam: [Default::default(); 64],
            secondary_oam: [Default::default(); 8],
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum SpritePriority {
    Foreground, Background,
}

///Reads the primary OAM table.
impl MemSegment for SpriteRenderer {
    fn read(&mut self, idx: u16) -> u8 {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].read( idx % 4 )
    }

    fn write(&mut self, idx: u16, val: u8) {
        if idx > 256 {
            invalid_address!(idx);
        }
        self.primary_oam[idx as usize / 4].write( idx % 4, val )
    }
}

impl PPU {
    pub fn visible_scanline_sprite(&mut self, pixel: u16, scanline: i16) {
        if pixel == 0 {
            self.sprite_eval(scanline);
        }
    }

    fn sprite_eval(&mut self, scanline: i16) {
        let mut n = 0;
        self.sprite_data.secondary_oam = [Default::default(); 8];
        for x in 0..64 {
            let oam = self.sprite_data.primary_oam[x];
            if self.is_on_scanline( oam, scanline ) {
                self.sprite_data.secondary_oam[n] = self.convert_oam_entry(oam, scanline);
                n += 1;
                if n == 8 {
                    return;
                }
            }
        }
    }

    fn is_on_scanline(&self, oam: OAMEntry, scanline: i16) -> bool {
        let y = oam.y as i16;
        y <= scanline && scanline < y + 8
    }

    fn convert_oam_entry(&mut self, oam: OAMEntry, sl: i16 ) -> SpriteDetails {
        let tile_id = oam.tile;
        let fine_y_scroll = sl as u16 - oam.y as u16;
        let tile_table = self.reg.ppuctrl.sprite_table();
        let tile = self.read_tile_pattern(tile_id, fine_y_scroll, tile_table);
        SpriteDetails {
            x: oam.x,
            attr: oam.attr,
            tile: tile,
        }
    }

    pub fn get_sprite_pixel(&mut self, x: u16, y: u16) -> (SpritePriority, PaletteIndex) {
        for n in 0..8 {
            let det_x = self.sprite_data.secondary_oam[n];
            if self.is_active( det_x, x ) {
                return self.do_get_pixel( det_x, x );
            }
        }
        return (SpritePriority::Background, TRANSPARENT);
    }

    fn is_active(&self, details: SpriteDetails, x: u16) -> bool {
        x.wrapping_sub(details.x as u16) < 8
    }

    fn do_get_pixel(&mut self, details: SpriteDetails, x: u16) -> (SpritePriority, PaletteIndex) {
        let fine_x = x - details.x as u16;
        let attr = details.attr;
        let color_id = self.get_color_in_pattern(details.tile, fine_x as u32);
        let idx = PaletteIndex {
            set: PaletteSet::Sprite,
            palette_id: attr.palette(),
            color_id: color_id,
        };
        return (attr.priority(), idx);
    }
}

#[cfg(tests)]
mod tests {
    use memory::MemSegment;
    use ppu::PPU;
    use super::*;

    #[test]
    fn reading_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        for x in 0u8..63u8 {
            ppu.oam[x as usize] = OAMEntry::new(x, x, x, x);
        }
        ppu.reg.oamaddr = 0;
        assert_eq!(ppu.read(0x2004), 0);
        ppu.reg.oamaddr = 10;
        assert_eq!(ppu.read(0x2004), 2);
    }

    #[test]
    fn writing_oamdata_uses_oamaddr_as_index_into_oam() {
        let mut ppu = create_test_ppu();
        ppu.reg.oamaddr = 0;
        ppu.write(0x2004, 12);
        assert_eq!(ppu.oam[0].y, 12);
        ppu.reg.oamaddr = 10;
        ppu.write(0x2004, 3);
        assert_eq!(ppu.oam[2].attr.bits(), 3);
    }

    #[test]
    fn test_sprite_on_scanline() {
    	let mut ppu = create_test_ppu();
    	let mut oam : OAMEntry = Default::default();
    	oam.y = 10;

    	assert!(!ppu.is_on_scanline(oam, 9));
    	for sl in 10..18 {
    	    assert!(ppu.is_on_scanline(oam, sl));
    	}
    	assert!(!ppu.is_on_scanline(oam, 18));
    }
}<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::mem;
use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};

use libc::{c_void, dup, eventfd, read, write};

use {errno_result, Result};

/// A safe wrapper around a Linux eventfd (man 2 eventfd).
///
/// An eventfd is useful because it is sendable across processes and can be used for signaling in
/// and out of the KVM API. They can also be polled like any other file descriptor.
pub struct EventFd {
    eventfd: File,
}

impl EventFd {
    /// Creates a new blocking EventFd with an initial value of 0.
    pub fn new() -> Result<EventFd> {
        // This is safe because eventfd merely allocated an eventfd for our process and we handle
        // the error case.
        let ret = unsafe { eventfd(0, 0) };
        if ret < 0 {
            return errno_result();
        }
        // This is safe because we checked ret for success and know the kernel gave us an fd that we
        // own.
        Ok(EventFd {
            eventfd: unsafe { File::from_raw_fd(ret) },
        })
    }

    /// Adds `v` to the eventfd's count, blocking until this won't overflow the count.
    pub fn write(&self, v: u64) -> Result<()> {
        // This is safe because we made this fd and the pointer we pass can not overflow because we
        // give the syscall's size parameter properly.
        let ret = unsafe {
            write(
                self.as_raw_fd(),
                &v as *const u64 as *const c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            return errno_result();
        }
        Ok(())
    }

    /// Blocks until the the eventfd's count is non-zero, then resets the count to zero.
    pub fn read(&self) -> Result<u64> {
        let mut buf: u64 = 0;
        let ret = unsafe {
            // This is safe because we made this fd and the pointer we pass can not overflow because
            // we give the syscall's size parameter properly.
            read(
                self.as_raw_fd(),
                &mut buf as *mut u64 as *mut c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            return errno_result();
        }
        Ok(buf)
    }

    /// Clones this EventFd, internally creating a new file descriptor. The new EventFd will share
    /// the same underlying count within the kernel.
    pub fn try_clone(&self) -> Result<EventFd> {
        // This is safe because we made this fd and properly check that it returns without error.
        let ret = unsafe { dup(self.as_raw_fd()) };
        if ret < 0 {
            return errno_result();
        }
        // This is safe because we checked ret for success and know the kernel gave us an fd that we
        // own.
        Ok(EventFd {
            eventfd: unsafe { File::from_raw_fd(ret) },
        })
    }
}

impl AsRawFd for EventFd {
    fn as_raw_fd(&self) -> RawFd {
        self.eventfd.as_raw_fd()
    }
}

unsafe impl ::Pollable for EventFd {
    fn pollable_fd(&self) -> RawFd {
        self.eventfd.as_raw_fd()
    }
}

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

    #[test]
    fn new() {
        EventFd::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = EventFd::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read(), Ok(55));
    }

    #[test]
    fn clone() {
        let evt = EventFd::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read(), Ok(923));
    }
}
<commit_msg>sys_util: eventfd: remove the Pollable trait impl as it is not used<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::mem;
use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};

use libc::{c_void, dup, eventfd, read, write};

use {errno_result, Result};

/// A safe wrapper around a Linux eventfd (man 2 eventfd).
///
/// An eventfd is useful because it is sendable across processes and can be used for signaling in
/// and out of the KVM API. They can also be polled like any other file descriptor.
pub struct EventFd {
    eventfd: File,
}

impl EventFd {
    /// Creates a new blocking EventFd with an initial value of 0.
    pub fn new() -> Result<EventFd> {
        // This is safe because eventfd merely allocated an eventfd for our process and we handle
        // the error case.
        let ret = unsafe { eventfd(0, 0) };
        if ret < 0 {
            return errno_result();
        }
        // This is safe because we checked ret for success and know the kernel gave us an fd that we
        // own.
        Ok(EventFd {
            eventfd: unsafe { File::from_raw_fd(ret) },
        })
    }

    /// Adds `v` to the eventfd's count, blocking until this won't overflow the count.
    pub fn write(&self, v: u64) -> Result<()> {
        // This is safe because we made this fd and the pointer we pass can not overflow because we
        // give the syscall's size parameter properly.
        let ret = unsafe {
            write(
                self.as_raw_fd(),
                &v as *const u64 as *const c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            return errno_result();
        }
        Ok(())
    }

    /// Blocks until the the eventfd's count is non-zero, then resets the count to zero.
    pub fn read(&self) -> Result<u64> {
        let mut buf: u64 = 0;
        let ret = unsafe {
            // This is safe because we made this fd and the pointer we pass can not overflow because
            // we give the syscall's size parameter properly.
            read(
                self.as_raw_fd(),
                &mut buf as *mut u64 as *mut c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            return errno_result();
        }
        Ok(buf)
    }

    /// Clones this EventFd, internally creating a new file descriptor. The new EventFd will share
    /// the same underlying count within the kernel.
    pub fn try_clone(&self) -> Result<EventFd> {
        // This is safe because we made this fd and properly check that it returns without error.
        let ret = unsafe { dup(self.as_raw_fd()) };
        if ret < 0 {
            return errno_result();
        }
        // This is safe because we checked ret for success and know the kernel gave us an fd that we
        // own.
        Ok(EventFd {
            eventfd: unsafe { File::from_raw_fd(ret) },
        })
    }
}

impl AsRawFd for EventFd {
    fn as_raw_fd(&self) -> RawFd {
        self.eventfd.as_raw_fd()
    }
}

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

    #[test]
    fn new() {
        EventFd::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = EventFd::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read(), Ok(55));
    }

    #[test]
    fn clone() {
        let evt = EventFd::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read(), Ok(923));
    }
}
<|endoftext|>


<commit_before>// This code is heavily inspired by hoedown.

//
// Performance observations:
//
// Note that buf.head() == Some(&ch) is much faster than buf[0] == ch.
//

static SP: u8 = ' ' as u8;
static NL: u8 = '\n' as u8;
static CR: u8 = '\r' as u8;
static TAB: u8 = '\t' as u8;
static STAR: u8 = '*' as u8;
static DASH: u8 = '-' as u8;
static UNDERSCORE: u8 = '_' as u8;
static TILDE: u8 = '~' as u8;
static BACKTICK: u8 = '`' as u8;

//
// Skip up to three leading spaces.
//
// Up to three leading spaces are allowed for many elements.
//
// We don't need to care about a TAB here as in this position it is equivalent
// to 4 spaces, which means that when we find a TAB here we would not parse the
// corresponding element.
//
fn skip_initial_three_spaces<'a>(buf: &'a[u8])-> &'a[u8] {
    let mut buf = buf;
    if buf.head() == Some(&SP) { buf = buf.tail();
    if buf.head() == Some(&SP) { buf = buf.tail();
    if buf.head() == Some(&SP) { buf = buf.tail(); } } }
    return buf;
}

//
// Return Some(`rem`) if the line is a horizontal rule, with `rem` being the
// buf after the hrule. Otherwise return None.
//
fn is_hrule<'a>(buf: &'a[u8]) -> Option<&'a[u8]> {
    let buf = skip_initial_three_spaces(buf);

    let item = match buf.head() {
        Some(&c) if c == STAR ||
                    c == DASH ||
                    c == UNDERSCORE => c,
        _                           => return None
    };

    // The count of '*', '-' or '_'
    let mut cnt: uint = 0;

    // Counts the consumed spaces (and the final NL)
    let mut spc: uint = 0;

    for &ch in buf.iter() {
        if      ch == item { cnt += 1; }
        else if ch == NL   { spc += 1; break; }
        else if ch == SP   { spc += 1; }
        else               { return None; }
    }

    if cnt >= 3 {
        Some(buf.slice_from(cnt + spc))
    } else {
        None
    }
}

#[test]
fn test_is_hrule() {
    // examples as given on the markdown homepage
    assert!(is_hrule(b"* * *\n").is_some());
    assert!(is_hrule(b"***\n").is_some());
    assert!(is_hrule(b"*****\n").is_some());
    assert!(is_hrule(b"- - -\n").is_some());
    assert!(is_hrule(b"---------------------------------------\n").is_some());

    // up to three spaces ignored
    assert!(is_hrule(b" ***\n").is_some());
    assert!(is_hrule(b"  ***\n").is_some());
    assert!(is_hrule(b"   ***\n").is_some());

    // but not four, or a tab which is equivalent to four spaces
    assert!(is_hrule(b"    ***\n").is_none());
    assert!(is_hrule(b"\t***\n").is_none());

    // need at least three
    assert!(is_hrule(b"*\n").is_none());
    assert!(is_hrule(b"**\n").is_none());
    assert!(is_hrule(b"* *\n").is_none());
    assert!(is_hrule(b"   * *\n").is_none());

    // Also works without newline at the end
    assert!(is_hrule(b"* * *").is_some());

    // And underscores also supported
    assert!(is_hrule(b"___").is_some());
    assert!(is_hrule(b"______________").is_some());
    assert!(is_hrule(b" ______________").is_some());

    // Test if the remaining buf actually works.
    let s = b"   * * *\nremaining";
    let res = is_hrule(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), b"remaining");
}

//
// Return Some(`rem`) if the line is an empty line, with `rem` being the buf
// after the empty line. Otherwise return None.
//
fn is_empty<'a>(buf: &'a[u8]) -> Option<&'a[u8]> {
    let mut cnt: uint = 0;

    for &ch in buf.iter() {
        if      ch == NL  { cnt += 1; break; }
        else if ch == CR  { cnt += 1; }
        else if ch == SP  { cnt += 1; }
        else if ch == TAB { cnt += 1; }
        else              { return None; }
    }

    if cnt > 0 {
        Some(buf.slice_from(cnt))
    } else {
        None
    }
}

#[test]
fn test_is_empty() {
    assert!(is_empty(b"\n").is_some());
    assert!(is_empty(b"    \n").is_some());
    assert!(is_empty(b"  \t  \n").is_some());
    assert!(is_empty(b"  \t  \r\n").is_some());
    assert!(is_empty(b"  \t  \nabc").is_some());
    assert!(is_empty(b"  \t  ").is_some());

    assert!(is_empty(b"a").is_none());
    assert!(is_empty(b" a").is_none());
    assert!(is_empty(b" a\n").is_none());
    assert!(is_empty(b" \ta\n").is_none());

    // Test if the remaining buf actually works.
    let s = b"   \t\r\nremaining";
    let res = is_empty(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), b"remaining");
}

fn is_codefence<'a>(buf: &'a[u8]) -> Option<(&'a[u8], uint, u8)> {
    let buf = skip_initial_three_spaces(buf);

    let item = match buf.head() {
        Some(&c) if c == TILDE ||
                    c == BACKTICK => c,
        _                         => return None
    };

    // The count of '~' or '`' characters
    let mut cnt: uint = 0;

    for &ch in buf.iter() {
        if ch == item { cnt += 1; }
        else          { break; }
    }

    if cnt >= 3 {
        Some((buf.slice_from(cnt), cnt, item))
    } else {
        None
    }
}

#[test]
fn test_is_codefence() {
    assert!(is_codefence(b"```").is_some());
    assert!(is_codefence(b"~~~").is_some());
    assert!(is_codefence(b"`````````").is_some());
    assert!(is_codefence(b"~~~~").is_some());
    assert!(is_codefence(b"   ```").is_some());
    assert!(is_codefence(b"  ~~~").is_some());

    assert!(is_codefence(b"  ~~").is_none());
    assert!(is_codefence(b" ``").is_none());
    assert!(is_codefence(b"    ```").is_none());
    assert!(is_codefence(b"\t```").is_none());

    // Test if the remaining buf actually works.
    let s = b"   ```remaining\n";
    let res = is_codefence(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), (b"remaining\n", 3, BACKTICK));

    let s = b"   ~~~~~~~~remaining\n";
    let res = is_codefence(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), (b"remaining\n", 8, TILDE));
}

<commit_msg>Use byte literals<commit_after>// This code is heavily inspired by hoedown.

//
// Performance observations:
//
// Note that buf.head() == Some(&ch) is much faster than buf[0] == ch.
//

static SP: u8 = b' ';
static NL: u8 = b'\n';
static CR: u8 = b'\r';
static TAB: u8 = b'\t';
static STAR: u8 = b'*';
static DASH: u8 = b'-';
static UNDERSCORE: u8 = b'_';
static TILDE: u8 = b'~';
static BACKTICK: u8 = b'`';

//
// Skip up to three leading spaces.
//
// Up to three leading spaces are allowed for many elements.
//
// We don't need to care about a TAB here as in this position it is equivalent
// to 4 spaces, which means that when we find a TAB here we would not parse the
// corresponding element.
//
fn skip_initial_three_spaces<'a>(buf: &'a[u8])-> &'a[u8] {
    let mut buf = buf;
    if buf.head() == Some(&SP) { buf = buf.tail();
    if buf.head() == Some(&SP) { buf = buf.tail();
    if buf.head() == Some(&SP) { buf = buf.tail(); } } }
    return buf;
}

//
// Return Some(`rem`) if the line is a horizontal rule, with `rem` being the
// buf after the hrule. Otherwise return None.
//
fn is_hrule<'a>(buf: &'a[u8]) -> Option<&'a[u8]> {
    let buf = skip_initial_three_spaces(buf);

    let item = match buf.head() {
        Some(&c) if c == STAR ||
                    c == DASH ||
                    c == UNDERSCORE => c,
        _                           => return None
    };

    // The count of '*', '-' or '_'
    let mut cnt: uint = 0;

    // Counts the consumed spaces (and the final NL)
    let mut spc: uint = 0;

    for &ch in buf.iter() {
        if      ch == item { cnt += 1; }
        else if ch == NL   { spc += 1; break; }
        else if ch == SP   { spc += 1; }
        else               { return None; }
    }

    if cnt >= 3 {
        Some(buf.slice_from(cnt + spc))
    } else {
        None
    }
}

#[test]
fn test_is_hrule() {
    // examples as given on the markdown homepage
    assert!(is_hrule(b"* * *\n").is_some());
    assert!(is_hrule(b"***\n").is_some());
    assert!(is_hrule(b"*****\n").is_some());
    assert!(is_hrule(b"- - -\n").is_some());
    assert!(is_hrule(b"---------------------------------------\n").is_some());

    // up to three spaces ignored
    assert!(is_hrule(b" ***\n").is_some());
    assert!(is_hrule(b"  ***\n").is_some());
    assert!(is_hrule(b"   ***\n").is_some());

    // but not four, or a tab which is equivalent to four spaces
    assert!(is_hrule(b"    ***\n").is_none());
    assert!(is_hrule(b"\t***\n").is_none());

    // need at least three
    assert!(is_hrule(b"*\n").is_none());
    assert!(is_hrule(b"**\n").is_none());
    assert!(is_hrule(b"* *\n").is_none());
    assert!(is_hrule(b"   * *\n").is_none());

    // Also works without newline at the end
    assert!(is_hrule(b"* * *").is_some());

    // And underscores also supported
    assert!(is_hrule(b"___").is_some());
    assert!(is_hrule(b"______________").is_some());
    assert!(is_hrule(b" ______________").is_some());

    // Test if the remaining buf actually works.
    let s = b"   * * *\nremaining";
    let res = is_hrule(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), b"remaining");
}

//
// Return Some(`rem`) if the line is an empty line, with `rem` being the buf
// after the empty line. Otherwise return None.
//
fn is_empty<'a>(buf: &'a[u8]) -> Option<&'a[u8]> {
    let mut cnt: uint = 0;

    for &ch in buf.iter() {
        if      ch == NL  { cnt += 1; break; }
        else if ch == CR  { cnt += 1; }
        else if ch == SP  { cnt += 1; }
        else if ch == TAB { cnt += 1; }
        else              { return None; }
    }

    if cnt > 0 {
        Some(buf.slice_from(cnt))
    } else {
        None
    }
}

#[test]
fn test_is_empty() {
    assert!(is_empty(b"\n").is_some());
    assert!(is_empty(b"    \n").is_some());
    assert!(is_empty(b"  \t  \n").is_some());
    assert!(is_empty(b"  \t  \r\n").is_some());
    assert!(is_empty(b"  \t  \nabc").is_some());
    assert!(is_empty(b"  \t  ").is_some());

    assert!(is_empty(b"a").is_none());
    assert!(is_empty(b" a").is_none());
    assert!(is_empty(b" a\n").is_none());
    assert!(is_empty(b" \ta\n").is_none());

    // Test if the remaining buf actually works.
    let s = b"   \t\r\nremaining";
    let res = is_empty(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), b"remaining");
}

fn is_codefence<'a>(buf: &'a[u8]) -> Option<(&'a[u8], uint, u8)> {
    let buf = skip_initial_three_spaces(buf);

    let item = match buf.head() {
        Some(&c) if c == TILDE ||
                    c == BACKTICK => c,
        _                         => return None
    };

    // The count of '~' or '`' characters
    let mut cnt: uint = 0;

    for &ch in buf.iter() {
        if ch == item { cnt += 1; }
        else          { break; }
    }

    if cnt >= 3 {
        Some((buf.slice_from(cnt), cnt, item))
    } else {
        None
    }
}

#[test]
fn test_is_codefence() {
    assert!(is_codefence(b"```").is_some());
    assert!(is_codefence(b"~~~").is_some());
    assert!(is_codefence(b"`````````").is_some());
    assert!(is_codefence(b"~~~~").is_some());
    assert!(is_codefence(b"   ```").is_some());
    assert!(is_codefence(b"  ~~~").is_some());

    assert!(is_codefence(b"  ~~").is_none());
    assert!(is_codefence(b" ``").is_none());
    assert!(is_codefence(b"    ```").is_none());
    assert!(is_codefence(b"\t```").is_none());

    // Test if the remaining buf actually works.
    let s = b"   ```remaining\n";
    let res = is_codefence(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), (b"remaining\n", 3, BACKTICK));

    let s = b"   ~~~~~~~~remaining\n";
    let res = is_codefence(s);
    assert!(res.is_some());
    assert_eq!(res.unwrap(), (b"remaining\n", 8, TILDE));
}

<|endoftext|>


<commit_before>// This code is heavily inspired by hoedown.

extern crate extra;
use extra::test::BenchHarness;

static SP: u8 = ' ' as u8;
static NL: u8 = '\n' as u8;

/// Returns true if the line is a horizontal rule.
fn is_hrule(buf: &[u8]) -> bool {
    let len = buf.len();

    if len < 3 { return false }

    let mut pos: uint = 0;

    // Skip up to three leading spaces.
    //
    // We don't need to care about tabs here as in this position they are
    // equivalent to 4 spaces.
    if buf[0] == SP {
        pos += 1;
        if buf[1] == SP {
            pos += 1;
            if buf[2] == SP {
                pos += 1;
            }
        }
    }

    // We need at least 3 items
    if pos + 2 >= len { return false }

    let item = buf[pos];

    if !(item == ('*' as u8) || item == ('-' as u8) || item == ('_' as u8)) {
        return false;
    }

    let mut cnt: uint = 0; // The count of '*', '-' or '_'
    for &ch in buf.slice_from(pos).iter() {
        if ch == item {
            cnt += 1;
        } else if ch == NL {
            break;
        } else if ch != SP {
            return false;
        }
    }

    return cnt >= 3;
}

#[test]
fn test_is_hrule() {
    // examples as given on the markdown homepage
    assert_eq!(is_hrule(bytes!("* * *\n")), true);
    assert_eq!(is_hrule(bytes!("***\n")), true);
    assert_eq!(is_hrule(bytes!("*****\n")), true);
    assert_eq!(is_hrule(bytes!("- - -\n")), true);
    assert_eq!(is_hrule(bytes!("---------------------------------------\n")), true);

    // up to three spaces ignored
    assert_eq!(is_hrule(bytes!(" ***\n")), true);
    assert_eq!(is_hrule(bytes!("  ***\n")), true);
    assert_eq!(is_hrule(bytes!("   ***\n")), true);

    // but not four, or a tab which is equivalent to four spaces
    assert_eq!(is_hrule(bytes!("    ***\n")), false);
    assert_eq!(is_hrule(bytes!("\t***\n")), false);

    // need at least three
    assert_eq!(is_hrule(bytes!("*\n")), false);
    assert_eq!(is_hrule(bytes!("**\n")), false);
    assert_eq!(is_hrule(bytes!("* *\n")), false);
    assert_eq!(is_hrule(bytes!("   * *\n")), false);

    // Also works without newline at the end
    assert_eq!(is_hrule(bytes!("* * *")), true);

    // And underscores also supported
    assert_eq!(is_hrule(bytes!("___")), true);
    assert_eq!(is_hrule(bytes!("______________")), true);
    assert_eq!(is_hrule(bytes!(" ______________")), true);
}

#[bench]
fn bench_is_hrule(b: &mut BenchHarness) {
    let s = bytes!("   * * * * * * * * * * * * * * * *\n");
    b.iter(|| is_hrule(s));
}
<commit_msg>Benchmark against using slices.<commit_after>// This code is heavily inspired by hoedown.

extern crate extra;
use extra::test::BenchHarness;

static SP: u8 = ' ' as u8;
static NL: u8 = '\n' as u8;

/// Returns true if the line is a horizontal rule.
fn is_hrule_old(buf: &[u8]) -> bool {
    let len = buf.len();

    if len < 3 { return false }

    let mut pos: uint = 0;

    // Skip up to three leading spaces.
    //
    // We don't need to care about tabs here as in this position they are
    // equivalent to 4 spaces.
    if buf[0] == SP {
        pos += 1;
        if buf[1] == SP {
            pos += 1;
            if buf[2] == SP {
                pos += 1;
            }
        }
    }

    // We need at least 3 items
    if pos + 2 >= len { return false }

    let item = buf[pos];

    if !(item == ('*' as u8) || item == ('-' as u8) || item == ('_' as u8)) {
        return false;
    }

    let mut cnt: uint = 0; // The count of '*', '-' or '_'
    for &ch in buf.slice_from(pos).iter() {
        if ch == item {
            cnt += 1;
        } else if ch == NL {
            break;
        } else if ch != SP {
            return false;
        }
    }

    return cnt >= 3;
}

/// Returns true if the line is a horizontal rule.
fn is_hrule(buf: &[u8]) -> bool {
    let mut buf = buf;

    if buf.len() < 3 { return false }

    // Skip up to three leading spaces.
    //
    // We don't need to care about tabs here as in this position they are
    // equivalent to 4 spaces.
    if buf.head() == Some(&SP) { buf = buf.tail(); }
    if buf.head() == Some(&SP) { buf = buf.tail(); }
    if buf.head() == Some(&SP) { buf = buf.tail(); }

    // We need at least 3 items
    if buf.len() < 3 { return false }

    let item = buf[0];

    if !(item == ('*' as u8) || item == ('-' as u8) || item == ('_' as u8)) {
        return false;
    }

    let mut cnt: uint = 0; // The count of '*', '-' or '_'
    for &ch in buf.iter() {
        if ch == item {
            cnt += 1;
        } else if ch == NL {
            break;
        } else if ch != SP {
            return false;
        }
    }

    return cnt >= 3;
}

#[test]
fn test_is_hrule() {
    // examples as given on the markdown homepage
    assert_eq!(is_hrule(bytes!("* * *\n")), true);
    assert_eq!(is_hrule(bytes!("***\n")), true);
    assert_eq!(is_hrule(bytes!("*****\n")), true);
    assert_eq!(is_hrule(bytes!("- - -\n")), true);
    assert_eq!(is_hrule(bytes!("---------------------------------------\n")), true);

    // up to three spaces ignored
    assert_eq!(is_hrule(bytes!(" ***\n")), true);
    assert_eq!(is_hrule(bytes!("  ***\n")), true);
    assert_eq!(is_hrule(bytes!("   ***\n")), true);

    // but not four, or a tab which is equivalent to four spaces
    assert_eq!(is_hrule(bytes!("    ***\n")), false);
    assert_eq!(is_hrule(bytes!("\t***\n")), false);

    // need at least three
    assert_eq!(is_hrule(bytes!("*\n")), false);
    assert_eq!(is_hrule(bytes!("**\n")), false);
    assert_eq!(is_hrule(bytes!("* *\n")), false);
    assert_eq!(is_hrule(bytes!("   * *\n")), false);

    // Also works without newline at the end
    assert_eq!(is_hrule(bytes!("* * *")), true);

    // And underscores also supported
    assert_eq!(is_hrule(bytes!("___")), true);
    assert_eq!(is_hrule(bytes!("______________")), true);
    assert_eq!(is_hrule(bytes!(" ______________")), true);
}

#[bench]
fn bench_is_hrule(b: &mut BenchHarness) {
    let s = bytes!("   * * * * * * * * * * * * * * * *\n");
    b.iter(|| is_hrule(s));
}

#[bench]
fn bench_is_hrule_old(b: &mut BenchHarness) {
    let s = bytes!("   * * * * * * * * * * * * * * * *\n");
    b.iter(|| is_hrule_old(s));
}

<|endoftext|>