<commit_before>use crate::ssz_encode;
use hashing::hash;
use int_to_bytes::int_to_bytes32;

const BYTES_PER_CHUNK: usize = 32;
const HASHSIZE: usize = 32;
const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;

pub struct TreeHashCache<'a> {
    chunk_offset: usize,
    cache: &'a mut [u8],
    chunk_modified: &'a mut [bool],
}

impl<'a> TreeHashCache<'a> {
    pub fn increment(&mut self) {
        self.chunk_offset += 1
    }

    pub fn modify_current_chunk(&mut self, to: &[u8]) -> Option<()> {
        self.modify_chunk(0, to)
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;
        self.cache.get_mut(start..end)?.copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Some(())
    }

    pub fn changed(&self, chunk: usize) -> Option<bool> {
        self.chunk_modified.get(chunk).cloned()
    }

    pub fn children_modified(&self, parent_chunk: usize) -> Option<bool> {
        let children = children(parent_chunk);

        Some(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, parent_chunk: usize) -> Option<Vec<u8>> {
        let children = children(parent_chunk);

        let start = children.0 * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK * 2;

        Some(hash(&self.cache.get(start..end)?))
    }

    pub fn just_the_leaves(&mut self, leaves: usize) -> Option<TreeHashCache> {
        let nodes = num_nodes(leaves);
        let internal = nodes - leaves;
        let leaves_start = (self.chunk_offset + internal) * HASHSIZE;

        Some(TreeHashCache {
            chunk_offset: self.chunk_offset + internal,
            cache: self.cache.get_mut(leaves_start..leaves * HASHSIZE)?,
            chunk_modified: self
                .chunk_modified
                .get_mut(self.chunk_offset..self.chunk_offset + leaves)?,
        })
    }
}

fn children(parent: usize) -> (usize, usize) {
    ((2 * parent + 1), (2 * parent + 2))
}

fn num_nodes(num_leaves: usize) -> usize {
    2 * num_leaves - 1
}

pub trait CachedTreeHash {
    fn build_cache_bytes(&self) -> Vec<u8>;

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>;
}

impl CachedTreeHash for u64 {
    fn build_cache_bytes(&self) -> Vec<u8> {
        merkleize(&ssz_encode(self))
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        if self != other {
            cache.modify_current_chunk(&hash(&ssz_encode(self)));
        }

        cache.increment();

        Some(())
    }
}

pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(&leaves)
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        let num_leaves = 4;

        let mut leaf_cache = cache.just_the_leaves(num_leaves)?;
        self.a.cached_hash_tree_root(&other.a, &mut leaf_cache)?;
        self.b.cached_hash_tree_root(&other.b, &mut leaf_cache)?;
        self.c.cached_hash_tree_root(&other.c, &mut leaf_cache)?;
        self.d.cached_hash_tree_root(&other.d, &mut leaf_cache)?;

        let nodes = num_nodes(num_leaves);
        let internal_chunks = nodes - num_leaves;

        for chunk in 0..internal_chunks {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(())
    }
}

/// A reference function to test against.
pub fn merkleize(values: &[u8]) -> Vec<u8> {
    let leaves = values.len() / HASHSIZE;

    if leaves == 0 || !leaves.is_power_of_two() {
        panic!("Handle bad leaf count");
    }

    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
    o.append(&mut values.to_vec());

    let mut i = o.len();
    let mut j = o.len() - values.len();

    while i >= MERKLE_HASH_CHUNCK {
        i -= MERKLE_HASH_CHUNCK;
        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);

        j -= HASHSIZE;
        o[j..j + HASHSIZE].copy_from_slice(&hash);
    }

    o
}

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

    fn join(many: Vec<&[u8]>) -> Vec<u8> {
        let mut all = vec![];
        for one in many {
            all.extend_from_slice(&mut one.clone())
        }
        all
    }

    #[test]
    fn merkleize_4_leaves() {
        let data1 = hash(&int_to_bytes32(1));
        let data2 = hash(&int_to_bytes32(2));
        let data3 = hash(&int_to_bytes32(3));
        let data4 = hash(&int_to_bytes32(4));

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = merkleize(&data);

        let hash_12 = {
            let mut joined = vec![];
            joined.append(&mut data1.clone());
            joined.append(&mut data2.clone());
            hash(&joined)
        };
        let hash_34 = {
            let mut joined = vec![];
            joined.append(&mut data3.clone());
            joined.append(&mut data4.clone());
            hash(&joined)
        };
        let hash_hash12_hash_34 = {
            let mut joined = vec![];
            joined.append(&mut hash_12.clone());
            joined.append(&mut hash_34.clone());
            hash(&joined)
        };

        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
            let expected = match i {
                0 => hash_hash12_hash_34.clone(),
                1 => hash_12.clone(),
                2 => hash_34.clone(),
                3 => data1.clone(),
                4 => data2.clone(),
                5 => data3.clone(),
                6 => data4.clone(),
                _ => vec![],
            };

            assert_eq!(chunk, &expected[..], "failed at {}", i);
        }
    }
}
<commit_msg>Implement (failing) cached tree hash tests<commit_after>use crate::ssz_encode;
use hashing::hash;
use int_to_bytes::int_to_bytes32;

const BYTES_PER_CHUNK: usize = 32;
const HASHSIZE: usize = 32;
const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;

pub struct TreeHashCache<'a> {
    chunk_offset: usize,
    cache: &'a mut [u8],
    chunk_modified: &'a mut [bool],
}

impl<'a> TreeHashCache<'a> {
    pub fn build_changes_vec(bytes: &[u8]) -> Vec<bool> {
        vec![false; bytes.len() / BYTES_PER_CHUNK]
    }

    pub fn from_mut_slice(bytes: &'a mut [u8], changes: &'a mut [bool]) -> Option<Self> {
        if bytes.len() % BYTES_PER_CHUNK > 0 {
            return None;
        }

        let chunk_modified = vec![false; bytes.len() / BYTES_PER_CHUNK];

        Some(Self {
            chunk_offset: 0,
            cache: bytes,
            chunk_modified: changes,
        })
    }

    pub fn increment(&mut self) {
        self.chunk_offset += 1
    }

    pub fn modify_current_chunk(&mut self, to: &[u8]) -> Option<()> {
        self.modify_chunk(0, to)
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;
        self.cache.get_mut(start..end)?.copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Some(())
    }

    pub fn changed(&self, chunk: usize) -> Option<bool> {
        self.chunk_modified.get(chunk).cloned()
    }

    pub fn children_modified(&self, parent_chunk: usize) -> Option<bool> {
        let children = children(parent_chunk);

        Some(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, parent_chunk: usize) -> Option<Vec<u8>> {
        let children = children(parent_chunk);

        let start = children.0 * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK * 2;

        Some(hash(&self.cache.get(start..end)?))
    }

    pub fn just_the_leaves(&mut self, leaves: usize) -> Option<TreeHashCache> {
        let nodes = num_nodes(leaves);
        let internal = nodes - leaves;
        let leaves_start = (self.chunk_offset + internal) * HASHSIZE;

        Some(TreeHashCache {
            chunk_offset: self.chunk_offset + internal,
            cache: self.cache.get_mut(leaves_start..leaves * HASHSIZE)?,
            chunk_modified: self
                .chunk_modified
                .get_mut(self.chunk_offset..self.chunk_offset + leaves)?,
        })
    }

    pub fn into_slice(self) -> &'a [u8] {
        self.cache
    }
}

fn children(parent: usize) -> (usize, usize) {
    ((2 * parent + 1), (2 * parent + 2))
}

fn num_nodes(num_leaves: usize) -> usize {
    2 * num_leaves - 1
}

pub trait CachedTreeHash {
    fn build_cache_bytes(&self) -> Vec<u8>;

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>;
}

impl CachedTreeHash for u64 {
    fn build_cache_bytes(&self) -> Vec<u8> {
        merkleize(&int_to_bytes32(*self))
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        if self != other {
            cache.modify_current_chunk(&hash(&ssz_encode(self)));
        }

        cache.increment();

        Some(())
    }
}

#[derive(Clone)]
pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(&leaves)
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        let num_leaves = 4;

        let mut leaf_cache = cache.just_the_leaves(num_leaves)?;
        self.a.cached_hash_tree_root(&other.a, &mut leaf_cache)?;
        self.b.cached_hash_tree_root(&other.b, &mut leaf_cache)?;
        self.c.cached_hash_tree_root(&other.c, &mut leaf_cache)?;
        self.d.cached_hash_tree_root(&other.d, &mut leaf_cache)?;

        dbg!(leaf_cache.into_slice());

        let nodes = num_nodes(num_leaves);
        let internal_chunks = nodes - num_leaves;

        for chunk in 0..internal_chunks {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(())
    }
}

/// A reference function to test against.
pub fn merkleize(values: &[u8]) -> Vec<u8> {
    let leaves = values.len() / HASHSIZE;

    if leaves == 0 {
        panic!("No full leaves");
    }

    if !leaves.is_power_of_two() {
        panic!("leaves is not power of two");
    }

    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
    o.append(&mut values.to_vec());

    let mut i = o.len();
    let mut j = o.len() - values.len();

    while i >= MERKLE_HASH_CHUNCK {
        i -= MERKLE_HASH_CHUNCK;
        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);

        j -= HASHSIZE;
        o[j..j + HASHSIZE].copy_from_slice(&hash);
    }

    o
}

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

    fn join(many: Vec<&[u8]>) -> Vec<u8> {
        let mut all = vec![];
        for one in many {
            all.extend_from_slice(&mut one.clone())
        }
        all
    }

    #[test]
    fn cached_hash_on_inner() {
        let inner = Inner {
            a: 1,
            b: 2,
            c: 3,
            d: 4,
        };

        let mut cache = inner.build_cache_bytes();

        let changed_inner = Inner {
            a: 42,
            ..inner.clone()
        };

        let mut changes = TreeHashCache::build_changes_vec(&cache);
        let mut cache_struct = TreeHashCache::from_mut_slice(&mut cache, &mut changes).unwrap();

        changed_inner.cached_hash_tree_root(&inner, &mut cache_struct);

        let new_cache = cache_struct.into_slice();

        let data1 = &int_to_bytes32(42);
        let data2 = &int_to_bytes32(2);
        let data3 = &int_to_bytes32(3);
        let data4 = &int_to_bytes32(4);

        let data = join(vec![&data1, &data2, &data3, &data4]);
        let expected = merkleize(&data);

        assert_eq!(expected, new_cache);
    }

    #[test]
    fn build_cache_matches_merkelize() {
        let data1 = &int_to_bytes32(1);
        let data2 = &int_to_bytes32(2);
        let data3 = &int_to_bytes32(3);
        let data4 = &int_to_bytes32(4);

        let data = join(vec![&data1, &data2, &data3, &data4]);
        let expected = merkleize(&data);

        let inner = Inner {
            a: 1,
            b: 2,
            c: 3,
            d: 4,
        };

        let cache = inner.build_cache_bytes();

        assert_eq!(expected, cache);
    }

    #[test]
    fn merkleize_4_leaves() {
        let data1 = hash(&int_to_bytes32(1));
        let data2 = hash(&int_to_bytes32(2));
        let data3 = hash(&int_to_bytes32(3));
        let data4 = hash(&int_to_bytes32(4));

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = merkleize(&data);

        let hash_12 = {
            let mut joined = vec![];
            joined.append(&mut data1.clone());
            joined.append(&mut data2.clone());
            hash(&joined)
        };
        let hash_34 = {
            let mut joined = vec![];
            joined.append(&mut data3.clone());
            joined.append(&mut data4.clone());
            hash(&joined)
        };
        let hash_hash12_hash_34 = {
            let mut joined = vec![];
            joined.append(&mut hash_12.clone());
            joined.append(&mut hash_34.clone());
            hash(&joined)
        };

        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
            let expected = match i {
                0 => hash_hash12_hash_34.clone(),
                1 => hash_12.clone(),
                2 => hash_34.clone(),
                3 => data1.clone(),
                4 => data2.clone(),
                5 => data3.clone(),
                6 => data4.clone(),
                _ => vec![],
            };

            assert_eq!(chunk, &expected[..], "failed at {}", i);
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate varlink;

use org_example_ping::*;
use std::env;
use std::io;
use std::process::exit;

mod org_example_ping;

fn run_app(address: String) -> io::Result<()> {
    let connection = varlink::Connection::new(&address)?;
    let call = VarlinkClient::new(connection);
    let ping: Option<String> = Some("Test".into());
    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());
    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());
    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());
    Ok(())
}

fn main() {
    let mut args: Vec<_> = env::args().collect();
    match args.len() {
        2 => {}
        _ => {
            eprintln!("Usage: {} <varlink address>", args[0]);
            exit(1);
        }
    };

    exit(match run_app(args.swap_remove(1)) {
        Ok(_) => 0,
        Err(err) => {
            eprintln!("error: {}", err);
            1
        }
    });
}
<commit_msg>beautify<commit_after>#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate varlink;

use org_example_ping::*;
use std::env;
use std::io;
use std::process::exit;

mod org_example_ping;

fn run_app(address: String) -> io::Result<()> {
    let connection = varlink::Connection::new(&address)?;
    let call = VarlinkClient::new(connection);
    let ping: Option<String> = Some("Test".into());

    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());

    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());

    let reply = call.more().ping(ping.clone())?.recv()?;
    assert_eq!(ping, reply.pong);
    println!("Pong: '{}'", reply.pong.unwrap());

    Ok(())
}

fn main() {
    let mut args: Vec<_> = env::args().collect();
    match args.len() {
        2 => {}
        _ => {
            eprintln!("Usage: {} <varlink address>", args[0]);
            exit(1);
        }
    };

    exit(match run_app(args.swap_remove(1)) {
        Ok(_) => 0,
        Err(err) => {
            eprintln!("error: {}", err);
            1
        }
    });
}
<|endoftext|>


<commit_before>// Copyright 2016 Alex Regueiro
//
// 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.
//

use ffi;
use {Error, DB};

use libc::{c_int, uint32_t};
use std::ffi::CString;
use std::path::Path;

pub struct BackupEngine {
    inner: *mut ffi::rocksdb_backup_engine_t,
}

pub struct BackupEngineOptions {
    inner: *mut ffi::rocksdb_options_t,
}

pub struct RestoreOptions {
    inner: *mut ffi::rocksdb_restore_options_t,
}

impl BackupEngine {
    /// Open a backup engine with the specified options.
    pub fn open<P: AsRef<Path>>(
        opts: &BackupEngineOptions,
        path: P,
    ) -> Result<BackupEngine, Error> {
        let path = path.as_ref();
        let cpath = match CString::new(path.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert path to CString \
                     when opening backup engine"
                        .to_owned(),
                ));
            }
        };

        let be: *mut ffi::rocksdb_backup_engine_t;
        unsafe { be = ffi_try!(ffi::rocksdb_backup_engine_open(opts.inner, cpath.as_ptr(),)) }

        if be.is_null() {
            return Err(Error::new("Could not initialize backup engine.".to_owned()));
        }

        Ok(BackupEngine { inner: be })
    }

    pub fn create_new_backup(&mut self, db: &DB) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_create_new_backup(
                self.inner, db.inner,
            ));
            Ok(())
        }
    }

    pub fn purge_old_backups(&mut self, num_backups_to_keep: usize) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_purge_old_backups(
                self.inner,
                num_backups_to_keep as uint32_t,
            ));
            Ok(())
        }
    }

    /// Restore from the latest backup
    ///
    /// # Arguments
    ///
    /// * `db_dir` - A path to the database directory
    /// * `wal_dir` - A path to the wal directory
    /// * `opts` - Restore options
    ///
    /// # Example
    ///
    /// ```ignore
    /// use rocksdb::backup::{BackupEngine, BackupEngineOptions};
    /// let backup_opts = BackupEngineOptions::default();
    /// let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();
    /// let mut restore_option = rocksdb::backup::RestoreOptions::default();
    /// restore_option.set_keep_log_files(true); /// true to keep log files
    /// if let Err(e) = backup_engine.restore_from_latest_backup(&db_path, &wal_dir, &restore_option) {
    ///     error!("Failed to restore from the backup. Error:{:?}", e);
    ///     return Err(e.to_string());
    ///  }
    /// ```

    pub fn restore_from_latest_backup<D: AsRef<Path>, W: AsRef<Path>>(
        &mut self,
        db_dir: D,
        wal_dir: W,
        opts: &RestoreOptions,
    ) -> Result<(), Error> {
        let db_dir = db_dir.as_ref();
        let c_db_dir = match CString::new(db_dir.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert db_dir to CString \
                     when restoring from latest backup"
                        .to_owned(),
                ));
            }
        };

        let wal_dir = wal_dir.as_ref();
        let c_wal_dir = match CString::new(wal_dir.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert wal_dir to CString \
                     when restoring from latest backup"
                        .to_owned(),
                ));
            }
        };

        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_restore_db_from_latest_backup(
                self.inner,
                c_db_dir.as_ptr(),
                c_wal_dir.as_ptr(),
                opts.inner,
            ));
        }
        Ok(())
    }
}

impl BackupEngineOptions {
    //
}

impl RestoreOptions {
    pub fn set_keep_log_files(&mut self, keep_log_files: bool) {
        unsafe {
            ffi::rocksdb_restore_options_set_keep_log_files(self.inner, keep_log_files as c_int);
        }
    }
}

impl Default for BackupEngineOptions {
    fn default() -> BackupEngineOptions {
        unsafe {
            let opts = ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB backup options".to_owned());
            }
            BackupEngineOptions { inner: opts }
        }
    }
}

impl Default for RestoreOptions {
    fn default() -> RestoreOptions {
        unsafe {
            let opts = ffi::rocksdb_restore_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB restore options".to_owned());
            }
            RestoreOptions { inner: opts }
        }
    }
}

impl Drop for BackupEngine {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_backup_engine_close(self.inner);
        }
    }
}

impl Drop for BackupEngineOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for RestoreOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_restore_options_destroy(self.inner);
        }
    }
}

#[test]
fn backup_restore() {
    use db::DBVector;
    // create backup
    let path = "_rust_rocksdb_backup_restore_test";
    {
        let db = DB::open_default(path).unwrap();
        let p = db.put(b"k1", b"v1111");
        assert!(p.is_ok());
        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");

        let backup_path = "_rust_rocksdb_backup_path";
        {
            let backup_opts = BackupEngineOptions::default();
            let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();

            let r = backup_engine.create_new_backup(&db);
            assert!(r.is_ok());

            let restore_path = "_rust_rocksdb_restore_from_backup_path";
            {
                let mut restore_option = RestoreOptions::default();
                restore_option.set_keep_log_files(true); // true to keep log files
                let restore_status = backup_engine.restore_from_latest_backup(
                    &restore_path,
                    &restore_path,
                    &restore_option,
                );
                assert!(restore_status.is_ok());

                let db = DB::open_default(restore_path).unwrap();

                let r: Result<Option<DBVector>, Error> = db.get(b"k1");
                assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");
            }
        }
    }
}
<commit_msg>adding DB::destroy()<commit_after>// Copyright 2016 Alex Regueiro
//
// 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.
//

use ffi;
use {Error, DB};

use libc::{c_int, uint32_t};
use std::ffi::CString;
use std::path::Path;

pub struct BackupEngine {
    inner: *mut ffi::rocksdb_backup_engine_t,
}

pub struct BackupEngineOptions {
    inner: *mut ffi::rocksdb_options_t,
}

pub struct RestoreOptions {
    inner: *mut ffi::rocksdb_restore_options_t,
}

impl BackupEngine {
    /// Open a backup engine with the specified options.
    pub fn open<P: AsRef<Path>>(
        opts: &BackupEngineOptions,
        path: P,
    ) -> Result<BackupEngine, Error> {
        let path = path.as_ref();
        let cpath = match CString::new(path.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert path to CString \
                     when opening backup engine"
                        .to_owned(),
                ));
            }
        };

        let be: *mut ffi::rocksdb_backup_engine_t;
        unsafe { be = ffi_try!(ffi::rocksdb_backup_engine_open(opts.inner, cpath.as_ptr(),)) }

        if be.is_null() {
            return Err(Error::new("Could not initialize backup engine.".to_owned()));
        }

        Ok(BackupEngine { inner: be })
    }

    pub fn create_new_backup(&mut self, db: &DB) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_create_new_backup(
                self.inner, db.inner,
            ));
            Ok(())
        }
    }

    pub fn purge_old_backups(&mut self, num_backups_to_keep: usize) -> Result<(), Error> {
        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_purge_old_backups(
                self.inner,
                num_backups_to_keep as uint32_t,
            ));
            Ok(())
        }
    }

    /// Restore from the latest backup
    ///
    /// # Arguments
    ///
    /// * `db_dir` - A path to the database directory
    /// * `wal_dir` - A path to the wal directory
    /// * `opts` - Restore options
    ///
    /// # Example
    ///
    /// ```ignore
    /// use rocksdb::backup::{BackupEngine, BackupEngineOptions};
    /// let backup_opts = BackupEngineOptions::default();
    /// let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();
    /// let mut restore_option = rocksdb::backup::RestoreOptions::default();
    /// restore_option.set_keep_log_files(true); /// true to keep log files
    /// if let Err(e) = backup_engine.restore_from_latest_backup(&db_path, &wal_dir, &restore_option) {
    ///     error!("Failed to restore from the backup. Error:{:?}", e);
    ///     return Err(e.to_string());
    ///  }
    /// ```

    pub fn restore_from_latest_backup<D: AsRef<Path>, W: AsRef<Path>>(
        &mut self,
        db_dir: D,
        wal_dir: W,
        opts: &RestoreOptions,
    ) -> Result<(), Error> {
        let db_dir = db_dir.as_ref();
        let c_db_dir = match CString::new(db_dir.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert db_dir to CString \
                     when restoring from latest backup"
                        .to_owned(),
                ));
            }
        };

        let wal_dir = wal_dir.as_ref();
        let c_wal_dir = match CString::new(wal_dir.to_string_lossy().as_bytes()) {
            Ok(c) => c,
            Err(_) => {
                return Err(Error::new(
                    "Failed to convert wal_dir to CString \
                     when restoring from latest backup"
                        .to_owned(),
                ));
            }
        };

        unsafe {
            ffi_try!(ffi::rocksdb_backup_engine_restore_db_from_latest_backup(
                self.inner,
                c_db_dir.as_ptr(),
                c_wal_dir.as_ptr(),
                opts.inner,
            ));
        }
        Ok(())
    }
}

impl BackupEngineOptions {
    //
}

impl RestoreOptions {
    pub fn set_keep_log_files(&mut self, keep_log_files: bool) {
        unsafe {
            ffi::rocksdb_restore_options_set_keep_log_files(self.inner, keep_log_files as c_int);
        }
    }
}

impl Default for BackupEngineOptions {
    fn default() -> BackupEngineOptions {
        unsafe {
            let opts = ffi::rocksdb_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB backup options".to_owned());
            }
            BackupEngineOptions { inner: opts }
        }
    }
}

impl Default for RestoreOptions {
    fn default() -> RestoreOptions {
        unsafe {
            let opts = ffi::rocksdb_restore_options_create();
            if opts.is_null() {
                panic!("Could not create RocksDB restore options".to_owned());
            }
            RestoreOptions { inner: opts }
        }
    }
}

impl Drop for BackupEngine {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_backup_engine_close(self.inner);
        }
    }
}

impl Drop for BackupEngineOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_options_destroy(self.inner);
        }
    }
}

impl Drop for RestoreOptions {
    fn drop(&mut self) {
        unsafe {
            ffi::rocksdb_restore_options_destroy(self.inner);
        }
    }
}

#[test]
fn backup_restore() {
    use db::DBVector;
    use Options;
    // create backup
    let path = "_rust_rocksdb_backup_restore_test";
    {
        let db = DB::open_default(path).unwrap();
        let p = db.put(b"k1", b"v1111");
        assert!(p.is_ok());
        let r: Result<Option<DBVector>, Error> = db.get(b"k1");
        assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");

        let backup_path = "_rust_rocksdb_backup_path";
        {
            let backup_opts = BackupEngineOptions::default();
            let mut backup_engine = BackupEngine::open(&backup_opts, &backup_path).unwrap();

            let r = backup_engine.create_new_backup(&db);
            assert!(r.is_ok());

            let restore_path = "_rust_rocksdb_restore_from_backup_path";
            {
                let mut restore_option = RestoreOptions::default();
                restore_option.set_keep_log_files(true); // true to keep log files
                let restore_status = backup_engine.restore_from_latest_backup(
                    &restore_path,
                    &restore_path,
                    &restore_option,
                );
                assert!(restore_status.is_ok());

                let db_restore = DB::open_default(restore_path).unwrap();

                let r: Result<Option<DBVector>, Error> = db_restore.get(b"k1");
                assert!(r.unwrap().unwrap().to_utf8().unwrap() == "v1111");

            }
            assert!(DB::destroy(&Options::default(), restore_path).is_ok());

        }
        assert!(DB::destroy(&Options::default(), backup_path).is_ok());

    }
    assert!(DB::destroy(&Options::default(), path).is_ok());

}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

//!
//! # Matcher
//!
//! The Matcher trait is implemented by objects that can match patterns against the left-hand side of a stream. It's a fairly
//! baseline implementation: it's up to the caller to implement things like rewinding in order to perform tokenisation. That is,
//! Matchers are greedy and may (indeed, are likely to) consume more characters than the longest match while trying to find
//! a longer one.
//!

///
/// Matcher that can read an input stream of type `Symbol` and find the longest matching pattern, which it will identify using
/// `OutputSymbol`
///
pub trait PatternMatcher<InputSymbol, OutputSymbol> {
}
<commit_msg>Add trait for matchers to use as a state machine<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

//!
//! # Matcher
//!
//! The Matcher trait is implemented by objects that can match patterns against the left-hand side of a stream. It's a fairly
//! baseline implementation: it's up to the caller to implement things like rewinding in order to perform tokenisation. That is,
//! Matchers are greedy and may (indeed, are likely to) consume more characters than the longest match while trying to find
//! a longer one.
//!

///
/// Matcher that can read an input stream of type `Symbol` and find the longest matching pattern, which it will identify using
/// `OutputSymbol`
///
pub trait PatternMatcher<InputSymbol, OutputSymbol> {
    // type State: MatchingState<InputSymbol, OutputSymbol>;

    //
    // Creates a state that begins matching this pattern
    //
    // fn start(&self) -> State;
}

///
/// Action to be taken after a matcher receives a symbol
///
pub enum MatchAction<OutputSymbol: ?Sized, State: ?Sized> {
    // State is also always: MatchingState<InputSymbol, OutputSymbol> (important to know that as its how More is used)
    //
    // However, rust complains that InputSymbol is unused if we declare it in MatchAction and that it is undeclared if we don't
    // If this wasn't an enum, could make it an associated type? But this is an enum so we can't do that.
    // Ie, MatchAction<InputSymbol, OutputSymbol, State: MatchingState<InputSymbol, OutputSymbol>> = Nope (InputSymbol is unused apparently even though it's required which seems to me to mean it's used)
    // and also MatchAction<OutputSymbol, State: MatchingState<DontCare, OutputSymbol> = Nope (DontCare is not declared)

    /// The pattern does not match
    Reject,

    /// The pattern matched a certain number of symbols (which may be fewer than were passed to the matcher)
    Accept(usize, OutputSymbol),

    /// The matcher needs more symbols to decide if the pattern matches
    More(State)
}

///
/// Represents a state during a pattern matching operation
///
pub trait MatchingState<InputSymbol, OutputSymbol: ?Sized> {
    ///
    /// Matches the next symbol
    ///
    fn next(self, symbol: InputSymbol) -> MatchAction<OutputSymbol, Self>;

    ///
    /// There are no more symbols available (this can only return `Reject` or `Accept`)
    ///
    fn finish(self) -> MatchAction<OutputSymbol, Self>;
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;
use std::collections::BTreeSet;
use std::str::FromStr;
use std::fmt;
use twiddle::Twiddle;
use enum_primitive::FromPrimitive;

use board::Board;
use board::PieceIter;
use board::PieceCount;
use board::board_from_str;
use board::str_from_board;
use piece::Piece;
use piece::Player;
use point::Point;

pub fn advance_piece_iterator(spot: &mut u16, extra: &mut [u16; 7])
                              -> Option<Piece> {
    for i in 7..0 {
        if extra[i].bits(15..11) != 0 {
            let piece = Piece::from_u8((((extra[i] & 1) << 2) | 1) as u8).unwrap();
            extra[i] = (extra[i] & 0xF800) | extra[i].bits(9..2);
            // We've shifted out everything from this continuation
            if extra[i] & 3 == 0 {
                extra[i] = 0;
            }
            return Some(piece);
        }
    }
    // Now the standard shift-out
    let temp = *spot;
    if temp & 3 != 0 {
        // Lower pieces
        let piece = Piece::from_u8((((temp & 1) << 2) | 1) as u8).unwrap();
        *spot = (temp & 0xF800) | temp.bits(11..2);
        Some(piece)
    } else if *spot != 0 {
        // Top piece
        let piece = Piece::from_u8(temp.bits(15..13) as u8).unwrap();
        *spot = 0;
        Some(piece)
    } else {
        None
    }
}

#[derive(Copy, Clone, Debug, RustcDecodable, RustcEncodable)]
pub struct Board5 {
    grid: [ [u16; 5]; 5],
    continuations: [u16; 7],
}

fn top_piece_bits(spot: u16) -> u8 {
    spot.bits(15..13) as u8
}

impl Board for Board5 {
    fn new(board_size: usize) -> Board5 {
        assert!(board_size == 5);
        Board5 {
            grid: [[0; 5]; 5],
            continuations: [0; 7],
        }
    }

    fn at(&self, point: &Point) -> Result<PieceIter, &str> {
        let mut extra = self.continuations;
        let location = ((point.x * 5 + point.y) as u16) << 11;
        for i in 0..7 {
            if (extra[i] & location) != location {
                extra[i] = 0;
            } else if extra[i] != 0 {
                while (extra[i] & 3) == 0 {
                    extra[i] = (extra[i] & location) | extra[i].bits(9..2);
                }
            }
        }

        let mut spot = self.grid[point.x][point.y];
        if (spot & 0x0FFF) != 0 {
            while (spot & 3) == 0 {
                spot = (spot & 0xE000) | spot.bits(11..2);
            }
        }

        Ok(PieceIter::Board5Iter {
            spot: spot,
            extra: extra,
        })
    }

    fn at_reset(&mut self, point: &Point) -> Result<PieceIter, &str> {
        let mut extra = self.continuations;
        let location = ((point.x * 5 + point.y) as u16) << 11;
        for i in 0..7 {
            if extra[i] & location != location {
                extra[i] = 0;
            } else if extra[i] != 0 {
                self.continuations[i] = 0;
                while extra[i] & 3 == 0 {
                    extra[i] = (extra[i] & location) | extra[i].bits(9..2);
                }
            }
        }

        let mut spot = self.grid[point.x][point.y];
        self.grid[point.x][point.y] = 0;
        if (spot & 0x0FFF) != 0 {
            while spot & 3 == 0 {
                spot = (spot & 0xE000) | spot.bits(11..2);
            }
        }

        Ok(PieceIter::Board5Iter {
            spot: spot,
            extra: extra,
        })
    }

    fn size(&self) -> usize {
        5
    }

    /// Checks to see if all spaces have at least one piece
    fn full(&self) -> bool {
        for row in self.grid.iter() {
            for square in row.iter() {
                if square.bits(15..13) == 0 {
                    return false
                }
            }
        }
        true
    }

    fn count(&self) -> PieceCount {
        let mut pieces = PieceCount::new(self.size());

        for row in self.grid.iter() {
            for spot in row.iter() {
                match Piece::from_u8(spot.bits(15..13) as u8) {
                    Some(Piece::OneFlat) => pieces.p1_flat += 1,
                    Some(Piece::OneStanding) => pieces.p1_flat += 1,
                    Some(Piece::OneCapstone) => pieces.p1_cap += 1,
                    Some(Piece::TwoFlat) => pieces.p2_flat += 1,
                    Some(Piece::TwoStanding) => pieces.p2_flat += 1,
                    Some(Piece::TwoCapstone) => pieces.p2_cap += 1,
                    None => {},
                }

                for i in 0..6 {
                    match spot.bits((2 * i + 1)..(2 * i)) {
                        1 => pieces.p1_flat += 1,
                        2 => pieces.p2_flat += 1,
                        _ => {},
                    }
                }
            }
        }
        //TODO: Count continuations
        pieces
    }

    fn place_piece(&mut self, point: &Point, piece: Piece) -> Result<(), String> {
        if top_piece_bits(self.grid[point.x][point.y]) == 0 {
            self.grid[point.x][point.y] = (piece as u16) << 13;
            Ok(())
        } else {
            return Err("Cannot place stone on top of existing stone.".into());
        }
    }

    fn add_piece(&mut self, point: &Point, piece: Piece) -> Result<(), String> {
        let spot = self.grid[point.x][point.y];
        let top_bits = top_piece_bits(spot);
        if top_bits == 0 {
            self.grid[point.x][point.y] = (piece as u16) << 13;
            return Ok(());
        }
        let mut top = Piece::from_u8(top_bits).unwrap();
        try!(piece.move_onto(&mut top));

        let mut lowest = spot.bits(1..0);
        let mut temp = spot.bits(11..2);
        if top.owner() == Player::One {
            temp = temp | 0x400;
        } else {
            temp = temp | 0x800;
        }
        self.grid[point.x][point.y] = temp | ((piece as u16) << 13);
        if lowest == 0 {
            Ok(())
        } else {
            // Need to add the bumped piece to an existing continuation
            for i in 0..7 {
                let cont = self.continuations[i];
                if cont.bits(15..11) == (point.x * 5 + point.y) as u16 {
                    let mut temp: u16 = cont & 0xF800; // Keep location
                    temp |= cont.bits(9..2); // Shift everything thats there
                    temp |= lowest << 9; // Add in the bumped piece
                    self.continuations[i] = temp;

                    lowest = cont & 3; // We may have to do this again...
                    if lowest == 0 {
                        return Ok(());
                    }
                }
            }

            // If we haven't returned yet, we need to start a new continuation
            for i in 0..7 {
                let cont = self.continuations[i];
                if cont.bits(15..11) == 0 {
                    // It's unclaimed - let's claim it by setting the location
                    let mut temp: u16 = ((point.x * 5 + point.y) as u16) << 11;
                    temp |= lowest << 9; // Add in the bumped piece
                    self.continuations[i] = temp;
                    return Ok(());
                }
            }

            // Need to figure out continuations
            Err("No available continuations. Did you play too many pieces?".into())
        }
    }
}

impl FromStr for Board5 {
    type Err=();
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        board_from_str::<Board5>(s)
    }
}

impl fmt::Display for Board5 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        str_from_board(self, f)
    }
}
<commit_msg>Use helper to determine location from point<commit_after>use std::collections::VecDeque;
use std::collections::BTreeSet;
use std::str::FromStr;
use std::fmt;
use twiddle::Twiddle;
use enum_primitive::FromPrimitive;

use board::Board;
use board::PieceIter;
use board::PieceCount;
use board::board_from_str;
use board::str_from_board;
use piece::Piece;
use piece::Player;
use point::Point;

pub fn advance_piece_iterator(spot: &mut u16, extra: &mut [u16; 7])
                              -> Option<Piece> {
    for i in 7..0 {
        if extra[i].bits(15..11) != 0 {
            let piece = Piece::from_u8((((extra[i] & 1) << 2) | 1) as u8).unwrap();
            extra[i] = (extra[i] & 0xF800) | extra[i].bits(9..2);
            // We've shifted out everything from this continuation
            if extra[i] & 3 == 0 {
                extra[i] = 0;
            }
            return Some(piece);
        }
    }
    // Now the standard shift-out
    let temp = *spot;
    if temp & 3 != 0 {
        // Lower pieces
        let piece = Piece::from_u8((((temp & 1) << 2) | 1) as u8).unwrap();
        *spot = (temp & 0xF800) | temp.bits(11..2);
        Some(piece)
    } else if *spot != 0 {
        // Top piece
        let piece = Piece::from_u8(temp.bits(15..13) as u8).unwrap();
        *spot = 0;
        Some(piece)
    } else {
        None
    }
}

#[derive(Copy, Clone, Debug, RustcDecodable, RustcEncodable)]
pub struct Board5 {
    grid: [ [u16; 5]; 5],
    continuations: [u16; 7],
}

fn top_piece_bits(spot: u16) -> u8 {
    spot.bits(15..13) as u8
}

impl Board5 {
    // Return shifted location value
    fn location(p: &Point) -> u16 {
        ((p.x * 5 + p.y + 1) as u16) << 11
    }
}

impl Board for Board5 {
    fn new(board_size: usize) -> Board5 {
        assert!(board_size == 5);
        Board5 {
            grid: [[0; 5]; 5],
            continuations: [0; 7],
        }
    }

    fn at(&self, point: &Point) -> Result<PieceIter, &str> {
        let mut extra = self.continuations;
        let location = Board5::location(point);
        for i in 0..7 {
            if (extra[i] & location) != location {
                extra[i] = 0;
            } else if extra[i] != 0 {
                while (extra[i] & 3) == 0 {
                    extra[i] = (extra[i] & location) | extra[i].bits(9..2);
                }
            }
        }

        let mut spot = self.grid[point.x][point.y];
        if (spot & 0x0FFF) != 0 {
            while (spot & 3) == 0 {
                spot = (spot & 0xE000) | spot.bits(11..2);
            }
        }

        Ok(PieceIter::Board5Iter {
            spot: spot,
            extra: extra,
        })
    }

    fn at_reset(&mut self, point: &Point) -> Result<PieceIter, &str> {
        let mut extra = self.continuations;
        let location = Board5::location(point);
        for i in 0..7 {
            if extra[i] & location != location {
                extra[i] = 0;
            } else if extra[i] != 0 {
                self.continuations[i] = 0;
                while extra[i] & 3 == 0 {
                    extra[i] = (extra[i] & location) | extra[i].bits(9..2);
                }
            }
        }

        let mut spot = self.grid[point.x][point.y];
        self.grid[point.x][point.y] = 0;
        if (spot & 0x0FFF) != 0 {
            while spot & 3 == 0 {
                spot = (spot & 0xE000) | spot.bits(11..2);
            }
        }

        Ok(PieceIter::Board5Iter {
            spot: spot,
            extra: extra,
        })
    }

    fn size(&self) -> usize {
        5
    }

    /// Checks to see if all spaces have at least one piece
    fn full(&self) -> bool {
        for row in self.grid.iter() {
            for square in row.iter() {
                if square.bits(15..13) == 0 {
                    return false
                }
            }
        }
        true
    }

    fn count(&self) -> PieceCount {
        let mut pieces = PieceCount::new(self.size());

        for row in self.grid.iter() {
            for spot in row.iter() {
                match Piece::from_u8(spot.bits(15..13) as u8) {
                    Some(Piece::OneFlat) => pieces.p1_flat += 1,
                    Some(Piece::OneStanding) => pieces.p1_flat += 1,
                    Some(Piece::OneCapstone) => pieces.p1_cap += 1,
                    Some(Piece::TwoFlat) => pieces.p2_flat += 1,
                    Some(Piece::TwoStanding) => pieces.p2_flat += 1,
                    Some(Piece::TwoCapstone) => pieces.p2_cap += 1,
                    None => {},
                }

                for i in 0..6 {
                    match spot.bits((2 * i + 1)..(2 * i)) {
                        1 => pieces.p1_flat += 1,
                        2 => pieces.p2_flat += 1,
                        _ => {},
                    }
                }
            }
        }
        //TODO: Count continuations
        pieces
    }

    fn place_piece(&mut self, point: &Point, piece: Piece) -> Result<(), String> {
        if top_piece_bits(self.grid[point.x][point.y]) == 0 {
            self.grid[point.x][point.y] = (piece as u16) << 13;
            Ok(())
        } else {
            return Err("Cannot place stone on top of existing stone.".into());
        }
    }

    fn add_piece(&mut self, point: &Point, piece: Piece) -> Result<(), String> {
        let spot = self.grid[point.x][point.y];
        let top_bits = top_piece_bits(spot);
        if top_bits == 0 {
            self.grid[point.x][point.y] = (piece as u16) << 13;
            return Ok(());
        }
        let mut top = Piece::from_u8(top_bits).unwrap();
        try!(piece.move_onto(&mut top));

        let mut lowest = spot.bits(1..0);
        let mut temp = spot.bits(11..2);
        if top.owner() == Player::One {
            temp = temp | 0x400;
        } else {
            temp = temp | 0x800;
        }
        self.grid[point.x][point.y] = temp | ((piece as u16) << 13);
        if lowest == 0 {
            Ok(())
        } else {
            // Need to add the bumped piece to an existing continuation
            for i in 0..7 {
                let cont = self.continuations[i];
                if (cont & 0xF800) == Board5::location(point) {
                    let mut temp: u16 = cont & 0xF800; // Keep location
                    temp |= cont.bits(9..2); // Shift everything thats there
                    temp |= lowest << 9; // Add in the bumped piece
                    self.continuations[i] = temp;

                    lowest = cont & 3; // We may have to do this again...
                    if lowest == 0 {
                        return Ok(());
                    }
                }
            }

            // If we haven't returned yet, we need to start a new continuation
            for i in 0..7 {
                let cont = self.continuations[i];
                if cont.bits(15..11) == 0 {
                    // It's unclaimed - let's claim it by setting the location
                    let mut temp: u16 = Board5::location(point);
                    temp |= lowest << 9; // Add in the bumped piece
                    self.continuations[i] = temp;
                    return Ok(());
                }
            }

            // Need to figure out continuations
            Err("No available continuations. Did you play too many pieces?".into())
        }
    }
}

impl FromStr for Board5 {
    type Err=();
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        board_from_str::<Board5>(s)
    }
}

impl fmt::Display for Board5 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        str_from_board(self, f)
    }
}
<|endoftext|>


<commit_before>use {Ipv4Network, Ipv6Network, IPV4_LENGTH, IPV6_LENGTH};
use std::net::{Ipv4Addr, Ipv6Addr};
use std::error::Error;

// TODO: These constants are true for Linux, but we have to check it for Windows and other systems
pub const IPV4_TYPE: u8 = 2;
pub const IPV6_TYPE: u8 = 3;

#[inline]
pub fn from_sql_ipv4_network(raw: &[u8]) -> Result<Ipv4Network, Box<Error + Sync + Send>> {
    if raw[0] != IPV4_TYPE {
        return Err("CIDR is not IP version 4".into())
    }

    if raw[2] != 1 {
        return Err("This field is not CIDR type, probably INET type".into())
    }

    if raw[3] != IPV4_LENGTH / 8 {
        return Err(format!("CIDR is IP version 4, but have bad length '{}'", raw[3]).into())
    }

    let network_address = Ipv4Addr::new(raw[4], raw[5], raw[6], raw[7]);
    let netmask = raw[1];
    Ok(Ipv4Network::from(network_address, netmask)?)
}

#[inline]
pub fn from_sql_ipv6_network(raw: &[u8]) -> Result<Ipv6Network, Box<Error + Sync + Send>> {
    if raw[0] != IPV6_TYPE {
        return Err("CIDR is not IP version 6".into())
    }

    if raw[2] != 1 {
        return Err("This field is not CIDR type, probably INET type".into())
    }

    if raw[3] != IPV6_LENGTH / 8 {
        return Err(format!("CIDR is IP version 6, but have bad length '{}'", raw[3]).into())
    }

    let mut octets = [0; 16];
    octets.copy_from_slice(&raw[4..]);
    let network_address = Ipv6Addr::from(octets);

    let netmask = raw[1];
    Ok(Ipv6Network::from(network_address, netmask)?)
}

#[inline]
pub fn to_sql_ipv4_network(network: &Ipv4Network) -> [u8; 8] {
    let ip_octets = network.network_address.octets();
    let mut bytes = [0; 8];
    bytes[0] = IPV4_TYPE;
    bytes[1] = network.netmask;
    bytes[2] = 1;
    bytes[3] = IPV4_LENGTH / 8;
    bytes[4..].copy_from_slice(&ip_octets);
    bytes
}

#[inline]
pub fn to_sql_ipv6_network(network: &Ipv6Network) -> [u8; 20] {
    let ip_octets = network.network_address.octets();
    let mut bytes = [0; 20];
    bytes[0] = IPV6_TYPE;
    bytes[1] = network.netmask;
    bytes[2] = 1;
    bytes[3] = IPV6_LENGTH / 8;
    bytes[4..].copy_from_slice(&ip_octets);
    bytes
}<commit_msg>Faster Postgres FromSql, bound check is made just once<commit_after>use {Ipv4Network, Ipv6Network, IPV4_LENGTH, IPV6_LENGTH};
use std::net::{Ipv4Addr, Ipv6Addr};
use std::error::Error;

// TODO: These constants are true for Linux, but we have to check it for Windows and other systems
pub const IPV4_TYPE: u8 = 2;
pub const IPV6_TYPE: u8 = 3;

#[inline]
pub fn from_sql_ipv4_network(raw: &[u8]) -> Result<Ipv4Network, Box<Error + Sync + Send>> {
    assert!(raw.len() >= 7);

    if raw[0] != IPV4_TYPE {
        return Err("CIDR is not IP version 4".into())
    }

    if raw[2] != 1 {
        return Err("This field is not CIDR type, probably INET type".into())
    }

    if raw[3] != IPV4_LENGTH / 8 {
        return Err(format!("CIDR is IP version 4, but have bad length '{}'", raw[3]).into())
    }

    let network_address = Ipv4Addr::new(raw[4], raw[5], raw[6], raw[7]);
    let netmask = raw[1];
    Ok(Ipv4Network::from(network_address, netmask)?)
}

#[inline]
pub fn from_sql_ipv6_network(raw: &[u8]) -> Result<Ipv6Network, Box<Error + Sync + Send>> {
    assert!(raw.len() >= 20);

    if raw[0] != IPV6_TYPE {
        return Err("CIDR is not IP version 6".into())
    }

    if raw[2] != 1 {
        return Err("This field is not CIDR type, probably INET type".into())
    }

    if raw[3] != IPV6_LENGTH / 8 {
        return Err(format!("CIDR is IP version 6, but have bad length '{}'", raw[3]).into())
    }

    let mut octets = [0; 16];
    octets.copy_from_slice(&raw[4..]);
    let network_address = Ipv6Addr::from(octets);

    let netmask = raw[1];
    Ok(Ipv6Network::from(network_address, netmask)?)
}

#[inline]
pub fn to_sql_ipv4_network(network: &Ipv4Network) -> [u8; 8] {
    let ip_octets = network.network_address.octets();
    let mut bytes = [0; 8];
    bytes[0] = IPV4_TYPE;
    bytes[1] = network.netmask;
    bytes[2] = 1;
    bytes[3] = IPV4_LENGTH / 8;
    bytes[4..].copy_from_slice(&ip_octets);
    bytes
}

#[inline]
pub fn to_sql_ipv6_network(network: &Ipv6Network) -> [u8; 20] {
    let ip_octets = network.network_address.octets();
    let mut bytes = [0; 20];
    bytes[0] = IPV6_TYPE;
    bytes[1] = network.netmask;
    bytes[2] = 1;
    bytes[3] = IPV6_LENGTH / 8;
    bytes[4..].copy_from_slice(&ip_octets);
    bytes
}<|endoftext|>


<commit_before>//! Performs autodetection of the host for the purposes of running
//! Cretonne to generate code to run on the same machine.

#![deny(missing_docs, trivial_numeric_casts, unused_extern_crates)]
#![warn(unused_import_braces, unstable_features)]
#![cfg_attr(feature = "clippy", plugin(clippy(conf_file = "../../clippy.toml")))]
#![cfg_attr(feature = "cargo-clippy",
            allow(new_without_default, new_without_default_derive))]
#![cfg_attr(feature="cargo-clippy", warn(
                float_arithmetic,
                mut_mut,
                nonminimal_bool,
                option_map_unwrap_or,
                option_map_unwrap_or_else,
                print_stdout,
                unicode_not_nfc,
                use_self,
                ))]
#![cfg_attr(not(feature = "std"), no_std)]

extern crate cretonne_codegen;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
extern crate raw_cpuid;

use cretonne_codegen::isa;
use cretonne_codegen::settings::{self, Configurable};

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
use raw_cpuid::CpuId;

/// Return `settings` and `isa` builders configured for the current host
/// machine, or `Err(())` if the host machine is not supported
/// in the current configuration.
pub fn builders() -> Result<(settings::Builder, isa::Builder), &'static str> {
    let mut flag_builder = settings::builder();

    if cfg!(unix) {
        flag_builder.set("call_conv", "system_v").unwrap();
    } else if cfg!(windows) {
        flag_builder.set("call_conv", "fastcall").unwrap();
    } else {
        return Err("unrecognized environment");
    }

    if cfg!(target_pointer_width = "64") {
        flag_builder.enable("is_64bit").unwrap();
    } else if !cfg!(target_pointer_width = "32") {
        return Err("unrecognized pointer size");
    }

    // TODO: Add RISC-V support once Rust supports it.
    let name = if cfg!(any(target_arch = "x86", target_arch = "x86_64")) {
        "x86"
    } else if cfg!(target_arch = "arm") {
        "arm32"
    } else if cfg!(target_arch = "aarch64") {
        "arm64"
    } else {
        return Err("unrecognized architecture");
    };

    let mut isa_builder = isa::lookup(name).map_err(|err| match err {
        isa::LookupError::Unknown => panic!(),
        isa::LookupError::Unsupported => "unsupported architecture",
    })?;

    if cfg!(any(target_arch = "x86", target_arch = "x86_64")) {
        parse_x86_cpuid(&mut isa_builder)?;
    }

    Ok((flag_builder, isa_builder))
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn parse_x86_cpuid(isa_builder: &mut isa::Builder) -> Result<(), &'static str> {
    let cpuid = CpuId::new();

    if let Some(info) = cpuid.get_feature_info() {
        if !info.has_sse2() {
            return Err("x86 support requires SSE2");
        }
        if info.has_sse3() {
            isa_builder.enable("has_sse3").unwrap();
        }
        if info.has_sse41() {
            isa_builder.enable("has_sse41").unwrap();
        }
        if info.has_sse42() {
            isa_builder.enable("has_sse42").unwrap();
        }
        if info.has_popcnt() {
            isa_builder.enable("has_popcnt").unwrap();
        }
        if info.has_avx() {
            isa_builder.enable("has_avx").unwrap();
        }
    }
    if let Some(info) = cpuid.get_extended_feature_info() {
        if info.has_bmi1() {
            isa_builder.enable("has_bmi1").unwrap();
        }
        if info.has_bmi2() {
            isa_builder.enable("has_bmi2").unwrap();
        }
    }
    if let Some(info) = cpuid.get_extended_function_info() {
        if info.has_lzcnt() {
            isa_builder.enable("has_lzcnt").unwrap();
        }
    }
    Ok(())
}
<commit_msg>Add support for target_os = "nebulet". (#319)<commit_after>//! Performs autodetection of the host for the purposes of running
//! Cretonne to generate code to run on the same machine.

#![deny(missing_docs, trivial_numeric_casts, unused_extern_crates)]
#![warn(unused_import_braces, unstable_features)]
#![cfg_attr(feature = "clippy", plugin(clippy(conf_file = "../../clippy.toml")))]
#![cfg_attr(feature = "cargo-clippy",
            allow(new_without_default, new_without_default_derive))]
#![cfg_attr(feature="cargo-clippy", warn(
                float_arithmetic,
                mut_mut,
                nonminimal_bool,
                option_map_unwrap_or,
                option_map_unwrap_or_else,
                print_stdout,
                unicode_not_nfc,
                use_self,
                ))]
#![cfg_attr(not(feature = "std"), no_std)]

extern crate cretonne_codegen;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
extern crate raw_cpuid;

use cretonne_codegen::isa;
use cretonne_codegen::settings::{self, Configurable};

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
use raw_cpuid::CpuId;

/// Return `settings` and `isa` builders configured for the current host
/// machine, or `Err(())` if the host machine is not supported
/// in the current configuration.
pub fn builders() -> Result<(settings::Builder, isa::Builder), &'static str> {
    let mut flag_builder = settings::builder();

    if cfg!(any(unix, target_os = "nebulet")) {
        flag_builder.set("call_conv", "system_v").unwrap();
    } else if cfg!(windows) {
        flag_builder.set("call_conv", "fastcall").unwrap();
    } else {
        return Err("unrecognized environment");
    }

    if cfg!(target_pointer_width = "64") {
        flag_builder.enable("is_64bit").unwrap();
    } else if !cfg!(target_pointer_width = "32") {
        return Err("unrecognized pointer size");
    }

    // TODO: Add RISC-V support once Rust supports it.
    let name = if cfg!(any(target_arch = "x86", target_arch = "x86_64")) {
        "x86"
    } else if cfg!(target_arch = "arm") {
        "arm32"
    } else if cfg!(target_arch = "aarch64") {
        "arm64"
    } else {
        return Err("unrecognized architecture");
    };

    let mut isa_builder = isa::lookup(name).map_err(|err| match err {
        isa::LookupError::Unknown => panic!(),
        isa::LookupError::Unsupported => "unsupported architecture",
    })?;

    if cfg!(any(target_arch = "x86", target_arch = "x86_64")) {
        parse_x86_cpuid(&mut isa_builder)?;
    }

    Ok((flag_builder, isa_builder))
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn parse_x86_cpuid(isa_builder: &mut isa::Builder) -> Result<(), &'static str> {
    let cpuid = CpuId::new();

    if let Some(info) = cpuid.get_feature_info() {
        if !info.has_sse2() {
            return Err("x86 support requires SSE2");
        }
        if info.has_sse3() {
            isa_builder.enable("has_sse3").unwrap();
        }
        if info.has_sse41() {
            isa_builder.enable("has_sse41").unwrap();
        }
        if info.has_sse42() {
            isa_builder.enable("has_sse42").unwrap();
        }
        if info.has_popcnt() {
            isa_builder.enable("has_popcnt").unwrap();
        }
        if info.has_avx() {
            isa_builder.enable("has_avx").unwrap();
        }
    }
    if let Some(info) = cpuid.get_extended_feature_info() {
        if info.has_bmi1() {
            isa_builder.enable("has_bmi1").unwrap();
        }
        if info.has_bmi2() {
            isa_builder.enable("has_bmi2").unwrap();
        }
    }
    if let Some(info) = cpuid.get_extended_function_info() {
        if info.has_lzcnt() {
            isa_builder.enable("has_lzcnt").unwrap();
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>//! Provide syntax highlighting via [Syntect].
//!
//! [Syntect]: https://docs.rs/syntect/1.0.0/syntect/

// Standard library
use std::default::Default;
use std::collections::HashMap;
use std::str;

// Third party
use quick_xml::{XmlReader, XmlWriter, Element, Event};
use syntect::html::highlighted_snippet_for_string;
use syntect::highlighting::Theme;
use syntect::parsing::{SyntaxDefinition, SyntaxSet};

/// A `Language` is a `String` representing a code highlighting language.
type Language = String;

/// Define XML parsing states of interest for highlighting.
///
/// We use these in conjunction with `ParseEvent`s to drive a state machine
/// which tracks whether a stream of `quick_xml::Event`s indicates the presence
/// of a code block to be highlighted.
#[derive(Clone, Debug, PartialEq)]
enum ParseState {
    /// Not in a code block at all.
    NotInBlock,
    /// Possibly the start of a highlight block, with the language as payload.
    MaybeStartBlock(Language),
    /// The start of the code block has been confirmed; highlight next text node.
    WillStartCodeBlock(Language),
    /// In a code block; use the content for highlighting.
    InCodeBlock(Language),
}

/// Define XML parsing events of interest for highlighting.
///
/// We don't care at all about any other items for this specific scenario; we
/// use these in conjunction with the `ParseState` to drive the state  machine
/// described there.
#[derive(Debug)]
enum ParseEvent {
    // Starting a `<pre>`; optionally specifying a language.
    StartPre(Option<Language>),
    // Starting a `<code>`.
    StartCode,
    // Ending a `<code>`.
    EndCode,
    Text,
    Whitespace,
    Other,
}

impl ParseState {
    /// Get the next `ParseState` given current `ParseState` and a `ParseEvent`.
    fn next(self, event: ParseEvent) -> ParseState {
        use self::ParseState::*;
        use self::ParseEvent::*;

        match (self, event) {
            (NotInBlock, StartPre(Some(language))) => MaybeStartBlock(language),
            (MaybeStartBlock(language), Whitespace) => MaybeStartBlock(language),
            (MaybeStartBlock(language), StartCode) => WillStartCodeBlock(language),
            (WillStartCodeBlock(language), Text) |
            (WillStartCodeBlock(language), Whitespace) => InCodeBlock(language),
            (InCodeBlock(_), EndCode) => NotInBlock,
            (_, _) => NotInBlock,
        }
    }
}

impl Default for ParseState {
    fn default() -> ParseState {
        ParseState::NotInBlock
    }
}

impl<'e> From<&'e Event> for ParseEvent {
    fn from(event: &'e Event) -> ParseEvent {
        // TODO @1.15: remove `'static`.
        const PRE: &'static [u8] = b"pre";
        const CODE: &'static [u8] = b"code";
        const CLASS: &'static [u8] = b"class";
        const WHITE_SPACE: &'static [u8] = b"";

        match *event {
            Event::Start(ref element) => {
                match element.name() {
                    PRE => {
                        let maybe_class_attr = element.attributes()
                            .map(|attr| attr.unwrap())
                            .filter(|&(attr, _value)| attr == CLASS)
                            .next();

                        if let Some((_attr, value)) = maybe_class_attr {
                            match str::from_utf8(value) {
                                Ok(lang) => ParseEvent::StartPre(Some(lang.to_string())),
                                Err(_) => ParseEvent::StartPre(None),
                            }
                        } else {
                            ParseEvent::StartPre(None)
                        }
                    }
                    CODE => ParseEvent::StartCode,
                    _ => ParseEvent::Other,
                }
            }
            Event::End(ref element) => {
                match element.name() {
                    CODE => ParseEvent::EndCode,
                    _ => ParseEvent::Other,
                }
            }
            Event::Text(ref element) => {
                match element.name() {
                    WHITE_SPACE => ParseEvent::Whitespace,
                    _ => ParseEvent::Text,
                }
            }
            _ => ParseEvent::Other,
        }
    }
}

/// Highlight all code blocks in a block of HTML.
///
/// Assumes that the blocks to be highlighted are in the following basic format:
///
/// ```html
/// <pre class="rust">
///   <code>
///     pub fn syntax_highlight(html_string: String) -> String {
///         // implementation details
///     }
///   </code>
/// </pre>
/// ```
///
/// The `class` attribute value from the `pre` tag defines the language used to
/// highlight the code.
///
/// Note that any `html_string` will do; if it cannot be parsed as XML, it will
/// simply be returned unchanged; and if there are no code blocks to highlight,
/// it will also be returned unchanged.
pub fn syntax_highlight(html_string: String, theme: &Theme) -> String {
    let ss = SyntaxSet::load_defaults_nonewlines();
    let mut syntax_definitions = HashMap::<Language, SyntaxDefinition>::new();

    let mut writer = XmlWriter::new(Vec::<u8>::new());
    let mut state = ParseState::default();

    for event in XmlReader::from(html_string.as_str()) {
        let event = match event {
            Ok(event) => event,
            Err(_) => continue
        };

        let parse_event = ParseEvent::from(&event);
        state = ParseState::next(state, parse_event);

        if let ParseState::InCodeBlock(ref language) = state {
            if let Ok(unescaped_content) = event.element().unescaped_content() {
                if let Ok(content_to_highlight) = str::from_utf8(&unescaped_content) {
                    if let Some(valid_syntax) = ss.find_syntax_by_token(&language) {
                        let syntax_definition =
                            syntax_definitions.entry(language.clone())
                                              .or_insert_with(|| valid_syntax.clone());

                        let highlighted = highlighted_snippet_for_string(content_to_highlight,
                                                                         syntax_definition,
                                                                         &theme);
                        let text = Element::new(highlighted);
                        assert!(writer.write(Event::Text(text)).is_ok());
                        continue;
                    }
                }
            }
        }

        // Syntax highlighting did not succeed, so just write the original event.
        assert!(writer.write(event).is_ok());
    }

    String::from_utf8(writer.into_inner()).unwrap_or(html_string)
}

#[cfg(test)]
mod tests {

    #[test]
    fn test_parse_state() {
        use super::ParseState;
        use super::ParseEvent;

        let lang = "rust";

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::StartPre(Some(lang.to_string()))),
                   ParseState::MaybeStartBlock(lang.to_string()));

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::EndCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::Other),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::StartCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::StartCode),
                   ParseState::WillStartCodeBlock(lang.to_string()));

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::Text),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::EndCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string())
                       .next(ParseEvent::StartPre(Some(lang.to_string()))),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::Other),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::WillStartCodeBlock(lang.to_string()).next(ParseEvent::Text),
                   ParseState::InCodeBlock(lang.to_string()));
    }
}
<commit_msg>Avoid cloning SyntaxDefinition values.<commit_after>//! Provide syntax highlighting via [Syntect].
//!
//! [Syntect]: https://docs.rs/syntect/1.0.0/syntect/

// Standard library
use std::default::Default;
use std::collections::HashMap;
use std::str;

// Third party
use quick_xml::{XmlReader, XmlWriter, Element, Event};
use syntect::html::highlighted_snippet_for_string;
use syntect::highlighting::Theme;
use syntect::parsing::{SyntaxDefinition, SyntaxSet};

/// A `Language` is a `String` representing a code highlighting language.
type Language = String;

/// Define XML parsing states of interest for highlighting.
///
/// We use these in conjunction with `ParseEvent`s to drive a state machine
/// which tracks whether a stream of `quick_xml::Event`s indicates the presence
/// of a code block to be highlighted.
#[derive(Clone, Debug, PartialEq)]
enum ParseState {
    /// Not in a code block at all.
    NotInBlock,
    /// Possibly the start of a highlight block, with the language as payload.
    MaybeStartBlock(Language),
    /// The start of the code block has been confirmed; highlight next text node.
    WillStartCodeBlock(Language),
    /// In a code block; use the content for highlighting.
    InCodeBlock(Language),
}

/// Define XML parsing events of interest for highlighting.
///
/// We don't care at all about any other items for this specific scenario; we
/// use these in conjunction with the `ParseState` to drive the state  machine
/// described there.
#[derive(Debug)]
enum ParseEvent {
    // Starting a `<pre>`; optionally specifying a language.
    StartPre(Option<Language>),
    // Starting a `<code>`.
    StartCode,
    // Ending a `<code>`.
    EndCode,
    Text,
    Whitespace,
    Other,
}

impl ParseState {
    /// Get the next `ParseState` given current `ParseState` and a `ParseEvent`.
    fn next(self, event: ParseEvent) -> ParseState {
        use self::ParseState::*;
        use self::ParseEvent::*;

        match (self, event) {
            (NotInBlock, StartPre(Some(language))) => MaybeStartBlock(language),
            (MaybeStartBlock(language), Whitespace) => MaybeStartBlock(language),
            (MaybeStartBlock(language), StartCode) => WillStartCodeBlock(language),
            (WillStartCodeBlock(language), Text) |
            (WillStartCodeBlock(language), Whitespace) => InCodeBlock(language),
            (InCodeBlock(_), EndCode) => NotInBlock,
            (_, _) => NotInBlock,
        }
    }
}

impl Default for ParseState {
    fn default() -> ParseState {
        ParseState::NotInBlock
    }
}

impl<'e> From<&'e Event> for ParseEvent {
    fn from(event: &'e Event) -> ParseEvent {
        // TODO @1.15: remove `'static`.
        const PRE: &'static [u8] = b"pre";
        const CODE: &'static [u8] = b"code";
        const CLASS: &'static [u8] = b"class";
        const WHITE_SPACE: &'static [u8] = b"";

        match *event {
            Event::Start(ref element) => {
                match element.name() {
                    PRE => {
                        let maybe_class_attr = element.attributes()
                            .map(|attr| attr.unwrap())
                            .filter(|&(attr, _value)| attr == CLASS)
                            .next();

                        if let Some((_attr, value)) = maybe_class_attr {
                            match str::from_utf8(value) {
                                Ok(lang) => ParseEvent::StartPre(Some(lang.to_string())),
                                Err(_) => ParseEvent::StartPre(None),
                            }
                        } else {
                            ParseEvent::StartPre(None)
                        }
                    }
                    CODE => ParseEvent::StartCode,
                    _ => ParseEvent::Other,
                }
            }
            Event::End(ref element) => {
                match element.name() {
                    CODE => ParseEvent::EndCode,
                    _ => ParseEvent::Other,
                }
            }
            Event::Text(ref element) => {
                match element.name() {
                    WHITE_SPACE => ParseEvent::Whitespace,
                    _ => ParseEvent::Text,
                }
            }
            _ => ParseEvent::Other,
        }
    }
}

/// Highlight all code blocks in a block of HTML.
///
/// Assumes that the blocks to be highlighted are in the following basic format:
///
/// ```html
/// <pre class="rust">
///   <code>
///     pub fn syntax_highlight(html_string: String) -> String {
///         // implementation details
///     }
///   </code>
/// </pre>
/// ```
///
/// The `class` attribute value from the `pre` tag defines the language used to
/// highlight the code.
///
/// Note that any `html_string` will do; if it cannot be parsed as XML, it will
/// simply be returned unchanged; and if there are no code blocks to highlight,
/// it will also be returned unchanged.
pub fn syntax_highlight(html_string: String, theme: &Theme) -> String {
    let ss = SyntaxSet::load_defaults_nonewlines();
    let mut syntax_definitions = HashMap::<Language, &SyntaxDefinition>::new();

    let mut writer = XmlWriter::new(Vec::<u8>::new());
    let mut state = ParseState::default();

    for event in XmlReader::from(html_string.as_str()) {
        let event = match event {
            Ok(event) => event,
            Err(_) => continue
        };

        let parse_event = ParseEvent::from(&event);
        state = ParseState::next(state, parse_event);

        if let ParseState::InCodeBlock(ref language) = state {
            if let Ok(unescaped_content) = event.element().unescaped_content() {
                if let Ok(content_to_highlight) = str::from_utf8(&unescaped_content) {
                    if let Some(valid_syntax) = ss.find_syntax_by_token(&language) {
                        let syntax_definition =
                            syntax_definitions.entry(language.clone()).or_insert(valid_syntax);

                        let highlighted = highlighted_snippet_for_string(content_to_highlight,
                                                                         syntax_definition,
                                                                         &theme);
                        let text = Element::new(highlighted);
                        assert!(writer.write(Event::Text(text)).is_ok());
                        continue;
                    }
                }
            }
        }

        // Syntax highlighting did not succeed, so just write the original event.
        assert!(writer.write(event).is_ok());
    }

    String::from_utf8(writer.into_inner()).unwrap_or(html_string)
}

#[cfg(test)]
mod tests {

    #[test]
    fn test_parse_state() {
        use super::ParseState;
        use super::ParseEvent;

        let lang = "rust";

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::StartPre(Some(lang.to_string()))),
                   ParseState::MaybeStartBlock(lang.to_string()));

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::EndCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::Other),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::NotInBlock.next(ParseEvent::StartCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::StartCode),
                   ParseState::WillStartCodeBlock(lang.to_string()));

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::Text),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::EndCode),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string())
                       .next(ParseEvent::StartPre(Some(lang.to_string()))),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::MaybeStartBlock(lang.to_string()).next(ParseEvent::Other),
                   ParseState::NotInBlock);

        assert_eq!(ParseState::WillStartCodeBlock(lang.to_string()).next(ParseEvent::Text),
                   ParseState::InCodeBlock(lang.to_string()));
    }
}
<|endoftext|>