<commit_before>use ::sodiumoxide::crypto::hash::sha256;
use ::sodiumoxide::crypto::sign;
use ::sodiumoxide::randombytes;
use ::bs58;

const PUB_ADDRESS_PREFIX: &str = "GOD";
const PRIV_BUF_PREFIX: u8 = 0x01;
const PUB_BUF_PREFIX: u8 = 0x02;

pub trait Wif<T> {
    fn from_wif(s: &str) -> Option<T>;
    fn to_wif(&self) -> Box<str>;
}

#[derive(Clone, Debug, PartialEq)]
pub struct PublicKey {
    key: sign::PublicKey
}

impl PublicKey {
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.key.as_ref()
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        sign::verify_detached(sig, msg, &self.key)
    }

    #[inline]
    pub fn from_bytes(bytes: &[u8]) -> Option<PublicKey> {
        let key = sign::PublicKey::from_slice(bytes)?;
        Some(PublicKey { key })
    }

    #[inline]
    pub fn verify_sig_pair(pair: &SigPair, msg: &[u8]) -> bool {
        sign::verify_detached(&pair.signature, msg, &pair.pub_key.key)
    }
}

impl Wif<PublicKey> for PublicKey {
    fn from_wif(s: &str) -> Option<PublicKey> {
        if s.len() < 3 || &s[0..3] != PUB_ADDRESS_PREFIX { return None }
        let raw = bs58::decode(&s[3..]).into_vec().ok()?;
        if raw.len() != 37 || raw[0] != PUB_BUF_PREFIX { return None }

        let prefixed_key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(prefixed_key)[0..4];
            if checksum_a != checksum_b { return None }
        }

        let key = &prefixed_key[1 .. prefixed_key.len()];
        Some(PublicKey {
            key: sign::PublicKey::from_slice(key)?
        })
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf: Vec<u8> = Vec::<u8>::with_capacity(37);
        buf.push(PUB_BUF_PREFIX);
        buf.extend_from_slice(self.key.as_ref());

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        let mut s = bs58::encode(buf).into_string();
        s.insert_str(0, PUB_ADDRESS_PREFIX);
        s.into_boxed_str()
    }
}

pub struct PrivateKey {
    seed: sign::Seed,
    key: sign::SecretKey
}

impl PrivateKey {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> sign::Signature {
        sign::sign_detached(msg, &self.key)
    }
}

impl Wif<KeyPair> for PrivateKey {
    fn from_wif(s: &str) -> Option<KeyPair> {
        let raw = bs58::decode(s).into_vec().ok()?;
        if raw.len() != 37 || raw[0] != PRIV_BUF_PREFIX { return None }

        let key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(key)[0..4];
            if checksum_a != checksum_b { return None }
        }

        let seed = sign::Seed::from_slice(&key[1..])?;
        let (pk, sk) = sign::keypair_from_seed(&seed);
        Some(KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        }))
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf = Vec::<u8>::with_capacity(37);
        buf.push(PRIV_BUF_PREFIX);
        buf.extend_from_slice(&self.seed.0);

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        bs58::encode(buf).into_string().into_boxed_str()
    }
}

pub struct KeyPair(pub PublicKey, pub PrivateKey);

impl KeyPair {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> SigPair {
        SigPair {
            pub_key: self.0.clone(),
            signature: self.1.sign(msg)
        }
    }

    pub fn gen_keypair() -> KeyPair {
        let mut raw_seed: [u8; sign::SEEDBYTES] = [0; sign::SEEDBYTES];
        randombytes::randombytes_into(&mut raw_seed);
        let seed = sign::Seed::from_slice(&raw_seed).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        })
    }
}

#[derive(Debug, PartialEq)]
pub struct SigPair {
    pub pub_key: PublicKey,
    pub signature: sign::Signature
}

#[inline]
pub fn double_sha256(buf: &[u8]) -> sha256::Digest {
    sha256::hash(sha256::hash(buf).as_ref())
}

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

    #[test]
    fn test_create_and_recover_keys() {
        let kp = KeyPair::gen_keypair();

        let pk = &*kp.0.to_wif();
        assert_eq!(&*PublicKey::from_wif(pk).unwrap().to_wif(), pk);

        let sk = &*kp.1.to_wif();
        assert_eq!(&*PrivateKey::from_wif(sk).unwrap().1.to_wif(), sk);
    }

    #[test]
    fn test_import_keys_from_wif() {
        let kp = PrivateKey::from_wif("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").unwrap();
        assert_eq!(&*kp.1.to_wif(), "3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM");
        assert_eq!(&*kp.0.to_wif(), "GOD52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8");
    }

    #[test]
    fn test_sign_message() {
        let msg = "Hello world!".as_bytes();
        let kp = KeyPair::gen_keypair();

        let sig = &kp.1.sign(msg);
        assert!(kp.0.verify(msg, sig));

        let pair = SigPair {
            pub_key: kp.0,
            signature: *sig
        };
        assert!(PublicKey::verify_sig_pair(&pair, msg));

        // Test bad keys
        let kp = KeyPair::gen_keypair();
        assert!(!kp.0.verify(msg, sig));
    }
}
<commit_msg>Add crypto convenience verification methods<commit_after>use ::sodiumoxide::crypto::hash::sha256;
use ::sodiumoxide::crypto::sign;
use ::sodiumoxide::randombytes;
use ::bs58;

const PUB_ADDRESS_PREFIX: &str = "GOD";
const PRIV_BUF_PREFIX: u8 = 0x01;
const PUB_BUF_PREFIX: u8 = 0x02;

pub trait Wif<T> {
    fn from_wif(s: &str) -> Option<T>;
    fn to_wif(&self) -> Box<str>;
}

#[derive(Clone, Debug, PartialEq)]
pub struct PublicKey {
    key: sign::PublicKey
}

impl PublicKey {
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.key.as_ref()
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        sign::verify_detached(sig, msg, &self.key)
    }

    #[inline]
    pub fn from_bytes(bytes: &[u8]) -> Option<PublicKey> {
        let key = sign::PublicKey::from_slice(bytes)?;
        Some(PublicKey { key })
    }
}

impl Wif<PublicKey> for PublicKey {
    fn from_wif(s: &str) -> Option<PublicKey> {
        if s.len() < 3 || &s[0..3] != PUB_ADDRESS_PREFIX { return None }
        let raw = bs58::decode(&s[3..]).into_vec().ok()?;
        if raw.len() != 37 || raw[0] != PUB_BUF_PREFIX { return None }

        let prefixed_key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(prefixed_key)[0..4];
            if checksum_a != checksum_b { return None }
        }

        let key = &prefixed_key[1 .. prefixed_key.len()];
        Some(PublicKey {
            key: sign::PublicKey::from_slice(key)?
        })
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf: Vec<u8> = Vec::<u8>::with_capacity(37);
        buf.push(PUB_BUF_PREFIX);
        buf.extend_from_slice(self.key.as_ref());

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        let mut s = bs58::encode(buf).into_string();
        s.insert_str(0, PUB_ADDRESS_PREFIX);
        s.into_boxed_str()
    }
}

pub struct PrivateKey {
    seed: sign::Seed,
    key: sign::SecretKey
}

impl PrivateKey {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> sign::Signature {
        sign::sign_detached(msg, &self.key)
    }
}

impl Wif<KeyPair> for PrivateKey {
    fn from_wif(s: &str) -> Option<KeyPair> {
        let raw = bs58::decode(s).into_vec().ok()?;
        if raw.len() != 37 || raw[0] != PRIV_BUF_PREFIX { return None }

        let key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(key)[0..4];
            if checksum_a != checksum_b { return None }
        }

        let seed = sign::Seed::from_slice(&key[1..])?;
        let (pk, sk) = sign::keypair_from_seed(&seed);
        Some(KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        }))
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf = Vec::<u8>::with_capacity(37);
        buf.push(PRIV_BUF_PREFIX);
        buf.extend_from_slice(&self.seed.0);

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        bs58::encode(buf).into_string().into_boxed_str()
    }
}

pub struct KeyPair(pub PublicKey, pub PrivateKey);

impl KeyPair {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> SigPair {
        SigPair {
            pub_key: self.0.clone(),
            signature: self.1.sign(msg)
        }
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        PublicKey::verify(&self.0, msg, sig)
    }

    pub fn gen_keypair() -> KeyPair {
        let mut raw_seed: [u8; sign::SEEDBYTES] = [0; sign::SEEDBYTES];
        randombytes::randombytes_into(&mut raw_seed);
        let seed = sign::Seed::from_slice(&raw_seed).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        })
    }
}

#[derive(Debug, PartialEq)]
pub struct SigPair {
    pub pub_key: PublicKey,
    pub signature: sign::Signature
}

impl SigPair {
    #[inline]
    pub fn verify(&self, msg: &[u8]) -> bool {
        sign::verify_detached(&self.signature, msg, &self.pub_key.key)
    }
}

#[inline]
pub fn double_sha256(buf: &[u8]) -> sha256::Digest {
    sha256::hash(sha256::hash(buf).as_ref())
}

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

    #[test]
    fn test_create_and_recover_keys() {
        let kp = KeyPair::gen_keypair();

        let pk = &*kp.0.to_wif();
        assert_eq!(&*PublicKey::from_wif(pk).unwrap().to_wif(), pk);

        let sk = &*kp.1.to_wif();
        assert_eq!(&*PrivateKey::from_wif(sk).unwrap().1.to_wif(), sk);
    }

    #[test]
    fn test_import_keys_from_wif() {
        let kp = PrivateKey::from_wif("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").unwrap();
        assert_eq!(&*kp.1.to_wif(), "3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM");
        assert_eq!(&*kp.0.to_wif(), "GOD52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8");
    }

    #[test]
    fn test_sign_message() {
        let msg = "Hello world!".as_bytes();
        let kp = KeyPair::gen_keypair();

        let sig = &kp.1.sign(msg);
        assert!(kp.0.verify(msg, sig));

        let pair = SigPair {
            pub_key: kp.0,
            signature: *sig
        };
        assert!(pair.verify(msg));

        // Test bad keys
        let kp = KeyPair::gen_keypair();
        assert!(!kp.verify(msg, sig));
    }
}
<|endoftext|>


<commit_before>#![feature(globs, phase, macro_rules)]

extern crate syntax;
extern crate rustc;

#[phase(link)]
extern crate regex;

#[phase(link, plugin)]
extern crate log;

#[phase(plugin)] extern crate regex_macros;

use std::collections::HashMap;
use std::io::File;

use syntax::parse;
use syntax::parse::lexer;
use rustc::driver::{session, config};

use syntax::ast;
use syntax::ast::Name;
use syntax::parse::token::*;
use syntax::parse::lexer::TokenAndSpan;

fn parse_token_list(file: &str) -> HashMap<String, Token> {
    fn id() -> Token {
        IDENT(ast::Ident { name: Name(0), ctxt: 0, }, false)
    }

    let mut res = HashMap::new();

    res.insert("-1".to_string(), EOF);

    for line in file.split('\n') {
        let eq = match line.trim().rfind('=') {
            Some(val) => val,
            None => continue
        };

        let val = line.slice_to(eq);
        let num = line.slice_from(eq + 1);

        let tok = match val {
            "SHR" => BINOP(SHR),
            "DOLLAR" => DOLLAR,
            "LT" => LT,
            "STAR" => BINOP(STAR),
            "FLOAT_SUFFIX" => id(),
            "INT_SUFFIX" => id(),
            "SHL" => BINOP(SHL),
            "LBRACE" => LBRACE,
            "RARROW" => RARROW,
            "LIT_STR" => LIT_STR(Name(0)),
            "DOTDOT" => DOTDOT,
            "MOD_SEP" => MOD_SEP,
            "DOTDOTDOT" => DOTDOTDOT,
            "NOT" => NOT,
            "AND" => BINOP(AND),
            "LPAREN" => LPAREN,
            "ANDAND" => ANDAND,
            "AT" => AT,
            "LBRACKET" => LBRACKET,
            "LIT_STR_RAW" => LIT_STR_RAW(Name(0), 0),
            "RPAREN" => RPAREN,
            "SLASH" => BINOP(SLASH),
            "COMMA" => COMMA,
            "LIFETIME" => LIFETIME(ast::Ident { name: Name(0), ctxt: 0 }),
            "CARET" => BINOP(CARET),
            "TILDE" => TILDE,
            "IDENT" => id(),
            "PLUS" => BINOP(PLUS),
            "LIT_CHAR" => LIT_CHAR(Name(0)),
            "LIT_BYTE" => LIT_BYTE(Name(0)),
            "EQ" => EQ,
            "RBRACKET" => RBRACKET,
            "COMMENT" => COMMENT,
            "DOC_COMMENT" => DOC_COMMENT(Name(0)),
            "DOT" => DOT,
            "EQEQ" => EQEQ,
            "NE" => NE,
            "GE" => GE,
            "PERCENT" => BINOP(PERCENT),
            "RBRACE" => RBRACE,
            "BINOP" => BINOP(PLUS),
            "POUND" => POUND,
            "OROR" => OROR,
            "LIT_INTEGER" => LIT_INTEGER(Name(0)),
            "BINOPEQ" => BINOPEQ(PLUS),
            "LIT_FLOAT" => LIT_FLOAT(Name(0)),
            "WHITESPACE" => WS,
            "UNDERSCORE" => UNDERSCORE,
            "MINUS" => BINOP(MINUS),
            "SEMI" => SEMI,
            "COLON" => COLON,
            "FAT_ARROW" => FAT_ARROW,
            "OR" => BINOP(OR),
            "GT" => GT,
            "LE" => LE,
            "LIT_BINARY" => LIT_BINARY(Name(0)),
            "LIT_BINARY_RAW" => LIT_BINARY_RAW(Name(0), 0),
            _ => continue
        };

        res.insert(num.to_string(), tok);
    }

    debug!("Token map: {}", res);
    res
}

fn str_to_binop(s: &str) -> BinOp {
    match s {
        "+" => PLUS,
        "/" => SLASH,
        "-" => MINUS,
        "*" => STAR,
        "%" => PERCENT,
        "^" => CARET,
        "&" => AND,
        "|" => OR,
        "<<" => SHL,
        ">>" => SHR,
        _ => fail!("Bad binop str `{}`", s)
    }
}

/// Assuming a string/binary literal, strip out the leading/trailing
/// hashes and surrounding quotes/raw/binary prefix.
fn fix(mut lit: &str) -> ast::Name {
    if lit.char_at(0) == 'r' {
        if lit.char_at(1) == 'b' {
            lit = lit.slice_from(2)
        } else {
            lit = lit.slice_from(1);
        }
    } else if lit.char_at(0) == 'b' {
        lit = lit.slice_from(1);
    }

    let leading_hashes = count(lit);

    // +1/-1 to adjust for single quotes
    parse::token::intern(lit.slice(leading_hashes + 1, lit.len() - leading_hashes - 1))
}

/// Assuming a char/byte literal, strip the 'b' prefix and the single quotes.
fn fixchar(mut lit: &str) -> ast::Name {
    if lit.char_at(0) == 'b' {
        lit = lit.slice_from(1);
    }

    parse::token::intern(lit.slice(1, lit.len() - 1))
}

fn count(lit: &str) -> uint {
    lit.chars().take_while(|c| *c == '#').count()
}

fn parse_antlr_token(s: &str, tokens: &HashMap<String, Token>) -> TokenAndSpan {
    let re = regex!(r"\[@(?P<seq>\d+),(?P<start>\d+):(?P<end>\d+)='(?P<content>.+?)',<(?P<toknum>-?\d+)>,\d+:\d+]");

    let m = re.captures(s).expect(format!("The regex didn't match {}", s).as_slice());
    let start = m.name("start");
    let end = m.name("end");
    let toknum = m.name("toknum");
    let content = m.name("content");

    let proto_tok = tokens.find_equiv(&toknum).expect(format!("didn't find token {} in the map", toknum).as_slice());

    let nm = parse::token::intern(content);

    debug!("What we got: content (`{}`), proto: {}", content, proto_tok);

    let real_tok = match *proto_tok {
        BINOP(..) => BINOP(str_to_binop(content)),
        BINOPEQ(..) => BINOPEQ(str_to_binop(content.slice_to(content.len() - 1))),
        LIT_STR(..) => LIT_STR(fix(content)),
        LIT_STR_RAW(..) => LIT_STR_RAW(fix(content), count(content)),
        LIT_CHAR(..) => LIT_CHAR(fixchar(content)),
        LIT_BYTE(..) => LIT_BYTE(fixchar(content)),
        DOC_COMMENT(..) => DOC_COMMENT(nm),
        LIT_INTEGER(..) => LIT_INTEGER(nm),
        LIT_FLOAT(..) => LIT_FLOAT(nm),
        LIT_BINARY(..) => LIT_BINARY(nm),
        LIT_BINARY_RAW(..) => LIT_BINARY_RAW(fix(content), count(content)),
        IDENT(..) => IDENT(ast::Ident { name: nm, ctxt: 0 }, true),
        LIFETIME(..) => LIFETIME(ast::Ident { name: nm, ctxt: 0 }),
        ref t => t.clone()
    };

    let offset = if real_tok == EOF {
        1
    } else {
        0
    };

    let sp = syntax::codemap::Span {
        lo: syntax::codemap::BytePos(from_str::<u32>(start).unwrap() - offset),
        hi: syntax::codemap::BytePos(from_str::<u32>(end).unwrap() + 1),
        expn_info: None
    };

    TokenAndSpan {
        tok: real_tok,
        sp: sp
    }
}

fn tok_cmp(a: &Token, b: &Token) -> bool {
    match a {
        &IDENT(id, _) => match b {
                &IDENT(id2, _) => id == id2,
                _ => false
        },
        _ => a == b
    }
}

fn main() {
    fn next(r: &mut lexer::StringReader) -> TokenAndSpan {
        use syntax::parse::lexer::Reader;
        r.next_token()
    }

    let args = std::os::args();

    let token_map = parse_token_list(File::open(&Path::new(args.get(2).as_slice())).unwrap().read_to_string().unwrap().as_slice());
    let mut stdin = std::io::stdin();
    let mut antlr_tokens = stdin.lines().map(|l| parse_antlr_token(l.unwrap().as_slice().trim(), &token_map));

    let code = File::open(&Path::new(args.get(1).as_slice())).unwrap().read_to_string().unwrap();
    let options = config::basic_options();
    let session = session::build_session(options, None,
                                         syntax::diagnostics::registry::Registry::new([]));
    let filemap = parse::string_to_filemap(&session.parse_sess,
                                           code,
                                           String::from_str("<n/a>"));
    let mut lexer = lexer::StringReader::new(session.diagnostic(), filemap);

    for antlr_tok in antlr_tokens {
        let rustc_tok = next(&mut lexer);
        if rustc_tok.tok == EOF && antlr_tok.tok == EOF {
            continue
        }

        assert!(rustc_tok.sp == antlr_tok.sp, "{} and {} have different spans", rustc_tok, antlr_tok);

        macro_rules! matches (
            ( $($x:pat),+ ) => (
                match rustc_tok.tok {
                    $($x => match antlr_tok.tok {
                        $x => {
                            if !tok_cmp(&rustc_tok.tok, &antlr_tok.tok) {
                                // FIXME #15677: needs more robust escaping in
                                // antlr
                                warn!("Different names for {} and {}", rustc_tok, antlr_tok);
                            }
                        }
                        _ => fail!("{} is not {}", antlr_tok, rustc_tok)
                    },)*
                    ref c => assert!(c == &antlr_tok.tok, "{} is not {}", rustc_tok, antlr_tok)
                }
            )
        )

        matches!(LIT_BYTE(..),
            LIT_CHAR(..),
            LIT_INTEGER(..),
            LIT_FLOAT(..),
            LIT_STR(..),
            LIT_STR_RAW(..),
            LIT_BINARY(..),
            LIT_BINARY_RAW(..),
            IDENT(..),
            LIFETIME(..),
            INTERPOLATED(..),
            DOC_COMMENT(..),
            SHEBANG(..)
        );
    }
}
<commit_msg>Break apart long lines in verify.rs<commit_after>#![feature(globs, phase, macro_rules)]

extern crate syntax;
extern crate rustc;

#[phase(link)]
extern crate regex;

#[phase(link, plugin)]
extern crate log;

#[phase(plugin)] extern crate regex_macros;

use std::collections::HashMap;
use std::io::File;

use syntax::parse;
use syntax::parse::lexer;
use rustc::driver::{session, config};

use syntax::ast;
use syntax::ast::Name;
use syntax::parse::token::*;
use syntax::parse::lexer::TokenAndSpan;

fn parse_token_list(file: &str) -> HashMap<String, Token> {
    fn id() -> Token {
        IDENT(ast::Ident { name: Name(0), ctxt: 0, }, false)
    }

    let mut res = HashMap::new();

    res.insert("-1".to_string(), EOF);

    for line in file.split('\n') {
        let eq = match line.trim().rfind('=') {
            Some(val) => val,
            None => continue
        };

        let val = line.slice_to(eq);
        let num = line.slice_from(eq + 1);

        let tok = match val {
            "SHR" => BINOP(SHR),
            "DOLLAR" => DOLLAR,
            "LT" => LT,
            "STAR" => BINOP(STAR),
            "FLOAT_SUFFIX" => id(),
            "INT_SUFFIX" => id(),
            "SHL" => BINOP(SHL),
            "LBRACE" => LBRACE,
            "RARROW" => RARROW,
            "LIT_STR" => LIT_STR(Name(0)),
            "DOTDOT" => DOTDOT,
            "MOD_SEP" => MOD_SEP,
            "DOTDOTDOT" => DOTDOTDOT,
            "NOT" => NOT,
            "AND" => BINOP(AND),
            "LPAREN" => LPAREN,
            "ANDAND" => ANDAND,
            "AT" => AT,
            "LBRACKET" => LBRACKET,
            "LIT_STR_RAW" => LIT_STR_RAW(Name(0), 0),
            "RPAREN" => RPAREN,
            "SLASH" => BINOP(SLASH),
            "COMMA" => COMMA,
            "LIFETIME" => LIFETIME(ast::Ident { name: Name(0), ctxt: 0 }),
            "CARET" => BINOP(CARET),
            "TILDE" => TILDE,
            "IDENT" => id(),
            "PLUS" => BINOP(PLUS),
            "LIT_CHAR" => LIT_CHAR(Name(0)),
            "LIT_BYTE" => LIT_BYTE(Name(0)),
            "EQ" => EQ,
            "RBRACKET" => RBRACKET,
            "COMMENT" => COMMENT,
            "DOC_COMMENT" => DOC_COMMENT(Name(0)),
            "DOT" => DOT,
            "EQEQ" => EQEQ,
            "NE" => NE,
            "GE" => GE,
            "PERCENT" => BINOP(PERCENT),
            "RBRACE" => RBRACE,
            "BINOP" => BINOP(PLUS),
            "POUND" => POUND,
            "OROR" => OROR,
            "LIT_INTEGER" => LIT_INTEGER(Name(0)),
            "BINOPEQ" => BINOPEQ(PLUS),
            "LIT_FLOAT" => LIT_FLOAT(Name(0)),
            "WHITESPACE" => WS,
            "UNDERSCORE" => UNDERSCORE,
            "MINUS" => BINOP(MINUS),
            "SEMI" => SEMI,
            "COLON" => COLON,
            "FAT_ARROW" => FAT_ARROW,
            "OR" => BINOP(OR),
            "GT" => GT,
            "LE" => LE,
            "LIT_BINARY" => LIT_BINARY(Name(0)),
            "LIT_BINARY_RAW" => LIT_BINARY_RAW(Name(0), 0),
            _ => continue
        };

        res.insert(num.to_string(), tok);
    }

    debug!("Token map: {}", res);
    res
}

fn str_to_binop(s: &str) -> BinOp {
    match s {
        "+" => PLUS,
        "/" => SLASH,
        "-" => MINUS,
        "*" => STAR,
        "%" => PERCENT,
        "^" => CARET,
        "&" => AND,
        "|" => OR,
        "<<" => SHL,
        ">>" => SHR,
        _ => fail!("Bad binop str `{}`", s)
    }
}

/// Assuming a string/binary literal, strip out the leading/trailing
/// hashes and surrounding quotes/raw/binary prefix.
fn fix(mut lit: &str) -> ast::Name {
    if lit.char_at(0) == 'r' {
        if lit.char_at(1) == 'b' {
            lit = lit.slice_from(2)
        } else {
            lit = lit.slice_from(1);
        }
    } else if lit.char_at(0) == 'b' {
        lit = lit.slice_from(1);
    }

    let leading_hashes = count(lit);

    // +1/-1 to adjust for single quotes
    parse::token::intern(lit.slice(leading_hashes + 1, lit.len() - leading_hashes - 1))
}

/// Assuming a char/byte literal, strip the 'b' prefix and the single quotes.
fn fixchar(mut lit: &str) -> ast::Name {
    if lit.char_at(0) == 'b' {
        lit = lit.slice_from(1);
    }

    parse::token::intern(lit.slice(1, lit.len() - 1))
}

fn count(lit: &str) -> uint {
    lit.chars().take_while(|c| *c == '#').count()
}

fn parse_antlr_token(s: &str, tokens: &HashMap<String, Token>) -> TokenAndSpan {
    let re = regex!(r"\[@(?P<seq>\d+),(?P<start>\d+):(?P<end>\d+)='(?P<content>.+?)',<(?P<toknum>-?\d+)>,\d+:\d+]");

    let m = re.captures(s).expect(format!("The regex didn't match {}", s).as_slice());
    let start = m.name("start");
    let end = m.name("end");
    let toknum = m.name("toknum");
    let content = m.name("content");

    let proto_tok = tokens.find_equiv(&toknum).expect(format!("didn't find token {} in the map", toknum).as_slice());

    let nm = parse::token::intern(content);

    debug!("What we got: content (`{}`), proto: {}", content, proto_tok);

    let real_tok = match *proto_tok {
        BINOP(..) => BINOP(str_to_binop(content)),
        BINOPEQ(..) => BINOPEQ(str_to_binop(content.slice_to(content.len() - 1))),
        LIT_STR(..) => LIT_STR(fix(content)),
        LIT_STR_RAW(..) => LIT_STR_RAW(fix(content), count(content)),
        LIT_CHAR(..) => LIT_CHAR(fixchar(content)),
        LIT_BYTE(..) => LIT_BYTE(fixchar(content)),
        DOC_COMMENT(..) => DOC_COMMENT(nm),
        LIT_INTEGER(..) => LIT_INTEGER(nm),
        LIT_FLOAT(..) => LIT_FLOAT(nm),
        LIT_BINARY(..) => LIT_BINARY(nm),
        LIT_BINARY_RAW(..) => LIT_BINARY_RAW(fix(content), count(content)),
        IDENT(..) => IDENT(ast::Ident { name: nm, ctxt: 0 }, true),
        LIFETIME(..) => LIFETIME(ast::Ident { name: nm, ctxt: 0 }),
        ref t => t.clone()
    };

    let offset = if real_tok == EOF {
        1
    } else {
        0
    };

    let sp = syntax::codemap::Span {
        lo: syntax::codemap::BytePos(from_str::<u32>(start).unwrap() - offset),
        hi: syntax::codemap::BytePos(from_str::<u32>(end).unwrap() + 1),
        expn_info: None
    };

    TokenAndSpan {
        tok: real_tok,
        sp: sp
    }
}

fn tok_cmp(a: &Token, b: &Token) -> bool {
    match a {
        &IDENT(id, _) => match b {
                &IDENT(id2, _) => id == id2,
                _ => false
        },
        _ => a == b
    }
}

fn main() {
    fn next(r: &mut lexer::StringReader) -> TokenAndSpan {
        use syntax::parse::lexer::Reader;
        r.next_token()
    }

    let args = std::os::args();

    let mut token_file = File::open(&Path::new(args.get(2).as_slice()));
    let token_map = parse_token_list(token_file.read_to_string().unwrap().as_slice());

    let mut stdin = std::io::stdin();
    let mut antlr_tokens = stdin.lines().map(|l| parse_antlr_token(l.unwrap().as_slice().trim(), &token_map));

    let code = File::open(&Path::new(args.get(1).as_slice())).unwrap().read_to_string().unwrap();
    let options = config::basic_options();
    let session = session::build_session(options, None,
                                         syntax::diagnostics::registry::Registry::new([]));
    let filemap = parse::string_to_filemap(&session.parse_sess,
                                           code,
                                           String::from_str("<n/a>"));
    let mut lexer = lexer::StringReader::new(session.diagnostic(), filemap);

    for antlr_tok in antlr_tokens {
        let rustc_tok = next(&mut lexer);
        if rustc_tok.tok == EOF && antlr_tok.tok == EOF {
            continue
        }

        assert!(rustc_tok.sp == antlr_tok.sp, "{} and {} have different spans", rustc_tok, antlr_tok);

        macro_rules! matches (
            ( $($x:pat),+ ) => (
                match rustc_tok.tok {
                    $($x => match antlr_tok.tok {
                        $x => {
                            if !tok_cmp(&rustc_tok.tok, &antlr_tok.tok) {
                                // FIXME #15677: needs more robust escaping in
                                // antlr
                                warn!("Different names for {} and {}", rustc_tok, antlr_tok);
                            }
                        }
                        _ => fail!("{} is not {}", antlr_tok, rustc_tok)
                    },)*
                    ref c => assert!(c == &antlr_tok.tok, "{} is not {}", rustc_tok, antlr_tok)
                }
            )
        )

        matches!(LIT_BYTE(..),
            LIT_CHAR(..),
            LIT_INTEGER(..),
            LIT_FLOAT(..),
            LIT_STR(..),
            LIT_STR_RAW(..),
            LIT_BINARY(..),
            LIT_BINARY_RAW(..),
            IDENT(..),
            LIFETIME(..),
            INTERPOLATED(..),
            DOC_COMMENT(..),
            SHEBANG(..)
        );
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::path;

use ::image;
use gfx;

use crate::context::{Context, DebugId};
use crate::error::GameError;
use crate::error::GameResult;
use crate::filesystem;
use crate::graphics;
use crate::graphics::shader::*;
use crate::graphics::*;

/// Generic in-GPU-memory image data available to be drawn on the screen.
/// You probably just want to look at the `Image` type.
#[derive(Clone, PartialEq)]
pub struct ImageGeneric<B>
where
    B: BackendSpec,
{
    pub(crate) texture: gfx::handle::RawShaderResourceView<B::Resources>,
    pub(crate) texture_handle: gfx::handle::RawTexture<B::Resources>,
    pub(crate) sampler_info: gfx::texture::SamplerInfo,
    pub(crate) blend_mode: Option<BlendMode>,
    pub(crate) width: u16,
    pub(crate) height: u16,

    pub(crate) debug_id: DebugId,
}

impl<B> ImageGeneric<B>
where
    B: BackendSpec,
{
    /// A helper function that just takes a factory directly so we can make an image
    /// without needing the full context object, so we can create an Image while still
    /// creating the GraphicsContext.
    pub(crate) fn make_raw(
        factory: &mut <B as BackendSpec>::Factory,
        sampler_info: &texture::SamplerInfo,
        width: u16,
        height: u16,
        rgba: &[u8],
        color_format: gfx::format::Format,
        debug_id: DebugId,
    ) -> GameResult<Self> {
        if width == 0 || height == 0 {
            let msg = format!(
                "Tried to create a texture of size {}x{}, each dimension must
                be >0",
                width, height
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        // Check for overflow, which might happen on 32-bit systems.
        // Textures can be max u16*u16, pixels, but then have 4 bytes per pixel.
        let uwidth = width as usize;
        let uheight = height as usize;
        let expected_bytes = uwidth
            .checked_mul(uheight)
            .and_then(|size| size.checked_mul(4))
            .ok_or_else(|| {
                let msg = format!(
                    "Integer overflow in Image::make_raw, image size: {} {}",
                    uwidth, uheight
                );
                GameError::ResourceLoadError(msg)
            })?;
        if expected_bytes != rgba.len() {
            let msg = format!(
                "Tried to create a texture of size {}x{}, but gave {} bytes of data (expected {})",
                width,
                height,
                rgba.len(),
                expected_bytes
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        let kind = gfx::texture::Kind::D2(width, height, gfx::texture::AaMode::Single);
        use gfx::memory::Bind;
        let gfx::format::Format(surface_format, channel_type) = color_format;
        let texinfo = gfx::texture::Info {
            kind,
            levels: 1,
            format: surface_format,
            bind: Bind::SHADER_RESOURCE
                | Bind::RENDER_TARGET
                | Bind::TRANSFER_SRC
                | Bind::TRANSFER_DST,
            usage: gfx::memory::Usage::Dynamic,
        };
        let raw_tex = factory.create_texture_raw(
            texinfo,
            Some(channel_type),
            Some((&[rgba], gfx::texture::Mipmap::Provided)),
        )?;
        let resource_desc = gfx::texture::ResourceDesc {
            channel: channel_type,
            layer: None,
            min: 0,
            max: raw_tex.get_info().levels - 1,
            swizzle: gfx::format::Swizzle::new(),
        };
        let raw_view = factory.view_texture_as_shader_resource_raw(&raw_tex, resource_desc)?;
        Ok(Self {
            texture: raw_view,
            texture_handle: raw_tex,
            sampler_info: *sampler_info,
            blend_mode: None,
            width,
            height,
            debug_id,
        })
    }
}

/// In-GPU-memory image data available to be drawn on the screen,
/// using the OpenGL backend.
///
/// Under the hood this is just an `Arc`'ed texture handle and
/// some metadata, so cloning it is fairly cheap; it doesn't
/// make another copy of the underlying image data.
pub type Image = ImageGeneric<GlBackendSpec>;

/// The supported formats for saving an image.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ImageFormat {
    /// .png image format (defaults to RGBA with 8-bit channels.)
    Png,
}

impl Image {
    /// Load a new image from the file at the given path. The documentation for the
    /// [`filesystem`](../filesystem/index.html) module explains how the path must be specified.
    pub fn new<P: AsRef<path::Path>>(context: &mut Context, path: P) -> GameResult<Self> {
        let img = {
            let mut buf = Vec::new();
            let mut reader = context.filesystem.open(path)?;
            let _ = reader.read_to_end(&mut buf)?;
            image::load_from_memory(&buf)?.to_rgba()
        };
        let (width, height) = img.dimensions();
        Self::from_rgba8(context, width as u16, height as u16, &img)
    }

    /// Creates a new `Image` from the given buffer of `u8` RGBA values.
    ///
    /// The pixel layout is row-major.  That is,
    /// the first 4 `u8` values make the top-left pixel in the `Image`, the
    /// next 4 make the next pixel in the same row, and so on to the end of
    /// the row.  The next `width * 4` values make up the second row, and so
    /// on.
    pub fn from_rgba8(
        context: &mut Context,
        width: u16,
        height: u16,
        rgba: &[u8],
    ) -> GameResult<Self> {
        let debug_id = DebugId::get(context);
        let color_format = context.gfx_context.color_format();
        Self::make_raw(
            &mut *context.gfx_context.factory,
            &context.gfx_context.default_sampler_info,
            width,
            height,
            rgba,
            color_format,
            debug_id,
        )
    }

    /// Dumps the `Image`'s data to a `Vec` of `u8` RGBA values.
    pub fn to_rgba8(&self, ctx: &mut Context) -> GameResult<Vec<u8>> {
        use gfx::memory::Typed;
        use gfx::traits::FactoryExt;

        let gfx = &mut ctx.gfx_context;
        let w = self.width;
        let h = self.height;

        // Note: In the GFX example, the download buffer is created ahead of time
        // and updated on screen resize events. This may be preferable, but then
        // the buffer also needs to be updated when we switch to/from a canvas.
        // Unsure of the performance impact of creating this as it is needed.
        // Probably okay for now though, since this probably won't be a super
        // common operation.
        let dl_buffer = gfx
            .factory
            .create_download_buffer::<[u8; 4]>(w as usize * h as usize)?;

        let factory = &mut *gfx.factory;
        let mut local_encoder = GlBackendSpec::encoder(factory);

        local_encoder.copy_texture_to_buffer_raw(
            &self.texture_handle,
            None,
            gfx::texture::RawImageInfo {
                xoffset: 0,
                yoffset: 0,
                zoffset: 0,
                width: w as u16,
                height: h as u16,
                depth: 0,
                format: gfx.color_format(),
                mipmap: 0,
            },
            dl_buffer.raw(),
            0,
        )?;
        local_encoder.flush(&mut *gfx.device);

        let reader = gfx.factory.read_mapping(&dl_buffer)?;

        // intermediary buffer to avoid casting.
        let mut data = Vec::with_capacity(self.width as usize * self.height as usize * 4);
        // Assuming OpenGL backend whose typical readback option (glReadPixels) has origin at bottom left.
        // Image formats on the other hand usually deal with top right.
        for y in (0..self.height as usize).rev() {
            data.extend(
                reader
                    .iter()
                    .skip(y * self.width as usize)
                    .take(self.width as usize)
                    .flatten(),
            );
        }
        Ok(data)
    }

    /// Encode the `Image` to the given file format and
    /// write it out to the given path.
    ///
    /// See the [`filesystem`](../filesystem/index.html) module docs for where exactly
    /// the file will end up.
    pub fn encode<P: AsRef<path::Path>>(
        &self,
        ctx: &mut Context,
        format: ImageFormat,
        path: P,
    ) -> GameResult {
        use std::io;
        let data = self.to_rgba8(ctx)?;
        let f = filesystem::create(ctx, path)?;
        let writer = &mut io::BufWriter::new(f);
        let color_format = image::ColorType::RGBA(8);
        match format {
            ImageFormat::Png => image::png::PNGEncoder::new(writer)
                .encode(
                    &data,
                    u32::from(self.width),
                    u32::from(self.height),
                    color_format,
                )
                .map_err(Into::into),
        }
    }

    /// A little helper function that creates a new `Image` that is just
    /// a solid square of the given size and color.  Mainly useful for
    /// debugging.
    pub fn solid(context: &mut Context, size: u16, color: Color) -> GameResult<Self> {
        let (r, g, b, a) = color.into();
        let pixel_array: [u8; 4] = [r, g, b, a];
        let size_squared = size as usize * size as usize;
        let mut buffer = Vec::with_capacity(size_squared);
        for _i in 0..size_squared {
            buffer.extend(&pixel_array[..]);
        }
        Image::from_rgba8(context, size, size, &buffer)
    }

    /// Return the width of the image.
    pub fn width(&self) -> u16 {
        self.width
    }

    /// Return the height of the image.
    pub fn height(&self) -> u16 {
        self.height
    }

    /// Get the filter mode for the image.
    pub fn filter(&self) -> FilterMode {
        self.sampler_info.filter.into()
    }

    /// Set the filter mode for the image.
    pub fn set_filter(&mut self, mode: FilterMode) {
        self.sampler_info.filter = mode.into();
    }

    /// Returns the dimensions of the image.
    pub fn dimensions(&self) -> Rect {
        Rect::new(0.0, 0.0, f32::from(self.width()), f32::from(self.height()))
    }

    /// Gets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn wrap(&self) -> (WrapMode, WrapMode) {
        (self.sampler_info.wrap_mode.0, self.sampler_info.wrap_mode.1)
    }

    /// Sets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn set_wrap(&mut self, wrap_x: WrapMode, wrap_y: WrapMode) {
        self.sampler_info.wrap_mode.0 = wrap_x;
        self.sampler_info.wrap_mode.1 = wrap_y;
    }
}

impl fmt::Debug for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "<Image: {}x{}, {:p}, texture address {:p}, sampler: {:?}>",
            self.width(),
            self.height(),
            self,
            &self.texture,
            &self.sampler_info
        )
    }
}

impl Drawable for Image {
    fn draw(&self, ctx: &mut Context, param: DrawParam) -> GameResult {
        self.debug_id.assert(ctx);

        let gfx = &mut ctx.gfx_context;
        let src_width = param.src.w;
        let src_height = param.src.h;
        // We have to mess with the scale to make everything
        // be its-unit-size-in-pixels.
        let real_scale = nalgebra::Vector2::new(
            param.scale.x * src_width * f32::from(self.width),
            param.scale.y * src_height * f32::from(self.height),
        );

        let mut new_param = param;
        new_param.scale = real_scale.into();

        gfx.update_instance_properties(new_param.into())?;
        let sampler = gfx
            .samplers
            .get_or_insert(self.sampler_info, gfx.factory.as_mut());
        gfx.data.vbuf = gfx.quad_vertex_buffer.clone();
        let typed_thingy = gfx
            .backend_spec
            .raw_to_typed_shader_resource(self.texture.clone());
        gfx.data.tex = (typed_thingy, sampler);
        let previous_mode: Option<BlendMode> = if let Some(mode) = self.blend_mode {
            let current_mode = gfx.blend_mode();
            if current_mode != mode {
                gfx.set_blend_mode(mode)?;
                Some(current_mode)
            } else {
                None
            }
        } else {
            None
        };

        gfx.draw(None)?;
        if let Some(mode) = previous_mode {
            gfx.set_blend_mode(mode)?;
        }
        Ok(())
    }

    fn dimensions(&self, _: &mut Context) -> Option<graphics::Rect> {
        Some(self.dimensions())
    }

    fn set_blend_mode(&mut self, mode: Option<BlendMode>) {
        self.blend_mode = mode;
    }

    fn blend_mode(&self) -> Option<BlendMode> {
        self.blend_mode
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ContextBuilder;
    #[test]
    fn test_invalid_image_size() {
        let (ctx, _) = &mut ContextBuilder::new("unittest", "unittest").build().unwrap();
        let _i = assert!(Image::from_rgba8(ctx, 0, 0, &vec![]).is_err());
        let _i = assert!(Image::from_rgba8(ctx, 3432, 432, &vec![]).is_err());
        let _i = Image::from_rgba8(ctx, 2, 2, &vec![99; 16]).unwrap();
    }
}
<commit_msg>Add Image::from_bytes constructor<commit_after>use std::io::Read;
use std::path;

use ::image;
use gfx;

use crate::context::{Context, DebugId};
use crate::error::GameError;
use crate::error::GameResult;
use crate::filesystem;
use crate::graphics;
use crate::graphics::shader::*;
use crate::graphics::*;

/// Generic in-GPU-memory image data available to be drawn on the screen.
/// You probably just want to look at the `Image` type.
#[derive(Clone, PartialEq)]
pub struct ImageGeneric<B>
where
    B: BackendSpec,
{
    pub(crate) texture: gfx::handle::RawShaderResourceView<B::Resources>,
    pub(crate) texture_handle: gfx::handle::RawTexture<B::Resources>,
    pub(crate) sampler_info: gfx::texture::SamplerInfo,
    pub(crate) blend_mode: Option<BlendMode>,
    pub(crate) width: u16,
    pub(crate) height: u16,

    pub(crate) debug_id: DebugId,
}

impl<B> ImageGeneric<B>
where
    B: BackendSpec,
{
    /// A helper function that just takes a factory directly so we can make an image
    /// without needing the full context object, so we can create an Image while still
    /// creating the GraphicsContext.
    pub(crate) fn make_raw(
        factory: &mut <B as BackendSpec>::Factory,
        sampler_info: &texture::SamplerInfo,
        width: u16,
        height: u16,
        rgba: &[u8],
        color_format: gfx::format::Format,
        debug_id: DebugId,
    ) -> GameResult<Self> {
        if width == 0 || height == 0 {
            let msg = format!(
                "Tried to create a texture of size {}x{}, each dimension must
                be >0",
                width, height
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        // Check for overflow, which might happen on 32-bit systems.
        // Textures can be max u16*u16, pixels, but then have 4 bytes per pixel.
        let uwidth = width as usize;
        let uheight = height as usize;
        let expected_bytes = uwidth
            .checked_mul(uheight)
            .and_then(|size| size.checked_mul(4))
            .ok_or_else(|| {
                let msg = format!(
                    "Integer overflow in Image::make_raw, image size: {} {}",
                    uwidth, uheight
                );
                GameError::ResourceLoadError(msg)
            })?;
        if expected_bytes != rgba.len() {
            let msg = format!(
                "Tried to create a texture of size {}x{}, but gave {} bytes of data (expected {})",
                width,
                height,
                rgba.len(),
                expected_bytes
            );
            return Err(GameError::ResourceLoadError(msg));
        }
        let kind = gfx::texture::Kind::D2(width, height, gfx::texture::AaMode::Single);
        use gfx::memory::Bind;
        let gfx::format::Format(surface_format, channel_type) = color_format;
        let texinfo = gfx::texture::Info {
            kind,
            levels: 1,
            format: surface_format,
            bind: Bind::SHADER_RESOURCE
                | Bind::RENDER_TARGET
                | Bind::TRANSFER_SRC
                | Bind::TRANSFER_DST,
            usage: gfx::memory::Usage::Dynamic,
        };
        let raw_tex = factory.create_texture_raw(
            texinfo,
            Some(channel_type),
            Some((&[rgba], gfx::texture::Mipmap::Provided)),
        )?;
        let resource_desc = gfx::texture::ResourceDesc {
            channel: channel_type,
            layer: None,
            min: 0,
            max: raw_tex.get_info().levels - 1,
            swizzle: gfx::format::Swizzle::new(),
        };
        let raw_view = factory.view_texture_as_shader_resource_raw(&raw_tex, resource_desc)?;
        Ok(Self {
            texture: raw_view,
            texture_handle: raw_tex,
            sampler_info: *sampler_info,
            blend_mode: None,
            width,
            height,
            debug_id,
        })
    }
}

/// In-GPU-memory image data available to be drawn on the screen,
/// using the OpenGL backend.
///
/// Under the hood this is just an `Arc`'ed texture handle and
/// some metadata, so cloning it is fairly cheap; it doesn't
/// make another copy of the underlying image data.
pub type Image = ImageGeneric<GlBackendSpec>;

/// The supported formats for saving an image.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ImageFormat {
    /// .png image format (defaults to RGBA with 8-bit channels.)
    Png,
}

impl Image {
    /// Load a new image from the file at the given path. The documentation for the
    /// [`filesystem`](../filesystem/index.html) module explains how the path must be specified.
    pub fn new<P: AsRef<path::Path>>(context: &mut Context, path: P) -> GameResult<Self> {
        let mut buf = Vec::new();
        let mut reader = context.filesystem.open(path)?;
        let _ = reader.read_to_end(&mut buf)?;
        Self::from_bytes(context, &buf)
    }

    /// Creates a new `Image` from the given buffer, which should contain an image encoded
    /// in a supported image file format.
    pub fn from_bytes(context: &mut Context, bytes: &[u8]) -> GameResult<Self> {
        let img = image::load_from_memory(&bytes)?.to_rgba();
        let (width, height) = img.dimensions();
        Self::from_rgba8(context, width as u16, height as u16, &img)
    }

    /// Creates a new `Image` from the given buffer of `u8` RGBA values.
    ///
    /// The pixel layout is row-major.  That is,
    /// the first 4 `u8` values make the top-left pixel in the `Image`, the
    /// next 4 make the next pixel in the same row, and so on to the end of
    /// the row.  The next `width * 4` values make up the second row, and so
    /// on.
    pub fn from_rgba8(
        context: &mut Context,
        width: u16,
        height: u16,
        rgba: &[u8],
    ) -> GameResult<Self> {
        let debug_id = DebugId::get(context);
        let color_format = context.gfx_context.color_format();
        Self::make_raw(
            &mut *context.gfx_context.factory,
            &context.gfx_context.default_sampler_info,
            width,
            height,
            rgba,
            color_format,
            debug_id,
        )
    }

    /// Dumps the `Image`'s data to a `Vec` of `u8` RGBA values.
    pub fn to_rgba8(&self, ctx: &mut Context) -> GameResult<Vec<u8>> {
        use gfx::memory::Typed;
        use gfx::traits::FactoryExt;

        let gfx = &mut ctx.gfx_context;
        let w = self.width;
        let h = self.height;

        // Note: In the GFX example, the download buffer is created ahead of time
        // and updated on screen resize events. This may be preferable, but then
        // the buffer also needs to be updated when we switch to/from a canvas.
        // Unsure of the performance impact of creating this as it is needed.
        // Probably okay for now though, since this probably won't be a super
        // common operation.
        let dl_buffer = gfx
            .factory
            .create_download_buffer::<[u8; 4]>(w as usize * h as usize)?;

        let factory = &mut *gfx.factory;
        let mut local_encoder = GlBackendSpec::encoder(factory);

        local_encoder.copy_texture_to_buffer_raw(
            &self.texture_handle,
            None,
            gfx::texture::RawImageInfo {
                xoffset: 0,
                yoffset: 0,
                zoffset: 0,
                width: w as u16,
                height: h as u16,
                depth: 0,
                format: gfx.color_format(),
                mipmap: 0,
            },
            dl_buffer.raw(),
            0,
        )?;
        local_encoder.flush(&mut *gfx.device);

        let reader = gfx.factory.read_mapping(&dl_buffer)?;

        // intermediary buffer to avoid casting.
        let mut data = Vec::with_capacity(self.width as usize * self.height as usize * 4);
        // Assuming OpenGL backend whose typical readback option (glReadPixels) has origin at bottom left.
        // Image formats on the other hand usually deal with top right.
        for y in (0..self.height as usize).rev() {
            data.extend(
                reader
                    .iter()
                    .skip(y * self.width as usize)
                    .take(self.width as usize)
                    .flatten(),
            );
        }
        Ok(data)
    }

    /// Encode the `Image` to the given file format and
    /// write it out to the given path.
    ///
    /// See the [`filesystem`](../filesystem/index.html) module docs for where exactly
    /// the file will end up.
    pub fn encode<P: AsRef<path::Path>>(
        &self,
        ctx: &mut Context,
        format: ImageFormat,
        path: P,
    ) -> GameResult {
        use std::io;
        let data = self.to_rgba8(ctx)?;
        let f = filesystem::create(ctx, path)?;
        let writer = &mut io::BufWriter::new(f);
        let color_format = image::ColorType::RGBA(8);
        match format {
            ImageFormat::Png => image::png::PNGEncoder::new(writer)
                .encode(
                    &data,
                    u32::from(self.width),
                    u32::from(self.height),
                    color_format,
                )
                .map_err(Into::into),
        }
    }

    /// A little helper function that creates a new `Image` that is just
    /// a solid square of the given size and color.  Mainly useful for
    /// debugging.
    pub fn solid(context: &mut Context, size: u16, color: Color) -> GameResult<Self> {
        let (r, g, b, a) = color.into();
        let pixel_array: [u8; 4] = [r, g, b, a];
        let size_squared = size as usize * size as usize;
        let mut buffer = Vec::with_capacity(size_squared);
        for _i in 0..size_squared {
            buffer.extend(&pixel_array[..]);
        }
        Image::from_rgba8(context, size, size, &buffer)
    }

    /// Return the width of the image.
    pub fn width(&self) -> u16 {
        self.width
    }

    /// Return the height of the image.
    pub fn height(&self) -> u16 {
        self.height
    }

    /// Get the filter mode for the image.
    pub fn filter(&self) -> FilterMode {
        self.sampler_info.filter.into()
    }

    /// Set the filter mode for the image.
    pub fn set_filter(&mut self, mode: FilterMode) {
        self.sampler_info.filter = mode.into();
    }

    /// Returns the dimensions of the image.
    pub fn dimensions(&self) -> Rect {
        Rect::new(0.0, 0.0, f32::from(self.width()), f32::from(self.height()))
    }

    /// Gets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn wrap(&self) -> (WrapMode, WrapMode) {
        (self.sampler_info.wrap_mode.0, self.sampler_info.wrap_mode.1)
    }

    /// Sets the `Image`'s `WrapMode` along the X and Y axes.
    pub fn set_wrap(&mut self, wrap_x: WrapMode, wrap_y: WrapMode) {
        self.sampler_info.wrap_mode.0 = wrap_x;
        self.sampler_info.wrap_mode.1 = wrap_y;
    }
}

impl fmt::Debug for Image {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "<Image: {}x{}, {:p}, texture address {:p}, sampler: {:?}>",
            self.width(),
            self.height(),
            self,
            &self.texture,
            &self.sampler_info
        )
    }
}

impl Drawable for Image {
    fn draw(&self, ctx: &mut Context, param: DrawParam) -> GameResult {
        self.debug_id.assert(ctx);

        let gfx = &mut ctx.gfx_context;
        let src_width = param.src.w;
        let src_height = param.src.h;
        // We have to mess with the scale to make everything
        // be its-unit-size-in-pixels.
        let real_scale = nalgebra::Vector2::new(
            param.scale.x * src_width * f32::from(self.width),
            param.scale.y * src_height * f32::from(self.height),
        );

        let mut new_param = param;
        new_param.scale = real_scale.into();

        gfx.update_instance_properties(new_param.into())?;
        let sampler = gfx
            .samplers
            .get_or_insert(self.sampler_info, gfx.factory.as_mut());
        gfx.data.vbuf = gfx.quad_vertex_buffer.clone();
        let typed_thingy = gfx
            .backend_spec
            .raw_to_typed_shader_resource(self.texture.clone());
        gfx.data.tex = (typed_thingy, sampler);
        let previous_mode: Option<BlendMode> = if let Some(mode) = self.blend_mode {
            let current_mode = gfx.blend_mode();
            if current_mode != mode {
                gfx.set_blend_mode(mode)?;
                Some(current_mode)
            } else {
                None
            }
        } else {
            None
        };

        gfx.draw(None)?;
        if let Some(mode) = previous_mode {
            gfx.set_blend_mode(mode)?;
        }
        Ok(())
    }

    fn dimensions(&self, _: &mut Context) -> Option<graphics::Rect> {
        Some(self.dimensions())
    }

    fn set_blend_mode(&mut self, mode: Option<BlendMode>) {
        self.blend_mode = mode;
    }

    fn blend_mode(&self) -> Option<BlendMode> {
        self.blend_mode
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ContextBuilder;
    #[test]
    fn test_invalid_image_size() {
        let (ctx, _) = &mut ContextBuilder::new("unittest", "unittest").build().unwrap();
        let _i = assert!(Image::from_rgba8(ctx, 0, 0, &vec![]).is_err());
        let _i = assert!(Image::from_rgba8(ctx, 3432, 432, &vec![]).is_err());
        let _i = Image::from_rgba8(ctx, 2, 2, &vec![99; 16]).unwrap();
    }
}
<|endoftext|>


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

use kythe::corpus::Corpus;
use kythe::schema::{VName, Fact, NodeKind, EdgeKind, Complete};
use kythe::writer::EntryWriter;
use rustc::lint::{LateContext, LintContext, LintPass, LateLintPass, LintArray};
use rustc::hir;
use std::io::prelude::*;
use std::io::stderr;
use std::fs::File;
use std::str;
use syntax::ast;
use syntax::codemap::{Pos, CodeMap, Span};

// Represents the shared data between all the check_* functions in rustc's LateLintPass
pub struct KytheLintPass {
    pub corpus: Corpus,
    pub writer: Box<EntryWriter>,
}

impl KytheLintPass {
    // Emits the appropriate node and facts for an anchor defined by a span
    // and returns the node's vname
    fn anchor_from_span(&self, span: Span, codemap: &CodeMap) -> VName {
        let start = codemap.lookup_byte_offset(span.lo);
        let end = codemap.lookup_byte_offset(span.hi);

        let start_byte = start.pos.to_usize();
        let end_byte = end.pos.to_usize();
        self.anchor(&start.fm.name, start_byte, end_byte)
    }
    // Emits the appropriate node and facts for an anchor defined as a substring within a span
    // and returns the node's VName.
    fn anchor_from_sub_span(&self,
                            span: Span,
                            sub: &str,
                            codemap: &CodeMap)
                            -> Result<VName, String> {
        let start = codemap.lookup_byte_offset(span.lo);
        let snippet = match codemap.span_to_snippet(span) {
            Ok(s) => s,
            Err(e) => return Err(format!("{:?}", e)),
        };
        let sub_start = match snippet.find(sub) {
            None => return Err(format!("Substring: '{}' not found in snippet '{}'", sub, snippet)),
            Some(s) => s,
        };
        let start_byte = start.pos.to_usize() + sub_start;
        let end_byte = start_byte + sub.len();

        Ok(self.anchor(&start.fm.name, start_byte, end_byte))
    }

    fn anchor(&self, file_name: &str, start_byte: usize, end_byte: usize) -> VName {
        let vname = self.corpus.anchor_vname(file_name, start_byte, end_byte);

        let start_str: String = start_byte.to_string();
        let end_str: String = end_byte.to_string();
        let file_vname = self.corpus.file_vname(&file_name);

        self.writer.node(&vname, Fact::NodeKind, &NodeKind::Anchor);
        self.writer.node(&vname, Fact::LocStart, &start_str);
        self.writer.node(&vname, Fact::LocEnd, &end_str);
        self.writer.edge(&vname, EdgeKind::ChildOf, &file_vname);
        vname
    }
}

// A no-op implementation since we aren't reporting linting errors
impl LintPass for KytheLintPass {
    fn get_lints(&self) -> LintArray {
        lint_array!()
    }
}

impl LateLintPass for KytheLintPass {
    fn check_crate(&mut self, cx: &LateContext, _: &hir::Crate) {
        for ref f in cx.sess().codemap().files.borrow().iter() {
            // The codemap contains references to virtual files all labeled <<std_macro>>
            // These are skipped as per this check
            if !f.is_real_file() {
                continue;
            }

            let vname = self.corpus.file_vname(&f.name);
            let mut contents = String::new();

            let res = File::open(&f.name).and_then(|mut f| f.read_to_string(&mut contents));

            match res {
                Err(e) => {
                    writeln!(stderr(), "Failed to read file {}\n{:?}", f.name, e).unwrap();
                }
                Ok(_) => {
                    self.writer.node(&vname, Fact::Text, &contents);
                    self.writer.node(&vname, Fact::NodeKind, &NodeKind::File);
                }
            }
        }
    }

    fn check_fn_post(&mut self,
                     ctx: &LateContext,
                     kind: hir::intravisit::FnKind,
                     _: &hir::FnDecl,
                     _: &hir::Block,
                     span: Span,
                     id: ast::NodeId) {
        match kind {
            // Module-level function definitions
            hir::intravisit::FnKind::ItemFn(n, _, _, _, _, _, _) => {
                let fn_name = n.to_string();
                let decl_vname = self.anchor_from_span(span, ctx.sess().codemap());
                let fn_path = ctx.tcx.node_path_str(id);
                let fn_vname = self.corpus.fn_vname(&fn_path);
                self.writer.node(&fn_vname, Fact::NodeKind, &NodeKind::Function);
                self.writer.node(&fn_vname, Fact::Complete, &Complete::Definition);
                self.writer.edge(&decl_vname, EdgeKind::Defines, &fn_vname);

                match self.anchor_from_sub_span(span, &fn_name, ctx.sess().codemap()) {
                    Ok(bind_vname) => {
                        self.writer.edge(&bind_vname, EdgeKind::DefinesBinding, &fn_vname)
                    }
                    // TODO(djrenren): use a logging library
                    Err(e) => writeln!(stderr(), "{}", e).unwrap(),
                }
            }
            _ => (),
        }
    }
}
<commit_msg>Read file contents from in-memory maps<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use kythe::corpus::Corpus;
use kythe::schema::{VName, Fact, NodeKind, EdgeKind, Complete};
use kythe::writer::EntryWriter;
use rustc::lint::{LateContext, LintContext, LintPass, LateLintPass, LintArray};
use rustc::hir;
use std::io::prelude::*;
use std::io::stderr;
use syntax::ast;
use syntax::codemap::{Pos, CodeMap, Span};

// Represents the shared data between all the check_* functions in rustc's LateLintPass
pub struct KytheLintPass {
    pub corpus: Corpus,
    pub writer: Box<EntryWriter>,
}

impl KytheLintPass {
    // Emits the appropriate node and facts for an anchor defined by a span
    // and returns the node's vname
    fn anchor_from_span(&self, span: Span, codemap: &CodeMap) -> VName {
        let start = codemap.lookup_byte_offset(span.lo);
        let end = codemap.lookup_byte_offset(span.hi);

        let start_byte = start.pos.to_usize();
        let end_byte = end.pos.to_usize();
        self.anchor(&start.fm.name, start_byte, end_byte)
    }
    // Emits the appropriate node and facts for an anchor defined as a substring within a span
    // and returns the node's VName.
    fn anchor_from_sub_span(&self,
                            span: Span,
                            sub: &str,
                            codemap: &CodeMap)
                            -> Result<VName, String> {
        let start = codemap.lookup_byte_offset(span.lo);
        let snippet = match codemap.span_to_snippet(span) {
            Ok(s) => s,
            Err(e) => return Err(format!("{:?}", e)),
        };
        let sub_start = match snippet.find(sub) {
            None => return Err(format!("Substring: '{}' not found in snippet '{}'", sub, snippet)),
            Some(s) => s,
        };
        let start_byte = start.pos.to_usize() + sub_start;
        let end_byte = start_byte + sub.len();

        Ok(self.anchor(&start.fm.name, start_byte, end_byte))
    }

    fn anchor(&self, file_name: &str, start_byte: usize, end_byte: usize) -> VName {
        let vname = self.corpus.anchor_vname(file_name, start_byte, end_byte);

        let start_str: String = start_byte.to_string();
        let end_str: String = end_byte.to_string();
        let file_vname = self.corpus.file_vname(&file_name);

        self.writer.node(&vname, Fact::NodeKind, &NodeKind::Anchor);
        self.writer.node(&vname, Fact::LocStart, &start_str);
        self.writer.node(&vname, Fact::LocEnd, &end_str);
        self.writer.edge(&vname, EdgeKind::ChildOf, &file_vname);
        vname
    }
}

// A no-op implementation since we aren't reporting linting errors
impl LintPass for KytheLintPass {
    fn get_lints(&self) -> LintArray {
        lint_array!()
    }
}

impl LateLintPass for KytheLintPass {
    fn check_crate(&mut self, cx: &LateContext, _: &hir::Crate) {
        for ref f in cx.sess().codemap().files.borrow().iter() {
            // The codemap contains references to virtual files all labeled <<std_macro>>
            // These are skipped as per this check
            if !f.is_real_file() {
                continue;
            }

            // TODO(djrenren): Find out if/when this will actually panic
            // Presumably, by this point the file is determined to have
            // actually been on disk and loaded into the codemap.
            let content = f.src.as_ref().unwrap();
            let vname = self.corpus.file_vname(&f.name);

            self.writer.node(&vname, Fact::NodeKind, &NodeKind::File);
            self.writer.node(&vname, Fact::Text, content);
        }
    }

    fn check_fn_post(&mut self,
                     ctx: &LateContext,
                     kind: hir::intravisit::FnKind,
                     _: &hir::FnDecl,
                     _: &hir::Block,
                     span: Span,
                     id: ast::NodeId) {
        match kind {
            // Module-level function definitions
            hir::intravisit::FnKind::ItemFn(n, _, _, _, _, _, _) => {
                let fn_name = n.to_string();
                let decl_vname = self.anchor_from_span(span, ctx.sess().codemap());
                let fn_path = ctx.tcx.node_path_str(id);
                let fn_vname = self.corpus.fn_vname(&fn_path);
                self.writer.node(&fn_vname, Fact::NodeKind, &NodeKind::Function);
                self.writer.node(&fn_vname, Fact::Complete, &Complete::Definition);
                self.writer.edge(&decl_vname, EdgeKind::Defines, &fn_vname);

                match self.anchor_from_sub_span(span, &fn_name, ctx.sess().codemap()) {
                    Ok(bind_vname) => {
                        self.writer.edge(&bind_vname, EdgeKind::DefinesBinding, &fn_vname)
                    }
                    // TODO(djrenren): use a logging library
                    Err(e) => writeln!(stderr(), "{}", e).unwrap(),
                }
            }
            _ => (),
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use crate::data::frame::ArcFrame;
use crate::data::packet::Packet;

use crate::common::CodecList;
use crate::error::*;

/// Used to interact with a decoder.
pub trait Decoder: Send + Sync {
    // TODO support codec configuration using set_option
    // fn open(&mut self) -> Result<()>;
    /// Saves the extra data contained in a codec.
    fn set_extradata(&mut self, extra: &[u8]);
    /// Sends to the decoder a packet to be decoded.
    fn send_packet(&mut self, pkt: &Packet) -> Result<()>;
    /// Returns a decoded frame.
    fn receive_frame(&mut self) -> Result<ArcFrame>;
    /// Configures the decoder.
    fn configure(&mut self) -> Result<()>;
    /// Tells decoder to clear its internal state.
    fn flush(&mut self) -> Result<()>;
}

/// Codec descriptor.
///
/// Contains information on a codec and its own decoder.
#[derive(Debug)]
pub struct Descr {
    /// The codec name.
    pub codec: &'static str,
    /// The extended codec name.
    pub name: &'static str,
    /// The codec description.
    pub desc: &'static str,
    /// The codec MIME.
    pub mime: &'static str,
    // TODO more fields regarding capabilities
}

/// Auxiliary structure to encapsulate a decoder object and
/// its additional data.
pub struct Context<D: Decoder> {
    dec: D,
    // TODO: Queue up packets/frames
}

impl<D: Decoder> Context<D> {
    // TODO: More constructors
    /// Retrieves a codec descriptor from a codec list through its name,
    /// creates the relative decoder, and encapsulates it into a new `Context`.
    pub fn by_name<T: Descriptor<OutputDecoder = D> + ?Sized>(
        codecs: &Codecs<T>,
        name: &str,
    ) -> Option<Self> {
        if let Some(builder) = codecs.by_name(name) {
            let dec = builder.create();
            Some(Context { dec })
        } else {
            None
        }
    }
    /// Saves the extra data contained in a codec.
    pub fn set_extradata(&mut self, extra: &[u8]) {
        self.dec.set_extradata(extra);
    }

    /// Sends to the decoder a packet to be decoded.
    pub fn send_packet(&mut self, pkt: &Packet) -> Result<()> {
        self.dec.send_packet(pkt)
    }
    /// Returns a decoded frame.
    pub fn receive_frame(&mut self) -> Result<ArcFrame> {
        self.dec.receive_frame()
    }
    /// Configures the decoder.
    pub fn configure(&mut self) -> Result<()> {
        self.dec.configure()
    }

    /// Tells decoder to clear its internal state.
    pub fn flush(&mut self) -> Result<()> {
        self.dec.flush()
    }
}

/// Used to get the descriptor of a codec and create its own decoder.
pub trait Descriptor {
    type OutputDecoder: Decoder;

    /// Creates a new decoder for the requested codec.
    fn create(&self) -> Self::OutputDecoder;
    /// Returns the codec descriptor.
    fn describe(&self) -> &Descr;
}

/// A list of codec descriptors.
pub struct Codecs<T: 'static + Descriptor + ?Sized> {
    list: HashMap<&'static str, Vec<&'static T>>,
}

impl<T: Descriptor + ?Sized> CodecList for Codecs<T> {
    type D = T;

    fn new() -> Self {
        Codecs {
            list: HashMap::new(),
        }
    }

    // TODO more lookup functions
    fn by_name(&self, name: &str) -> Option<&'static Self::D> {
        self.list.get(name).map(|descs| descs[0])
    }

    fn append(&mut self, desc: &'static Self::D) {
        let codec_name = desc.describe().codec;

        self.list
            .entry(codec_name)
            .or_insert_with(Vec::new)
            .push(desc);
    }
}

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

    mod dummy {
        use super::super::*;
        use crate::data::pixel::Formaton;
        use std::sync::Arc;

        pub struct Dec {
            state: usize,
            #[allow(dead_code)]
            format: Option<Arc<Formaton>>,
        }

        pub struct Des {
            descr: Descr,
        }

        impl Descriptor for Des {
            type OutputDecoder = Dec;

            fn create(&self) -> Self::OutputDecoder {
                Dec {
                    state: 0,
                    format: None,
                }
            }

            fn describe(&self) -> &Descr {
                &self.descr
            }
        }

        impl Decoder for Dec {
            fn configure(&mut self) -> Result<()> {
                Ok(())
            }
            fn set_extradata(&mut self, extra: &[u8]) {
                if extra.len() > 4 {
                    self.state = 42;
                } else {
                    self.state = 12;
                }
            }
            fn send_packet(&mut self, _packet: &Packet) -> Result<()> {
                self.state += 1;
                Ok(())
            }
            fn receive_frame(&mut self) -> Result<ArcFrame> {
                unimplemented!()
            }
            fn flush(&mut self) -> Result<()> {
                Ok(())
            }
        }

        pub const DUMMY_DESCR: &Des = &Des {
            descr: Descr {
                codec: "dummy",
                name: "dummy",
                desc: "Dummy decoder",
                mime: "x-application/dummy",
            },
        };
    }
    use self::dummy::DUMMY_DESCR;

    #[test]
    fn lookup() {
        let codecs = Codecs::from_list(&[DUMMY_DESCR]);

        let _dec = codecs.by_name("dummy").unwrap();
    }
}
<commit_msg>codec: Improve function code<commit_after>use std::collections::HashMap;

use crate::data::frame::ArcFrame;
use crate::data::packet::Packet;

use crate::common::CodecList;
use crate::error::*;

/// Used to interact with a decoder.
pub trait Decoder: Send + Sync {
    // TODO support codec configuration using set_option
    // fn open(&mut self) -> Result<()>;
    /// Saves the extra data contained in a codec.
    fn set_extradata(&mut self, extra: &[u8]);
    /// Sends to the decoder a packet to be decoded.
    fn send_packet(&mut self, pkt: &Packet) -> Result<()>;
    /// Returns a decoded frame.
    fn receive_frame(&mut self) -> Result<ArcFrame>;
    /// Configures the decoder.
    fn configure(&mut self) -> Result<()>;
    /// Tells decoder to clear its internal state.
    fn flush(&mut self) -> Result<()>;
}

/// Codec descriptor.
///
/// Contains information on a codec and its own decoder.
#[derive(Debug)]
pub struct Descr {
    /// The codec name.
    pub codec: &'static str,
    /// The extended codec name.
    pub name: &'static str,
    /// The codec description.
    pub desc: &'static str,
    /// The codec MIME.
    pub mime: &'static str,
    // TODO more fields regarding capabilities
}

/// Auxiliary structure to encapsulate a decoder object and
/// its additional data.
pub struct Context<D: Decoder> {
    dec: D,
    // TODO: Queue up packets/frames
}

impl<D: Decoder> Context<D> {
    // TODO: More constructors
    /// Retrieves a codec descriptor from a codec list through its name,
    /// creates the relative decoder, and encapsulates it into a new `Context`.
    pub fn by_name<T: Descriptor<OutputDecoder = D> + ?Sized>(
        codecs: &Codecs<T>,
        name: &str,
    ) -> Option<Self> {
        codecs.by_name(name).map(|builder| Context {
            dec: builder.create(),
        })
    }
    /// Saves the extra data contained in a codec.
    pub fn set_extradata(&mut self, extra: &[u8]) {
        self.dec.set_extradata(extra);
    }

    /// Sends to the decoder a packet to be decoded.
    pub fn send_packet(&mut self, pkt: &Packet) -> Result<()> {
        self.dec.send_packet(pkt)
    }
    /// Returns a decoded frame.
    pub fn receive_frame(&mut self) -> Result<ArcFrame> {
        self.dec.receive_frame()
    }
    /// Configures the decoder.
    pub fn configure(&mut self) -> Result<()> {
        self.dec.configure()
    }

    /// Tells decoder to clear its internal state.
    pub fn flush(&mut self) -> Result<()> {
        self.dec.flush()
    }
}

/// Used to get the descriptor of a codec and create its own decoder.
pub trait Descriptor {
    type OutputDecoder: Decoder;

    /// Creates a new decoder for the requested codec.
    fn create(&self) -> Self::OutputDecoder;
    /// Returns the codec descriptor.
    fn describe(&self) -> &Descr;
}

/// A list of codec descriptors.
pub struct Codecs<T: 'static + Descriptor + ?Sized> {
    list: HashMap<&'static str, Vec<&'static T>>,
}

impl<T: Descriptor + ?Sized> CodecList for Codecs<T> {
    type D = T;

    fn new() -> Self {
        Codecs {
            list: HashMap::new(),
        }
    }

    // TODO more lookup functions
    fn by_name(&self, name: &str) -> Option<&'static Self::D> {
        self.list.get(name).map(|descs| descs[0])
    }

    fn append(&mut self, desc: &'static Self::D) {
        let codec_name = desc.describe().codec;

        self.list
            .entry(codec_name)
            .or_insert_with(Vec::new)
            .push(desc);
    }
}

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

    mod dummy {
        use super::super::*;
        use crate::data::pixel::Formaton;
        use std::sync::Arc;

        pub struct Dec {
            state: usize,
            #[allow(dead_code)]
            format: Option<Arc<Formaton>>,
        }

        pub struct Des {
            descr: Descr,
        }

        impl Descriptor for Des {
            type OutputDecoder = Dec;

            fn create(&self) -> Self::OutputDecoder {
                Dec {
                    state: 0,
                    format: None,
                }
            }

            fn describe(&self) -> &Descr {
                &self.descr
            }
        }

        impl Decoder for Dec {
            fn configure(&mut self) -> Result<()> {
                Ok(())
            }
            fn set_extradata(&mut self, extra: &[u8]) {
                if extra.len() > 4 {
                    self.state = 42;
                } else {
                    self.state = 12;
                }
            }
            fn send_packet(&mut self, _packet: &Packet) -> Result<()> {
                self.state += 1;
                Ok(())
            }
            fn receive_frame(&mut self) -> Result<ArcFrame> {
                unimplemented!()
            }
            fn flush(&mut self) -> Result<()> {
                Ok(())
            }
        }

        pub const DUMMY_DESCR: &Des = &Des {
            descr: Descr {
                codec: "dummy",
                name: "dummy",
                desc: "Dummy decoder",
                mime: "x-application/dummy",
            },
        };
    }
    use self::dummy::DUMMY_DESCR;

    #[test]
    fn lookup() {
        let codecs = Codecs::from_list(&[DUMMY_DESCR]);

        let _dec = codecs.by_name("dummy").unwrap();
    }
}
<|endoftext|>


<commit_before>use super::merge_policy::{MergeCandidate, MergePolicy};
use crate::core::SegmentMeta;
use std::cmp;
use std::f64;

const DEFAULT_LEVEL_LOG_SIZE: f64 = 0.75;
const DEFAULT_MIN_LAYER_SIZE: u32 = 10_000;
const DEFAULT_MIN_MERGE_SIZE: usize = 8;
const DEFAULT_MAX_MERGE_SIZE: usize = 10_000_000;

/// `LogMergePolicy` tries to merge segments that have a similar number of
/// documents.
#[derive(Debug, Clone)]
pub struct LogMergePolicy {
    min_merge_size: usize,
    max_merge_size: usize,
    min_layer_size: u32,
    level_log_size: f64,
}

impl LogMergePolicy {
    fn clip_min_size(&self, size: u32) -> u32 {
        cmp::max(self.min_layer_size, size)
    }

    /// Set the minimum number of segment that may be merge together.
    pub fn set_min_merge_size(&mut self, min_merge_size: usize) {
        self.min_merge_size = min_merge_size;
    }

    /// Set the maximum number docs in a segment for it to be considered for
    /// merging.
    pub fn set_max_merge_size(&mut self, max_merge_size: usize) {
        self.max_merge_size = max_merge_size;
    }

    /// Set the minimum segment size under which all segment belong
    /// to the same level.
    pub fn set_min_layer_size(&mut self, min_layer_size: u32) {
        self.min_layer_size = min_layer_size;
    }

    /// Set the ratio between two consecutive levels.
    ///
    /// Segment are group in levels according to their sizes.
    /// These levels are defined as intervals of exponentially growing sizes.
    /// level_log_size define the factor by which one should multiply the limit
    /// to reach a level, in order to get the limit to reach the following
    /// level.
    pub fn set_level_log_size(&mut self, level_log_size: f64) {
        self.level_log_size = level_log_size;
    }
}

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        let mut size_sorted_tuples = segments
            .iter()
            .map(SegmentMeta::num_docs)
            .enumerate()
            .filter(|(_, s)| s <= &(self.max_merge_size as u32))
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x, y| y.1.cmp(&(x.1)));

        if size_sorted_tuples.len() <= 1 {
            return Vec::new();
        }

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples
            .into_iter()
            .map(|(ind, num_docs)| (ind, f64::from(self.clip_min_size(num_docs)).log2()))
            .collect();

        if let Some(&(first_ind, first_score)) = size_sorted_log_tuples.first() {
            let mut current_max_log_size = first_score;
            let mut levels = vec![vec![first_ind]];
            for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
                if score < (current_max_log_size - self.level_log_size) {
                    current_max_log_size = score;
                    levels.push(Vec::new());
                }
                levels.last_mut().unwrap().push(ind);
            }
            levels
                .iter()
                .filter(|level| level.len() >= self.min_merge_size)
                .map(|ind_vec| {
                    MergeCandidate(ind_vec.iter().map(|&ind| segments[ind].id()).collect())
                })
                .collect()
        } else {
            return vec![];
        }
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy {
            min_merge_size: DEFAULT_MIN_MERGE_SIZE,
            max_merge_size: DEFAULT_MAX_MERGE_SIZE,
            min_layer_size: DEFAULT_MIN_LAYER_SIZE,
            level_log_size: DEFAULT_LEVEL_LOG_SIZE,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{SegmentId, SegmentMeta, SegmentMetaInventory};
    use crate::indexer::merge_policy::MergePolicy;
    use once_cell::sync::Lazy;

    static INVENTORY: Lazy<SegmentMetaInventory> = Lazy::new(SegmentMetaInventory::default);

    fn test_merge_policy() -> LogMergePolicy {
        let mut log_merge_policy = LogMergePolicy::default();
        log_merge_policy.set_min_merge_size(3);
        log_merge_policy.set_max_merge_size(100_000);
        log_merge_policy.set_min_layer_size(2);
        log_merge_policy
    }

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = test_merge_policy().compute_merge_candidates(&y);
        assert!(result_list.is_empty());
    }

    fn create_random_segment_meta(num_docs: u32) -> SegmentMeta {
        INVENTORY.new_segment_meta(SegmentId::generate_random(), num_docs)
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        // 2 MergeCandidates expected:
        // * one with the 6 * 10-docs segments
        // * one with the 3 * 1000-docs segments
        // no MergeCandidate expected for the 2 * 10_000-docs segments as min_merge_size=3
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_within_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![
            create_random_segment_meta(10),   // log2(10) = ~3.32 (> 3.58 - 0.75)
            create_random_segment_meta(11),   // log2(11) = ~3.46
            create_random_segment_meta(12),   // log2(12) = ~3.58
            create_random_segment_meta(800),  // log2(800) = ~9.64 (> 9.97 - 0.75)
            create_random_segment_meta(1000), // log2(1000) = ~9.97
            create_random_segment_meta(1000),
        ]; // log2(1000) = ~9.97
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_small_segments() {
        // segments under min_layer_size are merged together
        let test_input = vec![
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_all_segments_too_large_to_merge() {
        let eight_large_segments: Vec<SegmentMeta> =
            std::iter::repeat_with(|| create_random_segment_meta(100_001))
                .take(8)
                .collect();
        assert!(test_merge_policy()
            .compute_merge_candidates(&eight_large_segments)
            .is_empty());
    }

    #[test]
    fn test_large_merge_segments() {
        let test_input = vec![
            create_random_segment_meta(1_000_000),
            create_random_segment_meta(100_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(1_000_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
            create_random_segment_meta(1_500_000),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        // Do not include large segments
        assert_eq!(result_list.len(), 1);
        assert_eq!(result_list[0].0.len(), 3);

        // Making sure merge policy points to the correct index of the original input
        assert_eq!(result_list[0].0[0], test_input[2].id());
        assert_eq!(result_list[0].0[1], test_input[4].id());
        assert_eq!(result_list[0].0[2], test_input[5].id());
    }

}
<commit_msg>applied cargo fmt<commit_after>use super::merge_policy::{MergeCandidate, MergePolicy};
use crate::core::SegmentMeta;
use std::cmp;
use std::f64;

const DEFAULT_LEVEL_LOG_SIZE: f64 = 0.75;
const DEFAULT_MIN_LAYER_SIZE: u32 = 10_000;
const DEFAULT_MIN_MERGE_SIZE: usize = 8;
const DEFAULT_MAX_MERGE_SIZE: usize = 10_000_000;

/// `LogMergePolicy` tries to merge segments that have a similar number of
/// documents.
#[derive(Debug, Clone)]
pub struct LogMergePolicy {
    min_merge_size: usize,
    max_merge_size: usize,
    min_layer_size: u32,
    level_log_size: f64,
}

impl LogMergePolicy {
    fn clip_min_size(&self, size: u32) -> u32 {
        cmp::max(self.min_layer_size, size)
    }

    /// Set the minimum number of segment that may be merge together.
    pub fn set_min_merge_size(&mut self, min_merge_size: usize) {
        self.min_merge_size = min_merge_size;
    }

    /// Set the maximum number docs in a segment for it to be considered for
    /// merging.
    pub fn set_max_merge_size(&mut self, max_merge_size: usize) {
        self.max_merge_size = max_merge_size;
    }

    /// Set the minimum segment size under which all segment belong
    /// to the same level.
    pub fn set_min_layer_size(&mut self, min_layer_size: u32) {
        self.min_layer_size = min_layer_size;
    }

    /// Set the ratio between two consecutive levels.
    ///
    /// Segment are group in levels according to their sizes.
    /// These levels are defined as intervals of exponentially growing sizes.
    /// level_log_size define the factor by which one should multiply the limit
    /// to reach a level, in order to get the limit to reach the following
    /// level.
    pub fn set_level_log_size(&mut self, level_log_size: f64) {
        self.level_log_size = level_log_size;
    }
}

impl MergePolicy for LogMergePolicy {
    fn compute_merge_candidates(&self, segments: &[SegmentMeta]) -> Vec<MergeCandidate> {
        let mut size_sorted_tuples = segments
            .iter()
            .map(SegmentMeta::num_docs)
            .enumerate()
            .filter(|(_, s)| s <= &(self.max_merge_size as u32))
            .collect::<Vec<(usize, u32)>>();

        size_sorted_tuples.sort_by(|x, y| y.1.cmp(&(x.1)));

        if size_sorted_tuples.len() <= 1 {
            return Vec::new();
        }

        let size_sorted_log_tuples: Vec<_> = size_sorted_tuples
            .into_iter()
            .map(|(ind, num_docs)| (ind, f64::from(self.clip_min_size(num_docs)).log2()))
            .collect();

        if let Some(&(first_ind, first_score)) = size_sorted_log_tuples.first() {
            let mut current_max_log_size = first_score;
            let mut levels = vec![vec![first_ind]];
            for &(ind, score) in (&size_sorted_log_tuples).iter().skip(1) {
                if score < (current_max_log_size - self.level_log_size) {
                    current_max_log_size = score;
                    levels.push(Vec::new());
                }
                levels.last_mut().unwrap().push(ind);
            }
            levels
                .iter()
                .filter(|level| level.len() >= self.min_merge_size)
                .map(|ind_vec| {
                    MergeCandidate(ind_vec.iter().map(|&ind| segments[ind].id()).collect())
                })
                .collect()
        } else {
            return vec![];
        }
    }
}

impl Default for LogMergePolicy {
    fn default() -> LogMergePolicy {
        LogMergePolicy {
            min_merge_size: DEFAULT_MIN_MERGE_SIZE,
            max_merge_size: DEFAULT_MAX_MERGE_SIZE,
            min_layer_size: DEFAULT_MIN_LAYER_SIZE,
            level_log_size: DEFAULT_LEVEL_LOG_SIZE,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::{SegmentId, SegmentMeta, SegmentMetaInventory};
    use crate::indexer::merge_policy::MergePolicy;
    use once_cell::sync::Lazy;

    static INVENTORY: Lazy<SegmentMetaInventory> = Lazy::new(SegmentMetaInventory::default);

    fn test_merge_policy() -> LogMergePolicy {
        let mut log_merge_policy = LogMergePolicy::default();
        log_merge_policy.set_min_merge_size(3);
        log_merge_policy.set_max_merge_size(100_000);
        log_merge_policy.set_min_layer_size(2);
        log_merge_policy
    }

    #[test]
    fn test_log_merge_policy_empty() {
        let y = Vec::new();
        let result_list = test_merge_policy().compute_merge_candidates(&y);
        assert!(result_list.is_empty());
    }

    fn create_random_segment_meta(num_docs: u32) -> SegmentMeta {
        INVENTORY.new_segment_meta(SegmentId::generate_random(), num_docs)
    }

    #[test]
    fn test_log_merge_policy_pair() {
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_levels() {
        // multiple levels all get merged correctly
        // 2 MergeCandidates expected:
        // * one with the 6 * 10-docs segments
        // * one with the 3 * 1000-docs segments
        // no MergeCandidate expected for the 2 * 10_000-docs segments as min_merge_size=3
        let test_input = vec![
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(1_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10_000),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
            create_random_segment_meta(10),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_within_levels() {
        // multiple levels all get merged correctly
        let test_input = vec![
            create_random_segment_meta(10),   // log2(10) = ~3.32 (> 3.58 - 0.75)
            create_random_segment_meta(11),   // log2(11) = ~3.46
            create_random_segment_meta(12),   // log2(12) = ~3.58
            create_random_segment_meta(800),  // log2(800) = ~9.64 (> 9.97 - 0.75)
            create_random_segment_meta(1000), // log2(1000) = ~9.97
            create_random_segment_meta(1000),
        ]; // log2(1000) = ~9.97
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 2);
    }

    #[test]
    fn test_log_merge_policy_small_segments() {
        // segments under min_layer_size are merged together
        let test_input = vec![
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(1),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
            create_random_segment_meta(2),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        assert_eq!(result_list.len(), 1);
    }

    #[test]
    fn test_log_merge_policy_all_segments_too_large_to_merge() {
        let eight_large_segments: Vec<SegmentMeta> =
            std::iter::repeat_with(|| create_random_segment_meta(100_001))
                .take(8)
                .collect();
        assert!(test_merge_policy()
            .compute_merge_candidates(&eight_large_segments)
            .is_empty());
    }

    #[test]
    fn test_large_merge_segments() {
        let test_input = vec![
            create_random_segment_meta(1_000_000),
            create_random_segment_meta(100_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(1_000_001),
            create_random_segment_meta(100_000),
            create_random_segment_meta(100_000),
            create_random_segment_meta(1_500_000),
        ];
        let result_list = test_merge_policy().compute_merge_candidates(&test_input);
        // Do not include large segments
        assert_eq!(result_list.len(), 1);
        assert_eq!(result_list[0].0.len(), 3);

        // Making sure merge policy points to the correct index of the original input
        assert_eq!(result_list[0].0[0], test_input[2].id());
        assert_eq!(result_list[0].0[1], test_input[4].id());
        assert_eq!(result_list[0].0[2], test_input[5].id());
    }
}
<|endoftext|>