<commit_before>use cpu::{Cpu};
use std::io::{File, BytesReader, Seek};
use std::io;

pub fn load_bin(file: &mut File, cpu: &mut Cpu) -> () {
	let mut i: uint = 0;
	for byte in file.bytes() {
		let byte = match byte {
			Ok(number) => number as i8,
			Err(e) => panic!("{}", e.desc),
		};
		cpu.memory.write_byte(i, byte);
		i += 1;
	}
}

pub fn load_c16(file: &mut File, cpu: &mut Cpu) -> () {
	let magic_number = match file.read_be_u64() {
		Ok(number) => number,
		Err(e) => panic!("{}", e.desc)
	};
	if magic_number == 0x43483135 {
		//Do verifications (checksum, size check)
		match file.seek(0x10, io::SeekSet){
			Ok(ok) => ok,
			Err(e) => panic!("{}", e.desc)
		};
		load_bin(file, cpu);
	} else {
		match file.seek(0, io::SeekSet){
			Ok(ok) => ok,
			Err(e) => panic!("{}", e.desc)
		};
		load_bin(file, cpu);
	}
}
<commit_msg>Added ROM size checking and fixed small error when loading C16 files<commit_after>use cpu::{Cpu};
use std::io::{File, BytesReader, Seek};
use std::io;

pub fn load_bin(file: &mut File, cpu: &mut Cpu) -> () {
	let mut i: uint = 0;
	for byte in file.bytes() {
		let byte = match byte {
			Ok(number) => number as i8,
			Err(e) => panic!("{}", e.desc),
		};
		cpu.memory.write_byte(i, byte);
		i += 1;
	}
}

pub fn load_c16(file: &mut File, cpu: &mut Cpu) -> () {
	let magic_number = match file.read_be_u64() {
		Ok(number) => number,
		Err(e) => panic!("{}", e.desc)
	};
	if magic_number == 0x43483135 {
		file.read_u8();
		file.read_u8();

		let rom_size: u32 = match file.read_be_u32() {
			Ok(rom) => rom,
			Err(e) => panic!("{}", e.desc)
		};
		let checksum: u32 = match file.read_be_u32(){
			Ok(sum) => sum,
			Err(e) => panic!("{}", e.desc)
		};
		cpu.pc = match file.read_be_i16() {
			Ok(ip) => ip,
			Err(e) => panic!("{}", e.desc)
		};

		check_rom_size(file, rom_size);
		crc32_checksum(file, checksum);

		match file.seek(0x10, io::SeekSet){
			Ok(ok) => ok,
			Err(e) => panic!("{}", e.desc)
		};
		load_bin(file, cpu);
	} else {
		match file.seek(0, io::SeekSet){
			Ok(ok) => ok,
			Err(e) => panic!("{}", e.desc)
		};
		load_bin(file, cpu);
	}
}

fn check_rom_size(file: &mut File, rom_size: u32) -> () {
	let file_size = match file.stat() {
		Ok(file_stat) => file_stat.size - 0x10,
		Err(e) => panic!("{}", e.desc)
	};
	if rom_size as u64 != file_size {
		panic!("The ROM size is not what the header says");
	}
}

fn crc32_checksum(file: &mut File, checksum: u32) -> () {
	()
}

<|endoftext|>


<commit_before>// TODO: COPYRIGHT, USE & AUTHORS
// TODO: RUST DOCS!

/// GA Random Numbers Util
///
/// Wrapper around the rand crate
///
use rand::{random, XorShiftRng, Rng, SeedableRng, thread_rng};
use rand::distributions::range::SampleRange;

const kGASeedSize : usize = 4;

/// GARandomCtx
///
/// Seeded and Stateful Random Number Generator.
/// Uses rand::XorShiftRng for speed purposes.
///
struct GARandomCtx
{
    seed: [u32; kGASeedSize],
    rng:  XorShiftRng,
    name: String,
    seeded: bool,
    values_genarated: u32
}

impl GARandomCtx
{
// Constructors
    pub fn new_unseeded(name: String) -> GARandomCtx
    {
        let std_rng = XorShiftRng::new_unseeded();
        GARandomCtx
        {
            seed: [0 ; kGASeedSize],
            rng: std_rng,
            name: name,
            seeded: false
        }
    }

    pub fn from_seed(seed: [u32; kGASeedSize], name : String) -> GARandomCtx
    {
        let std_rng = SeedableRng::from_seed(seed);
        GARandomCtx
        {
            seed: seed,
            rng:  std_rng,
            name: name,
            seeded: true
        }
    }

// Random Values
    pub fn random_float(&mut self) -> f32
    {
        self.rng.gen::<f32>()
    }

    pub fn random_float_test(&mut self, v: f32) -> bool
    {
        self.random_float() < v
    }

    pub fn random_range<T: PartialOrd + SampleRange>(&mut self, low: T, high: T) -> T
    {
        self.rng.gen_range(low, high)
    }

// Reset State
    pub fn reseed(&mut self, seed: [u32; kGASeedSize])
    {
        self.seed = seed;
        self.rng.reseed(self.seed);
    }

    pub fn reset(&mut self)
    {
        if self.seeded
        {
            self.rng = SeedableRng::from_seed(self.seed);
        }
        else
        {
            self.rng = XorShiftRng::new_unseeded();
        }
    }
}

impl fmt::Debug for GARandomCtx
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        write!(f, "GARandomCtx {{ x: {}, y: {} }}", self.x, self.y)
    }
}

pub fn ga_random_float() -> f32
{
    random()
}

// Return a random number from range [low, high).
pub fn ga_random_range<T: PartialOrd + SampleRange>(low: T, high: T) -> T
{
    thread_rng().gen_range(low, high)
}

pub fn ga_random_float_test(v: f32) -> bool
{
    ga_random_float() < v
}
<commit_msg>ga_random.rs buffed up<commit_after>// Copyright 2016 Revolution Solid & Contributors.
// author(s): sysnett
// rust-monster is licensed under an MIT License.

//! GA Random Numbers Util
//!
//! Wrapper around the rand crate that provides a Seeded
//! and Stateful Random Number Generator.
//!
//! Internally uses rand::XorShiftRng for speed purposes.
//!
//! # Examples
//!
//! GARandomCtx - Basic Use case
//!
//! ```rust
//! extern crate rust_monster;
//! use rust_monster::ga;
//! use rust_monster::ga::ga_random;
//!
//! fn main ()
//! {
//!     let seed : ga_random::GASeed = [1,2,3,4];
//!
//!     let mut ga_ctx = ga_random::GARandomCtx::from_seed(seed, String::from("MyRandomCtx"));
//!
//!     println!("{:?}", ga_ctx.gen::<f32>());
//!
//!     println!("{:?}", ga_ctx);
//! }
//! ```
//!
use rand::{random, thread_rng, Rng, Rand, SeedableRng, XorShiftRng};
use rand::distributions::range::SampleRange;

use std::fmt;

pub type GASeed = [u32; 4];
pub struct GARandomCtx
{
    seed: GASeed,
    rng:  XorShiftRng,
    name: String,
    seeded: bool,
    values_generated: u32
}

impl GARandomCtx
{
// Constructors
    pub fn new_unseeded(name: String) -> GARandomCtx
    {
        let std_rng = XorShiftRng::new_unseeded();
        GARandomCtx
        {
            seed: [0; 4],
            rng: std_rng,
            name: name,
            seeded: false,
            values_generated: 0
        }
    }

    pub fn from_seed(seed: GASeed, name: String) -> GARandomCtx
    {
        let std_rng = SeedableRng::from_seed(seed);
        GARandomCtx
        {
            seed: seed,
            rng:  std_rng,
            name: name,
            seeded: true,
            values_generated: 0
        }
    }

// Random Values - Subset of the RNG Trait
    pub fn gen<T: Rand>(&mut self) -> T where Self: Sized
    {
        self.values_generated += 1;
        self.rng.gen()
    }

    pub fn gen_range<T: PartialOrd + SampleRange>(&mut self, low: T, high: T) -> T
    {
        self.values_generated += 1;
        self.rng.gen_range(low, high)
    }

    pub fn next_u32(&mut self) -> u32 { self.gen::<u32>() }
    pub fn next_u64(&mut self) -> u64 { self.gen::<u64>() }
    pub fn next_f32(&mut self) -> f32 { self.gen::<f32>() }
    pub fn next_f64(&mut self) -> f64 { self.gen::<f64>() }

// Reset State
    pub fn reseed(&mut self, seed: GASeed)
    {
        self.seed = seed;
        self.seeded = true;
        self.reset();
    }

    pub fn reset(&mut self)
    {
        self.values_generated = 0;
        if self.seeded
        {
            self.rng.reseed(self.seed);
        }
        else
        {
            self.rng = XorShiftRng::new_unseeded();
        }
    }
}

impl fmt::Debug for GARandomCtx
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        let seeded_str = if self.seeded
            {
                "Seeded"
            }
            else
            {
                "Not Seeded"
            };

        write!(f, "GARandomCtx {} - {} {{ seed: {:?}, values_generated: {:?} }}",
               self.name,
               seeded_str,
               self.seed,
               self.values_generated)
    }
}

#[cfg(test)]
mod test
{
    extern crate env_logger;
    use super::{GASeed, GARandomCtx};

    #[test]
    fn same_seed()
    {
        let _ = env_logger::init();
        let seed : GASeed = [1,2,3,4];
        let mut ga_ctx = GARandomCtx::from_seed(seed, String::from("TestRandomCtx"));
        let mut ga_ctx_2 = GARandomCtx::from_seed(seed, String::from("TestRandomCtx2"));
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);

        for _ in 0..100
        {
            assert_eq!(ga_ctx.gen::<f64>(), ga_ctx_2.gen::<f64>());
        }
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);

        for _ in 0..100
        {
            assert_eq!(ga_ctx.gen::<u32>(), ga_ctx_2.gen::<u32>());
        }
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);
    }

    #[test]
    fn diff_seed()
    {
        let _ = env_logger::init();
        let seed_1 : GASeed = [1,2,3,4];
        let seed_2 : GASeed = [4,3,2,1];
        let mut ga_ctx = GARandomCtx::from_seed(seed_1, String::from("TestRandomCtx"));
        let mut ga_ctx_2 = GARandomCtx::from_seed(seed_2, String::from("TestRandomCtx2"));
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);

        for _ in 0..100
        {
            assert!(ga_ctx.gen::<f32>() != ga_ctx_2.gen::<f32>());
        }
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);

        for _ in 0..100
        {
            assert!(ga_ctx.gen::<u64>() != ga_ctx_2.gen::<u64>());
        }
        debug!("{:?}", ga_ctx);
        debug!("{:?}", ga_ctx_2);
    }
}

// DEPRECATED

pub fn ga_random_float() -> f32
{
    random()
}

// Return a random number from range [low, high).
pub fn ga_random_range<T: PartialOrd + SampleRange>(low: T, high: T) -> T
{
    thread_rng().gen_range(low, high)
}

pub fn ga_random_float_test(v: f32) -> bool
{
    ga_random_float() < v
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::Hash;
use std::rc::Rc;

struct CacheEntry<V> {
    val: V,
    instant: u64,
}

pub struct LRUCache<K: Eq + Hash, V> {
    capacity: usize,
    clock: u64,
    data: HashMap<Rc<K>, CacheEntry<V>>,
    order: BTreeMap<u64, Rc<K>>,
}

impl<K, V> LRUCache<K, V>
    where K: Eq + Hash
{
    pub fn new(capacity: usize) -> LRUCache<K, V> {
        LRUCache {
            capacity: capacity,
            clock: 0,
            data: HashMap::with_capacity(capacity),
            order: BTreeMap::new(),
        }
    }

    pub fn get(&mut self, key: K) -> Option<&V> {
        let now = self.clock;
        let key = Rc::new(key);
        let prev = match self.data.entry(key.clone()) {
            Entry::Occupied(mut e) => {
                self.clock += 1;
                let e = e.get_mut();
                let prev = e.instant;
                e.instant = now;
                Some(prev)
            }
            _ => None,
        };
        match prev {
            Some(t) => {
                let k = self.order.remove(&t);
                self.order.insert(now, k.unwrap());
                self.data.get(key.as_ref()).as_ref().map(|x| &x.val)
            }
            None => None,
        }
    }

    pub fn insert(&mut self, key: K, val: V) {
        let now = self.clock;
        self.clock += 1;
        let size = self.data.len();
        let key = Rc::new(key);
        let evict = match self.data.entry(key.clone()) {
            Entry::Occupied(mut e) => {
                let e = e.get_mut();
                let prev = e.instant;
                e.instant = now;
                e.val = val;
                self.order.remove(&prev);
                self.order.insert(now, key.clone());
                None
            }
            Entry::Vacant(e) => {
                let entry = CacheEntry {
                    val: val,
                    instant: now,
                };
                e.insert(entry);
                self.order.insert(now, key.clone());
                if size == self.capacity {
                    let oldest = self.order.keys().cloned().next().unwrap();
                    Some(self.order.remove(&oldest).unwrap())
                } else {
                    None
                }
            }
        };
        match evict {
            Some(k) => {
                self.data.remove(k.as_ref());
            }
            None => {}
        };
    }

    pub fn len(&self) -> usize {
        debug_assert!(self.data.len() == self.order.len());
        self.data.len()
    }
}

#[test]
fn lru_cache() {
    let mut cache = LRUCache::new(3);
    assert_eq!(0, cache.clock);
    assert_eq!(0, cache.len());

    cache.insert(1, 2);
    cache.insert(3, 4);
    cache.insert(5, 6);
    assert_eq!(3, cache.clock);
    assert_eq!(3, cache.len());
    assert_eq!(Some(&2), cache.get(1));
    assert_eq!(Some(&4), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(None, cache.get(7));
    assert_eq!(6, cache.clock);

    cache.insert(1, 1);
    cache.insert(3, 3);
    cache.insert(5, 6);
    assert_eq!(9, cache.clock);
    assert_eq!(3, cache.len());
    assert_eq!(Some(&1), cache.get(1));
    assert_eq!(Some(&3), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(None, cache.get(7));

    cache.insert(7, 8);
    assert_eq!(3, cache.len());
    assert_eq!(None, cache.get(1));
    assert_eq!(Some(&3), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(Some(&8), cache.get(7));
    assert_eq!(16, cache.clock);
}
<commit_msg>Add documentation.<commit_after>//! The 'lru' module implements a [least-recently used](
//! https://en.wikipedia.org/wiki/Cache_replacement_policies#Least_Recently_Used_.28LRU.29) cache.

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::Hash;
use std::rc::Rc;

struct CacheEntry<V> {
    // cache value
    val: V,
    // clock instant when entry was most recently accessed
    instant: u64,
}

pub struct LRUCache<K: Eq + Hash, V> {
    // maximum number of elements stored in the cache
    capacity: usize,
    // logical clock that is incremented on each operation
    clock: u64,
    // unordered map that stores (key, value) pairs
    data: HashMap<Rc<K>, CacheEntry<V>>,
    // ordered map sorted by clock instants. Used by eviction algorithm
    order: BTreeMap<u64, Rc<K>>,
}

impl<K, V> LRUCache<K, V>
    where K: Eq + Hash
{
    pub fn new(capacity: usize) -> LRUCache<K, V> {
        LRUCache {
            capacity: capacity,
            clock: 0,
            data: HashMap::with_capacity(capacity),
            order: BTreeMap::new(),
        }
    }

    pub fn get(&mut self, key: K) -> Option<&V> {
        let now = self.clock;
        let key = Rc::new(key);
        let prev = match self.data.entry(key.clone()) {
            // If the (key, value) pair is located,
            // then find the logical time instant associated
            // with that key. The instant must be set to now.
            Entry::Occupied(mut e) => {
                self.clock += 1;
                let e = e.get_mut();
                let prev = e.instant;
                e.instant = now;
                Some(prev)
            }
            _ => None,
        };
        match prev {
            Some(t) => {
                // If the (key, value) pair is located,
                // then delete the association with the old instant
                // and create an association to now.
                let k = self.order.remove(&t);
                self.order.insert(now, k.unwrap());
                self.data.get(key.as_ref()).as_ref().map(|x| &x.val)
            }
            None => None,
        }
    }

    pub fn insert(&mut self, key: K, val: V) {
        let now = self.clock;
        self.clock += 1;
        let size = self.data.len();
        let key = Rc::new(key);
        let evict = match self.data.entry(key.clone()) {
            // If the (key, value) pair is located,
            // then replace the previous value,
            // and update the logical time association
            // with the pair.
            Entry::Occupied(mut e) => {
                let e = e.get_mut();
                let prev = e.instant;
                e.instant = now;
                e.val = val;
                self.order.remove(&prev);
                self.order.insert(now, key.clone());
                None
            }
            // If the (key, value) pair is not located,
            // then insert the new association.
            Entry::Vacant(e) => {
                let evict = {
                    if size == self.capacity {
                        // Evict the oldest entry from the clock instant map
                        let oldest = self.order.keys().cloned().next().unwrap();
                        Some(self.order.remove(&oldest).unwrap())
                    } else {
                        None
                    }
                };
                let entry = CacheEntry {
                    val: val,
                    instant: now,
                };
                e.insert(entry);
                self.order.insert(now, key.clone());
                evict
            }
        };
        match evict {
            // Evict the oldest entry from the data map
            // Moved to end of function because of borrow checker
            Some(k) => {
                self.data.remove(k.as_ref());
            }
            None => {}
        };
    }

    pub fn len(&self) -> usize {
        debug_assert!(self.data.len() == self.order.len());
        self.data.len()
    }
}

#[test]
fn lru_cache() {
    let mut cache = LRUCache::new(3);
    assert_eq!(0, cache.clock);
    assert_eq!(0, cache.len());

    cache.insert(1, 2);
    cache.insert(3, 4);
    cache.insert(5, 6);
    assert_eq!(3, cache.clock);
    assert_eq!(3, cache.len());
    assert_eq!(Some(&2), cache.get(1));
    assert_eq!(Some(&4), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(None, cache.get(7));
    assert_eq!(6, cache.clock);

    cache.insert(1, 1);
    cache.insert(3, 3);
    cache.insert(5, 6);
    assert_eq!(9, cache.clock);
    assert_eq!(3, cache.len());
    assert_eq!(Some(&1), cache.get(1));
    assert_eq!(Some(&3), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(None, cache.get(7));

    cache.insert(7, 8);
    assert_eq!(3, cache.len());
    assert_eq!(None, cache.get(1));
    assert_eq!(Some(&3), cache.get(3));
    assert_eq!(Some(&6), cache.get(5));
    assert_eq!(Some(&8), cache.get(7));
    assert_eq!(16, cache.clock);
}
<|endoftext|>


<commit_before>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError, LuaTable};
use hlua::any::AnyLuaValue;

use rustc_serialize::json::Json;

use std::thread;
use std::fs::{File};
use std::path::Path;
use std::io::Write;
use std::fmt::{Debug, Formatter};
use std::fmt::Result as FmtResult;
use std::borrow::Borrow;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

#[macro_use]
mod funcs;
#[cfg(test)]
mod tests;

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Option<Sender<LuaMessage>>> = Mutex::new(None);

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Represents an identifier for dealing with nested tables.
///
/// To access foo.bar.baz, use vec!["foo", "bar", "baz"].
///
/// To access foo[2], use vec!["foo", 2].
pub type LuaIdent = Vec<String>;

/// Messages sent to the lua thread
#[derive(Debug)]
pub enum LuaQuery {
    /// Pings the lua thread
    Ping,
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,

    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecFile(String),

    /// Get a variable, expecting an AnyLuaValue
    GetValue(LuaIdent),
    /// Invoke a function found at the position,
    /// with the specified arguments.
    Invoke(LuaIdent, Vec<AnyLuaValue>),
    /// Set a value
    SetValue {
        /// The name of the thing to stuff
        name: LuaIdent,
        /// The value to store.
        val: Json
    },
    /// Create a new table
    NewTable(LuaIdent),
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// If the identifier had length 0
    InvalidName,
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>),
    /// Pong response from lua ping
    Pong,
}

/// Struct sent to the lua query
struct LuaMessage {
    reply: Sender<LuaResponse>,
    query: LuaQuery
}

unsafe impl Send for LuaQuery { }
unsafe impl Sync for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaResponse { }
unsafe impl Send for LuaMessage { }
unsafe impl Sync for LuaMessage { }

impl Debug for LuaResponse {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaResponse::InvalidName =>
                write!(f, "LuaResponse::InvalidName"),
            LuaResponse::Variable(ref var) =>
                write!(f, "LuaResponse::Variable({:?})", var),
            LuaResponse::Error(ref err) =>
                write!(f, "LuaResponse::Error({:?})", err),
            LuaResponse::Function(_) =>
                write!(f, "LuaResponse::Function"),
            LuaResponse::Pong =>
                write!(f, "LuaResponse::Pong")
        }
    }
}

impl Debug for LuaMessage {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        write!(f, "LuaMessage({:?})", self.query)
    }
}

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
#[derive(Debug)]
pub enum LuaSendError {
    /// The thread crashed, was shut down, or rebooted.
    ThreadClosed,
    /// The thread has not been initialized yet (maybe not used)
    ThreadUninitialized,
    /// The sender had an issue, most likey because the thread panicked.
    /// Following the `Sender` API, the original value sent is returned.
    Sender(LuaQuery)
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<Receiver<LuaResponse>,LuaSendError> {
    if !thread_running() {
        Err(LuaSendError::ThreadClosed)
    }
    else if let Some(sender) = SENDER.lock().unwrap().clone() {
        let (tx, rx) = channel();
        let message = LuaMessage { reply: tx, query: query };
        match sender.send(message) {
            Ok(_) => Ok(rx),
            Err(e) => Err(LuaSendError::Sender(e.0.query))
        }
    }
    else {
        Err(LuaSendError::ThreadUninitialized)
    }
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaMessage>();
    {
        let mut sender = SENDER.lock().unwrap();
        *sender = Some(query_tx);
    }

    thread::spawn(move || {
        thread_init(query_rx);
    });
    trace!("Created thread. Init finished.");
}

fn thread_init(receiver: Receiver<LuaMessage>) {
    trace!("thread: initializing.");
    let mut lua = Lua::new();
    //unsafe {
    //    hlua_ffi::lua_atpanic(&mut lua.as_mut_lua().0, thread_on_panic);
    //}
    debug!("thread: Loading Lua libraries...");
    lua.openlibs();
    trace!("thread: Loading way-cooler lua extensions...");
    // We should have some good file handling, read files from /usr by default,
    // but for now we're reading directly from the source.
    lua.execute_from_reader::<(), File>(
        File::open("lib/lua/init.lua").unwrap()
    ).unwrap();
    trace!("thread: loading way-cooler libraries...");
    funcs::register_libraries(&mut lua);
    // Only ready after loading libs
    *RUNNING.write().unwrap() = true;
    debug!("thread: entering main loop...");
    thread_main_loop(receiver, &mut lua);
}

fn thread_main_loop(receiver: Receiver<LuaMessage>, lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(message, lua);
            }
        }
    }
}

fn thread_handle_message(request: LuaMessage, lua: &mut Lua) {
    match request.query {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            thread_send(request.reply, LuaResponse::Pong);
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;
            thread_send(request.reply, LuaResponse::Pong);

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    thread_send(request.reply, LuaResponse::Error(error));
                }
                Ok(_) => {
                    trace!("thread: Code executed okay.");
                    thread_send(request.reply, LuaResponse::Pong);
                }
            }
        },

        LuaQuery::ExecFile(name) => {
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(request.reply, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                    thread_send(request.reply, LuaResponse::Pong);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(request.reply, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetValue(varname) => {
            trace!("thread: Received request to get variable {:?}", varname);

            if varname.len() == 0 {
                thread_send(request.reply, LuaResponse::InvalidName);
                return;
            }
            // Table[0] String had to be cloned, it'd be nice if Rust let us
            // borrow out parts of memory
            match lua.get::<AnyLuaValue, _>(varname[0].clone()) {
                Some(mut table) => {
                    let full_table = walk_table(table, &varname[1..]);
                    thread_send(request.reply,
                                LuaResponse::Variable(full_table));
                },
                None => thread_send(request.reply, LuaResponse::Variable(None))
            }
        },

        LuaQuery::SetValue { name: _name, val: _val } => {
            panic!("thread: unimplemented LuaQuery::SetValue!");
        },

        LuaQuery::NewTable(name_list) => {
            /*
            if name_list.len() == 0 {
                thread_send(request.reply, LuaResponse::Error(
                    LuaError::SyntaxError("Name cannot be empty".to_string())));
            }
            else {
                let mut curr_table = lua.empty_array(name_list[0]);
                for name in name_list.into_iter().skip(1) {
                    curr_table = lua.empty_array(name);
                }
                thread_send(request.reply, LuaResponse::Pong);
            }
            */
        },

        LuaQuery::Ping => {
            thread_send(request.reply, LuaResponse::Pong);
        },
        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: Sender<LuaResponse>, response: LuaResponse) {
    match sender.send(response) {
        Err(err) => {
            match err.0 {
                LuaResponse::Pong => {}, // Those are boring
                _ => {
                    warn!("thread: Someone ignored an important Lua response!");
                }
            }
        }
        Ok(_) => {}
    }
}

fn walk_table(mut table: AnyLuaValue, names: &[String]) -> Option<AnyLuaValue> {
    if let Some(name) = names.first() {
        if let AnyLuaValue::LuaArray(mut arr) = table {
            for (mut key, mut val) in arr {
                if let AnyLuaValue::LuaString(key_str) = key {
                    if *key_str == *name {
                        return walk_table(val, &names[1..]);
                    }
                }
            }
            return None;
        }
        else {
            return None;
        }
    }
    else {
        return Some(table); // ???
    }
}
<commit_msg>Clean up walk_table, but set_table types are too difficult<commit_after>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError, LuaTable, PushGuard};
use hlua::any::AnyLuaValue;

use rustc_serialize::json::Json;

use std::thread;
use std::fs::{File};
use std::path::Path;
use std::io::Write;
use std::fmt::{Debug, Formatter};
use std::fmt::Result as FmtResult;
use std::borrow::Borrow;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

#[macro_use]
mod funcs;
#[cfg(test)]
mod tests;

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Option<Sender<LuaMessage>>> = Mutex::new(None);

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Represents an identifier for dealing with nested tables.
///
/// To access foo.bar.baz, use vec!["foo", "bar", "baz"].
///
/// To access foo[2], use vec!["foo", 2].
pub type LuaIdent = Vec<String>;

/// Messages sent to the lua thread
#[derive(Debug)]
pub enum LuaQuery {
    /// Pings the lua thread
    Ping,
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,

    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecFile(String),

    /// Get a variable, expecting an AnyLuaValue
    GetValue(LuaIdent),
    /// Invoke a function found at the position,
    /// with the specified arguments.
    Invoke(LuaIdent, Vec<AnyLuaValue>),
    /// Set a value
    SetValue {
        /// The name of the thing to stuff
        name: LuaIdent,
        /// The value to store.
        val: Json
    },
    /// Create a new table
    NewTable(LuaIdent),
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// If the identifier had length 0
    InvalidName,
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>),
    /// Pong response from lua ping
    Pong,
}

/// Struct sent to the lua query
struct LuaMessage {
    reply: Sender<LuaResponse>,
    query: LuaQuery
}

unsafe impl Send for LuaQuery { }
unsafe impl Sync for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaResponse { }
unsafe impl Send for LuaMessage { }
unsafe impl Sync for LuaMessage { }

impl Debug for LuaResponse {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaResponse::InvalidName =>
                write!(f, "LuaResponse::InvalidName"),
            LuaResponse::Variable(ref var) =>
                write!(f, "LuaResponse::Variable({:?})", var),
            LuaResponse::Error(ref err) =>
                write!(f, "LuaResponse::Error({:?})", err),
            LuaResponse::Function(_) =>
                write!(f, "LuaResponse::Function"),
            LuaResponse::Pong =>
                write!(f, "LuaResponse::Pong")
        }
    }
}

impl Debug for LuaMessage {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        write!(f, "LuaMessage({:?})", self.query)
    }
}

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
#[derive(Debug)]
pub enum LuaSendError {
    /// The thread crashed, was shut down, or rebooted.
    ThreadClosed,
    /// The thread has not been initialized yet (maybe not used)
    ThreadUninitialized,
    /// The sender had an issue, most likey because the thread panicked.
    /// Following the `Sender` API, the original value sent is returned.
    Sender(LuaQuery)
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<Receiver<LuaResponse>,LuaSendError> {
    if !thread_running() {
        Err(LuaSendError::ThreadClosed)
    }
    else if let Some(sender) = SENDER.lock().unwrap().clone() {
        let (tx, rx) = channel();
        let message = LuaMessage { reply: tx, query: query };
        match sender.send(message) {
            Ok(_) => Ok(rx),
            Err(e) => Err(LuaSendError::Sender(e.0.query))
        }
    }
    else {
        Err(LuaSendError::ThreadUninitialized)
    }
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaMessage>();
    {
        let mut sender = SENDER.lock().unwrap();
        *sender = Some(query_tx);
    }

    thread::spawn(move || {
        thread_init(query_rx);
    });
    trace!("Created thread. Init finished.");
}

fn thread_init(receiver: Receiver<LuaMessage>) {
    trace!("thread: initializing.");
    let mut lua = Lua::new();
    //unsafe {
    //    hlua_ffi::lua_atpanic(&mut lua.as_mut_lua().0, thread_on_panic);
    //}
    debug!("thread: Loading Lua libraries...");
    lua.openlibs();
    trace!("thread: Loading way-cooler lua extensions...");
    // We should have some good file handling, read files from /usr by default,
    // but for now we're reading directly from the source.
    lua.execute_from_reader::<(), File>(
        File::open("lib/lua/init.lua").unwrap()
    ).unwrap();
    trace!("thread: loading way-cooler libraries...");
    funcs::register_libraries(&mut lua);
    // Only ready after loading libs
    *RUNNING.write().unwrap() = true;
    debug!("thread: entering main loop...");
    thread_main_loop(receiver, &mut lua);
}

fn thread_main_loop(receiver: Receiver<LuaMessage>, lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(message, lua);
            }
        }
    }
}

fn thread_handle_message(request: LuaMessage, lua: &mut Lua) {
    match request.query {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            thread_send(request.reply, LuaResponse::Pong);
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;
            thread_send(request.reply, LuaResponse::Pong);

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    thread_send(request.reply, LuaResponse::Error(error));
                }
                Ok(_) => {
                    trace!("thread: Code executed okay.");
                    thread_send(request.reply, LuaResponse::Pong);
                }
            }
        },

        LuaQuery::ExecFile(name) => {
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(request.reply, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                    thread_send(request.reply, LuaResponse::Pong);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(request.reply, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetValue(varname) => {
            trace!("thread: Received request to get variable {:?}", varname);

            if varname.len() == 0 {
                thread_send(request.reply, LuaResponse::InvalidName);
                return;
            }
            // Table[0] String had to be cloned, it'd be nice if Rust let us
            // borrow out parts of memory
            match lua.get::<AnyLuaValue, _>(varname[0].clone()) {
                Some(table) => {
                    let full_table = walk_table(table, &varname[1..]);
                    thread_send(request.reply,
                                LuaResponse::Variable(full_table));
                },
                None => thread_send(request.reply, LuaResponse::Variable(None))
            }
        },

        LuaQuery::SetValue { name: _name, val: _val } => {
            panic!("thread: unimplemented LuaQuery::SetValue!");
        },

        LuaQuery::NewTable(name_list) => {
            panic!("thread: unimplemented LuaQuery::NewTable!");
        },

        LuaQuery::Ping => {
            thread_send(request.reply, LuaResponse::Pong);
        },
        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: Sender<LuaResponse>, response: LuaResponse) {
    match sender.send(response) {
        Err(err) => {
            match err.0 {
                LuaResponse::Pong => {}, // Those are boring
                _ => {
                    warn!("thread: Someone dropped an important Lua response!");
                }
            }
        }
        Ok(_) => {}
    }
}

fn walk_table(table: AnyLuaValue, names: &[String]) -> Option<AnyLuaValue> {
    if let Some(name) = names.first() {
        if let AnyLuaValue::LuaArray(arr) = table {
            for (key, val) in arr {
                if let AnyLuaValue::LuaString(key_str) = key {
                    if *key_str == *name {
                        return walk_table(val, &names[1..]);
                    }
                }
            }
            return None;
        }
        else {
            return None;
        }
    }
    else {
        return Some(table); // ???
    }
}

/*
fn set_value<'l>(lua: &'l mut Lua, mut table: PushGuard<LuaTable<Lua>>,
             names: &[String], val: AnyLuaValue) {
    // Should not be reached!
    if names.len() == 0 {
        return;
    }
    // The last name is the name of the value
    else if names.len() == 1 {
        table.set(names[0].clone(), val);
    }
    else {
        let maybe_table = table.get::<_, _>(names[0]);
        match maybe_table {
            Some(new_table) => {
                set_value(&mut lua, new_table, &names[1..], val);
            }
            None => { return; }
        }
    }
}*/
<|endoftext|>


<commit_before>use crate::error::ParseMessageError;
use crate::structs;
use std::fmt;
use std::str::FromStr;

#[derive(Debug)]
pub enum Pw {
    Uid(u64),
    Nam(String),
}

#[derive(Debug)]
pub enum Sp {
    Nam(String),
}

#[derive(Debug)]
pub enum Gr {
    Gid(u64),
    Nam(String),
}

#[derive(Debug)]
pub enum ClientMessage {
    Key { user: String },
    Pam { user: String },
    CleanUp,
    RateLimit,
    SectorGroups,
    Pw(Pw),
    Sp(Sp),
    Gr(Gr),
}

#[allow(unused)]
#[derive(Debug)]
pub enum DaemonMessage {
    Error { message: String },
    Key { keys: String },
    Pam { result: bool },
    CleanUp,
    RateLimit { limit: usize },
    SectorGroups { sectors: Vec<structs::SectorGroup> },
    Pw { login: String, uid: u64, gid: u64 },
    Sp { login: String },
    Gr { sector: structs::SectorGroup },
}

impl fmt::Display for Pw {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Pw::Uid(uid) => write!(f, "uid={}", uid),
            Pw::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for Sp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Sp::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for Gr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Gr::Gid(gid) => write!(f, "gid={}", gid),
            Gr::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for ClientMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ClientMessage::Key { user } => write!(f, "c:key:{}", user),
            ClientMessage::Pam { user } => write!(f, "c:pam:{}", user),
            ClientMessage::CleanUp => write!(f, "c:cleanup"),
            ClientMessage::RateLimit => write!(f, "c:ratelimit"),
            ClientMessage::SectorGroups => write!(f, "c:sectors"),
            ClientMessage::Pw(pw) => write!(f, "c:pw:{}", pw),
            ClientMessage::Sp(sp) => write!(f, "c:sp:{}", sp),
            ClientMessage::Gr(gr) => write!(f, "c:gr:{}", gr),
        }
    }
}

impl fmt::Display for DaemonMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DaemonMessage::Error { message } => write!(f, "d:error:{}", message),
            DaemonMessage::Key { keys } => write!(f, "d:key:{}", keys),
            DaemonMessage::Pam { result } => write!(f, "d:pam:{}", result),
            DaemonMessage::CleanUp => write!(f, "d:cleanup"),
            DaemonMessage::RateLimit { limit } => write!(f, "d:ratelimit:{}", limit),
            DaemonMessage::SectorGroups { sectors } => {
                let ss: Vec<String> = sectors.iter().map(|s| s.to_string()).collect();
                write!(f, "d:sectors:{}", ss.join("\n"))
            }
            DaemonMessage::Pw { login, uid, gid } => write!(f, "d:pw:{}:{}:{}", login, uid, gid),
            DaemonMessage::Sp { login } => write!(f, "d:sp:{}", login),
            DaemonMessage::Gr { sector } => write!(f, "d:gr:{}", sector),
        }
    }
}

impl FromStr for Pw {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("uid=") {
            Ok(Pw::Uid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Pw::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Sp {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("name=") {
            Ok(Sp::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Gr {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("gid=") {
            Ok(Gr::Gid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Gr::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for ClientMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("c:key:") {
            Ok(ClientMessage::Key { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("c:pam:") {
            Ok(ClientMessage::Pam { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s == "c:cleanup" {
            Ok(ClientMessage::CleanUp)
        } else if s == "c:ratelimit" {
            Ok(ClientMessage::RateLimit)
        } else if s == "c:sectors" {
            Ok(ClientMessage::SectorGroups)
        } else if s.starts_with("c:pw:") {
            Ok(ClientMessage::Pw(s.get(5..).unwrap_or_default().parse::<Pw>()?))
        } else if s.starts_with("c:sp:") {
            Ok(ClientMessage::Sp(s.get(5..).unwrap_or_default().parse::<Sp>()?))
        } else if s.starts_with("c:gr:") {
            Ok(ClientMessage::Gr(s.get(5..).unwrap_or_default().parse::<Gr>()?))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for DaemonMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("d:key:") {
            Ok(DaemonMessage::Key { keys: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("d:pam:") {
            Ok(DaemonMessage::Pam { result: FromStr::from_str(s.get(6..).unwrap_or("false")).unwrap_or(false) })
        } else if s == "d:cleanup" {
            Ok(DaemonMessage::CleanUp)
        } else if s.starts_with("d:ratelimit:") {
            Ok(DaemonMessage::RateLimit { limit: s.get(12..).unwrap_or("0").parse::<usize>().unwrap_or(0) })
        } else if s.starts_with("d:sectors:") {
            let sectors = s.get(10..)
                           .unwrap_or_default()
                           .split('\n')
                           .filter_map(|l| l.parse::<structs::SectorGroup>().ok())
                           .collect();
            Ok(DaemonMessage::SectorGroups { sectors })
        } else if s.starts_with("d:pw:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(":")
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 3 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            let login: String = fields[0].clone();
            match (fields[1].parse::<u64>(), fields[2].parse::<u64>()) {
                (Ok(uid), Ok(gid)) => Ok(DaemonMessage::Pw { login, uid, gid }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:sp:") {
            match s.get(5..).unwrap_or_default().parse::<String>() {
                Ok(login) => Ok(DaemonMessage::Sp { login }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:gr:") {
            match s.get(5..).unwrap_or_default().parse::<structs::SectorGroup>() {
                Ok(sector) => Ok(DaemonMessage::Gr { sector }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else {
            Err(ParseMessageError::ParseDaemonMessageError)
        }
    }
}
<commit_msg>Add a struct for ent<commit_after>use crate::error::ParseMessageError;
use crate::structs;
use std::fmt;
use std::str::FromStr;

#[derive(Debug)]
pub enum Pw {
    Uid(u64),
    Nam(String),
}

#[derive(Debug)]
pub enum Sp {
    Nam(String),
}

#[derive(Debug)]
pub enum Gr {
    Gid(u64),
    Nam(String),
}

#[derive(Debug)]
pub enum Ent {
    Set(u32),
    Get(u32),
    End(u32),
}

#[derive(Debug)]
pub enum ClientMessage {
    Key { user: String },
    Pam { user: String },
    CleanUp,
    RateLimit,
    SectorGroups,
    Pw(Pw),
    Sp(Sp),
    Gr(Gr),
}

#[allow(unused)]
#[derive(Debug)]
pub enum DaemonMessage {
    Error { message: String },
    Key { keys: String },
    Pam { result: bool },
    CleanUp,
    RateLimit { limit: usize },
    SectorGroups { sectors: Vec<structs::SectorGroup> },
    Pw { login: String, uid: u64, gid: u64 },
    Sp { login: String },
    Gr { sector: structs::SectorGroup },
}

impl fmt::Display for Ent {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Ent::Set(pid) => write!(f, "set|{}", pid),
            Ent::Get(pid) => write!(f, "get|{}", pid),
            Ent::End(pid) => write!(f, "end|{}", pid),
        }
    }
}

impl fmt::Display for Pw {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Pw::Uid(uid) => write!(f, "uid={}", uid),
            Pw::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for Sp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Sp::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for Gr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Gr::Gid(gid) => write!(f, "gid={}", gid),
            Gr::Nam(name) => write!(f, "name={}", name),
        }
    }
}

impl fmt::Display for ClientMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ClientMessage::Key { user } => write!(f, "c:key:{}", user),
            ClientMessage::Pam { user } => write!(f, "c:pam:{}", user),
            ClientMessage::CleanUp => write!(f, "c:cleanup"),
            ClientMessage::RateLimit => write!(f, "c:ratelimit"),
            ClientMessage::SectorGroups => write!(f, "c:sectors"),
            ClientMessage::Pw(pw) => write!(f, "c:pw:{}", pw),
            ClientMessage::Sp(sp) => write!(f, "c:sp:{}", sp),
            ClientMessage::Gr(gr) => write!(f, "c:gr:{}", gr),
        }
    }
}

impl fmt::Display for DaemonMessage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DaemonMessage::Error { message } => write!(f, "d:error:{}", message),
            DaemonMessage::Key { keys } => write!(f, "d:key:{}", keys),
            DaemonMessage::Pam { result } => write!(f, "d:pam:{}", result),
            DaemonMessage::CleanUp => write!(f, "d:cleanup"),
            DaemonMessage::RateLimit { limit } => write!(f, "d:ratelimit:{}", limit),
            DaemonMessage::SectorGroups { sectors } => {
                let ss: Vec<String> = sectors.iter().map(|s| s.to_string()).collect();
                write!(f, "d:sectors:{}", ss.join("\n"))
            }
            DaemonMessage::Pw { login, uid, gid } => write!(f, "d:pw:{}:{}:{}", login, uid, gid),
            DaemonMessage::Sp { login } => write!(f, "d:sp:{}", login),
            DaemonMessage::Gr { sector } => write!(f, "d:gr:{}", sector),
        }
    }
}

impl FromStr for Ent {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("set|") {
            Ok(Ent::Set(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("get|") {
            Ok(Ent::Get(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else if s.starts_with("end|") {
            Ok(Ent::End(s.get(4..).unwrap_or_default().parse::<u32>().unwrap()))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Pw {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("uid=") {
            Ok(Pw::Uid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Pw::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Sp {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("name=") {
            Ok(Sp::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for Gr {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("gid=") {
            Ok(Gr::Gid(s.get(4..).unwrap_or_default().parse::<u64>().unwrap()))
        } else if s.starts_with("name=") {
            Ok(Gr::Nam(String::from(s.get(5..).unwrap_or_default())))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for ClientMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("c:key:") {
            Ok(ClientMessage::Key { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("c:pam:") {
            Ok(ClientMessage::Pam { user: String::from(s.get(6..).unwrap_or_default()) })
        } else if s == "c:cleanup" {
            Ok(ClientMessage::CleanUp)
        } else if s == "c:ratelimit" {
            Ok(ClientMessage::RateLimit)
        } else if s == "c:sectors" {
            Ok(ClientMessage::SectorGroups)
        } else if s.starts_with("c:pw:") {
            Ok(ClientMessage::Pw(s.get(5..).unwrap_or_default().parse::<Pw>()?))
        } else if s.starts_with("c:sp:") {
            Ok(ClientMessage::Sp(s.get(5..).unwrap_or_default().parse::<Sp>()?))
        } else if s.starts_with("c:gr:") {
            Ok(ClientMessage::Gr(s.get(5..).unwrap_or_default().parse::<Gr>()?))
        } else {
            Err(ParseMessageError::ParseClientMessageError)
        }
    }
}

impl FromStr for DaemonMessage {
    type Err = ParseMessageError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.starts_with("d:key:") {
            Ok(DaemonMessage::Key { keys: String::from(s.get(6..).unwrap_or_default()) })
        } else if s.starts_with("d:pam:") {
            Ok(DaemonMessage::Pam { result: FromStr::from_str(s.get(6..).unwrap_or("false")).unwrap_or(false) })
        } else if s == "d:cleanup" {
            Ok(DaemonMessage::CleanUp)
        } else if s.starts_with("d:ratelimit:") {
            Ok(DaemonMessage::RateLimit { limit: s.get(12..).unwrap_or("0").parse::<usize>().unwrap_or(0) })
        } else if s.starts_with("d:sectors:") {
            let sectors = s.get(10..)
                           .unwrap_or_default()
                           .split('\n')
                           .filter_map(|l| l.parse::<structs::SectorGroup>().ok())
                           .collect();
            Ok(DaemonMessage::SectorGroups { sectors })
        } else if s.starts_with("d:pw:") {
            let fields: Vec<String> = s.get(5..)
                                       .unwrap_or_default()
                                       .split(":")
                                       .map(|s| s.to_string())
                                       .collect();
            if fields.len() < 3 {
                return Err(ParseMessageError::ParseDaemonMessageError);
            }
            let login: String = fields[0].clone();
            match (fields[1].parse::<u64>(), fields[2].parse::<u64>()) {
                (Ok(uid), Ok(gid)) => Ok(DaemonMessage::Pw { login, uid, gid }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:sp:") {
            match s.get(5..).unwrap_or_default().parse::<String>() {
                Ok(login) => Ok(DaemonMessage::Sp { login }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else if s.starts_with("d:gr:") {
            match s.get(5..).unwrap_or_default().parse::<structs::SectorGroup>() {
                Ok(sector) => Ok(DaemonMessage::Gr { sector }),
                _ => Err(ParseMessageError::ParseDaemonMessageError),
            }
        } else {
            Err(ParseMessageError::ParseDaemonMessageError)
        }
    }
}
<|endoftext|>


<commit_before>use std::io::BufRead;

pub fn loc<R: BufRead>(input: R) -> usize {
    input.lines()
        // TODO: error handling
        .map(|result| result.unwrap())
        .map(|line| line.trim().to_string())
        .filter(|trimmed_line| trimmed_line.len() != 0)
        .filter(|trimmed_line| !trimmed_line.starts_with("//"))
        .count()
}

mod test {
    #[test]
    fn counts_lines() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\ncounts\r\ncounts");
        let loc = super::loc(cursor);
        assert_eq!(loc, 3);
    }

    #[test]
    fn ignores_empty_lines() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\n \r\ncounts");
        let loc = super::loc(cursor);
        assert_eq!(loc, 2);
    }

    #[test]
    fn ignores_single_line_comments() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\ncounts//\n//doesnt");
        let loc = super::loc(cursor);
        assert_eq!(loc, 2);
    }
}
<commit_msg>added missing compiler switch for test module<commit_after>use std::io::BufRead;

pub fn loc<R: BufRead>(input: R) -> usize {
    input.lines()
        // TODO: error handling
        .map(|result| result.unwrap())
        .map(|line| line.trim().to_string())
        .filter(|trimmed_line| trimmed_line.len() != 0)
        .filter(|trimmed_line| !trimmed_line.starts_with("//"))
        .count()
}

#[cfg(test)]
mod test {
    #[test]
    fn counts_lines() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\ncounts\r\ncounts");
        let loc = super::loc(cursor);
        assert_eq!(loc, 3);
    }

    #[test]
    fn ignores_empty_lines() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\n \r\ncounts");
        let loc = super::loc(cursor);
        assert_eq!(loc, 2);
    }

    #[test]
    fn ignores_single_line_comments() {
        use std::io::Cursor;
        let cursor = Cursor::new(b"counts\ncounts//\n//doesnt");
        let loc = super::loc(cursor);
        assert_eq!(loc, 2);
    }
}
<|endoftext|>


<commit_before>#![allow(unused_results)]

use std::{
    sync::atomic::AtomicUsize,
    time::{Duration, Instant},
};

#[cfg(feature = "no_metrics")]
use std::marker::PhantomData;

#[cfg(not(feature = "no_metrics"))]
use std::sync::atomic::Ordering::{Acquire, Relaxed};

use crate::Lazy;

use super::*;

/// A metric collector for all pagecache users running in this
/// process.
pub static M: Lazy<Metrics, fn() -> Metrics> = Lazy::new(Metrics::default);

pub(crate) fn clock() -> f64 {
    if cfg!(feature = "no_metrics") {
        0.
    } else {
        let u = uptime();
        (u.as_secs() * 1_000_000_000) as f64 + f64::from(u.subsec_nanos())
    }
}

// not correct, since it starts counting at the first observance...
pub(crate) fn uptime() -> Duration {
    static START: Lazy<Instant, fn() -> Instant> = Lazy::new(Instant::now);

    if cfg!(feature = "no_metrics") {
        Duration::new(0, 0)
    } else {
        START.elapsed()
    }
}

/// Measure the duration of an event, and call `Histogram::measure()`.
pub struct Measure<'h> {
    _start: f64,
    #[cfg(not(feature = "no_metrics"))]
    histo: &'h Histogram,
    #[cfg(feature = "no_metrics")]
    _pd: PhantomData<&'h ()>,
}

impl<'h> Measure<'h> {
    /// The time delta from ctor to dtor is recorded in `histo`.
    #[inline(always)]
    pub fn new(_histo: &'h Histogram) -> Measure<'h> {
        Measure {
            #[cfg(feature = "no_metrics")]
            _pd: PhantomData,
            #[cfg(not(feature = "no_metrics"))]
            histo: _histo,
            _start: clock(),
        }
    }
}

impl<'h> Drop for Measure<'h> {
    #[inline(always)]
    fn drop(&mut self) {
        #[cfg(not(feature = "no_metrics"))]
        self.histo.measure(clock() - self._start);
    }
}

/// Measure the time spent on calling a given function in a given `Histogram`.
#[cfg_attr(not(feature = "no_inline"), inline)]
pub(crate) fn measure<F: FnOnce() -> R, R>(_histo: &Histogram, f: F) -> R {
    #[cfg(not(feature = "no_metrics"))]
    let _measure = Measure::new(_histo);
    f()
}

#[derive(Default, Debug)]
pub struct Metrics {
    pub advance_snapshot: Histogram,
    pub tree_set: Histogram,
    pub tree_get: Histogram,
    pub tree_del: Histogram,
    pub tree_cas: Histogram,
    pub tree_scan: Histogram,
    pub tree_reverse_scan: Histogram,
    pub tree_merge: Histogram,
    pub tree_start: Histogram,
    pub tree_traverse: Histogram,
    pub tree_child_split_attempt: CachePadded<AtomicUsize>,
    pub tree_child_split_success: CachePadded<AtomicUsize>,
    pub tree_parent_split_attempt: CachePadded<AtomicUsize>,
    pub tree_parent_split_success: CachePadded<AtomicUsize>,
    pub tree_root_split_attempt: CachePadded<AtomicUsize>,
    pub tree_root_split_success: CachePadded<AtomicUsize>,
    pub get_page: Histogram,
    pub rewrite_page: Histogram,
    pub replace_page: Histogram,
    pub link_page: Histogram,
    pub merge_page: Histogram,
    pub page_out: Histogram,
    pub pull: Histogram,
    pub serialize: Histogram,
    pub deserialize: Histogram,
    pub compress: Histogram,
    pub decompress: Histogram,
    pub make_stable: Histogram,
    pub assign_offset: Histogram,
    pub assign_spinloop: Histogram,
    pub reserve_lat: Histogram,
    pub reserve_sz: Histogram,
    pub reserve_current_condvar_wait: Histogram,
    pub reserve_written_condvar_wait: Histogram,
    pub write_to_log: Histogram,
    pub written_bytes: Histogram,
    pub read: Histogram,
    pub tree_loops: CachePadded<AtomicUsize>,
    pub log_reservations: CachePadded<AtomicUsize>,
    pub log_reservation_attempts: CachePadded<AtomicUsize>,
    pub accountant_lock: Histogram,
    pub accountant_hold: Histogram,
    pub accountant_next: Histogram,
    pub accountant_mark_link: Histogram,
    pub accountant_mark_replace: Histogram,
    pub accountant_bump_tip: Histogram,
    #[cfg(feature = "measure_allocs")]
    pub allocations: CachePadded<AtomicUsize>,
    #[cfg(feature = "measure_allocs")]
    pub allocated_bytes: CachePadded<AtomicUsize>,
}

#[cfg(not(feature = "no_metrics"))]
impl Metrics {
    #[inline]
    pub fn tree_looped(&self) {
        self.tree_loops.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn log_reservation_attempted(&self) {
        self.log_reservation_attempts.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn log_reservation_success(&self) {
        self.log_reservations.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_child_split_attempt(&self) {
        self.tree_child_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_child_split_success(&self) {
        self.tree_child_split_success.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_parent_split_attempt(&self) {
        self.tree_parent_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_parent_split_success(&self) {
        self.tree_parent_split_success.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_root_split_attempt(&self) {
        self.tree_root_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_root_split_success(&self) {
        self.tree_root_split_success.fetch_add(1, Relaxed);
    }

    pub fn print_profile(&self) {
        println!(
            "pagecache profile:\n\
            {0: >17} | {1: >10} | {2: >10} | {3: >10} | {4: >10} | {5: >10} | {6: >10} | {7: >10} | {8: >10} | {9: >10}",
            "op",
            "min (us)",
            "med (us)",
            "90 (us)",
            "99 (us)",
            "99.9 (us)",
            "99.99 (us)",
            "max (us)",
            "count",
            "sum (s)"
        );
        println!("{}", std::iter::repeat("-").take(134).collect::<String>());

        let p = |mut tuples: Vec<(String, _, _, _, _, _, _, _, _, _)>| {
            tuples.sort_by_key(|t| (t.9 * -1. * 1e3) as i64);
            for v in tuples {
                println!(
                    "{0: >17} | {1: >10.1} | {2: >10.1} | {3: >10.1} \
                     | {4: >10.1} | {5: >10.1} | {6: >10.1} | {7: >10.1} \
                     | {8: >10.1} | {9: >10.3}",
                    v.0, v.1, v.2, v.3, v.4, v.5, v.6, v.7, v.8, v.9,
                );
            }
        };

        let lat = |name: &str, histo: &Histogram| {
            (
                name.to_string(),
                histo.percentile(0.) / 1e3,
                histo.percentile(50.) / 1e3,
                histo.percentile(90.) / 1e3,
                histo.percentile(99.) / 1e3,
                histo.percentile(99.9) / 1e3,
                histo.percentile(99.99) / 1e3,
                histo.percentile(100.) / 1e3,
                histo.count(),
                histo.sum() as f64 / 1e9,
            )
        };

        let sz = |name: &str, histo: &Histogram| {
            (
                name.to_string(),
                histo.percentile(0.),
                histo.percentile(50.),
                histo.percentile(90.),
                histo.percentile(99.),
                histo.percentile(99.9),
                histo.percentile(99.99),
                histo.percentile(100.),
                histo.count(),
                histo.sum() as f64,
            )
        };

        println!("tree:");
        p(vec![
            lat("start", &self.tree_start),
            lat("traverse", &self.tree_traverse),
            lat("get", &self.tree_get),
            lat("set", &self.tree_set),
            lat("merge", &self.tree_merge),
            lat("del", &self.tree_del),
            lat("cas", &self.tree_cas),
            lat("scan", &self.tree_scan),
            lat("rev scan", &self.tree_reverse_scan),
        ]);
        println!("tree contention loops: {}", self.tree_loops.load(Acquire));
        println!(
            "tree split success rates: child({}/{}) parent({}/{}) root({}/{})",
            self.tree_child_split_success.load(Acquire),
            self.tree_child_split_attempt.load(Acquire),
            self.tree_parent_split_success.load(Acquire),
            self.tree_parent_split_attempt.load(Acquire),
            self.tree_root_split_success.load(Acquire),
            self.tree_root_split_attempt.load(Acquire),
        );

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("pagecache:");
        p(vec![
            lat("snapshot", &self.advance_snapshot),
            lat("get", &self.get_page),
            lat("rewrite", &self.rewrite_page),
            lat("replace", &self.replace_page),
            lat("link", &self.link_page),
            lat("merge", &self.merge_page),
            lat("pull", &self.pull),
            lat("page_out", &self.page_out),
        ]);

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("serialization and compression:");
        p(vec![
            lat("serialize", &self.serialize),
            lat("deserialize", &self.deserialize),
            lat("compress", &self.compress),
            lat("decompress", &self.decompress),
        ]);

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("log:");
        p(vec![
            lat("make_stable", &self.make_stable),
            lat("read", &self.read),
            lat("write", &self.write_to_log),
            sz("written bytes", &self.written_bytes),
            lat("assign offset", &self.assign_offset),
            lat("assign spinloop", &self.assign_spinloop),
            lat("reserve lat", &self.reserve_lat),
            sz("reserve sz", &self.reserve_sz),
            lat("res cvar r", &self.reserve_current_condvar_wait),
            lat("res cvar w", &self.reserve_written_condvar_wait),
        ]);
        println!("log reservations: {}", self.log_reservations.load(Acquire));
        println!(
            "log res attempts: {}",
            self.log_reservation_attempts.load(Acquire)
        );

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("segment accountant:");
        p(vec![
            lat("acquire", &self.accountant_lock),
            lat("hold", &self.accountant_hold),
            lat("next", &self.accountant_next),
            lat("replace", &self.accountant_mark_replace),
            lat("link", &self.accountant_mark_link),
        ]);

        #[cfg(feature = "measure_allocs")]
        {
            println!(
                "{}",
                std::iter::repeat("-").take(134).collect::<String>()
            );
            println!("allocation statistics:");
            println!(
                "total allocations: {}",
                measure_allocs::ALLOCATIONS.load(Acquire)
            );
            println!(
                "allocated bytes: {}",
                measure_allocs::ALLOCATED_BYTES.load(Acquire)
            );
        }
    }
}

#[cfg(feature = "no_metrics")]
impl Metrics {
    pub fn log_reservation_attempted(&self) {}

    pub fn log_reservation_success(&self) {}

    pub fn tree_child_split_attempt(&self) {}

    pub fn tree_child_split_success(&self) {}

    pub fn tree_parent_split_attempt(&self) {}

    pub fn tree_parent_split_success(&self) {}

    pub fn tree_root_split_attempt(&self) {}

    pub fn tree_root_split_success(&self) {}

    pub fn tree_looped(&self) {}

    pub fn print_profile(&self) {}
}
<commit_msg>use rdtscp on x86_64 metrics<commit_after>#![allow(unused_results)]

use std::sync::atomic::AtomicUsize;

#[cfg(not(target_arch = "x86_64"))]
use time::{Duration, Instant};

#[cfg(feature = "no_metrics")]
use std::marker::PhantomData;

#[cfg(not(feature = "no_metrics"))]
use std::sync::atomic::Ordering::{Acquire, Relaxed};

use crate::Lazy;

use super::*;

/// A metric collector for all pagecache users running in this
/// process.
pub static M: Lazy<Metrics, fn() -> Metrics> = Lazy::new(Metrics::default);

pub(crate) fn clock() -> f64 {
    if cfg!(feature = "no_metrics") {
        0.
    } else {
        #[cfg(target_arch = "x86_64")]
        #[allow(unsafe_code)]
        unsafe {
            let mut _aux = 0;
            core::arch::x86_64::__rdtscp(&mut _aux) as f64
        }

        #[cfg(not(target_arch = "x86_64"))]
        {
            let u = uptime();
            (u.as_secs() * 1_000_000_000) as f64 + f64::from(u.subsec_nanos())
        }
    }
}

// not correct, since it starts counting at the first observance...
#[cfg(not(target_arch = "x86_64"))]
pub(crate) fn uptime() -> Duration {
    static START: Lazy<Instant, fn() -> Instant> = Lazy::new(Instant::now);

    if cfg!(feature = "no_metrics") {
        Duration::new(0, 0)
    } else {
        START.elapsed()
    }
}

/// Measure the duration of an event, and call `Histogram::measure()`.
pub struct Measure<'h> {
    _start: f64,
    #[cfg(not(feature = "no_metrics"))]
    histo: &'h Histogram,
    #[cfg(feature = "no_metrics")]
    _pd: PhantomData<&'h ()>,
}

impl<'h> Measure<'h> {
    /// The time delta from ctor to dtor is recorded in `histo`.
    #[inline(always)]
    pub fn new(_histo: &'h Histogram) -> Measure<'h> {
        Measure {
            #[cfg(feature = "no_metrics")]
            _pd: PhantomData,
            #[cfg(not(feature = "no_metrics"))]
            histo: _histo,
            _start: clock(),
        }
    }
}

impl<'h> Drop for Measure<'h> {
    #[inline(always)]
    fn drop(&mut self) {
        #[cfg(not(feature = "no_metrics"))]
        self.histo.measure(clock() - self._start);
    }
}

/// Measure the time spent on calling a given function in a given `Histogram`.
#[cfg_attr(not(feature = "no_inline"), inline)]
pub(crate) fn measure<F: FnOnce() -> R, R>(_histo: &Histogram, f: F) -> R {
    #[cfg(not(feature = "no_metrics"))]
    let _measure = Measure::new(_histo);
    f()
}

#[derive(Default, Debug)]
pub struct Metrics {
    pub advance_snapshot: Histogram,
    pub tree_set: Histogram,
    pub tree_get: Histogram,
    pub tree_del: Histogram,
    pub tree_cas: Histogram,
    pub tree_scan: Histogram,
    pub tree_reverse_scan: Histogram,
    pub tree_merge: Histogram,
    pub tree_start: Histogram,
    pub tree_traverse: Histogram,
    pub tree_child_split_attempt: CachePadded<AtomicUsize>,
    pub tree_child_split_success: CachePadded<AtomicUsize>,
    pub tree_parent_split_attempt: CachePadded<AtomicUsize>,
    pub tree_parent_split_success: CachePadded<AtomicUsize>,
    pub tree_root_split_attempt: CachePadded<AtomicUsize>,
    pub tree_root_split_success: CachePadded<AtomicUsize>,
    pub get_page: Histogram,
    pub rewrite_page: Histogram,
    pub replace_page: Histogram,
    pub link_page: Histogram,
    pub merge_page: Histogram,
    pub page_out: Histogram,
    pub pull: Histogram,
    pub serialize: Histogram,
    pub deserialize: Histogram,
    pub compress: Histogram,
    pub decompress: Histogram,
    pub make_stable: Histogram,
    pub assign_offset: Histogram,
    pub assign_spinloop: Histogram,
    pub reserve_lat: Histogram,
    pub reserve_sz: Histogram,
    pub reserve_current_condvar_wait: Histogram,
    pub reserve_written_condvar_wait: Histogram,
    pub write_to_log: Histogram,
    pub written_bytes: Histogram,
    pub read: Histogram,
    pub tree_loops: CachePadded<AtomicUsize>,
    pub log_reservations: CachePadded<AtomicUsize>,
    pub log_reservation_attempts: CachePadded<AtomicUsize>,
    pub accountant_lock: Histogram,
    pub accountant_hold: Histogram,
    pub accountant_next: Histogram,
    pub accountant_mark_link: Histogram,
    pub accountant_mark_replace: Histogram,
    pub accountant_bump_tip: Histogram,
    #[cfg(feature = "measure_allocs")]
    pub allocations: CachePadded<AtomicUsize>,
    #[cfg(feature = "measure_allocs")]
    pub allocated_bytes: CachePadded<AtomicUsize>,
}

#[cfg(not(feature = "no_metrics"))]
impl Metrics {
    #[inline]
    pub fn tree_looped(&self) {
        self.tree_loops.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn log_reservation_attempted(&self) {
        self.log_reservation_attempts.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn log_reservation_success(&self) {
        self.log_reservations.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_child_split_attempt(&self) {
        self.tree_child_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_child_split_success(&self) {
        self.tree_child_split_success.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_parent_split_attempt(&self) {
        self.tree_parent_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_parent_split_success(&self) {
        self.tree_parent_split_success.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_root_split_attempt(&self) {
        self.tree_root_split_attempt.fetch_add(1, Relaxed);
    }

    #[inline]
    pub fn tree_root_split_success(&self) {
        self.tree_root_split_success.fetch_add(1, Relaxed);
    }

    pub fn print_profile(&self) {
        println!(
            "pagecache profile:\n\
            {0: >17} | {1: >10} | {2: >10} | {3: >10} | {4: >10} | {5: >10} | {6: >10} | {7: >10} | {8: >10} | {9: >10}",
            "op",
            "min (us)",
            "med (us)",
            "90 (us)",
            "99 (us)",
            "99.9 (us)",
            "99.99 (us)",
            "max (us)",
            "count",
            "sum (s)"
        );
        println!("{}", std::iter::repeat("-").take(134).collect::<String>());

        let p = |mut tuples: Vec<(String, _, _, _, _, _, _, _, _, _)>| {
            tuples.sort_by_key(|t| (t.9 * -1. * 1e3) as i64);
            for v in tuples {
                println!(
                    "{0: >17} | {1: >10.1} | {2: >10.1} | {3: >10.1} \
                     | {4: >10.1} | {5: >10.1} | {6: >10.1} | {7: >10.1} \
                     | {8: >10.1} | {9: >10.3}",
                    v.0, v.1, v.2, v.3, v.4, v.5, v.6, v.7, v.8, v.9,
                );
            }
        };

        let lat = |name: &str, histo: &Histogram| {
            (
                name.to_string(),
                histo.percentile(0.) / 1e3,
                histo.percentile(50.) / 1e3,
                histo.percentile(90.) / 1e3,
                histo.percentile(99.) / 1e3,
                histo.percentile(99.9) / 1e3,
                histo.percentile(99.99) / 1e3,
                histo.percentile(100.) / 1e3,
                histo.count(),
                histo.sum() as f64 / 1e9,
            )
        };

        let sz = |name: &str, histo: &Histogram| {
            (
                name.to_string(),
                histo.percentile(0.),
                histo.percentile(50.),
                histo.percentile(90.),
                histo.percentile(99.),
                histo.percentile(99.9),
                histo.percentile(99.99),
                histo.percentile(100.),
                histo.count(),
                histo.sum() as f64,
            )
        };

        println!("tree:");
        p(vec![
            lat("start", &self.tree_start),
            lat("traverse", &self.tree_traverse),
            lat("get", &self.tree_get),
            lat("set", &self.tree_set),
            lat("merge", &self.tree_merge),
            lat("del", &self.tree_del),
            lat("cas", &self.tree_cas),
            lat("scan", &self.tree_scan),
            lat("rev scan", &self.tree_reverse_scan),
        ]);
        println!("tree contention loops: {}", self.tree_loops.load(Acquire));
        println!(
            "tree split success rates: child({}/{}) parent({}/{}) root({}/{})",
            self.tree_child_split_success.load(Acquire),
            self.tree_child_split_attempt.load(Acquire),
            self.tree_parent_split_success.load(Acquire),
            self.tree_parent_split_attempt.load(Acquire),
            self.tree_root_split_success.load(Acquire),
            self.tree_root_split_attempt.load(Acquire),
        );

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("pagecache:");
        p(vec![
            lat("snapshot", &self.advance_snapshot),
            lat("get", &self.get_page),
            lat("rewrite", &self.rewrite_page),
            lat("replace", &self.replace_page),
            lat("link", &self.link_page),
            lat("merge", &self.merge_page),
            lat("pull", &self.pull),
            lat("page_out", &self.page_out),
        ]);

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("serialization and compression:");
        p(vec![
            lat("serialize", &self.serialize),
            lat("deserialize", &self.deserialize),
            lat("compress", &self.compress),
            lat("decompress", &self.decompress),
        ]);

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("log:");
        p(vec![
            lat("make_stable", &self.make_stable),
            lat("read", &self.read),
            lat("write", &self.write_to_log),
            sz("written bytes", &self.written_bytes),
            lat("assign offset", &self.assign_offset),
            lat("assign spinloop", &self.assign_spinloop),
            lat("reserve lat", &self.reserve_lat),
            sz("reserve sz", &self.reserve_sz),
            lat("res cvar r", &self.reserve_current_condvar_wait),
            lat("res cvar w", &self.reserve_written_condvar_wait),
        ]);
        println!("log reservations: {}", self.log_reservations.load(Acquire));
        println!(
            "log res attempts: {}",
            self.log_reservation_attempts.load(Acquire)
        );

        println!("{}", std::iter::repeat("-").take(134).collect::<String>());
        println!("segment accountant:");
        p(vec![
            lat("acquire", &self.accountant_lock),
            lat("hold", &self.accountant_hold),
            lat("next", &self.accountant_next),
            lat("replace", &self.accountant_mark_replace),
            lat("link", &self.accountant_mark_link),
        ]);

        #[cfg(feature = "measure_allocs")]
        {
            println!(
                "{}",
                std::iter::repeat("-").take(134).collect::<String>()
            );
            println!("allocation statistics:");
            println!(
                "total allocations: {}",
                measure_allocs::ALLOCATIONS.load(Acquire)
            );
            println!(
                "allocated bytes: {}",
                measure_allocs::ALLOCATED_BYTES.load(Acquire)
            );
        }
    }
}

#[cfg(feature = "no_metrics")]
impl Metrics {
    pub fn log_reservation_attempted(&self) {}

    pub fn log_reservation_success(&self) {}

    pub fn tree_child_split_attempt(&self) {}

    pub fn tree_child_split_success(&self) {}

    pub fn tree_parent_split_attempt(&self) {}

    pub fn tree_parent_split_success(&self) {}

    pub fn tree_root_split_attempt(&self) {}

    pub fn tree_root_split_success(&self) {}

    pub fn tree_looped(&self) {}

    pub fn print_profile(&self) {}
}
<|endoftext|>


<commit_before>use crate::{build::Arg, util::Id, INTERNAL_ERROR_MSG};

use std::{ffi::OsStr, ffi::OsString, iter::Iterator, ops::Index};

#[derive(PartialEq, Eq, Debug, Clone)]
pub(crate) struct Key {
    key: KeyType,
    index: usize,
}

#[derive(Default, PartialEq, Eq, Debug, Clone)]
pub(crate) struct MKeyMap<'help> {
    /// All of the arguments.
    args: Vec<Arg<'help>>,

    // Cache part:
    /// Will be set after `_build()`.
    keys: Vec<Key>,
}

#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub(crate) enum KeyType {
    Short(char),
    Long(OsString),
    Position(usize),
}

impl KeyType {
    pub(crate) fn is_position(&self) -> bool {
        matches!(self, KeyType::Position(_))
    }
}

impl PartialEq<usize> for KeyType {
    fn eq(&self, rhs: &usize) -> bool {
        match self {
            KeyType::Position(x) => x == rhs,
            _ => false,
        }
    }
}

impl PartialEq<&str> for KeyType {
    fn eq(&self, rhs: &&str) -> bool {
        match self {
            KeyType::Long(l) => l == rhs,
            _ => false,
        }
    }
}

impl PartialEq<OsStr> for KeyType {
    fn eq(&self, rhs: &OsStr) -> bool {
        match self {
            KeyType::Long(l) => l == rhs,
            _ => false,
        }
    }
}

impl PartialEq<char> for KeyType {
    fn eq(&self, rhs: &char) -> bool {
        match self {
            KeyType::Short(c) => c == rhs,
            _ => false,
        }
    }
}

impl<'help> MKeyMap<'help> {
    /// If any arg has corresponding key in this map, we can search the key with
    /// u64(for positional argument), char(for short flag), &str and OsString
    /// (for long flag)
    pub(crate) fn contains<K>(&self, key: K) -> bool
    where
        KeyType: PartialEq<K>,
    {
        self.keys.iter().any(|x| x.key == key)
    }

    /// Push an argument in the map.
    pub(crate) fn push(&mut self, new_arg: Arg<'help>) {
        self.args.push(new_arg);
    }

    /// Find the arg have corresponding key in this map, we can search the key
    /// with u64(for positional argument), char(for short flag), &str and
    /// OsString (for long flag)
    pub(crate) fn get<K: ?Sized>(&self, key: &K) -> Option<&Arg<'help>>
    where
        KeyType: PartialEq<K>,
    {
        self.keys
            .iter()
            .find(|k| &k.key == key)
            .map(|k| &self.args[k.index])
    }

    /// Find out if the map have no arg.
    pub(crate) fn is_empty(&self) -> bool {
        self.args.is_empty()
    }

    /// Return iterators of all keys.
    pub(crate) fn keys(&self) -> impl Iterator<Item = &KeyType> {
        self.keys.iter().map(|x| &x.key)
    }

    /// Return iterators of all args.
    pub(crate) fn args(&self) -> impl Iterator<Item = &Arg<'help>> {
        self.args.iter()
    }

    /// Return mutable iterators of all args.
    pub(crate) fn args_mut<'map>(&'map mut self) -> impl Iterator<Item = &'map mut Arg<'help>> {
        self.args.iter_mut()
    }

    /// We need a lazy build here since some we may change args after creating
    /// the map, you can checkout who uses `args_mut`.
    pub(crate) fn _build(&mut self) {
        for (i, arg) in self.args.iter_mut().enumerate() {
            for k in _get_keys(arg) {
                self.keys.push(Key { key: k, index: i });
            }
        }
    }

    /// Remove an arg in the graph by Id, usually used by `mut_arg`. Return
    /// `Some(arg)` if removed.
    pub(crate) fn remove_by_name(&mut self, name: &Id) -> Option<Arg<'help>> {
        self.args
            .iter()
            .position(|arg| &arg.id == name)
            // since it's a cold function, using this wouldn't hurt much
            .map(|i| self.args.remove(i))
    }

    /// Remove an arg based on index
    pub(crate) fn remove(&mut self, index: usize) -> Arg<'help> {
        self.args.remove(index)
    }
}

impl<'help> Index<&'_ KeyType> for MKeyMap<'help> {
    type Output = Arg<'help>;

    fn index(&self, key: &KeyType) -> &Self::Output {
        self.get(key).expect(INTERNAL_ERROR_MSG)
    }
}

/// Generate key types for an specific Arg.
fn _get_keys(arg: &Arg) -> Vec<KeyType> {
    if let Some(index) = arg.index {
        return vec![KeyType::Position(index)];
    }

    let mut keys = vec![];

    if let Some(short) = arg.short {
        keys.push(KeyType::Short(short));
    }
    if let Some(long) = arg.long {
        keys.push(KeyType::Long(OsString::from(long)));
    }

    for (short, _) in arg.short_aliases.iter() {
        keys.push(KeyType::Short(*short));
    }
    for (long, _) in arg.aliases.iter() {
        keys.push(KeyType::Long(OsString::from(long)));
    }
    keys
}
<commit_msg>perf: Don't allocate a throw-away Vec<commit_after>use crate::{build::Arg, util::Id, INTERNAL_ERROR_MSG};

use std::{ffi::OsStr, ffi::OsString, iter::Iterator, ops::Index};

#[derive(PartialEq, Eq, Debug, Clone)]
pub(crate) struct Key {
    key: KeyType,
    index: usize,
}

#[derive(Default, PartialEq, Eq, Debug, Clone)]
pub(crate) struct MKeyMap<'help> {
    /// All of the arguments.
    args: Vec<Arg<'help>>,

    // Cache part:
    /// Will be set after `_build()`.
    keys: Vec<Key>,
}

#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub(crate) enum KeyType {
    Short(char),
    Long(OsString),
    Position(usize),
}

impl KeyType {
    pub(crate) fn is_position(&self) -> bool {
        matches!(self, KeyType::Position(_))
    }
}

impl PartialEq<usize> for KeyType {
    fn eq(&self, rhs: &usize) -> bool {
        match self {
            KeyType::Position(x) => x == rhs,
            _ => false,
        }
    }
}

impl PartialEq<&str> for KeyType {
    fn eq(&self, rhs: &&str) -> bool {
        match self {
            KeyType::Long(l) => l == rhs,
            _ => false,
        }
    }
}

impl PartialEq<OsStr> for KeyType {
    fn eq(&self, rhs: &OsStr) -> bool {
        match self {
            KeyType::Long(l) => l == rhs,
            _ => false,
        }
    }
}

impl PartialEq<char> for KeyType {
    fn eq(&self, rhs: &char) -> bool {
        match self {
            KeyType::Short(c) => c == rhs,
            _ => false,
        }
    }
}

impl<'help> MKeyMap<'help> {
    /// If any arg has corresponding key in this map, we can search the key with
    /// u64(for positional argument), char(for short flag), &str and OsString
    /// (for long flag)
    pub(crate) fn contains<K>(&self, key: K) -> bool
    where
        KeyType: PartialEq<K>,
    {
        self.keys.iter().any(|x| x.key == key)
    }

    /// Push an argument in the map.
    pub(crate) fn push(&mut self, new_arg: Arg<'help>) {
        self.args.push(new_arg);
    }

    /// Find the arg have corresponding key in this map, we can search the key
    /// with u64(for positional argument), char(for short flag), &str and
    /// OsString (for long flag)
    pub(crate) fn get<K: ?Sized>(&self, key: &K) -> Option<&Arg<'help>>
    where
        KeyType: PartialEq<K>,
    {
        self.keys
            .iter()
            .find(|k| &k.key == key)
            .map(|k| &self.args[k.index])
    }

    /// Find out if the map have no arg.
    pub(crate) fn is_empty(&self) -> bool {
        self.args.is_empty()
    }

    /// Return iterators of all keys.
    pub(crate) fn keys(&self) -> impl Iterator<Item = &KeyType> {
        self.keys.iter().map(|x| &x.key)
    }

    /// Return iterators of all args.
    pub(crate) fn args(&self) -> impl Iterator<Item = &Arg<'help>> {
        self.args.iter()
    }

    /// Return mutable iterators of all args.
    pub(crate) fn args_mut<'map>(&'map mut self) -> impl Iterator<Item = &'map mut Arg<'help>> {
        self.args.iter_mut()
    }

    /// We need a lazy build here since some we may change args after creating
    /// the map, you can checkout who uses `args_mut`.
    pub(crate) fn _build(&mut self) {
        for (i, arg) in self.args.iter().enumerate() {
            append_keys(&mut self.keys, arg, i);
        }
    }

    /// Remove an arg in the graph by Id, usually used by `mut_arg`. Return
    /// `Some(arg)` if removed.
    pub(crate) fn remove_by_name(&mut self, name: &Id) -> Option<Arg<'help>> {
        self.args
            .iter()
            .position(|arg| &arg.id == name)
            // since it's a cold function, using this wouldn't hurt much
            .map(|i| self.args.remove(i))
    }

    /// Remove an arg based on index
    pub(crate) fn remove(&mut self, index: usize) -> Arg<'help> {
        self.args.remove(index)
    }
}

impl<'help> Index<&'_ KeyType> for MKeyMap<'help> {
    type Output = Arg<'help>;

    fn index(&self, key: &KeyType) -> &Self::Output {
        self.get(key).expect(INTERNAL_ERROR_MSG)
    }
}

/// Generate key types for an specific Arg.
fn append_keys(keys: &mut Vec<Key>, arg: &Arg, index: usize) {
    if let Some(pos_index) = arg.index {
        let key = KeyType::Position(pos_index);
        keys.push(Key { key, index });
    } else {
        if let Some(short) = arg.short {
            let key = KeyType::Short(short);
            keys.push(Key { key, index });
        }
        if let Some(long) = arg.long {
            let key = KeyType::Long(OsString::from(long));
            keys.push(Key { key, index });
        }

        for (short, _) in arg.short_aliases.iter() {
            let key = KeyType::Short(*short);
            keys.push(Key { key, index });
        }
        for (long, _) in arg.aliases.iter() {
            let key = KeyType::Long(OsString::from(long));
            keys.push(Key { key, index });
        }
    }
}
<|endoftext|>