<commit_before>// Copyright 2013 Luis de Bethencourt <luis@debethencourt.com>
// Copyright 2013 The Rust Project Developers
// http://rust-lang.org

// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

/*!
 * Gregorian Calendar module for the datetime library of the Rust programming
 * language
 *
 * Implements a hybrid calendar that supports both the Julian and Gregorian
 * calendar systems.
 *
 * http://en.wikipedia.org/wiki/Gregorian_calendar
 * http://en.wikipedia.org/wiki/Julian_calendar
 */

static YEAR_BASE: int = 1900;
static DAYSPERLYEAR: uint = 366;
static DAYSPERNYEAR: uint = 365;
static DAYSPERWEEK: uint = 7;

pub fn is_leap_year(year: uint) -> bool {
    (year % 4 == 0) && ((year % 100 != 0) || (year % 400 == 0))
}

pub fn year_size(year: uint) -> uint {
    if is_leap_year(year) { DAYSPERLYEAR } else { DAYSPERNYEAR }
}

pub struct GCalendar {
    /**
    */
    sec: uint,
    min: uint,
    hour: uint,
    day_of_month: uint,
    month: uint,
    year: uint,
    day_of_week: uint,
    day_of_year: uint
}

impl GCalendar {
    /**
    * Allocates a GCalendar object
    */
    pub fn new_at_epoch() -> GCalendar {
        GCalendar {
            sec: 0,
            min: 0,
            hour: 0,
            day_of_month: 0,
            month: 0,
            year: 0,
            day_of_week: 0,
            day_of_year: 0,
        }
    }

    pub fn new(sec: uint, min: uint, hour: uint, day_of_month: uint, month: uint,
            year: uint, day_of_week: uint, day_of_year: uint) -> GCalendar {
        GCalendar {
            sec: sec,
            min: min,
            hour: hour,
            day_of_month: day_of_month,
            month: month,
            year: year,
            day_of_week: day_of_week,
            day_of_year: day_of_year,
        }
    }

    pub fn new_from_epoch(since_epoch: uint) -> GCalendar {
        let epoch_year = 1970;
        let mut year = epoch_year;

        let millisecs_day = 86400000;

        let mut dayclock = since_epoch % millisecs_day;
        let mut dayno = since_epoch / millisecs_day;

        let hour = dayclock / 3600000;
        dayclock = dayclock - (hour * 3600000);

        let min = dayclock / 60000;
        dayclock = dayclock - (min * 60000);

        let sec = dayclock / 1000;

        while (dayno >= year_size(year)) {
            dayno -= year_size(year);
            year += 1;
        }
        let year = year;

        GCalendar {
            sec: sec,
            min: min,
            hour: hour,
            day_of_month: 0,
            month: 0,
            year: year,
            day_of_week: 0,
            day_of_year: dayno,
        }
    }

    pub fn get_sec(&self) -> uint {
        self.sec
    }

    pub fn get_min(&self) -> uint {
        self.min
    }

    pub fn get_hour(&self) -> uint {
        self.hour
    }

    pub fn get_year(&self) -> uint {
        self.year
    }

    pub fn get_day_of_year(&self) -> uint {
        self.day_of_year
    }

    pub fn get_date(&self) -> ~str {
        // ToDo
        ~""
    }
}

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

    #[test]
    fn new() {
        let gc = GCalendar::new(21, 0, 12, 23, 9, 1983, 5, 265);
        assert_eq!(gc.get_sec(), 21);
        assert_eq!(gc.get_min(), 0);
        assert_eq!(gc.get_hour(), 12);
        assert_eq!(gc.get_year(), 1983);
        assert_eq!(gc.get_day_of_year(), 265);
    }

    fn new_from_epoch() {
        let gc = GCalendar::new_from_epoch(433166421023);
        assert_eq!(gc.get_sec(), 21);
        assert_eq!(gc.get_min(), 0);
        assert_eq!(gc.get_hour(), 12);
        assert_eq!(gc.get_year(), 1983);
        assert_eq!(gc.get_day_of_year(), 265);
    }
}
<commit_msg>gcalendar.rs: shorten names of GCalendar fields<commit_after>// Copyright 2013 Luis de Bethencourt <luis@debethencourt.com>
// Copyright 2013 The Rust Project Developers
// http://rust-lang.org

// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

/*!
 * Gregorian Calendar module for the datetime library of the Rust programming
 * language
 *
 * Implements a hybrid calendar that supports both the Julian and Gregorian
 * calendar systems.
 *
 * http://en.wikipedia.org/wiki/Gregorian_calendar
 * http://en.wikipedia.org/wiki/Julian_calendar
 */

static YEAR_BASE: int = 1900;
static DAYSPERLYEAR: uint = 366;
static DAYSPERNYEAR: uint = 365;
static DAYSPERWEEK: uint = 7;

pub fn is_leap_year(year: uint) -> bool {
    (year % 4 == 0) && ((year % 100 != 0) || (year % 400 == 0))
}

pub fn year_size(year: uint) -> uint {
    if is_leap_year(year) { DAYSPERLYEAR } else { DAYSPERNYEAR }
}

pub struct GCalendar {
    /*
     * Calendar object with date and time.
     */
    sec: uint,         /* Seconds       [0-59]  */
    min: uint,         /* Minutes       [0-59]  */
    hour: uint,        /* Hours         [0-23]  */
    mday: uint,        /* Day           [0-30]  */
    month: uint,       /* Month         [0-11]  */
    year: uint,        /* Year - 1900           */
    wday: uint,        /* Day of week   [0-6]   */
    yday: uint         /* Days in year  [0-365] */
}

impl GCalendar {
    /**
    * Allocates a GCalendar object at epoch.
    */
    pub fn new_at_epoch() -> GCalendar {
        GCalendar {
            sec: 0,
            min: 0,
            hour: 0,
            mday: 0,
            month: 0,
            year: 0,
            wday: 0,
            yday: 0,
        }
    }

    /**
    * Allocates a GCalendar object at the given date and time.
    */
    pub fn new(sec: uint, min: uint, hour: uint, mday: uint, month: uint,
            year: uint, wday: uint, yday: uint) -> GCalendar {
        GCalendar {
            sec: sec,
            min: min,
            hour: hour,
            mday: mday,
            month: month,
            year: year,
            wday: wday,
            yday: yday,
        }
    }

    /**
    * Allocates a GCalendar object from the milliseconds elapsed since epoch.
    */
    pub fn new_from_epoch(since_epoch: uint) -> GCalendar {
        let epoch_year = 1970;
        let mut year = epoch_year;

        let millisecs_day = 86400000;

        let mut dayclock = since_epoch % millisecs_day;
        let mut dayno = since_epoch / millisecs_day;

        let hour = dayclock / 3600000;
        dayclock = dayclock - (hour * 3600000);

        let min = dayclock / 60000;
        dayclock = dayclock - (min * 60000);

        let sec = dayclock / 1000;

        while (dayno >= year_size(year)) {
            dayno -= year_size(year);
            year += 1;
        }
        let year = year;

        GCalendar {
            sec: sec,
            min: min,
            hour: hour,
            mday: 0,
            month: 0,
            year: year,
            wday: 0,
            yday: dayno,
        }
    }

    pub fn get_sec(&self) -> uint {
        self.sec
    }

    pub fn get_min(&self) -> uint {
        self.min
    }

    pub fn get_hour(&self) -> uint {
        self.hour
    }

    pub fn get_year(&self) -> uint {
        self.year
    }

    pub fn get_day_of_year(&self) -> uint {
        self.yday
    }

    pub fn get_date(&self) -> ~str {
        // ToDo
        ~""
    }
}

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

    #[test]
    fn new() {
        let gc = GCalendar::new(21, 0, 12, 23, 9, 1983, 5, 265);
        assert_eq!(gc.get_sec(), 21);
        assert_eq!(gc.get_min(), 0);
        assert_eq!(gc.get_hour(), 12);
        assert_eq!(gc.get_year(), 1983);
        assert_eq!(gc.get_day_of_year(), 265);
    }

    fn new_from_epoch() {
        let gc = GCalendar::new_from_epoch(433166421023);
        assert_eq!(gc.get_sec(), 21);
        assert_eq!(gc.get_min(), 0);
        assert_eq!(gc.get_hour(), 12);
        assert_eq!(gc.get_year(), 1983);
        assert_eq!(gc.get_day_of_year(), 265);
    }
}
<|endoftext|>


<commit_before>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![desc = "Bindings for the GeoIP library."]
#![license = "BSD"]
#![crate_id = "geoip#0.1"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_uppercase_statics,
        unnecessary_qualification,
        managed_heap_memory)]

extern crate libc;

use libc::{c_void, c_char, c_int, c_ulong, c_float};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

type GeoIP_ = *c_void;
type In6Addr = [u8, ..16];

struct GeoIPLookup_ {
    netmask: c_int
}

impl GeoIPLookup_ {
    fn new() -> GeoIPLookup_ {
        GeoIPLookup_ {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *c_char, flags: c_int) -> GeoIP_;
    fn GeoIP_delete(db: GeoIP_);
    fn GeoIP_name_by_ipnum_gl(db: GeoIP_, ipnum: c_ulong, gl: &GeoIPLookup_) -> *c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: GeoIP_, ipnum: In6Addr, gl: &GeoIPLookup_) -> *c_char;
    fn GeoIP_record_by_ipnum(db: GeoIP_, ipnum: c_ulong) -> *GeoIPRecord_;
    fn GeoIP_record_by_ipnum_v6(db: GeoIP_, ipnum: In6Addr) -> *GeoIPRecord_;
    fn GeoIPRecord_delete(gir: *GeoIPRecord_);
    fn GeoIP_set_charset(db: GeoIP_, charset: c_int) -> c_int;
}

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIP {
    db: GeoIP_
}

pub struct GeoIPRecord_ {
    country_code: *c_char,
    country_code3: *c_char,
    country_name: *c_char,
    region: *c_char,
    city: *c_char,
    postal_code: *c_char,
    latitude: c_float,
    longitude: c_float,
    dma_code: c_int,
    area_code: c_int,
    charset: c_int,
    continent_code: *c_char,
    netmask: c_int
}

pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

impl CityInfo {
    fn from_geoiprecord(res: &GeoIPRecord_) -> CityInfo {
        let country_code = unsafe { if res.country_code.is_null() {
            None
        } else {
            CString::new(res.country_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let country_code3 = unsafe { if res.country_code3.is_null() {
            None
        } else {
            CString::new(res.country_code3, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let country_name = unsafe { if res.country_name.is_null() {
            None
        } else {
            CString::new(res.country_name, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let region = unsafe { if res.region.is_null() {
            None
        } else {
            CString::new(res.region, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let city = unsafe { if res.city.is_null() {
            None
        } else {
            CString::new(res.city, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let postal_code = unsafe { if res.postal_code.is_null() {
            None
        } else {
            CString::new(res.postal_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let continent_code = unsafe { if res.continent_code.is_null() {
            None
        } else {
            CString::new(res.continent_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        CityInfo {
            country_code: country_code,
            country_code3: country_code3,
            country_name: country_name,
            region: region,
            city: city,
            postal_code: postal_code,
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: continent_code,
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().unwrap(), options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { GeoIP_set_charset(db, UTF8 as c_int) } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIP { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_record_by_ipnum(self.db, ipnum)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_record_by_ipnum_v6(self.db, in6_addr)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let city_info = CityInfo::from_geoiprecord(&unsafe { *cres });
        unsafe { GeoIPRecord_delete(cres) };
        unsafe { std::mem::forget(cres) };
        Some(city_info)
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let gl = GeoIPLookup_::new();
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_name_by_ipnum_gl(self.db, ipnum, &gl)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_name_by_ipnum_v6_gl(self.db, in6_addr, &gl)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(' ', 1);
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIP {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<commit_msg>*constify<commit_after>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![desc = "Bindings for the GeoIP library."]
#![license = "BSD"]
#![crate_id = "geoip#0.1"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_uppercase_statics,
        unnecessary_qualification,
        managed_heap_memory)]

extern crate libc;

use libc::{c_void, c_char, c_int, c_ulong, c_float};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

type GeoIP_ = *const c_void;
type In6Addr = [u8, ..16];

struct GeoIPLookup_ {
    netmask: c_int
}

impl GeoIPLookup_ {
    fn new() -> GeoIPLookup_ {
        GeoIPLookup_ {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *const c_char, flags: c_int) -> GeoIP_;
    fn GeoIP_delete(db: GeoIP_);
    fn GeoIP_name_by_ipnum_gl(db: GeoIP_, ipnum: c_ulong, gl: &GeoIPLookup_) -> *const c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: GeoIP_, ipnum: In6Addr, gl: &GeoIPLookup_) -> *const c_char;
    fn GeoIP_record_by_ipnum(db: GeoIP_, ipnum: c_ulong) -> *const GeoIPRecord_;
    fn GeoIP_record_by_ipnum_v6(db: GeoIP_, ipnum: In6Addr) -> *const GeoIPRecord_;
    fn GeoIPRecord_delete(gir: *const GeoIPRecord_);
    fn GeoIP_set_charset(db: GeoIP_, charset: c_int) -> c_int;
}

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIP {
    db: GeoIP_
}

pub struct GeoIPRecord_ {
    country_code: *const c_char,
    country_code3: *const c_char,
    country_name: *const c_char,
    region: *const c_char,
    city: *const c_char,
    postal_code: *const c_char,
    latitude: c_float,
    longitude: c_float,
    dma_code: c_int,
    area_code: c_int,
    charset: c_int,
    continent_code: *const c_char,
    netmask: c_int
}

pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

impl CityInfo {
    fn from_geoiprecord(res: &GeoIPRecord_) -> CityInfo {
        let country_code = unsafe { if res.country_code.is_null() {
            None
        } else {
            CString::new(res.country_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let country_code3 = unsafe { if res.country_code3.is_null() {
            None
        } else {
            CString::new(res.country_code3, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let country_name = unsafe { if res.country_name.is_null() {
            None
        } else {
            CString::new(res.country_name, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let region = unsafe { if res.region.is_null() {
            None
        } else {
            CString::new(res.region, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let city = unsafe { if res.city.is_null() {
            None
        } else {
            CString::new(res.city, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let postal_code = unsafe { if res.postal_code.is_null() {
            None
        } else {
            CString::new(res.postal_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        let continent_code = unsafe { if res.continent_code.is_null() {
            None
        } else {
            CString::new(res.continent_code, false).as_str().
                and_then(|str| Some(String::from_str(str)))
        }};
        CityInfo {
            country_code: country_code,
            country_code3: country_code3,
            country_name: country_name,
            region: region,
            city: city,
            postal_code: postal_code,
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: continent_code,
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().unwrap(), options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { GeoIP_set_charset(db, UTF8 as c_int) } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIP { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_record_by_ipnum(self.db, ipnum)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_record_by_ipnum_v6(self.db, in6_addr)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let city_info = CityInfo::from_geoiprecord(&unsafe { *cres });
        unsafe { GeoIPRecord_delete(cres) };
        unsafe { std::mem::forget(cres) };
        Some(city_info)
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let gl = GeoIPLookup_::new();
        let cres = match ip {
            Ipv4Addr(a, b, c, d) => {
                let ipnum: c_ulong =
                    (a as c_ulong << 24) | (b as c_ulong << 16) |
                    (c as c_ulong << 8)  | (d as c_ulong);
                unsafe {
                    GeoIP_name_by_ipnum_gl(self.db, ipnum, &gl)
                }
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                let in6_addr: In6Addr = [(a >> 8) as u8, a as u8,
                                         (b >> 8) as u8, b as u8,
                                         (c >> 8) as u8, c as u8,
                                         (d >> 8) as u8, d as u8,
                                         (e >> 8) as u8, e as u8,
                                         (f >> 8) as u8, f as u8,
                                         (g >> 8) as u8, g as u8,
                                         (h >> 8) as u8, h as u8];
                unsafe {
                    GeoIP_name_by_ipnum_v6_gl(self.db, in6_addr, &gl)
                }
            }
        };
        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(' ', 1);
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIP {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIP::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), MemoryCache) {
        Err(err) => fail!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<|endoftext|>


<commit_before>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![desc = "Bindings for the GeoIP library."]
#![license = "BSD"]
#![crate_name = "geoip"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_upper_case_globals,
        unused_qualifications)]

extern crate libc;
extern crate serialize;

use libc::{c_void, c_char, c_int, c_ulong, c_float};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

type RawGeoIp = *const c_void;
type In6Addr = [u8, ..16];

#[repr(C)]
struct GeoIpLookup {
    netmask: c_int
}

impl GeoIpLookup {
    fn new() -> GeoIpLookup {
        GeoIpLookup {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *const c_char, flags: c_int) -> RawGeoIp;
    fn GeoIP_delete(db: RawGeoIp);
    fn GeoIP_name_by_ipnum_gl(db: RawGeoIp, ipnum: c_ulong, gl: *mut GeoIpLookup) -> *const c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: RawGeoIp, ipnum: In6Addr, gl: *mut GeoIpLookup) -> *const c_char;
    fn GeoIP_record_by_ipnum(db: RawGeoIp, ipnum: c_ulong) -> *const GeoIpRecord;
    fn GeoIP_record_by_ipnum_v6(db: RawGeoIp, ipnum: In6Addr) -> *const GeoIpRecord;
    fn GeoIPRecord_delete(gir: *const GeoIpRecord);
    fn GeoIP_set_charset(db: RawGeoIp, charset: c_int) -> c_int;
}

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIp {
    db: RawGeoIp
}

#[repr(C)]
pub struct GeoIpRecord {
    country_code: *const c_char,
    country_code3: *const c_char,
    country_name: *const c_char,
    region: *const c_char,
    city: *const c_char,
    postal_code: *const c_char,
    latitude: c_float,
    longitude: c_float,
    dma_code: c_int,
    area_code: c_int,
    charset: c_int,
    continent_code: *const c_char,
    netmask: c_int
}

#[deriving(Decodable, Encodable)]
pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

#[deriving(Decodable, Encodable)]
pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

unsafe fn maybe_string(c_str: *const c_char) -> Option<String> {
    c_str.as_ref().and_then(|opt| {
        CString::new(opt, false).as_str().map(|s| s.to_string())
    })
}

impl CityInfo {
    unsafe fn from_geoiprecord(res: &GeoIpRecord) -> CityInfo {
        CityInfo {
            country_code: maybe_string(res.country_code),
            country_code3: maybe_string(res.country_code3),
            country_name: maybe_string(res.country_name),
            region: maybe_string(res.region),
            city: maybe_string(res.city),
            postal_code: maybe_string(res.postal_code),
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: maybe_string(res.continent_code),
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

enum CNetworkIp {
    V4(c_ulong),
    V6(In6Addr)
}

impl CNetworkIp {
    fn new(ip: IpAddr) -> CNetworkIp {
        match ip {
            Ipv4Addr(a, b, c, d) => {
                CNetworkIp::V4((a as c_ulong << 24) | (b as c_ulong << 16) |
                               (c as c_ulong << 8)  | (d as c_ulong))
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                CNetworkIp::V6([(a >> 8) as u8, a as u8,
                                (b >> 8) as u8, b as u8,
                                (c >> 8) as u8, c as u8,
                                (d >> 8) as u8, d as u8,
                                (e >> 8) as u8, e as u8,
                                (f >> 8) as u8, f as u8,
                                (g >> 8) as u8, g as u8,
                                (h >> 8) as u8, h as u8])
            }
        }
    }
}

impl GeoIp {
    pub fn open(path: &Path, options: Options) -> Result<GeoIp, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().into_inner(), options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { GeoIP_set_charset(db, Charset::UTF8 as c_int) } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIp { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe { GeoIP_record_by_ipnum(self.db, ip) },
            CNetworkIp::V6(ip) => unsafe { GeoIP_record_by_ipnum_v6(self.db, ip) }
        };

        if cres.is_null() { return None; }

        unsafe {
            let city_info = CityInfo::from_geoiprecord(&*cres);
            GeoIPRecord_delete(cres);
            std::mem::forget(cres);
            Some(city_info)
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let mut gl = GeoIpLookup::new();
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe { GeoIP_name_by_ipnum_gl(self.db, ip, &mut gl) },
            CNetworkIp::V6(ip) => unsafe { GeoIP_name_by_ipnum_v6_gl(self.db, ip, &mut gl) }
        };

        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(1, ' ');
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIp {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<commit_msg>Remove unused attributes<commit_after>// (C)opyleft 2013,2014 Frank Denis

/*!
 * Bindings for the GeoIP library
 */
#![crate_name = "geoip"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_upper_case_globals,
        unused_qualifications)]

extern crate libc;
extern crate serialize;

use libc::{c_void, c_char, c_int, c_ulong, c_float};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

type RawGeoIp = *const c_void;
type In6Addr = [u8, ..16];

#[repr(C)]
struct GeoIpLookup {
    netmask: c_int
}

impl GeoIpLookup {
    fn new() -> GeoIpLookup {
        GeoIpLookup {
            netmask: 0
        }
    }
}

#[link(name = "GeoIP")]
extern {
    fn GeoIP_open(dbtype: *const c_char, flags: c_int) -> RawGeoIp;
    fn GeoIP_delete(db: RawGeoIp);
    fn GeoIP_name_by_ipnum_gl(db: RawGeoIp, ipnum: c_ulong, gl: *mut GeoIpLookup) -> *const c_char;
    fn GeoIP_name_by_ipnum_v6_gl(db: RawGeoIp, ipnum: In6Addr, gl: *mut GeoIpLookup) -> *const c_char;
    fn GeoIP_record_by_ipnum(db: RawGeoIp, ipnum: c_ulong) -> *const GeoIpRecord;
    fn GeoIP_record_by_ipnum_v6(db: RawGeoIp, ipnum: In6Addr) -> *const GeoIpRecord;
    fn GeoIPRecord_delete(gir: *const GeoIpRecord);
    fn GeoIP_set_charset(db: RawGeoIp, charset: c_int) -> c_int;
}

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

pub struct GeoIp {
    db: RawGeoIp
}

#[repr(C)]
pub struct GeoIpRecord {
    country_code: *const c_char,
    country_code3: *const c_char,
    country_name: *const c_char,
    region: *const c_char,
    city: *const c_char,
    postal_code: *const c_char,
    latitude: c_float,
    longitude: c_float,
    dma_code: c_int,
    area_code: c_int,
    charset: c_int,
    continent_code: *const c_char,
    netmask: c_int
}

#[deriving(Decodable, Encodable)]
pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

#[deriving(Decodable, Encodable)]
pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

unsafe fn maybe_string(c_str: *const c_char) -> Option<String> {
    c_str.as_ref().and_then(|opt| {
        CString::new(opt, false).as_str().map(|s| s.to_string())
    })
}

impl CityInfo {
    unsafe fn from_geoiprecord(res: &GeoIpRecord) -> CityInfo {
        CityInfo {
            country_code: maybe_string(res.country_code),
            country_code3: maybe_string(res.country_code3),
            country_name: maybe_string(res.country_name),
            region: maybe_string(res.region),
            city: maybe_string(res.city),
            postal_code: maybe_string(res.postal_code),
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: maybe_string(res.continent_code),
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

enum CNetworkIp {
    V4(c_ulong),
    V6(In6Addr)
}

impl CNetworkIp {
    fn new(ip: IpAddr) -> CNetworkIp {
        match ip {
            Ipv4Addr(a, b, c, d) => {
                CNetworkIp::V4((a as c_ulong << 24) | (b as c_ulong << 16) |
                               (c as c_ulong << 8)  | (d as c_ulong))
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                CNetworkIp::V6([(a >> 8) as u8, a as u8,
                                (b >> 8) as u8, b as u8,
                                (c >> 8) as u8, c as u8,
                                (d >> 8) as u8, d as u8,
                                (e >> 8) as u8, e as u8,
                                (f >> 8) as u8, f as u8,
                                (g >> 8) as u8, g as u8,
                                (h >> 8) as u8, h as u8])
            }
        }
    }
}

impl GeoIp {
    pub fn open(path: &Path, options: Options) -> Result<GeoIp, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            GeoIP_open(file.to_c_str().into_inner(), options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { GeoIP_set_charset(db, Charset::UTF8 as c_int) } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIp { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe { GeoIP_record_by_ipnum(self.db, ip) },
            CNetworkIp::V6(ip) => unsafe { GeoIP_record_by_ipnum_v6(self.db, ip) }
        };

        if cres.is_null() { return None; }

        unsafe {
            let city_info = CityInfo::from_geoiprecord(&*cres);
            GeoIPRecord_delete(cres);
            std::mem::forget(cres);
            Some(city_info)
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let mut gl = GeoIpLookup::new();
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe { GeoIP_name_by_ipnum_gl(self.db, ip, &mut gl) },
            CNetworkIp::V6(ip) => unsafe { GeoIP_name_by_ipnum_v6_gl(self.db, ip, &mut gl) }
        };

        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(1, ' ');
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIp {
    fn drop(&mut self) {
        unsafe {
            GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<|endoftext|>


<commit_before>use geom::{vertex, Tri, Vertex};
use math::EuclideanSpace;
use std::ops::{Deref, Index};

/// The number of vertices in a quad.
pub const NUM_VERTICES: u8 = 4;

/// A quad represented by its four vertices.
#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
pub struct Quad<V = vertex::Default>(pub [V; NUM_VERTICES as usize]);

/// An `Iterator` yielding the two triangles that make up a quad.
#[derive(Clone, Debug)]
pub struct Triangles<V> {
    a: Option<Tri<V>>,
    b: Option<Tri<V>>,
}

/// A simple iterator yielding each vertex in a `Quad`.
#[derive(Clone, Debug)]
pub struct Vertices<V> {
    quad: Quad<V>,
    index: u8,
}

impl<V> Quad<V>
where
    V: Vertex,
{
    /// Produce an iterator yielding each vertex in the `Quad`.
    pub fn vertices(self) -> Vertices<V> {
        vertices(self)
    }

    /// Produce the centroid of the quad, aka the "mean"/"average" vertex.
    pub fn centroid(&self) -> V
    where
        V: EuclideanSpace,
    {
        centroid(self)
    }

    /// Triangulates the given quad, represented by four points that describe its edges in either
    /// clockwise or anti-clockwise order.
    ///
    /// # Example
    ///
    /// The following rectangle
    ///
    /// ```ignore
    ///  a        b
    ///   --------
    ///   |      |
    ///   |      |
    ///   |      |
    ///   --------
    ///  d        c
    ///
    /// ```
    ///
    /// given as
    ///
    /// ```ignore
    /// triangles([a, b, c, d])
    /// ```
    ///
    /// returns
    ///
    /// ```ignore
    /// (Tri([a, b, c]), Tri([a, c, d]))
    /// ```
    ///
    /// Here's a basic code example:
    ///
    /// ```
    /// extern crate nannou;
    ///
    /// use nannou::geom::{self, Quad, Tri};
    /// use nannou::math::pt2;
    ///
    /// fn main() {
    ///     let a = pt2(0.0, 1.0);
    ///     let b = pt2(1.0, 1.0);
    ///     let c = pt2(1.0, 0.0);
    ///     let d = pt2(0.0, 0.0);
    ///     let quad = Quad([a, b, c, d]);
    ///     let triangles = geom::quad::triangles(&quad);
    ///     assert_eq!(triangles, (Tri([a, b, c]), Tri([a, c, d])));
    /// }
    /// ```
    #[inline]
    pub fn triangles(&self) -> (Tri<V>, Tri<V>) {
        triangles(self)
    }

    /// The same as `triangles` but provided as an `Iterator`.
    pub fn triangles_iter(&self) -> Triangles<V> {
        triangles_iter(self)
    }
}

/// Produce an iterator yielding each vertex in the given `Quad`.
pub fn vertices<V>(quad: Quad<V>) -> Vertices<V> {
    let index = 0;
    Vertices { quad, index }
}

/// Produce the centroid of the quad, aka the "mean"/"average" vertex.
pub fn centroid<V>(quad: &Quad<V>) -> V
where
    V: EuclideanSpace,
{
    EuclideanSpace::centroid(&quad[..])
}

/// Triangulates the given quad, represented by four points that describe its edges in either
/// clockwise or anti-clockwise order.
///
/// # Example
///
/// The following rectangle
///
/// ```ignore
///
///  a        b
///   --------
///   |      |
///   |      |
///   |      |
///   --------
///  d        c
///
/// ```
///
/// given as
///
/// ```ignore
/// triangles([a, b, c, d])
/// ```
///
/// returns
///
/// ```ignore
/// (Tri([a, b, c]), Tri([a, c, d]))
/// ```
///
/// Here's a basic code example:
///
/// ```
/// extern crate nannou;
///
/// use nannou::geom::{self, Quad, Tri};
/// use nannou::math::pt2;
///
/// fn main() {
///     let a = pt2(0.0, 1.0);
///     let b = pt2(1.0, 1.0);
///     let c = pt2(1.0, 0.0);
///     let d = pt2(0.0, 0.0);
///     let quad = Quad([a, b, c, d]);
///     let triangles = geom::quad::triangles(&quad);
///     assert_eq!(triangles, (Tri([a, b, c]), Tri([a, c, d])));
/// }
/// ```
#[inline]
pub fn triangles<V>(points: &Quad<V>) -> (Tri<V>, Tri<V>)
where
    V: Vertex,
{
    let (a, b, c, d) = (points[0], points[1], points[2], points[3]);
    (Tri([a, b, c]), Tri([a, c, d]))
}

/// The same as `triangles` but provided as an `Iterator`.
pub fn triangles_iter<V>(points: &Quad<V>) -> Triangles<V>
where
    V: Vertex,
{
    let (a, b) = triangles(points);
    Triangles {
        a: Some(a),
        b: Some(b),
    }
}

impl<V> Iterator for Triangles<V> {
    type Item = Tri<V>;
    fn next(&mut self) -> Option<Self::Item> {
        self.a.take().or_else(|| self.b.take())
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }
}

impl<V> DoubleEndedIterator for Triangles<V> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.b.take().or_else(|| self.a.take())
    }
}

impl<V> ExactSizeIterator for Triangles<V> {
    fn len(&self) -> usize {
        match (&self.a, &self.b) {
            (&Some(_), &Some(_)) => 2,
            (&None, &Some(_)) => 0,
            _ => 1,
        }
    }
}

impl<V> Iterator for Vertices<V>
where
    V: Clone,
{
    type Item = V;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index < NUM_VERTICES {
            let v = self.quad[self.index as usize].clone();
            self.index += 1;
            Some(v)
        } else {
            None
        }
    }
}

impl<V> Deref for Quad<V> {
    type Target = [V; NUM_VERTICES as usize];
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<V> From<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn from(points: [V; NUM_VERTICES as usize]) -> Self {
        Quad(points)
    }
}

impl<V> From<(V, V, V, V)> for Quad<V>
where
    V: Vertex,
{
    fn from((a, b, c, d): (V, V, V, V)) -> Self {
        Quad([a, b, c, d])
    }
}

impl<V> Into<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn into(self) -> [V; NUM_VERTICES as usize] {
        self.0
    }
}

impl<V> Into<(V, V, V, V)> for Quad<V>
where
    V: Vertex,
{
    fn into(self) -> (V, V, V, V) {
        (self[0], self[1], self[2], self[3])
    }
}

impl<V> AsRef<Quad<V>> for Quad<V>
where
    V: Vertex,
{
    fn as_ref(&self) -> &Quad<V> {
        self
    }
}

impl<V> AsRef<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn as_ref(&self) -> &[V; NUM_VERTICES as usize] {
        &self.0
    }
}

impl<V> Index<usize> for Quad<V>
where
    V: Vertex,
{
    type Output = V;
    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}
<commit_msg>Add `geom::quad` items for working with vertex indices.<commit_after>use geom::{tri, vertex, Tri, Vertex};
use math::EuclideanSpace;
use std::ops::{Deref, Index};

/// The number of vertices in a quad.
pub const NUM_VERTICES: u8 = 4;

/// The number of triangles that make up a quad.
pub const NUM_TRIANGLES: u8 = 2;

/// The same as `triangles`, but instead returns the vertex indices for each triangle.
pub const TRIANGLE_INDEX_TRIS: TrianglesIndexTris = [[0, 1, 2], [0, 2, 3]];
pub type TrianglesIndexTris = [[usize; tri::NUM_VERTICES as usize]; NUM_TRIANGLES as usize];

/// The number of indices used to describe each triangle in the quad.
pub const NUM_TRIANGLE_INDICES: u8 = 6;

/// The same as `triangles`, but instead returns the vertex indices for each triangle.
pub const TRIANGLE_INDICES: [usize; NUM_TRIANGLE_INDICES as usize] = [0, 1, 2, 0, 2, 3];

/// A quad represented by its four vertices.
#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
pub struct Quad<V = vertex::Default>(pub [V; NUM_VERTICES as usize]);

/// An `Iterator` yielding the two triangles that make up a quad.
#[derive(Clone, Debug)]
pub struct Triangles<V> {
    a: Option<Tri<V>>,
    b: Option<Tri<V>>,
}

/// A simple iterator yielding each vertex in a `Quad`.
#[derive(Clone, Debug)]
pub struct Vertices<V> {
    quad: Quad<V>,
    index: u8,
}

impl<V> Quad<V>
where
    V: Vertex,
{
    /// Produce an iterator yielding each vertex in the `Quad`.
    pub fn vertices(self) -> Vertices<V> {
        vertices(self)
    }

    /// Produce the centroid of the quad, aka the "mean"/"average" vertex.
    pub fn centroid(&self) -> V
    where
        V: EuclideanSpace,
    {
        centroid(self)
    }

    /// Triangulates the given quad, represented by four points that describe its edges in either
    /// clockwise or anti-clockwise order.
    ///
    /// # Example
    ///
    /// The following rectangle
    ///
    /// ```ignore
    ///  a        b
    ///   --------
    ///   |      |
    ///   |      |
    ///   |      |
    ///   --------
    ///  d        c
    ///
    /// ```
    ///
    /// given as
    ///
    /// ```ignore
    /// triangles([a, b, c, d])
    /// ```
    ///
    /// returns
    ///
    /// ```ignore
    /// (Tri([a, b, c]), Tri([a, c, d]))
    /// ```
    ///
    /// Here's a basic code example:
    ///
    /// ```
    /// extern crate nannou;
    ///
    /// use nannou::geom::{self, Quad, Tri};
    /// use nannou::math::pt2;
    ///
    /// fn main() {
    ///     let a = pt2(0.0, 1.0);
    ///     let b = pt2(1.0, 1.0);
    ///     let c = pt2(1.0, 0.0);
    ///     let d = pt2(0.0, 0.0);
    ///     let quad = Quad([a, b, c, d]);
    ///     let triangles = geom::quad::triangles(&quad);
    ///     assert_eq!(triangles, (Tri([a, b, c]), Tri([a, c, d])));
    /// }
    /// ```
    #[inline]
    pub fn triangles(&self) -> (Tri<V>, Tri<V>) {
        triangles(self)
    }

    /// The same as `triangles` but provided as an **Iterator**.
    pub fn triangles_iter(&self) -> Triangles<V> {
        triangles_iter(self)
    }

    /// Map the **Quad**'s vertices to a new type.
    pub fn map_vertices<F, V2>(self, mut map: F) -> Quad<V2>
    where
        F: FnMut(V) -> V2,
    {
        let (a, b, c, d) = self.into();
        Quad([map(a), map(b), map(c), map(d)])
    }
}

/// Produce an iterator yielding each vertex in the given **Quad**.
pub fn vertices<V>(quad: Quad<V>) -> Vertices<V> {
    let index = 0;
    Vertices { quad, index }
}

/// Produce the centroid of the quad, aka the "mean"/"average" vertex.
pub fn centroid<V>(quad: &Quad<V>) -> V
where
    V: EuclideanSpace,
{
    EuclideanSpace::centroid(&quad[..])
}

/// Triangulates the given quad, represented by four points that describe its edges in either
/// clockwise or anti-clockwise order.
///
/// # Example
///
/// The following rectangle
///
/// ```ignore
///
///  a        b
///   --------
///   |      |
///   |      |
///   |      |
///   --------
///  d        c
///
/// ```
///
/// given as
///
/// ```ignore
/// triangles([a, b, c, d])
/// ```
///
/// returns
///
/// ```ignore
/// (Tri([a, b, c]), Tri([a, c, d]))
/// ```
///
/// Here's a basic code example:
///
/// ```
/// extern crate nannou;
///
/// use nannou::geom::{self, Quad, Tri};
/// use nannou::math::pt2;
///
/// fn main() {
///     let a = pt2(0.0, 1.0);
///     let b = pt2(1.0, 1.0);
///     let c = pt2(1.0, 0.0);
///     let d = pt2(0.0, 0.0);
///     let quad = Quad([a, b, c, d]);
///     let triangles = geom::quad::triangles(&quad);
///     assert_eq!(triangles, (Tri([a, b, c]), Tri([a, c, d])));
/// }
/// ```
#[inline]
pub fn triangles<V>(q: &Quad<V>) -> (Tri<V>, Tri<V>)
where
    V: Vertex,
{
    let a = Tri::from_index_tri(&q.0, &TRIANGLE_INDEX_TRIS[0]);
    let b = Tri::from_index_tri(&q.0, &TRIANGLE_INDEX_TRIS[1]);
    (a, b)
}

/// The same as `triangles` but provided as an `Iterator`.
pub fn triangles_iter<V>(points: &Quad<V>) -> Triangles<V>
where
    V: Vertex,
{
    let (a, b) = triangles(points);
    Triangles {
        a: Some(a),
        b: Some(b),
    }
}

impl<V> Iterator for Triangles<V> {
    type Item = Tri<V>;
    fn next(&mut self) -> Option<Self::Item> {
        self.a.take().or_else(|| self.b.take())
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }
}

impl<V> DoubleEndedIterator for Triangles<V> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.b.take().or_else(|| self.a.take())
    }
}

impl<V> ExactSizeIterator for Triangles<V> {
    fn len(&self) -> usize {
        match (&self.a, &self.b) {
            (&Some(_), &Some(_)) => 2,
            (&None, &Some(_)) => 0,
            _ => 1,
        }
    }
}

impl<V> Iterator for Vertices<V>
where
    V: Clone,
{
    type Item = V;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index < NUM_VERTICES {
            let v = self.quad[self.index as usize].clone();
            self.index += 1;
            Some(v)
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.len();
        (len, Some(len))
    }
}

impl<V> ExactSizeIterator for Vertices<V>
where
    V: Clone,
{
    fn len(&self) -> usize {
        NUM_VERTICES as usize - self.index as usize
    }
}

impl<V> Deref for Quad<V> {
    type Target = [V; NUM_VERTICES as usize];
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<V> From<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn from(points: [V; NUM_VERTICES as usize]) -> Self {
        Quad(points)
    }
}

impl<V> From<(V, V, V, V)> for Quad<V>
where
    V: Vertex,
{
    fn from((a, b, c, d): (V, V, V, V)) -> Self {
        Quad([a, b, c, d])
    }
}

impl<V> Into<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn into(self) -> [V; NUM_VERTICES as usize] {
        self.0
    }
}

impl<V> Into<(V, V, V, V)> for Quad<V>
where
    V: Vertex,
{
    fn into(self) -> (V, V, V, V) {
        (self[0], self[1], self[2], self[3])
    }
}

impl<V> AsRef<Quad<V>> for Quad<V>
where
    V: Vertex,
{
    fn as_ref(&self) -> &Quad<V> {
        self
    }
}

impl<V> AsRef<[V; NUM_VERTICES as usize]> for Quad<V>
where
    V: Vertex,
{
    fn as_ref(&self) -> &[V; NUM_VERTICES as usize] {
        &self.0
    }
}

impl<V> Index<usize> for Quad<V>
where
    V: Vertex,
{
    type Output = V;
    fn index(&self, index: usize) -> &Self::Output {
        &self.0[index]
    }
}
<|endoftext|>


<commit_before>//! Functions for computing gradients of image intensities.

use image::{
    GenericImage,
    GrayImage,
    ImageBuffer,
    Luma
};

use definitions::Image;
use filter::filter3x3;

/// Sobel filter for detecting vertical gradients.
///
/// Used by the [`vertical_sobel`](fn.vertical_sobel.html) function.
pub static VERTICAL_SOBEL: [i32; 9] = [
    -1, -2, -1,
     0,  0,  0,
     1,  2,  1];

/// Sobel filter for detecting horizontal gradients.
///
/// Used by the [`horizontal_sobel`](fn.horizontal_sobel.html) function.
pub static HORIZONTAL_SOBEL: [i32; 9] = [
     -1, 0, 1,
     -2, 0, 2,
     -1, 0, 1];

/// Convolves an image with the [`HORIZONTAL_SOBEL`](static.HORIZONTAL_SOBEL.html)
/// kernel to detect horizontal gradients.
pub fn horizontal_sobel(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &HORIZONTAL_SOBEL)
}

/// Convolves an image with the [`VERTICAL_SOBEL`](static.VERTICAL_SOBEL.html)
/// kernel to detect vertical gradients.
pub fn vertical_sobel(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &VERTICAL_SOBEL)
}

/// Prewitt filter for detecting vertical gradients.
///
/// Used by the [`vertical_prewitt`](fn.vertical_prewitt.html) function.
pub static VERTICAL_PREWITT: [i32; 9] = [
    -1, -1, -1,
     0,  0,  0,
     1,  1,  1];

/// Prewitt filter for detecting horizontal gradients.
///
/// Used by the [`horizontal_prewitt`](fn.horizontal_prewitt.html) function.
pub static HORIZONTAL_PREWITT: [i32; 9] = [
     -1, 0, 1,
     -1, 0, 1,
     -1, 0, 1];

/// Convolves an image with the [`HORIZONTAL_PREWITT`](static.HORIZONTAL_PREWITT.html)
/// kernel to detect horizontal gradients.
pub fn horizontal_prewitt(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &HORIZONTAL_PREWITT)
}

/// Convolves an image with the [`VERTICAL_PREWITT`](static.VERTICAL_PREWITT.html)
/// kernel to detect vertical gradients.
pub fn vertical_prewitt(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &VERTICAL_PREWITT)
}

/// Returns the magnitudes of gradients in an image using Sobel filters.
pub fn sobel_gradients(image: &GrayImage) -> Image<Luma<u16>> {
    gradients(image, &HORIZONTAL_SOBEL, &VERTICAL_SOBEL)
}

/// Returns the magnitudes of gradients in an image using Prewitt filters.
pub fn prewitt_gradients(image: &GrayImage) -> Image<Luma<u16>> {
    gradients(image, &HORIZONTAL_PREWITT, &VERTICAL_PREWITT)
}

// TODO: Returns directions as well as magnitudes.
// TODO: Support filtering without allocating a fresh image - filtering functions could
// TODO: take some kind of pixel-sink. This would allow us to compute gradient magnitudes
// TODO: and directions without allocating intermediates for vertical and horizontal gradients.
fn gradients(image: &GrayImage, horizontal_kernel: &[i32; 9], vertical_kernel: &[i32; 9])
    -> Image<Luma<u16>> {
    let horizontal: ImageBuffer<Luma<i16>, Vec<i16>> = filter3x3(image, horizontal_kernel);
    let vertical: ImageBuffer<Luma<i16>, Vec<i16>> = filter3x3(image, vertical_kernel);

    let (width, height) = image.dimensions();
    let mut out = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            unsafe {
                let h = horizontal.unsafe_get_pixel(x, y)[0] as f32;
                let v = vertical.unsafe_get_pixel(x, y)[0] as f32;
                let m = (h.powi(2) + v.powi(2)).sqrt() as u16;
                out.unsafe_put_pixel(x, y, Luma([m]));
            }
        }
    }

    out
}

#[cfg(test)]
mod test {
    use super::{
        horizontal_sobel,
        vertical_sobel,
        horizontal_prewitt,
        vertical_prewitt
    };
    use image::{
        ImageBuffer,
        Luma
    };

    #[test]
    fn test_gradients_constant_image() {
        let image = ImageBuffer::from_pixel(5, 5, Luma([15u8]));
        let expected = ImageBuffer::from_pixel(5, 5, Luma([0i16]));
        assert_pixels_eq!(horizontal_sobel(&image), expected);
        assert_pixels_eq!(vertical_sobel(&image), expected);
        assert_pixels_eq!(horizontal_prewitt(&image), expected);
        assert_pixels_eq!(vertical_prewitt(&image), expected);
    }

    #[test]
    fn test_horizontal_sobel_gradient_image() {
        let image = gray_image!(
            3, 2, 1;
            6, 5, 4;
            9, 8, 7);

        let expected = gray_image_i16!(
            -4i16, -8i16, -4i16;
            -4i16, -8i16, -4i16;
            -4i16, -8i16, -4i16);

        let filtered = horizontal_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_sobel_gradient_image() {
        let image = gray_image!(
            3, 6, 9;
            2, 5, 8;
            1, 4, 7);

        let expected = gray_image_i16!(
            -4i16, -4i16, -4i16;
            -8i16, -8i16, -8i16;
            -4i16, -4i16, -4i16);

        let filtered = vertical_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_prewitt_gradient_image() {
        let image = gray_image!(
            3, 2, 1;
            6, 5, 4;
            9, 8, 7);

        let expected = gray_image_i16!(
            -3i16, -6i16, -3i16;
            -3i16, -6i16, -3i16;
            -3i16, -6i16, -3i16);

        let filtered = horizontal_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_prewitt_gradient_image() {
        let image = gray_image!(
            3, 6, 9;
            2, 5, 8;
            1, 4, 7);

        let expected = gray_image_i16!(
            -3i16, -3i16, -3i16;
            -6i16, -6i16, -6i16;
            -3i16, -3i16, -3i16);

        let filtered = vertical_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }
}
<commit_msg>Add benchmark for sobel_gradients<commit_after>//! Functions for computing gradients of image intensities.

use image::{
    GenericImage,
    GrayImage,
    ImageBuffer,
    Luma
};

use definitions::Image;
use filter::filter3x3;

/// Sobel filter for detecting vertical gradients.
///
/// Used by the [`vertical_sobel`](fn.vertical_sobel.html) function.
pub static VERTICAL_SOBEL: [i32; 9] = [
    -1, -2, -1,
     0,  0,  0,
     1,  2,  1];

/// Sobel filter for detecting horizontal gradients.
///
/// Used by the [`horizontal_sobel`](fn.horizontal_sobel.html) function.
pub static HORIZONTAL_SOBEL: [i32; 9] = [
     -1, 0, 1,
     -2, 0, 2,
     -1, 0, 1];

/// Convolves an image with the [`HORIZONTAL_SOBEL`](static.HORIZONTAL_SOBEL.html)
/// kernel to detect horizontal gradients.
pub fn horizontal_sobel(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &HORIZONTAL_SOBEL)
}

/// Convolves an image with the [`VERTICAL_SOBEL`](static.VERTICAL_SOBEL.html)
/// kernel to detect vertical gradients.
pub fn vertical_sobel(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &VERTICAL_SOBEL)
}

/// Prewitt filter for detecting vertical gradients.
///
/// Used by the [`vertical_prewitt`](fn.vertical_prewitt.html) function.
pub static VERTICAL_PREWITT: [i32; 9] = [
    -1, -1, -1,
     0,  0,  0,
     1,  1,  1];

/// Prewitt filter for detecting horizontal gradients.
///
/// Used by the [`horizontal_prewitt`](fn.horizontal_prewitt.html) function.
pub static HORIZONTAL_PREWITT: [i32; 9] = [
     -1, 0, 1,
     -1, 0, 1,
     -1, 0, 1];

/// Convolves an image with the [`HORIZONTAL_PREWITT`](static.HORIZONTAL_PREWITT.html)
/// kernel to detect horizontal gradients.
pub fn horizontal_prewitt(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &HORIZONTAL_PREWITT)
}

/// Convolves an image with the [`VERTICAL_PREWITT`](static.VERTICAL_PREWITT.html)
/// kernel to detect vertical gradients.
pub fn vertical_prewitt(image: &GrayImage) -> Image<Luma<i16>> {
    filter3x3(image, &VERTICAL_PREWITT)
}

/// Returns the magnitudes of gradients in an image using Sobel filters.
pub fn sobel_gradients(image: &GrayImage) -> Image<Luma<u16>> {
    gradients(image, &HORIZONTAL_SOBEL, &VERTICAL_SOBEL)
}

/// Returns the magnitudes of gradients in an image using Prewitt filters.
pub fn prewitt_gradients(image: &GrayImage) -> Image<Luma<u16>> {
    gradients(image, &HORIZONTAL_PREWITT, &VERTICAL_PREWITT)
}

// TODO: Returns directions as well as magnitudes.
// TODO: Support filtering without allocating a fresh image - filtering functions could
// TODO: take some kind of pixel-sink. This would allow us to compute gradient magnitudes
// TODO: and directions without allocating intermediates for vertical and horizontal gradients.
fn gradients(image: &GrayImage, horizontal_kernel: &[i32; 9], vertical_kernel: &[i32; 9]) -> Image<Luma<u16>> {
    let horizontal: Image<Luma<i16>> = filter3x3(image, horizontal_kernel);
    let vertical: Image<Luma<i16>> = filter3x3(image, vertical_kernel);

    let (width, height) = image.dimensions();
    let mut out = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            unsafe {
                let h = horizontal.unsafe_get_pixel(x, y)[0] as f32;
                let v = vertical.unsafe_get_pixel(x, y)[0] as f32;
                let m = (h.powi(2) + v.powi(2)).sqrt() as u16;
                out.unsafe_put_pixel(x, y, Luma([m]));
            }
        }
    }

    out
}

#[cfg(test)]
mod test {
    use super::*;
    use image::{
        ImageBuffer,
        Luma
    };
    use test::{Bencher, black_box};
    use utils::gray_bench_image;

    #[test]
    fn test_gradients_constant_image() {
        let image = ImageBuffer::from_pixel(5, 5, Luma([15u8]));
        let expected = ImageBuffer::from_pixel(5, 5, Luma([0i16]));
        assert_pixels_eq!(horizontal_sobel(&image), expected);
        assert_pixels_eq!(vertical_sobel(&image), expected);
        assert_pixels_eq!(horizontal_prewitt(&image), expected);
        assert_pixels_eq!(vertical_prewitt(&image), expected);
    }

    #[test]
    fn test_horizontal_sobel_gradient_image() {
        let image = gray_image!(
            3, 2, 1;
            6, 5, 4;
            9, 8, 7);

        let expected = gray_image_i16!(
            -4i16, -8i16, -4i16;
            -4i16, -8i16, -4i16;
            -4i16, -8i16, -4i16);

        let filtered = horizontal_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_sobel_gradient_image() {
        let image = gray_image!(
            3, 6, 9;
            2, 5, 8;
            1, 4, 7);

        let expected = gray_image_i16!(
            -4i16, -4i16, -4i16;
            -8i16, -8i16, -8i16;
            -4i16, -4i16, -4i16);

        let filtered = vertical_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_prewitt_gradient_image() {
        let image = gray_image!(
            3, 2, 1;
            6, 5, 4;
            9, 8, 7);

        let expected = gray_image_i16!(
            -3i16, -6i16, -3i16;
            -3i16, -6i16, -3i16;
            -3i16, -6i16, -3i16);

        let filtered = horizontal_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_prewitt_gradient_image() {
        let image = gray_image!(
            3, 6, 9;
            2, 5, 8;
            1, 4, 7);

        let expected = gray_image_i16!(
            -3i16, -3i16, -3i16;
            -6i16, -6i16, -6i16;
            -3i16, -3i16, -3i16);

        let filtered = vertical_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[bench]
    fn bench_sobel_gradients(b: &mut Bencher) {
        let image = gray_bench_image(500, 500);
        b.iter(|| {
            let gradients = sobel_gradients(&image);
            black_box(gradients);
            });
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2015, The Rust-GNOME Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! callback idle functions

pub mod idle {
    use std::cell::RefCell;
    use std::ops::DerefMut;
    use std::mem::transmute;
    use ffi::{gboolean, g_idle_add_full};
    use translate::ToGlib;

    pub struct Continue(pub bool);

    impl ToGlib for Continue {
        type GlibType = gboolean;

        #[inline]
        fn to_glib(&self) -> gboolean {
            self.0.to_glib()
        }
    }

    // Box::into_raw stability workaround
    unsafe fn into_raw<T>(b: Box<T>) -> *mut T { transmute(b) }

    extern "C" fn trampoline(func: &RefCell<Box<FnMut() -> Continue + 'static>>) -> gboolean {
        func.borrow_mut().deref_mut()().to_glib()
    }

    extern "C" fn destroy_closure(ptr: gpointer) {
        unsafe {
            // Box::from_raw API stability workaround
            let ptr = ptr as *mut RefCell<Box<FnMut() -> Continue + 'static>>;
            let _: Box<RefCell<Box<FnMut() -> Continue + 'static>>> = ::std::mem::transmute(ptr);
        }
    }

    pub fn add<F>(func: F) -> u32
        where F: FnMut() -> Continue + 'static {
        let f: Box<RefCell<Box<FnMut() -> Continue + 'static>>> = Box::new(RefCell::new(Box::new(func)));
        unsafe {
            g_idle_add_full(200 /* = G_PRIORITY_DEFAULT_IDLE */,
                transmute(trampoline), into_raw(f) as ffi::gpointer, destroy_closure)
        }
    }
}
<commit_msg>Idle: Fix: imports<commit_after>// Copyright 2013-2015, The Rust-GNOME Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

//! callback idle functions

pub mod idle {
    use std::cell::RefCell;
    use std::ops::DerefMut;
    use std::mem::transmute;
    use ffi::{gboolean, gpointer, g_idle_add_full};
    use translate::ToGlib;

    pub struct Continue(pub bool);

    impl ToGlib for Continue {
        type GlibType = gboolean;

        #[inline]
        fn to_glib(&self) -> gboolean {
            self.0.to_glib()
        }
    }

    // Box::into_raw stability workaround
    unsafe fn into_raw<T>(b: Box<T>) -> *mut T { transmute(b) }

    extern "C" fn trampoline(func: &RefCell<Box<FnMut() -> Continue + 'static>>) -> gboolean {
        func.borrow_mut().deref_mut()().to_glib()
    }

    extern "C" fn destroy_closure(ptr: gpointer) {
        unsafe {
            // Box::from_raw API stability workaround
            let ptr = ptr as *mut RefCell<Box<FnMut() -> Continue + 'static>>;
            let _: Box<RefCell<Box<FnMut() -> Continue + 'static>>> = ::std::mem::transmute(ptr);
        }
    }

    pub fn add<F>(func: F) -> u32
        where F: FnMut() -> Continue + 'static {
        let f: Box<RefCell<Box<FnMut() -> Continue + 'static>>> = Box::new(RefCell::new(Box::new(func)));
        unsafe {
            g_idle_add_full(200 /* = G_PRIORITY_DEFAULT_IDLE */,
                transmute(trampoline), into_raw(f) as gpointer, destroy_closure)
        }
    }
}
<|endoftext|>


<commit_before>// http://rosettacode.org/wiki/K-means%2B%2B_clustering

#![feature(test)]

extern crate csv;
extern crate getopts;
extern crate gnuplot;
extern crate nalgebra;
extern crate num;
extern crate rand;
extern crate rustc_serialize;
extern crate test;

use getopts::Options;
use gnuplot::{Axes2D, AxesCommon, Color, Figure, Fix, PointSize, PointSymbol};
use nalgebra::{DVec, Iterable};
use rand::{Rng, SeedableRng, StdRng};
use rand::distributions::{IndependentSample, Range};
use std::f64::consts::PI;
use std::env;

type Point = DVec<f64>;

struct Cluster<'a> {
    members: Vec<&'a Point>,
    center: Point,
}

struct Stats {
    centroids: Vec<Point>,
    mean_d_from_centroid: DVec<f64>,
}

// DVec doesn't implement BaseFloat, so I need my own distance fn
fn sqdist(p1: &Point, p2: &Point) -> f64 {
    (p1.clone() - p2.clone()).iter().map(|x| x * x).fold(0f64, |a, b| a + b)
}

// Returns (distance^2, index) tuple of winning point
fn nearest(p: &Point, candidates: &Vec<Point>) -> (f64, usize) {
    let (dsquared, the_index) = candidates.iter()
                                          .enumerate()
                                          .fold((sqdist(p, &candidates[0]), 0),
                                                |(d, index), next| {
                                                    let dprime = sqdist(p, &candidates[next.0]);
                                                    if dprime < d {
                                                        (dprime, next.0)
                                                    } else {
                                                        (d, index)
                                                    }
                                                });
    (dsquared, the_index)
}

// Computes starting centroids and makes initial assignments
fn kpp(points: &Vec<Point>, k: usize, rng: &mut StdRng) -> Stats {
    let mut centroids: Vec<Point> = Vec::new();
    // Random point for first centroid guess:
    centroids.push(points[rng.gen::<usize>() % points.len()].clone());
    let mut dists: Vec<f64> = vec![0f64; points.len()];

    for _ in 1..k {
        let mut sum = 0f64;
        for (j, p) in points.iter().enumerate() {
            let (dsquared, _) = nearest(&p, &centroids);
            dists[j] = dsquared;
            sum += dsquared;
        }

        // This bit chooses the next cluster center with a probability proportional to d^2
        sum *= rng.next_f64();
        for (j, d) in dists.iter().enumerate() {
            sum -= *d;
            if sum <= 0f64 {
                centroids.push(points[j].clone());
                break;
            }
        }
    }

    let clusters = assign_clusters(points, &centroids);
    compute_stats(&clusters)
}

fn assign_clusters<'a>(points: &'a Vec<Point>, centroids: &Vec<Point>) -> Vec<Cluster<'a>> {
    let mut clusters: Vec<Cluster> = Vec::new();

    for _ in 0..centroids.len() {
        clusters.push(Cluster {
            members: Vec::new(),
            center: DVec::new_zeros(points[0].len()),
        });
    }

    for p in points.iter() {
        let (_, nearest_index) = nearest(p, centroids);
        clusters[nearest_index].center = clusters[nearest_index].center.clone() + p.clone();
        clusters[nearest_index].members.push(p);
    }

    for i in 0..clusters.len() {
        clusters[i].center = clusters[i].center.clone() / clusters[i].members.len() as f64;
    }

    clusters
}

// Computes centroids and mean-distance-from-centroid for each cluster
fn compute_stats(clusters: &Vec<Cluster>) -> Stats {
    let mut centroids = Vec::new();
    let mut means_vec = Vec::new();

    for c in clusters.iter() {
        let pts = &c.members;
        let seed: DVec<f64> = DVec::new_zeros(pts[0].len());
        let centroid = pts.iter().fold(seed, |a, &b| a + b.clone()) / pts.len() as f64;
        means_vec.push(pts.iter().fold(0f64, |acc, pt| acc + sqdist(pt, &centroid).sqrt()) /
                       pts.len() as f64);
        centroids.push(centroid);
    }

    Stats {
        centroids: centroids,
        mean_d_from_centroid: DVec::from_slice(means_vec.len(), means_vec.as_slice()),
    }
}

fn lloyd<'a>(points: &'a Vec<Point>,
             k: usize,
             stoppage_delta: f64,
             max_iter: u32,
             rng: &mut StdRng)
             -> (Vec<Cluster<'a>>, Stats) {

    let mut clusters = Vec::new();
    // Choose starting centroids and make initial assignments
    let mut stats = kpp(points, k, rng);

    for i in 1..max_iter {
        let last_means: DVec<f64> = stats.mean_d_from_centroid.clone();
        clusters = assign_clusters(points, &stats.centroids);
        stats = compute_stats(&clusters);
        let err = sqdist(&stats.mean_d_from_centroid, &last_means).sqrt();
        if err < stoppage_delta {
            println!("Stoppage condition reached on iteration {}", i);
            return (clusters, stats);
        }
        // Console output
        print!("Iter {}: ", i);
        for (cen, mu) in stats.centroids.iter().zip(stats.mean_d_from_centroid.iter()) {
            print_dvec(cen);
            print!(" {:1.2} | ", mu);
        }
        print!("{:1.5}\n", err);
    }

    println!("Stoppage condition not reached by iteration {}", max_iter);
    (clusters, stats)
}

// Uniform sampling on the unit disk
fn generate_points(n: u32, rng: &mut StdRng) -> Vec<Point> {
    let r_range = Range::new(0f64, 1f64);
    let theta_range = Range::new(0f64, 2f64 * PI);
    let mut points: Vec<Point> = Vec::new();

    for _ in 0..n {
        let root_r = r_range.ind_sample(rng).sqrt();
        let theta = theta_range.ind_sample(rng);
        points.push(DVec::<f64>::from_slice(2, &[root_r * theta.cos(), root_r * theta.sin()]));
    }

    points
}

// Plot clusters (2d only). Wrapper idiom allows prep_axes to mutate the Axes2D.
fn viz(clusters: Vec<Cluster>, stats: Stats, k: usize, n: u32, e: f64) {
    let mut fg = Figure::new();
    prep_axes(&mut fg, clusters, stats, k, n, e);
    fg.show();
}

fn prep_axes(fg: &mut Figure, clusters: Vec<Cluster>, stats: Stats, k: usize, n: u32,
    e: f64) {
    let a: & mut Axes2D = fg.axes2d();
    let title: String = format!("k = {}, n = {}, e = {:4}", k, n, e);
    let centroids_x = stats.centroids.iter().map(|c| c[0]);
    let centroids_y = stats.centroids.iter().map(|c| c[1]);

    for cluster in clusters.iter() {
        a.points(cluster.members.iter().map(|p| p[0]), cluster.members.iter().map(|p| p[1]),
            &[PointSymbol('O'), PointSize(0.25)]);
    }

    a.set_aspect_ratio(Fix(1.0))
        .points(centroids_x, centroids_y, &[PointSymbol('o'), PointSize(1.5), Color("black")])
        .set_title(&title[..], &[]);
}

fn print_dvec(v: &DVec<f64>) {
    print!("(");
    for elem in v.at.iter().take(v.len() - 1) {
        print!("{:+1.2}, ", elem)
    }
    print!("{:+1.2})", v.at.iter().last().unwrap());
}

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

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut k: usize = 7;
    let mut n: u32 = 30000;
    let mut e: f64 = 1e-3;
    let max_iterations = 100u32;

    let mut opts = Options::new();
    opts.optflag("?", "help", "Print this help menu");
    opts.optopt("k", "", "Number of clusters to assign (default: 7)", "<clusters>");
    opts.optopt("n", "", "Operate on this many points on the unit disk (default: 30000)", "<pts>");
    opts.optopt("e", "", "Min delta in norm of successive cluster centroids to continue (default: \
        1e-3)", "<eps>");
    opts.optopt("f", "", "Read points from file (overrides -n)", "<csv>");

    let program = args[0].clone();
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("?") {
        print_usage(&program, opts);
        return;
    }
    match matches.opt_str("k") {
        None => {}
        Some(x) => k = x.parse::<usize>().unwrap(),
    };
    match matches.opt_str("n") {
        None => {}
        Some(x) => n = x.parse::<u32>().unwrap(),
    };
    match matches.opt_str("e") {
        None => {}
        Some(x) => e = x.parse::<f64>().unwrap(),
    };

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut points: Vec<Point>;

    match matches.opt_str("f") {
        None => {
            // Proceed with random 2d data
            points = generate_points(n, &mut rng)
        }
        Some(file) => {
            points = Vec::new();
            let mut rdr = csv::Reader::from_file(file.clone()).unwrap();
            for row in rdr.records().map(|r| r.unwrap()) {
                // row is Vec<String>
                let floats: Vec<f64> = row.iter().map(|s| s.parse::<f64>().unwrap()).collect();
                points.push(DVec::<f64>::from_slice(floats.len(), floats.as_slice()));
            }
            assert!(points.iter().all(|v| v.len() == points[0].len()));
            n = points.len() as u32;
            println!("Read {} points from {}", points.len(), file.clone());
        }
    };

    assert!(points.len() >= k);
    let (clusters, stats) = lloyd(&points, k, e, max_iterations, &mut rng);

    println!(" k       centroid{}mean dist    pop",
        std::iter::repeat(" ").take((points[0].len() - 2) * 7 + 7).collect::<String>());
    println!("===  {}  ===========  =====",
        std::iter::repeat("=").take(points[0].len()*7 + 2).collect::<String>());
    for i in 0..clusters.len() {
        print!(" {:>1}    ", i);
        print_dvec(&stats.centroids[i]);
        print!("      {:1.2}       {:>4}\n",
               stats.mean_d_from_centroid[i],
               clusters[i].members.len());
    }

    if points[0].len() == 2 {
        viz(clusters, stats, k, n, e)
    }
}

#[cfg(test)]
mod tests {
    use super::{generate_points, lloyd};
    use nalgebra::Iterable;
    use rand::{SeedableRng, StdRng};
    // use test::Bencher;

    #[test]
    fn test_lloyd2d() {
        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let points = generate_points(1000, &mut rng);

        let (clusters, stats) = lloyd(&points, 4, 0.001, 100, &mut rng);

        assert!(clusters.len() == 4);
        for i in 0..clusters.len() {
            assert!(clusters[i].members.len() > 0);
        }
        assert!(stats.mean_d_from_centroid.iter().all(|d| *d > 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] >= 0f64 && p[1] >= 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] >= 0f64 && p[1] < 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] < 0f64 && p[1] < 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] < 0f64 && p[1] >= 0f64));
    }

}
<commit_msg>changed plot prep function to closure, rustfmt<commit_after>// http://rosettacode.org/wiki/K-means%2B%2B_clustering

#![feature(test)]

extern crate csv;
extern crate getopts;
extern crate gnuplot;
extern crate nalgebra;
extern crate num;
extern crate rand;
extern crate rustc_serialize;
extern crate test;

use getopts::Options;
use gnuplot::{Axes2D, AxesCommon, Color, Figure, Fix, PointSize, PointSymbol};
use nalgebra::{DVec, Iterable};
use rand::{Rng, SeedableRng, StdRng};
use rand::distributions::{IndependentSample, Range};
use std::f64::consts::PI;
use std::env;

type Point = DVec<f64>;

struct Cluster<'a> {
    members: Vec<&'a Point>,
    center: Point,
}

struct Stats {
    centroids: Vec<Point>,
    mean_d_from_centroid: DVec<f64>,
}

// DVec doesn't implement BaseFloat, so I need my own distance fn
fn sqdist(p1: &Point, p2: &Point) -> f64 {
    (p1.clone() - p2.clone()).iter().map(|x| x * x).fold(0f64, |a, b| a + b)
}

// Returns (distance^2, index) tuple of winning point
fn nearest(p: &Point, candidates: &Vec<Point>) -> (f64, usize) {
    let (dsquared, the_index) = candidates.iter()
                                          .enumerate()
                                          .fold((sqdist(p, &candidates[0]), 0),
                                                |(d, index), next| {
                                                    let dprime = sqdist(p, &candidates[next.0]);
                                                    if dprime < d {
                                                        (dprime, next.0)
                                                    } else {
                                                        (d, index)
                                                    }
                                                });
    (dsquared, the_index)
}

// Computes starting centroids and makes initial assignments
fn kpp(points: &Vec<Point>, k: usize, rng: &mut StdRng) -> Stats {
    let mut centroids: Vec<Point> = Vec::new();
    // Random point for first centroid guess:
    centroids.push(points[rng.gen::<usize>() % points.len()].clone());
    let mut dists: Vec<f64> = vec![0f64; points.len()];

    for _ in 1..k {
        let mut sum = 0f64;
        for (j, p) in points.iter().enumerate() {
            let (dsquared, _) = nearest(&p, &centroids);
            dists[j] = dsquared;
            sum += dsquared;
        }

        // This bit chooses the next cluster center with a probability proportional to d^2
        sum *= rng.next_f64();
        for (j, d) in dists.iter().enumerate() {
            sum -= *d;
            if sum <= 0f64 {
                centroids.push(points[j].clone());
                break;
            }
        }
    }

    let clusters = assign_clusters(points, &centroids);
    compute_stats(&clusters)
}

fn assign_clusters<'a>(points: &'a Vec<Point>, centroids: &Vec<Point>) -> Vec<Cluster<'a>> {
    let mut clusters: Vec<Cluster> = Vec::new();

    for _ in 0..centroids.len() {
        clusters.push(Cluster {
            members: Vec::new(),
            center: DVec::new_zeros(points[0].len()),
        });
    }

    for p in points.iter() {
        let (_, nearest_index) = nearest(p, centroids);
        clusters[nearest_index].center = clusters[nearest_index].center.clone() + p.clone();
        clusters[nearest_index].members.push(p);
    }

    for i in 0..clusters.len() {
        clusters[i].center = clusters[i].center.clone() / clusters[i].members.len() as f64;
    }

    clusters
}

// Computes centroids and mean-distance-from-centroid for each cluster
fn compute_stats(clusters: &Vec<Cluster>) -> Stats {
    let mut centroids = Vec::new();
    let mut means_vec = Vec::new();

    for c in clusters.iter() {
        let pts = &c.members;
        let seed: DVec<f64> = DVec::new_zeros(pts[0].len());
        let centroid = pts.iter().fold(seed, |a, &b| a + b.clone()) / pts.len() as f64;
        means_vec.push(pts.iter().fold(0f64, |acc, pt| acc + sqdist(pt, &centroid).sqrt()) /
                       pts.len() as f64);
        centroids.push(centroid);
    }

    Stats {
        centroids: centroids,
        mean_d_from_centroid: DVec::from_slice(means_vec.len(), means_vec.as_slice()),
    }
}

fn lloyd<'a>(points: &'a Vec<Point>,
             k: usize,
             stoppage_delta: f64,
             max_iter: u32,
             rng: &mut StdRng)
             -> (Vec<Cluster<'a>>, Stats) {

    let mut clusters = Vec::new();
    // Choose starting centroids and make initial assignments
    let mut stats = kpp(points, k, rng);

    for i in 1..max_iter {
        let last_means: DVec<f64> = stats.mean_d_from_centroid.clone();
        clusters = assign_clusters(points, &stats.centroids);
        stats = compute_stats(&clusters);
        let err = sqdist(&stats.mean_d_from_centroid, &last_means).sqrt();
        if err < stoppage_delta {
            println!("Stoppage condition reached on iteration {}", i);
            return (clusters, stats);
        }
        // Console output
        print!("Iter {}: ", i);
        for (cen, mu) in stats.centroids.iter().zip(stats.mean_d_from_centroid.iter()) {
            print_dvec(cen);
            print!(" {:1.2} | ", mu);
        }
        print!("{:1.5}\n", err);
    }

    println!("Stoppage condition not reached by iteration {}", max_iter);
    (clusters, stats)
}

// Uniform sampling on the unit disk
fn generate_points(n: u32, rng: &mut StdRng) -> Vec<Point> {
    let r_range = Range::new(0f64, 1f64);
    let theta_range = Range::new(0f64, 2f64 * PI);
    let mut points: Vec<Point> = Vec::new();

    for _ in 0..n {
        let root_r = r_range.ind_sample(rng).sqrt();
        let theta = theta_range.ind_sample(rng);
        points.push(DVec::<f64>::from_slice(2, &[root_r * theta.cos(), root_r * theta.sin()]));
    }

    points
}

// Plot clusters (2d only). Closure idiom allows us to borrow and mutate the Axes2D.
fn viz(clusters: Vec<Cluster>, stats: Stats, k: usize, n: u32, e: f64) {
    let mut fg = Figure::new();
    {
        let prep = |fg: &mut Figure| {
            let axes: &mut Axes2D = fg.axes2d();
            let title: String = format!("k = {}, n = {}, e = {:4}", k, n, e);
            let centroids_x = stats.centroids.iter().map(|c| c[0]);
            let centroids_y = stats.centroids.iter().map(|c| c[1]);
            for cluster in clusters.iter() {
                axes.points(cluster.members.iter().map(|p| p[0]),
                            cluster.members
                                   .iter()
                                   .map(|p| p[1]),
                            &[PointSymbol('O'), PointSize(0.25)]);
            }
            axes.set_aspect_ratio(Fix(1.0))
                .points(centroids_x,
                        centroids_y,
                        &[PointSymbol('o'), PointSize(1.5), Color("black")])
                .set_title(&title[..], &[]);
        };
        prep(&mut fg);
    }
    fg.show();
}

fn print_dvec(v: &DVec<f64>) {
    print!("(");
    for elem in v.at.iter().take(v.len() - 1) {
        print!("{:+1.2}, ", elem)
    }
    print!("{:+1.2})", v.at.iter().last().unwrap());
}

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

fn main() {
    let args: Vec<String> = env::args().collect();
    let mut k: usize = 7;
    let mut n: u32 = 30000;
    let mut e: f64 = 1e-3;
    let max_iterations = 100u32;

    let mut opts = Options::new();
    opts.optflag("?", "help", "Print this help menu");
    opts.optopt("k",
                "",
                "Number of clusters to assign (default: 7)",
                "<clusters>");
    opts.optopt("n",
                "",
                "Operate on this many points on the unit disk (default: 30000)",
                "<pts>");
    opts.optopt("e",
                "",
                "Min delta in norm of successive cluster centroids to continue (default: 1e-3)",
                "<eps>");
    opts.optopt("f", "", "Read points from file (overrides -n)", "<csv>");

    let program = args[0].clone();
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => m,
        Err(f) => panic!(f.to_string()),
    };
    if matches.opt_present("?") {
        print_usage(&program, opts);
        return;
    }
    match matches.opt_str("k") {
        None => {}
        Some(x) => k = x.parse::<usize>().unwrap(),
    };
    match matches.opt_str("n") {
        None => {}
        Some(x) => n = x.parse::<u32>().unwrap(),
    };
    match matches.opt_str("e") {
        None => {}
        Some(x) => e = x.parse::<f64>().unwrap(),
    };

    let seed: &[_] = &[1, 2, 3, 4];
    let mut rng: StdRng = SeedableRng::from_seed(seed);

    let mut points: Vec<Point>;

    match matches.opt_str("f") {
        None => {
            // Proceed with random 2d data
            points = generate_points(n, &mut rng)
        }
        Some(file) => {
            points = Vec::new();
            let mut rdr = csv::Reader::from_file(file.clone()).unwrap();
            for row in rdr.records().map(|r| r.unwrap()) {
                // row is Vec<String>
                let floats: Vec<f64> = row.iter().map(|s| s.parse::<f64>().unwrap()).collect();
                points.push(DVec::<f64>::from_slice(floats.len(), floats.as_slice()));
            }
            assert!(points.iter().all(|v| v.len() == points[0].len()));
            n = points.len() as u32;
            println!("Read {} points from {}", points.len(), file.clone());
        }
    };

    assert!(points.len() >= k);
    let (clusters, stats) = lloyd(&points, k, e, max_iterations, &mut rng);

    println!(" k       centroid{}mean dist    pop",
             std::iter::repeat(" ").take((points[0].len() - 2) * 7 + 7).collect::<String>());
    println!("===  {}  ===========  =====",
             std::iter::repeat("=").take(points[0].len() * 7 + 2).collect::<String>());
    for i in 0..clusters.len() {
        print!(" {:>1}    ", i);
        print_dvec(&stats.centroids[i]);
        print!("      {:1.2}       {:>4}\n",
               stats.mean_d_from_centroid[i],
               clusters[i].members.len());
    }

    if points[0].len() == 2 {
        viz(clusters, stats, k, n, e)
    }
}

#[cfg(test)]
mod tests {
    use super::{generate_points, lloyd};
    use nalgebra::Iterable;
    use rand::{SeedableRng, StdRng};
    // use test::Bencher;

    #[test]
    fn test_lloyd2d() {
        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng: StdRng = SeedableRng::from_seed(seed);
        let points = generate_points(1000, &mut rng);

        let (clusters, stats) = lloyd(&points, 4, 0.001, 100, &mut rng);

        assert!(clusters.len() == 4);
        for i in 0..clusters.len() {
            assert!(clusters[i].members.len() > 0);
        }
        assert!(stats.mean_d_from_centroid.iter().all(|d| *d > 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] >= 0f64 && p[1] >= 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] >= 0f64 && p[1] < 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] < 0f64 && p[1] < 0f64));
        assert!(stats.centroids.iter().any(|p| p[0] < 0f64 && p[1] >= 0f64));
    }

}
<|endoftext|>