<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! An LRU cache designed for work with DNS lookups

use std::sync::Arc;
use std::time::{Duration, Instant};

use trust_dns_proto::op::Query;
use trust_dns_proto::rr::RData;

use error::*;
use lookup::Lookup;
use lru_cache::LruCache;

/// Maximum TTL as defined in https://tools.ietf.org/html/rfc2181
pub const MAX_TTL: u32 = 2147483647_u32;

#[derive(Debug)]
struct LruValue {
    // In the None case, this represents an NXDomain
    lookup: Option<Lookup>,
    valid_until: Instant,
}

impl LruValue {
    /// Returns true if this set of ips is still valid
    fn is_current(&self, now: Instant) -> bool {
        now <= self.valid_until
    }
}

#[derive(Debug)]
pub(crate) struct DnsLru {
    cache: LruCache<Query, LruValue>,
    /// An optional minimum TTL value.
    ///
    /// RDatas with TTLs under `cache_min_ttl` will use `cache_min_ttl` instead.
    cache_min_ttl: Duration,
}

impl DnsLru {
    pub(crate) fn new(capacity: usize) -> Self {
        Self::with_min_ttl(capacity, None)
    }

    pub(crate) fn with_min_ttl(capacity: usize, min_ttl: Option<Duration>) -> Self {
        let cache = LruCache::new(capacity);
        let cache_min_ttl = min_ttl.unwrap_or_else(|| Duration::from_secs(0));
        DnsLru {
            cache,
            cache_min_ttl,
        }
    }

    pub(crate) fn insert(
        &mut self,
        query: Query,
        rdatas_and_ttl: Vec<(RData, u32)>,
        now: Instant,
    ) -> Lookup {
        let len = rdatas_and_ttl.len();
        // collapse the values, we're going to take the Minimum TTL as the correct one
        let (rdatas, ttl): (Vec<RData>, u32) = rdatas_and_ttl.into_iter().fold(
            (Vec::with_capacity(len), MAX_TTL),
            |(mut rdatas, mut min_ttl), (rdata, ttl)| {
                rdatas.push(rdata);
                min_ttl = min_ttl.min(ttl);
                (rdatas, min_ttl)
            },
        );

        let ttl = Duration::from_secs(ttl as u64);
        // If the cache was configured with a minimum TTL, and that value is higher
        // than the minimum TTL in the values, use it instead.
        let ttl = self.cache_min_ttl.max(ttl);
        let valid_until = now + ttl;

        // insert into the LRU
        let lookup = Lookup::new_with_deadline(Arc::new(rdatas), valid_until);
        self.cache.insert(
            query,
            LruValue {
                lookup: Some(lookup.clone()),
                valid_until,
            },
        );

        lookup
    }

    /// Generally for inserting a set of records that have already been cached, but with a different Query.
    pub(crate) fn duplicate(
        &mut self,
        query: Query,
        lookup: Lookup,
        ttl: u32,
        now: Instant,
    ) -> Lookup {
        let ttl = Duration::from_secs(ttl as u64);
        let valid_until = now + ttl;

        self.cache.insert(
            query,
            LruValue {
                lookup: Some(lookup.clone()),
                valid_until,
            },
        );

        lookup
    }

    pub(crate) fn nx_error(query: Query, valid_until: Option<Instant>) -> ResolveError {
        ResolveErrorKind::NoRecordsFound { query, valid_until }.into()
    }

    pub(crate) fn negative(&mut self, query: Query, ttl: u32, now: Instant) -> ResolveError {
        // TODO: if we are getting a negative response, should we instead fallback to cache?
        //   this would cache indefinitely, probably not correct

        let ttl = Duration::from_secs(ttl as u64);
        let valid_until = now + ttl;

        self.cache.insert(
            query.clone(),
            LruValue {
                lookup: None,
                valid_until,
            },
        );

        Self::nx_error(query, Some(valid_until))
    }

    /// This needs to be mut b/c it's an LRU, meaning the ordering of elements will potentially change on retrieval...
    pub(crate) fn get(&mut self, query: &Query, now: Instant) -> Option<Lookup> {
        let mut out_of_date = false;
        let lookup = self.cache.get_mut(query).and_then(|value| {
            if value.is_current(now) {
                out_of_date = false;
                value.lookup.clone()
            } else {
                out_of_date = true;
                None
            }
        });

        // in this case, we can preemtively remove out of data elements
        // this assumes time is always moving forward, this would only not be true in contrived situations where now
        //  is not current time, like tests...
        if out_of_date {
            self.cache.remove(query);
        }

        lookup
    }
}

// see also the lookup_tests.rs in integration-tests crate
#[cfg(test)]
mod tests {
    use std::net::*;
    use std::str::FromStr;
    use std::time::*;

    use trust_dns_proto::op::Query;
    use trust_dns_proto::rr::{Name, RecordType};

    use super::*;

    #[test]
    fn test_is_current() {
        let now = Instant::now();
        let not_the_future = now + Duration::from_secs(4);
        let future = now + Duration::from_secs(5);
        let past_the_future = now + Duration::from_secs(6);

        let value = LruValue {
            lookup: None,
            valid_until: future,
        };

        assert!(value.is_current(now));
        assert!(value.is_current(not_the_future));
        assert!(value.is_current(future));
        assert!(!value.is_current(past_the_future));
    }

    #[test]
    fn test_insert() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        let ips_ttl = vec![(RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1)];
        let ips = vec![RData::A(Ipv4Addr::new(127, 0, 0, 1))];
        let mut lru = DnsLru::new(1);

        let rc_ips = lru.insert(name.clone(), ips_ttl, now);
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);

        let rc_ips = lru.get(&name, now).unwrap();
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);
    }

    #[test]
    fn test_insert_ttl() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        // TTL should be 1
        let ips_ttl = vec![
            (RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1),
            (RData::A(Ipv4Addr::new(127, 0, 0, 2)), 2),
        ];
        let ips = vec![
            RData::A(Ipv4Addr::new(127, 0, 0, 1)),
            RData::A(Ipv4Addr::new(127, 0, 0, 2)),
        ];
        let mut lru = DnsLru::new(1);

        lru.insert(name.clone(), ips_ttl, now);

        // still valid
        let rc_ips = lru.get(&name, now + Duration::from_secs(1)).unwrap();
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);

        // 2 should be one too far
        let rc_ips = lru.get(&name, now + Duration::from_secs(2));
        assert!(rc_ips.is_none());
    }

    #[test]
    fn test_insert_min_ttl() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        // TTL should be 1
        let ips_ttl = vec![
            (RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1),
            (RData::A(Ipv4Addr::new(127, 0, 0, 2)), 2),
        ];
        let ips = vec![
            RData::A(Ipv4Addr::new(127, 0, 0, 1)),
            RData::A(Ipv4Addr::new(127, 0, 0, 2)),
        ];

        // this cache should override the TTL of 1 seconds with the configured
        // minimum TTL of 3 seconds.
        let mut lru = DnsLru::with_min_ttl(1, Some(Duration::from_secs(3)));

        lru.insert(name.clone(), ips_ttl, now);

        // still valid
        let rc_ips = lru.get(&name, now + Duration::from_secs(1)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 1 second");
        }

        let rc_ips = lru.get(&name, now + Duration::from_secs(2)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 2 seconds");
        }

        let rc_ips = lru.get(&name, now + Duration::from_secs(3)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 3 seconds");
        }

        // after 4 seconds, the records should be invalid.
        let rc_ips = lru.get(&name, now + Duration::from_secs(4));
        assert!(rc_ips.is_none());
    }
}
<commit_msg>resolver: Test the TTLs of lookups returned by a cache with a min TTL<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! An LRU cache designed for work with DNS lookups

use std::sync::Arc;
use std::time::{Duration, Instant};

use trust_dns_proto::op::Query;
use trust_dns_proto::rr::RData;

use error::*;
use lookup::Lookup;
use lru_cache::LruCache;

/// Maximum TTL as defined in https://tools.ietf.org/html/rfc2181
pub const MAX_TTL: u32 = 2147483647_u32;

#[derive(Debug)]
struct LruValue {
    // In the None case, this represents an NXDomain
    lookup: Option<Lookup>,
    valid_until: Instant,
}

impl LruValue {
    /// Returns true if this set of ips is still valid
    fn is_current(&self, now: Instant) -> bool {
        now <= self.valid_until
    }
}

#[derive(Debug)]
pub(crate) struct DnsLru {
    cache: LruCache<Query, LruValue>,
    /// An optional minimum TTL value.
    ///
    /// RDatas with TTLs under `cache_min_ttl` will use `cache_min_ttl` instead.
    cache_min_ttl: Duration,
}

impl DnsLru {
    pub(crate) fn new(capacity: usize) -> Self {
        Self::with_min_ttl(capacity, None)
    }

    pub(crate) fn with_min_ttl(capacity: usize, min_ttl: Option<Duration>) -> Self {
        let cache = LruCache::new(capacity);
        let cache_min_ttl = min_ttl.unwrap_or_else(|| Duration::from_secs(0));
        DnsLru {
            cache,
            cache_min_ttl,
        }
    }

    pub(crate) fn insert(
        &mut self,
        query: Query,
        rdatas_and_ttl: Vec<(RData, u32)>,
        now: Instant,
    ) -> Lookup {
        let len = rdatas_and_ttl.len();
        // collapse the values, we're going to take the Minimum TTL as the correct one
        let (rdatas, ttl): (Vec<RData>, u32) = rdatas_and_ttl.into_iter().fold(
            (Vec::with_capacity(len), MAX_TTL),
            |(mut rdatas, mut min_ttl), (rdata, ttl)| {
                rdatas.push(rdata);
                min_ttl = min_ttl.min(ttl);
                (rdatas, min_ttl)
            },
        );

        let ttl = Duration::from_secs(ttl as u64);
        // If the cache was configured with a minimum TTL, and that value is higher
        // than the minimum TTL in the values, use it instead.
        let ttl = self.cache_min_ttl.max(ttl);
        let valid_until = now + ttl;

        // insert into the LRU
        let lookup = Lookup::new_with_deadline(Arc::new(rdatas), valid_until);
        self.cache.insert(
            query,
            LruValue {
                lookup: Some(lookup.clone()),
                valid_until,
            },
        );

        lookup
    }

    /// Generally for inserting a set of records that have already been cached, but with a different Query.
    pub(crate) fn duplicate(
        &mut self,
        query: Query,
        lookup: Lookup,
        ttl: u32,
        now: Instant,
    ) -> Lookup {
        let ttl = Duration::from_secs(ttl as u64);
        let valid_until = now + ttl;

        self.cache.insert(
            query,
            LruValue {
                lookup: Some(lookup.clone()),
                valid_until,
            },
        );

        lookup
    }

    pub(crate) fn nx_error(query: Query, valid_until: Option<Instant>) -> ResolveError {
        ResolveErrorKind::NoRecordsFound { query, valid_until }.into()
    }

    pub(crate) fn negative(&mut self, query: Query, ttl: u32, now: Instant) -> ResolveError {
        // TODO: if we are getting a negative response, should we instead fallback to cache?
        //   this would cache indefinitely, probably not correct

        let ttl = Duration::from_secs(ttl as u64);
        let valid_until = now + ttl;

        self.cache.insert(
            query.clone(),
            LruValue {
                lookup: None,
                valid_until,
            },
        );

        Self::nx_error(query, Some(valid_until))
    }

    /// This needs to be mut b/c it's an LRU, meaning the ordering of elements will potentially change on retrieval...
    pub(crate) fn get(&mut self, query: &Query, now: Instant) -> Option<Lookup> {
        let mut out_of_date = false;
        let lookup = self.cache.get_mut(query).and_then(|value| {
            if value.is_current(now) {
                out_of_date = false;
                value.lookup.clone()
            } else {
                out_of_date = true;
                None
            }
        });

        // in this case, we can preemtively remove out of data elements
        // this assumes time is always moving forward, this would only not be true in contrived situations where now
        //  is not current time, like tests...
        if out_of_date {
            self.cache.remove(query);
        }

        lookup
    }
}

// see also the lookup_tests.rs in integration-tests crate
#[cfg(test)]
mod tests {
    use std::net::*;
    use std::str::FromStr;
    use std::time::*;

    use trust_dns_proto::op::Query;
    use trust_dns_proto::rr::{Name, RecordType};

    use super::*;

    #[test]
    fn test_is_current() {
        let now = Instant::now();
        let not_the_future = now + Duration::from_secs(4);
        let future = now + Duration::from_secs(5);
        let past_the_future = now + Duration::from_secs(6);

        let value = LruValue {
            lookup: None,
            valid_until: future,
        };

        assert!(value.is_current(now));
        assert!(value.is_current(not_the_future));
        assert!(value.is_current(future));
        assert!(!value.is_current(past_the_future));
    }

    #[test]
    fn test_lookup_uses_cache_min_ttl() {
        let now = Instant::now();

        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        // record should have TTL of 1 second.
        let ips_ttl = vec![(RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1)];
        let ips = vec![RData::A(Ipv4Addr::new(127, 0, 0, 1))];

        // configure the cache with a minimum TTL of 2 seconds.
        let mut lru = DnsLru::with_min_ttl(1, Some(Duration::from_secs(2)));

        let rc_ips = lru.insert(name.clone(), ips_ttl, now);
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);
        // the returned lookup should use the cache's min TTL, since the
        // query's TTL was below the minimum.
        assert_eq!(rc_ips.valid_until(), now + Duration::from_secs(2));

        // record should have TTL of 3 seconds.
        let ips_ttl = vec![(RData::A(Ipv4Addr::new(127, 0, 0, 1)), 3)];

        let rc_ips = lru.insert(name.clone(), ips_ttl, now);
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);
        // the returned lookup should use the record's TTL, since it's
        // greater than the cache's minimum.
        assert_eq!(rc_ips.valid_until(), now + Duration::from_secs(3));
    }

    #[test]
    fn test_insert() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        let ips_ttl = vec![(RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1)];
        let ips = vec![RData::A(Ipv4Addr::new(127, 0, 0, 1))];
        let mut lru = DnsLru::new(1);

        let rc_ips = lru.insert(name.clone(), ips_ttl, now);
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);

        let rc_ips = lru.get(&name, now).unwrap();
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);
    }

    #[test]
    fn test_insert_ttl() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        // TTL should be 1
        let ips_ttl = vec![
            (RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1),
            (RData::A(Ipv4Addr::new(127, 0, 0, 2)), 2),
        ];
        let ips = vec![
            RData::A(Ipv4Addr::new(127, 0, 0, 1)),
            RData::A(Ipv4Addr::new(127, 0, 0, 2)),
        ];
        let mut lru = DnsLru::new(1);

        lru.insert(name.clone(), ips_ttl, now);

        // still valid
        let rc_ips = lru.get(&name, now + Duration::from_secs(1)).unwrap();
        assert_eq!(*rc_ips.iter().next().unwrap(), ips[0]);

        // 2 should be one too far
        let rc_ips = lru.get(&name, now + Duration::from_secs(2));
        assert!(rc_ips.is_none());
    }

    #[test]
    fn test_insert_min_ttl() {
        let now = Instant::now();
        let name = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
        // TTL should be 1
        let ips_ttl = vec![
            (RData::A(Ipv4Addr::new(127, 0, 0, 1)), 1),
            (RData::A(Ipv4Addr::new(127, 0, 0, 2)), 2),
        ];
        let ips = vec![
            RData::A(Ipv4Addr::new(127, 0, 0, 1)),
            RData::A(Ipv4Addr::new(127, 0, 0, 2)),
        ];

        // this cache should override the TTL of 1 seconds with the configured
        // minimum TTL of 3 seconds.
        let mut lru = DnsLru::with_min_ttl(1, Some(Duration::from_secs(3)));

        lru.insert(name.clone(), ips_ttl, now);

        // still valid
        let rc_ips = lru.get(&name, now + Duration::from_secs(1)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 1 second");
        }

        let rc_ips = lru.get(&name, now + Duration::from_secs(2)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 2 seconds");
        }

        let rc_ips = lru.get(&name, now + Duration::from_secs(3)).unwrap();
        for (rc_ip, ip) in rc_ips.iter().zip(ips.iter()) {
            assert_eq!(rc_ip, ip, "after 3 seconds");
        }

        // after 4 seconds, the records should be invalid.
        let rc_ips = lru.get(&name, now + Duration::from_secs(4));
        assert!(rc_ips.is_none());
    }
}
<|endoftext|>


<commit_before>use log::debug;
use serde::de::{self, DeserializeSeed, MapAccess, SeqAccess, Visitor};
use serde::forward_to_deserialize_any;
use serde::serde_if_integer128;
use std::convert::TryInto;

use super::MaxMindDBError;
use super::MaxMindDBError::DecodingError;

fn to_usize(base: u8, bytes: &[u8]) -> usize {
    bytes
        .iter()
        .fold(base as usize, |acc, &b| (acc << 8) | b as usize)
}

#[derive(Debug)]
pub struct Decoder<'de> {
    buf: &'de [u8],
    current_ptr: usize,
}

impl<'de> Decoder<'de> {
    pub fn new(buf: &'de [u8], start_ptr: usize) -> Decoder<'de> {
        Decoder {
            buf,
            current_ptr: start_ptr,
        }
    }

    fn eat_byte(&mut self) -> u8 {
        let b = self.buf[self.current_ptr];
        self.current_ptr += 1;
        b
    }

    fn size_from_ctrl_byte(&mut self, ctrl_byte: u8, type_num: u8) -> usize {
        let size = (ctrl_byte & 0x1f) as usize;
        // extended
        if type_num == 0 {
            return size;
        }

        let bytes_to_read = if size > 28 { size - 28 } else { 0 };

        let new_offset = self.current_ptr + bytes_to_read;
        let size_bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        match size {
            s if s < 29 => s,
            29 => 29_usize + size_bytes[0] as usize,
            30 => 285_usize + to_usize(0, size_bytes),
            _ => 65_821_usize + to_usize(0, size_bytes),
        }
    }

    fn decode_any<V: Visitor<'de>>(&mut self, visitor: V) -> DecodeResult<V::Value> {
        let ctrl_byte = self.eat_byte();
        let mut type_num = ctrl_byte >> 5;
        // Extended type
        if type_num == 0 {
            type_num = self.eat_byte() + 7;
        }
        let size = self.size_from_ctrl_byte(ctrl_byte, type_num);

        match type_num {
            1 => {
                let new_ptr = self.decode_pointer(size);
                let prev_ptr = self.current_ptr;
                self.current_ptr = new_ptr;

                let res = self.decode_any(visitor);
                self.current_ptr = prev_ptr;
                res
            }
            2 => self.decode_string(visitor, size),
            3 => self.decode_double(visitor, size),
            4 => self.decode_bytes(visitor, size),
            5 => self.decode_uint16(visitor, size),
            6 => self.decode_uint32(visitor, size),
            7 => self.decode_map(visitor, size),
            8 => self.decode_int(visitor, size),
            9 => self.decode_uint64(visitor, size),
            10 => {
                serde_if_integer128! {
                    return self.decode_uint128(visitor, size);
                }

                #[allow(unreachable_code)]
                self.decode_bytes(visitor, size)
            }
            11 => self.decode_array(visitor, size),
            14 => self.decode_bool(visitor, size),
            15 => self.decode_float(visitor, size),
            u => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "Unknown data type: {:?}",
                u
            ))),
        }
    }

    fn decode_array<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        visitor.visit_seq(ArrayAccess {
            de: self,
            count: size,
        })
    }

    fn decode_bool<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        match size {
            0 | 1 => visitor.visit_bool(size != 0),
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "bool of size {:?}",
                s
            ))),
        }
    }

    fn decode_bytes<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        let new_offset = self.current_ptr + size;
        let u8_slice = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        visitor.visit_borrowed_bytes(u8_slice)
    }

    fn decode_float<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        let new_offset = self.current_ptr + size;
        let value: [u8; 4] = self.buf[self.current_ptr..new_offset]
            .try_into()
            .map_err(|_| {
                MaxMindDBError::InvalidDatabaseError(format!(
                    "float of size {:?}",
                    new_offset - self.current_ptr
                ))
            })?;
        self.current_ptr = new_offset;
        let float_value = f32::from_be_bytes(value);
        visitor.visit_f32(float_value)
    }

    fn decode_double<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        let new_offset = self.current_ptr + size;
        let value: [u8; 8] = self.buf[self.current_ptr..new_offset]
            .try_into()
            .map_err(|_| {
                MaxMindDBError::InvalidDatabaseError(format!(
                    "double of size {:?}",
                    new_offset - self.current_ptr
                ))
            })?;
        self.current_ptr = new_offset;
        let float_value = f64::from_be_bytes(value);
        visitor.visit_f64(float_value)
    }

    fn decode_uint64<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        match size {
            s if s <= 8 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u64, |acc, &b| (acc << 8) | u64::from(b));
                self.current_ptr = new_offset;
                visitor.visit_u64(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u64 of size {:?}",
                s
            ))),
        }
    }

    serde_if_integer128! {
        fn decode_uint128<V: Visitor<'de>>(
            &mut self,
            visitor: V,
            size: usize,
        ) -> DecodeResult<V::Value> {
            match size {
                s if s <= 16 => {
                    let new_offset = self.current_ptr + size;

                    let value = self.buf[self.current_ptr..new_offset]
                        .iter()
                        .fold(0_u128, |acc, &b| (acc << 8) | u128::from(b));
                    self.current_ptr = new_offset;
                    visitor.visit_u128(value)
                }
                s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                    "u128 of size {:?}",
                    s
                ))),
            }
        }
    }

    fn decode_uint32<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        match size {
            s if s <= 4 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u32, |acc, &b| (acc << 8) | u32::from(b));
                self.current_ptr = new_offset;
                visitor.visit_u32(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u32 of size {:?}",
                s
            ))),
        }
    }

    fn decode_uint16<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        match size {
            s if s <= 2 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u16, |acc, &b| (acc << 8) | u16::from(b));
                self.current_ptr = new_offset;
                visitor.visit_u16(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u16 of size {:?}",
                s
            ))),
        }
    }

    fn decode_int<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        match size {
            s if s <= 4 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_i32, |acc, &b| (acc << 8) | i32::from(b));
                self.current_ptr = new_offset;
                visitor.visit_i32(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "int32 of size {:?}",
                s
            ))),
        }
    }

    fn decode_map<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        visitor.visit_map(MapAccessor {
            de: self,
            count: size * 2,
        })
    }

    fn decode_pointer(&mut self, size: usize) -> usize {
        let pointer_value_offset = [0, 0, 2048, 526_336, 0];
        let pointer_size = ((size >> 3) & 0x3) + 1;
        let new_offset = self.current_ptr + pointer_size;
        let pointer_bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        let base = if pointer_size == 4 {
            0
        } else {
            (size & 0x7) as u8
        };
        let unpacked = to_usize(base, pointer_bytes);

        unpacked + pointer_value_offset[pointer_size]
    }

    #[cfg(feature = "unsafe-str-decode")]
    fn decode_string<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        use std::str::from_utf8_unchecked;

        let new_offset: usize = self.current_ptr + size;
        let bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;
        // SAFETY:
        // A corrupt maxminddb will cause undefined behaviour.
        // If the caller has verified the integrity of their database and trusts their upstream
        // provider, they can opt-into the performance gains provided by this unsafe function via
        // the `unsafe-str-decode` feature flag.
        // This can provide around 20% performance increase in the lookup benchmark.
        let v = unsafe { from_utf8_unchecked(bytes) };
        visitor.visit_borrowed_str(v)
    }

    #[cfg(not(feature = "unsafe-str-decode"))]
    fn decode_string<V: Visitor<'de>>(
        &mut self,
        visitor: V,
        size: usize,
    ) -> DecodeResult<V::Value> {
        use std::str::from_utf8;

        let new_offset: usize = self.current_ptr + size;
        let bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;
        match from_utf8(bytes) {
            Ok(v) => visitor.visit_borrowed_str(v),
            Err(_) => Err(MaxMindDBError::InvalidDatabaseError(
                "error decoding string".to_owned(),
            )),
        }
    }
}

pub type DecodeResult<T> = Result<T, MaxMindDBError>;

impl<'de: 'a, 'a> de::Deserializer<'de> for &'a mut Decoder<'de> {
    type Error = MaxMindDBError;

    fn deserialize_any<V>(self, visitor: V) -> DecodeResult<V::Value>
    where
        V: Visitor<'de>,
    {
        debug!("deserialize_any");

        self.decode_any(visitor)
    }

    fn deserialize_option<V>(self, visitor: V) -> DecodeResult<V::Value>
    where
        V: Visitor<'de>,
    {
        debug!("deserialize_option");

        visitor.visit_some(self)
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 i128 u8 u16 u32 u64 u128 f32 f64 char str string
        bytes byte_buf unit unit_struct newtype_struct seq tuple
        tuple_struct map struct enum identifier ignored_any
    }
}

struct ArrayAccess<'a, 'de: 'a> {
    de: &'a mut Decoder<'de>,
    count: usize,
}

// `SeqAccess` is provided to the `Visitor` to give it the ability to iterate
// through elements of the sequence.
impl<'de, 'a> SeqAccess<'de> for ArrayAccess<'a, 'de> {
    type Error = MaxMindDBError;

    fn next_element_seed<T>(&mut self, seed: T) -> DecodeResult<Option<T::Value>>
    where
        T: DeserializeSeed<'de>,
    {
        // Check if there are no more elements.
        if self.count == 0 {
            return Ok(None);
        }
        self.count -= 1;

        // Deserialize an array element.
        seed.deserialize(&mut *self.de).map(Some)
    }
}

struct MapAccessor<'a, 'de: 'a> {
    de: &'a mut Decoder<'de>,
    count: usize,
}

// `MapAccess` is provided to the `Visitor` to give it the ability to iterate
// through entries of the map.
impl<'de, 'a> MapAccess<'de> for MapAccessor<'a, 'de> {
    type Error = MaxMindDBError;

    fn next_key_seed<K>(&mut self, seed: K) -> DecodeResult<Option<K::Value>>
    where
        K: DeserializeSeed<'de>,
    {
        // Check if there are no more entries.
        if self.count == 0 {
            return Ok(None);
        }
        self.count -= 1;

        // Deserialize a map key.
        seed.deserialize(&mut *self.de).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> DecodeResult<V::Value>
    where
        V: DeserializeSeed<'de>,
    {
        // Check if there are no more entries.
        if self.count == 0 {
            return Err(DecodingError("no more entries".to_owned()));
        }
        self.count -= 1;

        // Deserialize a map value.
        seed.deserialize(&mut *self.de)
    }
}
<commit_msg>perf: Reduce the amount of code generated<commit_after>use log::debug;
use serde::de::{self, DeserializeSeed, MapAccess, SeqAccess, Visitor};
use serde::forward_to_deserialize_any;
use serde::serde_if_integer128;
use std::convert::TryInto;

use super::MaxMindDBError;
use super::MaxMindDBError::DecodingError;

fn to_usize(base: u8, bytes: &[u8]) -> usize {
    bytes
        .iter()
        .fold(base as usize, |acc, &b| (acc << 8) | b as usize)
}

#[derive(Debug)]
pub struct Decoder<'de> {
    buf: &'de [u8],
    current_ptr: usize,
}

impl<'de> Decoder<'de> {
    pub fn new(buf: &'de [u8], start_ptr: usize) -> Decoder<'de> {
        Decoder {
            buf,
            current_ptr: start_ptr,
        }
    }

    fn eat_byte(&mut self) -> u8 {
        let b = self.buf[self.current_ptr];
        self.current_ptr += 1;
        b
    }

    fn size_from_ctrl_byte(&mut self, ctrl_byte: u8, type_num: u8) -> usize {
        let size = (ctrl_byte & 0x1f) as usize;
        // extended
        if type_num == 0 {
            return size;
        }

        let bytes_to_read = if size > 28 { size - 28 } else { 0 };

        let new_offset = self.current_ptr + bytes_to_read;
        let size_bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        match size {
            s if s < 29 => s,
            29 => 29_usize + size_bytes[0] as usize,
            30 => 285_usize + to_usize(0, size_bytes),
            _ => 65_821_usize + to_usize(0, size_bytes),
        }
    }

    fn decode_any<V: Visitor<'de>>(&mut self, visitor: V) -> DecodeResult<V::Value> {
        let ctrl_byte = self.eat_byte();
        let mut type_num = ctrl_byte >> 5;
        // Extended type
        if type_num == 0 {
            type_num = self.eat_byte() + 7;
        }
        let size = self.size_from_ctrl_byte(ctrl_byte, type_num);

        match type_num {
            1 => {
                let new_ptr = self.decode_pointer(size);
                let prev_ptr = self.current_ptr;
                self.current_ptr = new_ptr;

                let res = self.decode_any(visitor);
                self.current_ptr = prev_ptr;
                res
            }
            2 => visitor.visit_borrowed_str(self.decode_string(size)?),
            3 => visitor.visit_f64(self.decode_double(size)?),
            4 => visitor.visit_borrowed_bytes(self.decode_bytes(size)?),
            5 => visitor.visit_u16(self.decode_uint16(size)?),
            6 => visitor.visit_u32(self.decode_uint32(size)?),
            7 => self.decode_map(visitor, size),
            8 => visitor.visit_i32(self.decode_int(size)?),
            9 => visitor.visit_u64(self.decode_uint64(size)?),
            10 => {
                serde_if_integer128! {
                    return visitor.visit_u128(self.decode_uint128(size)?);
                }

                #[allow(unreachable_code)]
                visitor.visit_borrowed_bytes(self.decode_bytes(size)?)
            }
            11 => self.decode_array(visitor, size),
            14 => visitor.visit_bool(self.decode_bool(size)?),
            15 => visitor.visit_f32(self.decode_float(size)?),
            u => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "Unknown data type: {:?}",
                u
            ))),
        }
    }

    fn decode_array<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        visitor.visit_seq(ArrayAccess {
            de: self,
            count: size,
        })
    }

    fn decode_bool(&mut self, size: usize) -> DecodeResult<bool> {
        match size {
            0 | 1 => Ok(size != 0),
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "bool of size {:?}",
                s
            ))),
        }
    }

    fn decode_bytes(&mut self, size: usize) -> DecodeResult<&'de [u8]> {
        let new_offset = self.current_ptr + size;
        let u8_slice = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        Ok(u8_slice)
    }

    fn decode_float(&mut self, size: usize) -> DecodeResult<f32> {
        let new_offset = self.current_ptr + size;
        let value: [u8; 4] = self.buf[self.current_ptr..new_offset]
            .try_into()
            .map_err(|_| {
                MaxMindDBError::InvalidDatabaseError(format!(
                    "float of size {:?}",
                    new_offset - self.current_ptr
                ))
            })?;
        self.current_ptr = new_offset;
        let float_value = f32::from_be_bytes(value);
        Ok(float_value)
    }

    fn decode_double(&mut self, size: usize) -> DecodeResult<f64> {
        let new_offset = self.current_ptr + size;
        let value: [u8; 8] = self.buf[self.current_ptr..new_offset]
            .try_into()
            .map_err(|_| {
                MaxMindDBError::InvalidDatabaseError(format!(
                    "double of size {:?}",
                    new_offset - self.current_ptr
                ))
            })?;
        self.current_ptr = new_offset;
        let float_value = f64::from_be_bytes(value);
        Ok(float_value)
    }

    fn decode_uint64(&mut self, size: usize) -> DecodeResult<u64> {
        match size {
            s if s <= 8 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u64, |acc, &b| (acc << 8) | u64::from(b));
                self.current_ptr = new_offset;
                Ok(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u64 of size {:?}",
                s
            ))),
        }
    }

    serde_if_integer128! {
        fn decode_uint128(
            &mut self,
            size: usize,
        ) -> DecodeResult<u128> {
            match size {
                s if s <= 16 => {
                    let new_offset = self.current_ptr + size;

                    let value = self.buf[self.current_ptr..new_offset]
                        .iter()
                        .fold(0_u128, |acc, &b| (acc << 8) | u128::from(b));
                    self.current_ptr = new_offset;
                    Ok(value)
                }
                s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                    "u128 of size {:?}",
                    s
                ))),
            }
        }
    }

    fn decode_uint32(&mut self, size: usize) -> DecodeResult<u32> {
        match size {
            s if s <= 4 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u32, |acc, &b| (acc << 8) | u32::from(b));
                self.current_ptr = new_offset;
                Ok(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u32 of size {:?}",
                s
            ))),
        }
    }

    fn decode_uint16(&mut self, size: usize) -> DecodeResult<u16> {
        match size {
            s if s <= 2 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_u16, |acc, &b| (acc << 8) | u16::from(b));
                self.current_ptr = new_offset;
                Ok(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "u16 of size {:?}",
                s
            ))),
        }
    }

    fn decode_int(&mut self, size: usize) -> DecodeResult<i32> {
        match size {
            s if s <= 4 => {
                let new_offset = self.current_ptr + size;

                let value = self.buf[self.current_ptr..new_offset]
                    .iter()
                    .fold(0_i32, |acc, &b| (acc << 8) | i32::from(b));
                self.current_ptr = new_offset;
                Ok(value)
            }
            s => Err(MaxMindDBError::InvalidDatabaseError(format!(
                "int32 of size {:?}",
                s
            ))),
        }
    }

    fn decode_map<V: Visitor<'de>>(&mut self, visitor: V, size: usize) -> DecodeResult<V::Value> {
        visitor.visit_map(MapAccessor {
            de: self,
            count: size * 2,
        })
    }

    fn decode_pointer(&mut self, size: usize) -> usize {
        let pointer_value_offset = [0, 0, 2048, 526_336, 0];
        let pointer_size = ((size >> 3) & 0x3) + 1;
        let new_offset = self.current_ptr + pointer_size;
        let pointer_bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;

        let base = if pointer_size == 4 {
            0
        } else {
            (size & 0x7) as u8
        };
        let unpacked = to_usize(base, pointer_bytes);

        unpacked + pointer_value_offset[pointer_size]
    }

    #[cfg(feature = "unsafe-str-decode")]
    fn decode_string(&mut self, size: usize) -> DecodeResult<&'de str> {
        use std::str::from_utf8_unchecked;

        let new_offset: usize = self.current_ptr + size;
        let bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;
        // SAFETY:
        // A corrupt maxminddb will cause undefined behaviour.
        // If the caller has verified the integrity of their database and trusts their upstream
        // provider, they can opt-into the performance gains provided by this unsafe function via
        // the `unsafe-str-decode` feature flag.
        // This can provide around 20% performance increase in the lookup benchmark.
        let v = unsafe { from_utf8_unchecked(bytes) };
        Ok(v)
    }

    #[cfg(not(feature = "unsafe-str-decode"))]
    fn decode_string(&mut self, size: usize) -> DecodeResult<&'de str> {
        use std::str::from_utf8;

        let new_offset: usize = self.current_ptr + size;
        let bytes = &self.buf[self.current_ptr..new_offset];
        self.current_ptr = new_offset;
        match from_utf8(bytes) {
            Ok(v) => Ok(v),
            Err(_) => Err(MaxMindDBError::InvalidDatabaseError(
                "error decoding string".to_owned(),
            )),
        }
    }
}

pub type DecodeResult<T> = Result<T, MaxMindDBError>;

impl<'de: 'a, 'a> de::Deserializer<'de> for &'a mut Decoder<'de> {
    type Error = MaxMindDBError;

    fn deserialize_any<V>(self, visitor: V) -> DecodeResult<V::Value>
    where
        V: Visitor<'de>,
    {
        debug!("deserialize_any");

        self.decode_any(visitor)
    }

    fn deserialize_option<V>(self, visitor: V) -> DecodeResult<V::Value>
    where
        V: Visitor<'de>,
    {
        debug!("deserialize_option");

        visitor.visit_some(self)
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 i128 u8 u16 u32 u64 u128 f32 f64 char str string
        bytes byte_buf unit unit_struct newtype_struct seq tuple
        tuple_struct map struct enum identifier ignored_any
    }
}

struct ArrayAccess<'a, 'de: 'a> {
    de: &'a mut Decoder<'de>,
    count: usize,
}

// `SeqAccess` is provided to the `Visitor` to give it the ability to iterate
// through elements of the sequence.
impl<'de, 'a> SeqAccess<'de> for ArrayAccess<'a, 'de> {
    type Error = MaxMindDBError;

    fn next_element_seed<T>(&mut self, seed: T) -> DecodeResult<Option<T::Value>>
    where
        T: DeserializeSeed<'de>,
    {
        // Check if there are no more elements.
        if self.count == 0 {
            return Ok(None);
        }
        self.count -= 1;

        // Deserialize an array element.
        seed.deserialize(&mut *self.de).map(Some)
    }
}

struct MapAccessor<'a, 'de: 'a> {
    de: &'a mut Decoder<'de>,
    count: usize,
}

// `MapAccess` is provided to the `Visitor` to give it the ability to iterate
// through entries of the map.
impl<'de, 'a> MapAccess<'de> for MapAccessor<'a, 'de> {
    type Error = MaxMindDBError;

    fn next_key_seed<K>(&mut self, seed: K) -> DecodeResult<Option<K::Value>>
    where
        K: DeserializeSeed<'de>,
    {
        // Check if there are no more entries.
        if self.count == 0 {
            return Ok(None);
        }
        self.count -= 1;

        // Deserialize a map key.
        seed.deserialize(&mut *self.de).map(Some)
    }

    fn next_value_seed<V>(&mut self, seed: V) -> DecodeResult<V::Value>
    where
        V: DeserializeSeed<'de>,
    {
        // Check if there are no more entries.
        if self.count == 0 {
            return Err(DecodingError("no more entries".to_owned()));
        }
        self.count -= 1;

        // Deserialize a map value.
        seed.deserialize(&mut *self.de)
    }
}
<|endoftext|>


<commit_before>use std::ops::{Index, IndexMut};
use std::iter::{Iterator, repeat};

use graph::{Edge, Graph};

pub struct AdjacencyMatrix {
    rows: Vec<Vec<u32>>,
}

impl AdjacencyMatrix {
    fn new(size: usize) -> Self {
        let row = vec![0; size];
        let rows = repeat(row).take(size).collect();

        AdjacencyMatrix { rows: rows }
    }

    pub fn from_graph(g: &Graph) -> Self {
        let mut matrix = AdjacencyMatrix::new(g.vertex_count);

        for &Edge { x, y } in &g.edges {
            matrix[x][y] = 1;
            if !g.directed {
                matrix[y][x] = 1;
            }
        }

        matrix
    }

    pub fn size(&self) -> usize {
        self.rows.len()
    }
}

// For convenience, enable writing matrix[x] instead of matrix.rows[x].

impl Index<usize> for AdjacencyMatrix {
    type Output = Vec<u32>;

    fn index<'a>(&'a self, _index: usize) -> &'a Vec<u32> {
        &self.rows[_index]
    }
}

impl IndexMut<usize> for AdjacencyMatrix {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut Vec<u32> {
        &mut self.rows[_index]
    }
}

// Breadth-first search.

struct BreadthFirstSearch<'a> {
    matrix: &'a AdjacencyMatrix,
    visited: Vec<bool>,
    queue: Vec<usize>,
    y: usize,
}

impl AdjacencyMatrix {
    pub fn breadth_first_search(&self, start: usize) -> BreadthFirstSearch {
        let mut visited = vec!(false; self.size());
        visited[start] = true;

        BreadthFirstSearch {
            matrix: self,
            visited: visited,
            queue: vec![start],
            y: 0
        }
    }
}

impl<'a> Iterator for BreadthFirstSearch<'a> {
    type Item = Edge;

    fn next(&mut self) -> Option<Edge> {
        while !self.queue.is_empty() {
            let x = self.queue[0];

            for y in self.y..self.matrix.size() {
                if self.visited[y] || self.matrix[x][y] == 0 {
                    continue;
                }
                self.visited[y] = true;
                self.queue.push(y);
                self.y = y + 1;
                return Some(Edge { x: x, y: y });
            }

            // Process the next element in the queue.
            self.queue.remove(0);
            self.y = 0;
        }

        None
    }
}

// Breadth-first search using std::iter::Unfold. In this situation there's no reason to use it over
// implementing Iterator directly. Unfold seems to shine in situations where we can just pass it a
// closure, but here we cannot use a closure because we need to return an iterator (see
// http://is.gd/2MDM1r).

// struct BFS<'a> {
//     matrix: &'a AdjacencyMatrix,
//     visited: Vec<bool>,
//     queue: Vec<usize>,
//     y: usize,
// }

// impl AdjacencyMatrix {
//     pub fn breadth_first_search(&self, start: usize) ->
//         Unfold<BFS, fn(&mut BFS) -> Option<Edge>> {

//         let mut visited = vec!(false; self.size());
//         visited[start] = true;

//         let state = BFS {
//             matrix: self,
//             visited: visited,
//             queue: vec![start],
//             y: 0
//         };

//         Unfold::new(state, unfold_breadth_first_search)
//     }
// }

// fn unfold_breadth_first_search(state: &mut BFS) -> Option<Edge> {
//     while !state.queue.is_empty() {
//         let x = state.queue[0];

//         for y in state.y..state.matrix.size() {
//             if state.visited[y] || state.matrix[x][y] == 0 {
//                 continue;
//             }

//             state.visited[y] = true;
//             state.queue.push(y);
//             state.y = y + 1;

//             return Some(Edge { x: x, y: y });
//         }

//         // Process the next element in the queue.
//         state.queue.remove(0);
//         state.y = 0;
//     }

//     None
// }

// Breadth-first search using a callback closure. The implementation is much simpler than the
// iterator-based BFS, however the external API is not as nice (forced to use a callback closure
// instead of a classic for loop).

// impl AdjacencyMatrix {
//     pub fn breadth_first_search<C>(&self, start: usize, mut callback: C)
//        where C: FnMut(Edge) -> bool {

//         let mut visited = vec!(false; self.size());
//         visited[start] = true;

//         let mut queue = vec![start];

//         while queue.len() > 0 {
//             // Dequeue first element.
//             let x = queue.remove(0);

//             for y in 0..self.size() {
//                 if visited[y] || self[x][y] == 0 {
//                     continue;
//                 }
//                 visited[y] = true;
//                 queue.push(y);
//                 // The edge callback can return true to abort the current graph traversal.
//                 if callback(Edge { x: x, y: y }) {
//                     return;
//                 }
//             }
//         }
//     }
// }
<commit_msg>[edge_connectivity/rust] Fix compiler warning<commit_after>use std::ops::{Index, IndexMut};
use std::iter::{Iterator, repeat};

use graph::{Edge, Graph};

pub struct AdjacencyMatrix {
    rows: Vec<Vec<u32>>,
}

impl AdjacencyMatrix {
    fn new(size: usize) -> Self {
        let row = vec![0; size];
        let rows = repeat(row).take(size).collect();

        AdjacencyMatrix { rows: rows }
    }

    pub fn from_graph(g: &Graph) -> Self {
        let mut matrix = AdjacencyMatrix::new(g.vertex_count);

        for &Edge { x, y } in &g.edges {
            matrix[x][y] = 1;
            if !g.directed {
                matrix[y][x] = 1;
            }
        }

        matrix
    }

    pub fn size(&self) -> usize {
        self.rows.len()
    }
}

// For convenience, enable writing matrix[x] instead of matrix.rows[x].

impl Index<usize> for AdjacencyMatrix {
    type Output = Vec<u32>;

    fn index<'a>(&'a self, _index: usize) -> &'a Vec<u32> {
        &self.rows[_index]
    }
}

impl IndexMut<usize> for AdjacencyMatrix {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut Vec<u32> {
        &mut self.rows[_index]
    }
}

// Breadth-first search.

pub struct BreadthFirstSearch<'a> {
    matrix: &'a AdjacencyMatrix,
    visited: Vec<bool>,
    queue: Vec<usize>,
    y: usize,
}

impl AdjacencyMatrix {
    pub fn breadth_first_search(&self, start: usize) -> BreadthFirstSearch {
        let mut visited = vec!(false; self.size());
        visited[start] = true;

        BreadthFirstSearch {
            matrix: self,
            visited: visited,
            queue: vec![start],
            y: 0
        }
    }
}

impl<'a> Iterator for BreadthFirstSearch<'a> {
    type Item = Edge;

    fn next(&mut self) -> Option<Edge> {
        while !self.queue.is_empty() {
            let x = self.queue[0];

            for y in self.y..self.matrix.size() {
                if self.visited[y] || self.matrix[x][y] == 0 {
                    continue;
                }
                self.visited[y] = true;
                self.queue.push(y);
                self.y = y + 1;
                return Some(Edge { x: x, y: y });
            }

            // Process the next element in the queue.
            self.queue.remove(0);
            self.y = 0;
        }

        None
    }
}

// Breadth-first search using std::iter::Unfold. In this situation there's no reason to use it over
// implementing Iterator directly. Unfold seems to shine in situations where we can just pass it a
// closure, but here we cannot use a closure because we need to return an iterator (see
// http://is.gd/2MDM1r).

// struct BFS<'a> {
//     matrix: &'a AdjacencyMatrix,
//     visited: Vec<bool>,
//     queue: Vec<usize>,
//     y: usize,
// }

// impl AdjacencyMatrix {
//     pub fn breadth_first_search(&self, start: usize) ->
//         Unfold<BFS, fn(&mut BFS) -> Option<Edge>> {

//         let mut visited = vec!(false; self.size());
//         visited[start] = true;

//         let state = BFS {
//             matrix: self,
//             visited: visited,
//             queue: vec![start],
//             y: 0
//         };

//         Unfold::new(state, unfold_breadth_first_search)
//     }
// }

// fn unfold_breadth_first_search(state: &mut BFS) -> Option<Edge> {
//     while !state.queue.is_empty() {
//         let x = state.queue[0];

//         for y in state.y..state.matrix.size() {
//             if state.visited[y] || state.matrix[x][y] == 0 {
//                 continue;
//             }

//             state.visited[y] = true;
//             state.queue.push(y);
//             state.y = y + 1;

//             return Some(Edge { x: x, y: y });
//         }

//         // Process the next element in the queue.
//         state.queue.remove(0);
//         state.y = 0;
//     }

//     None
// }

// Breadth-first search using a callback closure. The implementation is much simpler than the
// iterator-based BFS, however the external API is not as nice (forced to use a callback closure
// instead of a classic for loop).

// impl AdjacencyMatrix {
//     pub fn breadth_first_search<C>(&self, start: usize, mut callback: C)
//        where C: FnMut(Edge) -> bool {

//         let mut visited = vec!(false; self.size());
//         visited[start] = true;

//         let mut queue = vec![start];

//         while queue.len() > 0 {
//             // Dequeue first element.
//             let x = queue.remove(0);

//             for y in 0..self.size() {
//                 if visited[y] || self[x][y] == 0 {
//                     continue;
//                 }
//                 visited[y] = true;
//                 queue.push(y);
//                 // The edge callback can return true to abort the current graph traversal.
//                 if callback(Edge { x: x, y: y }) {
//                     return;
//                 }
//             }
//         }
//     }
// }
<|endoftext|>


<commit_before>use util::{ZOPFLI_WINDOW_MASK, ZOPFLI_MIN_MATCH};

const HASH_SHIFT: i32 = 5;
const HASH_MASK: i32 = 32767;

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Which {
    Hash1,
    Hash2,
}

pub struct HashThing {
    head: Vec<i32>,  /* Hash value to index of its most recent occurrence. */
    prev: Vec<u16>,  /* Index to index of prev. occurrence of same hash. */
    hashval: Vec<i32>,  /* Index to hash value at this index. */
    val: i32,  /* Current hash value. */
}

impl HashThing {
    fn new(window_size: usize) -> HashThing {
        HashThing {
            head: vec![-1; 65536],
            prev: (0..window_size as u16).collect(),
            hashval: vec![-1; window_size],
            val: 0,
        }
    }

    fn reset(&mut self, window_size: usize) {
        self.val = 0;
        self.head = vec![-1; 65536];
        self.prev = (0..window_size as u16).collect();
        self.hashval = vec![-1; window_size];
    }

    fn update(&mut self, hpos: usize) {
        self.hashval[hpos] = self.val;

        let index = self.val as usize;
        if self.head[index] != -1 && self.hashval[self.head[index] as usize] == self.val {
            self.prev[hpos] = self.head[index] as u16;
        } else {
            self.prev[hpos] = hpos as u16;
        }
        self.head[index] = hpos as i32;
    }
}

pub struct ZopfliHash {
    hash1: HashThing,
    hash2: HashThing,
    pub same: Vec<u16>,  /* Amount of repetitions of same byte after this .*/
}

impl ZopfliHash {
    pub fn new(window_size: usize) -> ZopfliHash {
        ZopfliHash {
            hash1: HashThing::new(window_size),
            hash2: HashThing::new(window_size),
            same: vec![0; window_size],
        }
    }

    pub fn reset(&mut self, window_size: usize) {
        self.hash1.reset(window_size);
        self.hash2.reset(window_size);
        self.same = vec![0; window_size];
    }

    pub fn warmup(&mut self, arr: &[u8], pos: usize, end: usize) {
        let c = arr[pos];
        self.update_val(c);

        if pos + 1 < end {
            let c = arr[pos + 1];
            self.update_val(c);
        }
    }

    /// Update the sliding hash value with the given byte. All calls to this function
    /// must be made on consecutive input characters. Since the hash value exists out
    /// of multiple input bytes, a few warmups with this function are needed initially.
    fn update_val(&mut self, c: u8) {
        self.hash1.val = ((self.hash1.val << HASH_SHIFT) ^ c as i32) & HASH_MASK;
    }

    pub fn update(&mut self, array: &[u8], pos: usize) {
        let hash_value = array.get(pos + ZOPFLI_MIN_MATCH - 1).cloned().unwrap_or(0);
        self.update_val(hash_value);

        let hpos = pos & ZOPFLI_WINDOW_MASK;

        self.hash1.update(hpos);

        // Update "same".
        let mut amount = 0;
        let same_index = pos.wrapping_sub(1) & ZOPFLI_WINDOW_MASK;
        let same = self.same[same_index];
        if same > 1 {
            amount = same as i32 - 1;
        }

        let mut another_index = pos + amount as usize + 1;
        let array_pos = array[pos];
        while another_index < array.len() && array_pos == array[another_index] && amount < -1 {
            amount += 1;
            another_index += 1;
        }
        self.same[hpos] = amount as u16;

        self.hash2.val = ((self.same[hpos].wrapping_sub(ZOPFLI_MIN_MATCH as u16) & 255) ^ self.hash1.val as u16) as i32;

        self.hash2.update(hpos);
    }

    pub fn head_at(&self, index: usize, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.head[index],
            Which::Hash2 => self.hash2.head[index],
        }
    }

    pub fn prev_at(&self, index: usize, which: Which) -> u16 {
        match which {
            Which::Hash1 => self.hash1.prev[index],
            Which::Hash2 => self.hash2.prev[index],
        }
    }

    pub fn hash_val_at(&self, index: usize, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.hashval[index],
            Which::Hash2 => self.hash2.hashval[index],
        }
    }

    pub fn val(&self, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.val,
            Which::Hash2 => self.hash2.val,
        }
    }
}
<commit_msg>Keep amount in the right type; use instead of array indexing<commit_after>use util::{ZOPFLI_WINDOW_MASK, ZOPFLI_MIN_MATCH};

const HASH_SHIFT: i32 = 5;
const HASH_MASK: i32 = 32767;

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Which {
    Hash1,
    Hash2,
}

pub struct HashThing {
    head: Vec<i32>,  /* Hash value to index of its most recent occurrence. */
    prev: Vec<u16>,  /* Index to index of prev. occurrence of same hash. */
    hashval: Vec<i32>,  /* Index to hash value at this index. */
    val: i32,  /* Current hash value. */
}

impl HashThing {
    fn new(window_size: usize) -> HashThing {
        HashThing {
            head: vec![-1; 65536],
            prev: (0..window_size as u16).collect(),
            hashval: vec![-1; window_size],
            val: 0,
        }
    }

    fn reset(&mut self, window_size: usize) {
        self.val = 0;
        self.head = vec![-1; 65536];
        self.prev = (0..window_size as u16).collect();
        self.hashval = vec![-1; window_size];
    }

    fn update(&mut self, hpos: usize) {
        self.hashval[hpos] = self.val;

        let index = self.val as usize;
        if self.head[index] != -1 && self.hashval[self.head[index] as usize] == self.val {
            self.prev[hpos] = self.head[index] as u16;
        } else {
            self.prev[hpos] = hpos as u16;
        }
        self.head[index] = hpos as i32;
    }
}

pub struct ZopfliHash {
    hash1: HashThing,
    hash2: HashThing,
    pub same: Vec<u16>,  /* Amount of repetitions of same byte after this .*/
}

impl ZopfliHash {
    pub fn new(window_size: usize) -> ZopfliHash {
        ZopfliHash {
            hash1: HashThing::new(window_size),
            hash2: HashThing::new(window_size),
            same: vec![0; window_size],
        }
    }

    pub fn reset(&mut self, window_size: usize) {
        self.hash1.reset(window_size);
        self.hash2.reset(window_size);
        self.same = vec![0; window_size];
    }

    pub fn warmup(&mut self, arr: &[u8], pos: usize, end: usize) {
        let c = arr[pos];
        self.update_val(c);

        if pos + 1 < end {
            let c = arr[pos + 1];
            self.update_val(c);
        }
    }

    /// Update the sliding hash value with the given byte. All calls to this function
    /// must be made on consecutive input characters. Since the hash value exists out
    /// of multiple input bytes, a few warmups with this function are needed initially.
    fn update_val(&mut self, c: u8) {
        self.hash1.val = ((self.hash1.val << HASH_SHIFT) ^ c as i32) & HASH_MASK;
    }

    pub fn update(&mut self, array: &[u8], pos: usize) {
        let hash_value = array.get(pos + ZOPFLI_MIN_MATCH - 1).cloned().unwrap_or(0);
        self.update_val(hash_value);

        let hpos = pos & ZOPFLI_WINDOW_MASK;

        self.hash1.update(hpos);

        // Update "same".
        let mut amount = 0;
        let same_index = pos.wrapping_sub(1) & ZOPFLI_WINDOW_MASK;
        let same = self.same[same_index];
        if same > 1 {
            amount = same as i32 - 1;
        }

        let mut another_index = pos + amount as usize + 1;
        let array_pos = array[pos];
        while another_index < array.len() && array_pos == array[another_index] && amount < -1 {
            amount += 1;
            another_index += 1;
        }
        let amount_u16 = amount as u16;
        self.same[hpos] = amount_u16;

        self.hash2.val = ((amount_u16.wrapping_sub(ZOPFLI_MIN_MATCH as u16) & 255) ^ self.hash1.val as u16) as i32;

        self.hash2.update(hpos);
    }

    pub fn head_at(&self, index: usize, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.head[index],
            Which::Hash2 => self.hash2.head[index],
        }
    }

    pub fn prev_at(&self, index: usize, which: Which) -> u16 {
        match which {
            Which::Hash1 => self.hash1.prev[index],
            Which::Hash2 => self.hash2.prev[index],
        }
    }

    pub fn hash_val_at(&self, index: usize, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.hashval[index],
            Which::Hash2 => self.hash2.hashval[index],
        }
    }

    pub fn val(&self, which: Which) -> i32 {
        match which {
            Which::Hash1 => self.hash1.val,
            Which::Hash2 => self.hash2.val,
        }
    }
}
<|endoftext|>


<commit_before>
extern crate getopts;
extern crate toml;

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::path;
use std::collections::HashMap;

use getopts::Options;

const CONFIG: &'static str = ".plugwise.toml";
const CONFIG_HEAD: &'static str = "config";
const CONFIG_DEVICE: &'static str = "device";
const ALIAS_MAC: &'static str = "mac";

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options] [mac|alias]", program);
    println!("{}", opts.usage(&brief));
}

fn load_config(configfile: &path::PathBuf) -> toml::Table {
    let mut config = String::new();
    if let Ok(mut file) = File::open(configfile) {
        file.read_to_string(&mut config).ok().expect(
            &format!("unable to read file `{}`", configfile.display()));
    } // errors are silently ignored (assuming file didn't exist)
    toml::Parser::new(&config).parse().expect(&format!("unable to parse `{}`", configfile.display()))
}

fn write_config(configfile: &path::PathBuf, config: &toml::Table) {
    let mut file = File::create(configfile).ok().expect(
        &format!("unable to create `{}`", configfile.display()));
    write!(file, "{}", toml::Value::Table(config.clone())).ok().expect(
        &format!("unable to write to `{}`", configfile.display()));
}

fn get_device_from_config<'a>(config: &'a toml::Table) -> Option<String> {
    config.get(CONFIG_HEAD)
          .map_or(None, |item|item.as_table())
          .map_or(None, |table|table.get(CONFIG_DEVICE))
          .map_or(None, |string|string.as_str())
          .map(|string|string.to_string())
}

fn get_aliases<'a>(config: &'a toml::Table) -> HashMap<String, u64> {
    let mut aliases = HashMap::new();

    for key in config.keys().filter(|&k|k != CONFIG_HEAD) {
        let mac = config.get(key)
                        .map_or(None, |item|item.as_table())
                        .map_or(None, |table|table.get(ALIAS_MAC))
                        .map_or(None, |string|string.as_str());
        if let Some(mac) = mac {
            if mac.len() == 16 {
                let mac = u64::from_str_radix(&mac, 16);
                if let Ok(mac) = mac {
                    aliases.insert(key.to_string(), mac);
                }
            }
        }
    }

    aliases
}

fn update_device_from_config<'a>(config: &'a toml::Table, device: &'a str) -> toml::Table {
    let mut config_table = config.get(CONFIG_HEAD)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.insert(CONFIG_DEVICE.to_string(), toml::Value::String(device.to_string()));
    config_table
}

fn remove_device_from_config<'a>(config: &'a toml::Table) -> toml::Table {
    let mut config_table = config.get(CONFIG_HEAD)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.remove(CONFIG_DEVICE);
    config_table
}

fn update_mac_in_alias<'a>(config: &'a toml::Table, alias: &'a str, mac: u64) -> toml::Table {
    let mut config_table = config.get(alias)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.insert(ALIAS_MAC.to_string(), toml::Value::String(format!("{:016X}", mac)));
    config_table
}

fn plugwise_actions(matches: &getopts::Matches, serial: Option<String>, mac: u64) {
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();

    opts.optopt("s", "serial", "configure serial-port", "DEVICE")
        .optflag("t", "stub", "configure to use stub implementation")
        .optopt("a", "alias", "assign a alias to Mac", "NAME")
        .optflag("u", "unalias", "forget alias")
        .optflag("l", "list", "list aliassed circles")
        .optflag("r", "relaystatus", "print the relay status of a circle") // XXX
        .optflag("e", "enable", "enable the relay of a circle") // XXX
        .optflag("d", "disable", "disable the relay of a circle") // XXX
        .optflag("p", "powerusage", "print the actual power usage of a circle") // XXX
        .optopt("o", "powersince", "print the total power usage of a given number of days", "DAYS") // XXX
        .optflag("c", "clock", "print the internal clock value of a circle") // XXX
        .optflag("j", "updateclock", "update the internal clock of a circle using Internet time") // XXX
        .optflag("h", "help", "print this help menu") // XXX
        .optflag("v", "verbose", "print debug information"); // XXX

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(f) => { panic!(f.to_string()) }
    };

    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let mut configfile = env::home_dir().expect("unable to find home/user directory");
    configfile.push(CONFIG);

    let mut config = load_config(&configfile);
    let mut update_config = false;

    if let Some(new_device) = matches.opt_str("s") {
        // client has provided new device; update (any) loaded configuration
        let new_config = update_device_from_config(&config, &new_device);
        config.insert(CONFIG_HEAD.to_string(), toml::Value::Table(new_config));
        update_config = true;
    } else if matches.opt_present("t") {
        // client has indicated to use stub
        let new_config = remove_device_from_config(&config);
        config.insert(CONFIG_HEAD.to_string(), toml::Value::Table(new_config));
        update_config = true;
    }

    let serial = get_device_from_config(&config);
    let aliases = get_aliases(&config);
    let mac;

    if matches.opt_present("l") {
        for (alias, mac) in aliases {
            println!("{} (mac:{:016X})", alias, mac);
        }
        return;
    }

    // at least alias or mac must be specified
    if matches.free.len() == 1 {
        let free = &matches.free[0];
        // find mac by alias or try to decode mac address (16 digit hex)
        mac = aliases.get(free).map_or_else(|| {
            match free.len() {
                16 => u64::from_str_radix(free, 16).ok(),
                _ => None,
            }
        }, |&x| Some(x)).expect("unknown alias or MAC specified");

        //println!("{:?}", mac); // FIXME: remove this
        //println!("{:?}", serial); // FIXME: remove this
        // FIXME: implement more options
        if let Some(new_alias) = matches.opt_str("a") {
            let new_config = update_mac_in_alias(&config, &new_alias, mac);
            config.insert(new_alias, toml::Value::Table(new_config));
            update_config = true;
        } else if matches.opt_present("u") {
            config.remove(free).expect("cannot unalias MAC");
            update_config = true;
        } else {
            plugwise_actions(&matches, serial, mac);
        }
    } else if !update_config {
        panic!("only one alias or MAC must be specified");
    }

    if update_config {
        write_config(&configfile, &config);
    }
}
<commit_msg>circle: add relay related operations<commit_after>
extern crate getopts;
extern crate time;
extern crate toml;
extern crate plugwise;

use std::env;
use std::fs::File;
use std::io;
use std::io::prelude::*;
use std::path;
use std::collections::HashMap;

use getopts::Options;

use time::Duration;

use plugwise::Device;
use plugwise::ProtocolSnoop;
use plugwise::plugwise;

const CONFIG: &'static str = ".plugwise.toml";
const CONFIG_HEAD: &'static str = "config";
const CONFIG_DEVICE: &'static str = "device";
const ALIAS_MAC: &'static str = "mac";

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options] [mac|alias]", program);
    println!("{}", opts.usage(&brief));
}

fn load_config(configfile: &path::PathBuf) -> toml::Table {
    let mut config = String::new();
    if let Ok(mut file) = File::open(configfile) {
        file.read_to_string(&mut config).ok().expect(
            &format!("unable to read file `{}`", configfile.display()));
    } // errors are silently ignored (assuming file didn't exist)
    toml::Parser::new(&config).parse().expect(&format!("unable to parse `{}`", configfile.display()))
}

fn write_config(configfile: &path::PathBuf, config: &toml::Table) {
    let mut file = File::create(configfile).ok().expect(
        &format!("unable to create `{}`", configfile.display()));
    write!(file, "{}", toml::Value::Table(config.clone())).ok().expect(
        &format!("unable to write to `{}`", configfile.display()));
}

fn get_device_from_config<'a>(config: &'a toml::Table) -> Option<String> {
    config.get(CONFIG_HEAD)
          .map_or(None, |item|item.as_table())
          .map_or(None, |table|table.get(CONFIG_DEVICE))
          .map_or(None, |string|string.as_str())
          .map(|string|string.to_string())
}

fn get_aliases<'a>(config: &'a toml::Table) -> HashMap<String, u64> {
    let mut aliases = HashMap::new();

    for key in config.keys().filter(|&k|k != CONFIG_HEAD) {
        let mac = config.get(key)
                        .map_or(None, |item|item.as_table())
                        .map_or(None, |table|table.get(ALIAS_MAC))
                        .map_or(None, |string|string.as_str());
        if let Some(mac) = mac {
            if mac.len() == 16 {
                let mac = u64::from_str_radix(&mac, 16);
                if let Ok(mac) = mac {
                    aliases.insert(key.to_string(), mac);
                }
            }
        }
    }

    aliases
}

fn update_device_from_config<'a>(config: &'a toml::Table, device: &'a str) -> toml::Table {
    let mut config_table = config.get(CONFIG_HEAD)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.insert(CONFIG_DEVICE.to_string(), toml::Value::String(device.to_string()));
    config_table
}

fn remove_device_from_config<'a>(config: &'a toml::Table) -> toml::Table {
    let mut config_table = config.get(CONFIG_HEAD)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.remove(CONFIG_DEVICE);
    config_table
}

fn update_mac_in_alias<'a>(config: &'a toml::Table, alias: &'a str, mac: u64) -> toml::Table {
    let mut config_table = config.get(alias)
                                 .map_or(None, |item|item.as_table())
                                 .map_or(toml::Table::new(), |table|table.clone());
    config_table.insert(ALIAS_MAC.to_string(), toml::Value::String(format!("{:016X}", mac)));
    config_table
}

fn plugwise_actions(matches: &getopts::Matches, serial: Option<String>, mac: u64) {
    let mut debug = io::stdout();
    let snoop = match matches.opt_count("v") {
        0 => ProtocolSnoop::Nothing,
        1 => ProtocolSnoop::Debug(&mut debug),
        2 => ProtocolSnoop::Raw(&mut debug),
        _ => ProtocolSnoop::All(&mut debug)
    };
    let device = match serial {
        Some(ref serial) => Device::SerialExt{port: &serial,
                                              timeout: Duration::milliseconds(1000),
                                              retries: 3,
                                              snoop: snoop},
        None => Device::Simulator
    };
    let plugwise = plugwise(device).ok().expect("unable to connect to Plugwise device");
    let circle = plugwise.create_circle(mac).ok().expect("unable to connect to circle");

    if matches.opt_present("r") {
        let status = circle.is_switched_on().ok().expect("unable retrieve circle status");
        println!("circle {:016X} relay_status: {}", mac, status);
    } else if matches.opt_present("e") {
        circle.switch_on().ok().expect("unable to switch on circle");
        println!("circle {:016X} switched on", mac);
    } else if matches.opt_present("d") {
        circle.switch_off().ok().expect("unable to switch on circle");
        println!("circle {:016X} switched off", mac);
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();

    opts.optopt("s", "serial", "configure serial-port", "DEVICE")
        .optflag("t", "stub", "configure to use stub implementation")
        .optopt("a", "alias", "assign a alias to Mac", "NAME")
        .optflag("u", "unalias", "forget alias")
        .optflag("l", "list", "list aliassed circles")
        .optflag("r", "relaystatus", "print the relay status of a circle")
        .optflag("e", "enable", "enable the relay of a circle")
        .optflag("d", "disable", "disable the relay of a circle")
        .optflag("p", "powerusage", "print the actual power usage of a circle") // XXX
        .optopt("o", "powersince", "print the total power usage of a given number of days", "DAYS") // XXX
        .optflag("c", "clock", "print the internal clock value of a circle") // XXX
        .optflag("j", "updateclock", "update the internal clock of a circle using Internet time") // XXX
        .optflag("h", "help", "print this help menu")
        .optflagmulti("v", "verbose", "print debug information");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m },
        Err(f) => { panic!(f.to_string()) }
    };

    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }

    let mut configfile = env::home_dir().expect("unable to find home/user directory");
    configfile.push(CONFIG);

    let mut config = load_config(&configfile);
    let mut update_config = false;

    if let Some(new_device) = matches.opt_str("s") {
        // client has provided new device; update (any) loaded configuration
        let new_config = update_device_from_config(&config, &new_device);
        config.insert(CONFIG_HEAD.to_string(), toml::Value::Table(new_config));
        update_config = true;
    } else if matches.opt_present("t") {
        // client has indicated to use stub
        let new_config = remove_device_from_config(&config);
        config.insert(CONFIG_HEAD.to_string(), toml::Value::Table(new_config));
        update_config = true;
    }

    let serial = get_device_from_config(&config);
    let aliases = get_aliases(&config);
    let mac;

    if matches.opt_present("l") {
        for (alias, mac) in aliases {
            println!("{} (mac:{:016X})", alias, mac);
        }
        return;
    }

    // at least alias or mac must be specified
    if matches.free.len() == 1 {
        let free = &matches.free[0];
        // find mac by alias or try to decode mac address (16 digit hex)
        mac = aliases.get(free).map_or_else(|| {
            match free.len() {
                16 => u64::from_str_radix(free, 16).ok(),
                _ => None,
            }
        }, |&x| Some(x)).expect("unknown alias or MAC specified");

        //println!("{:?}", mac); // FIXME: remove this
        //println!("{:?}", serial); // FIXME: remove this
        // FIXME: implement more options
        if let Some(new_alias) = matches.opt_str("a") {
            let new_config = update_mac_in_alias(&config, &new_alias, mac);
            config.insert(new_alias, toml::Value::Table(new_config));
            update_config = true;
        } else if matches.opt_present("u") {
            config.remove(free).expect("cannot unalias MAC");
            update_config = true;
        } else {
            plugwise_actions(&matches, serial, mac);
        }
    } else if !update_config {
        panic!("only one alias or MAC must be specified");
    }

    if update_config {
        write_config(&configfile, &config);
    }
}
<|endoftext|>