<commit_before>use sodiumoxide::crypto::{
    hash::sha256::DIGESTBYTES,
    sign::{PUBLICKEYBYTES, SIGNATUREBYTES},
};
use std::io::{Cursor, Error, ErrorKind, Read};

use crate::asset::Asset;
use crate::crypto::{PublicKey, ScriptHash, SigPair, Signature};

macro_rules! read_exact_bytes {
    ($self:expr, $len:expr) => {{
        let mut buf = Vec::with_capacity($len);
        unsafe {
            buf.set_len($len);
        }
        $self.read_exact(&mut buf)?;
        buf
    }};
}

pub trait BufWrite {
    fn push_u16(&mut self, num: u16);
    fn push_u32(&mut self, num: u32);
    fn push_i64(&mut self, num: i64);
    fn push_u64(&mut self, num: u64);
    fn push_bytes(&mut self, slice: &[u8]);
    fn push_pub_key(&mut self, key: &PublicKey);
    fn push_script_hash(&mut self, hash: &ScriptHash);
    fn push_sig_pair(&mut self, pair: &SigPair);
    fn push_asset(&mut self, asset: Asset);
}

impl BufWrite for Vec<u8> {
    fn push_u16(&mut self, num: u16) {
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_u32(&mut self, num: u32) {
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_i64(&mut self, num: i64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_u64(&mut self, num: u64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_bytes(&mut self, other: &[u8]) {
        if other.is_empty() {
            self.push_u32(0);
            return;
        }
        self.push_u32(other.len() as u32);
        self.extend_from_slice(other);
    }

    fn push_pub_key(&mut self, key: &PublicKey) {
        self.extend_from_slice(key.as_ref());
    }

    fn push_script_hash(&mut self, hash: &ScriptHash) {
        self.extend_from_slice(hash.as_ref());
    }

    fn push_sig_pair(&mut self, pair: &SigPair) {
        self.push_pub_key(&pair.pub_key);
        self.extend_from_slice(pair.signature.as_ref());
    }

    fn push_asset(&mut self, asset: Asset) {
        self.push_i64(asset.amount);
    }
}

pub trait BufRead {
    fn take_u8(&mut self) -> Result<u8, Error>;
    fn take_u16(&mut self) -> Result<u16, Error>;
    fn take_u32(&mut self) -> Result<u32, Error>;
    fn take_i64(&mut self) -> Result<i64, Error>;
    fn take_u64(&mut self) -> Result<u64, Error>;
    fn take_bytes(&mut self) -> Result<Vec<u8>, Error>;
    fn take_pub_key(&mut self) -> Result<PublicKey, Error>;
    fn take_script_hash(&mut self) -> Result<ScriptHash, Error>;
    fn take_sig_pair(&mut self) -> Result<SigPair, Error>;
    fn take_asset(&mut self) -> Result<Asset, Error>;
}

impl<T: AsRef<[u8]> + Read> BufRead for Cursor<T> {
    fn take_u8(&mut self) -> Result<u8, Error> {
        let mut buf = [0u8; 1];
        self.read_exact(&mut buf)?;
        Ok(buf[0])
    }

    fn take_u16(&mut self) -> Result<u16, Error> {
        let mut buf = [0u8; 2];
        self.read_exact(&mut buf)?;
        Ok((u16::from(buf[0]) << 8) | u16::from(buf[1]))
    }

    fn take_u32(&mut self) -> Result<u32, Error> {
        let mut buf = [0u8; 4];
        self.read_exact(&mut buf)?;
        Ok(u32_from_buf!(buf))
    }

    fn take_i64(&mut self) -> Result<i64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok((i64::from(buf[0]) << 56)
            | (i64::from(buf[1]) << 48)
            | (i64::from(buf[2]) << 40)
            | (i64::from(buf[3]) << 32)
            | (i64::from(buf[4]) << 24)
            | (i64::from(buf[5]) << 16)
            | (i64::from(buf[6]) << 8)
            | i64::from(buf[7]))
    }

    fn take_u64(&mut self) -> Result<u64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok(u64_from_buf!(buf))
    }

    fn take_bytes(&mut self) -> Result<Vec<u8>, Error> {
        let len = self.take_u32()? as usize;
        let buf = read_exact_bytes!(self, len);
        Ok(buf)
    }

    fn take_pub_key(&mut self) -> Result<PublicKey, Error> {
        let buf = read_exact_bytes!(self, PUBLICKEYBYTES);
        PublicKey::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect public key length"))
    }

    fn take_script_hash(&mut self) -> Result<ScriptHash, Error> {
        let buf = read_exact_bytes!(self, DIGESTBYTES);
        ScriptHash::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect script hash length"))
    }

    fn take_sig_pair(&mut self) -> Result<SigPair, Error> {
        let pub_key = self.take_pub_key()?;
        let signature = {
            let buf = read_exact_bytes!(self, SIGNATUREBYTES);
            Signature::from_slice(&buf)
                .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect signature length"))?
        };
        Ok(SigPair { pub_key, signature })
    }

    fn take_asset(&mut self) -> Result<Asset, Error> {
        let amount = self.take_i64()?;
        Ok(Asset::new(amount))
    }
}

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

    #[test]
    fn u32_serialization() {
        let num: u32 = 0x0A0B0C0D;
        let mut v = Vec::with_capacity(4);
        v.push_u32(num);
        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D]);
    }

    #[test]
    fn u64_serialization() {
        let num: u64 = 0x0A0B0C0D_0A0B0C0D;
        let mut v = Vec::with_capacity(8);
        v.push_u64(num);

        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D, 0x0A, 0x0B, 0x0C, 0x0D]);
        let dec = Cursor::<&[u8]>::new(&v).take_u64().unwrap();
        assert_eq!(num, dec);
    }

    #[test]
    fn asset_serialization() {
        let a = "12.34567 GRAEL".parse().unwrap();
        let mut v = vec![];
        v.push_asset(a);

        let mut c = Cursor::<&[u8]>::new(&v);
        let b = c.take_asset().unwrap();
        assert_eq!(a.to_string(), b.to_string());
    }
}
<commit_msg>Implement variable i64 serialization<commit_after>use sodiumoxide::crypto::{
    hash::sha256::DIGESTBYTES,
    sign::{PUBLICKEYBYTES, SIGNATUREBYTES},
};
use std::io::{Cursor, Error, ErrorKind, Read};

use crate::asset::Asset;
use crate::crypto::{PublicKey, ScriptHash, SigPair, Signature};

macro_rules! read_exact_bytes {
    ($self:expr, $len:expr) => {{
        let mut buf = Vec::with_capacity($len);
        unsafe {
            buf.set_len($len);
        }
        $self.read_exact(&mut buf)?;
        buf
    }};
}

#[inline]
fn zigzag_encode(from: i64) -> u64 {
    ((from << 1) ^ (from >> 63)) as u64
}

#[inline]
fn zigzag_decode(from: u64) -> i64 {
    ((from >> 1) ^ (-((from & 1) as i64)) as u64) as i64
}

pub trait BufWrite {
    fn push_u16(&mut self, num: u16);
    fn push_u32(&mut self, num: u32);
    fn push_i64(&mut self, num: i64);
    fn push_var_i64(&mut self, num: i64);
    fn push_u64(&mut self, num: u64);
    fn push_bytes(&mut self, slice: &[u8]);
    fn push_pub_key(&mut self, key: &PublicKey);
    fn push_script_hash(&mut self, hash: &ScriptHash);
    fn push_sig_pair(&mut self, pair: &SigPair);
    fn push_asset(&mut self, asset: Asset);
}

impl BufWrite for Vec<u8> {
    fn push_u16(&mut self, num: u16) {
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_u32(&mut self, num: u32) {
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_i64(&mut self, num: i64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_var_i64(&mut self, num: i64) {
        let mut more = true;
        let mut num = zigzag_encode(num);

        while more {
            let mut byte: u8 = (num & 0x7F) as u8;
            num >>= 7;

            if num == 0 {
                more = false;
            } else {
                byte |= 0x80;
            }

            self.push(byte);
        }
    }

    fn push_u64(&mut self, num: u64) {
        self.push((num >> 56) as u8);
        self.push((num >> 48) as u8);
        self.push((num >> 40) as u8);
        self.push((num >> 32) as u8);
        self.push((num >> 24) as u8);
        self.push((num >> 16) as u8);
        self.push((num >> 8) as u8);
        self.push(num as u8);
    }

    fn push_bytes(&mut self, other: &[u8]) {
        if other.is_empty() {
            self.push_u32(0);
            return;
        }
        self.push_u32(other.len() as u32);
        self.extend_from_slice(other);
    }

    fn push_pub_key(&mut self, key: &PublicKey) {
        self.extend_from_slice(key.as_ref());
    }

    fn push_script_hash(&mut self, hash: &ScriptHash) {
        self.extend_from_slice(hash.as_ref());
    }

    fn push_sig_pair(&mut self, pair: &SigPair) {
        self.push_pub_key(&pair.pub_key);
        self.extend_from_slice(pair.signature.as_ref());
    }

    fn push_asset(&mut self, asset: Asset) {
        self.push_i64(asset.amount);
    }
}

pub trait BufRead {
    fn take_u8(&mut self) -> Result<u8, Error>;
    fn take_u16(&mut self) -> Result<u16, Error>;
    fn take_u32(&mut self) -> Result<u32, Error>;
    fn take_i64(&mut self) -> Result<i64, Error>;
    fn take_var_i64(&mut self) -> Result<i64, Error>;
    fn take_u64(&mut self) -> Result<u64, Error>;
    fn take_bytes(&mut self) -> Result<Vec<u8>, Error>;
    fn take_pub_key(&mut self) -> Result<PublicKey, Error>;
    fn take_script_hash(&mut self) -> Result<ScriptHash, Error>;
    fn take_sig_pair(&mut self) -> Result<SigPair, Error>;
    fn take_asset(&mut self) -> Result<Asset, Error>;
}

impl<T: AsRef<[u8]> + Read> BufRead for Cursor<T> {
    fn take_u8(&mut self) -> Result<u8, Error> {
        let mut buf = [0u8; 1];
        self.read_exact(&mut buf)?;
        Ok(buf[0])
    }

    fn take_u16(&mut self) -> Result<u16, Error> {
        let mut buf = [0u8; 2];
        self.read_exact(&mut buf)?;
        Ok((u16::from(buf[0]) << 8) | u16::from(buf[1]))
    }

    fn take_u32(&mut self) -> Result<u32, Error> {
        let mut buf = [0u8; 4];
        self.read_exact(&mut buf)?;
        Ok(u32_from_buf!(buf))
    }

    fn take_i64(&mut self) -> Result<i64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok((i64::from(buf[0]) << 56)
            | (i64::from(buf[1]) << 48)
            | (i64::from(buf[2]) << 40)
            | (i64::from(buf[3]) << 32)
            | (i64::from(buf[4]) << 24)
            | (i64::from(buf[5]) << 16)
            | (i64::from(buf[6]) << 8)
            | i64::from(buf[7]))
    }

    fn take_var_i64(&mut self) -> Result<i64, Error> {
        let mut result: u64 = 0;
        let mut shift = 0;
        let mut buf = [0u8; 1];
        loop {
            if shift > 63 {
                return Err(Error::new(ErrorKind::Other, "overflow taking varint"));
            }

            self.read_exact(&mut buf)?;
            let byte = buf[0];

            result |= ((byte & 0x7F) as u64) << shift;
            if byte & 0x80 == 0 {
                break;
            }

            shift += 7;
        }

        Ok(zigzag_decode(result))
    }

    fn take_u64(&mut self) -> Result<u64, Error> {
        let mut buf = [0u8; 8];
        self.read_exact(&mut buf)?;
        Ok(u64_from_buf!(buf))
    }

    fn take_bytes(&mut self) -> Result<Vec<u8>, Error> {
        let len = self.take_u32()? as usize;
        let buf = read_exact_bytes!(self, len);
        Ok(buf)
    }

    fn take_pub_key(&mut self) -> Result<PublicKey, Error> {
        let buf = read_exact_bytes!(self, PUBLICKEYBYTES);
        PublicKey::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect public key length"))
    }

    fn take_script_hash(&mut self) -> Result<ScriptHash, Error> {
        let buf = read_exact_bytes!(self, DIGESTBYTES);
        ScriptHash::from_slice(&buf)
            .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect script hash length"))
    }

    fn take_sig_pair(&mut self) -> Result<SigPair, Error> {
        let pub_key = self.take_pub_key()?;
        let signature = {
            let buf = read_exact_bytes!(self, SIGNATUREBYTES);
            Signature::from_slice(&buf)
                .ok_or_else(|| Error::new(ErrorKind::Other, "incorrect signature length"))?
        };
        Ok(SigPair { pub_key, signature })
    }

    fn take_asset(&mut self) -> Result<Asset, Error> {
        let amount = self.take_i64()?;
        Ok(Asset::new(amount))
    }
}

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

    #[test]
    fn u32_serialization() {
        let num: u32 = 0x0A0B0C0D;
        let mut v = Vec::with_capacity(4);
        v.push_u32(num);
        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D]);
    }

    #[test]
    fn u64_serialization() {
        let num: u64 = 0x0A0B0C0D_0A0B0C0D;
        let mut v = Vec::with_capacity(8);
        v.push_u64(num);

        assert_eq!(v, [0x0A, 0x0B, 0x0C, 0x0D, 0x0A, 0x0B, 0x0C, 0x0D]);
        let dec = Cursor::<&[u8]>::new(&v).take_u64().unwrap();
        assert_eq!(num, dec);
    }

    #[test]
    fn asset_serialization() {
        let a = "12.34567 GRAEL".parse().unwrap();
        let mut v = vec![];
        v.push_asset(a);

        let mut c = Cursor::<&[u8]>::new(&v);
        let b = c.take_asset().unwrap();
        assert_eq!(a.to_string(), b.to_string());
    }

    #[test]
    fn zigzag() {
        fn cmp(decoded: i64, encoded: u64) {
            assert_eq!(decoded, zigzag_decode(encoded));
            assert_eq!(encoded, zigzag_encode(decoded));
        }

        cmp(0, 0);
        cmp(-1, 1);
        cmp(1, 2);
        cmp(-2, 3);
        cmp(2147483647, 4294967294);
        cmp(-2147483648, 4294967295);
        cmp(9223372036854775807, 18446744073709551614);
        cmp(-9223372036854775808, 18446744073709551615);
    }

    #[test]
    fn var_i64_serialization() {
        let mut buf = vec![];
        buf.push_var_i64(0);
        buf.push_var_i64(300);
        buf.push_var_i64(-300);
        buf.push_var_i64(i64::max_value());
        buf.push_var_i64(i64::min_value());
        buf.extend(vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01,
        ]);

        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(c.take_var_i64().unwrap(), 0);
        assert_eq!(c.take_var_i64().unwrap(), 300);
        assert_eq!(c.take_var_i64().unwrap(), -300);
        assert_eq!(c.take_var_i64().unwrap(), i64::max_value());
        assert_eq!(c.take_var_i64().unwrap(), i64::min_value());
        assert_eq!(c.take_var_i64().unwrap(), 1 << 62);
    }

    #[test]
    fn var_i64_serialization_overflow() {
        use std::error;

        let buf = vec![
            0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0,
        ];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            error::Error::description(&c.take_var_i64().unwrap_err()),
            "overflow taking varint"
        );
    }

    #[test]
    fn var_i64_serialization_eof() {
        let buf = vec![0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80];
        let mut c = Cursor::<&[u8]>::new(&buf);
        assert_eq!(
            c.take_var_i64().unwrap_err().kind(),
            ErrorKind::UnexpectedEof
        );
    }
}
<|endoftext|>


<commit_before>use super::*;

pub(super) const PATTERN_FIRST: TokenSet = expressions::LITERAL_FIRST
    .union(paths::PATH_FIRST)
    .union(token_set![BOX_KW, REF_KW, MUT_KW, L_PAREN, L_BRACK, AMP, UNDERSCORE, MINUS, DOT]);

pub(crate) fn pattern(p: &mut Parser) {
    pattern_r(p, PAT_RECOVERY_SET);
}

/// Parses a pattern list separated by pipes `|`
pub(super) fn pattern_list(p: &mut Parser) {
    pattern_list_r(p, PAT_RECOVERY_SET)
}

/// Parses a pattern list separated by pipes `|`
/// using the given `recovery_set`
pub(super) fn pattern_list_r(p: &mut Parser, recovery_set: TokenSet) {
    p.eat(T![|]);
    pattern_r(p, recovery_set);

    while p.eat(T![|]) {
        pattern_r(p, recovery_set);
    }
}

pub(super) fn pattern_r(p: &mut Parser, recovery_set: TokenSet) {
    if let Some(lhs) = atom_pat(p, recovery_set) {
        // test range_pat
        // fn main() {
        //     match 92 {
        //         0 ... 100 => (),
        //         101 ..= 200 => (),
        //         200 .. 301=> (),
        //     }
        // }
        for &range_op in [T![...], T![..=], T![..]].iter() {
            if p.at(range_op) {
                let m = lhs.precede(p);
                p.bump(range_op);
                atom_pat(p, recovery_set);
                m.complete(p, RANGE_PAT);
                return;
            }
        }
        // test marco_pat
        // fn main() {
        //     let m!(x) = 0;
        // }
        if lhs.kind() == PATH_PAT && p.at(T![!]) {
            let m = lhs.precede(p);
            items::macro_call_after_excl(p);
            m.complete(p, MACRO_CALL);
        }
    }
}

const PAT_RECOVERY_SET: TokenSet =
    token_set![LET_KW, IF_KW, WHILE_KW, LOOP_KW, MATCH_KW, R_PAREN, COMMA];

fn atom_pat(p: &mut Parser, recovery_set: TokenSet) -> Option<CompletedMarker> {
    let m = match p.nth(0) {
        T![box] => box_pat(p),
        T![ref] | T![mut] => bind_pat(p, true),
        IDENT => match p.nth(1) {
            // Checks the token after an IDENT to see if a pattern is a path (Struct { .. }) or macro
            // (T![x]).
            T!['('] | T!['{'] | T![!] => path_pat(p),
            T![:] if p.nth_at(1, T![::]) => path_pat(p),
            _ => bind_pat(p, true),
        },

        _ if paths::is_use_path_start(p) => path_pat(p),
        _ if is_literal_pat_start(p) => literal_pat(p),

        T![.] if p.at(T![..]) => dot_dot_pat(p),
        T![_] => placeholder_pat(p),
        T![&] => ref_pat(p),
        T!['('] => tuple_pat(p),
        T!['['] => slice_pat(p),

        _ => {
            p.err_recover("expected pattern", recovery_set);
            return None;
        }
    };

    Some(m)
}

fn is_literal_pat_start(p: &Parser) -> bool {
    p.at(T![-]) && (p.nth(1) == INT_NUMBER || p.nth(1) == FLOAT_NUMBER)
        || p.at_ts(expressions::LITERAL_FIRST)
}

// test literal_pattern
// fn main() {
//     match () {
//         -1 => (),
//         92 => (),
//         'c' => (),
//         "hello" => (),
//     }
// }
fn literal_pat(p: &mut Parser) -> CompletedMarker {
    assert!(is_literal_pat_start(p));
    let m = p.start();
    if p.at(T![-]) {
        p.bump(T![-]);
    }
    expressions::literal(p);
    m.complete(p, LITERAL_PAT)
}

// test path_part
// fn foo() {
//     let foo::Bar = ();
//     let ::Bar = ();
//     let Bar { .. } = ();
//     let Bar(..) = ();
// }
fn path_pat(p: &mut Parser) -> CompletedMarker {
    assert!(paths::is_use_path_start(p));
    let m = p.start();
    paths::expr_path(p);
    let kind = match p.current() {
        T!['('] => {
            tuple_pat_fields(p);
            TUPLE_STRUCT_PAT
        }
        T!['{'] => {
            record_field_pat_list(p);
            RECORD_PAT
        }
        _ => PATH_PAT,
    };
    m.complete(p, kind)
}

// test tuple_pat_fields
// fn foo() {
//     let S() = ();
//     let S(_) = ();
//     let S(_,) = ();
//     let S(_, .. , x) = ();
// }
fn tuple_pat_fields(p: &mut Parser) {
    assert!(p.at(T!['(']));
    p.bump(T!['(']);
    pat_list(p, T![')']);
    p.expect(T![')']);
}

// test record_field_pat_list
// fn foo() {
//     let S {} = ();
//     let S { f, ref mut g } = ();
//     let S { h: _, ..} = ();
//     let S { h: _, } = ();
// }
fn record_field_pat_list(p: &mut Parser) {
    assert!(p.at(T!['{']));
    let m = p.start();
    p.bump(T!['{']);
    while !p.at(EOF) && !p.at(T!['}']) {
        match p.current() {
            // A trailing `..` is *not* treated as a DOT_DOT_PAT.
            T![.] if p.at(T![..]) => p.bump(T![..]),

            IDENT if p.nth(1) == T![:] => record_field_pat(p),
            INT_NUMBER if p.nth(1) == T![:] => record_field_pat(p),
            T!['{'] => error_block(p, "expected ident"),
            T![box] => {
                box_pat(p);
            }
            _ => {
                bind_pat(p, false);
            }
        }
        if !p.at(T!['}']) {
            p.expect(T![,]);
        }
    }
    p.expect(T!['}']);
    m.complete(p, RECORD_FIELD_PAT_LIST);
}

// test record_field_pat
// fn foo() {
//     let S { 0: 1 } = ();
//     let S { x: 1 } = ();
// }
fn record_field_pat(p: &mut Parser) {
    assert!(p.at(IDENT) || p.at(INT_NUMBER));
    assert!(p.nth(1) == T![:]);

    let m = p.start();

    match p.current() {
        IDENT => name(p),
        _ => p.bump_any(),
    }

    p.bump_any();
    pattern(p);
    m.complete(p, RECORD_FIELD_PAT);
}

// test placeholder_pat
// fn main() { let _ = (); }
fn placeholder_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![_]));
    let m = p.start();
    p.bump(T![_]);
    m.complete(p, PLACEHOLDER_PAT)
}

// test dot_dot_pat
// fn main() {
//     let .. = ();
//     //
//     // Tuples
//     //
//     let (a, ..) = ();
//     let (a, ..,) = ();
//     let Tuple(a, ..) = ();
//     let Tuple(a, ..,) = ();
//     let (.., ..) = ();
//     let Tuple(.., ..) = ();
//     let (.., a, ..) = ();
//     let Tuple(.., a, ..) = ();
//     //
//     // Slices
//     //
//     let [..] = ();
//     let [head, ..] = ();
//     let [head, tail @ ..] = ();
//     let [head, .., cons] = ();
//     let [head, mid @ .., cons] = ();
//     let [head, .., .., cons] = ();
//     let [head, .., mid, tail @ ..] = ();
//     let [head, .., mid, .., cons] = ();
// }
fn dot_dot_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![..]));
    let m = p.start();
    p.bump(T![..]);
    m.complete(p, DOT_DOT_PAT)
}

// test ref_pat
// fn main() {
//     let &a = ();
//     let &mut b = ();
// }
fn ref_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![&]));
    let m = p.start();
    p.bump(T![&]);
    p.eat(T![mut]);
    pattern(p);
    m.complete(p, REF_PAT)
}

// test tuple_pat
// fn main() {
//     let (a, b, ..) = ();
// }
fn tuple_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T!['(']));
    let m = p.start();
    tuple_pat_fields(p);
    m.complete(p, TUPLE_PAT)
}

// test slice_pat
// fn main() {
//     let [a, b, ..] = [];
// }
fn slice_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T!['[']));
    let m = p.start();
    p.bump(T!['[']);
    pat_list(p, T![']']);
    p.expect(T![']']);
    m.complete(p, SLICE_PAT)
}

fn pat_list(p: &mut Parser, ket: SyntaxKind) {
    while !p.at(EOF) && !p.at(ket) {
        if !p.at_ts(PATTERN_FIRST) {
            p.error("expected a pattern");
            break;
        }

        pattern(p);
        if !p.at(ket) {
            p.expect(T![,]);
        }
    }
}

// test bind_pat
// fn main() {
//     let a = ();
//     let mut b = ();
//     let ref c = ();
//     let ref mut d = ();
//     let e @ _ = ();
//     let ref mut f @ g @ _ = ();
// }
fn bind_pat(p: &mut Parser, with_at: bool) -> CompletedMarker {
    let m = p.start();
    p.eat(T![ref]);
    p.eat(T![mut]);
    name(p);
    if with_at && p.eat(T![@]) {
        pattern(p);
    }
    m.complete(p, BIND_PAT)
}

// test box_pat
// fn main() {
//     let box i = ();
//     let box Outer { box i, j: box Inner(box &x) } = ();
//     let box ref mut i = ();
// }
fn box_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![box]));
    let m = p.start();
    p.bump(T![box]);
    pattern(p);
    m.complete(p, BOX_PAT)
}
<commit_msg>Apply suggestion<commit_after>use super::*;

pub(super) const PATTERN_FIRST: TokenSet = expressions::LITERAL_FIRST
    .union(paths::PATH_FIRST)
    .union(token_set![BOX_KW, REF_KW, MUT_KW, L_PAREN, L_BRACK, AMP, UNDERSCORE, MINUS, DOT]);

pub(crate) fn pattern(p: &mut Parser) {
    pattern_r(p, PAT_RECOVERY_SET);
}

/// Parses a pattern list separated by pipes `|`
pub(super) fn pattern_list(p: &mut Parser) {
    pattern_list_r(p, PAT_RECOVERY_SET)
}

/// Parses a pattern list separated by pipes `|`
/// using the given `recovery_set`
pub(super) fn pattern_list_r(p: &mut Parser, recovery_set: TokenSet) {
    p.eat(T![|]);
    pattern_r(p, recovery_set);

    while p.eat(T![|]) {
        pattern_r(p, recovery_set);
    }
}

pub(super) fn pattern_r(p: &mut Parser, recovery_set: TokenSet) {
    if let Some(lhs) = atom_pat(p, recovery_set) {
        // test range_pat
        // fn main() {
        //     match 92 {
        //         0 ... 100 => (),
        //         101 ..= 200 => (),
        //         200 .. 301=> (),
        //     }
        // }
        for &range_op in [T![...], T![..=], T![..]].iter() {
            if p.at(range_op) {
                let m = lhs.precede(p);
                p.bump(range_op);
                atom_pat(p, recovery_set);
                m.complete(p, RANGE_PAT);
                return;
            }
        }
        // test marco_pat
        // fn main() {
        //     let m!(x) = 0;
        // }
        if lhs.kind() == PATH_PAT && p.at(T![!]) {
            let m = lhs.precede(p);
            items::macro_call_after_excl(p);
            m.complete(p, MACRO_CALL);
        }
    }
}

const PAT_RECOVERY_SET: TokenSet =
    token_set![LET_KW, IF_KW, WHILE_KW, LOOP_KW, MATCH_KW, R_PAREN, COMMA];

fn atom_pat(p: &mut Parser, recovery_set: TokenSet) -> Option<CompletedMarker> {
    let m = match p.nth(0) {
        T![box] => box_pat(p),
        T![ref] | T![mut] => bind_pat(p, true),
        IDENT => match p.nth(1) {
            // Checks the token after an IDENT to see if a pattern is a path (Struct { .. }) or macro
            // (T![x]).
            T!['('] | T!['{'] | T![!] => path_pat(p),
            T![:] if p.nth_at(1, T![::]) => path_pat(p),
            _ => bind_pat(p, true),
        },

        _ if paths::is_use_path_start(p) => path_pat(p),
        _ if is_literal_pat_start(p) => literal_pat(p),

        T![.] if p.at(T![..]) => dot_dot_pat(p),
        T![_] => placeholder_pat(p),
        T![&] => ref_pat(p),
        T!['('] => tuple_pat(p),
        T!['['] => slice_pat(p),

        _ => {
            p.err_recover("expected pattern", recovery_set);
            return None;
        }
    };

    Some(m)
}

fn is_literal_pat_start(p: &Parser) -> bool {
    p.at(T![-]) && (p.nth(1) == INT_NUMBER || p.nth(1) == FLOAT_NUMBER)
        || p.at_ts(expressions::LITERAL_FIRST)
}

// test literal_pattern
// fn main() {
//     match () {
//         -1 => (),
//         92 => (),
//         'c' => (),
//         "hello" => (),
//     }
// }
fn literal_pat(p: &mut Parser) -> CompletedMarker {
    assert!(is_literal_pat_start(p));
    let m = p.start();
    if p.at(T![-]) {
        p.bump(T![-]);
    }
    expressions::literal(p);
    m.complete(p, LITERAL_PAT)
}

// test path_part
// fn foo() {
//     let foo::Bar = ();
//     let ::Bar = ();
//     let Bar { .. } = ();
//     let Bar(..) = ();
// }
fn path_pat(p: &mut Parser) -> CompletedMarker {
    assert!(paths::is_use_path_start(p));
    let m = p.start();
    paths::expr_path(p);
    let kind = match p.current() {
        T!['('] => {
            tuple_pat_fields(p);
            TUPLE_STRUCT_PAT
        }
        T!['{'] => {
            record_field_pat_list(p);
            RECORD_PAT
        }
        _ => PATH_PAT,
    };
    m.complete(p, kind)
}

// test tuple_pat_fields
// fn foo() {
//     let S() = ();
//     let S(_) = ();
//     let S(_,) = ();
//     let S(_, .. , x) = ();
// }
fn tuple_pat_fields(p: &mut Parser) {
    assert!(p.at(T!['(']));
    p.bump(T!['(']);
    pat_list(p, T![')']);
    p.expect(T![')']);
}

// test record_field_pat_list
// fn foo() {
//     let S {} = ();
//     let S { f, ref mut g } = ();
//     let S { h: _, ..} = ();
//     let S { h: _, } = ();
// }
fn record_field_pat_list(p: &mut Parser) {
    assert!(p.at(T!['{']));
    let m = p.start();
    p.bump(T!['{']);
    while !p.at(EOF) && !p.at(T!['}']) {
        match p.current() {
            // A trailing `..` is *not* treated as a DOT_DOT_PAT.
            T![.] if p.at(T![..]) => p.bump(T![..]),

            IDENT if p.nth(1) == T![:] => record_field_pat(p),
            INT_NUMBER if p.nth(1) == T![:] => record_field_pat(p),
            T!['{'] => error_block(p, "expected ident"),
            T![box] => {
                box_pat(p);
            }
            _ => {
                bind_pat(p, false);
            }
        }
        if !p.at(T!['}']) {
            p.expect(T![,]);
        }
    }
    p.expect(T!['}']);
    m.complete(p, RECORD_FIELD_PAT_LIST);
}

// test record_field_pat
// fn foo() {
//     let S { 0: 1 } = ();
//     let S { x: 1 } = ();
// }
fn record_field_pat(p: &mut Parser) {
    assert!(p.at(IDENT) || p.at(INT_NUMBER));
    assert!(p.nth(1) == T![:]);

    let m = p.start();

    if !p.eat(INT_NUMBER) {
        name(p)
    }

    p.bump_any();
    pattern(p);
    m.complete(p, RECORD_FIELD_PAT);
}

// test placeholder_pat
// fn main() { let _ = (); }
fn placeholder_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![_]));
    let m = p.start();
    p.bump(T![_]);
    m.complete(p, PLACEHOLDER_PAT)
}

// test dot_dot_pat
// fn main() {
//     let .. = ();
//     //
//     // Tuples
//     //
//     let (a, ..) = ();
//     let (a, ..,) = ();
//     let Tuple(a, ..) = ();
//     let Tuple(a, ..,) = ();
//     let (.., ..) = ();
//     let Tuple(.., ..) = ();
//     let (.., a, ..) = ();
//     let Tuple(.., a, ..) = ();
//     //
//     // Slices
//     //
//     let [..] = ();
//     let [head, ..] = ();
//     let [head, tail @ ..] = ();
//     let [head, .., cons] = ();
//     let [head, mid @ .., cons] = ();
//     let [head, .., .., cons] = ();
//     let [head, .., mid, tail @ ..] = ();
//     let [head, .., mid, .., cons] = ();
// }
fn dot_dot_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![..]));
    let m = p.start();
    p.bump(T![..]);
    m.complete(p, DOT_DOT_PAT)
}

// test ref_pat
// fn main() {
//     let &a = ();
//     let &mut b = ();
// }
fn ref_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![&]));
    let m = p.start();
    p.bump(T![&]);
    p.eat(T![mut]);
    pattern(p);
    m.complete(p, REF_PAT)
}

// test tuple_pat
// fn main() {
//     let (a, b, ..) = ();
// }
fn tuple_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T!['(']));
    let m = p.start();
    tuple_pat_fields(p);
    m.complete(p, TUPLE_PAT)
}

// test slice_pat
// fn main() {
//     let [a, b, ..] = [];
// }
fn slice_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T!['[']));
    let m = p.start();
    p.bump(T!['[']);
    pat_list(p, T![']']);
    p.expect(T![']']);
    m.complete(p, SLICE_PAT)
}

fn pat_list(p: &mut Parser, ket: SyntaxKind) {
    while !p.at(EOF) && !p.at(ket) {
        if !p.at_ts(PATTERN_FIRST) {
            p.error("expected a pattern");
            break;
        }

        pattern(p);
        if !p.at(ket) {
            p.expect(T![,]);
        }
    }
}

// test bind_pat
// fn main() {
//     let a = ();
//     let mut b = ();
//     let ref c = ();
//     let ref mut d = ();
//     let e @ _ = ();
//     let ref mut f @ g @ _ = ();
// }
fn bind_pat(p: &mut Parser, with_at: bool) -> CompletedMarker {
    let m = p.start();
    p.eat(T![ref]);
    p.eat(T![mut]);
    name(p);
    if with_at && p.eat(T![@]) {
        pattern(p);
    }
    m.complete(p, BIND_PAT)
}

// test box_pat
// fn main() {
//     let box i = ();
//     let box Outer { box i, j: box Inner(box &x) } = ();
//     let box ref mut i = ();
// }
fn box_pat(p: &mut Parser) -> CompletedMarker {
    assert!(p.at(T![box]));
    let m = p.start();
    p.bump(T![box]);
    pattern(p);
    m.complete(p, BOX_PAT)
}
<|endoftext|>


<commit_before>extern crate futures;
extern crate uuid;

use futures::Stream;
use uuid::Uuid;

/// The sequence number within a single source of events
pub type SequenceNumber = u32;

/// An event with associated metadata that corresponds to how it was stored in
/// the event store.
#[derive(Debug, Clone, PartialEq)]
pub struct PersistedEvent<Offset, Event> {
    /// The global offset of the event in the event store
    pub offset: Offset,
    /// The unique identifier of a single source of events in the event store.
    /// This usually corresponds to an 'aggregate id' in DDD.
    pub source_id: Uuid,
    /// The sequence number within a single source of events
    pub sequence_number: SequenceNumber,
    /// The event payload that was stored by the client of the event store
    pub payload: Event,
}

impl<Offset, Event> PersistedEvent<Offset, Event> {
    /// Take the event data by reference, cloning the source id and offsets
    pub fn as_ref(&self) -> PersistedEvent<Offset, &Event>
        where Offset: Clone
    {
        PersistedEvent {
            offset: self.offset.clone(),
            source_id: self.source_id,
            sequence_number: self.sequence_number.clone(),
            payload: &self.payload,
        }
    }

    /// Apply a transformation to the event
    pub fn map<NewEvent, F>(self, f: F) -> PersistedEvent<Offset, NewEvent>
        where F: FnOnce(Event) -> NewEvent
    {
        PersistedEvent {
            offset: self.offset,
            source_id: self.source_id,
            sequence_number: self.sequence_number,
            payload: f(self.payload),
        }
    }
}

/// A repository for storing historical event logs.
pub trait EventStore {
    /// The type of the offests into the event store. Relational DBs
    /// will probably use a monotonically increasing sequence number,
    /// where as distributed data sources may use timestamps.
    type Offset;

    /// The event store's native serialized form
    type Event;

    /// A lazily loaded source of events
    type EventsStream: Stream<Item = PersistedEvent<Self::Offset, Self::Event>>;

    /// Append the events to the event store for the specified source id
    fn append_events(&self, source_id: Uuid, events: Vec<Self::Event>);

    /// Stream the events back from the event store for the specified source id
    fn events(&self, source_id: Uuid, offset: Self::Offset) -> Self::EventsStream;
}

#[cfg(test)]
mod tests {
    use uuid::Uuid;

    use super::*;

    #[test]
    fn persisted_event_map() {
        let event = PersistedEvent {
            offset: 123,
            source_id: Uuid::new_v4(),
            sequence_number: 354,
            payload: "hello",
        };

        let new_event = event.clone().map(String::from);

        assert_eq!(new_event.offset, event.offset);
        assert_eq!(new_event.source_id, event.source_id);
        assert_eq!(new_event.sequence_number, event.sequence_number);
        assert_eq!(new_event.payload, event.payload);
    }

    #[test]
    fn persisted_event_as_ref() {
        let event = PersistedEvent {
            offset: 123,
            source_id: Uuid::new_v4(),
            sequence_number: 354,
            payload: "hello",
        };

        let new_event = event.as_ref();

        assert_eq!(new_event.offset, event.offset);
        assert_eq!(new_event.source_id, event.source_id);
        assert_eq!(new_event.sequence_number, event.sequence_number);
        assert_eq!(*new_event.payload, event.payload);
    }
}
<commit_msg>Add PartialOrd constraint on EventStore::Offset<commit_after>extern crate futures;
extern crate uuid;

use futures::Stream;
use uuid::Uuid;

/// The sequence number within a single source of events
pub type SequenceNumber = u32;

/// An event with associated metadata that corresponds to how it was stored in
/// the event store.
#[derive(Debug, Clone, PartialEq)]
pub struct PersistedEvent<Offset, Event> {
    /// The global offset of the event in the event store
    pub offset: Offset,
    /// The unique identifier of a single source of events in the event store.
    /// This usually corresponds to an 'aggregate id' in DDD.
    pub source_id: Uuid,
    /// The sequence number within a single source of events
    pub sequence_number: SequenceNumber,
    /// The event payload that was stored by the client of the event store
    pub payload: Event,
}

impl<Offset, Event> PersistedEvent<Offset, Event> {
    /// Take the event data by reference, cloning the source id and offsets
    pub fn as_ref(&self) -> PersistedEvent<Offset, &Event>
        where Offset: Clone
    {
        PersistedEvent {
            offset: self.offset.clone(),
            source_id: self.source_id,
            sequence_number: self.sequence_number.clone(),
            payload: &self.payload,
        }
    }

    /// Apply a transformation to the event
    pub fn map<NewEvent, F>(self, f: F) -> PersistedEvent<Offset, NewEvent>
        where F: FnOnce(Event) -> NewEvent
    {
        PersistedEvent {
            offset: self.offset,
            source_id: self.source_id,
            sequence_number: self.sequence_number,
            payload: f(self.payload),
        }
    }
}

/// A repository for storing historical event logs.
pub trait EventStore {
    /// The type of the offests into the event store. Relational DBs
    /// will probably use a monotonically increasing sequence number,
    /// where as distributed data sources may use timestamps.
    type Offset: PartialOrd;

    /// The event store's native serialized form
    type Event;

    /// A lazily loaded source of events
    type EventsStream: Stream<Item = PersistedEvent<Self::Offset, Self::Event>>;

    /// Append the events to the event store for the specified source id
    fn append_events(&self, source_id: Uuid, events: Vec<Self::Event>);

    /// Stream the events back from the event store for the specified source id
    fn events(&self, source_id: Uuid, offset: Self::Offset) -> Self::EventsStream;
}

#[cfg(test)]
mod tests {
    use uuid::Uuid;

    use super::*;

    #[test]
    fn persisted_event_map() {
        let event = PersistedEvent {
            offset: 123,
            source_id: Uuid::new_v4(),
            sequence_number: 354,
            payload: "hello",
        };

        let new_event = event.clone().map(String::from);

        assert_eq!(new_event.offset, event.offset);
        assert_eq!(new_event.source_id, event.source_id);
        assert_eq!(new_event.sequence_number, event.sequence_number);
        assert_eq!(new_event.payload, event.payload);
    }

    #[test]
    fn persisted_event_as_ref() {
        let event = PersistedEvent {
            offset: 123,
            source_id: Uuid::new_v4(),
            sequence_number: 354,
            payload: "hello",
        };

        let new_event = event.as_ref();

        assert_eq!(new_event.offset, event.offset);
        assert_eq!(new_event.source_id, event.source_id);
        assert_eq!(new_event.sequence_number, event.sequence_number);
        assert_eq!(*new_event.payload, event.payload);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counters::{
        PROCESSED_STRUCT_LOG_COUNT, SENT_STRUCT_LOG_COUNT, STRUCT_LOG_PARSE_ERROR_COUNT,
        STRUCT_LOG_QUEUE_ERROR_COUNT, STRUCT_LOG_SEND_ERROR_COUNT,
    },
    logger::Logger,
    struct_log::TcpWriter,
    Event, Filter, Level, Metadata,
};
use chrono::{SecondsFormat, Utc};
use serde::Serialize;
use std::{
    collections::HashMap,
    env, fmt,
    io::Write,
    sync::{
        mpsc::{self, Receiver, SyncSender},
        Arc, RwLock,
    },
    thread,
};

const RUST_LOG: &str = "RUST_LOG";
pub const CHANNEL_SIZE: usize = 10000;
const NUM_SEND_RETRIES: u8 = 1;

#[derive(Debug, Serialize)]
pub struct LogEntry {
    #[serde(flatten)]
    metadata: Metadata,
    timestamp: String,
    #[serde(skip_serializing_if = "HashMap::is_empty")]
    data: HashMap<&'static str, serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    message: Option<String>,
}

impl LogEntry {
    fn new(event: &Event) -> Self {
        use crate::{Key, Value, Visitor};

        struct JsonVisitor<'a>(&'a mut HashMap<&'static str, serde_json::Value>);

        impl<'a> Visitor for JsonVisitor<'a> {
            fn visit_pair(&mut self, key: Key, value: Value<'_>) {
                let v = match value {
                    Value::Debug(d) => serde_json::Value::String(format!("{:?}", d)),
                    Value::Display(d) => serde_json::Value::String(d.to_string()),
                    Value::Serde(s) => serde_json::to_value(s).unwrap(),
                };

                self.0.insert(key.as_str(), v);
            }
        }

        let metadata = *event.metadata();
        let message = event.message().map(fmt::format);

        let mut data = HashMap::new();
        for schema in event.keys_and_values() {
            schema.visit(&mut JsonVisitor(&mut data));
        }

        Self {
            metadata,
            timestamp: Utc::now().to_rfc3339_opts(SecondsFormat::Micros, true),
            data,
            message,
        }
    }
}

pub struct LibraLoggerBuilder {
    channel_size: usize,
    level: Level,
    address: Option<String>,
    printer: Option<Box<dyn Writer>>,
    is_async: bool,
}

impl LibraLoggerBuilder {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            channel_size: CHANNEL_SIZE,
            level: Level::Debug,
            address: None,
            printer: Some(Box::new(StderrWriter)),
            is_async: false,
        }
    }

    pub fn address(&mut self, address: String) -> &mut Self {
        self.address = Some(address);
        self
    }

    pub fn read_env(&mut self) -> &mut Self {
        if let Ok(address) = env::var("STRUCT_LOG_TCP_ADDR") {
            self.address(address);
        }
        self
    }

    pub fn level(&mut self, level: Level) -> &mut Self {
        self.level = level;
        self
    }

    pub fn channel_size(&mut self, channel_size: usize) -> &mut Self {
        self.channel_size = channel_size;
        self
    }

    pub fn printer(&mut self, printer: Box<dyn Writer + Send + Sync + 'static>) -> &mut Self {
        self.printer = Some(printer);
        self
    }

    pub fn is_async(&mut self, is_async: bool) -> &mut Self {
        self.is_async = is_async;
        self
    }

    pub fn init(&mut self) {
        self.build();
    }

    pub fn build(&mut self) -> Arc<LibraLogger> {
        let filter = {
            let mut filter_builder = Filter::builder();

            if env::var(RUST_LOG).is_ok() {
                filter_builder.with_env(RUST_LOG);
            } else {
                filter_builder.filter_level(self.level.into());
            }

            filter_builder.build()
        };

        let logger = if self.is_async {
            let (sender, receiver) = mpsc::sync_channel(self.channel_size);
            let service = LoggerService {
                receiver,
                address: self.address.clone(),
                printer: self.printer.take(),
            };
            let logger = Arc::new(LibraLogger {
                sender: Some(sender),
                printer: None,
                filter: RwLock::new(filter),
            });

            thread::spawn(move || service.run());
            logger
        } else {
            Arc::new(LibraLogger {
                sender: None,
                printer: self.printer.take(),
                filter: RwLock::new(filter),
            })
        };

        crate::logger::set_global_logger(logger.clone());
        logger
    }
}

pub struct LibraLogger {
    sender: Option<SyncSender<LogEntry>>,
    printer: Option<Box<dyn Writer>>,
    filter: RwLock<Filter>,
}

impl LibraLogger {
    pub fn builder() -> LibraLoggerBuilder {
        LibraLoggerBuilder::new()
    }

    #[allow(clippy::new_ret_no_self)]
    pub fn new() -> LibraLoggerBuilder {
        Self::builder()
    }

    pub fn init_for_testing() {
        if env::var(RUST_LOG).is_err() {
            return;
        }

        Self::builder()
            .is_async(false)
            .printer(Box::new(StderrWriter))
            .build();
    }

    pub fn set_filter(&self, filter: Filter) {
        *self.filter.write().unwrap() = filter;
    }

    fn send_entry(&self, entry: LogEntry) {
        if let Some(printer) = &self.printer {
            let s = format(&entry).expect("Unable to format");
            printer.write(s);
        }

        if let Some(sender) = &self.sender {
            if let Err(e) = sender.try_send(entry) {
                STRUCT_LOG_QUEUE_ERROR_COUNT.inc();
                eprintln!("Failed to send structured log: {}", e);
            }
        }
    }
}

impl Logger for LibraLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        self.filter.read().unwrap().enabled(metadata)
    }

    fn record(&self, event: &Event) {
        let entry = LogEntry::new(event);

        self.send_entry(entry)
    }
}

struct LoggerService {
    receiver: Receiver<LogEntry>,
    address: Option<String>,
    printer: Option<Box<dyn Writer>>,
}

impl LoggerService {
    pub fn run(mut self) {
        let mut writer = self.address.take().map(TcpWriter::new);

        for entry in self.receiver {
            PROCESSED_STRUCT_LOG_COUNT.inc();

            if let Some(writer) = &mut writer {
                Self::write_to_logstash(writer, &entry);
            }

            if let Some(printer) = &self.printer {
                let s = format(&entry).expect("Unable to format");
                printer.write(s)
            }
        }
    }

    /// Writes a log line into json_lines logstash format, which has a newline at the end
    fn write_to_logstash(stream: &mut TcpWriter, entry: &LogEntry) {
        let message = if let Ok(json) = serde_json::to_string(entry) {
            json
        } else {
            STRUCT_LOG_PARSE_ERROR_COUNT.inc();
            return;
        };

        let message = message + "\n";
        let bytes = message.as_bytes();

        // Attempt to write the log up to NUM_SEND_RETRIES + 1, and then drop it
        // Each `write_all` call will attempt to open a connection if one isn't open
        let mut result = stream.write_all(bytes);
        for _ in 0..NUM_SEND_RETRIES {
            if result.is_ok() {
                break;
            } else {
                result = stream.write_all(bytes);
            }
        }

        if let Err(e) = result {
            STRUCT_LOG_SEND_ERROR_COUNT.inc();
            eprintln!(
                "[Logging] Error while sending data to logstash({}): {}",
                stream.endpoint(),
                e
            );
        } else {
            SENT_STRUCT_LOG_COUNT.inc()
        }
    }
}

/// An trait encapsulating the operations required for writing logs.
pub trait Writer: Send + Sync {
    /// Write the log.
    fn write(&self, log: String);
}

/// A struct for writing logs to stderr
struct StderrWriter;

impl Writer for StderrWriter {
    /// Write log to stderr
    fn write(&self, log: String) {
        eprintln!("{}", log);
    }
}

/// Converts a record into a string representation:
/// UNIX_TIMESTAMP LOG_LEVEL FILE:LINE MESSAGE
/// Example:
/// 2020-03-07 05:03:03 INFO common/libra-logger/src/lib.rs:261 Hello
fn format(entry: &LogEntry) -> Result<String, fmt::Error> {
    use std::fmt::Write;

    let mut w = String::new();
    write!(
        w,
        "{} {} {}",
        entry.metadata.level(),
        entry.timestamp,
        entry.metadata.location()
    )?;

    if let Some(message) = &entry.message {
        write!(w, " {}", message)?;
    }

    if !entry.data.is_empty() {
        write!(w, " {}", serde_json::to_string(&entry.data).unwrap())?;
    }

    Ok(w)
}

#[cfg(test)]
mod tests {
    use super::LogEntry;
    use crate::{
        debug, error, info, logger::Logger, trace, warn, Event, Key, KeyValue, Level, Metadata,
        Schema, Value, Visitor,
    };
    use chrono::{DateTime, Utc};
    use serde_json::Value as JsonValue;
    use std::sync::{
        mpsc::{self, Receiver, SyncSender},
        Arc,
    };

    #[derive(serde::Serialize)]
    #[serde(rename_all = "snake_case")]
    enum Enum {
        FooBar,
    }

    struct TestSchema<'a> {
        foo: usize,
        bar: &'a Enum,
    }

    impl Schema for TestSchema<'_> {
        fn visit(&self, visitor: &mut dyn Visitor) {
            visitor.visit_pair(Key::new("foo"), Value::from_serde(&self.foo));
            visitor.visit_pair(Key::new("bar"), Value::from_serde(&self.bar));
        }
    }

    struct LogStream(SyncSender<LogEntry>);

    impl LogStream {
        fn new() -> (Self, Receiver<LogEntry>) {
            let (sender, receiver) = mpsc::sync_channel(1024);
            (Self(sender), receiver)
        }
    }

    impl Logger for LogStream {
        fn enabled(&self, metadata: &Metadata) -> bool {
            metadata.level() <= Level::Debug
        }

        fn record(&self, event: &Event) {
            self.0.send(LogEntry::new(event)).unwrap();
        }
    }

    fn set_test_logger() -> Receiver<LogEntry> {
        let (logger, receiver) = LogStream::new();
        let logger = Arc::new(logger);
        crate::logger::set_global_logger(logger);
        receiver
    }

    // TODO: Find a better mechanism for testing that allows setting the logger not globally
    #[test]
    fn basic() {
        let receiver = set_test_logger();
        let number = 12345;

        // Send an info log
        let before = Utc::now();
        info!(
            TestSchema {
                foo: 5,
                bar: &Enum::FooBar
            },
            test = true,
            category = "name",
            KeyValue::new("display", Value::from_display(&number)),
            "This is a log"
        );
        let after = Utc::now();

        let entry = receiver.recv().unwrap();

        // Ensure standard fields are filled
        assert_eq!(entry.metadata.level(), Level::Info);
        assert_eq!(
            entry.metadata.target(),
            module_path!().split("::").next().unwrap()
        );
        assert_eq!(entry.metadata.module_path(), module_path!());
        assert_eq!(entry.metadata.file(), file!());
        assert_eq!(entry.message.as_deref(), Some("This is a log"));

        // Log time should be the time the structured log entry was created
        let timestamp = DateTime::parse_from_rfc3339(&entry.timestamp).unwrap();
        let timestamp: DateTime<Utc> = DateTime::from(timestamp);
        assert!(before <= timestamp && timestamp <= after);

        // Ensure data stored is the right type
        assert_eq!(entry.data.get("foo").and_then(JsonValue::as_u64), Some(5));
        assert_eq!(
            entry.data.get("bar").and_then(JsonValue::as_str),
            Some("foo_bar")
        );
        assert_eq!(
            entry.data.get("display").and_then(JsonValue::as_str),
            Some(format!("{}", number)).as_deref(),
        );
        assert_eq!(
            entry.data.get("test").and_then(JsonValue::as_bool),
            Some(true),
        );
        assert_eq!(
            entry.data.get("category").and_then(JsonValue::as_str),
            Some("name"),
        );

        // Test all log levels work properly
        // Tracing should be skipped because the Logger was setup to skip Tracing events
        trace!("trace");
        debug!("debug");
        info!("info");
        warn!("warn");
        error!("error");

        let levels = &[Level::Debug, Level::Info, Level::Warn, Level::Error];

        for level in levels {
            let entry = receiver.recv().unwrap();
            assert_eq!(entry.metadata.level(), *level);
        }
    }
}
<commit_msg>logger: change default log level to info<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    counters::{
        PROCESSED_STRUCT_LOG_COUNT, SENT_STRUCT_LOG_COUNT, STRUCT_LOG_PARSE_ERROR_COUNT,
        STRUCT_LOG_QUEUE_ERROR_COUNT, STRUCT_LOG_SEND_ERROR_COUNT,
    },
    logger::Logger,
    struct_log::TcpWriter,
    Event, Filter, Level, Metadata,
};
use chrono::{SecondsFormat, Utc};
use serde::Serialize;
use std::{
    collections::HashMap,
    env, fmt,
    io::Write,
    sync::{
        mpsc::{self, Receiver, SyncSender},
        Arc, RwLock,
    },
    thread,
};

const RUST_LOG: &str = "RUST_LOG";
pub const CHANNEL_SIZE: usize = 10000;
const NUM_SEND_RETRIES: u8 = 1;

#[derive(Debug, Serialize)]
pub struct LogEntry {
    #[serde(flatten)]
    metadata: Metadata,
    timestamp: String,
    #[serde(skip_serializing_if = "HashMap::is_empty")]
    data: HashMap<&'static str, serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    message: Option<String>,
}

impl LogEntry {
    fn new(event: &Event) -> Self {
        use crate::{Key, Value, Visitor};

        struct JsonVisitor<'a>(&'a mut HashMap<&'static str, serde_json::Value>);

        impl<'a> Visitor for JsonVisitor<'a> {
            fn visit_pair(&mut self, key: Key, value: Value<'_>) {
                let v = match value {
                    Value::Debug(d) => serde_json::Value::String(format!("{:?}", d)),
                    Value::Display(d) => serde_json::Value::String(d.to_string()),
                    Value::Serde(s) => serde_json::to_value(s).unwrap(),
                };

                self.0.insert(key.as_str(), v);
            }
        }

        let metadata = *event.metadata();
        let message = event.message().map(fmt::format);

        let mut data = HashMap::new();
        for schema in event.keys_and_values() {
            schema.visit(&mut JsonVisitor(&mut data));
        }

        Self {
            metadata,
            timestamp: Utc::now().to_rfc3339_opts(SecondsFormat::Micros, true),
            data,
            message,
        }
    }
}

pub struct LibraLoggerBuilder {
    channel_size: usize,
    level: Level,
    address: Option<String>,
    printer: Option<Box<dyn Writer>>,
    is_async: bool,
}

impl LibraLoggerBuilder {
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        Self {
            channel_size: CHANNEL_SIZE,
            level: Level::Info,
            address: None,
            printer: Some(Box::new(StderrWriter)),
            is_async: false,
        }
    }

    pub fn address(&mut self, address: String) -> &mut Self {
        self.address = Some(address);
        self
    }

    pub fn read_env(&mut self) -> &mut Self {
        if let Ok(address) = env::var("STRUCT_LOG_TCP_ADDR") {
            self.address(address);
        }
        self
    }

    pub fn level(&mut self, level: Level) -> &mut Self {
        self.level = level;
        self
    }

    pub fn channel_size(&mut self, channel_size: usize) -> &mut Self {
        self.channel_size = channel_size;
        self
    }

    pub fn printer(&mut self, printer: Box<dyn Writer + Send + Sync + 'static>) -> &mut Self {
        self.printer = Some(printer);
        self
    }

    pub fn is_async(&mut self, is_async: bool) -> &mut Self {
        self.is_async = is_async;
        self
    }

    pub fn init(&mut self) {
        self.build();
    }

    pub fn build(&mut self) -> Arc<LibraLogger> {
        let filter = {
            let mut filter_builder = Filter::builder();

            if env::var(RUST_LOG).is_ok() {
                filter_builder.with_env(RUST_LOG);
            } else {
                filter_builder.filter_level(self.level.into());
            }

            filter_builder.build()
        };

        let logger = if self.is_async {
            let (sender, receiver) = mpsc::sync_channel(self.channel_size);
            let service = LoggerService {
                receiver,
                address: self.address.clone(),
                printer: self.printer.take(),
            };
            let logger = Arc::new(LibraLogger {
                sender: Some(sender),
                printer: None,
                filter: RwLock::new(filter),
            });

            thread::spawn(move || service.run());
            logger
        } else {
            Arc::new(LibraLogger {
                sender: None,
                printer: self.printer.take(),
                filter: RwLock::new(filter),
            })
        };

        crate::logger::set_global_logger(logger.clone());
        logger
    }
}

pub struct LibraLogger {
    sender: Option<SyncSender<LogEntry>>,
    printer: Option<Box<dyn Writer>>,
    filter: RwLock<Filter>,
}

impl LibraLogger {
    pub fn builder() -> LibraLoggerBuilder {
        LibraLoggerBuilder::new()
    }

    #[allow(clippy::new_ret_no_self)]
    pub fn new() -> LibraLoggerBuilder {
        Self::builder()
    }

    pub fn init_for_testing() {
        if env::var(RUST_LOG).is_err() {
            return;
        }

        Self::builder()
            .is_async(false)
            .printer(Box::new(StderrWriter))
            .build();
    }

    pub fn set_filter(&self, filter: Filter) {
        *self.filter.write().unwrap() = filter;
    }

    fn send_entry(&self, entry: LogEntry) {
        if let Some(printer) = &self.printer {
            let s = format(&entry).expect("Unable to format");
            printer.write(s);
        }

        if let Some(sender) = &self.sender {
            if let Err(e) = sender.try_send(entry) {
                STRUCT_LOG_QUEUE_ERROR_COUNT.inc();
                eprintln!("Failed to send structured log: {}", e);
            }
        }
    }
}

impl Logger for LibraLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        self.filter.read().unwrap().enabled(metadata)
    }

    fn record(&self, event: &Event) {
        let entry = LogEntry::new(event);

        self.send_entry(entry)
    }
}

struct LoggerService {
    receiver: Receiver<LogEntry>,
    address: Option<String>,
    printer: Option<Box<dyn Writer>>,
}

impl LoggerService {
    pub fn run(mut self) {
        let mut writer = self.address.take().map(TcpWriter::new);

        for entry in self.receiver {
            PROCESSED_STRUCT_LOG_COUNT.inc();

            if let Some(writer) = &mut writer {
                Self::write_to_logstash(writer, &entry);
            }

            if let Some(printer) = &self.printer {
                let s = format(&entry).expect("Unable to format");
                printer.write(s)
            }
        }
    }

    /// Writes a log line into json_lines logstash format, which has a newline at the end
    fn write_to_logstash(stream: &mut TcpWriter, entry: &LogEntry) {
        let message = if let Ok(json) = serde_json::to_string(entry) {
            json
        } else {
            STRUCT_LOG_PARSE_ERROR_COUNT.inc();
            return;
        };

        let message = message + "\n";
        let bytes = message.as_bytes();

        // Attempt to write the log up to NUM_SEND_RETRIES + 1, and then drop it
        // Each `write_all` call will attempt to open a connection if one isn't open
        let mut result = stream.write_all(bytes);
        for _ in 0..NUM_SEND_RETRIES {
            if result.is_ok() {
                break;
            } else {
                result = stream.write_all(bytes);
            }
        }

        if let Err(e) = result {
            STRUCT_LOG_SEND_ERROR_COUNT.inc();
            eprintln!(
                "[Logging] Error while sending data to logstash({}): {}",
                stream.endpoint(),
                e
            );
        } else {
            SENT_STRUCT_LOG_COUNT.inc()
        }
    }
}

/// An trait encapsulating the operations required for writing logs.
pub trait Writer: Send + Sync {
    /// Write the log.
    fn write(&self, log: String);
}

/// A struct for writing logs to stderr
struct StderrWriter;

impl Writer for StderrWriter {
    /// Write log to stderr
    fn write(&self, log: String) {
        eprintln!("{}", log);
    }
}

/// Converts a record into a string representation:
/// UNIX_TIMESTAMP LOG_LEVEL FILE:LINE MESSAGE
/// Example:
/// 2020-03-07 05:03:03 INFO common/libra-logger/src/lib.rs:261 Hello
fn format(entry: &LogEntry) -> Result<String, fmt::Error> {
    use std::fmt::Write;

    let mut w = String::new();
    write!(
        w,
        "{} {} {}",
        entry.metadata.level(),
        entry.timestamp,
        entry.metadata.location()
    )?;

    if let Some(message) = &entry.message {
        write!(w, " {}", message)?;
    }

    if !entry.data.is_empty() {
        write!(w, " {}", serde_json::to_string(&entry.data).unwrap())?;
    }

    Ok(w)
}

#[cfg(test)]
mod tests {
    use super::LogEntry;
    use crate::{
        debug, error, info, logger::Logger, trace, warn, Event, Key, KeyValue, Level, Metadata,
        Schema, Value, Visitor,
    };
    use chrono::{DateTime, Utc};
    use serde_json::Value as JsonValue;
    use std::sync::{
        mpsc::{self, Receiver, SyncSender},
        Arc,
    };

    #[derive(serde::Serialize)]
    #[serde(rename_all = "snake_case")]
    enum Enum {
        FooBar,
    }

    struct TestSchema<'a> {
        foo: usize,
        bar: &'a Enum,
    }

    impl Schema for TestSchema<'_> {
        fn visit(&self, visitor: &mut dyn Visitor) {
            visitor.visit_pair(Key::new("foo"), Value::from_serde(&self.foo));
            visitor.visit_pair(Key::new("bar"), Value::from_serde(&self.bar));
        }
    }

    struct LogStream(SyncSender<LogEntry>);

    impl LogStream {
        fn new() -> (Self, Receiver<LogEntry>) {
            let (sender, receiver) = mpsc::sync_channel(1024);
            (Self(sender), receiver)
        }
    }

    impl Logger for LogStream {
        fn enabled(&self, metadata: &Metadata) -> bool {
            metadata.level() <= Level::Debug
        }

        fn record(&self, event: &Event) {
            self.0.send(LogEntry::new(event)).unwrap();
        }
    }

    fn set_test_logger() -> Receiver<LogEntry> {
        let (logger, receiver) = LogStream::new();
        let logger = Arc::new(logger);
        crate::logger::set_global_logger(logger);
        receiver
    }

    // TODO: Find a better mechanism for testing that allows setting the logger not globally
    #[test]
    fn basic() {
        let receiver = set_test_logger();
        let number = 12345;

        // Send an info log
        let before = Utc::now();
        info!(
            TestSchema {
                foo: 5,
                bar: &Enum::FooBar
            },
            test = true,
            category = "name",
            KeyValue::new("display", Value::from_display(&number)),
            "This is a log"
        );
        let after = Utc::now();

        let entry = receiver.recv().unwrap();

        // Ensure standard fields are filled
        assert_eq!(entry.metadata.level(), Level::Info);
        assert_eq!(
            entry.metadata.target(),
            module_path!().split("::").next().unwrap()
        );
        assert_eq!(entry.metadata.module_path(), module_path!());
        assert_eq!(entry.metadata.file(), file!());
        assert_eq!(entry.message.as_deref(), Some("This is a log"));

        // Log time should be the time the structured log entry was created
        let timestamp = DateTime::parse_from_rfc3339(&entry.timestamp).unwrap();
        let timestamp: DateTime<Utc> = DateTime::from(timestamp);
        assert!(before <= timestamp && timestamp <= after);

        // Ensure data stored is the right type
        assert_eq!(entry.data.get("foo").and_then(JsonValue::as_u64), Some(5));
        assert_eq!(
            entry.data.get("bar").and_then(JsonValue::as_str),
            Some("foo_bar")
        );
        assert_eq!(
            entry.data.get("display").and_then(JsonValue::as_str),
            Some(format!("{}", number)).as_deref(),
        );
        assert_eq!(
            entry.data.get("test").and_then(JsonValue::as_bool),
            Some(true),
        );
        assert_eq!(
            entry.data.get("category").and_then(JsonValue::as_str),
            Some("name"),
        );

        // Test all log levels work properly
        // Tracing should be skipped because the Logger was setup to skip Tracing events
        trace!("trace");
        debug!("debug");
        info!("info");
        warn!("warn");
        error!("error");

        let levels = &[Level::Debug, Level::Info, Level::Warn, Level::Error];

        for level in levels {
            let entry = receiver.recv().unwrap();
            assert_eq!(entry.metadata.level(), *level);
        }
    }
}
<|endoftext|>


<commit_before>extern crate test;
use super::TypedArena;
use std::cell::Cell;
use test::Bencher;

#[allow(dead_code)]
#[derive(Debug, Eq, PartialEq)]
struct Point {
    x: i32,
    y: i32,
    z: i32,
}

#[test]
pub fn test_unused() {
    let arena: TypedArena<Point> = TypedArena::default();
    assert!(arena.chunks.borrow().is_empty());
}

#[test]
fn test_arena_alloc_nested() {
    struct Inner {
        value: u8,
    }
    struct Outer<'a> {
        inner: &'a Inner,
    }
    enum EI<'e> {
        I(Inner),
        O(Outer<'e>),
    }

    struct Wrap<'a>(TypedArena<EI<'a>>);

    impl<'a> Wrap<'a> {
        fn alloc_inner<F: Fn() -> Inner>(&self, f: F) -> &Inner {
            let r: &EI<'_> = self.0.alloc(EI::I(f()));
            if let &EI::I(ref i) = r {
                i
            } else {
                panic!("mismatch");
            }
        }
        fn alloc_outer<F: Fn() -> Outer<'a>>(&self, f: F) -> &Outer<'_> {
            let r: &EI<'_> = self.0.alloc(EI::O(f()));
            if let &EI::O(ref o) = r {
                o
            } else {
                panic!("mismatch");
            }
        }
    }

    let arena = Wrap(TypedArena::default());

    let result = arena.alloc_outer(|| Outer { inner: arena.alloc_inner(|| Inner { value: 10 }) });

    assert_eq!(result.inner.value, 10);
}

#[test]
pub fn test_copy() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(Point { x: 1, y: 2, z: 3 });
    }
}

#[bench]
pub fn bench_copy(b: &mut Bencher) {
    let arena = TypedArena::default();
    b.iter(|| arena.alloc(Point { x: 1, y: 2, z: 3 }))
}

#[bench]
pub fn bench_copy_nonarena(b: &mut Bencher) {
    b.iter(|| {
        let _: Box<_> = Box::new(Point { x: 1, y: 2, z: 3 });
    })
}

#[allow(dead_code)]
struct Noncopy {
    string: String,
    array: Vec<i32>,
}

#[test]
pub fn test_noncopy() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] });
    }
}

#[test]
pub fn test_typed_arena_zero_sized() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(());
    }
}

#[test]
pub fn test_typed_arena_clear() {
    let mut arena = TypedArena::default();
    for _ in 0..10 {
        arena.clear();
        for _ in 0..10000 {
            arena.alloc(Point { x: 1, y: 2, z: 3 });
        }
    }
}

#[bench]
pub fn bench_typed_arena_clear(b: &mut Bencher) {
    let mut arena = TypedArena::default();
    b.iter(|| {
        arena.alloc(Point { x: 1, y: 2, z: 3 });
        arena.clear();
    })
}

// Drop tests

struct DropCounter<'a> {
    count: &'a Cell<u32>,
}

impl Drop for DropCounter<'_> {
    fn drop(&mut self) {
        self.count.set(self.count.get() + 1);
    }
}

#[test]
fn test_typed_arena_drop_count() {
    let counter = Cell::new(0);
    {
        let arena: TypedArena<DropCounter<'_>> = TypedArena::default();
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(DropCounter { count: &counter });
        }
    };
    assert_eq!(counter.get(), 100);
}

#[test]
fn test_typed_arena_drop_on_clear() {
    let counter = Cell::new(0);
    let mut arena: TypedArena<DropCounter<'_>> = TypedArena::default();
    for i in 0..10 {
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(DropCounter { count: &counter });
        }
        arena.clear();
        assert_eq!(counter.get(), i * 100 + 100);
    }
}

thread_local! {
    static DROP_COUNTER: Cell<u32> = Cell::new(0)
}

struct SmallDroppable;

impl Drop for SmallDroppable {
    fn drop(&mut self) {
        DROP_COUNTER.with(|c| c.set(c.get() + 1));
    }
}

#[test]
fn test_typed_arena_drop_small_count() {
    DROP_COUNTER.with(|c| c.set(0));
    {
        let arena: TypedArena<SmallDroppable> = TypedArena::default();
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(SmallDroppable);
        }
        // dropping
    };
    assert_eq!(DROP_COUNTER.with(|c| c.get()), 100);
}

#[bench]
pub fn bench_noncopy(b: &mut Bencher) {
    let arena = TypedArena::default();
    b.iter(|| {
        arena.alloc(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] })
    })
}

#[bench]
pub fn bench_noncopy_nonarena(b: &mut Bencher) {
    b.iter(|| {
        let _: Box<_> =
            Box::new(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] });
    })
}
<commit_msg>Add bench_typed_arena_clear_100 bench<commit_after>extern crate test;
use super::TypedArena;
use std::cell::Cell;
use test::Bencher;

#[allow(dead_code)]
#[derive(Debug, Eq, PartialEq)]
struct Point {
    x: i32,
    y: i32,
    z: i32,
}

#[test]
pub fn test_unused() {
    let arena: TypedArena<Point> = TypedArena::default();
    assert!(arena.chunks.borrow().is_empty());
}

#[test]
fn test_arena_alloc_nested() {
    struct Inner {
        value: u8,
    }
    struct Outer<'a> {
        inner: &'a Inner,
    }
    enum EI<'e> {
        I(Inner),
        O(Outer<'e>),
    }

    struct Wrap<'a>(TypedArena<EI<'a>>);

    impl<'a> Wrap<'a> {
        fn alloc_inner<F: Fn() -> Inner>(&self, f: F) -> &Inner {
            let r: &EI<'_> = self.0.alloc(EI::I(f()));
            if let &EI::I(ref i) = r {
                i
            } else {
                panic!("mismatch");
            }
        }
        fn alloc_outer<F: Fn() -> Outer<'a>>(&self, f: F) -> &Outer<'_> {
            let r: &EI<'_> = self.0.alloc(EI::O(f()));
            if let &EI::O(ref o) = r {
                o
            } else {
                panic!("mismatch");
            }
        }
    }

    let arena = Wrap(TypedArena::default());

    let result = arena.alloc_outer(|| Outer { inner: arena.alloc_inner(|| Inner { value: 10 }) });

    assert_eq!(result.inner.value, 10);
}

#[test]
pub fn test_copy() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(Point { x: 1, y: 2, z: 3 });
    }
}

#[bench]
pub fn bench_copy(b: &mut Bencher) {
    let arena = TypedArena::default();
    b.iter(|| arena.alloc(Point { x: 1, y: 2, z: 3 }))
}

#[bench]
pub fn bench_copy_nonarena(b: &mut Bencher) {
    b.iter(|| {
        let _: Box<_> = Box::new(Point { x: 1, y: 2, z: 3 });
    })
}

#[allow(dead_code)]
struct Noncopy {
    string: String,
    array: Vec<i32>,
}

#[test]
pub fn test_noncopy() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] });
    }
}

#[test]
pub fn test_typed_arena_zero_sized() {
    let arena = TypedArena::default();
    for _ in 0..100000 {
        arena.alloc(());
    }
}

#[test]
pub fn test_typed_arena_clear() {
    let mut arena = TypedArena::default();
    for _ in 0..10 {
        arena.clear();
        for _ in 0..10000 {
            arena.alloc(Point { x: 1, y: 2, z: 3 });
        }
    }
}

#[bench]
pub fn bench_typed_arena_clear(b: &mut Bencher) {
    let mut arena = TypedArena::default();
    b.iter(|| {
        arena.alloc(Point { x: 1, y: 2, z: 3 });
        arena.clear();
    })
}

#[bench]
pub fn bench_typed_arena_clear_100(b: &mut Bencher) {
    let mut arena = TypedArena::default();
    b.iter(|| {
        for _ in 0..100 {
            arena.alloc(Point { x: 1, y: 2, z: 3 });
        }
        arena.clear();
    })
}

// Drop tests

struct DropCounter<'a> {
    count: &'a Cell<u32>,
}

impl Drop for DropCounter<'_> {
    fn drop(&mut self) {
        self.count.set(self.count.get() + 1);
    }
}

#[test]
fn test_typed_arena_drop_count() {
    let counter = Cell::new(0);
    {
        let arena: TypedArena<DropCounter<'_>> = TypedArena::default();
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(DropCounter { count: &counter });
        }
    };
    assert_eq!(counter.get(), 100);
}

#[test]
fn test_typed_arena_drop_on_clear() {
    let counter = Cell::new(0);
    let mut arena: TypedArena<DropCounter<'_>> = TypedArena::default();
    for i in 0..10 {
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(DropCounter { count: &counter });
        }
        arena.clear();
        assert_eq!(counter.get(), i * 100 + 100);
    }
}

thread_local! {
    static DROP_COUNTER: Cell<u32> = Cell::new(0)
}

struct SmallDroppable;

impl Drop for SmallDroppable {
    fn drop(&mut self) {
        DROP_COUNTER.with(|c| c.set(c.get() + 1));
    }
}

#[test]
fn test_typed_arena_drop_small_count() {
    DROP_COUNTER.with(|c| c.set(0));
    {
        let arena: TypedArena<SmallDroppable> = TypedArena::default();
        for _ in 0..100 {
            // Allocate something with drop glue to make sure it doesn't leak.
            arena.alloc(SmallDroppable);
        }
        // dropping
    };
    assert_eq!(DROP_COUNTER.with(|c| c.get()), 100);
}

#[bench]
pub fn bench_noncopy(b: &mut Bencher) {
    let arena = TypedArena::default();
    b.iter(|| {
        arena.alloc(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] })
    })
}

#[bench]
pub fn bench_noncopy_nonarena(b: &mut Bencher) {
    b.iter(|| {
        let _: Box<_> =
            Box::new(Noncopy { string: "hello world".to_string(), array: vec![1, 2, 3, 4, 5] });
    })
}
<|endoftext|>