<commit_before>// (C)opyleft 2013 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)];

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

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;
}

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 {
    priv db: GeoIP_
}

pub struct ASInfo {
    asn: uint,
    name: ~str,
    netmask: uint
}

impl ToStr for ASInfo {
    fn to_str(&self) -> ~str {
        format!("{}\t{}", self.asn, self.name)
    }
}

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, ~str> {
        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)
        };
        match db.is_null() {
            true => Err(format!("Can't open {}", file)),
            false => Ok(GeoIP {
                    db: db
                })
        }
    }

    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 asinfo = ASInfo {
            asn: asn,
            name: name.to_owned(),
            netmask: gl.netmask as uint
        };
        Some(asinfo)
    }
}

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.contains("Telefun"));
    assert!(res.netmask == 22);
}
<commit_msg>ToStr -> Show<commit_after>// (C)opyleft 2013 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)];

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

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;
}

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 {
    priv db: GeoIP_
}

pub struct ASInfo {
    asn: uint,
    name: ~str,
    netmask: uint
}

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

impl GeoIP {
    pub fn open(path: &Path, options: Options) -> Result<GeoIP, ~str> {
        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)
        };
        match db.is_null() {
            true => Err(format!("Can't open {}", file)),
            false => Ok(GeoIP {
                    db: db
                })
        }
    }

    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 asinfo = ASInfo {
            asn: asn,
            name: name.to_owned(),
            netmask: gl.netmask as uint
        };
        Some(asinfo)
    }
}

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.contains("Telefun"));
    assert!(res.netmask == 22);
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Sorting_algorithms/Heapsort

// This is ported from the Dart heap sort implementation

fn heap_sort<T: Ord>(a: &mut [T]) {
    let count = a.len();

    // first place 'a' in max-heap order
    heapify(a, count);

    let mut end = count - 1;
    while end > 0 {
        // swap the root (maximum value) of the heap with the
        // last element of the heap
        a.swap(0, end);

        // put the heap back in max-heap order
        sift_down(a, 0, end - 1);

        // decrement the size of the heap so that the previous
        // max value will stay in its proper place
        end -= 1;
    }
}

fn heapify<T: Ord>(a: &mut [T], count: uint) {
    // start is assigned the index in 'a' of the last parent node
    let mut start:int = (count - 2 / 2) as int; // binary heap

    while start >= 0 {
        // sift down the node at index 'start' to the proper place
        // such that all nodes below the 'start' index are in heap order
        sift_down(a, start as uint, count - 1);
        start -= 1;
    }
}

fn sift_down<T: Ord>(a: &mut [T], start: uint, end: uint) {
    // end represents the limit of how far down the heap to shift
    let mut root = start;

    // while the root has at least one child
    while (root*2 + 1) <= end {
        // root*2+1 points to the left child
        let mut child:uint = root*2 + 1 as uint;

        // if the chile has a sibling and the child's value is less that its sibling's...
        if child + 1 <= end && a[child] < a[child + 1] {
            // .. then point to the right child instead
            child = child + 1;
        }

        // out of max-heap order
        if a[root] < a[child] {
            a.swap(root, child);
            // repeat to continue shifting down the child now
            root = child;
        } else {
            return;
        }
    }
}

#[cfg(not(test))]
pub fn main() {
    let mut arr = [1u,5,2,7,3,9,4,6,8];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());

    let mut arr = [1u,2,3,4,5,6,7,8,9];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());

    let mut arr = [9u,8,7,6,5,4,3,2,1];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());
}

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

    #[test]
    fn sorted() {
        let mut arr = [1u, 2, 3, 4, 6, 8];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [1u, 2, 3, 4, 6, 8].as_slice());
    }

    #[test]
    fn reverse() {
        let mut arr = [8i, 6, 4, 3, 2, 1];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [1i, 2, 3, 4, 6, 8].as_slice());
    }

    #[test]
    fn random() {
        let mut arr = [12u, 54, 2, 93, 13, 43, 15, 299, 234];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [2u, 12, 13, 15, 43, 54, 93, 234, 299].as_slice());
    }
}<commit_msg>Add tests for empty and 1-length arrays. Also add some checks <commit_after>// Implements http://rosettacode.org/wiki/Sorting_algorithms/Heapsort

// This is ported from the Dart heap sort implementation
fn heap_sort<T: Ord>(a: &mut [T]) {
    let count = a.len();

    if count == 0 {
        return;
    }

    // first place 'a' in max-heap order
    heapify(a, count);

    let mut end = count - 1;
    while end > 0 {
        // swap the root (maximum value) of the heap with the
        // last element of the heap
        a.swap(0, end);

        // put the heap back in max-heap order
        sift_down(a, 0, end - 1);

        // decrement the size of the heap so that the previous
        // max value will stay in its proper place
        end -= 1;
    }
}

fn heapify<T: Ord>(a: &mut [T], count: uint) {
    if count < 2 {
        return;
    }

    // start is assigned the index in 'a' of the last parent node
    let mut start:int = count as int - 2 / 2; // binary heap

    while start >= 0 {
        // sift down the node at index 'start' to the proper place
        // such that all nodes below the 'start' index are in heap order
        sift_down(a, start as uint, count - 1);
        start -= 1;
    }
}

fn sift_down<T: Ord>(a: &mut [T], start: uint, end: uint) {
    // end represents the limit of how far down the heap to shift
    let mut root = start;

    // while the root has at least one child
    while (root*2 + 1) <= end {
        // root*2+1 points to the left child
        let mut child:uint = root*2 + 1 as uint;

        // if the chile has a sibling and the child's value is less that its sibling's...
        if child + 1 <= end && a[child] < a[child + 1] {
            // .. then point to the right child instead
            child = child + 1;
        }

        // out of max-heap order
        if a[root] < a[child] {
            a.swap(root, child);
            // repeat to continue shifting down the child now
            root = child;
        } else {
            return;
        }
    }
}

#[cfg(not(test))]
pub fn main() {
    let mut arr = [1u,5,2,7,3,9,4,6,8];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());

    let mut arr = [1u,2,3,4,5,6,7,8,9];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());

    let mut arr = [9u,8,7,6,5,4,3,2,1];
    heap_sort(arr);
    println!("After sort: {}", arr.as_slice());
}

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

    #[test]
    fn sorted() {
        let mut arr = [1u, 2, 3, 4, 6, 8];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [1u, 2, 3, 4, 6, 8].as_slice());
    }

    #[test]
    fn reverse() {
        let mut arr = [8i, 6, 4, 3, 2, 1];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [1i, 2, 3, 4, 6, 8].as_slice());
    }

    #[test]
    fn random() {
        let mut arr = [12u, 54, 2, 93, 13, 43, 15, 299, 234];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [2u, 12, 13, 15, 43, 54, 93, 234, 299].as_slice());
    }

    #[test]
    fn one() {
        let mut arr = [9u];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), [9u].as_slice());
    }

    #[test]
    fn empty() {
        let mut arr: [int, ..0] = [];
        let empty: [int, ..0] = [];
        heap_sort(arr);
        assert_eq!(arr.as_slice(), empty.as_slice());
    }
}<|endoftext|>


<commit_before>mod page;
pub use self::page::Page;

use std::env;
use std::io::Read;
use std::time::Duration;
#[allow(unused_imports)]
use std::ascii::AsciiExt;
use std::default::Default;
use std::collections::BTreeMap;

use url::Url;
use mime;
use reqwest::Client as HttpClient;
use reqwest::header::{ContentType, UserAgent};
use reqwest::StatusCode;
use reqwest::Response as HttpResponse;
use robotparser::RobotFileParser;
use html5ever::tokenizer::Tokenizer;
use html5ever::tokenizer::BufferQueue;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use encoding::{DecoderTrap, Encoding};
use encoding::all::UTF_8;
use url_serde::Serde;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Serde<Url>>,
    pub unvisited_urls: Vec<Serde<Url>>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(
        redis_pool: RedisPool,
        base_url: Url,
        limit: isize,
        mime_types: Vec<mime::Mime>,
    ) -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts {
            namespace: Some(maman_env.to_string()),
        };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url: base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq: sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit: limit,
            mime_types: mime_types,
        }
    }

    #[allow(unknown_lints, needless_pass_by_value)]
    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser.can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for h in response.headers().iter() {
                headers.insert(h.name().to_ascii_lowercase(), h.value_string());
            }
        }
        let mut document = vec![];
        let _ = response.read_to_end(&mut document);
        match UTF_8.decode(&document, DecoderTrap::Ignore) {
            Ok(doc) => {
                let page = Page::new(
                    page_url.clone(),
                    doc.to_string(),
                    headers.clone(),
                    response.status().to_string(),
                );
                let read = Spider::read_page(page, &doc);
                Some(read.sink)
            }
            Err(_) => None,
        }
    }

    fn load_url(url: &str, mime_types: &[mime::Mime]) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .expect("HttpClient failed to construct");
        match client
            .get(url)
            .header(UserAgent::new(maman_user_agent!().to_owned()))
            .send()
        {
            Err(_) => None,
            Ok(response) => match response.status() {
                StatusCode::Ok | StatusCode::NotModified => {
                    if mime_types.is_empty()
                        || mime_types.contains(response.headers().get::<ContentType>().unwrap())
                    {
                        Some(response)
                    } else {
                        None
                    }
                }
                _ => None,
            },
        }
    }
}
<commit_msg>remove `use std::ascii::AsciiExt` for Rust 1.23<commit_after>mod page;
pub use self::page::Page;

use std::env;
use std::io::Read;
use std::time::Duration;
use std::default::Default;
use std::collections::BTreeMap;

use url::Url;
use mime;
use reqwest::Client as HttpClient;
use reqwest::header::{ContentType, UserAgent};
use reqwest::StatusCode;
use reqwest::Response as HttpResponse;
use robotparser::RobotFileParser;
use html5ever::tokenizer::Tokenizer;
use html5ever::tokenizer::BufferQueue;
use sidekiq::Client as SidekiqClient;
use sidekiq::ClientOpts as SidekiqClientOpts;
use sidekiq::RedisPool;
use encoding::{DecoderTrap, Encoding};
use encoding::all::UTF_8;
use url_serde::Serde;

const MAMAN_ENV: &str = "MAMAN_ENV";
const MAMAN_ENV_DEFAULT: &str = "development";

pub struct Spider<'a> {
    pub base_url: Url,
    pub visited_urls: Vec<Serde<Url>>,
    pub unvisited_urls: Vec<Serde<Url>>,
    pub env: String,
    pub limit: isize,
    pub mime_types: Vec<mime::Mime>,
    sidekiq: SidekiqClient,
    robot_parser: RobotFileParser<'a>,
}

impl<'a> Spider<'a> {
    pub fn new(
        redis_pool: RedisPool,
        base_url: Url,
        limit: isize,
        mime_types: Vec<mime::Mime>,
    ) -> Spider<'a> {
        let maman_env =
            env::var(&MAMAN_ENV.to_owned()).unwrap_or_else(|_| MAMAN_ENV_DEFAULT.to_owned());
        let robots_txt = base_url.join("/robots.txt").unwrap();
        let robot_file_parser = RobotFileParser::new(robots_txt);
        let client_opts = SidekiqClientOpts {
            namespace: Some(maman_env.to_string()),
        };
        let sidekiq = SidekiqClient::new(redis_pool, client_opts);
        Spider {
            base_url: base_url,
            visited_urls: Vec::new(),
            unvisited_urls: Vec::new(),
            sidekiq: sidekiq,
            env: maman_env,
            robot_parser: robot_file_parser,
            limit: limit,
            mime_types: mime_types,
        }
    }

    #[allow(unknown_lints, needless_pass_by_value)]
    pub fn visit_page(&mut self, page: Page) {
        self.visited_urls.push(page.url.clone());
        for u in &page.urls {
            self.unvisited_urls.push(u.clone());
        }
        if let Err(err) = self.sidekiq.push(page.to_job()) {
            error!("SidekiqClient push failed: {}", err);
        }
    }

    pub fn crawl(&mut self) {
        self.robot_parser.read();
        let base_url = self.base_url.clone();
        if let Some(response) = Spider::load_url(self.base_url.as_ref(), &self.mime_types) {
            self.visit(&base_url, response);
            while let Some(url) = self.unvisited_urls.pop() {
                if self.continue_to_crawl() {
                    if !self.visited_urls.contains(&url) {
                        if let Some(response) = Spider::load_url(url.as_ref(), &self.mime_types) {
                            self.visit(&url, response);
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    pub fn read_page(page: Page, document: &str) -> Tokenizer<Page> {
        let mut tok = Tokenizer::new(page, Default::default());
        let mut input = BufferQueue::new();
        input.push_back(String::from(document).into());
        let _ = tok.feed(&mut input);
        tok.end();
        tok
    }

    fn visit(&mut self, page_url: &Url, response: HttpResponse) {
        if self.can_visit(page_url) {
            info!("{}", page_url);
            if let Some(page) = Spider::read_response(page_url, response) {
                self.visit_page(page);
            }
        }
    }

    fn continue_to_crawl(&self) -> bool {
        self.limit == 0 || (self.visited_urls.len() as isize) < self.limit
    }

    fn can_visit(&self, page_url: &Url) -> bool {
        self.robot_parser.can_fetch(maman_name!(), page_url.path())
    }

    fn read_response(page_url: &Url, mut response: HttpResponse) -> Option<Page> {
        let mut headers = BTreeMap::new();
        {
            for h in response.headers().iter() {
                headers.insert(h.name().to_ascii_lowercase(), h.value_string());
            }
        }
        let mut document = vec![];
        let _ = response.read_to_end(&mut document);
        match UTF_8.decode(&document, DecoderTrap::Ignore) {
            Ok(doc) => {
                let page = Page::new(
                    page_url.clone(),
                    doc.to_string(),
                    headers.clone(),
                    response.status().to_string(),
                );
                let read = Spider::read_page(page, &doc);
                Some(read.sink)
            }
            Err(_) => None,
        }
    }

    fn load_url(url: &str, mime_types: &[mime::Mime]) -> Option<HttpResponse> {
        let client = HttpClient::builder()
            .timeout(Duration::from_secs(5))
            .build()
            .expect("HttpClient failed to construct");
        match client
            .get(url)
            .header(UserAgent::new(maman_user_agent!().to_owned()))
            .send()
        {
            Err(_) => None,
            Ok(response) => match response.status() {
                StatusCode::Ok | StatusCode::NotModified => {
                    if mime_types.is_empty()
                        || mime_types.contains(response.headers().get::<ContentType>().unwrap())
                    {
                        Some(response)
                    } else {
                        None
                    }
                }
                _ => None,
            },
        }
    }
}
<|endoftext|>


<commit_before>use crate::StrError;
use serde::{Deserialize, Serialize};
use std::ffi::OsStr;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::Path;

/// Aliases usize as Point ID
pub type PointId = usize;

/// Aliases usize as Cell ID
pub type CellId = usize;

/// Aliases usize as Cell's attribute ID
pub type CellAttributeId = usize;

/// Holds point data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Point {
    /// Identification number which equals the index of the point in the mesh
    pub id: PointId,

    /// Point coordinates (2D or 3D)
    pub coords: Vec<f64>,
}

/// Holds cell (aka geometric shape, polygon, polyhedra) data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Cell {
    /// Identification number which equals the index of the cell in the mesh
    pub id: CellId,

    /// Attribute identification number
    pub attribute_id: CellAttributeId,

    /// Space dimension of this cell
    ///
    /// The cell's ndim may be different than the space dimension of the mesh.
    /// For example, a 1D line in the 2D or 3D space or a 2D triangle in the 3D space.
    pub geo_ndim: usize,

    /// List of points defining this cell (nodes); in the right order (unsorted)
    ///
    /// Note: The list of nodes must follow a **counter-clockwise order**.
    pub points: Vec<PointId>,
}

/// Holds mesh data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Mesh {
    /// Space dimension of the mesh
    ///
    /// The mesh's ndim may be different that an cell's ndim.
    /// For example, a 3D mesh may contain 1D lines or 2D triangles.
    pub space_ndim: usize,

    /// All points in the mesh
    pub points: Vec<Point>,

    /// All cells (aka geometric shape, polygon, polyhedra) in the mesh
    pub cells: Vec<Cell>,
}

impl Mesh {
    /// Reads a binary file containing the mesh data
    ///
    /// # Input
    ///
    /// * `full_path` -- may be a String, &str, or Path
    pub fn read<P>(full_path: &P) -> Result<Self, StrError>
    where
        P: AsRef<OsStr> + ?Sized,
    {
        let path = Path::new(full_path).to_path_buf();
        let mut file = File::open(&path).map_err(|_| "file not found")?;
        let metadata = fs::metadata(&path).map_err(|_| "unable to read metadata")?;
        let mut bin = vec![0; metadata.len() as usize];
        file.read(&mut bin).expect("buffer overflow");
        let mut des = rmp_serde::Deserializer::new(&bin[..]);
        let mesh: Mesh = Deserialize::deserialize(&mut des).map_err(|_| "deserialize failed")?;
        Ok(mesh)
    }

    /// Writes a binary file with the mesh data
    ///
    /// # Input
    ///
    /// * `full_path` -- may be a String, &str, or Path
    pub fn write<P>(&self, full_path: &P) -> Result<(), StrError>
    where
        P: AsRef<OsStr> + ?Sized,
    {
        let path = Path::new(full_path).to_path_buf();
        if let Some(p) = path.parent() {
            fs::create_dir_all(p).map_err(|_| "cannot create directory")?;
        }
        let mut bin = Vec::new();
        let mut ser = rmp_serde::Serializer::new(&mut bin);
        self.serialize(&mut ser).map_err(|_| "serialize failed")?;
        let mut file = File::create(&path).map_err(|_| "cannot create file")?;
        file.write_all(&bin).map_err(|_| "cannot write file")?;
        Ok(())
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use crate::mesh::Mesh;
    use crate::StrError;

    #[test]
    fn read_and_write_capture_errors() -> Result<(), StrError> {
        assert_eq!(Mesh::read("/tmp/not_found").err(), Some("file not found"));
        assert_eq!(
            Mesh::read("./data/meshes/two_quads_horizontal.msh").err(),
            Some("deserialize failed")
        );
        let mesh = Mesh {
            space_ndim: 2,
            points: Vec::new(),
            cells: Vec::new(),
        };
        assert_eq!(mesh.write("/tmp/").err(), Some("cannot create file"));
        Ok(())
    }

    #[test]
    fn read_and_write_work() -> Result<(), StrError> {
        //
        //  3--------2--------5
        //  |        |        |
        //  |        |        |
        //  |        |        |
        //  0--------1--------4
        //
        let mesh = Mesh::from_text(
            r"# header
            # space_ndim npoint ncell
                       2      6     2

            # points
            # id   x   y
               0 0.0 0.0
               1 1.0 0.0
               2 1.0 1.0
               3 0.0 1.0
               4 2.0 0.0
               5 2.0 1.0

            # cells
            # id att geo_ndim nnode  point_ids...
               0   1        2     4  0 1 2 3
               1   0        2     4  1 4 5 2",
        )?;
        mesh.write("/tmp/gemlab/test.msh")?;
        let mesh_read = Mesh::read("/tmp/gemlab/test.msh")?;
        assert_eq!(format!("{:?}", mesh), format!("{:?}", mesh_read));
        Ok(())
    }
}
<commit_msg>Add test<commit_after>use crate::StrError;
use serde::{Deserialize, Serialize};
use std::ffi::OsStr;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::Path;

/// Aliases usize as Point ID
pub type PointId = usize;

/// Aliases usize as Cell ID
pub type CellId = usize;

/// Aliases usize as Cell's attribute ID
pub type CellAttributeId = usize;

/// Holds point data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Point {
    /// Identification number which equals the index of the point in the mesh
    pub id: PointId,

    /// Point coordinates (2D or 3D)
    pub coords: Vec<f64>,
}

/// Holds cell (aka geometric shape, polygon, polyhedra) data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Cell {
    /// Identification number which equals the index of the cell in the mesh
    pub id: CellId,

    /// Attribute identification number
    pub attribute_id: CellAttributeId,

    /// Space dimension of this cell
    ///
    /// The cell's ndim may be different than the space dimension of the mesh.
    /// For example, a 1D line in the 2D or 3D space or a 2D triangle in the 3D space.
    pub geo_ndim: usize,

    /// List of points defining this cell (nodes); in the right order (unsorted)
    ///
    /// Note: The list of nodes must follow a **counter-clockwise order**.
    pub points: Vec<PointId>,
}

/// Holds mesh data
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Mesh {
    /// Space dimension of the mesh
    ///
    /// The mesh's ndim may be different that an cell's ndim.
    /// For example, a 3D mesh may contain 1D lines or 2D triangles.
    pub space_ndim: usize,

    /// All points in the mesh
    pub points: Vec<Point>,

    /// All cells (aka geometric shape, polygon, polyhedra) in the mesh
    pub cells: Vec<Cell>,
}

impl Mesh {
    /// Reads a binary file containing the mesh data
    ///
    /// # Input
    ///
    /// * `full_path` -- may be a String, &str, or Path
    pub fn read<P>(full_path: &P) -> Result<Self, StrError>
    where
        P: AsRef<OsStr> + ?Sized,
    {
        let path = Path::new(full_path).to_path_buf();
        let mut file = File::open(&path).map_err(|_| "file not found")?;
        let metadata = fs::metadata(&path).map_err(|_| "unable to read metadata")?;
        let mut bin = vec![0; metadata.len() as usize];
        file.read(&mut bin).expect("buffer overflow");
        let mut des = rmp_serde::Deserializer::new(&bin[..]);
        let mesh: Mesh = Deserialize::deserialize(&mut des).map_err(|_| "deserialize failed")?;
        Ok(mesh)
    }

    /// Writes a binary file with the mesh data
    ///
    /// # Input
    ///
    /// * `full_path` -- may be a String, &str, or Path
    pub fn write<P>(&self, full_path: &P) -> Result<(), StrError>
    where
        P: AsRef<OsStr> + ?Sized,
    {
        let path = Path::new(full_path).to_path_buf();
        if let Some(p) = path.parent() {
            fs::create_dir_all(p).map_err(|_| "cannot create directory")?;
        }
        let mut bin = Vec::new();
        let mut ser = rmp_serde::Serializer::new(&mut bin);
        self.serialize(&mut ser).map_err(|_| "serialize failed")?;
        let mut file = File::create(&path).map_err(|_| "cannot create file")?;
        file.write_all(&bin).map_err(|_| "cannot write file")?;
        Ok(())
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use crate::mesh::{Mesh, Samples};
    use crate::StrError;

    #[test]
    fn read_and_write_capture_errors() -> Result<(), StrError> {
        assert_eq!(Mesh::read("/tmp/not_found").err(), Some("file not found"));
        assert_eq!(
            Mesh::read("./data/meshes/two_quads_horizontal.msh").err(),
            Some("deserialize failed")
        );
        let mesh = Mesh {
            space_ndim: 2,
            points: Vec::new(),
            cells: Vec::new(),
        };
        assert_eq!(mesh.write("/tmp/").err(), Some("cannot create file"));
        Ok(())
    }

    #[test]
    fn read_and_write_work() -> Result<(), StrError> {
        //
        //  3--------2--------5
        //  |        |        |
        //  |        |        |
        //  |        |        |
        //  0--------1--------4
        //
        let mesh = Mesh::from_text(
            r"# header
            # space_ndim npoint ncell
                       2      6     2

            # points
            # id   x   y
               0 0.0 0.0
               1 1.0 0.0
               2 1.0 1.0
               3 0.0 1.0
               4 2.0 0.0
               5 2.0 1.0

            # cells
            # id att geo_ndim nnode  point_ids...
               0   1        2     4  0 1 2 3
               1   0        2     4  1 4 5 2",
        )?;
        mesh.write("/tmp/gemlab/test.msh")?;
        let mesh_read = Mesh::read("/tmp/gemlab/test.msh")?;
        assert_eq!(format!("{:?}", mesh), format!("{:?}", mesh_read));
        Ok(())
    }

    #[test]
    fn derive_works() -> Result<(), StrError> {
        let mesh = Samples::two_quads_horizontal();
        let mesh_clone = mesh.clone();
        assert_eq!(format!("{:?}", mesh), "Mesh { space_ndim: 2, points: [Point { id: 0, coords: [0.0, 0.0] }, Point { id: 1, coords: [1.0, 0.0] }, Point { id: 2, coords: [1.0, 1.0] }, Point { id: 3, coords: [0.0, 1.0] }, Point { id: 4, coords: [2.0, 0.0] }, Point { id: 5, coords: [2.0, 1.0] }], cells: [Cell { id: 0, attribute_id: 1, geo_ndim: 2, points: [0, 1, 2, 3] }, Cell { id: 1, attribute_id: 2, geo_ndim: 2, points: [1, 4, 5, 2] }] }");
        assert_eq!(mesh_clone.space_ndim, mesh.space_ndim);
        assert_eq!(mesh_clone.points.len(), mesh.points.len());
        assert_eq!(mesh_clone.cells.len(), mesh.cells.len());
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Manual `GameState` manipulation, to be used when implementing a UI.
//!
//! For a UI implementing the use of this see the implementation of `Widgets::update()`.

use self::super::{Difficulty, GameState};
use self::super::super::util::FRUITS;
use rand::{Rng, thread_rng};

/// Press the Start button in MainMenu, transforming it into ChooseDifficulty.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let start_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if start_button_pressed {
///     state::press_start(&mut state);
/// }
/// ```
pub fn press_start(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::ChooseDifficulty,
        _ => {}
    }
}

/// Press the Display Highscores button in MainMenu, transforming it into LoadLeaderboard.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let display_highscores_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if display_highscores_button_pressed {
///     state::press_display_highscores(&mut state);
///     // Then probably load the leaderboard
/// }
/// ```
pub fn press_display_highscores(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::LoadLeaderboard,
        _ => {}
    }
}

/// Press the Exit button in MainMenu, transforming it into Exit.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let exit_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if exit_button_pressed {
///     state::press_exit(&mut state);
///     // Then probably close the window
/// }
/// ```
pub fn press_exit(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::Exit,
        _ => {}
    }
}

/// Press one of the Difficulty buttons in ChooseDifficulty, transforming it into Exit.
///
/// If the supplied argument is not ChooseDifficulty, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # let difficulty_easy_button_pressed = false;
/// # let difficulty_normal_button_pressed = true;
/// # let difficulty_hard_button_pressed = false;
/// let mut state = GameState::ChooseDifficulty;
/// if difficulty_easy_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Easy);
/// } else if difficulty_normal_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Normal);
/// } else if difficulty_hard_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Hard);
/// }
/// ```
pub fn select_difficulty(s: &mut GameState, difficulty: Difficulty) {
    match s {
        &mut GameState::ChooseDifficulty => {
            *s = GameState::Playing {
                difficulty: difficulty,
                score: 0,
                fruit: None,
            }
        }
        _ => {}
    }
}

/// Press the Back button in ChooseDifficulty, GameOver or DisplayLeaderboard, transforming them into MainMenu.
///
/// If the supplied argument is not ChooseDifficulty, GameOver or DisplayLeaderboard, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let back_button_pressed = false;
/// let mut state = GameState::ChooseDifficulty;
/// if back_button_pressed {
///     state::press_back(&mut state);
/// }
/// ```
pub fn press_back(s: &mut GameState) {
    match s {
        &mut GameState::ChooseDifficulty |
        &mut GameState::GameOver { difficulty: _, score: _, name: _ } |
        &mut GameState::DisplayLeaderboard(_) => *s = GameState::MainMenu,
        _ => {}
    }
}

/// Tick the mango in Playing state.
///
/// Call this before displaying/updating the label of the mango button. Returns the original fruit.
///
/// If the supplied argument is not Playing, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # use poke_a_mango::util::fruit_name;
/// # let mango_button = ();
/// # fn update_label(_: &(), _: &'static str) {}
/// let mut state = GameState::Playing {
///     difficulty: Difficulty::Hard,
///     score: 12,
///     fruit: None,
/// };
/// let fruit = state::tick_mango(&mut state);
/// update_label(&mango_button, fruit_name(&fruit));
/// ```
pub fn tick_mango(s: &mut GameState) -> Option<usize> {
    match s {
        &mut GameState::Playing { difficulty, score: _, ref mut fruit } => {
            let original_fruit = *fruit;

            // Difficulty's numeric value is inverted here
            let mut rng = thread_rng();
            let change_fruit = rng.gen_weighted_bool(25 * (4 - difficulty.numeric()));
            if change_fruit {
                *fruit = if fruit.is_none() {
                    Some(rng.gen_range(0, FRUITS.len()))
                } else {
                    None
                };
            }

            original_fruit
        }
        _ => None,
    }
}

/// End the mango processing sequence, given whether it was clicked and whether it was actually a mango.
///
/// If the supplied argument is not Playing or the button wasn't clicked, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # use poke_a_mango::util::fruit_name;
/// # let was_mango = true;
/// # let mango_button_clicked = true;
/// let mut state = GameState::Playing {
///     difficulty: Difficulty::Hard,
///     score: 12,
///     fruit: None,
/// };
/// state::end_mango(&mut state, mango_button_clicked, was_mango);
/// ```
pub fn end_mango(s: &mut GameState, clicked: bool, was_mango: bool) {
    if clicked {
        if was_mango {
            match s {
                &mut GameState::Playing { difficulty: _, ref mut score, fruit: _ } => *score += 1,
                _ => {}
            }
        } else {
            match s {
                &mut GameState::Playing { difficulty, score, fruit: _ } => {
                    *s = GameState::GameOver {
                        difficulty: difficulty,
                        score: score,
                        name: "Your name".to_string(),
                    };
                }
                _ => {}
            }
        }
    }
}

pub fn submit_name(s: &mut GameState) {
    match s {
        &mut GameState::GameOver { difficulty, score, name: _ } => {
            let name = match s {
                &mut GameState::GameOver { difficulty: _, score: _, ref name } => name.clone(),
                _ => "".to_string(),
            };
            *s = GameState::GameEnded {
                name: name,
                score: ((score as f64) * difficulty.point_weight()).floor() as u64,
            };
        }
        _ => {}
    }
}
<commit_msg>Fix doc to better reflect usage<commit_after>//! Manual `GameState` manipulation, to be used when implementing a UI.
//!
//! For a UI implementing the use of this see the implementation of `Widgets::update()`.

use self::super::{Difficulty, GameState};
use self::super::super::util::FRUITS;
use rand::{Rng, thread_rng};

/// Press the Start button in MainMenu, transforming it into ChooseDifficulty.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let start_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if start_button_pressed {
///     state::press_start(&mut state);
/// }
/// ```
pub fn press_start(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::ChooseDifficulty,
        _ => {}
    }
}

/// Press the Display Highscores button in MainMenu, transforming it into LoadLeaderboard.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let display_highscores_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if display_highscores_button_pressed {
///     state::press_display_highscores(&mut state);
///     // Then probably load the leaderboard
/// }
/// ```
pub fn press_display_highscores(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::LoadLeaderboard,
        _ => {}
    }
}

/// Press the Exit button in MainMenu, transforming it into Exit.
///
/// If the supplied argument is not MainMenu, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let exit_button_pressed = true;
/// let mut state = GameState::MainMenu;
/// if exit_button_pressed {
///     state::press_exit(&mut state);
///     // Then probably close the window
/// }
/// ```
pub fn press_exit(s: &mut GameState) {
    match s {
        &mut GameState::MainMenu => *s = GameState::Exit,
        _ => {}
    }
}

/// Press one of the Difficulty buttons in ChooseDifficulty, transforming it into Exit.
///
/// If the supplied argument is not ChooseDifficulty, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # let difficulty_easy_button_pressed = false;
/// # let difficulty_normal_button_pressed = true;
/// # let difficulty_hard_button_pressed = false;
/// let mut state = GameState::ChooseDifficulty;
/// if difficulty_easy_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Easy);
/// } else if difficulty_normal_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Normal);
/// } else if difficulty_hard_button_pressed {
///     state::select_difficulty(&mut state, Difficulty::Hard);
/// }
/// ```
pub fn select_difficulty(s: &mut GameState, difficulty: Difficulty) {
    match s {
        &mut GameState::ChooseDifficulty => {
            *s = GameState::Playing {
                difficulty: difficulty,
                score: 0,
                fruit: None,
            }
        }
        _ => {}
    }
}

/// Press the Back button in ChooseDifficulty, GameOver or DisplayLeaderboard, transforming them into MainMenu.
///
/// If the supplied argument is not ChooseDifficulty, GameOver or DisplayLeaderboard, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, GameState};
/// # let back_button_pressed = false;
/// let mut state = GameState::ChooseDifficulty;
/// if back_button_pressed {
///     state::press_back(&mut state);
/// }
/// ```
pub fn press_back(s: &mut GameState) {
    match s {
        &mut GameState::ChooseDifficulty |
        &mut GameState::GameOver { difficulty: _, score: _, name: _ } |
        &mut GameState::DisplayLeaderboard(_) => *s = GameState::MainMenu,
        _ => {}
    }
}

/// Tick the mango in Playing state.
///
/// Call this before displaying/updating the label of the mango button. Returns the original fruit.
///
/// If the supplied argument is not Playing, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # use poke_a_mango::util::fruit_name;
/// # let mango_button = ();
/// # fn update_label(_: &(), _: &'static str) {}
/// let mut state = GameState::Playing {
///     difficulty: Difficulty::Hard,
///     score: 12,
///     fruit: None,
/// };
/// let fruit = state::tick_mango(&mut state);
/// update_label(&mango_button, fruit_name(&fruit));
/// ```
pub fn tick_mango(s: &mut GameState) -> Option<usize> {
    match s {
        &mut GameState::Playing { difficulty, score: _, ref mut fruit } => {
            let original_fruit = *fruit;

            // Difficulty's numeric value is inverted here
            let mut rng = thread_rng();
            let change_fruit = rng.gen_weighted_bool(25 * (4 - difficulty.numeric()));
            if change_fruit {
                *fruit = if fruit.is_none() {
                    Some(rng.gen_range(0, FRUITS.len()))
                } else {
                    None
                };
            }

            original_fruit
        }
        _ => None,
    }
}

/// End the mango processing sequence, given whether it was clicked and whether it was actually a mango.
///
/// If the supplied argument is not Playing or the button wasn't clicked, it remains unchanged.
///
/// # Examples
///
/// ```
/// # use poke_a_mango::ops::{state, Difficulty, GameState};
/// # use poke_a_mango::util::fruit_name;
/// # let previous_fruit = Some(0);
/// # let mango_button_clicked = true;
/// let mut state = GameState::Playing {
///     difficulty: Difficulty::Hard,
///     score: 12,
///     fruit: None,
/// };
/// state::end_mango(&mut state, mango_button_clicked, previous_fruit.is_none());
/// ```
pub fn end_mango(s: &mut GameState, clicked: bool, was_mango: bool) {
    if clicked {
        if was_mango {
            match s {
                &mut GameState::Playing { difficulty: _, ref mut score, fruit: _ } => *score += 1,
                _ => {}
            }
        } else {
            match s {
                &mut GameState::Playing { difficulty, score, fruit: _ } => {
                    *s = GameState::GameOver {
                        difficulty: difficulty,
                        score: score,
                        name: "Your name".to_string(),
                    };
                }
                _ => {}
            }
        }
    }
}

pub fn submit_name(s: &mut GameState) {
    match s {
        &mut GameState::GameOver { difficulty, score, name: _ } => {
            let name = match s {
                &mut GameState::GameOver { difficulty: _, score: _, ref name } => name.clone(),
                _ => "".to_string(),
            };
            *s = GameState::GameEnded {
                name: name,
                score: ((score as f64) * difficulty.point_weight()).floor() as u64,
            };
        }
        _ => {}
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
use alloc::buddy;

pub use arch::memory::{PAddr, HEAP_BASE, HEAP_TOP};
pub mod vaddr_impls;
pub use self::vaddr_impls::*;

/// A virtual address is a machine-sized unsigned integer
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct VAddr(usize);

impl VAddr {
    #[inline] pub fn from_ptr(ptr: *mut u8) -> Self {
        VAddr(ptr as usize)
    }
    #[inline] pub const fn from_usize(u: usize) -> Self {
        VAddr(u)
    }
    #[inline] pub const fn as_usize(&self) -> usize {
        self.0
    }

    /// Calculate the index in the PML4 table corresponding to this address.
    #[inline] pub fn pml4_index(&self) -> usize {
        (self >> 39) & 0b111111111
    }

    /// Calculate the index in the PDPT table corresponding to this address.
    #[inline] pub fn pdpt_index(&self) -> usize {
        (self >> 30) & 0b111111111
    }

    /// Calculate the index in the PD table corresponding to this address.
    #[inline] pub fn pd_index(&self) -> usize {
        (self >> 21) & 0b111111111
    }

    /// Calculate the index in the PT table corresponding to this address.
    #[inline] pub fn pt_index(&self) -> usize {
        (self >> 12) & 0b111111111
    }
}

#[inline] pub fn heap_base_addr() -> usize {
    unsafe { (&mut HEAP_BASE as *mut _) as usize }
}

#[inline] pub fn heap_top_addr() -> usize {
    unsafe { (&mut HEAP_TOP as *mut _) as usize }
}

pub unsafe fn init_heap<'a>() -> Result<&'a str, &'a str> {
    let heap_base_ptr
        = &mut HEAP_BASE as *mut _;
    let heap_size
        = (&mut HEAP_TOP as *mut _) as usize - heap_base_ptr as usize;
    buddy::system::init_heap( heap_base_ptr, heap_size);
    Ok("[ OKAY ]")
}
<commit_msg>[kernel] Changes to VAddr<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
use alloc::buddy;

use core::ops;
use core::fmt;
use core::convert;

pub use arch::memory::{PAddr, HEAP_BASE, HEAP_TOP};
pub mod vaddr_impls;
pub use self::vaddr_impls::*;

/// A virtual address is a machine-sized unsigned integer
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct VAddr(usize);

pub trait Addr<R>: ops::Add<Self> + ops::Add<R>
                 + ops::Sub<Self> + ops::Sub<R>
                 + ops::Mul<Self> + ops::Mul<R>
                 + ops::Div<Self> + ops::Div<R>
                 + ops::Shl<Self> + ops::Shl<R>
                 + ops::Shr<Self> + ops::Shr<R>
                 + convert::From<R> + convert::Into<R>
                 + convert::From<*mut u8>
                 + Sized { }

impl Addr<usize> for VAddr { }

impl convert::Into<usize> for VAddr {
    #[inline] fn into(self) -> usize { self.as_usize() }
}

impl convert::From<usize> for VAddr {
    #[inline] fn from(u: usize) -> Self { VAddr::from_usize(u) }
}

impl convert::From<*mut u8> for VAddr {
    #[inline] fn from(ptr: *mut u8) -> Self { VAddr::from_ptr(ptr) }
}

impl VAddr {
    #[inline] pub fn from_ptr(ptr: *mut u8) -> Self {
        VAddr(ptr as usize)
    }
    #[inline] pub const fn from_usize(u: usize) -> Self {
        VAddr(u)
    }
    #[inline] pub const fn as_usize(&self) -> usize {
        self.0
    }

    /// Calculate the index in the PML4 table corresponding to this address.
    #[inline] pub fn pml4_index(&self) -> usize {
        (self >> 39) & 0b111111111
    }

    /// Calculate the index in the PDPT table corresponding to this address.
    #[inline] pub fn pdpt_index(&self) -> usize {
        (self >> 30) & 0b111111111
    }

    /// Calculate the index in the PD table corresponding to this address.
    #[inline] pub fn pd_index(&self) -> usize {
        (self >> 21) & 0b111111111
    }

    /// Calculate the index in the PT table corresponding to this address.
    #[inline] pub fn pt_index(&self) -> usize {
        (self >> 12) & 0b111111111
    }
}

#[inline] pub fn heap_base_addr() -> usize {
    unsafe { (&mut HEAP_BASE as *mut _) as usize }
}

#[inline] pub fn heap_top_addr() -> usize {
    unsafe { (&mut HEAP_TOP as *mut _) as usize }
}

pub unsafe fn init_heap<'a>() -> Result<&'a str, &'a str> {
    let heap_base_ptr
        = &mut HEAP_BASE as *mut _;
    let heap_size
        = (&mut HEAP_TOP as *mut _) as usize - heap_base_ptr as usize;
    buddy::system::init_heap( heap_base_ptr, heap_size);
    Ok("[ OKAY ]")
}

macro_rules! make_addr_range {
    $($name:ident, $addr:ty),+ => {$(

    )+}
}
<|endoftext|>


<commit_before>//! Implements a map from integer indices to data.
//! Rather than storing data for every index, internally, this maps entire ranges to the data.
//! To this end, the APIs all work on ranges, not on individual integers. Ranges are split as
//! necessary (e.g., when [0,5) is first associated with X, and then [1,2) is mutated).
//! Users must not depend on whether a range is coalesced or not, even though this is observable
//! via the iteration APIs.

use std::ops;

use rustc_target::abi::Size;

#[derive(Clone, Debug)]
struct Elem<T> {
    /// The range covered by this element; never empty.
    range: ops::Range<u64>,
    /// The data stored for this element.
    data: T,
}
#[derive(Clone, Debug)]
pub struct RangeMap<T> {
    v: Vec<Elem<T>>,
}

impl<T> RangeMap<T> {
    /// Creates a new `RangeMap` for the given size, and with the given initial value used for
    /// the entire range.
    #[inline(always)]
    pub fn new(size: Size, init: T) -> RangeMap<T> {
        let size = size.bytes();
        let mut map = RangeMap { v: Vec::new() };
        if size > 0 {
            map.v.push(Elem { range: 0..size, data: init });
        }
        map
    }

    /// Finds the index containing the given offset.
    fn find_offset(&self, offset: u64) -> usize {
        // We do a binary search.
        let mut left = 0usize; // inclusive
        let mut right = self.v.len(); // exclusive
        loop {
            debug_assert!(left < right, "find_offset: offset {} is out-of-bounds", offset);
            let candidate = left.checked_add(right).unwrap() / 2;
            let elem = &self.v[candidate];
            if offset < elem.range.start {
                // We are too far right (offset is further left).
                debug_assert!(candidate < right); // we are making progress
                right = candidate;
            } else if offset >= elem.range.end {
                // We are too far left (offset is further right).
                debug_assert!(candidate >= left); // we are making progress
                left = candidate + 1;
            } else {
                // This is it!
                return candidate;
            }
        }
    }

    /// Provides read-only iteration over everything in the given range. This does
    /// *not* split items if they overlap with the edges. Do not use this to mutate
    /// through interior mutability.
    ///
    /// The iterator also provides the offset of the given element.
    pub fn iter(&self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &T)> {
        let offset = offset.bytes();
        let len = len.bytes();
        // Compute a slice starting with the elements we care about.
        let slice: &[Elem<T>] = if len == 0 {
            // We just need any empty iterator. We don't even want to
            // yield the element that surrounds this position.
            &[]
        } else {
            let first_idx = self.find_offset(offset);
            &self.v[first_idx..]
        };
        // The first offset that is not included any more.
        let end = offset + len;
        slice
            .iter()
            .take_while(move |elem| elem.range.start < end)
            .map(|elem| (Size::from_bytes(elem.range.start), &elem.data))
    }

    pub fn iter_mut_all(&mut self) -> impl Iterator<Item = &mut T> {
        self.v.iter_mut().map(|elem| &mut elem.data)
    }

    // Splits the element situated at the given `index`, such that the 2nd one starts at offset
    // `split_offset`. Do nothing if the element already starts there.
    // Returns whether a split was necessary.
    fn split_index(&mut self, index: usize, split_offset: u64) -> bool
    where
        T: Clone,
    {
        let elem = &mut self.v[index];
        if split_offset == elem.range.start || split_offset == elem.range.end {
            // Nothing to do.
            return false;
        }
        debug_assert!(
            elem.range.contains(&split_offset),
            "the `split_offset` is not in the element to be split"
        );

        // Now we really have to split. Reduce length of first element.
        let second_range = split_offset..elem.range.end;
        elem.range.end = split_offset;
        // Copy the data, and insert second element.
        let second = Elem { range: second_range, data: elem.data.clone() };
        self.v.insert(index + 1, second);
        true
    }

    /// Provides mutable iteration over everything in the given range. As a side-effect,
    /// this will split entries in the map that are only partially hit by the given range,
    /// to make sure that when they are mutated, the effect is constrained to the given range.
    /// Moreover, this will opportunistically merge neighbouring equal blocks.
    ///
    /// The iterator also provides the offset of the given element.
    pub fn iter_mut(&mut self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &mut T)>
    where
        T: Clone + PartialEq,
    {
        let offset = offset.bytes();
        let len = len.bytes();
        // Compute a slice containing exactly the elements we care about
        let slice: &mut [Elem<T>] = if len == 0 {
            // We just need any empty iterator. We don't even want to
            // yield the element that surrounds this position, nor do
            // any splitting.
            &mut []
        } else {
            // Make sure we got a clear beginning
            let mut first_idx = self.find_offset(offset);
            if self.split_index(first_idx, offset) {
                // The newly created 2nd element is ours
                first_idx += 1;
            }
            // No more mutation.
            let first_idx = first_idx;
            // Find our end. Linear scan, but that's ok because the iteration
            // is doing the same linear scan anyway -- no increase in complexity.
            // We combine this scan with a scan for duplicates that we can merge, to reduce
            // the number of elements.
            // We stop searching after the first "block" of size 1, to avoid spending excessive
            // amounts of time on the merging.
            let mut equal_since_idx = first_idx;
            // Once we see too many non-mergeable blocks, we stop.
            // The initial value is chosen via... magic. Benchmarking and magic.
            let mut successful_merge_count = 3usize;
            // When the loop is done, this is the first excluded element.
            let mut end_idx = first_idx;
            loop {
                // Compute if `end` is the last element we need to look at.
                let done = self.v[end_idx].range.end >= offset + len;
                // We definitely need to include `end`, so move the index.
                end_idx += 1;
                debug_assert!(
                    done || end_idx < self.v.len(),
                    "iter_mut: end-offset {} is out-of-bounds",
                    offset + len
                );
                // see if we want to merge everything in `equal_since..end` (exclusive at the end!)
                if successful_merge_count > 0 {
                    if done || self.v[end_idx].data != self.v[equal_since_idx].data {
                        // Everything in `equal_since..end` was equal. Make them just one element covering
                        // the entire range.
                        let removed_elems = end_idx - equal_since_idx - 1; // number of elements that we would remove
                        if removed_elems > 0 {
                            // Adjust the range of the first element to cover all of them.
                            let equal_until = self.v[end_idx - 1].range.end; // end of range of last of the equal elements
                            self.v[equal_since_idx].range.end = equal_until;
                            // Delete the rest of them.
                            self.v.splice(equal_since_idx + 1..end_idx, std::iter::empty());
                            // Adjust `end_idx` because we made the list shorter.
                            end_idx -= removed_elems;
                            // Adjust the count for the cutoff.
                            successful_merge_count += removed_elems;
                        } else {
                            // Adjust the count for the cutoff.
                            successful_merge_count -= 1;
                        }
                        // Go on scanning for the next block starting here.
                        equal_since_idx = end_idx;
                    }
                }
                // Leave loop if this is the last element.
                if done {
                    break;
                }
            }
            // Move to last included instead of first excluded index.
            let end_idx = end_idx - 1;
            // We need to split the end as well. Even if this performs a
            // split, we don't have to adjust our index as we only care about
            // the first part of the split.
            self.split_index(end_idx, offset + len);
            // Now we yield the slice. `end` is inclusive.
            &mut self.v[first_idx..=end_idx]
        };
        slice.iter_mut().map(|elem| (Size::from_bytes(elem.range.start), &mut elem.data))
    }
}

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

    /// Query the map at every offset in the range and collect the results.
    fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64) -> Vec<T> {
        (offset..offset + len)
            .into_iter()
            .map(|i| {
                map.iter(Size::from_bytes(i), Size::from_bytes(1)).next().map(|(_, &t)| t).unwrap()
            })
            .collect()
    }

    #[test]
    fn basic_insert() {
        let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        // Insert.
        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
            *x = 42;
        }
        // Check.
        assert_eq!(to_vec(&map, 10, 1), vec![42]);
        assert_eq!(map.v.len(), 3);

        // Insert with size 0.
        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
            *x = 19;
        }
        for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
            *x = 19;
        }
        assert_eq!(to_vec(&map, 10, 2), vec![42, -1]);
        assert_eq!(map.v.len(), 3);
    }

    #[test]
    fn gaps() {
        let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
            *x = 42;
        }
        for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
            *x = 43;
        }
        assert_eq!(map.v.len(), 5);
        assert_eq!(to_vec(&map, 10, 10), vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]);

        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
            if *x < 42 {
                *x = 23;
            }
        }
        assert_eq!(map.v.len(), 6);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]);
        assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 43, 23, 23]);

        for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {
            *x = 19;
        }
        assert_eq!(map.v.len(), 6);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
        // Should be seeing two blocks with 19.
        assert_eq!(
            map.iter(Size::from_bytes(15), Size::from_bytes(2))
                .map(|(_, &t)| t)
                .collect::<Vec<_>>(),
            vec![19, 19]
        );

        // A NOP `iter_mut` should trigger merging.
        for _ in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {}
        assert_eq!(map.v.len(), 5);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
    }
}
<commit_msg>ensure that RangeMap panics on OOB<commit_after>//! Implements a map from integer indices to data.
//! Rather than storing data for every index, internally, this maps entire ranges to the data.
//! To this end, the APIs all work on ranges, not on individual integers. Ranges are split as
//! necessary (e.g., when [0,5) is first associated with X, and then [1,2) is mutated).
//! Users must not depend on whether a range is coalesced or not, even though this is observable
//! via the iteration APIs.

use std::ops;

use rustc_target::abi::Size;

#[derive(Clone, Debug)]
struct Elem<T> {
    /// The range covered by this element; never empty.
    range: ops::Range<u64>,
    /// The data stored for this element.
    data: T,
}
#[derive(Clone, Debug)]
pub struct RangeMap<T> {
    v: Vec<Elem<T>>,
}

impl<T> RangeMap<T> {
    /// Creates a new `RangeMap` for the given size, and with the given initial value used for
    /// the entire range.
    #[inline(always)]
    pub fn new(size: Size, init: T) -> RangeMap<T> {
        let size = size.bytes();
        let mut map = RangeMap { v: Vec::new() };
        if size > 0 {
            map.v.push(Elem { range: 0..size, data: init });
        }
        map
    }

    /// Finds the index containing the given offset.
    fn find_offset(&self, offset: u64) -> usize {
        // We do a binary search.
        let mut left = 0usize; // inclusive
        let mut right = self.v.len(); // exclusive
        loop {
            debug_assert!(left < right, "find_offset: offset {} is out-of-bounds", offset);
            let candidate = left.checked_add(right).unwrap() / 2;
            let elem = &self.v[candidate];
            if offset < elem.range.start {
                // We are too far right (offset is further left).
                debug_assert!(candidate < right); // we are making progress
                right = candidate;
            } else if offset >= elem.range.end {
                // We are too far left (offset is further right).
                debug_assert!(candidate >= left); // we are making progress
                left = candidate + 1;
            } else {
                // This is it!
                return candidate;
            }
        }
    }

    /// Provides read-only iteration over everything in the given range. This does
    /// *not* split items if they overlap with the edges. Do not use this to mutate
    /// through interior mutability.
    ///
    /// The iterator also provides the offset of the given element.
    pub fn iter(&self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &T)> {
        let offset = offset.bytes();
        let len = len.bytes();
        // Compute a slice starting with the elements we care about.
        let slice: &[Elem<T>] = if len == 0 {
            // We just need any empty iterator. We don't even want to
            // yield the element that surrounds this position.
            &[]
        } else {
            let first_idx = self.find_offset(offset);
            &self.v[first_idx..]
        };
        // The first offset that is not included any more.
        let end = offset + len;
        assert!(
            end <= self.v.last().unwrap().range.end,
            "iterating beyond the bounds of this RangeMap"
        );
        slice
            .iter()
            .take_while(move |elem| elem.range.start < end)
            .map(|elem| (Size::from_bytes(elem.range.start), &elem.data))
    }

    pub fn iter_mut_all(&mut self) -> impl Iterator<Item = &mut T> {
        self.v.iter_mut().map(|elem| &mut elem.data)
    }

    // Splits the element situated at the given `index`, such that the 2nd one starts at offset
    // `split_offset`. Do nothing if the element already starts there.
    // Returns whether a split was necessary.
    fn split_index(&mut self, index: usize, split_offset: u64) -> bool
    where
        T: Clone,
    {
        let elem = &mut self.v[index];
        if split_offset == elem.range.start || split_offset == elem.range.end {
            // Nothing to do.
            return false;
        }
        debug_assert!(
            elem.range.contains(&split_offset),
            "the `split_offset` is not in the element to be split"
        );

        // Now we really have to split. Reduce length of first element.
        let second_range = split_offset..elem.range.end;
        elem.range.end = split_offset;
        // Copy the data, and insert second element.
        let second = Elem { range: second_range, data: elem.data.clone() };
        self.v.insert(index + 1, second);
        true
    }

    /// Provides mutable iteration over everything in the given range. As a side-effect,
    /// this will split entries in the map that are only partially hit by the given range,
    /// to make sure that when they are mutated, the effect is constrained to the given range.
    /// Moreover, this will opportunistically merge neighbouring equal blocks.
    ///
    /// The iterator also provides the offset of the given element.
    pub fn iter_mut(&mut self, offset: Size, len: Size) -> impl Iterator<Item = (Size, &mut T)>
    where
        T: Clone + PartialEq,
    {
        let offset = offset.bytes();
        let len = len.bytes();
        // Compute a slice containing exactly the elements we care about
        let slice: &mut [Elem<T>] = if len == 0 {
            // We just need any empty iterator. We don't even want to
            // yield the element that surrounds this position, nor do
            // any splitting.
            &mut []
        } else {
            // Make sure we got a clear beginning
            let mut first_idx = self.find_offset(offset);
            if self.split_index(first_idx, offset) {
                // The newly created 2nd element is ours
                first_idx += 1;
            }
            // No more mutation.
            let first_idx = first_idx;
            // Find our end. Linear scan, but that's ok because the iteration
            // is doing the same linear scan anyway -- no increase in complexity.
            // We combine this scan with a scan for duplicates that we can merge, to reduce
            // the number of elements.
            // We stop searching after the first "block" of size 1, to avoid spending excessive
            // amounts of time on the merging.
            let mut equal_since_idx = first_idx;
            // Once we see too many non-mergeable blocks, we stop.
            // The initial value is chosen via... magic. Benchmarking and magic.
            let mut successful_merge_count = 3usize;
            // When the loop is done, this is the first excluded element.
            let mut end_idx = first_idx;
            loop {
                // Compute if `end` is the last element we need to look at.
                let done = self.v[end_idx].range.end >= offset + len;
                // We definitely need to include `end`, so move the index.
                end_idx += 1;
                debug_assert!(
                    done || end_idx < self.v.len(),
                    "iter_mut: end-offset {} is out-of-bounds",
                    offset + len
                );
                // see if we want to merge everything in `equal_since..end` (exclusive at the end!)
                if successful_merge_count > 0 {
                    if done || self.v[end_idx].data != self.v[equal_since_idx].data {
                        // Everything in `equal_since..end` was equal. Make them just one element covering
                        // the entire range.
                        let removed_elems = end_idx - equal_since_idx - 1; // number of elements that we would remove
                        if removed_elems > 0 {
                            // Adjust the range of the first element to cover all of them.
                            let equal_until = self.v[end_idx - 1].range.end; // end of range of last of the equal elements
                            self.v[equal_since_idx].range.end = equal_until;
                            // Delete the rest of them.
                            self.v.splice(equal_since_idx + 1..end_idx, std::iter::empty());
                            // Adjust `end_idx` because we made the list shorter.
                            end_idx -= removed_elems;
                            // Adjust the count for the cutoff.
                            successful_merge_count += removed_elems;
                        } else {
                            // Adjust the count for the cutoff.
                            successful_merge_count -= 1;
                        }
                        // Go on scanning for the next block starting here.
                        equal_since_idx = end_idx;
                    }
                }
                // Leave loop if this is the last element.
                if done {
                    break;
                }
            }
            // Move to last included instead of first excluded index.
            let end_idx = end_idx - 1;
            // We need to split the end as well. Even if this performs a
            // split, we don't have to adjust our index as we only care about
            // the first part of the split.
            self.split_index(end_idx, offset + len);
            // Now we yield the slice. `end` is inclusive.
            &mut self.v[first_idx..=end_idx]
        };
        slice.iter_mut().map(|elem| (Size::from_bytes(elem.range.start), &mut elem.data))
    }
}

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

    /// Query the map at every offset in the range and collect the results.
    fn to_vec<T: Copy>(map: &RangeMap<T>, offset: u64, len: u64) -> Vec<T> {
        (offset..offset + len)
            .into_iter()
            .map(|i| {
                map.iter(Size::from_bytes(i), Size::from_bytes(1)).next().map(|(_, &t)| t).unwrap()
            })
            .collect()
    }

    #[test]
    fn basic_insert() {
        let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        // Insert.
        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(1)) {
            *x = 42;
        }
        // Check.
        assert_eq!(to_vec(&map, 10, 1), vec![42]);
        assert_eq!(map.v.len(), 3);

        // Insert with size 0.
        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(0)) {
            *x = 19;
        }
        for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(0)) {
            *x = 19;
        }
        assert_eq!(to_vec(&map, 10, 2), vec![42, -1]);
        assert_eq!(map.v.len(), 3);
    }

    #[test]
    fn gaps() {
        let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        for (_, x) in map.iter_mut(Size::from_bytes(11), Size::from_bytes(1)) {
            *x = 42;
        }
        for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(1)) {
            *x = 43;
        }
        assert_eq!(map.v.len(), 5);
        assert_eq!(to_vec(&map, 10, 10), vec![-1, 42, -1, -1, -1, 43, -1, -1, -1, -1]);

        for (_, x) in map.iter_mut(Size::from_bytes(10), Size::from_bytes(10)) {
            if *x < 42 {
                *x = 23;
            }
        }
        assert_eq!(map.v.len(), 6);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 43, 23, 23, 23, 23]);
        assert_eq!(to_vec(&map, 13, 5), vec![23, 23, 43, 23, 23]);

        for (_, x) in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {
            *x = 19;
        }
        assert_eq!(map.v.len(), 6);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
        // Should be seeing two blocks with 19.
        assert_eq!(
            map.iter(Size::from_bytes(15), Size::from_bytes(2))
                .map(|(_, &t)| t)
                .collect::<Vec<_>>(),
            vec![19, 19]
        );

        // A NOP `iter_mut` should trigger merging.
        for _ in map.iter_mut(Size::from_bytes(15), Size::from_bytes(5)) {}
        assert_eq!(map.v.len(), 5);
        assert_eq!(to_vec(&map, 10, 10), vec![23, 42, 23, 23, 23, 19, 19, 19, 19, 19]);
    }

    #[test]
    #[should_panic]
    fn out_of_range_iter_mut() {
        let mut map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        let _ = map.iter_mut(Size::from_bytes(11), Size::from_bytes(11));
    }

    #[test]
    #[should_panic]
    fn out_of_range_iter() {
        let map = RangeMap::<i32>::new(Size::from_bytes(20), -1);
        let _ = map.iter(Size::from_bytes(11), Size::from_bytes(11));
    }
}
<|endoftext|>