<commit_before>use std::fs::File;
use std::io::{self, Read};
use std::thread::sleep;
use std::time::{Duration, Instant};

use rand::{thread_rng, Rng};
use redis::Value::Okay;
use redis::{Client, IntoConnectionInfo, RedisResult, Value};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
#[derive(Debug, Clone)]
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

pub struct RedLockGuard<'a> {
    pub lock: Lock<'a>,
}

impl Drop for RedLockGuard<'_> {
    fn drop(&mut self) {
        self.lock.lock_manager.unlock(&self.lock);
    }
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new<T: AsRef<str> + IntoConnectionInfo>(uris: Vec<T>) -> RedLock {
        let quorum = (uris.len() as u32) / 2 + 1;

        let servers: Vec<Client> = uris
            .into_iter()
            .map(|uri| Client::open(uri).unwrap())
            .collect();

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open("/dev/urandom")?;
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&mut con);
        match result {
            Ok(Okay) => true,
            Ok(_) | Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let mut rng = thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = rng.gen_range(0..self.retry_delay);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    /// Acquire the lock for the given resource and the requested TTL. \
    /// Will wait and yield current task (tokio runtime) until the lock \
    /// is acquired
    ///
    /// Returns a `RedLockGuard` instance which is a RAII wrapper for \
    /// the old `Lock` object
    pub async fn acquire_async(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => { lock = l; break; }
                None => { tokio::task::yield_now().await }
            }
        }
        RedLockGuard {
            lock: lock
        }
    }

    pub fn acquire(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => { lock = l; break; }
                None => {  }
            }
        }
        RedLockGuard {
            lock: lock
        }
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&mut con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use once_cell::sync::Lazy;
    use once_cell::unsync::Lazy as UnsyncLazy;
    use testcontainers::clients::Cli;
    use testcontainers::images::redis::Redis;
    use testcontainers::Container;

    use super::*;

    static DOCKER: Lazy<Cli> = Lazy::new(Cli::default);
    // Actual containers are not thread-safe, but they don't have to be
    thread_local! {
        static CONTAINERS: UnsyncLazy<Vec<Container<'static, Redis>>> = UnsyncLazy::new(|| {
            (0..3)
                .map(|_| DOCKER.run(Redis::default().with_tag("6-alpine")))
                .collect()
        })
    }

    static ADDRESSES: Lazy<Vec<String>> = Lazy::new(|| match std::env::var("ADDRESSES") {
        Ok(addresses) => addresses.split(',').map(String::from).collect(),
        Err(_) => CONTAINERS.with(|containers| {
            containers
                .iter()
                .map(|c| format!("redis://localhost:{}", c.get_host_port(6379)))
                .collect()
        }),
    });

    #[test]
    fn test_redlock_get_unique_id() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());
        assert_eq!(rl.get_unique_lock_id()?.len(), 20);
        Ok(())
    }

    #[test]
    fn test_redlock_get_unique_id_uniqueness() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());

        let id1 = rl.get_unique_lock_id()?;
        let id2 = rl.get_unique_lock_id()?;

        assert_eq!(20, id1.len());
        assert_eq!(20, id2.len());
        assert_ne!(id1, id2);
        Ok(())
    }

    #[test]
    fn test_redlock_valid_instance() {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        assert_eq!(3, rl.servers.len());
        assert_eq!(2, rl.quorum);
    }

    #[test]
    fn test_redlock_direct_unlock_fails() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        assert!(!rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_unlock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        redis::cmd("SET").arg(&*key).arg(&*val).execute(&mut con);

        assert!(rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_lock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;

        redis::cmd("DEL").arg(&*key).execute(&mut con);
        assert!(rl.lock_instance(&rl.servers[0], &*key, &*val, 1000));
        Ok(())
    }

    #[test]
    fn test_redlock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        let _: () = redis::cmd("SET")
            .arg(&*key)
            .arg(&*val)
            .query(&mut con)
            .unwrap();

        let lock = Lock {
            lock_manager: &rl,
            resource: key,
            val,
            validity_time: 0,
        };
        rl.unlock(&lock);
        Ok(())
    }

    #[test]
    fn test_redlock_lock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        match rl.lock(&key, 1000) {
            Some(lock) => {
                assert_eq!(key, lock.resource);
                assert_eq!(20, lock.val.len());
                assert!(lock.validity_time > 900);
                assert!(
                    lock.validity_time > 900,
                    "validity time: {}",
                    lock.validity_time
                );
            }
            None => panic!("Lock failed"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;

        let lock = rl.lock(&key, 1000).unwrap();
        assert!(
            lock.validity_time > 900,
            "validity time: {}",
            lock.validity_time
        );

        if let Some(_l) = rl2.lock(&key, 1000) {
            panic!("Lock acquired, even though it should be locked")
        }

        rl.unlock(&lock);

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }

	#[test]
    fn test_redlock_lock_unlock_raii() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
		{
			let lock_guard = rl.acquire(&key, 1000);
			let lock = &lock_guard.lock;
			assert!(
				lock.validity_time > 900,
				"validity time: {}",
				lock.validity_time
			);

			if let Some(_l) = rl2.lock(&key, 1000) {
				panic!("Lock acquired, even though it should be locked")
			}
		}

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }
}
<commit_msg>Format Rust code using rustfmt<commit_after>use std::fs::File;
use std::io::{self, Read};
use std::thread::sleep;
use std::time::{Duration, Instant};

use rand::{thread_rng, Rng};
use redis::Value::Okay;
use redis::{Client, IntoConnectionInfo, RedisResult, Value};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
#[derive(Debug, Clone)]
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

pub struct RedLockGuard<'a> {
    pub lock: Lock<'a>,
}

impl Drop for RedLockGuard<'_> {
    fn drop(&mut self) {
        self.lock.lock_manager.unlock(&self.lock);
    }
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new<T: AsRef<str> + IntoConnectionInfo>(uris: Vec<T>) -> RedLock {
        let quorum = (uris.len() as u32) / 2 + 1;

        let servers: Vec<Client> = uris
            .into_iter()
            .map(|uri| Client::open(uri).unwrap())
            .collect();

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open("/dev/urandom")?;
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&mut con);
        match result {
            Ok(Okay) => true,
            Ok(_) | Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let mut rng = thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = rng.gen_range(0..self.retry_delay);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    /// Acquire the lock for the given resource and the requested TTL. \
    /// Will wait and yield current task (tokio runtime) until the lock \
    /// is acquired
    ///
    /// Returns a `RedLockGuard` instance which is a RAII wrapper for \
    /// the old `Lock` object
    pub async fn acquire_async(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => {
                    lock = l;
                    break;
                }
                None => tokio::task::yield_now().await,
            }
        }
        RedLockGuard { lock: lock }
    }

    pub fn acquire(&self, resource: &[u8], ttl: usize) -> RedLockGuard<'_> {
        let lock;
        loop {
            match self.lock(resource, ttl) {
                Some(l) => {
                    lock = l;
                    break;
                }
                None => {}
            }
        }
        RedLockGuard { lock: lock }
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let mut con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&mut con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use once_cell::sync::Lazy;
    use once_cell::unsync::Lazy as UnsyncLazy;
    use testcontainers::clients::Cli;
    use testcontainers::images::redis::Redis;
    use testcontainers::Container;

    use super::*;

    static DOCKER: Lazy<Cli> = Lazy::new(Cli::default);
    // Actual containers are not thread-safe, but they don't have to be
    thread_local! {
        static CONTAINERS: UnsyncLazy<Vec<Container<'static, Redis>>> = UnsyncLazy::new(|| {
            (0..3)
                .map(|_| DOCKER.run(Redis::default().with_tag("6-alpine")))
                .collect()
        })
    }

    static ADDRESSES: Lazy<Vec<String>> = Lazy::new(|| match std::env::var("ADDRESSES") {
        Ok(addresses) => addresses.split(',').map(String::from).collect(),
        Err(_) => CONTAINERS.with(|containers| {
            containers
                .iter()
                .map(|c| format!("redis://localhost:{}", c.get_host_port(6379)))
                .collect()
        }),
    });

    #[test]
    fn test_redlock_get_unique_id() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());
        assert_eq!(rl.get_unique_lock_id()?.len(), 20);
        Ok(())
    }

    #[test]
    fn test_redlock_get_unique_id_uniqueness() -> Result<()> {
        let rl = RedLock::new(Vec::<String>::new());

        let id1 = rl.get_unique_lock_id()?;
        let id2 = rl.get_unique_lock_id()?;

        assert_eq!(20, id1.len());
        assert_eq!(20, id2.len());
        assert_ne!(id1, id2);
        Ok(())
    }

    #[test]
    fn test_redlock_valid_instance() {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        assert_eq!(3, rl.servers.len());
        assert_eq!(2, rl.quorum);
    }

    #[test]
    fn test_redlock_direct_unlock_fails() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        assert!(!rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_unlock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        redis::cmd("SET").arg(&*key).arg(&*val).execute(&mut con);

        assert!(rl.unlock_instance(&rl.servers[0], &key, &val));
        Ok(())
    }

    #[test]
    fn test_redlock_direct_lock_succeeds() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;

        redis::cmd("DEL").arg(&*key).execute(&mut con);
        assert!(rl.lock_instance(&rl.servers[0], &*key, &*val, 1000));
        Ok(())
    }

    #[test]
    fn test_redlock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let key = rl.get_unique_lock_id()?;

        let val = rl.get_unique_lock_id()?;
        let mut con = rl.servers[0].get_connection()?;
        let _: () = redis::cmd("SET")
            .arg(&*key)
            .arg(&*val)
            .query(&mut con)
            .unwrap();

        let lock = Lock {
            lock_manager: &rl,
            resource: key,
            val,
            validity_time: 0,
        };
        rl.unlock(&lock);
        Ok(())
    }

    #[test]
    fn test_redlock_lock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        match rl.lock(&key, 1000) {
            Some(lock) => {
                assert_eq!(key, lock.resource);
                assert_eq!(20, lock.val.len());
                assert!(lock.validity_time > 900);
                assert!(
                    lock.validity_time > 900,
                    "validity time: {}",
                    lock.validity_time
                );
            }
            None => panic!("Lock failed"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;

        let lock = rl.lock(&key, 1000).unwrap();
        assert!(
            lock.validity_time > 900,
            "validity time: {}",
            lock.validity_time
        );

        if let Some(_l) = rl2.lock(&key, 1000) {
            panic!("Lock acquired, even though it should be locked")
        }

        rl.unlock(&lock);

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }

    #[test]
    fn test_redlock_lock_unlock_raii() -> Result<()> {
        println!("{}", ADDRESSES.join(","));
        let rl = RedLock::new(ADDRESSES.clone());
        let rl2 = RedLock::new(ADDRESSES.clone());

        let key = rl.get_unique_lock_id()?;
        {
            let lock_guard = rl.acquire(&key, 1000);
            let lock = &lock_guard.lock;
            assert!(
                lock.validity_time > 900,
                "validity time: {}",
                lock.validity_time
            );

            if let Some(_l) = rl2.lock(&key, 1000) {
                panic!("Lock acquired, even though it should be locked")
            }
        }

        match rl2.lock(&key, 1000) {
            Some(l) => assert!(l.validity_time > 900),
            None => panic!("Lock couldn't be acquired"),
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use super::ssz::{Decodable, DecodeError, Encodable, SszStream};
use super::Hash256;

#[derive(Debug, Clone, PartialEq)]
pub struct AttestationData {
    pub slot: u64,
    pub shard: u64,
    pub beacon_block_hash: Hash256,
    pub epoch_boundary_hash: Hash256,
    pub shard_block_hash: Hash256,
    pub latest_crosslink_hash: Hash256,
    pub justified_slot: u64,
    pub justified_block_hash: Hash256,
}

impl Encodable for AttestationData {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append(&self.slot);
        s.append(&self.shard);
        s.append(&self.beacon_block_hash);
        s.append(&self.epoch_boundary_hash);
        s.append(&self.shard_block_hash);
        s.append(&self.latest_crosslink_hash);
        s.append(&self.justified_slot);
        s.append(&self.justified_block_hash);
    }
}

impl Decodable for AttestationData {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (slot, i) = u64::ssz_decode(bytes, i)?;
        let (shard, i) = u64::ssz_decode(bytes, i)?;
        let (beacon_block_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (epoch_boundary_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (shard_block_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (latest_crosslink_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (justified_slot, i) = u64::ssz_decode(bytes, i)?;
        let (justified_block_hash, i) = Hash256::ssz_decode(bytes, i)?;

        let attestation_data = AttestationData {
            slot,
            shard,
            beacon_block_hash,
            epoch_boundary_hash,
            shard_block_hash,
            latest_crosslink_hash,
            justified_slot,
            justified_block_hash,
        };
        Ok((attestation_data, i))
    }
}

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

    #[test]
    pub fn test_attestation_record_ssz_round_trip() {
        let original = AttestationData {
            slot: 42,
            shard: 16,
            beacon_block_hash: Hash256::from("beacon".as_bytes()),
            epoch_boundary_hash: Hash256::from("epoch".as_bytes()),
            shard_block_hash: Hash256::from("shard".as_bytes()),
            latest_crosslink_hash: Hash256::from("xlink".as_bytes()),
            justified_slot: 8,
            justified_block_hash: Hash256::from("justified".as_bytes()),
        };

        let ssz = ssz_encode(&original);

        let (decoded, _) = AttestationData::ssz_decode(&ssz, 0).unwrap();

        assert_eq!(original, decoded);
    }
}
<commit_msg>Update AttestationData struct<commit_after>use super::ssz::{Decodable, DecodeError, Encodable, SszStream};
use super::Hash256;

pub const SSZ_ATTESTION_DATA_LENGTH: usize = {
    8 +             // slot
    8 +             // shard
    32 +            // beacon_block_hash
    32 +            // epoch_boundary_hash
    32 +            // shard_block_hash
    32 +            // latest_crosslink_hash
    8 +             // justified_slot
    32 // justified_block_hash
};

#[derive(Debug, Clone, PartialEq)]
pub struct AttestationData {
    pub slot: u64,
    pub shard: u64,
    pub beacon_block_hash: Hash256,
    pub epoch_boundary_hash: Hash256,
    pub shard_block_hash: Hash256,
    pub latest_crosslink_hash: Hash256,
    pub justified_slot: u64,
    pub justified_block_hash: Hash256,
}

impl AttestationData {
    pub fn zero() -> Self {
        Self {
            slot: 0,
            shard: 0,
            beacon_block_hash: Hash256::zero(),
            epoch_boundary_hash: Hash256::zero(),
            shard_block_hash: Hash256::zero(),
            latest_crosslink_hash: Hash256::zero(),
            justified_slot: 0,
            justified_block_hash: Hash256::zero(),
        }
    }
}

impl Encodable for AttestationData {
    fn ssz_append(&self, s: &mut SszStream) {
        s.append(&self.slot);
        s.append(&self.shard);
        s.append(&self.beacon_block_hash);
        s.append(&self.epoch_boundary_hash);
        s.append(&self.shard_block_hash);
        s.append(&self.latest_crosslink_hash);
        s.append(&self.justified_slot);
        s.append(&self.justified_block_hash);
    }
}

impl Decodable for AttestationData {
    fn ssz_decode(bytes: &[u8], i: usize) -> Result<(Self, usize), DecodeError> {
        let (slot, i) = u64::ssz_decode(bytes, i)?;
        let (shard, i) = u64::ssz_decode(bytes, i)?;
        let (beacon_block_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (epoch_boundary_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (shard_block_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (latest_crosslink_hash, i) = Hash256::ssz_decode(bytes, i)?;
        let (justified_slot, i) = u64::ssz_decode(bytes, i)?;
        let (justified_block_hash, i) = Hash256::ssz_decode(bytes, i)?;

        let attestation_data = AttestationData {
            slot,
            shard,
            beacon_block_hash,
            epoch_boundary_hash,
            shard_block_hash,
            latest_crosslink_hash,
            justified_slot,
            justified_block_hash,
        };
        Ok((attestation_data, i))
    }
}

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

    #[test]
    pub fn test_attestation_record_ssz_round_trip() {
        let original = AttestationData {
            slot: 42,
            shard: 16,
            beacon_block_hash: Hash256::from("beacon".as_bytes()),
            epoch_boundary_hash: Hash256::from("epoch".as_bytes()),
            shard_block_hash: Hash256::from("shard".as_bytes()),
            latest_crosslink_hash: Hash256::from("xlink".as_bytes()),
            justified_slot: 8,
            justified_block_hash: Hash256::from("justified".as_bytes()),
        };

        let ssz = ssz_encode(&original);

        let (decoded, _) = AttestationData::ssz_decode(&ssz, 0).unwrap();

        assert_eq!(original, decoded);
    }
}
<|endoftext|>


<commit_before>use tournament;

#[test]
fn just_the_header_if_no_input() {
    let input = "";
    let expected = "Team                           | MP |  W |  D |  L |  P";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_win_is_three_points_a_loss_is_zero_points() {
    let input = "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  1 |  1 |  0 |  0 |  3\n"
        + "Blithering Badgers             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_win_can_also_be_expressed_as_a_loss() {
    let input = "Blithering Badgers;Allegoric Alaskans;loss";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  1 |  1 |  0 |  0 |  3\n"
        + "Blithering Badgers             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_different_team_can_win() {
    let input = "Blithering Badgers;Allegoric Alaskans;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Blithering Badgers             |  1 |  1 |  0 |  0 |  3\n"
        + "Allegoric Alaskans             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_draw_is_one_point_each() {
    let input = "Allegoric Alaskans;Blithering Badgers;draw";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  1 |  0 |  1 |  0 |  1\n"
        + "Blithering Badgers             |  1 |  0 |  1 |  0 |  1";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_one_match() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  2 |  2 |  0 |  0 |  6\n"
        + "Blithering Badgers             |  2 |  0 |  0 |  2 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_one_winner() {
    let input = "Allegoric Alaskans;Blithering Badgers;loss\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  2 |  1 |  0 |  1 |  3\n"
        + "Blithering Badgers             |  2 |  1 |  0 |  1 |  3";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_two_teams() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Blithering Badgers;Courageous Californians;win\n"
        + "Courageous Californians;Allegoric Alaskans;loss";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  2 |  2 |  0 |  0 |  6\n"
        + "Blithering Badgers             |  2 |  1 |  0 |  1 |  3\n"
        + "Courageous Californians        |  2 |  0 |  0 |  2 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn typical_input() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Devastating Donkeys;Courageous Californians;draw\n"
        + "Devastating Donkeys;Allegoric Alaskans;win\n"
        + "Courageous Californians;Blithering Badgers;loss\n"
        + "Blithering Badgers;Devastating Donkeys;loss\n"
        + "Allegoric Alaskans;Courageous Californians;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Devastating Donkeys            |  3 |  2 |  1 |  0 |  7\n"
        + "Allegoric Alaskans             |  3 |  2 |  0 |  1 |  6\n"
        + "Blithering Badgers             |  3 |  1 |  0 |  2 |  3\n"
        + "Courageous Californians        |  3 |  0 |  1 |  2 |  1";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn incomplete_competition_not_all_pairs_have_played() {
    let input = "Allegoric Alaskans;Blithering Badgers;loss\n".to_string()
        + "Devastating Donkeys;Allegoric Alaskans;loss\n"
        + "Courageous Californians;Blithering Badgers;draw\n"
        + "Allegoric Alaskans;Courageous Californians;win";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  3 |  2 |  0 |  1 |  6\n"
        + "Blithering Badgers             |  2 |  1 |  1 |  0 |  4\n"
        + "Courageous Californians        |  2 |  0 |  1 |  1 |  1\n"
        + "Devastating Donkeys            |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn ties_broken_alphabetically() {
    let input = "Courageous Californians;Devastating Donkeys;win\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win\n"
        + "Devastating Donkeys;Allegoric Alaskans;loss\n"
        + "Courageous Californians;Blithering Badgers;win\n"
        + "Blithering Badgers;Devastating Donkeys;draw\n"
        + "Allegoric Alaskans;Courageous Californians;draw";
    let expected = "Team                           | MP |  W |  D |  L |  P\n".to_string()
        + "Allegoric Alaskans             |  3 |  2 |  1 |  0 |  7\n"
        + "Courageous Californians        |  3 |  2 |  1 |  0 |  7\n"
        + "Blithering Badgers             |  3 |  0 |  1 |  2 |  1\n"
        + "Devastating Donkeys            |  3 |  0 |  1 |  2 |  1";

    assert_eq!(tournament::tally(&input), expected);
}
<commit_msg>tournament: line up the expected table header with the rest of the table<commit_after>use tournament;

#[test]
fn just_the_header_if_no_input() {
    let input = "";
    let expected = "Team                           | MP |  W |  D |  L |  P";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_win_is_three_points_a_loss_is_zero_points() {
    let input = "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  1 |  1 |  0 |  0 |  3\n"
        + "Blithering Badgers             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_win_can_also_be_expressed_as_a_loss() {
    let input = "Blithering Badgers;Allegoric Alaskans;loss";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  1 |  1 |  0 |  0 |  3\n"
        + "Blithering Badgers             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_different_team_can_win() {
    let input = "Blithering Badgers;Allegoric Alaskans;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Blithering Badgers             |  1 |  1 |  0 |  0 |  3\n"
        + "Allegoric Alaskans             |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn a_draw_is_one_point_each() {
    let input = "Allegoric Alaskans;Blithering Badgers;draw";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  1 |  0 |  1 |  0 |  1\n"
        + "Blithering Badgers             |  1 |  0 |  1 |  0 |  1";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_one_match() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  2 |  2 |  0 |  0 |  6\n"
        + "Blithering Badgers             |  2 |  0 |  0 |  2 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_one_winner() {
    let input = "Allegoric Alaskans;Blithering Badgers;loss\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  2 |  1 |  0 |  1 |  3\n"
        + "Blithering Badgers             |  2 |  1 |  0 |  1 |  3";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn there_can_be_more_than_two_teams() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Blithering Badgers;Courageous Californians;win\n"
        + "Courageous Californians;Allegoric Alaskans;loss";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  2 |  2 |  0 |  0 |  6\n"
        + "Blithering Badgers             |  2 |  1 |  0 |  1 |  3\n"
        + "Courageous Californians        |  2 |  0 |  0 |  2 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn typical_input() {
    let input = "Allegoric Alaskans;Blithering Badgers;win\n".to_string()
        + "Devastating Donkeys;Courageous Californians;draw\n"
        + "Devastating Donkeys;Allegoric Alaskans;win\n"
        + "Courageous Californians;Blithering Badgers;loss\n"
        + "Blithering Badgers;Devastating Donkeys;loss\n"
        + "Allegoric Alaskans;Courageous Californians;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Devastating Donkeys            |  3 |  2 |  1 |  0 |  7\n"
        + "Allegoric Alaskans             |  3 |  2 |  0 |  1 |  6\n"
        + "Blithering Badgers             |  3 |  1 |  0 |  2 |  3\n"
        + "Courageous Californians        |  3 |  0 |  1 |  2 |  1";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn incomplete_competition_not_all_pairs_have_played() {
    let input = "Allegoric Alaskans;Blithering Badgers;loss\n".to_string()
        + "Devastating Donkeys;Allegoric Alaskans;loss\n"
        + "Courageous Californians;Blithering Badgers;draw\n"
        + "Allegoric Alaskans;Courageous Californians;win";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  3 |  2 |  0 |  1 |  6\n"
        + "Blithering Badgers             |  2 |  1 |  1 |  0 |  4\n"
        + "Courageous Californians        |  2 |  0 |  1 |  1 |  1\n"
        + "Devastating Donkeys            |  1 |  0 |  0 |  1 |  0";

    assert_eq!(tournament::tally(&input), expected);
}

#[test]
#[ignore]
fn ties_broken_alphabetically() {
    let input = "Courageous Californians;Devastating Donkeys;win\n".to_string()
        + "Allegoric Alaskans;Blithering Badgers;win\n"
        + "Devastating Donkeys;Allegoric Alaskans;loss\n"
        + "Courageous Californians;Blithering Badgers;win\n"
        + "Blithering Badgers;Devastating Donkeys;draw\n"
        + "Allegoric Alaskans;Courageous Californians;draw";
    let expected = "".to_string()
        + "Team                           | MP |  W |  D |  L |  P\n"
        + "Allegoric Alaskans             |  3 |  2 |  1 |  0 |  7\n"
        + "Courageous Californians        |  3 |  2 |  1 |  0 |  7\n"
        + "Blithering Badgers             |  3 |  0 |  1 |  2 |  1\n"
        + "Devastating Donkeys            |  3 |  0 |  1 |  2 |  1";

    assert_eq!(tournament::tally(&input), expected);
}
<|endoftext|>


<commit_before>
/// An iterator that produces *n* repetitions of an element.
///
/// See [`repeat_n()`](../fn.repeat_n.html) for more information.
pub struct RepeatN {
    elt: Option,
    n: usize,
}

/// Create an iterator that produces `n` repetitions of `element`.
pub fn repeat_n(element: A, n: usize) -> RepeatN
    where A: Clone,
{
    if n == 0 {
        RepeatN { elt: None, n: n, }
    } else {
        RepeatN { elt: Some(element), n: n, }
    }
}

impl RepeatN {
    #[deprecated(note = "The ::new constructor is deprecated. Use `repeat_n`")]
    ///
    pub fn new(elt: A, n: usize) -> Self {
        // The code is duplicated here because the new version uses
        // the proper A: Clone bound.
        if n == 0 {
            RepeatN { elt: None, n: n }
        } else {
            RepeatN {
                elt: Some(elt),
                n: n,
            }
        }
    }
}

impl Iterator for RepeatN
    where A: Clone
{
    type Item = A;

    fn next(&mut self) -> Option<Self::Item> {
        if self.n > 1 {
            self.n -= 1;
            self.elt.as_ref().cloned()
        } else {
            self.n = 0;
            self.elt.take()
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.n, Some(self.n))
    }
}

impl DoubleEndedIterator for RepeatN
    where A: Clone
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        self.next()
    }
}

impl ExactSizeIterator for RepeatN
    where A: Clone
{}
<commit_msg>Rm deprecated RepeatN::new (use repeat_n)<commit_after>
/// An iterator that produces *n* repetitions of an element.
///
/// See [`repeat_n()`](../fn.repeat_n.html) for more information.
pub struct RepeatN {
    elt: Option,
    n: usize,
}

/// Create an iterator that produces `n` repetitions of `element`.
pub fn repeat_n(element: A, n: usize) -> RepeatN
    where A: Clone,
{
    if n == 0 {
        RepeatN { elt: None, n: n, }
    } else {
        RepeatN { elt: Some(element), n: n, }
    }
}

impl Iterator for RepeatN
    where A: Clone
{
    type Item = A;

    fn next(&mut self) -> Option<Self::Item> {
        if self.n > 1 {
            self.n -= 1;
            self.elt.as_ref().cloned()
        } else {
            self.n = 0;
            self.elt.take()
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.n, Some(self.n))
    }
}

impl DoubleEndedIterator for RepeatN
    where A: Clone
{
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        self.next()
    }
}

impl ExactSizeIterator for RepeatN
    where A: Clone
{}
<|endoftext|>


<commit_before>#![allow(dead_code)]
#[macro_use]
extern crate enum_primitive;
#[macro_use]
extern crate conrod;
extern crate rustc_serialize;
extern crate time;
extern crate piston_window;

mod arena;
mod db;
mod graphic;
mod player;

use player::Player;
use player::PlayerType;
use arena::Arena;

fn main() {
    println!("");
    println!("Hello and Welcome to the Pokemon Battle Arena");
    println!("");
    testing();
}

// This function is for testing. Pls uncommend before commiting!
fn testing() {
    println!("Testing:");
    // for entry in db::movedex::Movedex::new().get_entries() {
    //     if entry.get_category() == enums::Move_Category::Swagger {
    //         println!("{:?}", entry.get_name());
    //     }
    // }
    // test_players();
    // test_arena();
    graphic::gui::draw_window();
}

fn test_players() {

    println!("The Player Section");
    let mut human = Player::new_by_id(&[5, 3, 17], PlayerType::Human);

    println!("Creating a player and testing count and getting the name");
    println!("Count: {}", human.get_pokemon_count());
    println!("Place 1: {}", human.get_pokemon_list()[0].get_name());
    println!("Place 2: {}", human.get_pokemon_list()[1].get_name());
    println!("Place 3: {}", human.get_pokemon_list()[2].get_name());

    println!("Alive list:");
    println!("Number: {}", human.get_alive_count());
    println!("List: {:#?}", human.get_alive_list());

    println!("Current:");
    println!("Default: {}", human.get_current());
    human.set_current(2);
    println!("Custom: {}", human.get_current());
}

fn test_arena2() {
    // let mut human1 = Player::new_by_id(&[5, 3, 17], PlayerType::Human);
    // let mut human2 = Player::new_by_id(&[18, 19, 122], PlayerType::Human);
    // let mut arena = Arena::new(&mut human1, &mut human2,
    //     db::enums::Types::Normal, db::enums::Weather::ClearSky);
    // let movedex = db::movedex::Movedex::new();
    // let attackid = 1;
    // let attack = movedex.move_by_id(attackid).unwrap();
    // // println!("{}", arena.get_player_one().get_pokemon_count());
    // // println!("{}", arena.get_player_two().get_pokemon_count());
    // println!("P 1: {:#?}", arena.get_player_one().get_pokemon_list()[0].get_name());
    // println!("Evasion: {:#?}", arena.get_player_one().get_pokemon_list()[0].get_current()
    //     .get_stat(db::enums::Stats::Evasion));
    // println!("P 2:{:#?}", arena.get_player_two().get_pokemon_list()[0].get_name());
    // println!("Attack:{:?}", attack);
    // println!("HPAnfang:{:#?}", arena.get_player_one().get_pokemon_list()[0].get_current()
    //     .get_stat(db::enums::Stats::Hp));

    // let ref mut player_one = arena.get_player_two().get_pokemon_list()[0];
    // let ref mut player_two = arena.get_player_one().get_pokemon_list()[0];
    // attack.resolve(&mut player_one,
    //     vec![&mut player_two],
    //     arena.get_player_one(),
    //     arena.get_player_two(),
    //     arena);
    // println!("HPEnde:{:#?}", arena.get_player_one().get_pokemon_list()[0].get_current()
    //     .get_stat(db::enums::Stats::Hp));

    // // println!("{:#?}", arena.get_player_two().get_pokemon_list());
    // // println!("{:?}", arena.get_weather());

}

fn test_arena() {
    // Arena erstellen
    let mut p1 = Player::new_by_id(&[5], PlayerType::Human);
    let mut p2 = Player::new_by_id(&[8], PlayerType::Human);
    let mut arena = Arena::new(&mut p1,
                               &mut p2,
                               db::enums::Types::Normal,
                               db::enums::Weather::ClearSky);
    println!("Player One: {:#?}", arena.get_player_one());
    println!("Player Two: {:#?}", arena.get_player_two());
    // Attacke erstellen und "Kampf"
    let movedex = db::movedex::Movedex::new();
    let attack = movedex.move_by_id(151).unwrap();
    println!("Attack: {}", attack.get_name());

    println!("HP1 vorher: {}",
             arena.get_player_one().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("HP2 vorher: {}",
             arena.get_player_two().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));

    attack.resolve(&mut arena, 2);
    println!("HP1 nachher: {}",
             arena.get_player_one().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("HP2 nachher: {}",
             arena.get_player_two().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("Player One: {:#?}", arena.get_player_one());

    // use std::cell::RefCell;
    // use std::rc::Rc;
    // let mut p1_clone = p1.clone();
    // let mut p2_clone = p2.clone();
    // let mut arena = Arena::new(&mut p1_clone, &mut p2_clone,
    //     db::enums::Types::Normal,
    //     db::enums::Weather::ClearSky);
    // let cell: Rc<RefCell<_>> = Rc::new(RefCell::new(Arena::new(&mut p1, &mut p2,
    //     db::enums::Types::Normal,
    //     db::enums::Weather::ClearSky)));

    // println!("Player One: {:#?}", cell.borrow_mut().get_player_one());
    // // Attacke erstellen und "Kampf"
    // let movedex = db::movedex::Movedex::new();
    // let attack = movedex.move_by_id(1).unwrap();

    // attack.resolve(&mut cell.borrow_mut().get_player_one().get_pokemon_list()[0],
    //     vec![&mut cell.borrow_mut().get_player_two().get_pokemon_list()[0]],
    //     cell.borrow_mut().get_player_one(),
    //     cell.borrow_mut().get_player_two(),
    //     &mut arena);
}
<commit_msg>deleted unused code in main.rs (commit for testing travis)<commit_after>#![allow(dead_code)]
#[macro_use]
extern crate enum_primitive;
#[macro_use]
extern crate conrod;
extern crate rustc_serialize;
extern crate time;
extern crate piston_window;

mod arena;
mod db;
mod graphic;
mod player;

use player::Player;
use player::PlayerType;
use arena::Arena;

fn main() {
    println!("");
    println!("Hello and Welcome to the Pokemon Battle Arena");
    println!("");
    testing();
}

// This function is for testing. Pls uncommend before commiting!
fn testing() {
    println!("Testing:");
    // for entry in db::movedex::Movedex::new().get_entries() {
    //     if entry.get_category() == enums::Move_Category::Swagger {
    //         println!("{:?}", entry.get_name());
    //     }
    // }
    // test_players();
    // test_arena();
    // graphic::gui::draw_window();
}

fn test_players() {

    println!("The Player Section");
    let mut human = Player::new_by_id(&[5, 3, 17], PlayerType::Human);

    println!("Creating a player and testing count and getting the name");
    println!("Count: {}", human.get_pokemon_count());
    println!("Place 1: {}", human.get_pokemon_list()[0].get_name());
    println!("Place 2: {}", human.get_pokemon_list()[1].get_name());
    println!("Place 3: {}", human.get_pokemon_list()[2].get_name());

    println!("Alive list:");
    println!("Number: {}", human.get_alive_count());
    println!("List: {:#?}", human.get_alive_list());

    println!("Current:");
    println!("Default: {}", human.get_current());
    human.set_current(2);
    println!("Custom: {}", human.get_current());
}

fn test_arena() {
    // Arena erstellen
    let mut p1 = Player::new_by_id(&[5], PlayerType::Human);
    let mut p2 = Player::new_by_id(&[8], PlayerType::Human);
    let mut arena = Arena::new(&mut p1,
                               &mut p2,
                               db::enums::Types::Normal,
                               db::enums::Weather::ClearSky);
    println!("Player One: {:#?}", arena.get_player_one());
    println!("Player Two: {:#?}", arena.get_player_two());
    // Attacke erstellen und "Kampf"
    let movedex = db::movedex::Movedex::new();
    let attack = movedex.move_by_id(151).unwrap();
    println!("Attack: {}", attack.get_name());

    println!("HP1 vorher: {}",
             arena.get_player_one().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("HP2 vorher: {}",
             arena.get_player_two().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));

    attack.resolve(&mut arena, 2);
    println!("HP1 nachher: {}",
             arena.get_player_one().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("HP2 nachher: {}",
             arena.get_player_two().get_pokemon_list()[0]
                 .get_current()
                 .get_stat(&db::enums::Stats::Hp));
    println!("Player One: {:#?}", arena.get_player_one());
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Handlers and types for agent's actions.
//!
//! The basic functionality that a GRR agent exposes is called an _action_.
//! Actions are invoked by the server (when running a _flow_), should gather
//! requested information and report back to the server.
//!
//! In RRG each action consists of three components: a request type, a response
//! type and an action handler. Request and response types wrap lower-level
//! Protocol Buffer messages sent by and to the GRR server. Handlers accept one
//! instance of the corresponding request type and send some (zero or more)
//! instances of the corresponding response type.

#[cfg(feature = "action-filesystems")]
#[cfg(target_os = "linux")]
pub mod filesystems;

#[cfg(feature = "action-interfaces")]
#[cfg(target_family = "unix")]
pub mod interfaces;

#[cfg(feature = "action-metadata")]
pub mod metadata;

#[cfg(feature = "action-listdir")]
pub mod listdir;

#[cfg(feature = "action-timeline")]
pub mod timeline;

#[cfg(feature = "action-network")]
pub mod network;

#[cfg(feature = "action-stat")]
pub mod stat;

#[cfg(feature = "action-insttime")]
pub mod insttime;

#[cfg(feature = "action-memsize")]
pub mod memsize;

#[cfg(feature = "action-finder")]
pub mod finder;

use crate::session::{self, Session, Task};

/// Abstraction for action-specific requests.
///
/// Protocol Buffer messages received from the GRR server are not necessarily
/// easy to work with and are hardly idiomatic to Rust. For this reason, actions
/// should define more structured data types to represent their input and should
/// be able to parse raw messages into them.
pub trait Request: Sized {

    /// A type of the corresponding raw proto message.
    type Proto: protobuf::Message + Default;

    /// A method for converting raw proto messages into structured requests.
    fn from_proto(proto: Self::Proto) -> Result<Self, session::ParseError>;
}

/// Abstraction for action-specific responses.
///
/// Like with the [`Request`] type, Protocol Buffer messages sent to the GRR
/// server are very idiomatic to Rust. For this reason, actions should define
/// more structured data types to represent responses and provide a way to
/// convert them into the wire format.
///
/// Note that because of the design flaws in the protocol, actions also need to
/// specify a name of the wrapper RDF class from the Python implementation.
/// Hopefully, one day this issue would be fixed and class names will not leak
/// into the specification.
///
/// [`Request`]: trait.Request.html
pub trait Response: Sized {

    /// A name of the corresponding RDF class.
    const RDF_NAME: Option<&'static str>;

    /// A type of the corresponding raw proto message.
    type Proto: protobuf::Message + Default;

    /// A method for converting structured responses into raw proto messages.
    fn into_proto(self) -> Self::Proto;
}

impl Request for () {

    type Proto = protobuf::well_known_types::Empty;

    fn from_proto(_: protobuf::well_known_types::Empty) -> Result<(), session::ParseError> {
        Ok(())
    }
}

impl Response for () {

    const RDF_NAME: Option<&'static str> = None;

    type Proto = protobuf::well_known_types::Empty;

    fn into_proto(self) -> protobuf::well_known_types::Empty {
        protobuf::well_known_types::Empty::new()
    }
}

/// Dispatches `task` to a handler appropriate for the given `action`.
///
/// This method is a mapping between action names (as specified in the protocol)
/// and action handlers (implemented on the agent).
///
/// If the given action is unknown (or not yet implemented), this function will
/// return an error.
pub fn dispatch<'s, S>(action: &str, task: Task<'s, S>) -> session::Result<()>
where
    S: Session,
{
    match action {
        #[cfg(feature = "action-metadata")]
        "GetClientInfo" => task.execute(self::metadata::handle),

        #[cfg(feature = "action-listdir")]
        "ListDirectory" => task.execute(self::listdir::handle),

        #[cfg(feature = "action-timeline")]
        "Timeline" => task.execute(self::timeline::handle),

        #[cfg(feature = "action-network")]
        "ListNetworkConnections" => task.execute(self::network::handle),

        #[cfg(feature = "action-stat")]
        "GetFileStat" => task.execute(self::stat::handle),

        #[cfg(feature = "action-insttime")]
        "GetInstallDate" => task.execute(self::insttime::handle),

        #[cfg(feature = "action-interfaces")]
        #[cfg(target_family = "unix")]
        "EnumerateInterfaces" => task.execute(self::interfaces::handle),

        #[cfg(feature = "action-filesystems")]
        #[cfg(target_os = "linux")]
        "EnumerateFilesystems" => task.execute(self::filesystems::handle),

        #[cfg(feature = "action-memsize")]
        "GetMemorySize" => task.execute(self::memsize::handle),

        action => return Err(session::Error::Dispatch(String::from(action))),
    }
}
<commit_msg>Define new action argument and result types.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Handlers and types for agent's actions.
//!
//! The basic functionality that a GRR agent exposes is called an _action_.
//! Actions are invoked by the server (when running a _flow_), should gather
//! requested information and report back to the server.
//!
//! In RRG each action consists of three components: a request type, a response
//! type and an action handler. Request and response types wrap lower-level
//! Protocol Buffer messages sent by and to the GRR server. Handlers accept one
//! instance of the corresponding request type and send some (zero or more)
//! instances of the corresponding response type.

#[cfg(feature = "action-filesystems")]
#[cfg(target_os = "linux")]
pub mod filesystems;

#[cfg(feature = "action-interfaces")]
#[cfg(target_family = "unix")]
pub mod interfaces;

#[cfg(feature = "action-metadata")]
pub mod metadata;

#[cfg(feature = "action-listdir")]
pub mod listdir;

#[cfg(feature = "action-timeline")]
pub mod timeline;

#[cfg(feature = "action-network")]
pub mod network;

#[cfg(feature = "action-stat")]
pub mod stat;

#[cfg(feature = "action-insttime")]
pub mod insttime;

#[cfg(feature = "action-memsize")]
pub mod memsize;

#[cfg(feature = "action-finder")]
pub mod finder;

use crate::session::{self, Session, Task};

/// Abstraction for action-specific requests.
///
/// Protocol Buffer messages received from the GRR server are not necessarily
/// easy to work with and are hardly idiomatic to Rust. For this reason, actions
/// should define more structured data types to represent their input and should
/// be able to parse raw messages into them.
pub trait Request: Sized {

    /// A type of the corresponding raw proto message.
    type Proto: protobuf::Message + Default;

    /// A method for converting raw proto messages into structured requests.
    fn from_proto(proto: Self::Proto) -> Result<Self, session::ParseError>;
}

/// Abstraction for action-specific responses.
///
/// Like with the [`Request`] type, Protocol Buffer messages sent to the GRR
/// server are very idiomatic to Rust. For this reason, actions should define
/// more structured data types to represent responses and provide a way to
/// convert them into the wire format.
///
/// Note that because of the design flaws in the protocol, actions also need to
/// specify a name of the wrapper RDF class from the Python implementation.
/// Hopefully, one day this issue would be fixed and class names will not leak
/// into the specification.
///
/// [`Request`]: trait.Request.html
pub trait Response: Sized {

    /// A name of the corresponding RDF class.
    const RDF_NAME: Option<&'static str>;

    /// A type of the corresponding raw proto message.
    type Proto: protobuf::Message + Default;

    /// A method for converting structured responses into raw proto messages.
    fn into_proto(self) -> Self::Proto;
}

impl Request for () {

    type Proto = protobuf::well_known_types::Empty;

    fn from_proto(_: protobuf::well_known_types::Empty) -> Result<(), session::ParseError> {
        Ok(())
    }
}

impl Response for () {

    const RDF_NAME: Option<&'static str> = None;

    type Proto = protobuf::well_known_types::Empty;

    fn into_proto(self) -> protobuf::well_known_types::Empty {
        protobuf::well_known_types::Empty::new()
    }
}

/// Dispatches `task` to a handler appropriate for the given `action`.
///
/// This method is a mapping between action names (as specified in the protocol)
/// and action handlers (implemented on the agent).
///
/// If the given action is unknown (or not yet implemented), this function will
/// return an error.
pub fn dispatch<'s, S>(action: &str, task: Task<'s, S>) -> session::Result<()>
where
    S: Session,
{
    match action {
        #[cfg(feature = "action-metadata")]
        "GetClientInfo" => task.execute(self::metadata::handle),

        #[cfg(feature = "action-listdir")]
        "ListDirectory" => task.execute(self::listdir::handle),

        #[cfg(feature = "action-timeline")]
        "Timeline" => task.execute(self::timeline::handle),

        #[cfg(feature = "action-network")]
        "ListNetworkConnections" => task.execute(self::network::handle),

        #[cfg(feature = "action-stat")]
        "GetFileStat" => task.execute(self::stat::handle),

        #[cfg(feature = "action-insttime")]
        "GetInstallDate" => task.execute(self::insttime::handle),

        #[cfg(feature = "action-interfaces")]
        #[cfg(target_family = "unix")]
        "EnumerateInterfaces" => task.execute(self::interfaces::handle),

        #[cfg(feature = "action-filesystems")]
        #[cfg(target_os = "linux")]
        "EnumerateFilesystems" => task.execute(self::filesystems::handle),

        #[cfg(feature = "action-memsize")]
        "GetMemorySize" => task.execute(self::memsize::handle),

        action => return Err(session::Error::Dispatch(String::from(action))),
    }
}

// TODO(panhania@): Remove all usages of the `Request` trait and replace it with
// `Args`.
pub trait Args: Sized {
    /// Low-level Protocol Buffers type representing the action arguments.
    type Proto: protobuf::Message + Default;

    /// Converts a low-level type to a structured request arguments.
    fn from_proto(proto: Self::Proto) -> Result<Self, ParseArgsError>;
}

pub trait Item: Sized {
    /// Low-level Protocol Buffers type representing the action results.
    type Proto: protobuf::Message + Default;

    /// A name of the corresponding RDF class in GRR.
    const RDF_NAME: &'static str;

    /// Converts an action result ot its low-level representation.
    fn into_proto(self) -> Self::Proto;
}

/// The error type for cases when action argument parsing fails.
#[derive(Debug)]
pub struct ParseArgsError {
    /// A corresponding [`ParseArgsErrorKind`] of this error.
    kind: ParseArgsErrorKind,
    /// A detailed payload associated with the error.
    error: Box<dyn std::error::Error + Send + Sync>,
}

/// Kinds of errors that can happen when parsing action arguments.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ParseArgsErrorKind {
    /// The serialized message with arguments was impossible to deserialize.
    InvalidProto,
}

impl ParseArgsErrorKind {

    fn as_str(&self) -> &'static str {
        use ParseArgsErrorKind::*;

        match *self {
            InvalidProto => "invalid serialized protobuf message",
        }
    }
}

impl std::fmt::Display for ParseArgsError {

    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(fmt, "{}: {}", self.kind.as_str(), self.error)
    }
}

impl std::error::Error for ParseArgsError {

    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.error.source()
    }
}

impl From<protobuf::ProtobufError> for ParseArgsError {

    fn from(error: protobuf::ProtobufError) -> Self {
        ParseArgsError {
            kind: ParseArgsErrorKind::InvalidProto,
            error: Box::new(error),
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::{mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

impl Node {
    pub fn new(weight: size_t, initial_count: c_int, capacity: usize) -> Node {
        let mut n = Node {
            weight: weight,
            leaf_counts: Vec::with_capacity(capacity),
        };
        n.leaf_counts.push(initial_count);
        n
    }
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<c_uint> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in &mut leaves {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in &mut leaves {
        leaf.weight >>= 9;
    }

    let max_num_leaves = 2 * leaves.len() - 2;

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node::new(leaves[0].weight, 1, max_num_leaves),
            lookahead2: Node::new(leaves[1].weight, 2, max_num_leaves),
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = max_num_leaves - 2;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &[Leaf]) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let next_leaf = &leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend(current_list.lookahead1.leaf_counts.iter());
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let num_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.len();
    let previous_list_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.as_ptr();
    {
        let mut current_list = &mut lists[current_list_index];

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();
        for i in 0..num_leaf_counts {
            current_list.lookahead2.leaf_counts.push(unsafe {
                *previous_list_leaf_counts.offset(i as isize)
            });
        }
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<commit_msg>Extend from slice, thanks clippy!!<commit_after>use std::cmp::Ordering;
use std::{mem};

use libc::{size_t, c_int, c_uint};

// Bounded package merge algorithm, based on the paper
// "A Fast and Space-Economical Algorithm for Length-Limited Coding
// Jyrki Katajainen, Alistair Moffat, Andrew Turpin".

#[derive(Debug)]
struct Node {
    weight: size_t,
    leaf_counts: Vec<c_int>,
}

impl Node {
    pub fn new(weight: size_t, initial_count: c_int, capacity: usize) -> Node {
        let mut n = Node {
            weight: weight,
            leaf_counts: Vec::with_capacity(capacity),
        };
        n.leaf_counts.push(initial_count);
        n
    }
}

#[derive(Debug)]
struct Leaf {
    pub weight: size_t,
    pub index: size_t,
}
impl PartialEq for Leaf {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}
impl Eq for Leaf { }
impl Ord for Leaf {
    fn cmp(&self, other: &Self) -> Ordering {
        self.weight.cmp(&other.weight)
    }
}
impl PartialOrd for Leaf {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Debug)]
struct List {
    lookahead1: Node,
    lookahead2: Node,
    next_leaf_index: size_t,
}

/// Calculates the bitlengths for the Huffman tree, based on the counts of each
/// symbol.
pub fn length_limited_code_lengths(frequencies: &[size_t], maxbits: c_int) -> Vec<c_uint> {
    let mut leaves = vec![];

    // Count used symbols and place them in the leaves.
    for (i, &freq) in frequencies.iter().enumerate() {
        if freq != 0 {
            leaves.push(Leaf { weight: freq, index: i });
        }
    }

    // Short circuit some special cases
    if leaves.is_empty() {
        // There are no non-zero frequencies.
        return vec![0; frequencies.len()];
    }
    if leaves.len() == 1 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        return result;
    }
    if leaves.len() == 2 {
        let mut result = vec![0; frequencies.len()];
        result[leaves[0].index] = 1;
        result[leaves[1].index] = 1;
        return result;
    }

    // Sort the leaves from least frequent to most frequent.
    // Add index into the same variable for stable sorting.
    for leaf in &mut leaves {
        leaf.weight = (leaf.weight << 9) | leaf.index;
    }
    leaves.sort();
    for leaf in &mut leaves {
        leaf.weight >>= 9;
    }

    let max_num_leaves = 2 * leaves.len() - 2;

    let mut lists = Vec::with_capacity(maxbits as usize);
    for _ in 0..maxbits {
        lists.push(List {
            lookahead1: Node::new(leaves[0].weight, 1, max_num_leaves),
            lookahead2: Node::new(leaves[1].weight, 2, max_num_leaves),
            next_leaf_index: 2,
        });
    }

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each boundary_pm run creates one.
    let num_boundary_pm_runs = max_num_leaves - 2;
    for _ in 0..num_boundary_pm_runs {
        boundary_pm_toplevel(&mut lists[..], &leaves);
    }

    let n = frequencies.len();
    let mut result = vec![0; n];

    let mut a = lists.pop().unwrap().lookahead2.leaf_counts.into_iter().rev().peekable();

    let mut bitlength_value = 1;
    while let Some(leaf_count) = a.next() {
        let next_count = *a.peek().unwrap_or(&0);
        for i in next_count..leaf_count {
            result[leaves[i as usize].index] = bitlength_value;
        }
        bitlength_value += 1;
    }

    result
}

fn lowest_list(lists: &mut [List], leaves: &[Leaf]) {
    // We're in the lowest list, just add another leaf to the lookaheads
    // There will always be more leaves to be added on level 0 so this is safe.
    let mut current_list = lists.get_mut(0).unwrap();
    let next_leaf = &leaves[current_list.next_leaf_index];
    current_list.lookahead2.weight = next_leaf.weight;

    current_list.lookahead2.leaf_counts[0] = current_list.lookahead1.leaf_counts.last().unwrap() + 1;
    current_list.next_leaf_index += 1;
}

fn next_leaf(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let mut current_list = lists.get_mut(current_list_index).unwrap();

    // The next leaf goes next; counting itself makes the leaf_count increase by one.
    current_list.lookahead2.weight = leaves[current_list.next_leaf_index].weight;
    current_list.lookahead2.leaf_counts.clear();
    current_list.lookahead2.leaf_counts.extend_from_slice(&current_list.lookahead1.leaf_counts);
    let last_index = current_list.lookahead2.leaf_counts.len() - 1;
    current_list.lookahead2.leaf_counts[last_index] += 1;
    current_list.next_leaf_index += 1;
}

fn next_tree(weight_sum: size_t, lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let num_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.len();
    let previous_list_leaf_counts = lists[current_list_index - 1].lookahead2.leaf_counts.as_ptr();
    {
        let mut current_list = &mut lists[current_list_index];

        // Make a tree from the lookaheads from the previous list; that goes next.
        // This is not a leaf node, so the leaf count stays the same.
        current_list.lookahead2.weight = weight_sum;
        current_list.lookahead2.leaf_counts.clear();
        for i in 0..num_leaf_counts {
            current_list.lookahead2.leaf_counts.push(unsafe {
                *previous_list_leaf_counts.offset(i as isize)
            });
        }
        current_list.lookahead2.leaf_counts.push(*current_list.lookahead1.leaf_counts.last().unwrap());
    }

    // The previous list needs two new lookahead nodes.
    boundary_pm(lists, leaves, current_list_index - 1);
    boundary_pm(lists, leaves, current_list_index - 1);
}

fn boundary_pm_toplevel(lists: &mut [List], leaves: &[Leaf]) {
    let last_index = lists.len() - 1;
    boundary_pm(lists, leaves, last_index);
}

fn boundary_pm(lists: &mut [List], leaves: &[Leaf], current_list_index: usize) {
    let next_leaf_index = lists[current_list_index].next_leaf_index;

    if current_list_index == 0 && next_leaf_index == leaves.len() {
        // We've added all the leaves to the lowest list, so we're done here
        return;
    }

    mem::swap(&mut lists[current_list_index].lookahead1, &mut lists[current_list_index].lookahead2);

    if current_list_index == 0 {
        lowest_list(lists, leaves);
    } else {
        // We're at a list other than the lowest list.
        let weight_sum = {
            let previous_list = lists.get(current_list_index - 1).unwrap();
            previous_list.lookahead1.weight + previous_list.lookahead2.weight
        };

        if next_leaf_index < leaves.len() && weight_sum > leaves[next_leaf_index].weight {
            next_leaf(lists, leaves, current_list_index);
        } else {
            next_tree(weight_sum, lists, leaves, current_list_index);
        }
    }
}

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

    #[test]
    fn test_from_paper_3() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 3);
        let answer = vec![3, 3, 3, 3, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn test_from_paper_4() {
        let input = [1, 1, 5, 7, 10, 14];
        let output = length_limited_code_lengths(&input, 4);
        let answer = vec![4, 4, 3, 2, 2, 2];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_7() {
        let input = [252, 0, 1, 6, 9, 10, 6, 3, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![1, 0, 6, 4, 3, 3, 3, 5, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn maxbits_15() {
        let input = [0, 0, 0, 0, 0, 0, 18, 0, 6, 0, 12, 2, 14, 9, 27, 15, 23, 15, 17, 8, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 15);
        let answer = vec! [0, 0, 0, 0, 0, 0, 3, 0, 5, 0, 4, 6, 4, 4, 3, 4, 3, 3, 3, 4, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn no_frequencies() {
        let input = [0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_one_frequency() {
        let input = [0, 10, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = vec![0, 1, 0];
        assert_eq!(output, answer);
    }

    #[test]
    fn only_two_frequencies() {
        let input = [0, 0, 0, 0, 252, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let output = length_limited_code_lengths(&input, 7);
        let answer = [0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(output, answer);
    }
}
<|endoftext|>