<commit_before>#![feature(old_io, core)]
#![cfg_attr(test, deny(warnings))]

extern crate time;
extern crate conduit;
extern crate "conduit-middleware" as middleware;

use conduit::{Request, Method};
use middleware::Middleware;
use std::borrow::Cow;
use std::error::Error;
use std::old_io::util::NullReader;
use std::string::CowString;
use time::{Tm, strptime, ParseError};

type Response = Result<conduit::Response, Box<Error+Send>>;

#[allow(missing_copy_implementations)]
pub struct ConditionalGet;

impl Middleware for ConditionalGet {
    fn after(&self, req: &mut Request, res: Response) -> Response {
        let mut res = try!(res);

        match req.method() {
            Method::Get | Method::Head => {
                if is_ok(&res) && is_fresh(req, &res) {
                    res.status = (304, "Not Modified");
                    res.headers.remove("Content-Type");
                    res.headers.remove("Content-Length");
                    res.body = Box::new(NullReader);
                }
            },
            _ => ()
        }

        Ok(res)
    }
}

fn is_ok(response: &conduit::Response) -> bool {
    match *response {
        conduit::Response { status: (200, _), .. } => true,
        _ => false
    }
}

fn is_fresh(req: &Request, res: &conduit::Response) -> bool {
    let modified_since = req.headers().find("If-Modified-Since").map(header_val);
    let none_match     = req.headers().find("If-None-Match").map(header_val);

    if modified_since.as_ref().or_else(|| none_match.as_ref()).is_none() {
        return false;
    }

    let mut success = true;

    modified_since.and_then(|modified_since| {
        parse_http_date(modified_since).map_err(|_| { success = false; }).ok()
    }).map(|parsed| {
        success = success && is_modified_since(parsed, res);
    });

    none_match.map(|none_match| {
        success = success && etag_matches(none_match, res);
    });

    success
}

fn etag_matches<S: Str>(none_match: S, res: &conduit::Response) -> bool {
    res.headers.get("ETag").map(|etag| {
        res_header_val(etag).as_slice() == none_match.as_slice()
    }).unwrap_or(false)
}

fn is_modified_since(modified_since: Tm, res: &conduit::Response) -> bool {
    res.headers.get("Last-Modified").and_then(|last_modified| {
        parse_http_date(res_header_val(last_modified).as_slice()).ok()
    }).map(|last_modified| {
        modified_since.to_timespec() >= last_modified.to_timespec()
    }).unwrap_or(false)
}

fn header_val<'a>(header: Vec<&'a str>) -> CowString<'a> {
    if header.len() == 1 {
        Cow::Borrowed(header[0])
    } else {
        Cow::Owned(header.concat())
    }
}

fn res_header_val<'a>(header: &'a Vec<String>) -> CowString<'a> {
    if header.len() == 1 {
        Cow::Borrowed(header[0].as_slice())
    } else {
        Cow::Owned(header.concat())
    }
}

fn parse_http_date<S: Str>(string: S) -> Result<Tm, ()> {
    let string = string.as_slice();

    parse_rfc1123(string)
        .or_else(|_| parse_rfc850(string))
        .or_else(|_| parse_asctime(string))
        .map_err(|_| ())
}

fn parse_rfc1123(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a, %d %b %Y %T GMT")
}

fn parse_rfc850(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a, %d-%m-%y %T GMT")
}

fn parse_asctime(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a %m%t%d %T %Y")
}

#[cfg(test)]
mod tests {
    extern crate "conduit-test" as test;

    use std::collections::HashMap;
    use std::error::Error;
    use std::old_io::MemReader;
    use time;
    use time::Tm;
    use conduit;
    use conduit::{Request, Response, Handler, Method};
    use middleware::MiddlewareBuilder;

    use super::ConditionalGet;

    macro_rules! returning {
        ($code:expr, $($header:expr => $value:expr),+) => ({
            let mut headers = HashMap::new();
            $(headers.insert($header.to_string(), vec!($value.to_string()));)+
            let handler = SimpleHandler::new(headers, $code, "hello");
            let mut stack = MiddlewareBuilder::new(handler);
            stack.add(ConditionalGet);
            stack
        });
        ($($header:expr => $value:expr),+) => ({
            returning!((200, "OK"), $($header => $value),+)
        })
    }

    macro_rules! request {
        ($($header:expr => $value:expr),+) => ({
            let mut req = test::MockRequest::new(Method::Get, "/");
            $(req.header($header, $value.to_string().as_slice());)+
            req
        })
    }

    #[test]
    fn test_sends_304() {
        let handler = returning!("Last-Modified" => httpdate(time::now()));
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => httpdate(time::now())
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_if_older_than_now() {
        let handler = returning!("Last-Modified" => before_now());
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => httpdate(time::now())
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_with_etag() {
        let handler = returning!("ETag" => "1234");
        expect_304(handler.call(&mut request!(
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_time_but_not_etag() {
        let handler = returning!("Last-Modified" => before_now(), "ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "4321"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_etag_but_not_time() {
        let handler = returning!("Last-Modified" => now(), "ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => before_now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_etag() {
        let handler = returning!("ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-None-Match" => "4321"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_time() {
        let handler = returning!("Last-Modified" => now());
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => before_now()
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_with_fresh_time_and_etag() {
        let handler = returning!("Last-Modified" => before_now(), "ETag" => "1234");
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_does_not_affect_non_200() {
        let code = (302, "Found");
        let handler = returning!(code, "Last-Modified" => before_now(), "ETag" => "1234");
        expect(code, handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_does_not_affect_malformed_timestamp() {
        let bad_stamp = time::now().strftime("%Y-%m-%d %H:%M:%S %z").unwrap().to_string();
        let handler = returning!("Last-Modified" => before_now());
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => bad_stamp
        ) as &mut conduit::Request));
    }

    fn expect_304(response: Result<Response, Box<Error+Send>>) {
        let mut response = response.ok().expect("No response");
        let body = response.body.read_to_string().ok().expect("No body");

        assert_eq!(response.status, (304, "Not Modified"));
        assert_eq!(body.as_slice(), "");
    }

    fn expect_200(response: Result<Response, Box<Error+Send>>) {
        expect((200, "OK"), response);
    }

    fn expect(status: (u32, &'static str), response: Result<Response, Box<Error+Send>>) {
        let mut response = response.ok().expect("No response");
        let body = response.body.read_to_string().ok().expect("No body");

        assert_eq!(response.status, status);
        assert_eq!(body.as_slice(), "hello");
    }

    struct SimpleHandler {
        map: HashMap<String, Vec<String>>,
        status: Status,
        body: &'static str
    }

    type Status = (u32, &'static str);

    impl SimpleHandler {
        fn new(map: HashMap<String, Vec<String>>, status: Status, body: &'static str) -> SimpleHandler {
            SimpleHandler { map: map, status: status, body: body }
        }
    }

    impl Handler for SimpleHandler {
        fn call(&self, _: &mut Request) -> Result<Response, Box<Error+Send>> {
            Ok(Response {
                status: self.status,
                headers: self.map.clone(),
                body: Box::new(MemReader::new(self.body.to_string().into_bytes())),
            })
        }
    }

    fn before_now() -> String {
        let mut now = time::now();
        now.tm_year -= 1;
        httpdate(now)
    }

    fn now() -> String {
        httpdate(time::now())
    }

    fn httpdate(time: Tm) -> String {
        time.strftime("%a, %d-%m-%y %T GMT").unwrap().to_string()
    }
}
<commit_msg>Update to rust master<commit_after>#![feature(old_io, core)]
#![cfg_attr(test, deny(warnings))]

extern crate time;
extern crate conduit;
extern crate "conduit-middleware" as middleware;

use conduit::{Request, Method};
use middleware::Middleware;
use std::borrow::Cow;
use std::error::Error;
use std::old_io::util::NullReader;
use time::{Tm, strptime, ParseError};

type Response = Result<conduit::Response, Box<Error+Send>>;

#[allow(missing_copy_implementations)]
pub struct ConditionalGet;

impl Middleware for ConditionalGet {
    fn after(&self, req: &mut Request, res: Response) -> Response {
        let mut res = try!(res);

        match req.method() {
            Method::Get | Method::Head => {
                if is_ok(&res) && is_fresh(req, &res) {
                    res.status = (304, "Not Modified");
                    res.headers.remove("Content-Type");
                    res.headers.remove("Content-Length");
                    res.body = Box::new(NullReader);
                }
            },
            _ => ()
        }

        Ok(res)
    }
}

fn is_ok(response: &conduit::Response) -> bool {
    match *response {
        conduit::Response { status: (200, _), .. } => true,
        _ => false
    }
}

fn is_fresh(req: &Request, res: &conduit::Response) -> bool {
    let modified_since = req.headers().find("If-Modified-Since").map(header_val);
    let none_match     = req.headers().find("If-None-Match").map(header_val);

    if modified_since.as_ref().or_else(|| none_match.as_ref()).is_none() {
        return false;
    }

    let mut success = true;

    modified_since.and_then(|modified_since| {
        parse_http_date(modified_since).map_err(|_| { success = false; }).ok()
    }).map(|parsed| {
        success = success && is_modified_since(parsed, res);
    });

    none_match.map(|none_match| {
        success = success && etag_matches(none_match, res);
    });

    success
}

fn etag_matches<S: Str>(none_match: S, res: &conduit::Response) -> bool {
    res.headers.get("ETag").map(|etag| {
        res_header_val(etag).as_slice() == none_match.as_slice()
    }).unwrap_or(false)
}

fn is_modified_since(modified_since: Tm, res: &conduit::Response) -> bool {
    res.headers.get("Last-Modified").and_then(|last_modified| {
        parse_http_date(res_header_val(last_modified).as_slice()).ok()
    }).map(|last_modified| {
        modified_since.to_timespec() >= last_modified.to_timespec()
    }).unwrap_or(false)
}

fn header_val<'a>(header: Vec<&'a str>) -> Cow<'a, str> {
    if header.len() == 1 {
        Cow::Borrowed(header[0])
    } else {
        Cow::Owned(header.concat())
    }
}

fn res_header_val<'a>(header: &'a Vec<String>) -> Cow<'a, str> {
    if header.len() == 1 {
        Cow::Borrowed(header[0].as_slice())
    } else {
        Cow::Owned(header.concat())
    }
}

fn parse_http_date<S: Str>(string: S) -> Result<Tm, ()> {
    let string = string.as_slice();

    parse_rfc1123(string)
        .or_else(|_| parse_rfc850(string))
        .or_else(|_| parse_asctime(string))
        .map_err(|_| ())
}

fn parse_rfc1123(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a, %d %b %Y %T GMT")
}

fn parse_rfc850(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a, %d-%m-%y %T GMT")
}

fn parse_asctime(string: &str) -> Result<Tm, ParseError> {
    time::strptime(string, "%a %m%t%d %T %Y")
}

#[cfg(test)]
mod tests {
    extern crate "conduit-test" as test;

    use std::collections::HashMap;
    use std::error::Error;
    use std::old_io::MemReader;
    use time;
    use time::Tm;
    use conduit;
    use conduit::{Request, Response, Handler, Method};
    use middleware::MiddlewareBuilder;

    use super::ConditionalGet;

    macro_rules! returning {
        ($code:expr, $($header:expr => $value:expr),+) => ({
            let mut headers = HashMap::new();
            $(headers.insert($header.to_string(), vec!($value.to_string()));)+
            let handler = SimpleHandler::new(headers, $code, "hello");
            let mut stack = MiddlewareBuilder::new(handler);
            stack.add(ConditionalGet);
            stack
        });
        ($($header:expr => $value:expr),+) => ({
            returning!((200, "OK"), $($header => $value),+)
        })
    }

    macro_rules! request {
        ($($header:expr => $value:expr),+) => ({
            let mut req = test::MockRequest::new(Method::Get, "/");
            $(req.header($header, $value.to_string().as_slice());)+
            req
        })
    }

    #[test]
    fn test_sends_304() {
        let handler = returning!("Last-Modified" => httpdate(time::now()));
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => httpdate(time::now())
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_if_older_than_now() {
        let handler = returning!("Last-Modified" => before_now());
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => httpdate(time::now())
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_with_etag() {
        let handler = returning!("ETag" => "1234");
        expect_304(handler.call(&mut request!(
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_time_but_not_etag() {
        let handler = returning!("Last-Modified" => before_now(), "ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "4321"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_etag_but_not_time() {
        let handler = returning!("Last-Modified" => now(), "ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => before_now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_etag() {
        let handler = returning!("ETag" => "1234");
        expect_200(handler.call(&mut request!(
            "If-None-Match" => "4321"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_200_with_fresh_time() {
        let handler = returning!("Last-Modified" => now());
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => before_now()
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_sends_304_with_fresh_time_and_etag() {
        let handler = returning!("Last-Modified" => before_now(), "ETag" => "1234");
        expect_304(handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_does_not_affect_non_200() {
        let code = (302, "Found");
        let handler = returning!(code, "Last-Modified" => before_now(), "ETag" => "1234");
        expect(code, handler.call(&mut request!(
            "If-Modified-Since" => now(),
            "If-None-Match" => "1234"
        ) as &mut conduit::Request));
    }

    #[test]
    fn test_does_not_affect_malformed_timestamp() {
        let bad_stamp = time::now().strftime("%Y-%m-%d %H:%M:%S %z").unwrap().to_string();
        let handler = returning!("Last-Modified" => before_now());
        expect_200(handler.call(&mut request!(
            "If-Modified-Since" => bad_stamp
        ) as &mut conduit::Request));
    }

    fn expect_304(response: Result<Response, Box<Error+Send>>) {
        let mut response = response.ok().expect("No response");
        let body = response.body.read_to_string().ok().expect("No body");

        assert_eq!(response.status, (304, "Not Modified"));
        assert_eq!(body.as_slice(), "");
    }

    fn expect_200(response: Result<Response, Box<Error+Send>>) {
        expect((200, "OK"), response);
    }

    fn expect(status: (u32, &'static str), response: Result<Response, Box<Error+Send>>) {
        let mut response = response.ok().expect("No response");
        let body = response.body.read_to_string().ok().expect("No body");

        assert_eq!(response.status, status);
        assert_eq!(body.as_slice(), "hello");
    }

    struct SimpleHandler {
        map: HashMap<String, Vec<String>>,
        status: Status,
        body: &'static str
    }

    type Status = (u32, &'static str);

    impl SimpleHandler {
        fn new(map: HashMap<String, Vec<String>>, status: Status, body: &'static str) -> SimpleHandler {
            SimpleHandler { map: map, status: status, body: body }
        }
    }

    impl Handler for SimpleHandler {
        fn call(&self, _: &mut Request) -> Result<Response, Box<Error+Send>> {
            Ok(Response {
                status: self.status,
                headers: self.map.clone(),
                body: Box::new(MemReader::new(self.body.to_string().into_bytes())),
            })
        }
    }

    fn before_now() -> String {
        let mut now = time::now();
        now.tm_year -= 1;
        httpdate(now)
    }

    fn now() -> String {
        httpdate(time::now())
    }

    fn httpdate(time: Tm) -> String {
        time.strftime("%a, %d-%m-%y %T GMT").unwrap().to_string()
    }
}
<|endoftext|>


<commit_before>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter { inner: self.hashmap1.iter().chain(self.hashmap2.iter()) }
    }

    pub fn keys(&self) -> Keys<K, V> {
        Keys { inner: self.hashmap1.keys().chain(self.hashmap2.keys()) }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Keys<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Keys<'a, K, V>, hash_map::Keys<'a, K, V>>
}

impl<'a, K, V> Iterator for Keys<'a, K, V> {
    type Item = &'a K;

    #[inline] fn next(&mut self) -> Option<&'a K> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn keys() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.keys() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<commit_msg>Add values()<commit_after>use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::Chain;
use std::collections::hash_map;

pub struct RehashingHashMap<K, V> {
    // NOTE: I tried to make an array of 2 elements, but run into borrowing problems
    hashmap1: HashMap<K, V>,
    hashmap2: HashMap<K, V>,
    is1main: bool,
    rehashing: bool,
}

impl<K, V> RehashingHashMap<K, V>
    where K: Eq + Hash + Clone
{
    pub fn new() -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::new(),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    pub fn with_capacity(capacity: usize) -> RehashingHashMap<K, V> {
        RehashingHashMap {
            hashmap1: HashMap::with_capacity(capacity),
            hashmap2: HashMap::new(),
            is1main: true,
            rehashing: false,
        }
    }

    fn get_main(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap1 } else { &self.hashmap2 }
    }

    fn get_mut_main(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap1 } else { &mut self.hashmap2 }
    }

    fn get_secondary(&self) -> &HashMap<K, V> {
        if self.is1main { &self.hashmap2 } else { &self.hashmap1 }
    }

    fn get_mut_secondary(&mut self) -> &mut HashMap<K, V> {
        if self.is1main { &mut self.hashmap2 } else { &mut self.hashmap1 }
    }

    pub fn rehash(&mut self) {
        if self.rehashing {
            if self.get_secondary().len() == 0 {
                self.drop_secondary();
                return;
            }
            let (mut main, mut sec) = if self.is1main {
                (&mut self.hashmap1, &mut self.hashmap2)
            } else {
                (&mut self.hashmap2, &mut self.hashmap1)
            };
            // unwrap is safe, checked len() > 0 already
            let k: K = sec.keys().take(1).next().unwrap().clone();
            // FIXME: I wish I did not have to clone they key
            // unwrap is safe, we know the key exists in the hashmap
            let val = sec.remove(&k).unwrap();
            main.insert(k, val);
        }
    }

    pub fn capacity(&self) -> usize {
        self.get_main().capacity()
    }

    pub fn reserve(&mut self, additional: usize) {
        self.rehash();
        self.get_mut_main().reserve(additional)
    }

    pub fn is_rehashing(&self) -> bool {
        if !self.rehashing {
            assert_eq!(self.get_secondary().len(), 0);
        }
        self.rehashing
    }

    pub fn shrink_to_fit(&mut self) {
        if !self.rehashing {
            self.rehashing = true;
            self.is1main = !self.is1main;
            let len = self.len();
            self.get_mut_main().reserve(len)
        }
    }

    pub fn len(&self) -> usize {
        self.get_main().len() + self.get_secondary().len()
    }

    pub fn is_empty(&self) -> bool {
        self.get_main().is_empty() && self.get_secondary().is_empty()
    }

    fn drop_secondary(&mut self) {
        self.rehashing = false;
        assert_eq!(self.get_secondary().len(), 0);
        if self.is1main {
            self.hashmap2 = HashMap::new();
        } else {
            self.hashmap1 = HashMap::new();
        }
    }

    fn assert_state(&self) {
        #![allow(dead_code)]
        if self.rehashing {
            assert!(self.get_secondary().capacity() > 0);
        } else {
            assert!(self.get_secondary().capacity() == 0);
        }
    }

    pub fn clear(&mut self) {
        self.get_mut_main().clear();
        self.drop_secondary();
    }

    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
        // while rehashing, they key can be in either hashmap1 or hashmap2
        // but we want to remove them from wherever it is and add it to main
        let mut ret = None;
        if self.rehashing || self.is1main {
            ret = self.hashmap1.remove(&k);
        }
        if ret.is_none() && (self.rehashing || !self.is1main) {
            ret = self.hashmap2.remove(&k);
        }
        self.get_mut_main().insert(k, v);
        self.rehash();
        ret
    }

    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
            where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            match self.get_main().get(k) {
                Some(ref v) => Some(v),
                None => self.get_secondary().get(k),
            }
        } else {
            self.get_main().get(k)
        }
    }

    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
        where K: Borrow<Q>, Q: Hash + Eq {
        if self.rehashing {
            self.rehash();
            match self.get_mut_main().remove(k) {
                Some(v) => Some(v),
                None => self.get_mut_secondary().remove(k),
            }
        } else {
            self.get_mut_main().remove(k)
        }
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter { inner: self.hashmap1.iter().chain(self.hashmap2.iter()) }
    }

    pub fn keys(&self) -> Keys<K, V> {
        Keys { inner: self.hashmap1.keys().chain(self.hashmap2.keys()) }
    }

    pub fn values(&self) -> Values<K, V> {
        Values { inner: self.hashmap1.values().chain(self.hashmap2.values()) }
    }
}

#[derive(Clone)]
pub struct Iter<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Iter<'a, K, V>, hash_map::Iter<'a, K, V>>
}

impl<'a, K, V> Iterator for Iter<'a, K, V> {
    type Item = (&'a K, &'a V);

    #[inline] fn next(&mut self) -> Option<(&'a K, &'a V)> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Keys<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Keys<'a, K, V>, hash_map::Keys<'a, K, V>>
}

impl<'a, K, V> Iterator for Keys<'a, K, V> {
    type Item = &'a K;

    #[inline] fn next(&mut self) -> Option<&'a K> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[derive(Clone)]
pub struct Values<'a, K: 'a, V: 'a> {
    inner: Chain<hash_map::Values<'a, K, V>, hash_map::Values<'a, K, V>>
}

impl<'a, K, V> Iterator for Values<'a, K, V> {
    type Item = &'a V;

    #[inline] fn next(&mut self) -> Option<&'a V> { self.inner.next() }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
}

#[test]
fn capacity() {
    let mut hash:RehashingHashMap<u8, u8> = RehashingHashMap::with_capacity(20);
    assert!(hash.capacity() >= 20);
    hash.reserve(40);
    assert!(hash.capacity() >= 40);
}

#[test]
fn insert() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value1 = 2;
    let value2 = 3;

    assert_eq!(hash.insert(key.clone(), value1.clone()), None);
    assert_eq!(hash.insert(key.clone(), value2.clone()), Some(value1.clone()));
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.insert(key.clone(), value1.clone()), Some(value2.clone()));
    assert!(!hash.is_rehashing());
    hash.assert_state();
}

#[test]
fn insert_many_rehash_get() {
    let mut hash = RehashingHashMap::new();

    let len = 1000;

    for i in 0..len {
        hash.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2){
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }

    for _ in 0..(len / 2 + 1){
        hash.rehash();
    }
    assert!(!hash.is_rehashing());
    hash.assert_state();

    assert_eq!(hash.len(), len);

    for i in 0..len {
        assert_eq!(hash.get(&i).unwrap(), &i);
    }
    for i in len..(len * 2) {
        assert!(hash.get(&i).is_none());
    }
}

#[test]
fn is_empty() {
    let mut hash = RehashingHashMap::new();
    assert!(hash.is_empty());

    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    assert!(!hash.is_empty());
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert!(!hash.is_empty());
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert!(!hash.is_empty());
}

#[test]
fn clear() {
    let mut hash = RehashingHashMap::with_capacity(1000);
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.clear();
    hash.assert_state();

    assert!(hash.capacity() >= 1000);
}

#[test]
fn remove0() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove1() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    assert!(hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn remove2() {
    let mut hash = RehashingHashMap::new();
    let key = 0;
    let value = 2;
    assert_eq!(hash.insert(key.clone(), value.clone()), None);
    hash.shrink_to_fit();
    hash.rehash();
    hash.rehash();
    assert!(!hash.is_rehashing());
    assert_eq!(hash.remove(&key).unwrap(), value);
}

#[test]
fn iterator() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for (_, i) in hash.iter() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn keys() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.keys() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}

#[test]
fn values() {
    let len = 100;
    let mut hash = RehashingHashMap::with_capacity(len);
    let mut control = HashMap::new();
    for i in 0..len {
        hash.insert(i.clone(), i.clone());
        control.insert(i.clone(), i.clone());
    }
    hash.shrink_to_fit();
    for _ in 0..(len / 2) {
        hash.rehash();
    }
    assert!(hash.is_rehashing());

    for i in hash.values() {
        control.remove(&i);
    }
    assert!(control.is_empty());
}
<|endoftext|>


<commit_before>static DIVISOR_IS_ZERO: &'static str              = "Divisor is zero";
static DIVISOR_LARGER_THAN_DIVIDEND: &'static str = "Divisor is larger than the dividend";

pub trait SplitInto {
    fn split_into(&self, divisor: uint) -> Result<Vec<int>, &'static str>;
}

impl SplitInto for int {
    fn split_into(&self, divisor: uint) -> Result<Vec<int>, &'static str> {
        let dividend = *self;
        if divisor == 0                     { return Err(DIVISOR_IS_ZERO) };
        if dividend == 0                    { return Ok(Vec::from_elem(divisor,0i)) };
        if divisor > dividend.abs() as uint { return Err(DIVISOR_LARGER_THAN_DIVIDEND) }

        let quotient  = dividend / divisor as int;
        let remainder = dividend.abs() as uint % divisor;

        let mut parts: Vec<int> = Vec::from_elem(divisor, quotient);

        let adjustment = if dividend > 0 { 1i } else { -1i };

        for x in range(0, remainder as uint) {
            *parts.get_mut(x) = parts[x] + adjustment;
        }

        parts.reverse();
        return Ok(parts);
    }
}
<commit_msg>Clean up<commit_after>static DIVISOR_IS_ZERO: &'static str              = "Divisor is zero";
static DIVISOR_LARGER_THAN_DIVIDEND: &'static str = "Divisor is larger than the dividend";

pub trait SplitInto {
    fn split_into(&self, divisor: uint) -> Result<Vec<int>, &'static str>;
}

impl SplitInto for int {
    fn split_into(&self, divisor: uint) -> Result<Vec<int>, &'static str> {
        let dividend = *self;
        if divisor == 0                     { return Err(DIVISOR_IS_ZERO) };
        if dividend == 0                    { return Ok(Vec::from_elem(divisor,0i)) };
        if divisor > dividend.abs() as uint { return Err(DIVISOR_LARGER_THAN_DIVIDEND) }

        let quotient  = dividend / divisor as int;
        let remainder = dividend.abs() as uint % divisor;
        let adjustment = if dividend > 0 { 1i } else { -1i };

        let mut parts: Vec<int> = Vec::from_elem(divisor, quotient);

        for x in range(0, remainder as uint) {
            *parts.get_mut(x) = parts[x] + adjustment;
        }

        parts.reverse();
        return Ok(parts);
    }
}
<|endoftext|>


<commit_before>//! `sled` is a high-performance embedded database with
//! an API that is similar to a `BTreeMap<[u8], [u8]>`,
//! but with several additional capabilities for
//! assisting creators of stateful systems.
//!
//! It is fully thread-safe, and all operations are
//! atomic. Multiple `Tree`s with isolated keyspaces
//! are supported with the
//! [`Db::open_tree`](struct.Db.html#method.open_tree) method.
//!
//! ACID transactions involving reads and writes to
//! multiple items are supported with the
//! [`Tree::transaction`](struct.Tree.html#method.transaction)
//! method. Transactions may also operate over
//! multiple `Tree`s (see
//! [`Tree::transaction`](struct.Tree.html#method.transaction)
//! docs for more info).
//!
//! Users may also subscribe to updates on individual
//! `Tree`s by using the
//! [`Tree::watch_prefix`](struct.Tree.html#method.watch_prefix)
//! method, which returns a blocking `Iterator` over
//! updates to keys that begin with the provided
//! prefix. You may supply an empty prefix to subscribe
//! to everything.
//!
//! [Merge operators](https://github.com/spacejam/sled/wiki/merge-operators)
//! (aka read-modify-write operators) are supported. A
//! merge operator is a function that specifies
//! how new data can be merged into an existing value
//! without requiring both a read and a write.
//! Using the
//! [`Tree::merge`](struct.Tree.html#method.merge)
//! method, you may "push" data to a `Tree` value
//! and have the provided merge operator combine
//! it with the existing value, if there was one.
//! They are set on a per-`Tree` basis, and essentially
//! allow any sort of data structure to be built
//! using merges as an atomic high-level operation.
//!
//! `sled` is built by experienced database engineers
//! who think users should spend less time tuning and
//! working against high-friction APIs. Expect
//! significant ergonomic and performance improvements
//! over time. Most surprises are bugs, so please
//! [let us know](mailto:t@jujit.su?subject=sled%20sucks!!!) if something
//! is high friction.
//!
//! # Examples
//!
//! ```
//! # let _ = std::fs::remove_dir_all("my_db");
//! let t = sled::open("my_db").unwrap();
//!
//! // insert and get
//! t.insert(b"yo!", b"v1");
//! assert_eq!(&t.get(b"yo!").unwrap().unwrap(), b"v1");
//!
//! // Atomic compare-and-swap.
//! t.compare_and_swap(
//!     b"yo!",      // key
//!     Some(b"v1"), // old value, None for not present
//!     Some(b"v2"), // new value, None for delete
//! )
//! .unwrap();
//!
//! // Iterates over key-value pairs, starting at the given key.
//! let scan_key: &[u8] = b"a non-present key before yo!";
//! let mut iter = t.range(scan_key..);
//! assert_eq!(&iter.next().unwrap().unwrap().0, b"yo!");
//! assert_eq!(iter.next(), None);
//!
//! t.remove(b"yo!");
//! assert_eq!(t.get(b"yo!"), Ok(None));
//! # let _ = std::fs::remove_dir_all("my_db");
//! ```
#![doc(
    html_logo_url = "https://raw.githubusercontent.com/spacejam/sled/master/art/tree_face_anti-transphobia.png"
)]
#![cfg_attr(test, deny(warnings))]
#![deny(
    missing_docs,
    future_incompatible,
    nonstandard_style,
    rust_2018_idioms,
    missing_copy_implementations,
    trivial_casts,
    trivial_numeric_casts,
    unsafe_code,
    unused_qualifications
)]
#![deny(
    // over time, consider enabling the commented-out lints below
    clippy::cast_lossless,
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss,
    clippy::checked_conversions,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    // clippy::else_if_without_else,
    clippy::empty_enum,
    clippy::explicit_into_iter_loop,
    clippy::explicit_iter_loop,
    clippy::expl_impl_clone_on_copy,
    clippy::fallible_impl_from,
    clippy::filter_map,
    clippy::filter_map_next,
    clippy::find_map,
    clippy::float_arithmetic,
    clippy::get_unwrap,
    clippy::if_not_else,
    // clippy::indexing_slicing,
    clippy::inline_always,
    //clippy::integer_arithmetic,
    clippy::invalid_upcast_comparisons,
    clippy::items_after_statements,
    clippy::map_flatten,
    clippy::match_same_arms,
    clippy::maybe_infinite_iter,
    clippy::mem_forget,
    // clippy::missing_const_for_fn,
    // clippy::missing_docs_in_private_items,
    clippy::module_name_repetitions,
    clippy::multiple_crate_versions,
    clippy::multiple_inherent_impl,
    clippy::mut_mut,
    clippy::needless_borrow,
    clippy::needless_continue,
    clippy::needless_pass_by_value,
    clippy::non_ascii_literal,
    clippy::option_map_unwrap_or,
    clippy::option_map_unwrap_or_else,
    clippy::path_buf_push_overwrite,
    clippy::print_stdout,
    clippy::pub_enum_variant_names,
    clippy::redundant_closure_for_method_calls,
    clippy::replace_consts,
    clippy::result_map_unwrap_or_else,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::string_add,
    clippy::string_add_assign,
    clippy::type_repetition_in_bounds,
    clippy::unicode_not_nfc,
    // clippy::unimplemented,
    clippy::unseparated_literal_suffix,
    clippy::used_underscore_binding,
    clippy::wildcard_dependencies,
    // clippy::wildcard_enum_match_arm,
    clippy::wrong_pub_self_convention,
)]
#![recursion_limit = "128"]

macro_rules! io_fail {
    ($config:expr, $e:expr) => {
        #[cfg(feature = "failpoints")]
        {
            if fail::is_active($e) {
                $config.set_global_error(Error::FailPoint);
                return Err(Error::FailPoint).into();
            }
        }
    };
}

macro_rules! testing_assert {
    ($($e:expr),*) => {
        #[cfg(feature = "lock_free_delays")]
        assert!($($e),*)
    };
}

mod batch;
mod binary_search;
mod concurrency_control;
mod config;
mod context;
mod db;
mod dll;
mod fastcmp;
mod fastlock;
mod histogram;
mod iter;
mod ivec;
mod lazy;
mod lru;
mod meta;
mod metrics;
mod node;
mod oneshot;
mod pagecache;
mod prefix;
mod result;
mod serialization;
mod stack;
mod stackvec;
mod subscription;
mod sys_limits;
pub mod transaction;
mod tree;

/// Functionality for conditionally triggering failpoints under test.
#[cfg(feature = "failpoints")]
pub mod fail;

#[cfg(feature = "docs")]
pub mod doc;

#[cfg(not(any(windows, target_os = "linux", target_os = "macos")))]
mod threadpool {
    use super::OneShot;

    /// Just execute a task without involving threads.
    pub fn spawn<F, R>(work: F) -> OneShot<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let (promise_filler, promise) = OneShot::pair();
        promise_filler.fill((work)());
        return promise;
    }
}

#[cfg(any(windows, target_os = "linux", target_os = "macos"))]
mod threadpool;

#[cfg(any(windows, target_os = "linux", target_os = "macos"))]
mod flusher;

#[cfg(feature = "event_log")]
/// The event log helps debug concurrency issues.
pub mod event_log;

#[cfg(feature = "measure_allocs")]
mod measure_allocs;

#[cfg(feature = "measure_allocs")]
#[global_allocator]
static ALLOCATOR: measure_allocs::TrackingAllocator =
    measure_allocs::TrackingAllocator;

const DEFAULT_TREE_ID: &[u8] = b"__sled__default";

/// hidden re-export of items for testing purposes
#[doc(hidden)]
pub use {
    self::{
        config::RunningConfig,
        lazy::Lazy,
        pagecache::{
            constants::{
                MAX_MSG_HEADER_LEN, MAX_SPACE_AMPLIFICATION,
                MINIMUM_ITEMS_PER_SEGMENT, SEG_HEADER_LEN,
            },
            BatchManifest, DiskPtr, Log, LogKind, LogOffset, LogRead, Lsn,
            PageCache, PageId,
        },
        serialization::Serialize,
    },
    crossbeam_epoch::{pin, Atomic, Guard, Owned, Shared},
};

pub use self::{
    batch::Batch,
    config::{Config, Mode},
    db::{open, Db},
    iter::Iter,
    ivec::IVec,
    result::{Error, Result},
    subscription::{Event, Subscriber},
    transaction::Transactional,
    tree::{CompareAndSwapError, Tree},
};

use {
    self::{
        binary_search::binary_search_lub,
        concurrency_control::{ConcurrencyControl, Protector},
        context::Context,
        fastcmp::fastcmp,
        histogram::Histogram,
        lru::Lru,
        meta::Meta,
        metrics::{clock, Measure, M},
        node::{Data, Node},
        oneshot::{OneShot, OneShotFiller},
        result::CasResult,
        stackvec::StackVec,
        subscription::Subscriptions,
        tree::TreeInner,
    },
    crossbeam_utils::{Backoff, CachePadded},
    log::{debug, error, trace, warn},
    pagecache::RecoveryGuard,
    parking_lot::{Condvar, Mutex, RwLock},
    std::{
        collections::BTreeMap,
        convert::TryFrom,
        fmt::{self, Debug},
        io::{Read, Write},
        sync::{
            atomic::{
                AtomicI64 as AtomicLsn, AtomicU64, AtomicUsize,
                Ordering::{Acquire, Relaxed, Release, SeqCst},
            },
            Arc,
        },
    },
};

fn crc32(buf: &[u8]) -> u32 {
    let mut hasher = crc32fast::Hasher::new();
    hasher.update(buf);
    hasher.finalize()
}

fn calculate_message_crc32(header: &[u8], body: &[u8]) -> u32 {
    let mut hasher = crc32fast::Hasher::new();
    hasher.update(body);
    hasher.update(&header[4..]);
    let crc32 = hasher.finalize();
    crc32 ^ 0xFFFF_FFFF
}

#[cfg(any(test, feature = "lock_free_delays"))]
mod debug_delay;

#[cfg(any(test, feature = "lock_free_delays"))]
use debug_delay::debug_delay;

/// This function is useful for inducing random jitter into our atomic
/// operations, shaking out more possible interleavings quickly. It gets
/// fully eliminated by the compiler in non-test code.
#[cfg(not(any(test, feature = "lock_free_delays")))]
const fn debug_delay() {}

/// Link denotes a tree node or its modification fragment such as
/// key addition or removal.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Link {
    /// A new value is set for a given key
    Set(IVec, IVec),
    /// The associated value is removed for a given key
    Del(IVec),
    /// A child of this Index node is marked as mergable
    ParentMergeIntention(PageId),
    /// The merging child has been completely merged into its left sibling
    ParentMergeConfirm,
    /// A Node is marked for being merged into its left sibling
    ChildMergeCap,
}

/// A fast map that is not resistant to collision attacks. Works
/// on 8 bytes at a time.
pub(crate) type FastMap8<K, V> = std::collections::HashMap<
    K,
    V,
    std::hash::BuildHasherDefault<fxhash::FxHasher64>,
>;

/// A fast set that is not resistant to collision attacks. Works
/// on 8 bytes at a time.
pub(crate) type FastSet8<V> = std::collections::HashSet<
    V,
    std::hash::BuildHasherDefault<fxhash::FxHasher64>,
>;

/// Allows arbitrary logic to be injected into mere operations of the
/// `PageCache`.
pub trait MergeOperator: Fn(&[u8], Option<&[u8]>, &[u8]) -> Option<Vec<u8>> {}
impl<F> MergeOperator for F where F: Fn(&[u8], Option<&[u8]>, &[u8]) -> Option<Vec<u8>> {}
<commit_msg>small format change<commit_after>//! `sled` is a high-performance embedded database with
//! an API that is similar to a `BTreeMap<[u8], [u8]>`,
//! but with several additional capabilities for
//! assisting creators of stateful systems.
//!
//! It is fully thread-safe, and all operations are
//! atomic. Multiple `Tree`s with isolated keyspaces
//! are supported with the
//! [`Db::open_tree`](struct.Db.html#method.open_tree) method.
//!
//! ACID transactions involving reads and writes to
//! multiple items are supported with the
//! [`Tree::transaction`](struct.Tree.html#method.transaction)
//! method. Transactions may also operate over
//! multiple `Tree`s (see
//! [`Tree::transaction`](struct.Tree.html#method.transaction)
//! docs for more info).
//!
//! Users may also subscribe to updates on individual
//! `Tree`s by using the
//! [`Tree::watch_prefix`](struct.Tree.html#method.watch_prefix)
//! method, which returns a blocking `Iterator` over
//! updates to keys that begin with the provided
//! prefix. You may supply an empty prefix to subscribe
//! to everything.
//!
//! [Merge operators](https://github.com/spacejam/sled/wiki/merge-operators)
//! (aka read-modify-write operators) are supported. A
//! merge operator is a function that specifies
//! how new data can be merged into an existing value
//! without requiring both a read and a write.
//! Using the
//! [`Tree::merge`](struct.Tree.html#method.merge)
//! method, you may "push" data to a `Tree` value
//! and have the provided merge operator combine
//! it with the existing value, if there was one.
//! They are set on a per-`Tree` basis, and essentially
//! allow any sort of data structure to be built
//! using merges as an atomic high-level operation.
//!
//! `sled` is built by experienced database engineers
//! who think users should spend less time tuning and
//! working against high-friction APIs. Expect
//! significant ergonomic and performance improvements
//! over time. Most surprises are bugs, so please
//! [let us know](mailto:t@jujit.su?subject=sled%20sucks!!!) if something
//! is high friction.
//!
//! # Examples
//!
//! ```
//! # let _ = std::fs::remove_dir_all("my_db");
//! let t = sled::open("my_db").unwrap();
//!
//! // insert and get
//! t.insert(b"yo!", b"v1");
//! assert_eq!(&t.get(b"yo!").unwrap().unwrap(), b"v1");
//!
//! // Atomic compare-and-swap.
//! t.compare_and_swap(
//!     b"yo!",      // key
//!     Some(b"v1"), // old value, None for not present
//!     Some(b"v2"), // new value, None for delete
//! )
//! .unwrap();
//!
//! // Iterates over key-value pairs, starting at the given key.
//! let scan_key: &[u8] = b"a non-present key before yo!";
//! let mut iter = t.range(scan_key..);
//! assert_eq!(&iter.next().unwrap().unwrap().0, b"yo!");
//! assert_eq!(iter.next(), None);
//!
//! t.remove(b"yo!");
//! assert_eq!(t.get(b"yo!"), Ok(None));
//! # let _ = std::fs::remove_dir_all("my_db");
//! ```
#![doc(
    html_logo_url = "https://raw.githubusercontent.com/spacejam/sled/master/art/tree_face_anti-transphobia.png"
)]
#![cfg_attr(test, deny(warnings))]
#![deny(
    missing_docs,
    future_incompatible,
    nonstandard_style,
    rust_2018_idioms,
    missing_copy_implementations,
    trivial_casts,
    trivial_numeric_casts,
    unsafe_code,
    unused_qualifications
)]
#![deny(
    // over time, consider enabling the commented-out lints below
    clippy::cast_lossless,
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss,
    clippy::checked_conversions,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    // clippy::else_if_without_else,
    clippy::empty_enum,
    clippy::explicit_into_iter_loop,
    clippy::explicit_iter_loop,
    clippy::expl_impl_clone_on_copy,
    clippy::fallible_impl_from,
    clippy::filter_map,
    clippy::filter_map_next,
    clippy::find_map,
    clippy::float_arithmetic,
    clippy::get_unwrap,
    clippy::if_not_else,
    // clippy::indexing_slicing,
    clippy::inline_always,
    //clippy::integer_arithmetic,
    clippy::invalid_upcast_comparisons,
    clippy::items_after_statements,
    clippy::map_flatten,
    clippy::match_same_arms,
    clippy::maybe_infinite_iter,
    clippy::mem_forget,
    // clippy::missing_const_for_fn,
    // clippy::missing_docs_in_private_items,
    clippy::module_name_repetitions,
    clippy::multiple_crate_versions,
    clippy::multiple_inherent_impl,
    clippy::mut_mut,
    clippy::needless_borrow,
    clippy::needless_continue,
    clippy::needless_pass_by_value,
    clippy::non_ascii_literal,
    clippy::option_map_unwrap_or,
    clippy::option_map_unwrap_or_else,
    clippy::path_buf_push_overwrite,
    clippy::print_stdout,
    clippy::pub_enum_variant_names,
    clippy::redundant_closure_for_method_calls,
    clippy::replace_consts,
    clippy::result_map_unwrap_or_else,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::string_add,
    clippy::string_add_assign,
    clippy::type_repetition_in_bounds,
    clippy::unicode_not_nfc,
    // clippy::unimplemented,
    clippy::unseparated_literal_suffix,
    clippy::used_underscore_binding,
    clippy::wildcard_dependencies,
    // clippy::wildcard_enum_match_arm,
    clippy::wrong_pub_self_convention,
)]
#![recursion_limit = "128"]

macro_rules! io_fail {
    ($config:expr, $e:expr) => {
        #[cfg(feature = "failpoints")]
        {
            if fail::is_active($e) {
                $config.set_global_error(Error::FailPoint);
                return Err(Error::FailPoint).into();
            }
        }
    };
}

macro_rules! testing_assert {
    ($($e:expr),*) => {
        #[cfg(feature = "lock_free_delays")]
        assert!($($e),*)
    };
}

mod batch;
mod binary_search;
mod concurrency_control;
mod config;
mod context;
mod db;
mod dll;
mod fastcmp;
mod fastlock;
mod histogram;
mod iter;
mod ivec;
mod lazy;
mod lru;
mod meta;
mod metrics;
mod node;
mod oneshot;
mod pagecache;
mod prefix;
mod result;
mod serialization;
mod stack;
mod stackvec;
mod subscription;
mod sys_limits;
pub mod transaction;
mod tree;

/// Functionality for conditionally triggering failpoints under test.
#[cfg(feature = "failpoints")]
pub mod fail;

#[cfg(feature = "docs")]
pub mod doc;

#[cfg(not(any(windows, target_os = "linux", target_os = "macos")))]
mod threadpool {
    use super::OneShot;

    /// Just execute a task without involving threads.
    pub fn spawn<F, R>(work: F) -> OneShot<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let (promise_filler, promise) = OneShot::pair();
        promise_filler.fill((work)());
        return promise;
    }
}

#[cfg(any(windows, target_os = "linux", target_os = "macos"))]
mod threadpool;

#[cfg(any(windows, target_os = "linux", target_os = "macos"))]
mod flusher;

#[cfg(feature = "event_log")]
/// The event log helps debug concurrency issues.
pub mod event_log;

#[cfg(feature = "measure_allocs")]
mod measure_allocs;

#[cfg(feature = "measure_allocs")]
#[global_allocator]
static ALLOCATOR: measure_allocs::TrackingAllocator =
    measure_allocs::TrackingAllocator;

const DEFAULT_TREE_ID: &[u8] = b"__sled__default";

/// hidden re-export of items for testing purposes
#[doc(hidden)]
pub use {
    self::{
        config::RunningConfig,
        lazy::Lazy,
        pagecache::{
            constants::{
                MAX_MSG_HEADER_LEN, MAX_SPACE_AMPLIFICATION,
                MINIMUM_ITEMS_PER_SEGMENT, SEG_HEADER_LEN,
            },
            BatchManifest, DiskPtr, Log, LogKind, LogOffset, LogRead, Lsn,
            PageCache, PageId,
        },
        serialization::Serialize,
    },
    crossbeam_epoch::{pin, Atomic, Guard, Owned, Shared},
};

pub use self::{
    batch::Batch,
    config::{Config, Mode},
    db::{open, Db},
    iter::Iter,
    ivec::IVec,
    result::{Error, Result},
    subscription::{Event, Subscriber},
    transaction::Transactional,
    tree::{CompareAndSwapError, Tree},
};

use {
    self::{
        binary_search::binary_search_lub,
        concurrency_control::{ConcurrencyControl, Protector},
        context::Context,
        fastcmp::fastcmp,
        histogram::Histogram,
        lru::Lru,
        meta::Meta,
        metrics::{clock, Measure, M},
        node::{Data, Node},
        oneshot::{OneShot, OneShotFiller},
        result::CasResult,
        stackvec::StackVec,
        subscription::Subscriptions,
        tree::TreeInner,
    },
    crossbeam_utils::{Backoff, CachePadded},
    log::{debug, error, trace, warn},
    pagecache::RecoveryGuard,
    parking_lot::{Condvar, Mutex, RwLock},
    std::{
        collections::BTreeMap,
        convert::TryFrom,
        fmt::{self, Debug},
        io::{Read, Write},
        sync::{
            atomic::{
                AtomicI64 as AtomicLsn, AtomicU64, AtomicUsize,
                Ordering::{Acquire, Relaxed, Release, SeqCst},
            },
            Arc,
        },
    },
};

fn crc32(buf: &[u8]) -> u32 {
    let mut hasher = crc32fast::Hasher::new();
    hasher.update(buf);
    hasher.finalize()
}

fn calculate_message_crc32(header: &[u8], body: &[u8]) -> u32 {
    let mut hasher = crc32fast::Hasher::new();
    hasher.update(body);
    hasher.update(&header[4..]);
    let crc32 = hasher.finalize();
    crc32 ^ 0xFFFF_FFFF
}

#[cfg(any(test, feature = "lock_free_delays"))]
mod debug_delay;

#[cfg(any(test, feature = "lock_free_delays"))]
use debug_delay::debug_delay;

/// This function is useful for inducing random jitter into our atomic
/// operations, shaking out more possible interleavings quickly. It gets
/// fully eliminated by the compiler in non-test code.
#[cfg(not(any(test, feature = "lock_free_delays")))]
const fn debug_delay() {}

/// Link denotes a tree node or its modification fragment such as
/// key addition or removal.
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum Link {
    /// A new value is set for a given key
    Set(IVec, IVec),
    /// The associated value is removed for a given key
    Del(IVec),
    /// A child of this Index node is marked as mergable
    ParentMergeIntention(PageId),
    /// The merging child has been completely merged into its left sibling
    ParentMergeConfirm,
    /// A Node is marked for being merged into its left sibling
    ChildMergeCap,
}

/// A fast map that is not resistant to collision attacks. Works
/// on 8 bytes at a time.
pub(crate) type FastMap8<K, V> = std::collections::HashMap<
    K,
    V,
    std::hash::BuildHasherDefault<fxhash::FxHasher64>,
>;

/// A fast set that is not resistant to collision attacks. Works
/// on 8 bytes at a time.
pub(crate) type FastSet8<V> = std::collections::HashSet<
    V,
    std::hash::BuildHasherDefault<fxhash::FxHasher64>,
>;

/// Allows arbitrary logic to be injected into mere operations of the
/// `PageCache`.
pub trait MergeOperator:
    Fn(&[u8], Option<&[u8]>, &[u8]) -> Option<Vec<u8>>
{
}
impl<F> MergeOperator for F where
    F: Fn(&[u8], Option<&[u8]>, &[u8]) -> Option<Vec<u8>>
{
}
<|endoftext|>


<commit_before>//! Slow and basic handling of primes.

use std::{iter, cmp};
use std::collections::{Bitv, bitv};

/// Stores information about primes.
pub struct Primes {
    // This only stores odd numbers, since even numbers are mostly
    // non-prime.
    v: Bitv
}

/// Iterator over the primes stored in a sieve.
pub struct PrimeIterator<'a> {
    two: bool,
    iter: iter::Enumerate<bitv::Bits<'a>>,
}

impl Primes {
    /// Construct a `Primes` via a sieve, stopping at `upto`.
    pub fn sieve(upto: uint) -> Primes {
        // bad stuff happens for very small bounds.
        let upto = cmp::max(10, upto);

        let mut is_prime = Bitv::with_capacity((upto + 1) / 2, true);
        // 1 isn't prime
        is_prime.set(0, false);

        // multiples of 3 aren't prime (3 is handled separately, so
        // the ticking works properly)
        for i in iter::range_step(4, upto / 2, 3) {
            is_prime.set(i, false);
        }

        let mut check = 2; // 5
        let mut tick = 1; // step by 2 to get 7
        let bound = (upto as f64).sqrt() as uint + 1;
        while check <= bound {
            if is_prime[check] {
                let p = 2 * check + 1;
                for zero in iter::range_step(2 * check * (check + 1), upto / 2, p) {
                    is_prime.set(zero, false);
                }
            }

            check += tick;
            tick = 3 - tick;
        }

        Primes { v: is_prime }
    }

    /// The largest number stored.
    pub fn upper_bound(&self) -> uint {
        self.v.len() * 2 + 1
    }

    /// Check if `n` is prime, failing if `n` is larger than the upper
    /// bound of this Primes instance.
    pub fn is_prime(&self, n: uint) -> bool {
        assert!(n <= self.upper_bound());
        if n % 2 == 0 {
            // 2 is the evenest prime.
            n == 2
        } else {
            self.v[n / 2]
        }
    }

    /// Iterator over the primes stored in this map.
    pub fn primes<'a>(&'a self) -> PrimeIterator<'a> {
        PrimeIterator {
            two: true,
            iter: self.v.iter().enumerate()
        }
    }

    /// Factorise `n` into (prime, exponent) pairs.
    ///
    /// Returns None if `n` cannot be factored, specifically if `n` is
    /// zero, or if `n` has primes factors larger than the largest
    /// stored in this sieve.
    pub fn factor(&self, mut n: uint) -> Option<Vec<(uint, u32)>> {
        if n == 0 { return None }

        let mut ret = Vec::new();

        for p in self.primes() {
            if n == 1 { break }

            let mut count = 0;
            while n % p == 0 {
                n /= p;
                count += 1;
            }
            if count > 0 {
                ret.push((p,count));
            }
        }

        if n == 1 {
            Some(ret)
        } else {
            // large factors! :(
            None
        }
    }
}

/// Returns estimated bounds for π(n), the number of primes less than
/// or equal to `n`.
///
/// That is, if `(a, b) = estimate_prime_pi(n)`, `a ≤ π(n) ≤ b`.
pub fn estimate_prime_pi(n: uint) -> (uint, uint) {
    if n >= 55 {
        let n = n as f64;
        let lg = n.ln();

        ((n / (lg + 2.0)) as uint, (n / (lg - 4.0)) as uint)
    } else {
        static SMALL_PRIME_PI: [uint, .. 55] =
            [0, 0, /*2*/1, /*3*/2, 2, /*5*/3, 3, /*7*/ 4, 4, 4,
             4, /*11*/5, 5, /*13*/ 6, 6, 6, 6, /*17*/ 7, 7, /*19*/8,
             8, 8, 8, /*23*/9, 9, 9, 9, 9, 9, /*29*/10,
             10, /*31*/11, 11, 11, 11, 11, 11, /*37*/12, 12, 12,
             12, /*41*/13, 13, /*43*/14, 14, 14, 14, /*47*/15, 15, 15,
             15, 15, 15, /*53*/16, 16];
        let x = SMALL_PRIME_PI[n];
        (x, x)
    }
}

impl<'a> Iterator<uint> for PrimeIterator<'a> {
    #[inline]
    fn next(&mut self) -> Option<uint> {
        if self.two {
            self.two = false;
            Some(2)
        } else {
            for (i, is_prime) in self.iter {
                if is_prime {
                    return Some(2 * i + 1)
                }
            }
            None
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        let mut iter = *self;
        // TODO: this doesn't run in constant time, is it super-bad?
        match (iter.next(), iter.next_back()) {
            (Some(lo), Some(hi)) => {
                let (below_hi, above_hi) = estimate_prime_pi(hi);
                let (below_lo, above_lo) = estimate_prime_pi(lo);

                (below_hi - cmp::min(above_lo, below_hi),
                 Some(above_hi - below_lo + 1))
            }
            (Some(_), None) => (1, Some(1)),
            (None, _) => (0, Some(0))
        }
    }
}

impl<'a> DoubleEndedIterator<uint> for PrimeIterator<'a> {
    #[inline]
    fn next_back(&mut self) -> Option<uint> {
        loop {
            match self.iter.next_back() {
                Some((i, true)) => return Some(2 * i + 1),
                Some((_, false)) => {/* continue */}
                None if self.two => {
                    self.two = false;
                    return Some(2)
                }
                None => return None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate test;

    use std::iter::range_step;
    use super::{Primes, estimate_prime_pi};
    use self::test::Bencher;

    #[test]
    fn is_prime() {
        let primes = Primes::sieve(1000);
        let tests = [
            (0, false),
            (1, false),
            (2, true),
            (3, true),
            (4, false),
            (5, true),
            (6, false),
            (7, true),
            (8, false),
            (9, false),
            (10, false),
            (11, true)
                ];

        for &(n, expected) in tests.iter() {
            assert_eq!(primes.is_prime(n), expected);
        }
    }

    #[test]
    fn primes_iterator() {
        let primes = Primes::sieve(50);
        let mut expected = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47];

        assert_eq!(primes.primes().collect::<Vec<uint>>().as_slice(),
                   expected.as_slice());

        expected.reverse();
        assert_eq!(primes.primes().rev().collect::<Vec<uint>>().as_slice(),
                   expected.as_slice());
    }

    #[test]
    fn factor() {
        let primes = Primes::sieve(1000);

        let tests: &[(uint, &[(uint, u32)])] = &[
            (1, &[]),
            (2, &[(2u, 1u32)]),
            (3, &[(3, 1)]),
            (4, &[(2, 2)]),
            (5, &[(5, 1)]),
            (6, &[(2, 1), (3, 1)]),
            (7, &[(7, 1)]),
            (8, &[(2, 3)]),
            (9, &[(3, 2)]),
            (10, &[(2, 1), (5, 1)]),
            ];
        for &(n, expected) in tests.iter() {
            assert_eq!(primes.factor(n), Some(expected.to_vec()));
        }
    }

    #[test]
    fn factor_failures() {
        let primes = Primes::sieve(30);

        assert_eq!(primes.factor(0), None);
        assert_eq!(primes.factor(97), None);
    }

    #[test]
    fn size_hint() {
        for i in range_step(0, 1000, 100) {
            let sieve = Primes::sieve(i);

            let mut primes = sieve.primes();

            // check the size hint at each and every iteration
            loop {
                let (lo, hi) = primes.size_hint();

                let mut copy = primes;
                let len = copy.count();

                let next = primes.next();

                assert!(lo <= len && len <= hi.unwrap(),
                        "found failing size_hint for {} to {}, should satisfy: {} <= {} <= {}",
                        next, i, lo, len, hi);

                if next.is_none() {
                    break
                }
            }
        }
    }

    #[test]
    fn prime_pi() {
        let primes = Primes::sieve(1_000_000);

        let mut last = 0;
        for (i, p) in primes.primes().enumerate() {
            for j in range(last, p) {
                let (lo, hi) = estimate_prime_pi(j);
                assert!(lo <= i && i <= hi,
                        "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                        j, lo, i, hi)
            }
            last = p;
        }
    }

    #[bench]
    fn sieve_small(b: &mut Bencher) {
        b.iter(|| Primes::sieve(100))
    }
    #[bench]
    fn sieve_medium(b: &mut Bencher) {
        b.iter(|| Primes::sieve(10_000))
    }
    #[bench]
    fn sieve_large(b: &mut Bencher) {
        b.iter(|| Primes::sieve(100_000))
    }

    fn bench_iterate(b: &mut Bencher, upto: uint) {
        let sieve = Primes::sieve(upto);

        b.iter(|| {
            sieve.primes().count()
        })
    }

    #[bench]
    fn iterate_small(b: &mut Bencher) { bench_iterate(b, 100) }
    #[bench]
    fn iterate_large(b: &mut Bencher) { bench_iterate(b, 100_000) }
}
<commit_msg>Add "sporadic" prime_pi tests, testing powers of 10 upto 10^17.<commit_after>//! Slow and basic handling of primes.

use std::{iter, cmp};
use std::collections::{Bitv, bitv};

/// Stores information about primes.
pub struct Primes {
    // This only stores odd numbers, since even numbers are mostly
    // non-prime.
    v: Bitv
}

/// Iterator over the primes stored in a sieve.
pub struct PrimeIterator<'a> {
    two: bool,
    iter: iter::Enumerate<bitv::Bits<'a>>,
}

impl Primes {
    /// Construct a `Primes` via a sieve, stopping at `upto`.
    pub fn sieve(upto: uint) -> Primes {
        // bad stuff happens for very small bounds.
        let upto = cmp::max(10, upto);

        let mut is_prime = Bitv::with_capacity((upto + 1) / 2, true);
        // 1 isn't prime
        is_prime.set(0, false);

        // multiples of 3 aren't prime (3 is handled separately, so
        // the ticking works properly)
        for i in iter::range_step(4, upto / 2, 3) {
            is_prime.set(i, false);
        }

        let mut check = 2; // 5
        let mut tick = 1; // step by 2 to get 7
        let bound = (upto as f64).sqrt() as uint + 1;
        while check <= bound {
            if is_prime[check] {
                let p = 2 * check + 1;
                for zero in iter::range_step(2 * check * (check + 1), upto / 2, p) {
                    is_prime.set(zero, false);
                }
            }

            check += tick;
            tick = 3 - tick;
        }

        Primes { v: is_prime }
    }

    /// The largest number stored.
    pub fn upper_bound(&self) -> uint {
        self.v.len() * 2 + 1
    }

    /// Check if `n` is prime, failing if `n` is larger than the upper
    /// bound of this Primes instance.
    pub fn is_prime(&self, n: uint) -> bool {
        assert!(n <= self.upper_bound());
        if n % 2 == 0 {
            // 2 is the evenest prime.
            n == 2
        } else {
            self.v[n / 2]
        }
    }

    /// Iterator over the primes stored in this map.
    pub fn primes<'a>(&'a self) -> PrimeIterator<'a> {
        PrimeIterator {
            two: true,
            iter: self.v.iter().enumerate()
        }
    }

    /// Factorise `n` into (prime, exponent) pairs.
    ///
    /// Returns None if `n` cannot be factored, specifically if `n` is
    /// zero, or if `n` has primes factors larger than the largest
    /// stored in this sieve.
    pub fn factor(&self, mut n: uint) -> Option<Vec<(uint, u32)>> {
        if n == 0 { return None }

        let mut ret = Vec::new();

        for p in self.primes() {
            if n == 1 { break }

            let mut count = 0;
            while n % p == 0 {
                n /= p;
                count += 1;
            }
            if count > 0 {
                ret.push((p,count));
            }
        }

        if n == 1 {
            Some(ret)
        } else {
            // large factors! :(
            None
        }
    }
}

/// Returns estimated bounds for π(n), the number of primes less than
/// or equal to `n`.
///
/// That is, if `(a, b) = estimate_prime_pi(n)`, `a ≤ π(n) ≤ b`.
pub fn estimate_prime_pi(n: uint) -> (uint, uint) {
    if n >= 55 {
        let n = n as f64;
        let lg = n.ln();

        ((n / (lg + 2.0)) as uint, (n / (lg - 4.0)) as uint)
    } else {
        static SMALL_PRIME_PI: [uint, .. 55] =
            [0, 0, /*2*/1, /*3*/2, 2, /*5*/3, 3, /*7*/ 4, 4, 4,
             4, /*11*/5, 5, /*13*/ 6, 6, 6, 6, /*17*/ 7, 7, /*19*/8,
             8, 8, 8, /*23*/9, 9, 9, 9, 9, 9, /*29*/10,
             10, /*31*/11, 11, 11, 11, 11, 11, /*37*/12, 12, 12,
             12, /*41*/13, 13, /*43*/14, 14, 14, 14, /*47*/15, 15, 15,
             15, 15, 15, /*53*/16, 16];
        let x = SMALL_PRIME_PI[n];
        (x, x)
    }
}

impl<'a> Iterator<uint> for PrimeIterator<'a> {
    #[inline]
    fn next(&mut self) -> Option<uint> {
        if self.two {
            self.two = false;
            Some(2)
        } else {
            for (i, is_prime) in self.iter {
                if is_prime {
                    return Some(2 * i + 1)
                }
            }
            None
        }
    }

    fn size_hint(&self) -> (uint, Option<uint>) {
        let mut iter = *self;
        // TODO: this doesn't run in constant time, is it super-bad?
        match (iter.next(), iter.next_back()) {
            (Some(lo), Some(hi)) => {
                let (below_hi, above_hi) = estimate_prime_pi(hi);
                let (below_lo, above_lo) = estimate_prime_pi(lo);

                (below_hi - cmp::min(above_lo, below_hi),
                 Some(above_hi - below_lo + 1))
            }
            (Some(_), None) => (1, Some(1)),
            (None, _) => (0, Some(0))
        }
    }
}

impl<'a> DoubleEndedIterator<uint> for PrimeIterator<'a> {
    #[inline]
    fn next_back(&mut self) -> Option<uint> {
        loop {
            match self.iter.next_back() {
                Some((i, true)) => return Some(2 * i + 1),
                Some((_, false)) => {/* continue */}
                None if self.two => {
                    self.two = false;
                    return Some(2)
                }
                None => return None
            }
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate test;

    use std::iter::range_step;
    use super::{Primes, estimate_prime_pi};
    use self::test::Bencher;

    #[test]
    fn is_prime() {
        let primes = Primes::sieve(1000);
        let tests = [
            (0, false),
            (1, false),
            (2, true),
            (3, true),
            (4, false),
            (5, true),
            (6, false),
            (7, true),
            (8, false),
            (9, false),
            (10, false),
            (11, true)
                ];

        for &(n, expected) in tests.iter() {
            assert_eq!(primes.is_prime(n), expected);
        }
    }

    #[test]
    fn primes_iterator() {
        let primes = Primes::sieve(50);
        let mut expected = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47];

        assert_eq!(primes.primes().collect::<Vec<uint>>().as_slice(),
                   expected.as_slice());

        expected.reverse();
        assert_eq!(primes.primes().rev().collect::<Vec<uint>>().as_slice(),
                   expected.as_slice());
    }

    #[test]
    fn factor() {
        let primes = Primes::sieve(1000);

        let tests: &[(uint, &[(uint, u32)])] = &[
            (1, &[]),
            (2, &[(2u, 1u32)]),
            (3, &[(3, 1)]),
            (4, &[(2, 2)]),
            (5, &[(5, 1)]),
            (6, &[(2, 1), (3, 1)]),
            (7, &[(7, 1)]),
            (8, &[(2, 3)]),
            (9, &[(3, 2)]),
            (10, &[(2, 1), (5, 1)]),
            ];
        for &(n, expected) in tests.iter() {
            assert_eq!(primes.factor(n), Some(expected.to_vec()));
        }
    }

    #[test]
    fn factor_failures() {
        let primes = Primes::sieve(30);

        assert_eq!(primes.factor(0), None);
        assert_eq!(primes.factor(97), None);
    }

    #[test]
    fn size_hint() {
        for i in range_step(0, 1000, 100) {
            let sieve = Primes::sieve(i);

            let mut primes = sieve.primes();

            // check the size hint at each and every iteration
            loop {
                let (lo, hi) = primes.size_hint();

                let mut copy = primes;
                let len = copy.count();

                let next = primes.next();

                assert!(lo <= len && len <= hi.unwrap(),
                        "found failing size_hint for {} to {}, should satisfy: {} <= {} <= {}",
                        next, i, lo, len, hi);

                if next.is_none() {
                    break
                }
            }
        }
    }

    #[test]
    fn prime_pi() {
        fn check(n: u64, pi: u64) {
            let (lo, hi) = estimate_prime_pi(n);
            assert!(lo <= pi && pi <= hi,
                    "found failing estimate at {}, should satisfy: {} <= {} <= {}",
                    n, lo, pi, hi)
        }
        let primes = Primes::sieve(1_000_000);

        let mut last = 0;
        for (i, p) in primes.primes().enumerate() {
            for j in range(last, p) {
                check(j as u64, i as u64);
            }
            last = p;
        }

        let sporadic = [
            (1, 4),
            (2, 25),
            (3, 168),
            (4, 1229),
            (5, 9592),
            (6, 78498),
            (7, 664579),
            (8, 5761455),
            (9, 50847534),
            (10, 455052511),
            (11, 4118054813),
            (12, 37607912018),
            (13, 346065536839),
            (14, 3204941750802),
            (15, 29844570422669),
            (16, 279238341033925),
            (17, 2623557157654233),
            ];
        for &(exponent, real) in sporadic.iter() {
            let n = num::pow(10, exponent);
            check(n, real);
        }
    }

    #[bench]
    fn sieve_small(b: &mut Bencher) {
        b.iter(|| Primes::sieve(100))
    }
    #[bench]
    fn sieve_medium(b: &mut Bencher) {
        b.iter(|| Primes::sieve(10_000))
    }
    #[bench]
    fn sieve_large(b: &mut Bencher) {
        b.iter(|| Primes::sieve(100_000))
    }

    fn bench_iterate(b: &mut Bencher, upto: uint) {
        let sieve = Primes::sieve(upto);

        b.iter(|| {
            sieve.primes().count()
        })
    }

    #[bench]
    fn iterate_small(b: &mut Bencher) { bench_iterate(b, 100) }
    #[bench]
    fn iterate_large(b: &mut Bencher) { bench_iterate(b, 100_000) }
}
<|endoftext|>


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

    #[test]
    fn push_get() {
        let mut v = PVec::new();
        for i in 0..(BRANCH_FACTOR * 4 + 1) {
            v = v.push(i);
        }
        println!("{:#?}", v);
        for i in 0..(BRANCH_FACTOR * 4 + 1) {
            println!("{}:", i);
            assert_eq!(v.get(i), Some(&i));
        }
        assert_eq!(v.get(BRANCH_FACTOR * 4 + 1), None);
    }
}

use std::sync::Arc;
use std::mem;

#[cfg(feature = "narrow_branching")]
const BRANCH_EXPONENT: usize = 2;
#[cfg(feature = "narrow_branching")]
const BRANCH_FACTOR: usize = 4; // 2^BRANCH_EXPONENT

#[cfg(not(feature = "narrow_branching"))]
const BRANCH_EXPONENT: usize = 5;
#[cfg(not(feature = "narrow_branching"))]
const BRANCH_FACTOR: usize = 32; // 2^BRANCH_EXPONENT

#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Default)]
pub struct PVec<T> {
    root: Node<T>,
    len: usize,
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
enum Node<T> {
    Branch {
        children: [Option<Arc<PVec<T>>>; BRANCH_FACTOR],
        depth: usize,
    },
    Leaf { elements: [Option<T>; BRANCH_FACTOR], },
}

#[cfg(feature = "narrow_branching")]
macro_rules! empty_arr {
    () => {
        [None, None,
         None, None]
    }
}

#[cfg(not(feature = "narrow_branching"))]
macro_rules! empty_arr {
    () => {
        [None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None]
    }
}

impl<T> Default for Node<T> {
    fn default() -> Self {
        Node::Leaf { elements: empty_arr!() }
    }
}

impl<T: Clone> Clone for Node<T> {
    fn clone(&self) -> Self {
        match *self {
            Node::Branch {
                ref children,
                ref depth,
            } => {
                let mut kids = empty_arr!();
                for i in 0..children.len() {
                    kids[i] = children[i].clone()
                }
                Node::Branch {
                    children: kids,
                    depth: *depth,
                }
            }
            Node::Leaf { ref elements } => {
                let mut elems = empty_arr!();
                for i in 0..elements.len() {
                    elems[i] = elements[i].clone()
                }
                Node::Leaf { elements: elems }
            }
        }
    }
}

use std::fmt::Debug;

impl<T: Clone + Debug> PVec<T> {
    pub fn new() -> Self {
        PVec {
            root: Node::default(),
            len: 0,
        }
    }

    pub fn push(self, element: T) -> Self {
        if self.len < BRANCH_FACTOR {
            let mut new = self.clone();
            match new.root {
                Node::Leaf { ref mut elements } => {
                    elements[new.len] = Some(element);
                    new.len += 1;
                }
                Node::Branch { ref mut children, .. } => {
                    let i = self.len;
                    let old = mem::replace(&mut children[i], None);
                    match old {
                        Some(child_ref) => {
                            let child = (*child_ref).clone().push(element);
                            if child.len == BRANCH_FACTOR {
                                new.len += 1;
                            }
                            mem::replace(&mut children[i], Some(Arc::new(child)));
                        }
                        None => {
                            let child = PVec::new().push(element);
                            mem::replace(&mut children[i], Some(Arc::new(child)));
                        }
                    }
                }
            }
            new
        } else {
            let mut children = empty_arr!();
            let depth = match self.root {
                Node::Branch { ref depth, .. } => *depth,
                Node::Leaf { .. } => 0,
            };
            children[0] = Some(Arc::new(self));
            children[1] = Some(Arc::new(PVec::new().push(element)));
            PVec {
                root: Node::Branch {
                    children: children,
                    depth: depth + 1,
                },
                len: 1,
            }
        }
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        match self.root {
            Node::Branch {
                ref children,
                ref depth,
            } => {
                let i = (index >> (BRANCH_EXPONENT * depth)) & (BRANCH_FACTOR - 1);
                children[i].as_ref().and_then(|c| c.get(index))
            }
            Node::Leaf { ref elements } => {
                let i = index & (BRANCH_FACTOR - 1);
                elements[i].as_ref()
            }
        }
    }
}
<commit_msg>Slightly larger simple test<commit_after>#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn push_get() {
        let mut v = PVec::new();
        for i in 0..(BRANCH_FACTOR.pow(2) + 1) {
            v = v.push(i);
        }
        println!("{:#?}", v);
        for i in 0..(BRANCH_FACTOR.pow(2) + 1) {
            println!("{}:", i);
            assert_eq!(v.get(i), Some(&i));
        }
        assert_eq!(v.get(BRANCH_FACTOR.pow(2) + 1), None);
    }
}

use std::sync::Arc;
use std::mem;

#[cfg(feature = "narrow_branching")]
const BRANCH_EXPONENT: usize = 2;
#[cfg(feature = "narrow_branching")]
const BRANCH_FACTOR: usize = 4; // 2^BRANCH_EXPONENT

#[cfg(not(feature = "narrow_branching"))]
const BRANCH_EXPONENT: usize = 5;
#[cfg(not(feature = "narrow_branching"))]
const BRANCH_FACTOR: usize = 32; // 2^BRANCH_EXPONENT

#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Default)]
pub struct PVec<T> {
    root: Node<T>,
    len: usize,
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
enum Node<T> {
    Branch {
        children: [Option<Arc<PVec<T>>>; BRANCH_FACTOR],
        depth: usize,
    },
    Leaf { elements: [Option<T>; BRANCH_FACTOR], },
}

#[cfg(feature = "narrow_branching")]
macro_rules! empty_arr {
    () => {
        [None, None,
         None, None]
    }
}

#[cfg(not(feature = "narrow_branching"))]
macro_rules! empty_arr {
    () => {
        [None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None,
         None, None, None, None]
    }
}

impl<T> Default for Node<T> {
    fn default() -> Self {
        Node::Leaf { elements: empty_arr!() }
    }
}

impl<T: Clone> Clone for Node<T> {
    fn clone(&self) -> Self {
        match *self {
            Node::Branch {
                ref children,
                ref depth,
            } => {
                let mut kids = empty_arr!();
                for i in 0..children.len() {
                    kids[i] = children[i].clone()
                }
                Node::Branch {
                    children: kids,
                    depth: *depth,
                }
            }
            Node::Leaf { ref elements } => {
                let mut elems = empty_arr!();
                for i in 0..elements.len() {
                    elems[i] = elements[i].clone()
                }
                Node::Leaf { elements: elems }
            }
        }
    }
}

use std::fmt::Debug;

impl<T: Clone + Debug> PVec<T> {
    pub fn new() -> Self {
        PVec {
            root: Node::default(),
            len: 0,
        }
    }

    pub fn push(self, element: T) -> Self {
        if self.len < BRANCH_FACTOR {
            let mut new = self.clone();
            match new.root {
                Node::Leaf { ref mut elements } => {
                    elements[new.len] = Some(element);
                    new.len += 1;
                }
                Node::Branch { ref mut children, .. } => {
                    let i = self.len;
                    let old = mem::replace(&mut children[i], None);
                    match old {
                        Some(child_ref) => {
                            let child = (*child_ref).clone().push(element);
                            if child.len == BRANCH_FACTOR {
                                new.len += 1;
                            }
                            mem::replace(&mut children[i], Some(Arc::new(child)));
                        }
                        None => {
                            let child = PVec::new().push(element);
                            mem::replace(&mut children[i], Some(Arc::new(child)));
                        }
                    }
                }
            }
            new
        } else {
            let mut children = empty_arr!();
            let depth = match self.root {
                Node::Branch { ref depth, .. } => *depth,
                Node::Leaf { .. } => 0,
            };
            children[0] = Some(Arc::new(self));
            children[1] = Some(Arc::new(PVec::new().push(element)));
            PVec {
                root: Node::Branch {
                    children: children,
                    depth: depth + 1,
                },
                len: 1,
            }
        }
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        match self.root {
            Node::Branch {
                ref children,
                ref depth,
            } => {
                let i = (index >> (BRANCH_EXPONENT * depth)) & (BRANCH_FACTOR - 1);
                children[i].as_ref().and_then(|c| c.get(index))
            }
            Node::Leaf { ref elements } => {
                let i = index & (BRANCH_FACTOR - 1);
                elements[i].as_ref()
            }
        }
    }
}
<|endoftext|>