<commit_before>// Copyright 2019 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![feature(thread_spawn_unchecked)]

extern crate libc;
extern crate rocksdb;

mod util;

use libc::size_t;

use rocksdb::{DBVector, Error, IteratorMode, Options, WriteBatch, DB};
use std::thread;
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { libc::calloc(len, mem::size_of::<u8>()) as *mut u8 };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
    let path = DBPath::new("_rust_rocksdb_externaltest");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");

        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

#[test]
fn db_vector_as_ref_byte_slice() {
    let path = DBPath::new("_rust_rocksdb_db_vector_as_ref_byte_slice");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
        let vector = r.unwrap().unwrap();

        assert!(get_byte_slice(&vector) == b"v1111");
    }
}

fn get_byte_slice<T: AsRef<[u8]>>(source: &'_ T) -> &'_ [u8] {
    source.as_ref()
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            assert!(db.write(batch).is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.write(batch).is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let data = [(b"k1", b"v1111"), (b"k2", b"v2222"), (b"k3", b"v3333")];
        let db = DB::open_default(&path).unwrap();

        for (key, value) in &data {
            assert!(db.put(key, value).is_ok());
        }

        let iter = db.iterator(IteratorMode::Start);

        for (idx, (db_key, db_value)) in iter.enumerate() {
            let (key, value) = data[idx];
            assert_eq!((&key[..], &value[..]), (db_key.as_ref(), db_value.as_ref()));
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let snap = db.snapshot();
        assert!(snap.get(b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1111");

        assert!(db.put(b"k2", b"v2222").is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

#[test]
fn sync_snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_sync_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1").is_ok());
        assert!(db.put(b"k2", b"v2").is_ok());

        let snapshot = db.snapshot();

        // Unsafe here is safe, because `handler.join()` is called at the end of the
        // method to ensure that snapshot will not outlive database.
        let handler_1 = unsafe {
            thread::Builder::new()
                .spawn_unchecked(|| {
                    assert_eq!(
                        snapshot.get(b"k1").unwrap().unwrap().to_utf8().unwrap(),
                        "v1"
                    );
                })
                .unwrap()
        };

        let handler_2 = unsafe {
            thread::Builder::new()
                .spawn_unchecked(|| {
                    assert_eq!(
                        snapshot.get(b"k2").unwrap().unwrap().to_utf8().unwrap(),
                        "v2"
                    );
                })
                .unwrap()
        };

        handler_1.join().unwrap();
        handler_2.join().unwrap();
    }
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<commit_msg>Rewrite sync_snapshot_test without unchecked threads.<commit_after>// Copyright 2019 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate libc;
extern crate rocksdb;

mod util;

use libc::size_t;

use rocksdb::{DBVector, Error, IteratorMode, Options, Snapshot, WriteBatch, DB};
use std::{mem, thread};
use util::DBPath;

#[test]
fn test_db_vector() {
    use std::mem;
    let len: size_t = 4;
    let data: *mut u8 = unsafe { libc::calloc(len, mem::size_of::<u8>()) as *mut u8 };
    let v = unsafe { DBVector::from_c(data, len) };
    let ctrl = [0u8, 0, 0, 0];
    assert_eq!(&*v, &ctrl[..]);
}

#[test]
fn external() {
    let path = DBPath::new("_rust_rocksdb_externaltest");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");

        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

#[test]
fn db_vector_as_ref_byte_slice() {
    let path = DBPath::new("_rust_rocksdb_db_vector_as_ref_byte_slice");

    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
        let vector = r.unwrap().unwrap();

        assert!(get_byte_slice(&vector) == b"v1111");
    }
}

fn get_byte_slice<T: AsRef<[u8]>>(source: &'_ T) -> &'_ [u8] {
    source.as_ref()
}

#[test]
fn errors_do_stuff() {
    let path = DBPath::new("_rust_rocksdb_error");
    let _db = DB::open_default(&path).unwrap();
    let opts = Options::default();
    // The DB will still be open when we try to destroy it and the lock should fail.
    match DB::destroy(&opts, &path) {
        Err(s) => {
            let message = s.to_string();
            assert!(message.find("IO error:").is_some());
            assert!(message.find("_rust_rocksdb_error").is_some());
            assert!(message.find("/LOCK:").is_some());
        }
        Ok(_) => panic!("should fail"),
    }
}

#[test]
fn writebatch_works() {
    let path = DBPath::new("_rust_rocksdb_writebacktest");
    {
        let db = DB::open_default(&path).unwrap();
        {
            // test put
            let mut batch = WriteBatch::default();
            assert!(db.get(b"k1").unwrap().is_none());
            assert_eq!(batch.len(), 0);
            assert!(batch.is_empty());
            let _ = batch.put(b"k1", b"v1111");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.get(b"k1").unwrap().is_none());
            assert!(db.write(batch).is_ok());
            let r: Result<Option<DBVector>, Error> = db.get(b"k1");
            assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
        }
        {
            // test delete
            let mut batch = WriteBatch::default();
            let _ = batch.delete(b"k1");
            assert_eq!(batch.len(), 1);
            assert!(!batch.is_empty());
            assert!(db.write(batch).is_ok());
            assert!(db.get(b"k1").unwrap().is_none());
        }
        {
            // test size_in_bytes
            let mut batch = WriteBatch::default();
            let before = batch.size_in_bytes();
            let _ = batch.put(b"k1", b"v1234567890");
            let after = batch.size_in_bytes();
            assert!(before + 10 <= after);
        }
    }
}

#[test]
fn iterator_test() {
    let path = DBPath::new("_rust_rocksdb_iteratortest");
    {
        let data = [(b"k1", b"v1111"), (b"k2", b"v2222"), (b"k3", b"v3333")];
        let db = DB::open_default(&path).unwrap();

        for (key, value) in &data {
            assert!(db.put(key, value).is_ok());
        }

        let iter = db.iterator(IteratorMode::Start);

        for (idx, (db_key, db_value)) in iter.enumerate() {
            let (key, value) = data[idx];
            assert_eq!((&key[..], &value[..]), (db_key.as_ref(), db_value.as_ref()));
        }
    }
}

#[test]
fn snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_snapshottest");
    {
        let db = DB::open_default(&path).unwrap();

        assert!(db.put(b"k1", b"v1111").is_ok());

        let snap = db.snapshot();
        assert!(snap.get(b"k1").unwrap().unwrap().to_utf8().unwrap() == "v1111");

        assert!(db.put(b"k2", b"v2222").is_ok());

        assert!(db.get(b"k2").unwrap().is_some());
        assert!(snap.get(b"k2").unwrap().is_none());
    }
}

struct SnapshotWrapper {
    snapshot: Snapshot<'static>,
}

impl SnapshotWrapper {
    fn new(db: &DB) -> Self {
        Self {
            snapshot: unsafe { mem::transmute(db.snapshot()) },
        }
    }

    fn check<K>(&self, key: K, value: &str) -> bool
    where
        K: AsRef<[u8]>,
    {
        self.snapshot.get(key).unwrap().unwrap().to_utf8().unwrap() == value
    }
}

#[test]
fn sync_snapshot_test() {
    let path = DBPath::new("_rust_rocksdb_sync_snapshottest");
    let db = DB::open_default(&path).unwrap();

    assert!(db.put(b"k1", b"v1").is_ok());
    assert!(db.put(b"k2", b"v2").is_ok());

    let wrapper = SnapshotWrapper::new(&db);
    let handler_1 = thread::spawn(move || wrapper.check("k1", "v1"));

    let wrapper = SnapshotWrapper::new(&db);
    let handler_2 = thread::spawn(move || wrapper.check("k2", "v2"));

    assert!(handler_1.join().unwrap());
    assert!(handler_2.join().unwrap());
}

#[test]
fn set_option_test() {
    let path = DBPath::new("_rust_rocksdb_set_optionstest");
    {
        let db = DB::open_default(&path).unwrap();
        // set an option to valid values
        assert!(db
            .set_options(&[("disable_auto_compactions", "true")])
            .is_ok());
        assert!(db
            .set_options(&[("disable_auto_compactions", "false")])
            .is_ok());
        // invalid names/values should result in an error
        assert!(db
            .set_options(&[("disable_auto_compactions", "INVALID_VALUE")])
            .is_err());
        assert!(db
            .set_options(&[("INVALID_NAME", "INVALID_VALUE")])
            .is_err());
        // option names/values must not contain NULLs
        assert!(db
            .set_options(&[("disable_auto_compactions", "true\0")])
            .is_err());
        assert!(db
            .set_options(&[("disable_auto_compactions\0", "true")])
            .is_err());
        // empty options are not allowed
        assert!(db.set_options(&[]).is_err());
        // multiple options can be set in a single API call
        let multiple_options = [
            ("paranoid_file_checks", "true"),
            ("report_bg_io_stats", "true"),
        ];
        db.set_options(&multiple_options).unwrap();
    }
}
<|endoftext|>


<commit_before>mod array;
pub mod date_and_time;
pub mod floats;
mod integers;
mod primitives;
#[cfg(feature = "uuid")]
mod uuid;

#[doc(hidden)]
pub mod sql_types {
    #[derive(Debug, Clone, Copy, Default)] pub struct Oid;
    #[derive(Debug, Clone, Copy, Default)] pub struct Array<T>(T);
    pub type SmallSerial = ::types::SmallInt;
    pub type Serial = ::types::Integer;
    pub type BigSerial = ::types::BigInt;
    #[cfg(feature = "uuid")]
    #[derive(Debug, Clone, Copy, Default)] pub struct Uuid;
    pub type Bytea = ::types::Binary;
}
<commit_msg>Add Bpchar type aliased to VarChar<commit_after>mod array;
pub mod date_and_time;
pub mod floats;
mod integers;
mod primitives;
#[cfg(feature = "uuid")]
mod uuid;

#[doc(hidden)]
pub mod sql_types {
    #[derive(Debug, Clone, Copy, Default)] pub struct Oid;
    #[derive(Debug, Clone, Copy, Default)] pub struct Array<T>(T);
    pub type SmallSerial = ::types::SmallInt;
    pub type Serial = ::types::Integer;
    pub type BigSerial = ::types::BigInt;
    #[cfg(feature = "uuid")]
    #[derive(Debug, Clone, Copy, Default)] pub struct Uuid;
    pub type Bytea = ::types::Binary;
    pub type Bpchar = ::types::VarChar;
}
<|endoftext|>


<commit_before>use std::num;

pub trait Addressable<ADDR: Int> {
	pub fn get (&self, addr: ADDR) -> u8;
	pub fn set (&mut self, addr: ADDR, data: u8);
}

// FIXME: With default methods, we won't need this anymore
pub trait AddressableUtil<ADDR: Int> {
	pub fn getx (&self, addr:ADDR, offset: int) -> u8;
	pub fn get_be_n (&self, addr: ADDR, nbytes: uint) -> u64;
	pub fn get_be<T: Int> (&self, addr: ADDR) -> T;
	pub fn get_le_n (&self, addr: ADDR, nbytes: uint) -> u64;
	pub fn get_le<T: Int> (&self, addr: ADDR) -> T;

	pub fn setx (&mut self, addr: ADDR, offset: int, data: u8);
	pub fn set_be_n (&mut self, addr: ADDR, nbytes: uint, val: u64);
	pub fn set_be<T: Int> (&mut self, addr: ADDR, val: T);
	pub fn set_le_n (&mut self, addr: ADDR, nbytes: uint, val: u64);
	pub fn set_le<T: Int> (&mut self, addr: ADDR, val: T);
}

impl<ADDR: Int+Copy, A: Addressable<ADDR>> AddressableUtil<ADDR> for A {
	pub fn getx (&self, addr:ADDR, offset: int) -> u8 {
		self.get(addr + num::cast(offset))
	}

	pub fn get_be_n (&self, addr: ADDR, nbytes: uint) -> u64 {
		assert!(nbytes > 0 && nbytes < 8);
		let mut val = 0u64;
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			val |= self.get(addr + num::cast(i)) as u64 << (nbytes-i-1) * 8;
		}
		val
	}

	pub fn get_be<T: Int> (&self, addr: ADDR) -> T {
		num::cast(self.get_be_n(addr, num::Primitive::bytes::<T>()))
	}

	pub fn get_le_n (&self, addr: ADDR, nbytes: uint) -> u64 {
		assert!(nbytes > 0 && nbytes < 8);
		let mut val = 0u64;
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			val |= self.get(addr + num::cast(i)) as u64 << i * 8;
		}
		val
	}

	pub fn get_le<T: Int> (&self, addr: ADDR) -> T {
		num::cast(self.get_le_n(addr, num::Primitive::bytes::<T>()))
	}

	pub fn setx (&mut self, addr: ADDR, offset: int, data: u8) {
		self.set(addr + num::cast(offset), data);
	}

	pub fn set_be_n (&mut self, addr: ADDR, nbytes: uint, val: u64) {
		assert!(nbytes > 0 && nbytes < 8);
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			self.set(addr + num::cast(i), (val >> (nbytes-i-1) * 8) as u8);
		}
	}

	pub fn set_be<T: Int> (&mut self, addr: ADDR, val: T) {
		self.set_be_n(addr, num::Primitive::bytes::<T>(), num::cast(val));
	}

	pub fn set_le_n (&mut self, addr: ADDR, nbytes: uint, val: u64) {
		assert!(nbytes > 0 && nbytes < 8);
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			self.set(addr + num::cast(i), (val >> i * 8) as u8);
		}
	}

	pub fn set_le<T: Int> (&mut self, addr: ADDR, val: T) {
		self.set_le_n(addr, num::Primitive::bytes::<T>(), num::cast(val));
	}
}

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

	struct DummyData;
	impl Addressable<u16> for DummyData {
		pub fn get (&self, addr: u16) -> u8 {
			addr as u8
		}
		pub fn set (&mut self, addr: u16, data: u8) {
			assert_eq!(data, addr as u8);
		}
	}

	#[test]
	fn test_get () {
		let data = DummyData;
		assert_eq!(data.get(0x0012), 0x12);
		assert_eq!(data.get(0x1234), 0x34);
	}

	#[test]
	fn test_getx () {
		let data = DummyData;
		assert_eq!(data.getx(0x0012_u16, 5), 0x17);
		assert_eq!(data.getx(0x1234_u16, 5), 0x39);
	}

	#[test]
	fn test_get_big_endian () {
		let data = DummyData;
		assert_eq!(      0x02_u8 , data.get_be(0x0002_u16));
		assert_eq!(      0x54_u8 , data.get_be(0x0054_u16));
		assert_eq!(    0x0203_u16, data.get_be(0x0002_u16));
		assert_eq!(    0x5455_u16, data.get_be(0x0054_u16));
		assert_eq!(0x02030405_u32, data.get_be(0x0002_u16));
		assert_eq!(0x54555657_u32, data.get_be(0x0054_u16));
	}

	#[test]
	fn test_get_little_endian () {
		let data = DummyData;
		assert_eq!(      0x02_u8 , data.get_le(0x0002_u16));
		assert_eq!(      0x54_u8 , data.get_le(0x0054_u16));
		assert_eq!(    0x0302_u16, data.get_le(0x0002_u16));
		assert_eq!(    0x5554_u16, data.get_le(0x0054_u16));
		assert_eq!(0x05040302_u32, data.get_le(0x0002_u16));
		assert_eq!(0x57565554_u32, data.get_le(0x0054_u16));
	}

	#[test]
	fn test_set () {
		let mut data = DummyData;
		data.set(0x0012, 0x12);
		data.set(0x1234, 0x34);
	}

	#[test]
	fn test_setx () {
		let mut data = DummyData;
		data.setx(0x0012_u16, 5, 0x17);
		data.setx(0x1234_u16, 5, 0x39);
	}

	#[test]
	fn test_set_big_endian () {
		let mut data = DummyData;
		data.set_be(0x0002_u16, 0x02_u8);
		data.set_be(0x0054_u16, 0x54_u8);
		data.set_be(0x0002_u16, 0x0203_u16);
		data.set_be(0x0054_u16, 0x5455_u16);
		data.set_be(0x0002_u16, 0x02030405_u32);
		data.set_be(0x0054_u16, 0x54555657_u32);
	}

	#[test]
	fn test_set_little_endian () {
		let mut data = DummyData;
		data.set_le(0x0002_u16, 0x02_u8);
		data.set_le(0x0054_u16, 0x54_u8);
		data.set_le(0x0002_u16, 0x0302_u16);
		data.set_le(0x0054_u16, 0x5554_u16);
		data.set_le(0x0002_u16, 0x05040302_u32);
		data.set_le(0x0054_u16, 0x57565554_u32);
	}
}
<commit_msg>Added tests for signed big/little endian numbers<commit_after>use std::num;

pub trait Addressable<ADDR: Int> {
	pub fn get (&self, addr: ADDR) -> u8;
	pub fn set (&mut self, addr: ADDR, data: u8);
}

// FIXME: With default methods, we won't need this anymore
pub trait AddressableUtil<ADDR: Int> {
	pub fn getx (&self, addr:ADDR, offset: int) -> u8;
	pub fn get_be_n (&self, addr: ADDR, nbytes: uint) -> u64;
	pub fn get_be<T: Int> (&self, addr: ADDR) -> T;
	pub fn get_le_n (&self, addr: ADDR, nbytes: uint) -> u64;
	pub fn get_le<T: Int> (&self, addr: ADDR) -> T;

	pub fn setx (&mut self, addr: ADDR, offset: int, data: u8);
	pub fn set_be_n (&mut self, addr: ADDR, nbytes: uint, val: u64);
	pub fn set_be<T: Int> (&mut self, addr: ADDR, val: T);
	pub fn set_le_n (&mut self, addr: ADDR, nbytes: uint, val: u64);
	pub fn set_le<T: Int> (&mut self, addr: ADDR, val: T);
}

impl<ADDR: Int+Copy, A: Addressable<ADDR>> AddressableUtil<ADDR> for A {
	pub fn getx (&self, addr:ADDR, offset: int) -> u8 {
		self.get(addr + num::cast(offset))
	}

	pub fn get_be_n (&self, addr: ADDR, nbytes: uint) -> u64 {
		assert!(nbytes > 0 && nbytes < 8);
		let mut val = 0u64;
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			val |= self.get(addr + num::cast(i)) as u64 << (nbytes-i-1) * 8;
		}
		val
	}

	pub fn get_be<T: Int> (&self, addr: ADDR) -> T {
		num::cast(self.get_be_n(addr, num::Primitive::bytes::<T>()))
	}

	pub fn get_le_n (&self, addr: ADDR, nbytes: uint) -> u64 {
		assert!(nbytes > 0 && nbytes < 8);
		let mut val = 0u64;
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			val |= self.get(addr + num::cast(i)) as u64 << i * 8;
		}
		val
	}

	pub fn get_le<T: Int> (&self, addr: ADDR) -> T {
		num::cast(self.get_le_n(addr, num::Primitive::bytes::<T>()))
	}

	pub fn setx (&mut self, addr: ADDR, offset: int, data: u8) {
		self.set(addr + num::cast(offset), data);
	}

	pub fn set_be_n (&mut self, addr: ADDR, nbytes: uint, val: u64) {
		assert!(nbytes > 0 && nbytes < 8);
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			self.set(addr + num::cast(i), (val >> (nbytes-i-1) * 8) as u8);
		}
	}

	pub fn set_be<T: Int> (&mut self, addr: ADDR, val: T) {
		self.set_be_n(addr, num::Primitive::bytes::<T>(), num::cast(val));
	}

	pub fn set_le_n (&mut self, addr: ADDR, nbytes: uint, val: u64) {
		assert!(nbytes > 0 && nbytes < 8);
		let mut i = nbytes;
		while i > 0 {
			i -= 1;
			self.set(addr + num::cast(i), (val >> i * 8) as u8);
		}
	}

	pub fn set_le<T: Int> (&mut self, addr: ADDR, val: T) {
		self.set_le_n(addr, num::Primitive::bytes::<T>(), num::cast(val));
	}
}

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

	struct DummyData;
	impl Addressable<u16> for DummyData {
		pub fn get (&self, addr: u16) -> u8 {
			addr as u8
		}
		pub fn set (&mut self, addr: u16, data: u8) {
			assert_eq!(data, addr as u8);
		}
	}

	#[test]
	fn test_get () {
		let data = DummyData;
		assert_eq!(data.get(0x0012), 0x12);
		assert_eq!(data.get(0x1234), 0x34);
	}

	#[test]
	fn test_getx () {
		let data = DummyData;
		assert_eq!(data.getx(0x0012_u16, 5), 0x17);
		assert_eq!(data.getx(0x1234_u16, 5), 0x39);
	}

	#[test]
	fn test_get_big_endian () {
		let data = DummyData;
		assert_eq!(      0x02_u8 , data.get_be(0x0002_u16));
		assert_eq!(      0x54_u8 , data.get_be(0x0054_u16));
		assert_eq!(    0x0203_u16, data.get_be(0x0002_u16));
		assert_eq!(    0x5455_u16, data.get_be(0x0054_u16));
		assert_eq!(0x02030405_u32, data.get_be(0x0002_u16));
		assert_eq!(0x54555657_u32, data.get_be(0x0054_u16));
	}

	#[test]
	fn test_get_signed_big_endian () {
		let data = DummyData;
		assert_eq!(       0x54_i8 , data.get_be(0x0054_u16));
		assert_eq!(      -0x5b_i8 , data.get_be(0x00a5_u16));
		assert_eq!(     0x5455_i16, data.get_be(0x0054_u16));
		assert_eq!(    -0x5a5a_i16, data.get_be(0x00a5_u16));
		assert_eq!( 0x54555657_i32, data.get_be(0x0054_u16));
		assert_eq!(-0x5a595858_i32, data.get_be(0x00a5_u16));
	}

	#[test]
	fn test_get_little_endian () {
		let data = DummyData;
		assert_eq!(      0x02_u8 , data.get_le(0x0002_u16));
		assert_eq!(      0x54_u8 , data.get_le(0x0054_u16));
		assert_eq!(    0x0302_u16, data.get_le(0x0002_u16));
		assert_eq!(    0x5554_u16, data.get_le(0x0054_u16));
		assert_eq!(0x05040302_u32, data.get_le(0x0002_u16));
		assert_eq!(0x57565554_u32, data.get_le(0x0054_u16));
	}

	#[test]
	fn test_get_signed_little_endian () {
		let data = DummyData;
		assert_eq!(       0x54_i8 , data.get_le(0x0054_u16));
		assert_eq!(      -0x5b_i8 , data.get_le(0x00a5_u16));
		assert_eq!(     0x5554_i16, data.get_le(0x0054_u16));
		assert_eq!(    -0x595b_i16, data.get_le(0x00a5_u16));
		assert_eq!( 0x57565554_i32, data.get_le(0x0054_u16));
		assert_eq!(-0x5758595b_i32, data.get_le(0x00a5_u16));
	}

	#[test]
	fn test_set () {
		let mut data = DummyData;
		data.set(0x0012, 0x12);
		data.set(0x1234, 0x34);
	}

	#[test]
	fn test_setx () {
		let mut data = DummyData;
		data.setx(0x0012_u16, 5, 0x17);
		data.setx(0x1234_u16, 5, 0x39);
	}

	#[test]
	fn test_set_big_endian () {
		let mut data = DummyData;
		data.set_be(0x0002_u16,       0x02_u8 );
		data.set_be(0x0054_u16,       0x54_u8 );
		data.set_be(0x0002_u16,     0x0203_u16);
		data.set_be(0x0054_u16,     0x5455_u16);
		data.set_be(0x0002_u16, 0x02030405_u32);
		data.set_be(0x0054_u16, 0x54555657_u32);
	}

	#[test]
	fn test_set_signed_big_endian () {
		let mut data = DummyData;
		data.set_be(0x0054_u16,        0x54_i8 );
		data.set_be(0x00a5_u16,       -0x5b_i8 );
		data.set_be(0x0054_u16,      0x5455_i16);
		data.set_be(0x00a5_u16,     -0x5a5a_i16);
		data.set_be(0x0054_u16,  0x54555657_i32);
		data.set_be(0x00a5_u16, -0x5a595858_i32);
	}

	#[test]
	fn test_set_little_endian () {
		let mut data = DummyData;
		data.set_le(0x0002_u16,       0x02_u8 );
		data.set_le(0x0054_u16,       0x54_u8 );
		data.set_le(0x0002_u16,     0x0302_u16);
		data.set_le(0x0054_u16,     0x5554_u16);
		data.set_le(0x0002_u16, 0x05040302_u32);
		data.set_le(0x0054_u16, 0x57565554_u32);
	}

	fn test_set_signed_little_endian () {
		let mut data = DummyData;
		data.set_le(0x0054_u16,        0x54_i8 );
		data.set_le(0x00a5_u16,       -0x5b_i8 );
		data.set_le(0x0054_u16,      0x5554_i16);
		data.set_le(0x00a5_u16,     -0x595b_i16);
		data.set_le(0x0054_u16,  0x57565554_i32);
		data.set_le(0x00a5_u16, -0x5758595b_i32);
	}
}
<|endoftext|>


<commit_before>#![doc(test(attr(deny(warnings))))]
#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(not(feature = "futures-v0_3"), allow(dead_code, unused_imports))]

//! A crate for integrating signal handling with the Tokio runtime.
//!
//! This provides the [`Signals`][Signals] struct which acts as a
//! [`Stream`] of signals.
//!
//! Note that the `futures-v0_3` feature of this crate must be
//! enabled for `Signals` to implement the `Stream` trait.
//!
//! # Example
//!
//! ```rust
//! # #[cfg(feature = "futures-v0_3")]
//! # mod test {
//! use std::io::Error;
//!
//! use signal_hook::consts::signal::*;
//! use signal_hook_tokio::Signals;
//!
//! use futures::stream::StreamExt;
//!
//! async fn handle_signals(signals: Signals) {
//!     let mut signals = signals.fuse();
//!     while let Some(signal) = signals.next().await {
//!         match signal {
//!             SIGHUP => {
//!                 // Reload configuration
//!                 // Reopen the log file
//!             }
//!             SIGTERM | SIGINT | SIGQUIT => {
//!                 // Shutdown the system;
//!             },
//!             _ => unreachable!(),
//!         }
//!     }
//! }
//!
//! #[tokio::main]
//! # pub
//! async fn main() -> Result<(), Error> {
//!     let signals = Signals::new(&[
//!         SIGHUP,
//!         SIGTERM,
//!         SIGINT,
//!         SIGQUIT,
//!     ])?;
//!     let handle = signals.handle();
//!
//!     let signals_task = tokio::spawn(handle_signals(signals));
//!
//!     // Execute your main program logic
//!
//!     // Terminate the signal stream.
//!     handle.close();
//!     signals_task.await?;
//!
//!     Ok(())
//! }
//! }
//! # fn main() -> Result<(), std::io::Error> {
//! #    #[cfg(feature = "futures-v0_3")]
//! #    test::main()?;
//! #    Ok(())
//! # }
//! ```

macro_rules! implement_signals_with_pipe {
    ($pipe:ty) => {
        use std::borrow::Borrow;
        use std::io::Error;

        use libc::c_int;

        pub use signal_hook::iterator::backend::Handle;
        use signal_hook::iterator::backend::{OwningSignalIterator, SignalDelivery};
        use signal_hook::iterator::exfiltrator::{Exfiltrator, SignalOnly};

        /// An asynchronous [`Stream`] of arriving signals.
        ///
        /// The stream doesn't return the signals in the order they were recieved by
        /// the process and may merge signals received multiple times.
        pub struct SignalsInfo<E: Exfiltrator = SignalOnly>(OwningSignalIterator<$pipe, E>);

        impl<E: Exfiltrator> SignalsInfo<E> {
            /// Create a `SignalsInfo` instance.
            ///
            /// This registers all the signals listed. The same restrictions (panics, errors) apply
            /// as with [`Handle::add_signal`].
            pub fn new<I, S>(signals: I) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
                E: Default,
            {
                Self::with_exfiltrator(signals, E::default())
            }

            /// A constructor with explicit exfiltrator.
            pub fn with_exfiltrator<I, S>(signals: I, exfiltrator: E) -> Result<Self, Error>
            where
                I: IntoIterator<Item = S>,
                S: Borrow<c_int>,
            {
                let (read, write) = UnixStream::pair()?;
                let inner = SignalDelivery::with_pipe(read, write, exfiltrator, signals)?;
                Ok(Self(OwningSignalIterator::new(inner)))
            }

            /// Get a shareable [`Handle`] for this `Signals` instance.
            ///
            /// This can be used to add further signals or close the [`Signals`] instance
            /// which terminates the whole signal stream.
            pub fn handle(&self) -> Handle {
                self.0.handle()
            }
        }

        /// Simplified version of the signals stream.
        ///
        /// This one simply returns the signal numbers, while [`SignalsInfo`] can provide additional
        /// information.
        pub type Signals = SignalsInfo<SignalOnly>;
    };
}

use tokio::net::UnixStream;
implement_signals_with_pipe!(UnixStream);

use std::pin::Pin;
use std::task::{Context, Poll};

#[cfg(feature = "futures-v0_3")]
use futures_core_0_3::Stream;

use signal_hook::iterator::backend::PollResult;
use tokio::io::{AsyncRead, ReadBuf};

impl Signals {
    fn has_signals(read: &mut UnixStream, ctx: &mut Context<'_>) -> Result<bool, Error> {
        let mut buf = [0u8];
        let mut read_buf = ReadBuf::new(&mut buf);
        match Pin::new(read).poll_read(ctx, &mut read_buf) {
            Poll::Pending => Ok(false),
            Poll::Ready(Ok(())) => Ok(true),
            Poll::Ready(Err(error)) => Err(error),
        }
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "futures-v0_3")))]
#[cfg(feature = "futures-v0_3")]
impl Stream for Signals {
    type Item = c_int;

    fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.0.poll_signal(&mut |read| Self::has_signals(read, ctx)) {
            PollResult::Signal(sig) => Poll::Ready(Some(sig)),
            PollResult::Closed => Poll::Ready(None),
            PollResult::Pending => Poll::Pending,
            PollResult::Err(error) => panic!("Unexpected error: {}", error),
        }
    }
}
<commit_msg>Remove the implement_signals_with_pipe macro<commit_after>#![doc(test(attr(deny(warnings))))]
#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(not(feature = "futures-v0_3"), allow(dead_code, unused_imports))]

//! A crate for integrating signal handling with the Tokio runtime.
//!
//! This provides the [`Signals`][Signals] struct which acts as a
//! [`Stream`] of signals.
//!
//! Note that the `futures-v0_3` feature of this crate must be
//! enabled for `Signals` to implement the `Stream` trait.
//!
//! # Example
//!
//! ```rust
//! # #[cfg(feature = "futures-v0_3")]
//! # mod test {
//! use std::io::Error;
//!
//! use signal_hook::consts::signal::*;
//! use signal_hook_tokio::Signals;
//!
//! use futures::stream::StreamExt;
//!
//! async fn handle_signals(signals: Signals) {
//!     let mut signals = signals.fuse();
//!     while let Some(signal) = signals.next().await {
//!         match signal {
//!             SIGHUP => {
//!                 // Reload configuration
//!                 // Reopen the log file
//!             }
//!             SIGTERM | SIGINT | SIGQUIT => {
//!                 // Shutdown the system;
//!             },
//!             _ => unreachable!(),
//!         }
//!     }
//! }
//!
//! #[tokio::main]
//! # pub
//! async fn main() -> Result<(), Error> {
//!     let signals = Signals::new(&[
//!         SIGHUP,
//!         SIGTERM,
//!         SIGINT,
//!         SIGQUIT,
//!     ])?;
//!     let handle = signals.handle();
//!
//!     let signals_task = tokio::spawn(handle_signals(signals));
//!
//!     // Execute your main program logic
//!
//!     // Terminate the signal stream.
//!     handle.close();
//!     signals_task.await?;
//!
//!     Ok(())
//! }
//! }
//! # fn main() -> Result<(), std::io::Error> {
//! #    #[cfg(feature = "futures-v0_3")]
//! #    test::main()?;
//! #    Ok(())
//! # }
//! ```

use std::borrow::Borrow;
use std::io::Error;
use std::pin::Pin;
use std::task::{Context, Poll};

use libc::c_int;

use tokio::io::{AsyncRead, ReadBuf};
use tokio::net::UnixStream;

pub use signal_hook::iterator::backend::Handle;
use signal_hook::iterator::backend::{OwningSignalIterator, PollResult, SignalDelivery};
use signal_hook::iterator::exfiltrator::{Exfiltrator, SignalOnly};

#[cfg(feature = "futures-v0_3")]
use futures_core_0_3::Stream;

/// An asynchronous [`Stream`] of arriving signals.
///
/// The stream doesn't return the signals in the order they were recieved by
/// the process and may merge signals received multiple times.
pub struct SignalsInfo<E: Exfiltrator = SignalOnly>(OwningSignalIterator<UnixStream, E>);

impl<E: Exfiltrator> SignalsInfo<E> {
    /// Create a `SignalsInfo` instance.
    ///
    /// This registers all the signals listed. The same restrictions (panics, errors) apply
    /// as with [`Handle::add_signal`].
    pub fn new<I, S>(signals: I) -> Result<Self, Error>
    where
        I: IntoIterator<Item = S>,
        S: Borrow<c_int>,
        E: Default,
    {
        Self::with_exfiltrator(signals, E::default())
    }

    /// A constructor with explicit exfiltrator.
    pub fn with_exfiltrator<I, S>(signals: I, exfiltrator: E) -> Result<Self, Error>
    where
        I: IntoIterator<Item = S>,
        S: Borrow<c_int>,
    {
        let (read, write) = UnixStream::pair()?;
        let inner = SignalDelivery::with_pipe(read, write, exfiltrator, signals)?;
        Ok(Self(OwningSignalIterator::new(inner)))
    }

    /// Get a shareable [`Handle`] for this `Signals` instance.
    ///
    /// This can be used to add further signals or close the [`Signals`] instance
    /// which terminates the whole signal stream.
    pub fn handle(&self) -> Handle {
        self.0.handle()
    }
}

/// Simplified version of the signals stream.
///
/// This one simply returns the signal numbers, while [`SignalsInfo`] can provide additional
/// information.
pub type Signals = SignalsInfo<SignalOnly>;

impl Signals {
    fn has_signals(read: &mut UnixStream, ctx: &mut Context<'_>) -> Result<bool, Error> {
        let mut buf = [0u8];
        let mut read_buf = ReadBuf::new(&mut buf);
        match Pin::new(read).poll_read(ctx, &mut read_buf) {
            Poll::Pending => Ok(false),
            Poll::Ready(Ok(())) => Ok(true),
            Poll::Ready(Err(error)) => Err(error),
        }
    }
}

#[cfg_attr(docsrs, doc(cfg(feature = "futures-v0_3")))]
#[cfg(feature = "futures-v0_3")]
impl Stream for Signals {
    type Item = c_int;

    fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.0.poll_signal(&mut |read| Self::has_signals(read, ctx)) {
            PollResult::Signal(sig) => Poll::Ready(Some(sig)),
            PollResult::Closed => Poll::Ready(None),
            PollResult::Pending => Poll::Pending,
            PollResult::Err(error) => panic!("Unexpected error: {}", error),
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>Count Matadors jack strait (German: Spitzen).<commit_after><|endoftext|>


<commit_before><commit_msg>Add solution to problem 26<commit_after><|endoftext|>


<commit_before><commit_msg>add GameType constant<commit_after><|endoftext|>


<commit_before><commit_msg>Added mutable cost value to scheduler entry<commit_after><|endoftext|>


<commit_before><commit_msg>Change doc for Nginxs UNIX interconnector<commit_after><|endoftext|>


<commit_before>use consts::*;
use proto::*;
use io::ZkIo;
use listeners::{ListenerSet, Subscription};
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::io::{Read, Result};
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{Sender, sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>,
    listeners: ListenerSet<ZkState>,
}

impl ZooKeeper {

    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
    where F: FnOnce() + Send + 'static {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_|ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
    where W: Watcher + 'static {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher, chroot.clone());
        let listeners = ListenerSet::<ZkState>::new();
        let listeners1 = listeners.clone();
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender(), listeners1);
        let sender = io.sender();

        try!(Self::zk_thread("event", move || { watch.run().unwrap() }));
        try!(Self::zk_thread("io", move || { io.run().unwrap() }));

        Ok(ZooKeeper{chroot: chroot,
                     xid: AtomicIsize::new(1),
                     io: sender,
                     listeners: listeners,
        })
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start)
                }
            },
            None => (None, connect_string.len())
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.to_socket_addrs() {
                Ok(mut addrs) => match addrs.nth(0) {
                    Some(addr) => addr,
                    None => return Err(ZkError::BadArguments)
                },
                Err(_) => return Err(ZkError::BadArguments)
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self, opcode: OpCode, xid: i32, req: Req, watch: Option<Watch>) -> ZkResult<Resp> {
        let rh = RequestHeader{xid: xid, opcode: opcode};
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_|ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest{opcode: opcode, data: buf, listener: Some(resp_tx), watch: watch};

        try!(self.io.send(request).map_err(|err| {
            warn!("error sending request: {:?}", err);
            ZkError::ConnectionLoss
        }));

        let mut response = try!(resp_rx.recv().map_err(|err| {
            warn!("error receiving response: {:?}", err);
            ZkError::ConnectionLoss
        }));

        match response.header.err {
            0 => Ok(try!(ReadFrom::read_from(&mut response.data).map_err(|_|ZkError::MarshallingError))),
            e => Err(FromPrimitive::from_i32(e).unwrap())
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => match path {
                "/" => Ok(chroot.clone()),
                path => Ok(chroot.clone() + try!(Self::validate_path(path)))
            },
            None => Ok(try!(Self::validate_path(path)).to_owned())
        }
    }

    fn cut_chroot(&self, path: String) -> String {
        if let Some(ref chroot) = self.chroot {
            path[chroot.len()..].to_owned()
        } else {
            path
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest{typ: 0, scheme: scheme.to_owned(), auth: auth};

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self, path: &str, data: Vec<u8>, acl: Vec<Acl>, mode: CreateMode) -> ZkResult<String> {
        let req = CreateRequest{path: try!(self.path(path)), data: data, acl: acl, flags: mode as i32};

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(self.cut_chroot(response.path))
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest{path: try!(self.path(path)), version: version};

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: watch};

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e)
        }
    }

    pub fn exists_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Stat> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Exist,
                          watcher: Box::new(watcher)};

        let response: ExistsResponse = try!(self.request(OpCode::Exists, self.xid(), req, Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest{path: try!(self.path(path))};

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Child,
                          watcher: Box::new(watcher)};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: watch};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: watch};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Data,
                          watcher: Box::new(watcher)};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest{path: try!(self.path(path)), acl: acl, version: version};

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest{path: try!(self.path(path)), data: data, version: version};

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn add_listener(&self, chan: Sender<ZkState>) -> Subscription {
        self.listeners.subscribe(chan)
    }

    pub fn remove_listener(&self, sub: Subscription) {
        self.listeners.unsubscribe(sub);
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        if let Err(err) = self.close() {
            info!("error closing zookeeper connection in drop: {:?}", err);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky, host resolving happens in a hectic order, sort?
    #[test]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos").unwrap();
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot.unwrap(), "/mesos");

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<commit_msg>Update zookeeper.rs<commit_after>use consts::*;
use proto::*;
use io::ZkIo;
use listeners::{ListenerSet, Subscription};
use mio;
use num::FromPrimitive;
use watch::{Watch, Watcher, WatchType, ZkWatch};
use std::io::{Read, Result};
use std::net::{SocketAddr, ToSocketAddrs};
use std::result;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::sync::mpsc::{Sender, sync_channel, SyncSender};
use std::time::Duration;
use std::thread;

pub type ZkResult<T> = result::Result<T, ZkError>;

pub struct RawRequest {
    pub opcode: OpCode,
    pub data: ByteBuf,
    pub listener: Option<SyncSender<RawResponse>>,
    pub watch: Option<Watch>
}

pub struct RawResponse {
    pub header: ReplyHeader,
    pub data: ByteBuf
}

pub struct ZooKeeper {
    chroot: Option<String>,
    xid: AtomicIsize,
    io: mio::Sender<RawRequest>,
    listeners: ListenerSet<ZkState>,
}

impl ZooKeeper {

    fn zk_thread<F>(name: &str, task: F) -> ZkResult<thread::JoinHandle<()>>
    where F: FnOnce() + Send + 'static {
        thread::Builder::new()
            .name(name.to_owned())
            .spawn(task)
            .map_err(|_|ZkError::SystemError)
    }

    pub fn connect<W>(connect_string: &str, timeout: Duration, watcher: W) -> ZkResult<ZooKeeper>
    where W: Watcher + 'static {

        let (addrs, chroot) = try!(Self::parse_connect_string(connect_string));

        debug!("Initiating connection to {}", connect_string);

        let watch = ZkWatch::new(watcher, chroot.clone());
        let listeners = ListenerSet::<ZkState>::new();
        let listeners1 = listeners.clone();
        let io = ZkIo::new(addrs.clone(), timeout, watch.sender(), listeners1);
        let sender = io.sender();

        try!(Self::zk_thread("event", move || { watch.run().unwrap() }));
        try!(Self::zk_thread("io", move || { io.run().unwrap() }));

        Ok(ZooKeeper{chroot: chroot,
                     xid: AtomicIsize::new(1),
                     io: sender,
                     listeners: listeners,
        })
    }

    fn parse_connect_string(connect_string: &str) -> ZkResult<(Vec<SocketAddr>, Option<String>)> {
        let (chroot, end) = match connect_string.find('/') {
            Some(start) => {
                match &connect_string[start..connect_string.len()] {
                    "" | "/" => (None, start),
                    chroot => (Some(try!(Self::validate_path(chroot)).to_owned()), start)
                }
            },
            None => (None, connect_string.len())
        };

        let mut addrs = Vec::new();
        for addr_str in connect_string[..end].split(',') {
            let addr = match addr_str.trim().to_socket_addrs() {
                Ok(mut addrs) => match addrs.nth(0) {
                    Some(addr) => addr,
                    None => return Err(ZkError::BadArguments)
                },
                Err(_) => return Err(ZkError::BadArguments)
            };
            addrs.push(addr);
        }

        Ok((addrs, chroot))
    }

    fn xid(&self) -> i32 {
        self.xid.fetch_add(1, Ordering::Relaxed) as i32
    }

    fn request<Req: WriteTo, Resp: ReadFrom>(&self, opcode: OpCode, xid: i32, req: Req, watch: Option<Watch>) -> ZkResult<Resp> {
        let rh = RequestHeader{xid: xid, opcode: opcode};
        let buf = try!(to_len_prefixed_buf(rh, req).map_err(|_|ZkError::MarshallingError));

        let (resp_tx, resp_rx) = sync_channel(0);
        let request = RawRequest{opcode: opcode, data: buf, listener: Some(resp_tx), watch: watch};

        try!(self.io.send(request).map_err(|err| {
            warn!("error sending request: {:?}", err);
            ZkError::ConnectionLoss
        }));

        let mut response = try!(resp_rx.recv().map_err(|err| {
            warn!("error receiving response: {:?}", err);
            ZkError::ConnectionLoss
        }));

        match response.header.err {
            0 => Ok(try!(ReadFrom::read_from(&mut response.data).map_err(|_|ZkError::MarshallingError))),
            e => Err(FromPrimitive::from_i32(e).unwrap())
        }
    }

    fn validate_path(path: &str) -> ZkResult<&str> {
        match path {
            "" => Err(ZkError::BadArguments),
            path => if path.len() > 1 && path.chars().last() == Some('/') {
                    Err(ZkError::BadArguments)
                } else {
                    Ok(path)
                }
        }
    }

    fn path(&self, path: &str) -> ZkResult<String> {
        match self.chroot {
            Some(ref chroot) => match path {
                "/" => Ok(chroot.clone()),
                path => Ok(chroot.clone() + try!(Self::validate_path(path)))
            },
            None => Ok(try!(Self::validate_path(path)).to_owned())
        }
    }

    fn cut_chroot(&self, path: String) -> String {
        if let Some(ref chroot) = self.chroot {
            path[chroot.len()..].to_owned()
        } else {
            path
        }
    }

    pub fn add_auth(&self, scheme: &str, auth: Vec<u8>) -> ZkResult<()> {
        let req = AuthRequest{typ: 0, scheme: scheme.to_owned(), auth: auth};

        let _: EmptyResponse = try!(self.request(OpCode::Auth, -4, req, None));

        Ok(())
    }

    pub fn create(&self, path: &str, data: Vec<u8>, acl: Vec<Acl>, mode: CreateMode) -> ZkResult<String> {
        let req = CreateRequest{path: try!(self.path(path)), data: data, acl: acl, flags: mode as i32};

        let response: CreateResponse = try!(self.request(OpCode::Create, self.xid(), req, None));

        Ok(self.cut_chroot(response.path))
    }

    pub fn delete(&self, path: &str, version: i32) -> ZkResult<()> {
        let req = DeleteRequest{path: try!(self.path(path)), version: version};

        let _: EmptyResponse = try!(self.request(OpCode::Delete, self.xid(), req, None));

        Ok(())
    }

    pub fn exists(&self, path: &str, watch: bool) -> ZkResult<Option<Stat>> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: watch};

        match self.request::<ExistsRequest, ExistsResponse>(OpCode::Exists, self.xid(), req, None) {
            Ok(response) => Ok(Some(response.stat)),
            Err(ZkError::NoNode) => Ok(None),
            Err(e) => Err(e)
        }
    }

    pub fn exists_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Stat> {
        let req = ExistsRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Exist,
                          watcher: Box::new(watcher)};

        let response: ExistsResponse = try!(self.request(OpCode::Exists, self.xid(), req, Some(watch)));

        Ok(response.stat)
    }

    pub fn get_acl(&self, path: &str) -> ZkResult<(Vec<Acl>, Stat)> {
        let req = GetAclRequest{path: try!(self.path(path))};

        let response: GetAclResponse = try!(self.request(OpCode::GetAcl, self.xid(), req, None));

        Ok(response.acl_stat)
    }

    pub fn get_children_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Child,
                          watcher: Box::new(watcher)};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, Some(watch)));

        Ok(response.children)
    }

    pub fn get_children(&self, path: &str, watch: bool) -> ZkResult<Vec<String>> {
        let req = GetChildrenRequest{path: try!(self.path(path)), watch: watch};

        let response: GetChildrenResponse = try!(self.request(OpCode::GetChildren, self.xid(), req, None));

        Ok(response.children)
    }

    pub fn get_data(&self, path: &str, watch: bool) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: watch};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, None));

        Ok(response.data_stat)
    }

    pub fn get_data_w<W: FnOnce(&WatchedEvent) + Send + 'static>(&self, path: &str, watcher: W) -> ZkResult<(Vec<u8>, Stat)> {
        let req = GetDataRequest{path: try!(self.path(path)), watch: true};

        let watch = Watch{path: path.to_owned(),
                          watch_type: WatchType::Data,
                          watcher: Box::new(watcher)};

        let response: GetDataResponse = try!(self.request(OpCode::GetData, self.xid(), req, Some(watch)));

        Ok(response.data_stat)
    }

    pub fn set_acl(&self, path: &str, acl: Vec<Acl>, version: i32) -> ZkResult<Stat> {
        let req = SetAclRequest{path: try!(self.path(path)), acl: acl, version: version};

        let response: SetAclResponse = try!(self.request(OpCode::SetAcl, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn set_data(&self, path: &str, data: Vec<u8>, version: i32) -> ZkResult<Stat> {
        let req = SetDataRequest{path: try!(self.path(path)), data: data, version: version};

        let response: SetDataResponse = try!(self.request(OpCode::SetData, self.xid(), req, None));

        Ok(response.stat)
    }

    pub fn add_listener(&self, chan: Sender<ZkState>) -> Subscription {
        self.listeners.subscribe(chan)
    }

    pub fn remove_listener(&self, sub: Subscription) {
        self.listeners.unsubscribe(sub);
    }

    pub fn close(&self) -> ZkResult<()> {
        let _: EmptyResponse = try!(self.request(OpCode::CloseSession, 0, EmptyRequest, None));

        Ok(())
    }
}

impl Drop for ZooKeeper {
    fn drop(&mut self) {
        if let Err(err) = self.close() {
            info!("error closing zookeeper connection in drop: {:?}", err);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ZooKeeper;
    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};

    // TODO This is flaky, host resolving happens in a hectic order, sort?
    #[test]
    fn parse_connect_string() {
        let (addrs, chroot) = ZooKeeper::parse_connect_string("127.0.0.1:2181,::1:2181/mesos").unwrap();
        assert_eq!(addrs,
                   vec![SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 2181)),
                        SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot.unwrap(), "/mesos");

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);

        let (addrs, chroot) = ZooKeeper::parse_connect_string("::1:2181/").unwrap();
        assert_eq!(addrs, vec![SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0,0,0,0,0,0,0,1), 2181, 0, 0))]);
        assert_eq!(chroot, None);
    }

    #[test]
    #[should_panic(expected = "BadArguments")]
    fn parse_connect_string_fails() {
        // This fails with ZooKeeper.java: Path must not end with / character
        ZooKeeper::parse_connect_string("127.0.0.1:2181/mesos/").unwrap();
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(test, feature(test))]

extern crate random_choice;
extern crate test;

#[cfg(test)]
mod benches {
    use test::Bencher;
    use random_choice::RandomChoice;

    #[bench]
    fn bench_random_choice_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            RandomChoice::random_choice_f64(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            RandomChoice::random_choice_in_place_f64(&mut samples, &weights);
        });
    }

    #[bench]
    fn bench_random_choice_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            RandomChoice::random_choice_f32(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            RandomChoice::random_choice_in_place_f32(&mut samples, &weights);
        });
    }
}
<commit_msg>formatting<commit_after>#![cfg_attr(test, feature(test))]

extern crate random_choice;
extern crate test;

#[cfg(test)]
mod benches {
    use test::Bencher;
    use random_choice::RandomChoice;

    #[bench]
    fn bench_random_choice_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            RandomChoice::random_choice_f64(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_64(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f64> = Vec::with_capacity(capacity);
        let mut weights: Vec<f64> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f64);
            weights.push((i + 1usize) as f64);
        }
        b.iter(|| {
            RandomChoice::random_choice_in_place_f64(&mut samples, &weights);
        });
    }

    #[bench]
    fn bench_random_choice_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            RandomChoice::random_choice_f32(&samples, &weights, 1200 as usize);
        });
    }

    #[bench]
    fn bench_random_choice_in_place_32(b: &mut Bencher) {
        let capacity: usize = 500;
        let mut samples: Vec<f32> = Vec::with_capacity(capacity);
        let mut weights: Vec<f32> = Vec::with_capacity(capacity);

        for i in 0..capacity {
            samples.push((i + 1usize) as f32);
            weights.push((i + 1usize) as f32);
        }
        b.iter(|| {
            RandomChoice::random_choice_in_place_f32(&mut samples, &weights);
        });
    }
}<|endoftext|>


<commit_before>use std::cmp;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::fmt;

use utf8_ranges::{Utf8Range, Utf8Sequences};

use crate::automaton::Automaton;

const DEFAULT_STATE_LIMIT: usize = 10_000; // currently at least 20MB >_<

/// An error that occurred while building a Levenshtein automaton.
///
/// This error is only defined when the `levenshtein` crate feature is enabled.
#[derive(Debug)]
pub enum LevenshteinError {
    /// If construction of the automaton reaches some hard-coded limit
    /// on the number of states, then this error is returned.
    ///
    /// The number given is the limit that was exceeded.
    TooManyStates(usize),
}

impl fmt::Display for LevenshteinError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            LevenshteinError::TooManyStates(size_limit) => write!(
                f,
                "Levenshtein automaton exceeds size limit of \
                           {} states",
                size_limit
            ),
        }
    }
}

impl std::error::Error for LevenshteinError {}

/// A Unicode aware Levenshtein automaton for running efficient fuzzy queries.
///
/// This is only defined when the `levenshtein` crate feature is enabled.
///
/// A Levenshtein automata is one way to search any finite state transducer
/// for keys that *approximately* match a given query. A Levenshtein automaton
/// approximates this by returning all keys within a certain edit distance of
/// the query. The edit distance is defined by the number of insertions,
/// deletions and substitutions required to turn the query into the key.
/// Insertions, deletions and substitutions are based on
/// **Unicode characters** (where each character is a single Unicode scalar
/// value).
///
/// # Example
///
/// This example shows how to find all keys within an edit distance of `1`
/// from `foo`.
///
/// ```rust
/// use fst::automaton::Levenshtein;
/// use fst::{IntoStreamer, Streamer, Set};
///
/// fn main() {
///     let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
///     let set = Set::from_iter(keys).unwrap();
///
///     let lev = Levenshtein::new("foo", 1).unwrap();
///     let mut stream = set.search(&lev).into_stream();
///
///     let mut keys = vec![];
///     while let Some(key) = stream.next() {
///         keys.push(key.to_vec());
///     }
///     assert_eq!(keys, vec![
///         "fo".as_bytes(),   // 1 deletion
///         "fob".as_bytes(),  // 1 substitution
///         "foo".as_bytes(),  // 0 insertions/deletions/substitutions
///         "food".as_bytes(), // 1 insertion
///     ]);
/// }
/// ```
///
/// This example only uses ASCII characters, but it will work equally well
/// on Unicode characters.
///
/// # Warning: experimental
///
/// While executing this Levenshtein automaton against a finite state
/// transducer will be very fast, *constructing* an automaton may not be.
/// Namely, this implementation is a proof of concept. While I believe the
/// algorithmic complexity is not exponential, the implementation is not speedy
/// and it can use enormous amounts of memory (tens of MB before a hard-coded
/// limit will cause an error to be returned).
///
/// This is important functionality, so one should count on this implementation
/// being vastly improved in the future.
pub struct Levenshtein {
    prog: DynamicLevenshtein,
    dfa: Dfa,
}

impl Levenshtein {
    /// Create a new Levenshtein query.
    ///
    /// The query finds all matching terms that are at most `distance`
    /// edit operations from `query`. (An edit operation may be an insertion,
    /// a deletion or a substitution.)
    ///
    /// If the underlying automaton becomes too big, then an error is returned.
    /// Use `new_with_limit` to raise the limit dynamically.
    ///
    /// A `Levenshtein` value satisfies the `Automaton` trait, which means it
    /// can be used with the `search` method of any finite state transducer.
    #[inline]
    pub fn new(
        query: &str,
        distance: u32,
    ) -> Result<Levenshtein, LevenshteinError> {
        let lev = DynamicLevenshtein {
            query: query.to_owned(),
            dist: distance as usize,
        };
        let dfa = DfaBuilder::new(lev.clone()).build()?;
        Ok(Levenshtein { prog: lev, dfa })
    }

    /// Create a new Levenshtein query, but pass the state limit yourself.
    ///
    /// The query finds all matching terms that are at most `distance`
    /// edit operations from `query`. (An edit operation may be an insertion,
    /// a deletion or a substitution.)
    ///
    /// If the underlying automaton becomes too big, then an error is returned.
    /// This limit can be configured with `state_limit`.
    ///
    /// A `Levenshtein` value satisfies the `Automaton` trait, which means it
    /// can be used with the `search` method of any finite state transducer.
    #[inline]
    pub fn new_with_limit(
        query: &str,
        distance: u32,
        state_limit: usize,
    ) -> Result<Levenshtein, LevenshteinError> {
        let lev = DynamicLevenshtein {
            query: query.to_owned(),
            dist: distance as usize,
        };
        let dfa =
            DfaBuilder::new(lev.clone()).build_with_limit(state_limit)?;
        Ok(Levenshtein { prog: lev, dfa })
    }
}

impl fmt::Debug for Levenshtein {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Levenshtein(query: {:?}, distance: {:?})",
            self.prog.query, self.prog.dist
        )
    }
}

#[derive(Clone)]
struct DynamicLevenshtein {
    query: String,
    dist: usize,
}

impl DynamicLevenshtein {
    fn start(&self) -> Vec<usize> {
        (0..self.query.chars().count() + 1).collect()
    }

    fn is_match(&self, state: &[usize]) -> bool {
        state.last().map(|&n| n <= self.dist).unwrap_or(false)
    }

    fn can_match(&self, state: &[usize]) -> bool {
        state.iter().min().map(|&n| n <= self.dist).unwrap_or(false)
    }

    fn accept(&self, state: &[usize], chr: Option<char>) -> Vec<usize> {
        let mut next = vec![state[0] + 1];
        for (i, c) in self.query.chars().enumerate() {
            let cost = if Some(c) == chr { 0 } else { 1 };
            let v = cmp::min(
                cmp::min(next[i] + 1, state[i + 1] + 1),
                state[i] + cost,
            );
            next.push(cmp::min(v, self.dist + 1));
        }
        next
    }
}

impl Automaton for Levenshtein {
    type State = Option<usize>;

    #[inline]
    fn start(&self) -> Option<usize> {
        Some(0)
    }

    #[inline]
    fn is_match(&self, state: &Option<usize>) -> bool {
        state.map(|state| self.dfa.states[state].is_match).unwrap_or(false)
    }

    #[inline]
    fn can_match(&self, state: &Option<usize>) -> bool {
        state.is_some()
    }

    #[inline]
    fn accept(&self, state: &Option<usize>, byte: u8) -> Option<usize> {
        state.and_then(|state| self.dfa.states[state].next[byte as usize])
    }
}

#[derive(Debug)]
struct Dfa {
    states: Vec<State>,
}

struct State {
    next: [Option<usize>; 256],
    is_match: bool,
}

impl fmt::Debug for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "State {{")?;
        writeln!(f, "  is_match: {:?}", self.is_match)?;
        for i in 0..256 {
            if let Some(si) = self.next[i] {
                writeln!(f, "  {:?}: {:?}", i, si)?;
            }
        }
        write!(f, "}}")
    }
}

struct DfaBuilder {
    dfa: Dfa,
    lev: DynamicLevenshtein,
    cache: HashMap<Vec<usize>, usize>,
}

impl DfaBuilder {
    fn new(lev: DynamicLevenshtein) -> DfaBuilder {
        DfaBuilder {
            dfa: Dfa { states: Vec::with_capacity(16) },
            lev,
            cache: HashMap::with_capacity(1024),
        }
    }

    fn build_with_limit(
        mut self,
        state_limit: usize,
    ) -> Result<Dfa, LevenshteinError> {
        let mut stack = vec![self.lev.start()];
        let mut seen = HashSet::new();
        let query = self.lev.query.clone(); // temp work around of borrowck
        while let Some(lev_state) = stack.pop() {
            let dfa_si = self.cached_state(&lev_state).unwrap();
            let mismatch = self.add_mismatch_utf8_states(dfa_si, &lev_state);
            if let Some((next_si, lev_next)) = mismatch {
                if !seen.contains(&next_si) {
                    seen.insert(next_si);
                    stack.push(lev_next);
                }
            }
            for (i, c) in query.chars().enumerate() {
                if lev_state[i] > self.lev.dist {
                    continue;
                }
                let lev_next = self.lev.accept(&lev_state, Some(c));
                let next_si = self.cached_state(&lev_next);
                if let Some(next_si) = next_si {
                    self.add_utf8_sequences(true, dfa_si, next_si, c, c);
                    if !seen.contains(&next_si) {
                        seen.insert(next_si);
                        stack.push(lev_next);
                    }
                }
            }
            if self.dfa.states.len() > state_limit {
                return Err(LevenshteinError::TooManyStates(state_limit));
            }
        }
        Ok(self.dfa)
    }

    fn build(mut self) -> Result<Dfa, LevenshteinError> {
        self.build_with_limit(DEFAULT_STATE_LIMIT)
    }

    fn cached_state(&mut self, lev_state: &[usize]) -> Option<usize> {
        self.cached(lev_state).map(|(si, _)| si)
    }

    fn cached(&mut self, lev_state: &[usize]) -> Option<(usize, bool)> {
        if !self.lev.can_match(lev_state) {
            return None;
        }
        Some(match self.cache.entry(lev_state.to_vec()) {
            Entry::Occupied(v) => (*v.get(), true),
            Entry::Vacant(v) => {
                let is_match = self.lev.is_match(lev_state);
                self.dfa.states.push(State { next: [None; 256], is_match });
                (*v.insert(self.dfa.states.len() - 1), false)
            }
        })
    }

    fn add_mismatch_utf8_states(
        &mut self,
        from_si: usize,
        lev_state: &[usize],
    ) -> Option<(usize, Vec<usize>)> {
        let mismatch_state = self.lev.accept(lev_state, None);
        let to_si = match self.cached(&mismatch_state) {
            None => return None,
            Some((si, _)) => si,
            // Some((si, true)) => return Some((si, mismatch_state)),
            // Some((si, false)) => si,
        };
        self.add_utf8_sequences(false, from_si, to_si, '\u{0}', '\u{10FFFF}');
        return Some((to_si, mismatch_state));
    }

    fn add_utf8_sequences(
        &mut self,
        overwrite: bool,
        from_si: usize,
        to_si: usize,
        from_chr: char,
        to_chr: char,
    ) {
        for seq in Utf8Sequences::new(from_chr, to_chr) {
            let mut fsi = from_si;
            for range in &seq.as_slice()[0..seq.len() - 1] {
                let tsi = self.new_state(false);
                self.add_utf8_range(overwrite, fsi, tsi, range);
                fsi = tsi;
            }
            self.add_utf8_range(
                overwrite,
                fsi,
                to_si,
                &seq.as_slice()[seq.len() - 1],
            );
        }
    }

    fn add_utf8_range(
        &mut self,
        overwrite: bool,
        from: usize,
        to: usize,
        range: &Utf8Range,
    ) {
        for b in range.start as usize..range.end as usize + 1 {
            if overwrite || self.dfa.states[from].next[b].is_none() {
                self.dfa.states[from].next[b] = Some(to);
            }
        }
    }

    fn new_state(&mut self, is_match: bool) -> usize {
        self.dfa.states.push(State { next: [None; 256], is_match });
        self.dfa.states.len() - 1
    }
}
<commit_msg>levenshtein: remove 'inline' annotations<commit_after>use std::cmp;
use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::fmt;

use utf8_ranges::{Utf8Range, Utf8Sequences};

use crate::automaton::Automaton;

const DEFAULT_STATE_LIMIT: usize = 10_000; // currently at least 20MB >_<

/// An error that occurred while building a Levenshtein automaton.
///
/// This error is only defined when the `levenshtein` crate feature is enabled.
#[derive(Debug)]
pub enum LevenshteinError {
    /// If construction of the automaton reaches some hard-coded limit
    /// on the number of states, then this error is returned.
    ///
    /// The number given is the limit that was exceeded.
    TooManyStates(usize),
}

impl fmt::Display for LevenshteinError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            LevenshteinError::TooManyStates(size_limit) => write!(
                f,
                "Levenshtein automaton exceeds size limit of \
                           {} states",
                size_limit
            ),
        }
    }
}

impl std::error::Error for LevenshteinError {}

/// A Unicode aware Levenshtein automaton for running efficient fuzzy queries.
///
/// This is only defined when the `levenshtein` crate feature is enabled.
///
/// A Levenshtein automata is one way to search any finite state transducer
/// for keys that *approximately* match a given query. A Levenshtein automaton
/// approximates this by returning all keys within a certain edit distance of
/// the query. The edit distance is defined by the number of insertions,
/// deletions and substitutions required to turn the query into the key.
/// Insertions, deletions and substitutions are based on
/// **Unicode characters** (where each character is a single Unicode scalar
/// value).
///
/// # Example
///
/// This example shows how to find all keys within an edit distance of `1`
/// from `foo`.
///
/// ```rust
/// use fst::automaton::Levenshtein;
/// use fst::{IntoStreamer, Streamer, Set};
///
/// fn main() {
///     let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
///     let set = Set::from_iter(keys).unwrap();
///
///     let lev = Levenshtein::new("foo", 1).unwrap();
///     let mut stream = set.search(&lev).into_stream();
///
///     let mut keys = vec![];
///     while let Some(key) = stream.next() {
///         keys.push(key.to_vec());
///     }
///     assert_eq!(keys, vec![
///         "fo".as_bytes(),   // 1 deletion
///         "fob".as_bytes(),  // 1 substitution
///         "foo".as_bytes(),  // 0 insertions/deletions/substitutions
///         "food".as_bytes(), // 1 insertion
///     ]);
/// }
/// ```
///
/// This example only uses ASCII characters, but it will work equally well
/// on Unicode characters.
///
/// # Warning: experimental
///
/// While executing this Levenshtein automaton against a finite state
/// transducer will be very fast, *constructing* an automaton may not be.
/// Namely, this implementation is a proof of concept. While I believe the
/// algorithmic complexity is not exponential, the implementation is not speedy
/// and it can use enormous amounts of memory (tens of MB before a hard-coded
/// limit will cause an error to be returned).
///
/// This is important functionality, so one should count on this implementation
/// being vastly improved in the future.
pub struct Levenshtein {
    prog: DynamicLevenshtein,
    dfa: Dfa,
}

impl Levenshtein {
    /// Create a new Levenshtein query.
    ///
    /// The query finds all matching terms that are at most `distance`
    /// edit operations from `query`. (An edit operation may be an insertion,
    /// a deletion or a substitution.)
    ///
    /// If the underlying automaton becomes too big, then an error is returned.
    /// Use `new_with_limit` to raise the limit dynamically.
    ///
    /// A `Levenshtein` value satisfies the `Automaton` trait, which means it
    /// can be used with the `search` method of any finite state transducer.
    pub fn new(
        query: &str,
        distance: u32,
    ) -> Result<Levenshtein, LevenshteinError> {
        let lev = DynamicLevenshtein {
            query: query.to_owned(),
            dist: distance as usize,
        };
        let dfa = DfaBuilder::new(lev.clone()).build()?;
        Ok(Levenshtein { prog: lev, dfa })
    }

    /// Create a new Levenshtein query, but pass the state limit yourself.
    ///
    /// The query finds all matching terms that are at most `distance`
    /// edit operations from `query`. (An edit operation may be an insertion,
    /// a deletion or a substitution.)
    ///
    /// If the underlying automaton becomes too big, then an error is returned.
    /// This limit can be configured with `state_limit`.
    ///
    /// A `Levenshtein` value satisfies the `Automaton` trait, which means it
    /// can be used with the `search` method of any finite state transducer.
    pub fn new_with_limit(
        query: &str,
        distance: u32,
        state_limit: usize,
    ) -> Result<Levenshtein, LevenshteinError> {
        let lev = DynamicLevenshtein {
            query: query.to_owned(),
            dist: distance as usize,
        };
        let dfa =
            DfaBuilder::new(lev.clone()).build_with_limit(state_limit)?;
        Ok(Levenshtein { prog: lev, dfa })
    }
}

impl fmt::Debug for Levenshtein {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Levenshtein(query: {:?}, distance: {:?})",
            self.prog.query, self.prog.dist
        )
    }
}

#[derive(Clone)]
struct DynamicLevenshtein {
    query: String,
    dist: usize,
}

impl DynamicLevenshtein {
    fn start(&self) -> Vec<usize> {
        (0..self.query.chars().count() + 1).collect()
    }

    fn is_match(&self, state: &[usize]) -> bool {
        state.last().map(|&n| n <= self.dist).unwrap_or(false)
    }

    fn can_match(&self, state: &[usize]) -> bool {
        state.iter().min().map(|&n| n <= self.dist).unwrap_or(false)
    }

    fn accept(&self, state: &[usize], chr: Option<char>) -> Vec<usize> {
        let mut next = vec![state[0] + 1];
        for (i, c) in self.query.chars().enumerate() {
            let cost = if Some(c) == chr { 0 } else { 1 };
            let v = cmp::min(
                cmp::min(next[i] + 1, state[i + 1] + 1),
                state[i] + cost,
            );
            next.push(cmp::min(v, self.dist + 1));
        }
        next
    }
}

impl Automaton for Levenshtein {
    type State = Option<usize>;

    #[inline]
    fn start(&self) -> Option<usize> {
        Some(0)
    }

    #[inline]
    fn is_match(&self, state: &Option<usize>) -> bool {
        state.map(|state| self.dfa.states[state].is_match).unwrap_or(false)
    }

    #[inline]
    fn can_match(&self, state: &Option<usize>) -> bool {
        state.is_some()
    }

    #[inline]
    fn accept(&self, state: &Option<usize>, byte: u8) -> Option<usize> {
        state.and_then(|state| self.dfa.states[state].next[byte as usize])
    }
}

#[derive(Debug)]
struct Dfa {
    states: Vec<State>,
}

struct State {
    next: [Option<usize>; 256],
    is_match: bool,
}

impl fmt::Debug for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "State {{")?;
        writeln!(f, "  is_match: {:?}", self.is_match)?;
        for i in 0..256 {
            if let Some(si) = self.next[i] {
                writeln!(f, "  {:?}: {:?}", i, si)?;
            }
        }
        write!(f, "}}")
    }
}

struct DfaBuilder {
    dfa: Dfa,
    lev: DynamicLevenshtein,
    cache: HashMap<Vec<usize>, usize>,
}

impl DfaBuilder {
    fn new(lev: DynamicLevenshtein) -> DfaBuilder {
        DfaBuilder {
            dfa: Dfa { states: Vec::with_capacity(16) },
            lev,
            cache: HashMap::with_capacity(1024),
        }
    }

    fn build_with_limit(
        mut self,
        state_limit: usize,
    ) -> Result<Dfa, LevenshteinError> {
        let mut stack = vec![self.lev.start()];
        let mut seen = HashSet::new();
        let query = self.lev.query.clone(); // temp work around of borrowck
        while let Some(lev_state) = stack.pop() {
            let dfa_si = self.cached_state(&lev_state).unwrap();
            let mismatch = self.add_mismatch_utf8_states(dfa_si, &lev_state);
            if let Some((next_si, lev_next)) = mismatch {
                if !seen.contains(&next_si) {
                    seen.insert(next_si);
                    stack.push(lev_next);
                }
            }
            for (i, c) in query.chars().enumerate() {
                if lev_state[i] > self.lev.dist {
                    continue;
                }
                let lev_next = self.lev.accept(&lev_state, Some(c));
                let next_si = self.cached_state(&lev_next);
                if let Some(next_si) = next_si {
                    self.add_utf8_sequences(true, dfa_si, next_si, c, c);
                    if !seen.contains(&next_si) {
                        seen.insert(next_si);
                        stack.push(lev_next);
                    }
                }
            }
            if self.dfa.states.len() > state_limit {
                return Err(LevenshteinError::TooManyStates(state_limit));
            }
        }
        Ok(self.dfa)
    }

    fn build(mut self) -> Result<Dfa, LevenshteinError> {
        self.build_with_limit(DEFAULT_STATE_LIMIT)
    }

    fn cached_state(&mut self, lev_state: &[usize]) -> Option<usize> {
        self.cached(lev_state).map(|(si, _)| si)
    }

    fn cached(&mut self, lev_state: &[usize]) -> Option<(usize, bool)> {
        if !self.lev.can_match(lev_state) {
            return None;
        }
        Some(match self.cache.entry(lev_state.to_vec()) {
            Entry::Occupied(v) => (*v.get(), true),
            Entry::Vacant(v) => {
                let is_match = self.lev.is_match(lev_state);
                self.dfa.states.push(State { next: [None; 256], is_match });
                (*v.insert(self.dfa.states.len() - 1), false)
            }
        })
    }

    fn add_mismatch_utf8_states(
        &mut self,
        from_si: usize,
        lev_state: &[usize],
    ) -> Option<(usize, Vec<usize>)> {
        let mismatch_state = self.lev.accept(lev_state, None);
        let to_si = match self.cached(&mismatch_state) {
            None => return None,
            Some((si, _)) => si,
            // Some((si, true)) => return Some((si, mismatch_state)),
            // Some((si, false)) => si,
        };
        self.add_utf8_sequences(false, from_si, to_si, '\u{0}', '\u{10FFFF}');
        return Some((to_si, mismatch_state));
    }

    fn add_utf8_sequences(
        &mut self,
        overwrite: bool,
        from_si: usize,
        to_si: usize,
        from_chr: char,
        to_chr: char,
    ) {
        for seq in Utf8Sequences::new(from_chr, to_chr) {
            let mut fsi = from_si;
            for range in &seq.as_slice()[0..seq.len() - 1] {
                let tsi = self.new_state(false);
                self.add_utf8_range(overwrite, fsi, tsi, range);
                fsi = tsi;
            }
            self.add_utf8_range(
                overwrite,
                fsi,
                to_si,
                &seq.as_slice()[seq.len() - 1],
            );
        }
    }

    fn add_utf8_range(
        &mut self,
        overwrite: bool,
        from: usize,
        to: usize,
        range: &Utf8Range,
    ) {
        for b in range.start as usize..range.end as usize + 1 {
            if overwrite || self.dfa.states[from].next[b].is_none() {
                self.dfa.states[from].next[b] = Some(to);
            }
        }
    }

    fn new_state(&mut self, is_match: bool) -> usize {
        self.dfa.states.push(State { next: [None; 256], is_match });
        self.dfa.states.len() - 1
    }
}
<|endoftext|>