<commit_before>#[macro_use]
extern crate serde_derive;

extern crate serde;

extern crate toml;
extern crate clap;

use clap::{App, Arg, ArgMatches};
use std::path::Path;
use std::fs::{create_dir, File};
use std::io::prelude::*;
use std::io::Result;

#[derive(Serialize, Deserialize)]
struct Config {
    data_dir: Option<String>,
}

const CONFIG_PATH: &'static str = "/etc/blochs/";
const CONFIG_FILE_NAME: &'static str = "server.toml";
const DEFAULT_DATA_DIR: &'static str = "/var/lib/blochs/";

fn main() {
    let options = get_provided_options();

    let config_path = Path::new(CONFIG_PATH);
    ensure_dir_exists(&config_path);

    let config_file_path = config_path.join(CONFIG_FILE_NAME);

    let actual_content = get_file_content(&config_file_path);

    let mut config_values: Config = toml::from_str(&actual_content).unwrap();

    let config_or_default_data_dir = config_values.data_dir.unwrap_or(DEFAULT_DATA_DIR.to_string());

    let new_data_dir = options.value_of("data.dir").unwrap_or(&config_or_default_data_dir);

    config_values.data_dir = Some(new_data_dir.to_string());

    let new_config_content = toml::to_string(&config_values).unwrap();

    match rewrite_file(&config_file_path, &new_config_content) {
        Ok(_) => println!("New config saved at {}:\n\n{}", config_file_path.display(), new_config_content),
        Err(why) => panic!("Could not write config file {:?}: {}", config_file_path.display(), why),
    };
}

fn rewrite_file(path: &Path, new_content: &String) -> Result<()> {
    let mut file = match File::create(path) {
        Ok(file) => file,
        Err(err) => panic!("Could not open config file {:?}: {}", path.display(), err),
    };

    return file.write_all(new_content.as_bytes());
}

fn get_file_content(path: &Path) -> String {
    let mut content = String::new();
    if path.exists() {
        let mut file = match File::open(&path) {
            Ok(file) => file,
            Err(err) => panic!("Could not open file {:?}: {}", path.display(), err),
        };

        match file.read_to_string(&mut content) {
            Ok(_) => {},
            Err(err) => panic!("Could not read file {:?}: {}", path.display(), err),
        };
    }
    return content;
}

fn ensure_dir_exists(path: &Path) {
    if !path.exists() {
        match create_dir(path) {
            Err(err) => panic!("Could not create directory under {}: {}", path.display(), err),
            Ok(_) => {}
        };
    }
}

fn get_provided_options<'a>() -> ArgMatches<'a> {
    return App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::with_name("data.dir")
            .long("data-dir")
            .value_name("DIR")
            .help(&format!("Sets where database data will be stored (default {})", DEFAULT_DATA_DIR))
        ).get_matches();
}
<commit_msg>extract options settings to a new function<commit_after>#[macro_use]
extern crate serde_derive;

extern crate serde;

extern crate toml;
extern crate clap;

use clap::{App, Arg, ArgMatches};
use std::path::Path;
use std::fs::{create_dir, File};
use std::io::prelude::*;
use std::io::Result;

#[derive(Serialize, Deserialize)]
struct Config {
    data_dir: Option<String>,
}

const CONFIG_PATH: &'static str = "/etc/blochs/";
const CONFIG_FILE_NAME: &'static str = "server.toml";
const DEFAULT_DATA_DIR: &'static str = "/var/lib/blochs/";

fn main() {
    let options = get_provided_options();

    let config_path = Path::new(CONFIG_PATH);
    ensure_dir_exists(&config_path);

    let config_file_path = config_path.join(CONFIG_FILE_NAME);
    let actual_content = get_file_content(&config_file_path);

    let mut config_values: Config = toml::from_str(&actual_content).unwrap();
    set_options(&mut config_values, &options);

    let new_config_content = toml::to_string(&config_values).unwrap();
    match rewrite_file(&config_file_path, &new_config_content) {
        Ok(_) => println!("New config saved at {}:\n\n{}", config_file_path.display(), new_config_content),
        Err(why) => panic!("Could not write config file {:?}: {}", config_file_path.display(), why),
    };
}

fn set_options(config_values: &mut Config, options: &ArgMatches) {
    let config_or_default_data_dir = config_values.data_dir.unwrap_or(DEFAULT_DATA_DIR.to_string());
    let new_data_dir = options.value_of("data.dir").unwrap_or(&config_or_default_data_dir);

    config_values.data_dir = Some(new_data_dir.to_string());
}

fn rewrite_file(path: &Path, new_content: &String) -> Result<()> {
    let mut file = match File::create(path) {
        Ok(file) => file,
        Err(err) => panic!("Could not open config file {:?}: {}", path.display(), err),
    };

    return file.write_all(new_content.as_bytes());
}

fn get_file_content(path: &Path) -> String {
    let mut content = String::new();
    if path.exists() {
        let mut file = match File::open(&path) {
            Ok(file) => file,
            Err(err) => panic!("Could not open file {:?}: {}", path.display(), err),
        };

        match file.read_to_string(&mut content) {
            Ok(_) => {},
            Err(err) => panic!("Could not read file {:?}: {}", path.display(), err),
        };
    }
    return content;
}

fn ensure_dir_exists(path: &Path) {
    if !path.exists() {
        match create_dir(path) {
            Err(err) => panic!("Could not create directory under {}: {}", path.display(), err),
            Ok(_) => {}
        };
    }
}

fn get_provided_options<'a>() -> ArgMatches<'a> {
    return App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .author(env!("CARGO_PKG_AUTHORS"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .arg(Arg::with_name("data.dir")
            .long("data-dir")
            .value_name("DIR")
            .help(&format!("Sets where database data will be stored (default {})", DEFAULT_DATA_DIR))
        ).get_matches();
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Servo's compiler plugin/macro crate
//!
//! Attributes this crate provides:
//!
//!  - `#[derive(DenyPublicFields)]` : Forces all fields in a struct/enum to be private
//!  - `#[derive(JSTraceable)]` : Auto-derives an implementation of `JSTraceable` for a struct in the script crate
//!  - `#[must_root]` : Prevents data of the marked type from being used on the stack.
//!                     See the lints module for more details
//!  - `#[dom_struct]` : Implies #[derive(JSTraceable, DenyPublicFields)]`, and `#[must_root]`.
//!                       Use this for structs that correspond to a DOM type

#![deny(unsafe_code)]
#![feature(plugin)]
#![feature(plugin_registrar)]
#![feature(rustc_private)]

#[cfg(feature = "unrooted_must_root_lint")]
#[macro_use]
extern crate rustc;

extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::Registry;
use syntax::feature_gate::AttributeType::Whitelisted;

#[cfg(feature = "unrooted_must_root_lint")]
mod unrooted_must_root;

/// Utilities for writing plugins
#[cfg(feature = "unrooted_must_root_lint")]
mod utils;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    #[cfg(feature = "unrooted_must_root_lint")]
    reg.register_late_lint_pass(Box::new(unrooted_must_root::UnrootedPass::new()));

    reg.register_attribute("allow_unrooted_interior".to_string(), Whitelisted);
    reg.register_attribute("must_root".to_string(), Whitelisted);
}
<commit_msg>Fix symbol internment in script-servo<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Servo's compiler plugin/macro crate
//!
//! Attributes this crate provides:
//!
//!  - `#[derive(DenyPublicFields)]` : Forces all fields in a struct/enum to be private
//!  - `#[derive(JSTraceable)]` : Auto-derives an implementation of `JSTraceable` for a struct in the script crate
//!  - `#[must_root]` : Prevents data of the marked type from being used on the stack.
//!                     See the lints module for more details
//!  - `#[dom_struct]` : Implies #[derive(JSTraceable, DenyPublicFields)]`, and `#[must_root]`.
//!                       Use this for structs that correspond to a DOM type

#![deny(unsafe_code)]
#![feature(plugin)]
#![feature(plugin_registrar)]
#![feature(rustc_private)]

#[cfg(feature = "unrooted_must_root_lint")]
#[macro_use]
extern crate rustc;

extern crate rustc_plugin;
extern crate syntax;
extern crate syntax_pos;

use rustc_plugin::Registry;
use syntax::feature_gate::AttributeType::Whitelisted;
use syntax_pos::symbol::Symbol;

#[cfg(feature = "unrooted_must_root_lint")]
mod unrooted_must_root;

/// Utilities for writing plugins
#[cfg(feature = "unrooted_must_root_lint")]
mod utils;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    #[cfg(feature = "unrooted_must_root_lint")]
    reg.register_late_lint_pass(Box::new(unrooted_must_root::UnrootedPass::new()));

    reg.register_attribute(Symbol::intern("allow_unrooted_interior"), Whitelisted);
    reg.register_attribute(Symbol::intern("must_root"), Whitelisted);
}

<|endoftext|>


<commit_before>extern crate gnutls;

use gnutls::creds::{Cert, CredType};
use gnutls::{
    init,
    deinit,
    CertVerifyFlags,
    Session
};

use std::io::prelude::*;
use std::os::unix::io::AsRawFd;
use std::net::TcpStream;

fn main() {
    match init() {
        Ok(_) => {},
        Err(_) => panic!("could not initialize.")
    };

    let hostname = "localhost";
    let mut creds = Cert::new().unwrap();
    let mut session = Session::new(gnutls::GNUTLS_CLIENT).unwrap();

    if creds.x509_set_system_trust().err() != None {
        panic!("Error: couldn't set system trust.");
    }

    if creds.x509_set_key_file("../../gnutls/tests/ca.cert.pem",
                               "../../gnutls/tests/ca.noenckey.pem", None).err() != None {
        panic!("Error: couldn't set the keyfile.");
    }

    match session.set_priority(None) {
        Err(e) => panic!("{}", e),
        Ok(_) => {}
    };

    if session.set_creds(CredType::GNUTLS_CRD_CERTIFICATE, &mut creds) .err() != None {
        panic!("Error: could not set the session credentials.");
    }

    session.set_verify_cert(hostname,
                            Some(CertVerifyFlags::GNUTLS_VERIFY_ALLOW_ANY_X509_V1_CA_CRT));

    let mut stream = match TcpStream::connect("bheesham.com:443") {
        Err(e) => panic!("{}", e),
        Ok(s) => s
    };

    session.set_fd(stream.as_raw_fd());
    session.handshake_timeout(None);

    match session.handshake() {
        Ok(_) => println!("Established a connection!"),
        Err(e) => panic!("Error: {}", e)
    };

    let _ = stream.write("GET / HTTP/1.0\n\n".as_bytes()).unwrap();

    let mut res = String::new();
    let _ = stream.read_to_string(&mut res).unwrap();
    println!("{}", res);

    deinit();
}
<commit_msg>Actually use hostname verification in example<commit_after>extern crate gnutls;

use gnutls::creds::{Cert, CredType};
use gnutls::{init, deinit, Session};

use std::io::prelude::*;
use std::os::unix::io::AsRawFd;
use std::net::TcpStream;

fn main() {
    match init() {
        Ok(_) => {},
        Err(_) => panic!("could not initialize.")
    };

    let host = ("bheesham.com", 443);
    let mut creds = Cert::new().unwrap();
    let mut session = Session::new(gnutls::GNUTLS_CLIENT).unwrap();

    if creds.x509_set_system_trust().err() != None {
        panic!("Error: couldn't set system trust.");
    }

    if creds.x509_set_key_file("../../gnutls/tests/ca.cert.pem",
                               "../../gnutls/tests/ca.noenckey.pem", None).err() != None {
        panic!("Error: couldn't set the keyfile.");
    }

    match session.set_priority(None) {
        Err(e) => panic!("{}", e),
        Ok(_) => {}
    };

    if session.set_creds(CredType::GNUTLS_CRD_CERTIFICATE, &mut creds) .err() != None {
        panic!("Error: could not set the session credentials.");
    }

    session.set_verify_cert(host.0, None);

    let mut stream = match TcpStream::connect(host) {
        Err(e) => panic!("{}", e),
        Ok(s) => s
    };

    session.set_fd(stream.as_raw_fd());
    session.handshake_timeout(None);

    match session.handshake() {
        Ok(_) => println!("Established a connection!"),
        Err(e) => panic!("Error: {}", e)
    };

    let _ = stream.write("GET / HTTP/1.0\n\n".as_bytes()).unwrap();

    let mut res = String::new();
    let _ = stream.read_to_string(&mut res).unwrap();
    println!("{}", res);

    deinit();
}
<|endoftext|>


<commit_before>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use crate::codec::{Error, Result};
use codec::prelude::*;
use num_traits::PrimInt;
use std::cmp::Ordering::{Equal, Greater, Less};
use std::marker::PhantomData;

pub enum RowSlice<'a> {
    Small {
        non_null_ids: LEBytes<'a, u8>,
        null_ids: LEBytes<'a, u8>,
        offsets: LEBytes<'a, u16>,
        values: LEBytes<'a, u8>,
    },
    Big {
        non_null_ids: LEBytes<'a, u32>,
        null_ids: LEBytes<'a, u32>,
        offsets: LEBytes<'a, u32>,
        values: LEBytes<'a, u8>,
    },
}

impl RowSlice<'_> {
    /// # Panics
    ///
    /// Panics if the value of first byte is not 128(v2 version code)
    pub fn from_bytes(mut data: &[u8]) -> Result<RowSlice> {
        assert_eq!(data.read_u8()?, super::CODEC_VERSION);
        let is_big = super::Flags::from_bits_truncate(data.read_u8()?) == super::Flags::BIG;

        // read ids count
        let non_null_cnt = data.read_u16_le()? as usize;
        let null_cnt = data.read_u16_le()? as usize;
        let row = if is_big {
            RowSlice::Big {
                non_null_ids: read_le_bytes(&mut data, non_null_cnt)?,
                null_ids: read_le_bytes(&mut data, null_cnt)?,
                offsets: read_le_bytes(&mut data, non_null_cnt)?,
                values: LEBytes::new(data),
            }
        } else {
            RowSlice::Small {
                non_null_ids: read_le_bytes(&mut data, non_null_cnt)?,
                null_ids: read_le_bytes(&mut data, null_cnt)?,
                offsets: read_le_bytes(&mut data, non_null_cnt)?,
                values: LEBytes::new(data),
            }
        };
        Ok(row)
    }

    /// Search `id` in non-null ids
    ///
    /// Returns the `start` position and `offset` in `values` field if found, otherwise returns `None`
    ///
    /// # Errors
    ///
    /// If the id is found with no offset(It will only happen when the row data is broken),
    /// `Error::ColumnOffset` will be returned.
    pub fn search_in_non_null_ids(&self, id: i64) -> Result<Option<(usize, usize)>> {
        if !self.id_valid(id) {
            return Ok(None);
        }
        match self {
            RowSlice::Big {
                non_null_ids,
                offsets,
                ..
            } => {
                if let Ok(idx) = non_null_ids.binary_search(&(id as u32)) {
                    let offset = offsets.get(idx).ok_or(Error::ColumnOffset(idx))?;
                    let start = if idx > 0 {
                        // Previous `offsets.get(idx)` indicates it's ok to index `idx - 1`
                        unsafe { offsets.get_unchecked(idx - 1) as usize }
                    } else {
                        0usize
                    };
                    return Ok(Some((start, (offset as usize))));
                }
            }
            RowSlice::Small {
                non_null_ids,
                offsets,
                ..
            } => {
                if let Ok(idx) = non_null_ids.binary_search(&(id as u8)) {
                    let offset = offsets.get(idx).ok_or(Error::ColumnOffset(idx))?;
                    let start = if idx > 0 {
                        // Previous `offsets.get(idx)` indicates it's ok to index `idx - 1`
                        unsafe { offsets.get_unchecked(idx - 1) as usize }
                    } else {
                        0usize
                    };
                    return Ok(Some((start, (offset as usize))));
                }
            }
        }
        Ok(None)
    }

    /// Search `id` in null ids
    ///
    /// Returns true if found
    pub fn search_in_null_ids(&self, id: i64) -> bool {
        match self {
            RowSlice::Big { null_ids, .. } => null_ids.binary_search(&(id as u32)).is_ok(),
            RowSlice::Small { null_ids, .. } => null_ids.binary_search(&(id as u8)).is_ok(),
        }
    }

    #[inline]
    fn id_valid(&self, id: i64) -> bool {
        let upper: i64 = if self.is_big() {
            i64::from(u32::max_value())
        } else {
            i64::from(u8::max_value())
        };
        id > 0 && id <= upper
    }

    #[inline]
    fn is_big(&self) -> bool {
        match self {
            RowSlice::Big { .. } => true,
            RowSlice::Small { .. } => false,
        }
    }

    #[inline]
    pub fn values(&self) -> &[u8] {
        match self {
            RowSlice::Big { values, .. } => values.slice,
            RowSlice::Small { values, .. } => values.slice,
        }
    }
}

/// Decodes `len` number of ints from `buf` in little endian
///
/// Note:
/// This method is only implemented on little endianness currently, since x86 use little endianness.
#[cfg(target_endian = "little")]
#[inline]
fn read_le_bytes<'a, T>(buf: &mut &'a [u8], len: usize) -> Result<LEBytes<'a, T>>
where
    T: PrimInt,
{
    let bytes_len = std::mem::size_of::<T>() * len;
    if buf.len() < bytes_len {
        return Err(Error::unexpected_eof());
    }
    let slice = &buf[..bytes_len];
    buf.advance(bytes_len);
    Ok(LEBytes::new(slice))
}

#[cfg(target_endian = "little")]
pub struct LEBytes<'a, T: PrimInt> {
    slice: &'a [u8],
    _marker: PhantomData<T>,
}

#[cfg(target_endian = "little")]
impl<'a, T: PrimInt> LEBytes<'a, T> {
    fn new(slice: &'a [u8]) -> Self {
        Self {
            slice,
            _marker: PhantomData::default(),
        }
    }

    #[inline]
    fn get(&self, index: usize) -> Option<T> {
        if std::mem::size_of::<T>() * index >= self.slice.len() {
            None
        } else {
            unsafe { Some(self.get_unchecked(index)) }
        }
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> T {
        let ptr = self.slice.as_ptr() as *const T;
        let ptr = ptr.add(index);
        std::ptr::read_unaligned(ptr)
    }

    #[inline]
    fn binary_search(&self, value: &T) -> std::result::Result<usize, usize> {
        let mut size = self.slice.len() / std::mem::size_of::<T>();
        if size == 0 {
            return Err(0);
        }
        let mut base = 0usize;

        // Note that the count of ids is not greater than `u16::MAX`. The number
        // of binary search steps will not over 16 unless the data is corrupted.
        // Let's relex to 20.
        let mut steps = 20usize;

        while steps > 0 && size > 1 {
            let half = size / 2;
            let mid = base + half;
            let cmp = unsafe { self.get_unchecked(mid) }.cmp(value);
            base = if cmp == Greater { base } else { mid };
            size -= half;
            steps -= 1;
        }

        let cmp = unsafe { self.get_unchecked(base) }.cmp(value);
        if cmp == Equal {
            Ok(base)
        } else {
            Err(base + (cmp == Less) as usize)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::encoder_for_test::{Column, RowEncoder};
    use super::{read_le_bytes, RowSlice};
    use crate::codec::data_type::ScalarValue;
    use crate::expr::EvalContext;
    use codec::prelude::NumberEncoder;
    use std::u16;

    #[test]
    fn test_read_le_bytes() {
        let data = vec![1, 128, 512, u16::MAX, 256];
        let mut buf = vec![];
        for n in &data {
            buf.write_u16_le(*n).unwrap();
        }

        for i in 1..=data.len() {
            let le_bytes = read_le_bytes::<u16>(&mut buf.as_slice(), i).unwrap();
            for j in 0..i {
                assert_eq!(unsafe { le_bytes.get_unchecked(j) }, data[j]);
            }
        }
    }

    fn encoded_data_big() -> Vec<u8> {
        let cols = vec![
            Column::new(1, 1000),
            Column::new(356, 2),
            Column::new(33, ScalarValue::Int(None)),
            Column::new(3, 3),
        ];
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    fn encoded_data() -> Vec<u8> {
        let cols = vec![
            Column::new(1, 1000),
            Column::new(33, ScalarValue::Int(None)),
            Column::new(3, 3),
        ];
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    #[test]
    fn test_search_in_non_null_ids() {
        let data = encoded_data_big();
        let big_row = RowSlice::from_bytes(&data).unwrap();
        assert!(big_row.is_big());
        assert_eq!(big_row.search_in_non_null_ids(33).unwrap(), None);
        assert_eq!(big_row.search_in_non_null_ids(333).unwrap(), None);
        assert_eq!(
            big_row
                .search_in_non_null_ids(i64::from(u32::max_value()) + 2)
                .unwrap(),
            None
        );
        assert_eq!(Some((0, 2)), big_row.search_in_non_null_ids(1).unwrap());
        assert_eq!(Some((3, 4)), big_row.search_in_non_null_ids(356).unwrap());

        let data = encoded_data();
        let row = RowSlice::from_bytes(&data).unwrap();
        assert!(!row.is_big());
        assert_eq!(row.search_in_non_null_ids(33).unwrap(), None);
        assert_eq!(row.search_in_non_null_ids(35).unwrap(), None);
        assert_eq!(
            row.search_in_non_null_ids(i64::from(u8::max_value()) + 2)
                .unwrap(),
            None
        );
        assert_eq!(Some((0, 2)), row.search_in_non_null_ids(1).unwrap());
        assert_eq!(Some((2, 3)), row.search_in_non_null_ids(3).unwrap());
    }

    #[test]
    fn test_search_in_null_ids() {
        let data = encoded_data_big();
        let row = RowSlice::from_bytes(&data).unwrap();
        assert!(row.search_in_null_ids(33));
        assert!(!row.search_in_null_ids(3));
        assert!(!row.search_in_null_ids(333));
    }
}

#[cfg(test)]
mod benches {
    use super::super::encoder_for_test::{Column, RowEncoder};
    use super::RowSlice;
    use crate::codec::data_type::ScalarValue;
    use crate::expr::EvalContext;
    use test::black_box;

    fn encoded_data(len: usize) -> Vec<u8> {
        let mut cols = vec![];
        for i in 0..(len as i64) {
            if i % 10 == 0 {
                cols.push(Column::new(i, ScalarValue::Int(None)))
            } else {
                cols.push(Column::new(i, i))
            }
        }
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    #[bench]
    fn bench_search_in_non_null_ids(b: &mut test::Bencher) {
        let data = encoded_data(10);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(3))
        });
    }

    #[bench]
    fn bench_search_in_non_null_ids_middle(b: &mut test::Bencher) {
        let data = encoded_data(100);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(89))
        });
    }

    #[bench]
    fn bench_search_in_null_ids_middle(b: &mut test::Bencher) {
        let data = encoded_data(100);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(20))
        });
    }

    #[bench]
    fn bench_search_in_non_null_ids_big(b: &mut test::Bencher) {
        let data = encoded_data(350);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(257))
        });
    }

    #[bench]
    fn bench_from_bytes_big(b: &mut test::Bencher) {
        let data = encoded_data(350);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(&row);
        });
    }
}
<commit_msg>copr: Fix search_null incorrect result when scanning small rows for column id > 255 (#9128)<commit_after>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use crate::codec::{Error, Result};
use codec::prelude::*;
use num_traits::PrimInt;
use std::cmp::Ordering::{Equal, Greater, Less};
use std::marker::PhantomData;

pub enum RowSlice<'a> {
    Small {
        non_null_ids: LEBytes<'a, u8>,
        null_ids: LEBytes<'a, u8>,
        offsets: LEBytes<'a, u16>,
        values: LEBytes<'a, u8>,
    },
    Big {
        non_null_ids: LEBytes<'a, u32>,
        null_ids: LEBytes<'a, u32>,
        offsets: LEBytes<'a, u32>,
        values: LEBytes<'a, u8>,
    },
}

impl RowSlice<'_> {
    /// # Panics
    ///
    /// Panics if the value of first byte is not 128(v2 version code)
    pub fn from_bytes(mut data: &[u8]) -> Result<RowSlice> {
        assert_eq!(data.read_u8()?, super::CODEC_VERSION);
        let is_big = super::Flags::from_bits_truncate(data.read_u8()?) == super::Flags::BIG;

        // read ids count
        let non_null_cnt = data.read_u16_le()? as usize;
        let null_cnt = data.read_u16_le()? as usize;
        let row = if is_big {
            RowSlice::Big {
                non_null_ids: read_le_bytes(&mut data, non_null_cnt)?,
                null_ids: read_le_bytes(&mut data, null_cnt)?,
                offsets: read_le_bytes(&mut data, non_null_cnt)?,
                values: LEBytes::new(data),
            }
        } else {
            RowSlice::Small {
                non_null_ids: read_le_bytes(&mut data, non_null_cnt)?,
                null_ids: read_le_bytes(&mut data, null_cnt)?,
                offsets: read_le_bytes(&mut data, non_null_cnt)?,
                values: LEBytes::new(data),
            }
        };
        Ok(row)
    }

    /// Search `id` in non-null ids
    ///
    /// Returns the `start` position and `offset` in `values` field if found, otherwise returns `None`
    ///
    /// # Errors
    ///
    /// If the id is found with no offset(It will only happen when the row data is broken),
    /// `Error::ColumnOffset` will be returned.
    pub fn search_in_non_null_ids(&self, id: i64) -> Result<Option<(usize, usize)>> {
        if !self.id_valid(id) {
            return Ok(None);
        }
        match self {
            RowSlice::Big {
                non_null_ids,
                offsets,
                ..
            } => {
                if let Ok(idx) = non_null_ids.binary_search(&(id as u32)) {
                    let offset = offsets.get(idx).ok_or(Error::ColumnOffset(idx))?;
                    let start = if idx > 0 {
                        // Previous `offsets.get(idx)` indicates it's ok to index `idx - 1`
                        unsafe { offsets.get_unchecked(idx - 1) as usize }
                    } else {
                        0usize
                    };
                    return Ok(Some((start, (offset as usize))));
                }
            }
            RowSlice::Small {
                non_null_ids,
                offsets,
                ..
            } => {
                if let Ok(idx) = non_null_ids.binary_search(&(id as u8)) {
                    let offset = offsets.get(idx).ok_or(Error::ColumnOffset(idx))?;
                    let start = if idx > 0 {
                        // Previous `offsets.get(idx)` indicates it's ok to index `idx - 1`
                        unsafe { offsets.get_unchecked(idx - 1) as usize }
                    } else {
                        0usize
                    };
                    return Ok(Some((start, (offset as usize))));
                }
            }
        }
        Ok(None)
    }

    /// Search `id` in null ids
    ///
    /// Returns true if found
    pub fn search_in_null_ids(&self, id: i64) -> bool {
        if !self.id_valid(id) {
            return false;
        }
        match self {
            RowSlice::Big { null_ids, .. } => null_ids.binary_search(&(id as u32)).is_ok(),
            RowSlice::Small { null_ids, .. } => null_ids.binary_search(&(id as u8)).is_ok(),
        }
    }

    #[inline]
    fn id_valid(&self, id: i64) -> bool {
        let upper: i64 = if self.is_big() {
            i64::from(u32::max_value())
        } else {
            i64::from(u8::max_value())
        };
        id > 0 && id <= upper
    }

    #[inline]
    fn is_big(&self) -> bool {
        match self {
            RowSlice::Big { .. } => true,
            RowSlice::Small { .. } => false,
        }
    }

    #[inline]
    pub fn values(&self) -> &[u8] {
        match self {
            RowSlice::Big { values, .. } => values.slice,
            RowSlice::Small { values, .. } => values.slice,
        }
    }
}

/// Decodes `len` number of ints from `buf` in little endian
///
/// Note:
/// This method is only implemented on little endianness currently, since x86 use little endianness.
#[cfg(target_endian = "little")]
#[inline]
fn read_le_bytes<'a, T>(buf: &mut &'a [u8], len: usize) -> Result<LEBytes<'a, T>>
where
    T: PrimInt,
{
    let bytes_len = std::mem::size_of::<T>() * len;
    if buf.len() < bytes_len {
        return Err(Error::unexpected_eof());
    }
    let slice = &buf[..bytes_len];
    buf.advance(bytes_len);
    Ok(LEBytes::new(slice))
}

#[cfg(target_endian = "little")]
pub struct LEBytes<'a, T: PrimInt> {
    slice: &'a [u8],
    _marker: PhantomData<T>,
}

#[cfg(target_endian = "little")]
impl<'a, T: PrimInt> LEBytes<'a, T> {
    fn new(slice: &'a [u8]) -> Self {
        Self {
            slice,
            _marker: PhantomData::default(),
        }
    }

    #[inline]
    fn get(&self, index: usize) -> Option<T> {
        if std::mem::size_of::<T>() * index >= self.slice.len() {
            None
        } else {
            unsafe { Some(self.get_unchecked(index)) }
        }
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> T {
        let ptr = self.slice.as_ptr() as *const T;
        let ptr = ptr.add(index);
        std::ptr::read_unaligned(ptr)
    }

    #[inline]
    fn binary_search(&self, value: &T) -> std::result::Result<usize, usize> {
        let mut size = self.slice.len() / std::mem::size_of::<T>();
        if size == 0 {
            return Err(0);
        }
        let mut base = 0usize;

        // Note that the count of ids is not greater than `u16::MAX`. The number
        // of binary search steps will not over 16 unless the data is corrupted.
        // Let's relex to 20.
        let mut steps = 20usize;

        while steps > 0 && size > 1 {
            let half = size / 2;
            let mid = base + half;
            let cmp = unsafe { self.get_unchecked(mid) }.cmp(value);
            base = if cmp == Greater { base } else { mid };
            size -= half;
            steps -= 1;
        }

        let cmp = unsafe { self.get_unchecked(base) }.cmp(value);
        if cmp == Equal {
            Ok(base)
        } else {
            Err(base + (cmp == Less) as usize)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::encoder_for_test::{Column, RowEncoder};
    use super::{read_le_bytes, RowSlice};
    use crate::codec::data_type::ScalarValue;
    use crate::expr::EvalContext;
    use codec::prelude::NumberEncoder;
    use std::u16;

    #[test]
    fn test_read_le_bytes() {
        let data = vec![1, 128, 512, u16::MAX, 256];
        let mut buf = vec![];
        for n in &data {
            buf.write_u16_le(*n).unwrap();
        }

        for i in 1..=data.len() {
            let le_bytes = read_le_bytes::<u16>(&mut buf.as_slice(), i).unwrap();
            for j in 0..i {
                assert_eq!(unsafe { le_bytes.get_unchecked(j) }, data[j]);
            }
        }
    }

    fn encoded_data_big() -> Vec<u8> {
        let cols = vec![
            Column::new(1, 1000),
            Column::new(356, 2),
            Column::new(33, ScalarValue::Int(None)), //0x21
            Column::new(3, 3),
            Column::new(64123, 5),
        ];
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    fn encoded_data() -> Vec<u8> {
        let cols = vec![
            Column::new(1, 1000),
            Column::new(33, ScalarValue::Int(None)),
            Column::new(3, 3),
        ];
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    #[test]
    fn test_search_in_non_null_ids() {
        let data = encoded_data_big();
        let big_row = RowSlice::from_bytes(&data).unwrap();
        assert!(big_row.is_big());
        assert_eq!(big_row.search_in_non_null_ids(33).unwrap(), None);
        assert_eq!(big_row.search_in_non_null_ids(333).unwrap(), None);
        assert_eq!(
            big_row
                .search_in_non_null_ids(i64::from(u32::max_value()) + 2)
                .unwrap(),
            None
        );
        assert_eq!(Some((0, 2)), big_row.search_in_non_null_ids(1).unwrap());
        assert_eq!(Some((3, 4)), big_row.search_in_non_null_ids(356).unwrap());
        assert_eq!(Some((4, 5)), big_row.search_in_non_null_ids(64123).unwrap());
        assert_eq!(None, big_row.search_in_non_null_ids(64124).unwrap());

        let data = encoded_data();
        let row = RowSlice::from_bytes(&data).unwrap();
        assert!(!row.is_big());
        assert_eq!(row.search_in_non_null_ids(33).unwrap(), None);
        assert_eq!(row.search_in_non_null_ids(35).unwrap(), None);
        assert_eq!(
            row.search_in_non_null_ids(i64::from(u8::max_value()) + 2)
                .unwrap(),
            None
        );
        assert_eq!(Some((0, 2)), row.search_in_non_null_ids(1).unwrap());
        assert_eq!(Some((2, 3)), row.search_in_non_null_ids(3).unwrap());
    }

    #[test]
    fn test_search_in_null_ids() {
        let data = encoded_data_big();
        let row = RowSlice::from_bytes(&data).unwrap();
        assert!(row.search_in_null_ids(0x21));
        assert!(!row.search_in_null_ids(3));
        assert!(!row.search_in_null_ids(333));
        assert!(!row.search_in_null_ids(0xCC21));
        assert!(!row.search_in_null_ids(0xFF0021));
        assert!(!row.search_in_null_ids(0xFF00000021));

        let data = encoded_data();
        let row = RowSlice::from_bytes(&data).unwrap();
        assert!(row.search_in_null_ids(0x21));
        assert!(!row.search_in_null_ids(3));
        assert!(!row.search_in_null_ids(333));
        assert!(!row.search_in_null_ids(0xCC21));
        assert!(!row.search_in_null_ids(0xFF0021));
        assert!(!row.search_in_null_ids(0xFF00000021));
    }
}

#[cfg(test)]
mod benches {
    use super::super::encoder_for_test::{Column, RowEncoder};
    use super::RowSlice;
    use crate::codec::data_type::ScalarValue;
    use crate::expr::EvalContext;
    use test::black_box;

    fn encoded_data(len: usize) -> Vec<u8> {
        let mut cols = vec![];
        for i in 0..(len as i64) {
            if i % 10 == 0 {
                cols.push(Column::new(i, ScalarValue::Int(None)))
            } else {
                cols.push(Column::new(i, i))
            }
        }
        let mut buf = vec![];
        buf.write_row(&mut EvalContext::default(), cols).unwrap();
        buf
    }

    #[bench]
    fn bench_search_in_non_null_ids(b: &mut test::Bencher) {
        let data = encoded_data(10);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(3))
        });
    }

    #[bench]
    fn bench_search_in_non_null_ids_middle(b: &mut test::Bencher) {
        let data = encoded_data(100);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(89))
        });
    }

    #[bench]
    fn bench_search_in_null_ids_middle(b: &mut test::Bencher) {
        let data = encoded_data(100);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(20))
        });
    }

    #[bench]
    fn bench_search_in_non_null_ids_big(b: &mut test::Bencher) {
        let data = encoded_data(350);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(row.search_in_non_null_ids(257))
        });
    }

    #[bench]
    fn bench_from_bytes_big(b: &mut test::Bencher) {
        let data = encoded_data(350);

        b.iter(|| {
            let row = RowSlice::from_bytes(black_box(&data)).unwrap();
            black_box(&row);
        });
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(test)]

extern crate test;
extern crate rand;
extern crate rand_isaac;
extern crate rand_chacha;
extern crate rand_hc;
extern crate rand_pcg;
extern crate rand_xorshift;
extern crate rand_xoshiro;

const RAND_BENCH_N: u64 = 1000;
const BYTES_LEN: usize = 1024;

use std::mem::size_of;
use test::{black_box, Bencher};

use rand::prelude::*;
use rand::rngs::adapter::ReseedingRng;
use rand::rngs::OsRng;
use rand_isaac::{IsaacRng, Isaac64Rng};
use rand_chacha::ChaChaRng;
use rand_hc::{Hc128Rng, Hc128Core};
use rand_pcg::{Lcg64Xsh32, Mcg128Xsl64};
use rand_xorshift::XorShiftRng;
use rand_xoshiro::{Xoshiro256StarStar, Xoshiro256Plus, Xoshiro128StarStar,
    Xoshiro128Plus, Xoroshiro128StarStar, Xoroshiro128Plus, SplitMix64,
    Xoroshiro64StarStar, Xoroshiro64Star};

macro_rules! gen_bytes {
    ($fnn:ident, $gen:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = $gen;
            let mut buf = [0u8; BYTES_LEN];
            b.iter(|| {
                for _ in 0..RAND_BENCH_N {
                    rng.fill_bytes(&mut buf);
                    black_box(buf);
                }
            });
            b.bytes = BYTES_LEN as u64 * RAND_BENCH_N;
        }
    }
}

gen_bytes!(gen_bytes_xorshift, XorShiftRng::from_entropy());
gen_bytes!(gen_bytes_xoshiro256starstar, Xoshiro256StarStar::from_entropy());
gen_bytes!(gen_bytes_xoshiro256plus, Xoshiro256Plus::from_entropy());
gen_bytes!(gen_bytes_xoshiro128starstar, Xoshiro128StarStar::from_entropy());
gen_bytes!(gen_bytes_xoshiro128plus, Xoshiro128Plus::from_entropy());
gen_bytes!(gen_bytes_xoroshiro128starstar, Xoroshiro128StarStar::from_entropy());
gen_bytes!(gen_bytes_xoroshiro128plus, Xoroshiro128Plus::from_entropy());
gen_bytes!(gen_bytes_xoroshiro64starstar, Xoroshiro64StarStar::from_entropy());
gen_bytes!(gen_bytes_xoroshiro64star, Xoroshiro64Star::from_entropy());
gen_bytes!(gen_bytes_splitmix64, SplitMix64::from_entropy());
gen_bytes!(gen_bytes_lcg64_xsh32, Lcg64Xsh32::from_entropy());
gen_bytes!(gen_bytes_mcg128_xsh64, Mcg128Xsl64::from_entropy());
gen_bytes!(gen_bytes_chacha20, ChaChaRng::from_entropy());
gen_bytes!(gen_bytes_hc128, Hc128Rng::from_entropy());
gen_bytes!(gen_bytes_isaac, IsaacRng::from_entropy());
gen_bytes!(gen_bytes_isaac64, Isaac64Rng::from_entropy());
gen_bytes!(gen_bytes_std, StdRng::from_entropy());
gen_bytes!(gen_bytes_small, SmallRng::from_entropy());
gen_bytes!(gen_bytes_os, OsRng);

macro_rules! gen_uint {
    ($fnn:ident, $ty:ty, $gen:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = $gen;
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

gen_uint!(gen_u32_xorshift, u32, XorShiftRng::from_entropy());
gen_uint!(gen_u32_xoshiro256starstar, u32, Xoshiro256StarStar::from_entropy());
gen_uint!(gen_u32_xoshiro256plus, u32, Xoshiro256Plus::from_entropy());
gen_uint!(gen_u32_xoshiro128starstar, u32, Xoshiro128StarStar::from_entropy());
gen_uint!(gen_u32_xoshiro128plus, u32, Xoshiro128Plus::from_entropy());
gen_uint!(gen_u32_xoroshiro128starstar, u32, Xoroshiro128StarStar::from_entropy());
gen_uint!(gen_u32_xoroshiro128plus, u32, Xoroshiro128Plus::from_entropy());
gen_uint!(gen_u32_xoroshiro64starstar, u32, Xoroshiro64StarStar::from_entropy());
gen_uint!(gen_u32_xoroshiro64star, u32, Xoroshiro64Star::from_entropy());
gen_uint!(gen_u32_splitmix64, u32, SplitMix64::from_entropy());
gen_uint!(gen_u32_lcg64_xsh32, u32, Lcg64Xsh32::from_entropy());
gen_uint!(gen_u32_mcg128_xsh64, u32, Mcg128Xsl64::from_entropy());
gen_uint!(gen_u32_chacha20, u32, ChaChaRng::from_entropy());
gen_uint!(gen_u32_hc128, u32, Hc128Rng::from_entropy());
gen_uint!(gen_u32_isaac, u32, IsaacRng::from_entropy());
gen_uint!(gen_u32_isaac64, u32, Isaac64Rng::from_entropy());
gen_uint!(gen_u32_std, u32, StdRng::from_entropy());
gen_uint!(gen_u32_small, u32, SmallRng::from_entropy());
gen_uint!(gen_u32_os, u32, OsRng);

gen_uint!(gen_u64_xorshift, u64, XorShiftRng::from_entropy());
gen_uint!(gen_u64_xoshiro256starstar, u64, Xoshiro256StarStar::from_entropy());
gen_uint!(gen_u64_xoshiro256plus, u64, Xoshiro256Plus::from_entropy());
gen_uint!(gen_u64_xoshiro128starstar, u64, Xoshiro128StarStar::from_entropy());
gen_uint!(gen_u64_xoshiro128plus, u64, Xoshiro128Plus::from_entropy());
gen_uint!(gen_u64_xoroshiro128starstar, u64, Xoroshiro128StarStar::from_entropy());
gen_uint!(gen_u64_xoroshiro128plus, u64, Xoroshiro128Plus::from_entropy());
gen_uint!(gen_u64_xoroshiro64starstar, u64, Xoroshiro64StarStar::from_entropy());
gen_uint!(gen_u64_xoroshiro64star, u64, Xoroshiro64Star::from_entropy());
gen_uint!(gen_u64_splitmix64, u64, SplitMix64::from_entropy());
gen_uint!(gen_u64_lcg64_xsh32, u64, Lcg64Xsh32::from_entropy());
gen_uint!(gen_u64_mcg128_xsh64, u64, Mcg128Xsl64::from_entropy());
gen_uint!(gen_u64_chacha20, u64, ChaChaRng::from_entropy());
gen_uint!(gen_u64_hc128, u64, Hc128Rng::from_entropy());
gen_uint!(gen_u64_isaac, u64, IsaacRng::from_entropy());
gen_uint!(gen_u64_isaac64, u64, Isaac64Rng::from_entropy());
gen_uint!(gen_u64_std, u64, StdRng::from_entropy());
gen_uint!(gen_u64_small, u64, SmallRng::from_entropy());
gen_uint!(gen_u64_os, u64, OsRng);

macro_rules! init_gen {
    ($fnn:ident, $gen:ident) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = XorShiftRng::from_entropy();
            b.iter(|| {
                let r2 = $gen::from_rng(&mut rng).unwrap();
                r2
            });
        }
    }
}

init_gen!(init_xorshift, XorShiftRng);
init_gen!(init_xoshiro256starstar, Xoshiro256StarStar);
init_gen!(init_xoshiro256plus, Xoshiro256Plus);
init_gen!(init_xoshiro128starstar, Xoshiro128StarStar);
init_gen!(init_xoshiro128plus, Xoshiro128Plus);
init_gen!(init_xoroshiro128starstar, Xoroshiro128StarStar);
init_gen!(init_xoroshiro128plus, Xoroshiro128Plus);
init_gen!(init_xoroshiro64starstar, Xoroshiro64StarStar);
init_gen!(init_xoroshiro64star, Xoroshiro64Star);
init_gen!(init_splitmix64, SplitMix64);
init_gen!(init_lcg64_xsh32, Lcg64Xsh32);
init_gen!(init_mcg128_xsh64, Mcg128Xsl64);
init_gen!(init_hc128, Hc128Rng);
init_gen!(init_isaac, IsaacRng);
init_gen!(init_isaac64, Isaac64Rng);
init_gen!(init_chacha, ChaChaRng);

const RESEEDING_THRESHOLD: u64 = 1024*1024*1024; // something high enough to get
                                                 // deterministic measurements

#[bench]
fn reseeding_hc128_bytes(b: &mut Bencher) {
    let mut rng = ReseedingRng::new(Hc128Core::from_entropy(),
                                    RESEEDING_THRESHOLD,
                                    OsRng);
    let mut buf = [0u8; BYTES_LEN];
    b.iter(|| {
        for _ in 0..RAND_BENCH_N {
            rng.fill_bytes(&mut buf);
            black_box(buf);
        }
    });
    b.bytes = BYTES_LEN as u64 * RAND_BENCH_N;
}

macro_rules! reseeding_uint {
    ($fnn:ident, $ty:ty) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = ReseedingRng::new(Hc128Core::from_entropy(),
                                            RESEEDING_THRESHOLD,
                                            OsRng);
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

reseeding_uint!(reseeding_hc128_u32, u32);
reseeding_uint!(reseeding_hc128_u64, u64);

macro_rules! threadrng_uint {
    ($fnn:ident, $ty:ty) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = thread_rng();
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

threadrng_uint!(thread_rng_u32, u32);
threadrng_uint!(thread_rng_u64, u64);
<commit_msg>Add benchmarks for new chacha variants<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![feature(test)]

extern crate test;
extern crate rand;
extern crate rand_isaac;
extern crate rand_chacha;
extern crate rand_hc;
extern crate rand_pcg;
extern crate rand_xorshift;
extern crate rand_xoshiro;

const RAND_BENCH_N: u64 = 1000;
const BYTES_LEN: usize = 1024;

use std::mem::size_of;
use test::{black_box, Bencher};

use rand::prelude::*;
use rand::rngs::adapter::ReseedingRng;
use rand::rngs::OsRng;
use rand_isaac::{IsaacRng, Isaac64Rng};
use rand_chacha::{ChaCha8Rng, ChaCha12Rng, ChaCha20Rng};
use rand_hc::{Hc128Rng, Hc128Core};
use rand_pcg::{Lcg64Xsh32, Mcg128Xsl64};
use rand_xorshift::XorShiftRng;
use rand_xoshiro::{Xoshiro256StarStar, Xoshiro256Plus, Xoshiro128StarStar,
    Xoshiro128Plus, Xoroshiro128StarStar, Xoroshiro128Plus, SplitMix64,
    Xoroshiro64StarStar, Xoroshiro64Star};

macro_rules! gen_bytes {
    ($fnn:ident, $gen:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = $gen;
            let mut buf = [0u8; BYTES_LEN];
            b.iter(|| {
                for _ in 0..RAND_BENCH_N {
                    rng.fill_bytes(&mut buf);
                    black_box(buf);
                }
            });
            b.bytes = BYTES_LEN as u64 * RAND_BENCH_N;
        }
    }
}

gen_bytes!(gen_bytes_xorshift, XorShiftRng::from_entropy());
gen_bytes!(gen_bytes_xoshiro256starstar, Xoshiro256StarStar::from_entropy());
gen_bytes!(gen_bytes_xoshiro256plus, Xoshiro256Plus::from_entropy());
gen_bytes!(gen_bytes_xoshiro128starstar, Xoshiro128StarStar::from_entropy());
gen_bytes!(gen_bytes_xoshiro128plus, Xoshiro128Plus::from_entropy());
gen_bytes!(gen_bytes_xoroshiro128starstar, Xoroshiro128StarStar::from_entropy());
gen_bytes!(gen_bytes_xoroshiro128plus, Xoroshiro128Plus::from_entropy());
gen_bytes!(gen_bytes_xoroshiro64starstar, Xoroshiro64StarStar::from_entropy());
gen_bytes!(gen_bytes_xoroshiro64star, Xoroshiro64Star::from_entropy());
gen_bytes!(gen_bytes_splitmix64, SplitMix64::from_entropy());
gen_bytes!(gen_bytes_lcg64_xsh32, Lcg64Xsh32::from_entropy());
gen_bytes!(gen_bytes_mcg128_xsh64, Mcg128Xsl64::from_entropy());
gen_bytes!(gen_bytes_chacha8, ChaCha8Rng::from_entropy());
gen_bytes!(gen_bytes_chacha12, ChaCha12Rng::from_entropy());
gen_bytes!(gen_bytes_chacha20, ChaCha20Rng::from_entropy());
gen_bytes!(gen_bytes_hc128, Hc128Rng::from_entropy());
gen_bytes!(gen_bytes_isaac, IsaacRng::from_entropy());
gen_bytes!(gen_bytes_isaac64, Isaac64Rng::from_entropy());
gen_bytes!(gen_bytes_std, StdRng::from_entropy());
gen_bytes!(gen_bytes_small, SmallRng::from_entropy());
gen_bytes!(gen_bytes_os, OsRng);

macro_rules! gen_uint {
    ($fnn:ident, $ty:ty, $gen:expr) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = $gen;
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

gen_uint!(gen_u32_xorshift, u32, XorShiftRng::from_entropy());
gen_uint!(gen_u32_xoshiro256starstar, u32, Xoshiro256StarStar::from_entropy());
gen_uint!(gen_u32_xoshiro256plus, u32, Xoshiro256Plus::from_entropy());
gen_uint!(gen_u32_xoshiro128starstar, u32, Xoshiro128StarStar::from_entropy());
gen_uint!(gen_u32_xoshiro128plus, u32, Xoshiro128Plus::from_entropy());
gen_uint!(gen_u32_xoroshiro128starstar, u32, Xoroshiro128StarStar::from_entropy());
gen_uint!(gen_u32_xoroshiro128plus, u32, Xoroshiro128Plus::from_entropy());
gen_uint!(gen_u32_xoroshiro64starstar, u32, Xoroshiro64StarStar::from_entropy());
gen_uint!(gen_u32_xoroshiro64star, u32, Xoroshiro64Star::from_entropy());
gen_uint!(gen_u32_splitmix64, u32, SplitMix64::from_entropy());
gen_uint!(gen_u32_lcg64_xsh32, u32, Lcg64Xsh32::from_entropy());
gen_uint!(gen_u32_mcg128_xsh64, u32, Mcg128Xsl64::from_entropy());
gen_uint!(gen_u32_chacha8, u32, ChaCha8Rng::from_entropy());
gen_uint!(gen_u32_chacha12, u32, ChaCha12Rng::from_entropy());
gen_uint!(gen_u32_chacha20, u32, ChaCha20Rng::from_entropy());
gen_uint!(gen_u32_hc128, u32, Hc128Rng::from_entropy());
gen_uint!(gen_u32_isaac, u32, IsaacRng::from_entropy());
gen_uint!(gen_u32_isaac64, u32, Isaac64Rng::from_entropy());
gen_uint!(gen_u32_std, u32, StdRng::from_entropy());
gen_uint!(gen_u32_small, u32, SmallRng::from_entropy());
gen_uint!(gen_u32_os, u32, OsRng);

gen_uint!(gen_u64_xorshift, u64, XorShiftRng::from_entropy());
gen_uint!(gen_u64_xoshiro256starstar, u64, Xoshiro256StarStar::from_entropy());
gen_uint!(gen_u64_xoshiro256plus, u64, Xoshiro256Plus::from_entropy());
gen_uint!(gen_u64_xoshiro128starstar, u64, Xoshiro128StarStar::from_entropy());
gen_uint!(gen_u64_xoshiro128plus, u64, Xoshiro128Plus::from_entropy());
gen_uint!(gen_u64_xoroshiro128starstar, u64, Xoroshiro128StarStar::from_entropy());
gen_uint!(gen_u64_xoroshiro128plus, u64, Xoroshiro128Plus::from_entropy());
gen_uint!(gen_u64_xoroshiro64starstar, u64, Xoroshiro64StarStar::from_entropy());
gen_uint!(gen_u64_xoroshiro64star, u64, Xoroshiro64Star::from_entropy());
gen_uint!(gen_u64_splitmix64, u64, SplitMix64::from_entropy());
gen_uint!(gen_u64_lcg64_xsh32, u64, Lcg64Xsh32::from_entropy());
gen_uint!(gen_u64_mcg128_xsh64, u64, Mcg128Xsl64::from_entropy());
gen_uint!(gen_u64_chacha8, u64, ChaCha8Rng::from_entropy());
gen_uint!(gen_u64_chacha12, u64, ChaCha12Rng::from_entropy());
gen_uint!(gen_u64_chacha20, u64, ChaCha20Rng::from_entropy());
gen_uint!(gen_u64_hc128, u64, Hc128Rng::from_entropy());
gen_uint!(gen_u64_isaac, u64, IsaacRng::from_entropy());
gen_uint!(gen_u64_isaac64, u64, Isaac64Rng::from_entropy());
gen_uint!(gen_u64_std, u64, StdRng::from_entropy());
gen_uint!(gen_u64_small, u64, SmallRng::from_entropy());
gen_uint!(gen_u64_os, u64, OsRng);

macro_rules! init_gen {
    ($fnn:ident, $gen:ident) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = XorShiftRng::from_entropy();
            b.iter(|| {
                let r2 = $gen::from_rng(&mut rng).unwrap();
                r2
            });
        }
    }
}

init_gen!(init_xorshift, XorShiftRng);
init_gen!(init_xoshiro256starstar, Xoshiro256StarStar);
init_gen!(init_xoshiro256plus, Xoshiro256Plus);
init_gen!(init_xoshiro128starstar, Xoshiro128StarStar);
init_gen!(init_xoshiro128plus, Xoshiro128Plus);
init_gen!(init_xoroshiro128starstar, Xoroshiro128StarStar);
init_gen!(init_xoroshiro128plus, Xoroshiro128Plus);
init_gen!(init_xoroshiro64starstar, Xoroshiro64StarStar);
init_gen!(init_xoroshiro64star, Xoroshiro64Star);
init_gen!(init_splitmix64, SplitMix64);
init_gen!(init_lcg64_xsh32, Lcg64Xsh32);
init_gen!(init_mcg128_xsh64, Mcg128Xsl64);
init_gen!(init_hc128, Hc128Rng);
init_gen!(init_isaac, IsaacRng);
init_gen!(init_isaac64, Isaac64Rng);
init_gen!(init_chacha, ChaCha20Rng);

const RESEEDING_THRESHOLD: u64 = 1024*1024*1024; // something high enough to get
                                                 // deterministic measurements

#[bench]
fn reseeding_hc128_bytes(b: &mut Bencher) {
    let mut rng = ReseedingRng::new(Hc128Core::from_entropy(),
                                    RESEEDING_THRESHOLD,
                                    OsRng);
    let mut buf = [0u8; BYTES_LEN];
    b.iter(|| {
        for _ in 0..RAND_BENCH_N {
            rng.fill_bytes(&mut buf);
            black_box(buf);
        }
    });
    b.bytes = BYTES_LEN as u64 * RAND_BENCH_N;
}

macro_rules! reseeding_uint {
    ($fnn:ident, $ty:ty) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = ReseedingRng::new(Hc128Core::from_entropy(),
                                            RESEEDING_THRESHOLD,
                                            OsRng);
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

reseeding_uint!(reseeding_hc128_u32, u32);
reseeding_uint!(reseeding_hc128_u64, u64);

macro_rules! threadrng_uint {
    ($fnn:ident, $ty:ty) => {
        #[bench]
        fn $fnn(b: &mut Bencher) {
            let mut rng = thread_rng();
            b.iter(|| {
                let mut accum: $ty = 0;
                for _ in 0..RAND_BENCH_N {
                    accum = accum.wrapping_add(rng.gen::<$ty>());
                }
                accum
            });
            b.bytes = size_of::<$ty>() as u64 * RAND_BENCH_N;
        }
    }
}

threadrng_uint!(thread_rng_u32, u32);
threadrng_uint!(thread_rng_u64, u64);
<|endoftext|>


<commit_before>#![feature(test)]
extern crate asio;
extern crate test;

use asio::IoService;

use std::thread;
use std::sync::Arc;
use std::sync::atomic::*;
use test::Bencher;

#[bench]
fn bench_thrd01_1000(b: &mut Bencher) {
    let io = IoService::new();
    b.iter(|| {
        io.reset();
        fn repeat(io: &IoService, count: usize) {
            if count > 0 {
                io.post(move |io| repeat(io, count-1));
            }
        }
        repeat(&io, 1000);
        io.run();
    })
}

#[bench]
fn bench_thrd10_1000(b: &mut Bencher) {
    static STOP_BENCH: AtomicBool = ATOMIC_BOOL_INIT;
    let io = IoService::new();
    let mut thrd = Vec::new();
    for _ in 0..10 {
        let ios = io.clone();
        thrd.push(thread::spawn(move || while !STOP_BENCH.load(Ordering::Relaxed) {
            ios.run();
            thread::yield_now();
        }));
    }
    b.iter(|| {
        io.work(|io| {
            io.reset();
            let count = Arc::new(AtomicUsize::new(1000));
            for _ in 0..count.load(Ordering::Relaxed) {
                let count = count.clone();
                io.post(move |io| {
                    if count.fetch_sub(1, Ordering::SeqCst) == 1 {
                        io.stop();
                    }
                });
            }
        });
    });

    STOP_BENCH.store(true, Ordering::Relaxed);
    for th in thrd {
        th.join().unwrap();
    }
}
<commit_msg>fix: bench-test<commit_after>#![feature(test)]
extern crate asio;
extern crate test;

use asio::IoService;

use std::thread;
use std::sync::Arc;
use std::sync::atomic::*;
use test::Bencher;

#[bench]
fn bench_thrd01_1000(b: &mut Bencher) {
    let io = IoService::new();
    b.iter(|| {
        io.reset();
        fn repeat(io: &IoService, count: usize) {
            if count > 0 {
                io.post(move |io| repeat(io, count-1));
            }
        }
        repeat(&io, 1000);
        io.run();
    })
}

#[bench]
fn bench_thrd10_1000(b: &mut Bencher) {
    static STOP_BENCH: AtomicBool = ATOMIC_BOOL_INIT;
    let io = IoService::new();
    let mut thrd = Vec::new();
    for _ in 0..10 {
        let ios = io.clone();
        thrd.push(thread::spawn(move || while !STOP_BENCH.load(Ordering::Relaxed) {
            ios.run();
            thread::yield_now();
        }));
    }
    b.iter(|| {
        io.reset();
        io.work(|io| {
            let count = Arc::new(AtomicUsize::new(1000));
            for _ in 0..count.load(Ordering::Relaxed) {
                let count = count.clone();
                io.post(move |io| {
                    if count.fetch_sub(1, Ordering::SeqCst) == 1 {
                        io.stop();
                    }
                });
            }
        });
    });

    STOP_BENCH.store(true, Ordering::Relaxed);
    for th in thrd {
        th.join().unwrap();
    }
}
<|endoftext|>


<commit_before>use parser::{Conds, Expr, Node, Nodes, Target};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
}

impl Generator {

    fn new() -> Generator {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &[u8]) {
        let s = str::from_utf8(s).unwrap();
        if self.locals.contains(s) {
            self.write(&format!("{}", s));
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_compare(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, &val),
            Expr::Compare(op, ref left, ref right) =>
                self.visit_compare(op, &left, &right),
        }
    }

    fn visit_target_single(&mut self, name: &[u8]) -> Vec<String> {
        vec![str::from_utf8(name).unwrap().to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn write_lit(&mut self, s: &[u8]) {
        self.write("writer.write_str(");
        self.write(&format!("{:#?}", str::from_utf8(s).unwrap()));
        self.writeln(").unwrap();");
    }

    fn write_expr(&mut self, s: &Expr) {
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &Conds) {
        for (i, &(ref cond, ref nodes)) in conds.iter().enumerate() {
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.writeln("}");
    }

    fn write_loop(&mut self, var: &Target, iter: &Expr, body: &Nodes) {

        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(&format!("{}", name));
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, name: &str) {
        self.writeln(&format!("self.render_block_{}_into(writer);", name));
    }

    fn handle(&mut self, nodes: &Vec<Node>) {
        for n in nodes {
            match *n {
                Node::Lit(val) => { self.write_lit(val); },
                Node::Expr(ref val) => { self.write_expr(&val); },
                Node::Cond(ref conds) => { self.write_cond(&conds); },
                Node::Loop(ref var, ref iter, ref body) => {
                    self.write_loop(&var, &iter, &body);
                },
                Node::Block(ref name) => { self.write_block(name) },
                Node::Extends(_) | Node::BlockDef(_, _) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn block_methods(&mut self, blocks: &Vec<Node>) {
        for b in blocks {
            if let &Node::BlockDef(name, ref nodes) = b {
                self.writeln("#[allow(unused_variables)]");
                self.writeln(&format!(
                    "fn render_block_{}_into(&self, writer: &mut std::fmt::Write) {{",
                    name));
                self.indent();
                self.handle(&nodes);
                self.dedent();
                self.writeln("}");
            } else {
                panic!("only block definitions allowed");
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn template_impl(&mut self, ast: &syn::DeriveInput, nodes: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn trait_impl(&mut self, ast: &syn::DeriveInput, base: &str, blocks: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();

        self.block_methods(blocks);

        self.dedent();
        self.writeln("}");
    }

    fn trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.writeln("self.render_trait_into(writer);");
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      blocks: &Vec<Node>, nodes: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

        self.block_methods(blocks);

        self.writeln("fn render_trait_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(name, _) => {
                blocks.push(n);
                content.push(Node::Block(name));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if blocks.len() > 0 {
        if let Some(extends) = base {
            if let Expr::StrLit(base_path) = extends {
                gen.trait_impl(ast, &base_path, &blocks);
            }
        } else {
            gen.template_trait(ast, path, &blocks, &content);
            gen.trait_impl(ast, path, &Vec::new());
        }
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<commit_msg>Purge useless format!() calls<commit_after>use parser::{Conds, Expr, Node, Nodes, Target};
use std::str;
use std::collections::HashSet;
use syn;

fn annotations(generics: &syn::Generics) -> String {
    if generics.lifetimes.len() < 1 {
        return String::new();
    }
    let mut res = String::new();
    res.push('<');
    for lt in &generics.lifetimes {
        res.push_str(lt.lifetime.ident.as_ref());
    }
    res.push('>');
    res
}

fn path_as_identifier(s: &str) -> String {
    let mut res = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() {
            res.push(c);
        } else {
            res.push_str(&format!("{:x}", c as u32));
        }
    }
    res
}

struct Generator {
    buf: String,
    indent: u8,
    start: bool,
    locals: HashSet<String>,
}

impl Generator {

    fn new() -> Generator {
        Generator {
            buf: String::new(),
            indent: 0,
            start: true,
            locals: HashSet::new(),
        }
    }

    fn indent(&mut self) {
        self.indent += 1;
    }

    fn dedent(&mut self) {
        self.indent -= 1;
    }

    fn write(&mut self, s: &str) {
        if self.start {
            for _ in 0..(self.indent * 4) {
                self.buf.push(' ');
            }
            self.start = false;
        }
        self.buf.push_str(s);
    }

    fn writeln(&mut self, s: &str) {
        if s.is_empty() {
            return;
        }
        self.write(s);
        self.buf.push('\n');
        self.start = true;
    }

    fn visit_str_lit(&mut self, s: &str) {
        self.write(&format!("\"{}\"", s));
    }

    fn visit_var(&mut self, s: &[u8]) {
        let s = str::from_utf8(s).unwrap();
        if self.locals.contains(s) {
            self.write(&s);
        } else {
            self.write(&format!("self.{}", s));
        }
    }

    fn visit_filter(&mut self, name: &str, val: &Expr) {
        self.write(&format!("askama::filters::{}(&", name));
        self.visit_expr(val);
        self.write(")");
    }

    fn visit_compare(&mut self, op: &str, left: &Expr, right: &Expr) {
        self.visit_expr(left);
        self.write(&format!(" {} ", op));
        self.visit_expr(right);
    }

    fn visit_expr(&mut self, expr: &Expr) {
        match *expr {
            Expr::StrLit(s) => self.visit_str_lit(s),
            Expr::Var(s) => self.visit_var(s),
            Expr::Filter(name, ref val) => self.visit_filter(name, &val),
            Expr::Compare(op, ref left, ref right) =>
                self.visit_compare(op, &left, &right),
        }
    }

    fn visit_target_single(&mut self, name: &[u8]) -> Vec<String> {
        vec![str::from_utf8(name).unwrap().to_string()]
    }

    fn visit_target(&mut self, target: &Target) -> Vec<String> {
        match *target {
            Target::Name(s) => { self.visit_target_single(s) },
        }
    }

    fn write_lit(&mut self, s: &[u8]) {
        self.write("writer.write_str(");
        self.write(&format!("{:#?}", str::from_utf8(s).unwrap()));
        self.writeln(").unwrap();");
    }

    fn write_expr(&mut self, s: &Expr) {
        self.write("writer.write_fmt(format_args!(\"{}\", ");
        self.visit_expr(s);
        self.writeln(")).unwrap();");
    }

    fn write_cond(&mut self, conds: &Conds) {
        for (i, &(ref cond, ref nodes)) in conds.iter().enumerate() {
            match *cond {
                Some(ref expr) => {
                    if i == 0 {
                        self.write("if ");
                    } else {
                        self.write("} else if ");
                    }
                    self.visit_expr(expr);
                },
                None => { self.writeln("} else"); },
            }
            self.writeln(" {");
            self.indent();
            self.handle(nodes);
            self.dedent();
        }
        self.writeln("}");
    }

    fn write_loop(&mut self, var: &Target, iter: &Expr, body: &Nodes) {

        self.write("for ");
        let targets = self.visit_target(var);
        for name in &targets {
            self.locals.insert(name.clone());
            self.write(name);
        }
        self.write(" in &");
        self.visit_expr(iter);
        self.writeln(" {");

        self.indent();
        self.handle(body);
        self.dedent();
        self.writeln("}");
        for name in &targets {
            self.locals.remove(name);
        }
    }

    fn write_block(&mut self, name: &str) {
        self.writeln(&format!("self.render_block_{}_into(writer);", name));
    }

    fn handle(&mut self, nodes: &Vec<Node>) {
        for n in nodes {
            match *n {
                Node::Lit(val) => { self.write_lit(val); },
                Node::Expr(ref val) => { self.write_expr(&val); },
                Node::Cond(ref conds) => { self.write_cond(&conds); },
                Node::Loop(ref var, ref iter, ref body) => {
                    self.write_loop(&var, &iter, &body);
                },
                Node::Block(ref name) => { self.write_block(name) },
                Node::Extends(_) | Node::BlockDef(_, _) => {
                    panic!("no extends or block definition allowed in content");
                },
            }
        }
    }

    fn block_methods(&mut self, blocks: &Vec<Node>) {
        for b in blocks {
            if let &Node::BlockDef(name, ref nodes) = b {
                self.writeln("#[allow(unused_variables)]");
                self.writeln(&format!(
                    "fn render_block_{}_into(&self, writer: &mut std::fmt::Write) {{",
                    name));
                self.indent();
                self.handle(&nodes);
                self.dedent();
                self.writeln("}");
            } else {
                panic!("only block definitions allowed");
            }
        }
    }

    fn path_based_name(&mut self, ast: &syn::DeriveInput, path: &str) {
        let encoded = path_as_identifier(path);
        let original = ast.ident.as_ref();
        let anno = annotations(&ast.generics);
        self.writeln("#[allow(dead_code, non_camel_case_types)]");
        let s = format!("type TemplateFrom{}{} = {}{};",
                        encoded, &anno, original, &anno);
        self.writeln(&s);
    }

    fn template_impl(&mut self, ast: &syn::DeriveInput, nodes: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn trait_impl(&mut self, ast: &syn::DeriveInput, base: &str, blocks: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} TraitFrom{} for {}{} {{",
                              anno, path_as_identifier(base),
                              ast.ident.as_ref(), anno));
        self.indent();

        self.block_methods(blocks);

        self.dedent();
        self.writeln("}");
    }

    fn trait_based_impl(&mut self, ast: &syn::DeriveInput) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("impl{} askama::Template for {}{} {{",
                              anno, ast.ident.as_ref(), anno));
        self.indent();

        self.writeln("fn render_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.writeln("self.render_trait_into(writer);");
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn template_trait(&mut self, ast: &syn::DeriveInput, path: &str,
                      blocks: &Vec<Node>, nodes: &Vec<Node>) {
        let anno = annotations(&ast.generics);
        self.writeln(&format!("trait{} TraitFrom{}{} {{", anno,
                              path_as_identifier(path), anno));
        self.indent();

        self.block_methods(blocks);

        self.writeln("fn render_trait_into(&self, writer: &mut std::fmt::Write) {");
        self.indent();
        self.handle(nodes);
        self.dedent();
        self.writeln("}");

        self.dedent();
        self.writeln("}");
    }

    fn result(self) -> String {
        self.buf
    }

}

pub fn generate(ast: &syn::DeriveInput, path: &str, mut nodes: Vec<Node>) -> String {
    let mut base: Option<Expr> = None;
    let mut blocks = Vec::new();
    let mut content = Vec::new();
    for n in nodes.drain(..) {
        match n {
            Node::Extends(path) => {
                match base {
                    Some(_) => panic!("multiple extend blocks found"),
                    None => { base = Some(path); },
                }
            },
            Node::BlockDef(name, _) => {
                blocks.push(n);
                content.push(Node::Block(name));
            },
            _ => { content.push(n); },
        }
    }

    let mut gen = Generator::new();
    gen.path_based_name(ast, path);
    if blocks.len() > 0 {
        if let Some(extends) = base {
            if let Expr::StrLit(base_path) = extends {
                gen.trait_impl(ast, &base_path, &blocks);
            }
        } else {
            gen.template_trait(ast, path, &blocks, &content);
            gen.trait_impl(ast, path, &Vec::new());
        }
        gen.trait_based_impl(ast);
    } else {
        gen.template_impl(ast, &content);
    }
    gen.result()
}
<|endoftext|>