<commit_before>use address::{LENGTH, Address, Addressable};
use node::Node;
use num::bigint::{BigUint, ToBigUint};
use std::cmp;
use std::collections::HashMap;
use std::fmt;
use time;
use num;

const MINUTES_UNTIL_NEEDS_REFRESH: i64 = 15;

#[derive(Debug, PartialEq)]
pub enum InsertOutcome {
    Inserted,   // Inserted new node
    Updated,    // Updated existing node
    Discarded   // Bucket is full
}

pub type InsertionResult = Result<InsertOutcome, String>;

pub struct NodeBucket {
    k: usize,
    min: BigUint,
    max: BigUint,
    addresses: Vec<Address>,
    nodes: HashMap<Address, Box<Node>>,
    last_inserted: time::Tm
}

impl fmt::Debug for NodeBucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "NodeBucket {{ [{:040x} - {:040x}], contains: {:?} last_changed: {:?}}}",
               self.min, self.max, self.addresses, self.last_changed())
    }
}

impl NodeBucket {
    pub fn new(k: usize) -> NodeBucket {
        let min = 0.to_biguint().unwrap();
        let max = num::pow(2.to_biguint().unwrap(), LENGTH);
        NodeBucket {
            k: k,
            min: min,
            max: max,
            addresses: Vec::with_capacity(k),
            nodes: HashMap::with_capacity(k),
            last_inserted: time::empty_tm()
        }
    }

    pub fn contains(&self, address: &Address) -> bool {
        self.addresses.contains(address)
    }

    pub fn covers(&self, address: &Address) -> bool {
        let numeric = address.as_numeric();
        self.min <= numeric && numeric < self.max
    }

    pub fn find_node(&mut self, address: &Address) -> Option<&mut Box<Node>> {
        self.nodes.get_mut(address)
    }

    pub fn get_nodes(&mut self) -> Vec<&mut Box<Node>> {
        self.nodes.iter_mut().map(|(_, node)| node).collect()
    }

    pub fn questionable_nodes(&mut self) -> Vec<&mut Box<Node>> {
        self.nodes.iter_mut()
            .map(|(_, node)| node)
            .filter(|n| n.is_questionable())
            .collect()
    }

    pub fn insert(&mut self, node: Box<Node>) -> InsertionResult {
        let address = node.get_address();
        if self.covers(&address) {
            if let Some(pos) = self.addresses.iter().position(|a| a == &address) {
                // TODO: Should this update the node's socket address incase a device has changed
                // IPs?
                self.addresses.remove(pos);
                self.addresses.insert(0, address);
                self.last_inserted = time::now_utc();
                Ok(InsertOutcome::Updated)
            } else if !self.is_full() {
                self.addresses.insert(0, address);
                self.nodes.insert(address, node);
                self.last_inserted = time::now_utc();
                Ok(InsertOutcome::Inserted)
            } else {
                Ok(InsertOutcome::Discarded)
            }
        } else {
            Err(format!("Bucket {:?}  does not cover {:?}", self, address))
        }
    }

    pub fn is_full(&self) -> bool {
        self.nodes.len() >= self.k
    }

    pub fn last_changed(&self) -> time::Tm {
        if let Some((_, node)) = self.nodes.iter().max_by_key(|&(_, node)| node.last_seen()) {
            cmp::max(self.last_inserted, node.last_seen())
        } else {
            self.last_inserted
        }
    }

    pub fn needs_refresh(&self) -> bool {
        let time_since_changed = time::now_utc() - self.last_changed();
        time_since_changed > time::Duration::minutes(MINUTES_UNTIL_NEEDS_REFRESH)
    }

    pub fn random_address_in_space(&self) -> Address {
        Address::random(&self.min, &self.max)
    }

    fn remove(&mut self, address: &Address) {
        self.nodes.remove(address);
        let pos = self.addresses.iter().position(|a| a == address).unwrap();
        self.addresses.remove(pos);
    }

    pub fn remove_bad_nodes(&mut self) {
        let mut to_remove = vec![];
        for (address, node) in self.nodes.iter() {
            if node.is_bad() {
                to_remove.push(address.clone());
            }
        }
        for address in to_remove {
            self.remove(&address);
        }
    }

    pub fn split(self) -> (Self, Self) {
        let difference = &self.max - &self.min;
        let partition = difference / 2.to_biguint().unwrap() + &self.min;
        let (a_addresses, b_addresses) = self.addresses
            .into_iter()
            .partition(|&a| a.as_numeric() < partition);
        let mut a_nodes = HashMap::new();
        let mut b_nodes = HashMap::new();
        for (address, node) in self.nodes.into_iter() {
            if address.as_numeric() < partition {
                a_nodes.insert(address, node);
            } else {
                b_nodes.insert(address, node);
            }
        }

        let a = NodeBucket {
            k: self.k,
            min: self.min,
            max: partition.clone(),
            addresses: a_addresses,
            nodes: a_nodes,
            last_inserted: self.last_inserted
        };
        let b = NodeBucket {
            k: self.k,
            min: partition,
            max: self.max,
            addresses: b_addresses,
            nodes: b_nodes,
            last_inserted: self.last_inserted
        };
        (a, b)
    }
}

#[cfg(test)]
mod tests {
    use address::{Address, Addressable};
    use super::{InsertOutcome, NodeBucket};
    use tests::TestNode;
    use time;

    #[test]
    fn test_insert() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let node = Box::new(TestNode::new(Address::for_content("some string")));
        assert_eq!(bucket.insert(node).unwrap(), InsertOutcome::Inserted);
        assert_eq!(bucket.addresses.len(), 1);
    }

    #[test]
    fn test_insert_duplicate() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let addr_1 = Address::for_content("node 1");
        let addr_2 = Address::for_content("node 2");
        let addr_3 = Address::for_content("node 1");
        let a = Box::new(TestNode::new(addr_1));
        let b = Box::new(TestNode::new(addr_2));
        let c = Box::new(TestNode::new(addr_3));

        bucket.insert(a).unwrap();
        bucket.insert(b).unwrap();
        assert_eq!(bucket.insert(c).unwrap(), InsertOutcome::Updated);

        assert_eq!(bucket.addresses.len(), 2);
        assert_eq!(bucket.nodes[&Address::for_content("node 1")].get_address(), addr_1);
    }

    #[test]
    fn test_insert_full() {
        let mut bucket: NodeBucket = NodeBucket::new(2);
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 2")))).unwrap();
        let result = bucket.insert(Box::new(TestNode::new(Address::for_content("node 3")))).unwrap();
        assert_eq!(result, InsertOutcome::Discarded);
        assert_eq!(bucket.addresses.len(), 2);
        assert_eq!(bucket.nodes.get(&Address::for_content("node 3")), None);
    }

    #[test]
    fn insert_node_outside_of_address_space() {
        use num::bigint::ToBigUint;
        use std::collections::HashMap;

        let mut bucket = NodeBucket {
            k: 8,
            min: 0.to_biguint().unwrap(),
            max: 1.to_biguint().unwrap(),
            addresses: vec![],
            nodes: HashMap::new(),
            last_inserted: time::empty_tm()
        };
        let result = bucket.insert(Box::new(TestNode::new(Address::for_content("node 3"))));
        assert!(result.is_err());
    }

    #[test]
    fn test_is_full() {
        let mut bucket: NodeBucket = NodeBucket::new(2);
        assert!(!bucket.is_full());

        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 2")))).unwrap();
        assert!(bucket.is_full());
    }

    #[test]
    fn test_contains() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        assert!(bucket.contains(&Address::for_content("node 1")));
    }

    #[test]
    fn test_covers() {
        let bucket: NodeBucket = NodeBucket::new(8);
        let address_1 = Address::from_str("0000000000000000000000000000000000000000");
        let address_2 = Address::from_str("ffffffffffffffffffffffffffffffffffffffff");
        assert!(bucket.covers(&address_1));
        assert!(bucket.covers(&address_2));
    }

    #[test]
    fn test_split() {
        let mut bucket: NodeBucket = NodeBucket::new(4);
        let addr_1 = Address::from_str("0000000000000000000000000000000000000000");
        let addr_2 = Address::from_str("7fffffffffffffffffffffffffffffffffffffff");
        let addr_3 = Address::from_str("8000000000000000000000000000000000000000");
        let addr_4 = Address::from_str("ffffffffffffffffffffffffffffffffffffffff");
        let node_1 = Box::new(TestNode::new(addr_1));
        let node_2 = Box::new(TestNode::new(addr_2));
        let node_3 = Box::new(TestNode::new(addr_3));
        let node_4 = Box::new(TestNode::new(addr_4));

        bucket.insert(node_1).unwrap();
        bucket.insert(node_2).unwrap();
        bucket.insert(node_3).unwrap();
        bucket.insert(node_4).unwrap();

        let (a, b) = bucket.split();

        // Splits up known address
        assert_eq!(a.addresses.len(), 2);
        assert_eq!(b.addresses.len(), 2);

        // Splits up known nodes
        assert_eq!(a.nodes[&Address::from_str("0000000000000000000000000000000000000000")].get_address(), addr_1);
        assert_eq!(a.nodes[&Address::from_str("0000000000000000000000000000000000000000")].get_address(), addr_1);
        assert_eq!(a.nodes[&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")].get_address(), addr_2);
        assert_eq!(b.nodes[&Address::from_str("8000000000000000000000000000000000000000")].get_address(), addr_3);
        assert_eq!(b.nodes[&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")].get_address(), addr_4);

        // Equitably covers address space
        assert!(a.covers(&Address::from_str("0000000000000000000000000000000000000000")));
        assert!(!b.covers(&Address::from_str("0000000000000000000000000000000000000000")));
        assert!(a.covers(&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")));
        assert!(!b.covers(&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")));
        assert!(!a.covers(&Address::from_str("8000000000000000000000000000000000000000")));
        assert!(b.covers(&Address::from_str("8000000000000000000000000000000000000000")));
        assert!(!a.covers(&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")));
        assert!(b.covers(&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")));
    }

    #[test]
    fn test_last_changed() {
        let mut bucket: NodeBucket = NodeBucket::new(4);

        // it updates upon inserting a node
        let last_changed_before_insert = bucket.last_changed();
        let addr = Address::from_str("0000000000000000000000000000000000000000");
        let node = Box::new(TestNode::new(addr));
        bucket.insert(node).unwrap();
        let last_changed_after_insert = bucket.last_changed();
        assert!(last_changed_before_insert < last_changed_after_insert);

        // it updates when one of its nodes are heard from
        let last_changed_before_received = bucket.last_changed();
        bucket.find_node(&addr).unwrap().received_response(1);
        let last_changed_after_received = bucket.last_changed();
        assert!(last_changed_before_received < last_changed_after_received);
    }

    #[test]
    fn test_remove_bad_nodes() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let node_1 = Box::new(TestNode::new(Address::for_content("good node")));
        let node_2 = Box::new(TestNode::questionable(Address::for_content("questionable node")));
        let node_3 = Box::new(TestNode::bad(Address::for_content("bad node")));
        bucket.insert(node_1).unwrap();
        bucket.insert(node_2).unwrap();
        bucket.insert(node_3).unwrap();
        assert_eq!(bucket.addresses.len(), 3);

        bucket.remove_bad_nodes();
        assert_eq!(bucket.addresses.len(), 2);
    }
}
<commit_msg>Log when removing a bad node from bucket<commit_after>use address::{LENGTH, Address, Addressable};
use node::Node;
use num::bigint::{BigUint, ToBigUint};
use std::cmp;
use std::collections::HashMap;
use std::fmt;
use time;
use num;

const MINUTES_UNTIL_NEEDS_REFRESH: i64 = 15;

#[derive(Debug, PartialEq)]
pub enum InsertOutcome {
    Inserted,   // Inserted new node
    Updated,    // Updated existing node
    Discarded   // Bucket is full
}

pub type InsertionResult = Result<InsertOutcome, String>;

pub struct NodeBucket {
    k: usize,
    min: BigUint,
    max: BigUint,
    addresses: Vec<Address>,
    nodes: HashMap<Address, Box<Node>>,
    last_inserted: time::Tm
}

impl fmt::Debug for NodeBucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "NodeBucket {{ [{:040x} - {:040x}], contains: {:?} last_changed: {:?}}}",
               self.min, self.max, self.addresses, self.last_changed())
    }
}

impl NodeBucket {
    pub fn new(k: usize) -> NodeBucket {
        let min = 0.to_biguint().unwrap();
        let max = num::pow(2.to_biguint().unwrap(), LENGTH);
        NodeBucket {
            k: k,
            min: min,
            max: max,
            addresses: Vec::with_capacity(k),
            nodes: HashMap::with_capacity(k),
            last_inserted: time::empty_tm()
        }
    }

    pub fn contains(&self, address: &Address) -> bool {
        self.addresses.contains(address)
    }

    pub fn covers(&self, address: &Address) -> bool {
        let numeric = address.as_numeric();
        self.min <= numeric && numeric < self.max
    }

    pub fn find_node(&mut self, address: &Address) -> Option<&mut Box<Node>> {
        self.nodes.get_mut(address)
    }

    pub fn get_nodes(&mut self) -> Vec<&mut Box<Node>> {
        self.nodes.iter_mut().map(|(_, node)| node).collect()
    }

    pub fn questionable_nodes(&mut self) -> Vec<&mut Box<Node>> {
        self.nodes.iter_mut()
            .map(|(_, node)| node)
            .filter(|n| n.is_questionable())
            .collect()
    }

    pub fn insert(&mut self, node: Box<Node>) -> InsertionResult {
        let address = node.get_address();
        if self.covers(&address) {
            if let Some(pos) = self.addresses.iter().position(|a| a == &address) {
                // TODO: Should this update the node's socket address incase a device has changed
                // IPs?
                self.addresses.remove(pos);
                self.addresses.insert(0, address);
                self.last_inserted = time::now_utc();
                Ok(InsertOutcome::Updated)
            } else if !self.is_full() {
                self.addresses.insert(0, address);
                self.nodes.insert(address, node);
                self.last_inserted = time::now_utc();
                Ok(InsertOutcome::Inserted)
            } else {
                Ok(InsertOutcome::Discarded)
            }
        } else {
            Err(format!("Bucket {:?}  does not cover {:?}", self, address))
        }
    }

    pub fn is_full(&self) -> bool {
        self.nodes.len() >= self.k
    }

    pub fn last_changed(&self) -> time::Tm {
        if let Some((_, node)) = self.nodes.iter().max_by_key(|&(_, node)| node.last_seen()) {
            cmp::max(self.last_inserted, node.last_seen())
        } else {
            self.last_inserted
        }
    }

    pub fn needs_refresh(&self) -> bool {
        let time_since_changed = time::now_utc() - self.last_changed();
        time_since_changed > time::Duration::minutes(MINUTES_UNTIL_NEEDS_REFRESH)
    }

    pub fn random_address_in_space(&self) -> Address {
        Address::random(&self.min, &self.max)
    }

    fn remove(&mut self, address: &Address) {
        self.nodes.remove(address);
        let pos = self.addresses.iter().position(|a| a == address).unwrap();
        self.addresses.remove(pos);
    }

    pub fn remove_bad_nodes(&mut self) {
        let mut to_remove = vec![];
        for (address, node) in self.nodes.iter() {
            if node.is_bad() {
                debug!("Removing bad node {:?}", node);
                to_remove.push(address.clone());
            }
        }
        for address in to_remove {
            self.remove(&address);
        }
    }

    pub fn split(self) -> (Self, Self) {
        let difference = &self.max - &self.min;
        let partition = difference / 2.to_biguint().unwrap() + &self.min;
        let (a_addresses, b_addresses) = self.addresses
            .into_iter()
            .partition(|&a| a.as_numeric() < partition);
        let mut a_nodes = HashMap::new();
        let mut b_nodes = HashMap::new();
        for (address, node) in self.nodes.into_iter() {
            if address.as_numeric() < partition {
                a_nodes.insert(address, node);
            } else {
                b_nodes.insert(address, node);
            }
        }

        let a = NodeBucket {
            k: self.k,
            min: self.min,
            max: partition.clone(),
            addresses: a_addresses,
            nodes: a_nodes,
            last_inserted: self.last_inserted
        };
        let b = NodeBucket {
            k: self.k,
            min: partition,
            max: self.max,
            addresses: b_addresses,
            nodes: b_nodes,
            last_inserted: self.last_inserted
        };
        (a, b)
    }
}

#[cfg(test)]
mod tests {
    use address::{Address, Addressable};
    use super::{InsertOutcome, NodeBucket};
    use tests::TestNode;
    use time;

    #[test]
    fn test_insert() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let node = Box::new(TestNode::new(Address::for_content("some string")));
        assert_eq!(bucket.insert(node).unwrap(), InsertOutcome::Inserted);
        assert_eq!(bucket.addresses.len(), 1);
    }

    #[test]
    fn test_insert_duplicate() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let addr_1 = Address::for_content("node 1");
        let addr_2 = Address::for_content("node 2");
        let addr_3 = Address::for_content("node 1");
        let a = Box::new(TestNode::new(addr_1));
        let b = Box::new(TestNode::new(addr_2));
        let c = Box::new(TestNode::new(addr_3));

        bucket.insert(a).unwrap();
        bucket.insert(b).unwrap();
        assert_eq!(bucket.insert(c).unwrap(), InsertOutcome::Updated);

        assert_eq!(bucket.addresses.len(), 2);
        assert_eq!(bucket.nodes[&Address::for_content("node 1")].get_address(), addr_1);
    }

    #[test]
    fn test_insert_full() {
        let mut bucket: NodeBucket = NodeBucket::new(2);
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 2")))).unwrap();
        let result = bucket.insert(Box::new(TestNode::new(Address::for_content("node 3")))).unwrap();
        assert_eq!(result, InsertOutcome::Discarded);
        assert_eq!(bucket.addresses.len(), 2);
        assert_eq!(bucket.nodes.get(&Address::for_content("node 3")), None);
    }

    #[test]
    fn insert_node_outside_of_address_space() {
        use num::bigint::ToBigUint;
        use std::collections::HashMap;

        let mut bucket = NodeBucket {
            k: 8,
            min: 0.to_biguint().unwrap(),
            max: 1.to_biguint().unwrap(),
            addresses: vec![],
            nodes: HashMap::new(),
            last_inserted: time::empty_tm()
        };
        let result = bucket.insert(Box::new(TestNode::new(Address::for_content("node 3"))));
        assert!(result.is_err());
    }

    #[test]
    fn test_is_full() {
        let mut bucket: NodeBucket = NodeBucket::new(2);
        assert!(!bucket.is_full());

        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 2")))).unwrap();
        assert!(bucket.is_full());
    }

    #[test]
    fn test_contains() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        bucket.insert(Box::new(TestNode::new(Address::for_content("node 1")))).unwrap();
        assert!(bucket.contains(&Address::for_content("node 1")));
    }

    #[test]
    fn test_covers() {
        let bucket: NodeBucket = NodeBucket::new(8);
        let address_1 = Address::from_str("0000000000000000000000000000000000000000");
        let address_2 = Address::from_str("ffffffffffffffffffffffffffffffffffffffff");
        assert!(bucket.covers(&address_1));
        assert!(bucket.covers(&address_2));
    }

    #[test]
    fn test_split() {
        let mut bucket: NodeBucket = NodeBucket::new(4);
        let addr_1 = Address::from_str("0000000000000000000000000000000000000000");
        let addr_2 = Address::from_str("7fffffffffffffffffffffffffffffffffffffff");
        let addr_3 = Address::from_str("8000000000000000000000000000000000000000");
        let addr_4 = Address::from_str("ffffffffffffffffffffffffffffffffffffffff");
        let node_1 = Box::new(TestNode::new(addr_1));
        let node_2 = Box::new(TestNode::new(addr_2));
        let node_3 = Box::new(TestNode::new(addr_3));
        let node_4 = Box::new(TestNode::new(addr_4));

        bucket.insert(node_1).unwrap();
        bucket.insert(node_2).unwrap();
        bucket.insert(node_3).unwrap();
        bucket.insert(node_4).unwrap();

        let (a, b) = bucket.split();

        // Splits up known address
        assert_eq!(a.addresses.len(), 2);
        assert_eq!(b.addresses.len(), 2);

        // Splits up known nodes
        assert_eq!(a.nodes[&Address::from_str("0000000000000000000000000000000000000000")].get_address(), addr_1);
        assert_eq!(a.nodes[&Address::from_str("0000000000000000000000000000000000000000")].get_address(), addr_1);
        assert_eq!(a.nodes[&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")].get_address(), addr_2);
        assert_eq!(b.nodes[&Address::from_str("8000000000000000000000000000000000000000")].get_address(), addr_3);
        assert_eq!(b.nodes[&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")].get_address(), addr_4);

        // Equitably covers address space
        assert!(a.covers(&Address::from_str("0000000000000000000000000000000000000000")));
        assert!(!b.covers(&Address::from_str("0000000000000000000000000000000000000000")));
        assert!(a.covers(&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")));
        assert!(!b.covers(&Address::from_str("7fffffffffffffffffffffffffffffffffffffff")));
        assert!(!a.covers(&Address::from_str("8000000000000000000000000000000000000000")));
        assert!(b.covers(&Address::from_str("8000000000000000000000000000000000000000")));
        assert!(!a.covers(&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")));
        assert!(b.covers(&Address::from_str("ffffffffffffffffffffffffffffffffffffffff")));
    }

    #[test]
    fn test_last_changed() {
        let mut bucket: NodeBucket = NodeBucket::new(4);

        // it updates upon inserting a node
        let last_changed_before_insert = bucket.last_changed();
        let addr = Address::from_str("0000000000000000000000000000000000000000");
        let node = Box::new(TestNode::new(addr));
        bucket.insert(node).unwrap();
        let last_changed_after_insert = bucket.last_changed();
        assert!(last_changed_before_insert < last_changed_after_insert);

        // it updates when one of its nodes are heard from
        let last_changed_before_received = bucket.last_changed();
        bucket.find_node(&addr).unwrap().received_response(1);
        let last_changed_after_received = bucket.last_changed();
        assert!(last_changed_before_received < last_changed_after_received);
    }

    #[test]
    fn test_remove_bad_nodes() {
        let mut bucket: NodeBucket = NodeBucket::new(8);
        let node_1 = Box::new(TestNode::new(Address::for_content("good node")));
        let node_2 = Box::new(TestNode::questionable(Address::for_content("questionable node")));
        let node_3 = Box::new(TestNode::bad(Address::for_content("bad node")));
        bucket.insert(node_1).unwrap();
        bucket.insert(node_2).unwrap();
        bucket.insert(node_3).unwrap();
        assert_eq!(bucket.addresses.len(), 3);

        bucket.remove_bad_nodes();
        assert_eq!(bucket.addresses.len(), 2);
    }
}
<|endoftext|>


<commit_before>use std::sync::atomic::{AtomicUsize, Ordering};
use std::ptr;

use ctxt::SemContext;
use gc::Address;
use gc::root::IndirectObj;
use gc::swiper::{CARD_SIZE, PROMOTION_AGE, Region};
use gc::swiper::card::CardTable;
use gc::swiper::crossing::CrossingMap;
use gc::swiper::old::OldGen;
use mem;
use object::Obj;
use timer::{in_ms, Timer};

pub struct YoungGen {
    // bounds of from- & to-space
    total: Region,

    // address that separates from & to-space
    separator: Address,

    // address of next free memory
    free: AtomicUsize,

    // end of free memory (either separator or total.end)
    end: AtomicUsize,
}

impl YoungGen {
    pub fn new(young_start: Address, young_end: Address) -> YoungGen {
        let half_size = (young_end.to_usize() - young_start.to_usize()) / 2;
        let half_address = young_start.offset(half_size);

        YoungGen {
            total: Region::new(young_start, young_end),
            separator: half_address,
            free: AtomicUsize::new(young_start.to_usize()),
            end: AtomicUsize::new(half_address.to_usize()),
        }
    }

    pub fn alloc(&self, size: usize) -> *const u8 {
        let mut old = self.free.load(Ordering::Relaxed);
        let mut new;

        loop {
            new = old + size;

            if new >= self.end.load(Ordering::Relaxed) {
                return ptr::null();
            }

            let res = self.free.compare_exchange_weak(
                old,
                new,
                Ordering::SeqCst,
                Ordering::Relaxed,
            );

            match res {
                Ok(_) => break,
                Err(x) => old = x,
            }
        }

        old as *const u8
    }

    pub fn collect(
        &self,
        ctxt: &SemContext,
        rootset: Vec<IndirectObj>,
        card_table: &CardTable,
        crossing_map: &CrossingMap,
        old: &OldGen,
    ) {
        let mut timer = Timer::new(ctxt.args.flag_gc_events);

        let mut scan;

        let mut free;
        let new_end;

        if self.end.load(Ordering::Relaxed) == self.separator.to_usize() {
            scan = self.total.start;
            free = self.separator;
            new_end = self.total.end;

        } else {
            scan = self.separator;
            free = self.total.start;
            new_end = self.separator;
        }

        // detect all references from roots into young generation
        for &root in &rootset {
            let root_ptr = root.get();

            if self.total.includes(Address::from_ptr(root_ptr)) {
                root.set(copy(root_ptr, &mut free, old));
            }
        }

        // detect references from old generation (dirty cards) into young generation
        copy_dirty_cards(card_table, crossing_map, &mut free, &self.total, old);

        // visit all fields in copied objects
        while scan < free {
            let object = unsafe { &mut *scan.to_mut_ptr::<Obj>() };

            object.visit_reference_fields(|child| {
                let child_ptr = child.get();

                if self.total.includes(Address::from_ptr(child_ptr)) {
                    child.set(copy(child_ptr, &mut free, old));
                }
            });

            scan = scan.offset(object.size());
        }

        // swap spaces
        self.end.store(new_end.to_usize(), Ordering::Relaxed);

        // update start of free memory
        self.free.store(free.to_usize(), Ordering::SeqCst);

        // memset from-space to garbage data for debug builds
        // makes sure that no pointer into from-space is left
        if cfg!(debug_assertions) {
            let start = if new_end == self.separator {
                self.total.start
            } else {
                self.separator
            };

            let size = self.separator.to_usize() - self.total.start.to_usize();

            unsafe {
                ptr::write_bytes(start.to_mut_ptr::<u8>(), 0xcc, size);
            }
        }

        timer.stop_with(|dur| if ctxt.args.flag_gc_events {
            println!("GC minor: collect garbage ({} ms)", in_ms(dur));
        });
    }
}

// copy all references from old- into young-generation.
fn copy_dirty_cards(
    card_table: &CardTable,
    crossing_map: &CrossingMap,
    mut free: &mut Address,
    young: &Region,
    old: &OldGen,
) {
    card_table.visit_dirty(|card| {
        let crossing_entry = crossing_map.get(card);

        // card contains: any data but no references, then first object
        if crossing_entry.is_first_object() {
            let card_address = crossing_map.address_of_card(card);

            // first_object() returns number of words before end of card
            let offset_from_end = crossing_entry.first_object() as usize * ptr_width();
            let ptr = card_address.offset(CARD_SIZE - offset_from_end);

            let end = card_address.offset(CARD_SIZE);

            // copy all objects from this card
            copy_card(ptr, end, free, young, old);

        // card contains: references, then first object
        } else if crossing_entry.is_references_at_start() {
            let number_references = crossing_entry.references_at_start() as usize;
            let card_address = crossing_map.address_of_card(card);
            let mut ptr = card_address;

            for _ in 0..number_references {
                let ind_ptr = IndirectObj::from_address(ptr);
                let dir_ptr = ind_ptr.get();

                if young.includes(Address::from_ptr(dir_ptr)) {
                    ind_ptr.set(copy(dir_ptr, &mut free, old));
                }

                ptr = ptr.offset(ptr_width());
            }

            let end = card_address.offset(CARD_SIZE);

            // copy all objects from this card
            copy_card(ptr, end, free, young, old);

        // object spans multiple cards
        } else if crossing_entry.is_previous_card() {
            let mut previous = card - crossing_entry.previous_card() as usize;
            let mut crossing_entry = crossing_map.get(previous);

            while crossing_entry.is_previous_card() {
                previous = previous - crossing_entry.previous_card() as usize;
                crossing_entry = crossing_map.get(previous);
            }

            let previous_address = crossing_map.address_of_card(previous);
            let mut ptr;

            if crossing_entry.is_references_at_start() {
                let number_references = crossing_entry.references_at_start() as usize;
                ptr = previous_address.offset(number_references * ptr_width());

            } else if crossing_entry.is_first_object() {
                let offset_from_end = crossing_entry.first_object() as usize * ptr_width();
                ptr = previous_address.offset(CARD_SIZE - offset_from_end);

            } else {
                assert!(crossing_entry.is_no_references());
                panic!("dirty card without references: is this allowed to happen?");
            }

            let card_address = crossing_map.address_of_card(card);
            let mut last_object = ptr;

            while ptr < card_address {
                let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

                last_object = ptr;
                ptr = ptr.offset(object.size());
            }

            // last_object now stores pointer that reaches into current card
            // TODO: handle fields in last_object

            copy_card(ptr, card_address.offset(CARD_SIZE), free, young, old);

        } else {
            assert!(crossing_entry.is_no_references());
            panic!("dirty card without references: is this allowed to happen?");
        }
    });
}

fn copy_card(mut ptr: Address, end: Address, mut free: &mut Address, young: &Region, old: &OldGen) {
    while ptr < end {
        let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

        object.visit_reference_fields(|child| {
            let child_ptr = child.get();

            if young.includes(Address::from_ptr(child_ptr)) {
                child.set(copy(child_ptr, &mut free, old));
            }
        });

        ptr = ptr.offset(object.size());
    }
}

pub fn copy(obj: *mut Obj, free: &mut Address, old: &OldGen) -> *mut Obj {
    let obj = unsafe { &mut *obj };
    let addr = get_forwarding_address(obj);

    if is_forwarding_address(addr) {
        unmark_forwarding_address(addr).to_mut_ptr::<Obj>()
    } else {
        let obj_size = obj.size();
        let age = obj.header_mut().increase_age();

        let addr: Address;

        // if object is old enough we copy it into the old generation
        if age >= PROMOTION_AGE {
            addr = Address::from_ptr(old.alloc(obj_size));
            copy_object(obj, addr, obj_size);

        // otherwise the object remains in the young generation for now
        } else {
            addr = *free;
            copy_object(obj, addr, obj_size);
            *free = addr.offset(obj_size);
        }

        set_forwarding_address(obj, addr);

        addr.to_mut_ptr::<Obj>()
    }
}

fn copy_object(obj: &Obj, addr: Address, size: usize) {
    unsafe {
        ptr::copy_nonoverlapping(
            obj as *const Obj as *const u8,
            addr.to_mut_ptr::<u8>(),
            size,
        );
    }
}

pub fn is_forwarding_address(obj: Address) -> bool {
    (obj.to_usize() & 1) == 1
}

pub fn mark_forwarding_address(obj: Address) -> Address {
    (obj.to_usize() | 1).into()
}

pub fn unmark_forwarding_address(obj: Address) -> Address {
    (obj.to_usize() & !1).into()
}

pub fn get_forwarding_address(obj: &Obj) -> Address {
    unsafe { *(obj as *const Obj as *const Address) }
}

pub fn set_forwarding_address(obj: &mut Obj, addr: Address) {
    unsafe {
        *(obj as *mut Obj as *mut Address) = mark_forwarding_address(addr);
    }
}

#[inline(always)]
fn ptr_width() -> usize {
    mem::ptr_width() as usize
}<commit_msg>use mprotect() for protecting regions<commit_after>use std::sync::atomic::{AtomicUsize, Ordering};
use std::ptr;

use ctxt::SemContext;
use gc::Address;
use gc::root::IndirectObj;
use gc::swiper::{CARD_SIZE, PROMOTION_AGE, Region};
use gc::swiper::card::CardTable;
use gc::swiper::crossing::CrossingMap;
use gc::swiper::old::OldGen;
use mem;
use os::{self, ProtType};
use object::Obj;
use timer::{in_ms, Timer};

pub struct YoungGen {
    // bounds of from- & to-space
    total: Region,

    // size of both from- or to-space.
    // Not combined, use total.size() for that.
    size: usize,

    // address that separates from & to-space
    separator: Address,

    // address of next free memory
    free: AtomicUsize,

    // end of free memory (either separator or total.end)
    end: AtomicUsize,
}

impl YoungGen {
    pub fn new(young_start: Address, young_end: Address) -> YoungGen {
        let half_size = (young_end.to_usize() - young_start.to_usize()) / 2;
        let half_address = young_start.offset(half_size);

        YoungGen {
            total: Region::new(young_start, young_end),
            size: half_size,
            separator: half_address,
            free: AtomicUsize::new(young_start.to_usize()),
            end: AtomicUsize::new(half_address.to_usize()),
        }
    }

    pub fn alloc(&self, size: usize) -> *const u8 {
        let mut old = self.free.load(Ordering::Relaxed);
        let mut new;

        loop {
            new = old + size;

            if new >= self.end.load(Ordering::Relaxed) {
                return ptr::null();
            }

            let res = self.free.compare_exchange_weak(
                old,
                new,
                Ordering::SeqCst,
                Ordering::Relaxed,
            );

            match res {
                Ok(_) => break,
                Err(x) => old = x,
            }
        }

        old as *const u8
    }

    pub fn collect(
        &self,
        ctxt: &SemContext,
        rootset: Vec<IndirectObj>,
        card_table: &CardTable,
        crossing_map: &CrossingMap,
        old: &OldGen,
    ) {
        let mut timer = Timer::new(ctxt.args.flag_gc_events);

        let mut scan;

        let mut free;
        let new_end;

        if self.end.load(Ordering::Relaxed) == self.separator.to_usize() {
            scan = self.total.start;
            free = self.separator;
            new_end = self.total.end;

        } else {
            scan = self.separator;
            free = self.total.start;
            new_end = self.separator;
        }

        if cfg!(debug_assertions) {
            // make memory readable & writable again, so that we
            // can copy objects to the to-space.
            // Since this has some overhead, do it only in debug builds.
            os::mprotect(free.to_ptr::<u8>(), self.size, ProtType::Writable);
        }

        // detect all references from roots into young generation
        for &root in &rootset {
            let root_ptr = root.get();

            if self.total.includes(Address::from_ptr(root_ptr)) {
                root.set(copy(root_ptr, &mut free, old));
            }
        }

        // detect references from old generation (dirty cards) into young generation
        copy_dirty_cards(card_table, crossing_map, &mut free, &self.total, old);

        // visit all fields in copied objects
        while scan < free {
            let object = unsafe { &mut *scan.to_mut_ptr::<Obj>() };

            object.visit_reference_fields(|child| {
                let child_ptr = child.get();

                if self.total.includes(Address::from_ptr(child_ptr)) {
                    child.set(copy(child_ptr, &mut free, old));
                }
            });

            scan = scan.offset(object.size());
        }

        // swap spaces
        self.end.store(new_end.to_usize(), Ordering::Relaxed);

        // update start of free memory
        self.free.store(free.to_usize(), Ordering::SeqCst);

        // Make from-space unaccessible both from read/write.
        // Since this has some overhead, do it only in debug builds.
        if cfg!(debug_assertions) {
            let start = if new_end == self.separator {
                self.total.start
            } else {
                self.separator
            };

            os::mprotect(start.to_ptr::<u8>(), self.size, ProtType::None);
        }

        timer.stop_with(|dur| if ctxt.args.flag_gc_events {
            println!("GC minor: collect garbage ({} ms)", in_ms(dur));
        });
    }
}

// copy all references from old- into young-generation.
fn copy_dirty_cards(
    card_table: &CardTable,
    crossing_map: &CrossingMap,
    mut free: &mut Address,
    young: &Region,
    old: &OldGen,
) {
    card_table.visit_dirty(|card| {
        let crossing_entry = crossing_map.get(card);

        // card contains: any data but no references, then first object
        if crossing_entry.is_first_object() {
            let card_address = crossing_map.address_of_card(card);

            // first_object() returns number of words before end of card
            let offset_from_end = crossing_entry.first_object() as usize * ptr_width();
            let ptr = card_address.offset(CARD_SIZE - offset_from_end);

            let end = card_address.offset(CARD_SIZE);

            // copy all objects from this card
            copy_card(ptr, end, free, young, old);

        // card contains: references, then first object
        } else if crossing_entry.is_references_at_start() {
            let number_references = crossing_entry.references_at_start() as usize;
            let card_address = crossing_map.address_of_card(card);
            let mut ptr = card_address;

            for _ in 0..number_references {
                let ind_ptr = IndirectObj::from_address(ptr);
                let dir_ptr = ind_ptr.get();

                if young.includes(Address::from_ptr(dir_ptr)) {
                    ind_ptr.set(copy(dir_ptr, &mut free, old));
                }

                ptr = ptr.offset(ptr_width());
            }

            let end = card_address.offset(CARD_SIZE);

            // copy all objects from this card
            copy_card(ptr, end, free, young, old);

        // object spans multiple cards
        } else if crossing_entry.is_previous_card() {
            let mut previous = card - crossing_entry.previous_card() as usize;
            let mut crossing_entry = crossing_map.get(previous);

            while crossing_entry.is_previous_card() {
                previous = previous - crossing_entry.previous_card() as usize;
                crossing_entry = crossing_map.get(previous);
            }

            let previous_address = crossing_map.address_of_card(previous);
            let mut ptr;

            if crossing_entry.is_references_at_start() {
                let number_references = crossing_entry.references_at_start() as usize;
                ptr = previous_address.offset(number_references * ptr_width());

            } else if crossing_entry.is_first_object() {
                let offset_from_end = crossing_entry.first_object() as usize * ptr_width();
                ptr = previous_address.offset(CARD_SIZE - offset_from_end);

            } else {
                assert!(crossing_entry.is_no_references());
                panic!("dirty card without references: is this allowed to happen?");
            }

            let card_address = crossing_map.address_of_card(card);
            let mut last_object = ptr;

            while ptr < card_address {
                let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

                last_object = ptr;
                ptr = ptr.offset(object.size());
            }

            // last_object now stores pointer that reaches into current card
            // TODO: handle fields in last_object

            copy_card(ptr, card_address.offset(CARD_SIZE), free, young, old);

        } else {
            assert!(crossing_entry.is_no_references());
            panic!("dirty card without references: is this allowed to happen?");
        }
    });
}

fn copy_card(mut ptr: Address, end: Address, mut free: &mut Address, young: &Region, old: &OldGen) {
    while ptr < end {
        let object = unsafe { &mut *ptr.to_mut_ptr::<Obj>() };

        object.visit_reference_fields(|child| {
            let child_ptr = child.get();

            if young.includes(Address::from_ptr(child_ptr)) {
                child.set(copy(child_ptr, &mut free, old));
            }
        });

        ptr = ptr.offset(object.size());
    }
}

pub fn copy(obj: *mut Obj, free: &mut Address, old: &OldGen) -> *mut Obj {
    let obj = unsafe { &mut *obj };
    let addr = get_forwarding_address(obj);

    if is_forwarding_address(addr) {
        unmark_forwarding_address(addr).to_mut_ptr::<Obj>()
    } else {
        let obj_size = obj.size();
        let age = obj.header_mut().increase_age();

        let addr: Address;

        // if object is old enough we copy it into the old generation
        if age >= PROMOTION_AGE {
            addr = Address::from_ptr(old.alloc(obj_size));
            copy_object(obj, addr, obj_size);

        // otherwise the object remains in the young generation for now
        } else {
            addr = *free;
            copy_object(obj, addr, obj_size);
            *free = addr.offset(obj_size);
        }

        set_forwarding_address(obj, addr);

        addr.to_mut_ptr::<Obj>()
    }
}

fn copy_object(obj: &Obj, addr: Address, size: usize) {
    unsafe {
        ptr::copy_nonoverlapping(
            obj as *const Obj as *const u8,
            addr.to_mut_ptr::<u8>(),
            size,
        );
    }
}

pub fn is_forwarding_address(obj: Address) -> bool {
    (obj.to_usize() & 1) == 1
}

pub fn mark_forwarding_address(obj: Address) -> Address {
    (obj.to_usize() | 1).into()
}

pub fn unmark_forwarding_address(obj: Address) -> Address {
    (obj.to_usize() & !1).into()
}

pub fn get_forwarding_address(obj: &Obj) -> Address {
    unsafe { *(obj as *const Obj as *const Address) }
}

pub fn set_forwarding_address(obj: &mut Obj, addr: Address) {
    unsafe {
        *(obj as *mut Obj as *mut Address) = mark_forwarding_address(addr);
    }
}

#[inline(always)]
fn ptr_width() -> usize {
    mem::ptr_width() as usize
}<|endoftext|>


<commit_before>// http://rosettacode.org/wiki/Hamming_numbers
extern crate num;
use num::bigint::{BigUint, ToBigUint};
use std::cmp::min;
use std::num::{One, one};
use std::collections::{RingBuf, Deque};

// needed because hamming_numbers_alt uses this as a library
#[allow(dead_code)]
#[cfg(not(test))]
fn main() {
    static HAMMING_1_MILLION: &'static str =
"519312780448388736089589843750000000000000000000000000000000000000000000000000000000";
    // capacity of the queue currently needs to be a power of 2 because of a bug with RingBuf
    let hamming : Hamming<BigUint> = Hamming::new(128);

    for (idx, h) in hamming.enumerate().take(1_000_000) {
        match idx + 1 {
            1...20 => print!("{} ", h.to_biguint().unwrap()),
            i @ 1691 => println!("\n{}th number: {}", i, h.to_biguint().unwrap()),
            i @ 1000000 => {
                let n = h.to_biguint().unwrap();
                println!("\n{}th number: {}", i, n);
                assert_eq!(n, from_str(HAMMING_1_MILLION).unwrap())
            },
            _ =>  continue
        }
    }
}

//representing a Hamming number as a BigUint
impl HammingNumber for BigUint {
    // returns the multipliers 2, 3 and 5 in the representation for the HammingNumber
    fn multipliers()-> (BigUint, BigUint, BigUint) {
        (2u.to_biguint().unwrap(),
        3u.to_biguint().unwrap(),
        5u.to_biguint().unwrap())
    }
}

/// representation of a Hamming number
/// allows to abstract on how the hamming number is stored
/// i.e. as BigUint directly or just as the powers of 2, 3 and 5 used to build it
pub trait HammingNumber : Eq + Ord + ToBigUint + Mul<Self, Self> + One {
    fn multipliers() -> (Self, Self, Self);
}

/// Hamming numbers are multiples of 2, 3 or 5.
///
/// We keep them on three queues and extract the lowest (leftmost) value from
/// the three queues at each iteration.
pub struct Hamming<T> {
    // Using a RingBuf as a queue, push to the back, pop from the front
    q2: RingBuf<T>,
    q3: RingBuf<T>,
    q5: RingBuf<T>
}

impl<T: HammingNumber> Hamming<T> {
    /// Static constructor method
    /// `n` initializes the capacity of the queues
    pub fn new(n: uint) -> Hamming<T> {
        let mut h = Hamming {
            q2: RingBuf::with_capacity(n),
            q3: RingBuf::with_capacity(n),
            q5: RingBuf::with_capacity(n)
        };

        h.q2.push(one());
        h.q3.push(one());
        h.q5.push(one());

        h
    }

    /// Pushes the next multiple of `n` (x2, x3, x5) to the queues
    pub fn enqueue(&mut self, n: &T) {
        let (two, three, five) = HammingNumber::multipliers();
        self.q2.push(*n * two);
        self.q3.push(*n * three);
        self.q5.push(*n * five);
    }
}

// Implements the `Iterator` trait, so we can generate Hamming numbers lazily
impl<T: HammingNumber> Iterator<T> for Hamming<T> {
    // The core of the work is done in the `next` method.
    // We check which of the 3 queues has the lowest candidate and extract it
    // as the next Hamming number.
    fn next(&mut self) -> Option<T> {
        // Return `pop_targets` so the borrow from `front()` will be finished
        let (two, three, five) = match (self.q2.front(),
                                        self.q3.front(),
                                        self.q5.front()) {
            (Some(head2), Some(head3), Some(head5)) => {
                let n = min(head2, min(head3, head5));
                (head2 == n, head3 == n, head5 == n)
            },
            _ => unreachable!()
        };

        let h2 = if two { self.q2.pop_front() } else { None };
        let h3 = if three { self.q3.pop_front() } else { None };
        let h5 = if five { self.q5.pop_front() } else { None };

        match h2.or(h3).or(h5) {
            Some(n) => {
                self.enqueue(&n);
                Some(n)
            }
            None => unreachable!()
        }
    }
}

#[test]
fn create() {
    let mut h = Hamming::<BigUint>::new(5);
    h.q2.push(one::<BigUint>());
    h.q2.push(one::<BigUint>() * 3u.to_biguint().unwrap());

    assert_eq!(h.q2.pop_front().unwrap(), one::<BigUint>());
}

#[test]
fn try_enqueue() {
    let mut h = Hamming::<BigUint>::new(5);
    let (two, three, five) = HammingNumber::multipliers();
    h.enqueue(&one::<BigUint>());
    h.enqueue(&(&one::<BigUint>() * two));

    assert!(h.q2.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q3.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q5.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q2.pop_front().unwrap() == one::<BigUint>() * two);
    assert!(h.q3.pop_front().unwrap() == one::<BigUint>() * three);
    assert!(h.q5.pop_front().unwrap() == one::<BigUint>() * five);
 }

#[test]
fn hamming_iter() {
    let mut hamming = Hamming::<BigUint>::new(20);
    assert!(hamming.nth(19).unwrap().to_biguint() == 36u.to_biguint());
}
<commit_msg>Added clarifying comment about large hamming number test.<commit_after>// http://rosettacode.org/wiki/Hamming_numbers
// NOTE: If you change this file please make sure you run main().  It has a larger test that can
// trip up on large floating point values, but is too expensive to run every time we run the tests.
extern crate num;
use num::bigint::{BigUint, ToBigUint};
use std::cmp::min;
use std::num::{One, one};
use std::collections::{RingBuf, Deque};

// needed because hamming_numbers_alt uses this as a library
#[allow(dead_code)]
#[cfg(not(test))]
fn main() {
    static HAMMING_1_MILLION: &'static str =
"519312780448388736089589843750000000000000000000000000000000000000000000000000000000";
    // capacity of the queue currently needs to be a power of 2 because of a bug with RingBuf
    let hamming : Hamming<BigUint> = Hamming::new(128);

    for (idx, h) in hamming.enumerate().take(1_000_000) {
        match idx + 1 {
            1...20 => print!("{} ", h.to_biguint().unwrap()),
            i @ 1691 => println!("\n{}th number: {}", i, h.to_biguint().unwrap()),
            i @ 1000000 => {
                let n = h.to_biguint().unwrap();
                println!("\n{}th number: {}", i, n);
                assert_eq!(n, from_str(HAMMING_1_MILLION).unwrap())
            },
            _ =>  continue
        }
    }
}

//representing a Hamming number as a BigUint
impl HammingNumber for BigUint {
    // returns the multipliers 2, 3 and 5 in the representation for the HammingNumber
    fn multipliers()-> (BigUint, BigUint, BigUint) {
        (2u.to_biguint().unwrap(),
        3u.to_biguint().unwrap(),
        5u.to_biguint().unwrap())
    }
}

/// representation of a Hamming number
/// allows to abstract on how the hamming number is stored
/// i.e. as BigUint directly or just as the powers of 2, 3 and 5 used to build it
pub trait HammingNumber : Eq + Ord + ToBigUint + Mul<Self, Self> + One {
    fn multipliers() -> (Self, Self, Self);
}

/// Hamming numbers are multiples of 2, 3 or 5.
///
/// We keep them on three queues and extract the lowest (leftmost) value from
/// the three queues at each iteration.
pub struct Hamming<T> {
    // Using a RingBuf as a queue, push to the back, pop from the front
    q2: RingBuf<T>,
    q3: RingBuf<T>,
    q5: RingBuf<T>
}

impl<T: HammingNumber> Hamming<T> {
    /// Static constructor method
    /// `n` initializes the capacity of the queues
    pub fn new(n: uint) -> Hamming<T> {
        let mut h = Hamming {
            q2: RingBuf::with_capacity(n),
            q3: RingBuf::with_capacity(n),
            q5: RingBuf::with_capacity(n)
        };

        h.q2.push(one());
        h.q3.push(one());
        h.q5.push(one());

        h
    }

    /// Pushes the next multiple of `n` (x2, x3, x5) to the queues
    pub fn enqueue(&mut self, n: &T) {
        let (two, three, five) = HammingNumber::multipliers();
        self.q2.push(*n * two);
        self.q3.push(*n * three);
        self.q5.push(*n * five);
    }
}

// Implements the `Iterator` trait, so we can generate Hamming numbers lazily
impl<T: HammingNumber> Iterator<T> for Hamming<T> {
    // The core of the work is done in the `next` method.
    // We check which of the 3 queues has the lowest candidate and extract it
    // as the next Hamming number.
    fn next(&mut self) -> Option<T> {
        // Return `pop_targets` so the borrow from `front()` will be finished
        let (two, three, five) = match (self.q2.front(),
                                        self.q3.front(),
                                        self.q5.front()) {
            (Some(head2), Some(head3), Some(head5)) => {
                let n = min(head2, min(head3, head5));
                (head2 == n, head3 == n, head5 == n)
            },
            _ => unreachable!()
        };

        let h2 = if two { self.q2.pop_front() } else { None };
        let h3 = if three { self.q3.pop_front() } else { None };
        let h5 = if five { self.q5.pop_front() } else { None };

        match h2.or(h3).or(h5) {
            Some(n) => {
                self.enqueue(&n);
                Some(n)
            }
            None => unreachable!()
        }
    }
}

#[test]
fn create() {
    let mut h = Hamming::<BigUint>::new(5);
    h.q2.push(one::<BigUint>());
    h.q2.push(one::<BigUint>() * 3u.to_biguint().unwrap());

    assert_eq!(h.q2.pop_front().unwrap(), one::<BigUint>());
}

#[test]
fn try_enqueue() {
    let mut h = Hamming::<BigUint>::new(5);
    let (two, three, five) = HammingNumber::multipliers();
    h.enqueue(&one::<BigUint>());
    h.enqueue(&(&one::<BigUint>() * two));

    assert!(h.q2.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q3.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q5.pop_front().unwrap() == one::<BigUint>());
    assert!(h.q2.pop_front().unwrap() == one::<BigUint>() * two);
    assert!(h.q3.pop_front().unwrap() == one::<BigUint>() * three);
    assert!(h.q5.pop_front().unwrap() == one::<BigUint>() * five);
 }

#[test]
fn hamming_iter() {
    let mut hamming = Hamming::<BigUint>::new(20);
    assert!(hamming.nth(19).unwrap().to_biguint() == 36u.to_biguint());
}
<|endoftext|>


<commit_before>use std::sync::Mutex;
use std::collections::HashMap;
use circular_queue::CircularQueue;
use regex::{Regex, RegexBuilder};

use irc::client::prelude::*;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

lazy_static! {
    static ref RE: Regex = Regex::new(r"^s/((?:\\/|[^/])+)/((?:\\/|[^/])*)/(?:(\w+))?\s*$").unwrap();
}

#[derive(PluginName, Debug)]
pub struct Sed {
    per_channel: usize,
    channel_messages: Mutex<HashMap<String, CircularQueue<String>>>,
}

macro_rules! try_lock {
    ( $m:expr ) => {
        match $m.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

impl Sed {
    pub fn new(per_channel: usize) -> Sed {
        Sed {
            per_channel: per_channel,
            channel_messages: Mutex::new(HashMap::new()),
        }
    }

    fn add_message(&self, channel: String, message: String) {
        let mut channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages
            .entry(channel)
            .or_insert(CircularQueue::with_capacity(self.per_channel));
        messages.push(message);
    }

    fn run_regex(&self, channel: &str, message: &str) -> Result<String, SedError> {
        let mut global_match = false;
        let mut case_insens = false;
        let mut ign_whitespace = false;
        let mut swap_greed = false;
        let mut disable_unicode = false;

        let captures = RE.captures(message).unwrap();
        debug!("{:?}", captures);

        let first = captures.get(1).unwrap().as_str().replace(r"\/", "/");
        let second = captures.get(2).unwrap().as_str().replace(r"\/", "/");

        if let Some(flags) = captures.get(3) {
            let flags = flags.as_str();

            global_match = flags.contains('g');
            case_insens = flags.contains('i');
            ign_whitespace = flags.contains('x');
            swap_greed = flags.contains('U');
            disable_unicode = !flags.contains('u');
        }

        let user_re = RegexBuilder::new(&first)
            .case_insensitive(case_insens)
            .ignore_whitespace(ign_whitespace)
            .unicode(disable_unicode)
            .swap_greed(swap_greed)
            .build()
            .context(ErrorKind::InvalidRegex)?;

        let channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages.get(channel).ok_or(ErrorKind::NoMessages)?;

        for message in messages.iter() {
            if user_re.is_match(message) {
                let response = if global_match {
                    user_re.replace_all(message, &second[..])
                } else {
                    user_re.replace(message, &second[..])
                };

                return Ok(response.to_string());
            }
        }

        Err(ErrorKind::NoMatch)?
    }
}

impl Plugin for Sed {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                let channel = message.response_target().unwrap();
                if channel == message.source_nickname().unwrap() {
                    return ExecutionStatus::Done;
                }

                if RE.is_match(content) {
                    let result = match self.run_regex(channel, content) {
                        Ok(msg) => client.send_privmsg(channel, &msg),
                        Err(e) => client.send_notice(channel, &e.to_string()),
                    };

                    match result {
                        Err(e) => {
                            ExecutionStatus::Err(e.context(FrippyErrorKind::Connection).into())
                        }
                        Ok(_) => ExecutionStatus::Done,
                    }
                } else {
                    self.add_message(channel.to_string(), content.to_string());

                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Sed should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        Ok(client
            .send_notice(
                &command.source,
                "Currently this Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for sed at this time",
        ))
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "SedError"]
    pub enum ErrorKind {
        /// Invalid regex error
        #[fail(display = "Invalid regex")]
        InvalidRegex,

        /// No messages found error
        #[fail(display = "No messages were found for this channel")]
        NoMessages,

        /// No match found error
        #[fail(display = "No recent messages match this regex")]
        NoMatch,
    }
}
<commit_msg>Fix escaped backslashes in sed strings<commit_after>use std::sync::Mutex;
use std::collections::HashMap;
use circular_queue::CircularQueue;
use regex::{Regex, RegexBuilder};

use irc::client::prelude::*;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

lazy_static! {
    static ref RE: Regex = Regex::new(r"^s/((?:\\/|[^/])+)/((?:\\/|[^/])*)/(?:(\w+))?\s*$").unwrap();
}

#[derive(PluginName, Debug)]
pub struct Sed {
    per_channel: usize,
    channel_messages: Mutex<HashMap<String, CircularQueue<String>>>,
}

macro_rules! try_lock {
    ( $m:expr ) => {
        match $m.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

impl Sed {
    pub fn new(per_channel: usize) -> Sed {
        Sed {
            per_channel: per_channel,
            channel_messages: Mutex::new(HashMap::new()),
        }
    }

    fn add_message(&self, channel: String, message: String) {
        let mut channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages
            .entry(channel)
            .or_insert(CircularQueue::with_capacity(self.per_channel));
        messages.push(message);
    }

    fn format_escaped(&self, input: &str) -> String {
        let mut output = String::with_capacity(input.len());
        let mut escape = false;

        for c in input.chars() {
            if !escape && c == '\\' {
                escape = true;
                continue;
            }
            escape = false;

            output.push(c);
        }

        output
    }

    fn run_regex(&self, channel: &str, message: &str) -> Result<String, SedError> {
        let mut global_match = false;
        let mut case_insens = false;
        let mut ign_whitespace = false;
        let mut swap_greed = false;
        let mut disable_unicode = false;

        let captures = RE.captures(message).unwrap();
        debug!("{:?}", captures);

        let first = self.format_escaped(captures.get(1).unwrap().as_str());
        let second = self.format_escaped(captures.get(2).unwrap().as_str());

        if let Some(flags) = captures.get(3) {
            let flags = flags.as_str();

            global_match = flags.contains('g');
            case_insens = flags.contains('i');
            ign_whitespace = flags.contains('x');
            swap_greed = flags.contains('U');
            disable_unicode = !flags.contains('u');
        }

        let user_re = RegexBuilder::new(&first)
            .case_insensitive(case_insens)
            .ignore_whitespace(ign_whitespace)
            .unicode(disable_unicode)
            .swap_greed(swap_greed)
            .build()
            .context(ErrorKind::InvalidRegex)?;

        let channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages.get(channel).ok_or(ErrorKind::NoMessages)?;

        for message in messages.iter() {
            if user_re.is_match(message) {
                let response = if global_match {
                    user_re.replace_all(message, &second[..])
                } else {
                    user_re.replace(message, &second[..])
                };

                return Ok(response.to_string());
            }
        }

        Err(ErrorKind::NoMatch)?
    }
}

impl Plugin for Sed {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                let channel = message.response_target().unwrap();
                if channel == message.source_nickname().unwrap() {
                    return ExecutionStatus::Done;
                }

                if RE.is_match(content) {
                    let result = match self.run_regex(channel, content) {
                        Ok(msg) => client.send_privmsg(channel, &msg),
                        Err(e) => client.send_notice(channel, &e.to_string()),
                    };

                    match result {
                        Err(e) => {
                            ExecutionStatus::Err(e.context(FrippyErrorKind::Connection).into())
                        }
                        Ok(_) => ExecutionStatus::Done,
                    }
                } else {
                    self.add_message(channel.to_string(), content.to_string());

                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Sed should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        Ok(client
            .send_notice(
                &command.source,
                "Currently this Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for sed at this time",
        ))
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "SedError"]
    pub enum ErrorKind {
        /// Invalid regex error
        #[fail(display = "Invalid regex")]
        InvalidRegex,

        /// No messages found error
        #[fail(display = "No messages were found for this channel")]
        NoMessages,

        /// No match found error
        #[fail(display = "No recent messages match this regex")]
        NoMatch,
    }
}
<|endoftext|>


<commit_before>extern crate htmlescape;

use std::marker::PhantomData;
use std::time::Duration;

use irc::client::prelude::*;

use regex::Regex;

use plugin::*;
use utils::Url;
use FrippyClient;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

lazy_static! {
    static ref URL_RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap();
    static ref WORD_RE: Regex = Regex::new(r"(\w+)").unwrap();
}

#[derive(PluginName, Debug)]
pub struct UrlTitles<C> {
    max_kib: usize,
    phantom: PhantomData<C>,
}

#[derive(Clone, Debug)]
struct Title(String, Option<usize>);

impl From<String> for Title {
    fn from(title: String) -> Self {
        Title(title, None)
    }
}

impl From<Title> for String {
    fn from(title: Title) -> Self {
        title.0
    }
}

impl Title {
    fn find_by_delimiters(body: &str, delimiters: [&str; 3]) -> Result<Self, UrlError> {
        let title = body.find(delimiters[0])
            .map(|tag| {
                body[tag..]
                    .find(delimiters[1])
                    .map(|offset| tag + offset + delimiters[1].len())
                    .map(|start| {
                        body[start..]
                            .find(delimiters[2])
                            .map(|offset| start + offset)
                            .map(|end| &body[start..end])
                    })
            })
            .and_then(|s| s.and_then(|s| s))
            .ok_or(ErrorKind::MissingTitle)?;

        debug!("Found title {:?} with delimiters {:?}", title, delimiters);

        htmlescape::decode_html(title)
            .map(|t| t.into())
            .map_err(|_| ErrorKind::HtmlDecoding.into())
    }

    fn find_ogtitle(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["property=\"og:title\"", "content=\"", "\""])
    }

    fn find_title(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["<title", ">", "</title>"])
    }

    // TODO Improve logic
    fn get_usefulness(self, url: &str) -> Self {
        let mut usefulness = 0;
        for word in WORD_RE.find_iter(&self.0) {
            let w = word.as_str().to_lowercase();
            if w.len() > 2 && !url.to_lowercase().contains(&w) {
                usefulness += 1;
            }
        }

        Title(self.0, Some(usefulness))
    }

    pub fn usefulness(&self) -> usize {
        self.1.expect("Usefulness should be calculated already")
    }

    fn clean_up(self) -> Self {
        Title(self.0.trim().replace('\n', "|").replace('\r', "|"), self.1)
    }

    pub fn find_clean_ogtitle(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_ogtitle(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }

    pub fn find_clean_title(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_title(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }
}

impl<C: FrippyClient> UrlTitles<C> {
    /// If a file is larger than `max_kib` KiB the download is stopped
    pub fn new(max_kib: usize) -> Self {
        UrlTitles {
            max_kib,
            phantom: PhantomData,
        }
    }

    fn grep_url<'a>(&self, msg: &'a str) -> Option<Url<'a>> {
        let captures = URL_RE.captures(msg)?;
        debug!("Url captures: {:?}", captures);

        Some(captures.get(2)?.as_str().into())
    }

    fn url(&self, text: &str) -> Result<String, UrlError> {
        let url = self.grep_url(text)
            .ok_or(ErrorKind::MissingUrl)?
            .max_kib(self.max_kib)
            .timeout(Duration::from_secs(5));
        let body = url.request().context(ErrorKind::Download)?;

        let title = Title::find_clean_title(&body, url.as_str());
        let og_title = Title::find_clean_ogtitle(&body, url.as_str());

        let title = match (title, og_title) {
            (Ok(title), Ok(og_title)) => {
                if title.usefulness() > og_title.usefulness() {
                    title
                } else {
                    og_title
                }
            }
            (Ok(title), _) => title,
            (_, Ok(title)) => title,
            (Err(e), _) => Err(e)?,
        };

        if title.usefulness() > 1 {
            Ok(title.into())
        } else {
            Err(ErrorKind::UselessTitle.into())
        }
    }
}

impl<C: FrippyClient> Plugin for UrlTitles<C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref msg) => if URL_RE.is_match(msg) {
                ExecutionStatus::RequiresThread
            } else {
                ExecutionStatus::Done
            },
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, ref content) = message.command {
            let title = self.url(content).context(FrippyErrorKind::Url)?;
            let response = format!("[URL] {}", title);

            client
                .send_privmsg(message.response_target().unwrap(), &response)
                .context(FrippyErrorKind::Connection)?;
        }

        Ok(())
    }

    fn command(&self, client: &Self::Client, command: PluginCommand) -> Result<(), FrippyError> {
        client
            .send_notice(
                &command.source,
                "This Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?;

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, command: PluginCommand) -> Result<String, String> {
        self.url(&command.tokens[0])
            .map_err(|e| e.cause().unwrap().to_string())
    }
}

pub mod error {
    /// A URL plugin error
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "UrlError"]
    pub enum ErrorKind {
        /// A download error
        #[fail(display = "A download error occured")]
        Download,

        /// Missing URL error
        #[fail(display = "No URL was found")]
        MissingUrl,

        /// Missing title error
        #[fail(display = "No title was found")]
        MissingTitle,

        /// Useless title error
        #[fail(display = "The titles found were not useful enough")]
        UselessTitle,

        /// Html decoding error
        #[fail(display = "Failed to decode Html characters")]
        HtmlDecoding,
    }
}
<commit_msg>Url: Send titles with a sigle useful word again<commit_after>extern crate htmlescape;

use std::marker::PhantomData;
use std::time::Duration;

use irc::client::prelude::*;

use regex::Regex;

use plugin::*;
use utils::Url;
use FrippyClient;

use self::error::*;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use failure::Fail;
use failure::ResultExt;

lazy_static! {
    static ref URL_RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap();
    static ref WORD_RE: Regex = Regex::new(r"(\w+)").unwrap();
}

#[derive(PluginName, Debug)]
pub struct UrlTitles<C> {
    max_kib: usize,
    phantom: PhantomData<C>,
}

#[derive(Clone, Debug)]
struct Title(String, Option<usize>);

impl From<String> for Title {
    fn from(title: String) -> Self {
        Title(title, None)
    }
}

impl From<Title> for String {
    fn from(title: Title) -> Self {
        title.0
    }
}

impl Title {
    fn find_by_delimiters(body: &str, delimiters: [&str; 3]) -> Result<Self, UrlError> {
        let title = body
            .find(delimiters[0])
            .map(|tag| {
                body[tag..]
                    .find(delimiters[1])
                    .map(|offset| tag + offset + delimiters[1].len())
                    .map(|start| {
                        body[start..]
                            .find(delimiters[2])
                            .map(|offset| start + offset)
                            .map(|end| &body[start..end])
                    })
            })
            .and_then(|s| s.and_then(|s| s))
            .ok_or(ErrorKind::MissingTitle)?;

        debug!("Found title {:?} with delimiters {:?}", title, delimiters);

        htmlescape::decode_html(title)
            .map(|t| t.into())
            .map_err(|_| ErrorKind::HtmlDecoding.into())
    }

    fn find_ogtitle(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["property=\"og:title\"", "content=\"", "\""])
    }

    fn find_title(body: &str) -> Result<Self, UrlError> {
        Self::find_by_delimiters(body, ["<title", ">", "</title>"])
    }

    // TODO Improve logic
    fn get_usefulness(self, url: &str) -> Self {
        let mut usefulness = 0;
        for word in WORD_RE.find_iter(&self.0) {
            let w = word.as_str().to_lowercase();
            if w.len() > 2 && !url.to_lowercase().contains(&w) {
                usefulness += 1;
            }
        }

        Title(self.0, Some(usefulness))
    }

    pub fn usefulness(&self) -> usize {
        self.1.expect("Usefulness should be calculated already")
    }

    fn clean_up(self) -> Self {
        Title(self.0.trim().replace('\n', "|").replace('\r', "|"), self.1)
    }

    pub fn find_clean_ogtitle(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_ogtitle(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }

    pub fn find_clean_title(body: &str, url: &str) -> Result<Self, UrlError> {
        let title = Self::find_title(body)?;
        Ok(title.get_usefulness(url).clean_up())
    }
}

impl<C: FrippyClient> UrlTitles<C> {
    /// If a file is larger than `max_kib` KiB the download is stopped
    pub fn new(max_kib: usize) -> Self {
        UrlTitles {
            max_kib,
            phantom: PhantomData,
        }
    }

    fn grep_url<'a>(&self, msg: &'a str) -> Option<Url<'a>> {
        let captures = URL_RE.captures(msg)?;
        debug!("Url captures: {:?}", captures);

        Some(captures.get(2)?.as_str().into())
    }

    fn url(&self, text: &str) -> Result<String, UrlError> {
        let url = self
            .grep_url(text)
            .ok_or(ErrorKind::MissingUrl)?
            .max_kib(self.max_kib)
            .timeout(Duration::from_secs(5));
        let body = url.request().context(ErrorKind::Download)?;

        let title = Title::find_clean_title(&body, url.as_str());
        let og_title = Title::find_clean_ogtitle(&body, url.as_str());

        let title = match (title, og_title) {
            (Ok(title), Ok(og_title)) => {
                if title.usefulness() > og_title.usefulness() {
                    title
                } else {
                    og_title
                }
            }
            (Ok(title), _) => title,
            (_, Ok(title)) => title,
            (Err(e), _) => Err(e)?,
        };

        Ok(title.into())
    }
}

impl<C: FrippyClient> Plugin for UrlTitles<C> {
    type Client = C;
    fn execute(&self, _: &Self::Client, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref msg) => {
                if URL_RE.is_match(msg) {
                    ExecutionStatus::RequiresThread
                } else {
                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(
        &self,
        client: &Self::Client,
        message: &Message,
    ) -> Result<(), FrippyError> {
        if let Command::PRIVMSG(_, ref content) = message.command {
            let title = self.url(content).context(FrippyErrorKind::Url)?;
            let response = format!("[URL] {}", title);

            client
                .send_privmsg(message.response_target().unwrap(), &response)
                .context(FrippyErrorKind::Connection)?;
        }

        Ok(())
    }

    fn command(&self, client: &Self::Client, command: PluginCommand) -> Result<(), FrippyError> {
        client
            .send_notice(
                &command.source,
                "This Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?;

        Ok(())
    }

    fn evaluate(&self, _: &Self::Client, command: PluginCommand) -> Result<String, String> {
        self.url(&command.tokens[0])
            .map_err(|e| e.cause().unwrap().to_string())
    }
}

pub mod error {
    /// A URL plugin error
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "UrlError"]
    pub enum ErrorKind {
        /// A download error
        #[fail(display = "A download error occured")]
        Download,

        /// Missing URL error
        #[fail(display = "No URL was found")]
        MissingUrl,

        /// Missing title error
        #[fail(display = "No title was found")]
        MissingTitle,

        /// Useless title error
        #[fail(display = "The titles found were not useful enough")]
        UselessTitle,

        /// Html decoding error
        #[fail(display = "Failed to decode Html characters")]
        HtmlDecoding,
    }
}
<|endoftext|>


<commit_before>use crate::{
    model::{GameId, Mode, PersonalEntry, PersonalGroup, Speed, Stats, UserName},
    opening::Opening,
};
use serde::{Deserialize, Serialize};
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, FromInto, StringWithSeparator};
use shakmaty::{
    fen::{Fen, ParseFenError},
    san::SanPlus,
    uci::Uci,
    variant::VariantPosition,
    Color,
};
use std::str::FromStr;

mod error;
mod variant;

pub use error::Error;
pub use variant::LilaVariant;

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PersonalQuery {
    #[serde(default)]
    pub variant: LilaVariant,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub fen: Option<LaxFen>,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    pub play: Vec<Uci>,
    #[serde_as(as = "DisplayFromStr")]
    pub player: UserName,
    #[serde_as(as = "FromInto<ColorProxy>")]
    pub color: Color,
    #[serde(default)]
    pub since: u32, // year
    #[serde(flatten)]
    pub filter: PersonalQueryFilter,
    #[serde(default)]
    pub update: bool,
}

#[derive(Deserialize, Debug)]
pub struct PersonalQueryFilter {
    #[serde(default)]
    pub modes: Option<Vec<Mode>>,
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
}

#[derive(Serialize, Debug)]
pub struct PersonalResponse {
    #[serde(flatten)]
    total: Stats,
    moves: Vec<PersonalMoveRow>,
    opening: Option<&'static Opening>,
}

#[serde_as]
#[derive(Serialize, Debug)]
struct PersonalMoveRow {
    #[serde_as(as = "DisplayFromStr")]
    uci: Uci,
    #[serde_as(as = "DisplayFromStr")]
    san: SanPlus,
    #[serde(flatten)]
    stats: Stats,
    #[serde_as(as = "Option<DisplayFromStr>")]
    game: Option<GameId>,
}

impl PersonalQueryFilter {
    pub fn respond(
        &self,
        pos: VariantPosition,
        entry: PersonalEntry,
        opening: Option<&'static Opening>,
    ) -> PersonalResponse {
        let mut total = Stats::default();

        let mut moves = Vec::with_capacity(entry.sub_entries.len());

        for (uci, sub_entry) in entry.sub_entries {
            let m = uci.to_move(&pos).expect("legal uci in personal entry");
            let san = SanPlus::from_move(pos.clone(), &m);

            let mut group = PersonalGroup::default();

            for speed in Speed::ALL {
                if self
                    .speeds
                    .as_ref()
                    .map_or(true, |speeds| speeds.contains(&speed))
                {
                    for mode in Mode::ALL {
                        if self
                            .modes
                            .as_ref()
                            .map_or(true, |modes| modes.contains(&mode))
                        {
                            group += sub_entry.by_speed(speed).by_mode(mode).to_owned();
                            // TODO: Optimize?
                        }
                    }
                }
            }

            moves.push(PersonalMoveRow {
                uci,
                san,
                stats: group.stats.to_owned(),
                game: group
                    .games
                    .into_iter()
                    .max_by_key(|(idx, _id)| *idx)
                    .map(|(_, id)| id),
            });

            total += group.stats;
        }

        moves.sort_by_key(|row| row.stats.total());

        PersonalResponse {
            total,
            moves,
            opening,
        }
    }
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum ColorProxy {
    White,
    Black,
}

impl From<Color> for ColorProxy {
    fn from(color: Color) -> ColorProxy {
        match color {
            Color::White => ColorProxy::White,
            Color::Black => ColorProxy::Black,
        }
    }
}

impl From<ColorProxy> for Color {
    fn from(color: ColorProxy) -> Color {
        match color {
            ColorProxy::White => Color::White,
            ColorProxy::Black => Color::Black,
        }
    }
}

#[derive(Debug)]
pub struct LaxFen(Fen);

impl From<Fen> for LaxFen {
    fn from(fen: Fen) -> LaxFen {
        LaxFen(fen)
    }
}

impl From<LaxFen> for Fen {
    fn from(LaxFen(fen): LaxFen) -> Fen {
        fen
    }
}

impl FromStr for LaxFen {
    type Err = ParseFenError;

    fn from_str(s: &str) -> Result<LaxFen, ParseFenError> {
        s.replace("_", " ").parse().map(LaxFen)
    }
}
<commit_msg>optimize selection of latest game per move<commit_after>use crate::{
    model::{GameId, Mode, PersonalEntry, Speed, Stats, UserName},
    opening::Opening,
};
use serde::{Deserialize, Serialize};
use serde_with::{serde_as, CommaSeparator, DisplayFromStr, FromInto, StringWithSeparator};
use shakmaty::{
    fen::{Fen, ParseFenError},
    san::SanPlus,
    uci::Uci,
    variant::VariantPosition,
    Color,
};
use std::str::FromStr;

mod error;
mod variant;

pub use error::Error;
pub use variant::LilaVariant;

#[serde_as]
#[derive(Deserialize, Debug)]
pub struct PersonalQuery {
    #[serde(default)]
    pub variant: LilaVariant,
    #[serde_as(as = "Option<DisplayFromStr>")]
    #[serde(default)]
    pub fen: Option<LaxFen>,
    #[serde_as(as = "StringWithSeparator<CommaSeparator, Uci>")]
    #[serde(default)]
    pub play: Vec<Uci>,
    #[serde_as(as = "DisplayFromStr")]
    pub player: UserName,
    #[serde_as(as = "FromInto<ColorProxy>")]
    pub color: Color,
    #[serde(default)]
    pub since: u32, // year
    #[serde(flatten)]
    pub filter: PersonalQueryFilter,
    #[serde(default)]
    pub update: bool,
}

#[derive(Deserialize, Debug)]
pub struct PersonalQueryFilter {
    #[serde(default)]
    pub modes: Option<Vec<Mode>>,
    #[serde(default)]
    pub speeds: Option<Vec<Speed>>,
}

#[derive(Serialize, Debug)]
pub struct PersonalResponse {
    #[serde(flatten)]
    total: Stats,
    moves: Vec<PersonalMoveRow>,
    opening: Option<&'static Opening>,
}

#[serde_as]
#[derive(Serialize, Debug)]
struct PersonalMoveRow {
    #[serde_as(as = "DisplayFromStr")]
    uci: Uci,
    #[serde_as(as = "DisplayFromStr")]
    san: SanPlus,
    #[serde(flatten)]
    stats: Stats,
    #[serde_as(as = "Option<DisplayFromStr>")]
    game: Option<GameId>,
}

impl PersonalQueryFilter {
    pub fn respond(
        &self,
        pos: VariantPosition,
        entry: PersonalEntry,
        opening: Option<&'static Opening>,
    ) -> PersonalResponse {
        let mut total = Stats::default();

        let mut moves = Vec::with_capacity(entry.sub_entries.len());

        for (uci, sub_entry) in entry.sub_entries {
            let m = uci.to_move(&pos).expect("legal uci in personal entry");
            let san = SanPlus::from_move(pos.clone(), &m);

            let mut latest_game: Option<(u64, GameId)> = None;
            let mut stats = Stats::default();

            for speed in Speed::ALL {
                if self
                    .speeds
                    .as_ref()
                    .map_or(true, |speeds| speeds.contains(&speed))
                {
                    for mode in Mode::ALL {
                        if self
                            .modes
                            .as_ref()
                            .map_or(true, |modes| modes.contains(&mode))
                        {
                            let group = sub_entry.by_speed(speed).by_mode(mode);
                            stats += group.stats.to_owned();

                            for (idx, game) in group.games.iter().copied() {
                                if latest_game.map_or(true, |(latest_idx, _game)| latest_idx < idx)
                                {
                                    latest_game = Some((idx, game))
                                }
                            }
                        }
                    }
                }
            }

            total += stats.clone();

            moves.push(PersonalMoveRow {
                uci,
                san,
                stats,
                game: latest_game.map(|(_, id)| id),
            });
        }

        moves.sort_by_key(|row| row.stats.total());

        PersonalResponse {
            total,
            moves,
            opening,
        }
    }
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum ColorProxy {
    White,
    Black,
}

impl From<Color> for ColorProxy {
    fn from(color: Color) -> ColorProxy {
        match color {
            Color::White => ColorProxy::White,
            Color::Black => ColorProxy::Black,
        }
    }
}

impl From<ColorProxy> for Color {
    fn from(color: ColorProxy) -> Color {
        match color {
            ColorProxy::White => Color::White,
            ColorProxy::Black => Color::Black,
        }
    }
}

#[derive(Debug)]
pub struct LaxFen(Fen);

impl From<Fen> for LaxFen {
    fn from(fen: Fen) -> LaxFen {
        LaxFen(fen)
    }
}

impl From<LaxFen> for Fen {
    fn from(LaxFen(fen): LaxFen) -> Fen {
        fen
    }
}

impl FromStr for LaxFen {
    type Err = ParseFenError;

    fn from_str(s: &str) -> Result<LaxFen, ParseFenError> {
        s.replace("_", " ").parse().map(LaxFen)
    }
}
<|endoftext|>