<commit_before>use crate::alloc::{host_page_size, instance_heap_offset, Alloc, Limits, Slot};
use crate::embed_ctx::CtxMap;
use crate::error::Error;
use crate::instance::{new_instance_handle, Instance, InstanceHandle};
use crate::module::Module;
use crate::region::{Region, RegionCreate, RegionInternal};
use libc::{c_void, memset, SIGSTKSZ};
use nix::sys::mman::{madvise, mmap, munmap, MapFlags, MmapAdvise, ProtFlags};
use std::ptr;
use std::sync::{Arc, Mutex, Weak};

/// A [`Region`](trait.Region.html) backed by `mmap`.
pub struct MmapRegion {
    capacity: usize,
    freelist: Mutex<Vec<Slot>>,
    limits: Limits,
}

impl Region for MmapRegion {}

impl RegionInternal for MmapRegion {
    fn new_instance_with(
        &self,
        module: Arc<dyn Module>,
        embed_ctx: CtxMap,
    ) -> Result<InstanceHandle, Error> {
        let slot = self
            .freelist
            .lock()
            .unwrap()
            .pop()
            .ok_or(Error::RegionFull(self.capacity))?;

        if slot.heap as usize % host_page_size() != 0 {
            lucet_bail!("heap is not page-aligned; this is a bug");
        }

        let limits = &slot.limits;
        module.validate_runtime_spec(limits)?;

        for (ptr, len) in [
            // make the stack read/writable
            (slot.stack, limits.stack_size),
            // make the globals read/writable
            (slot.globals, limits.globals_size),
            // make the sigstack read/writable
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting r/w {:p}[{:x}]", *ptr, len);
            unsafe { mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? };
        }

        // note: the initial heap will be made read/writable when `new_instance_handle` calls `reset`

        let inst_ptr = slot.start as *mut Instance;

        // upgrade the slot's weak region pointer so the region can't get dropped while the instance
        // exists
        let region = slot
            .region
            .upgrade()
            // if this precondition isn't met, something is deeply wrong as some other region's slot
            // ended up in our freelist
            .expect("backing region of slot (`self`) exists");

        let alloc = Alloc {
            heap_accessible_size: 0, // the `reset` call in `new_instance_handle` will set this
            heap_inaccessible_size: slot.limits.heap_address_space_size,
            slot: Some(slot),
            region,
        };

        let inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;

        Ok(inst)
    }

    fn drop_alloc(&self, alloc: &mut Alloc) {
        let slot = alloc
            .slot
            .take()
            .expect("alloc didn't have a slot during drop; dropped twice?");

        if slot.heap as usize % host_page_size() != 0 {
            panic!("heap is not page-aligned");
        }

        // clear and disable access to the heap, stack, globals, and sigstack
        for (ptr, len) in [
            (slot.heap, slot.limits.heap_address_space_size),
            (slot.stack, slot.limits.stack_size),
            (slot.globals, slot.limits.globals_size),
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting none {:p}[{:x}]", *ptr, len);
            unsafe {
                mprotect(*ptr, *len, ProtFlags::PROT_NONE).expect("mprotect succeeds during drop");
                madvise(*ptr, *len, MmapAdvise::MADV_DONTNEED)
                    .expect("madvise succeeds during drop");
            }
        }

        self.freelist.lock().unwrap().push(slot);
    }

    fn expand_heap(&self, slot: &Slot, start: u32, len: u32) -> Result<(), Error> {
        unsafe {
            mprotect(
                (slot.heap as usize + start as usize) as *mut c_void,
                len as usize,
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?;
        }
        Ok(())
    }

    fn reset_heap(&self, alloc: &mut Alloc, module: &dyn Module) -> Result<(), Error> {
        let heap = alloc.slot().heap;

        if alloc.heap_accessible_size > 0 {
            // zero the whole heap, if any of it is currently accessible
            let heap_size = alloc.slot().limits.heap_address_space_size;

            unsafe {
                // `mprotect()` and `madvise()` are sufficient to zero a page on Linux,
                // but not necessarily on all POSIX operating systems, and on macOS in particular.
                #[cfg(not(target_os = "linux"))]
                {
                    mprotect(
                        heap,
                        alloc.heap_accessible_size,
                        ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                    )?;
                    memset(heap, 0, alloc.heap_accessible_size);
                }
                mprotect(heap, heap_size, ProtFlags::PROT_NONE)?;
                madvise(heap, heap_size, MmapAdvise::MADV_DONTNEED)?;
            }
        }

        let initial_size = module
            .heap_spec()
            .map(|h| h.initial_size as usize)
            .unwrap_or(0);

        // reset the heap to the initial size, and mprotect those pages appropriately
        if initial_size > 0 {
            unsafe {
                mprotect(
                    heap,
                    initial_size,
                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                )?
            };
        }
        alloc.heap_accessible_size = initial_size;
        alloc.heap_inaccessible_size = alloc.slot().limits.heap_address_space_size - initial_size;

        // Initialize the heap using the module sparse page data. There cannot be more pages in the
        // sparse page data than will fit in the initial heap size.
        //
        // Pages with a corresponding Some entry in the sparse page data are initialized with
        // the contents of that data.
        //
        // Any pages which don't have an entry in the sparse page data, either because their entry
        // is None, or because the sparse data has fewer pages than the initial heap, are zeroed.
        let heap = unsafe { alloc.heap_mut() };
        let initial_pages =
            initial_size
                .checked_div(host_page_size())
                .ok_or(lucet_incorrect_module!(
                    "initial heap size {} is not divisible by host page size ({})",
                    initial_size,
                    host_page_size()
                ))?;
        for page_num in 0..initial_pages {
            let page_base = page_num * host_page_size();
            if heap.len() < page_base {
                return Err(lucet_incorrect_module!(
                    "sparse page data length exceeded initial heap size"
                ));
            }
            if let Some(contents) = module.get_sparse_page_data(page_num) {
                // otherwise copy in the page data
                heap[page_base..page_base + host_page_size()].copy_from_slice(contents);
            }
        }

        Ok(())
    }

    fn as_dyn_internal(&self) -> &dyn RegionInternal {
        self
    }
}

impl Drop for MmapRegion {
    fn drop(&mut self) {
        for slot in self.freelist.get_mut().unwrap().drain(0..) {
            Self::free_slot(slot);
        }
    }
}

impl RegionCreate for MmapRegion {
    const TYPE_NAME: &'static str = "MmapRegion";

    fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        MmapRegion::create(instance_capacity, limits)
    }
}

impl MmapRegion {
    /// Create a new `MmapRegion` that can support a given number instances, each subject to the
    /// same runtime limits.
    ///
    /// The region is returned in an `Arc`, because any instances created from it carry a reference
    /// back to the region.
    pub fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        assert!(
            SIGSTKSZ % host_page_size() == 0,
            "signal stack size is a multiple of host page size"
        );
        limits.validate()?;

        let region = Arc::new(MmapRegion {
            capacity: instance_capacity,
            freelist: Mutex::new(Vec::with_capacity(instance_capacity)),
            limits: limits.clone(),
        });
        {
            let mut freelist = region.freelist.lock().unwrap();
            for _ in 0..instance_capacity {
                freelist.push(MmapRegion::create_slot(&region)?);
            }
        }

        Ok(region)
    }

    fn create_slot(region: &Arc<MmapRegion>) -> Result<Slot, Error> {
        // get the chunk of virtual memory that the `Slot` will manage
        let mem = unsafe {
            mmap(
                ptr::null_mut(),
                region.limits.total_memory_size(),
                ProtFlags::PROT_NONE,
                MapFlags::MAP_ANON | MapFlags::MAP_PRIVATE,
                0,
                0,
            )?
        };

        // set the first part of the memory to read/write so that the `Instance` can be stored there
        // TODO: post slot refactor, is this necessary/desirable?
        unsafe {
            mprotect(
                mem,
                instance_heap_offset(),
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?
        };

        // lay out the other sections in memory
        let heap = mem as usize + instance_heap_offset();
        let stack = heap + region.limits.heap_address_space_size;
        let globals = stack + region.limits.stack_size + host_page_size();
        let sigstack = globals + host_page_size();

        Ok(Slot {
            start: mem,
            heap: heap as *mut c_void,
            stack: stack as *mut c_void,
            globals: globals as *mut c_void,
            sigstack: sigstack as *mut c_void,
            limits: region.limits.clone(),
            region: Arc::downgrade(region) as Weak<dyn RegionInternal>,
        })
    }

    fn free_slot(slot: Slot) {
        // eprintln!(
        //     "unmapping {:p}[{:x}]",
        //     slot.start,
        //     slot.limits.total_memory_size()
        // );
        let res = unsafe { munmap(slot.start, slot.limits.total_memory_size()) };
        res.expect("munmap succeeded");
    }
}

// TODO: remove this once `nix` PR https://github.com/nix-rust/nix/pull/991 is merged
unsafe fn mprotect(addr: *mut c_void, length: libc::size_t, prot: ProtFlags) -> nix::Result<()> {
    nix::errno::Errno::result(libc::mprotect(addr, length, prot.bits())).map(drop)
}
<commit_msg>Just remove a warning about an unused use statement on linux<commit_after>use crate::alloc::{host_page_size, instance_heap_offset, Alloc, Limits, Slot};
use crate::embed_ctx::CtxMap;
use crate::error::Error;
use crate::instance::{new_instance_handle, Instance, InstanceHandle};
use crate::module::Module;
use crate::region::{Region, RegionCreate, RegionInternal};
#[cfg(not(target_os = "linux"))]
use libc::memset;
use libc::{c_void, SIGSTKSZ};
use nix::sys::mman::{madvise, mmap, munmap, MapFlags, MmapAdvise, ProtFlags};
use std::ptr;
use std::sync::{Arc, Mutex, Weak};

/// A [`Region`](trait.Region.html) backed by `mmap`.
pub struct MmapRegion {
    capacity: usize,
    freelist: Mutex<Vec<Slot>>,
    limits: Limits,
}

impl Region for MmapRegion {}

impl RegionInternal for MmapRegion {
    fn new_instance_with(
        &self,
        module: Arc<dyn Module>,
        embed_ctx: CtxMap,
    ) -> Result<InstanceHandle, Error> {
        let slot = self
            .freelist
            .lock()
            .unwrap()
            .pop()
            .ok_or(Error::RegionFull(self.capacity))?;

        if slot.heap as usize % host_page_size() != 0 {
            lucet_bail!("heap is not page-aligned; this is a bug");
        }

        let limits = &slot.limits;
        module.validate_runtime_spec(limits)?;

        for (ptr, len) in [
            // make the stack read/writable
            (slot.stack, limits.stack_size),
            // make the globals read/writable
            (slot.globals, limits.globals_size),
            // make the sigstack read/writable
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting r/w {:p}[{:x}]", *ptr, len);
            unsafe { mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? };
        }

        // note: the initial heap will be made read/writable when `new_instance_handle` calls `reset`

        let inst_ptr = slot.start as *mut Instance;

        // upgrade the slot's weak region pointer so the region can't get dropped while the instance
        // exists
        let region = slot
            .region
            .upgrade()
            // if this precondition isn't met, something is deeply wrong as some other region's slot
            // ended up in our freelist
            .expect("backing region of slot (`self`) exists");

        let alloc = Alloc {
            heap_accessible_size: 0, // the `reset` call in `new_instance_handle` will set this
            heap_inaccessible_size: slot.limits.heap_address_space_size,
            slot: Some(slot),
            region,
        };

        let inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;

        Ok(inst)
    }

    fn drop_alloc(&self, alloc: &mut Alloc) {
        let slot = alloc
            .slot
            .take()
            .expect("alloc didn't have a slot during drop; dropped twice?");

        if slot.heap as usize % host_page_size() != 0 {
            panic!("heap is not page-aligned");
        }

        // clear and disable access to the heap, stack, globals, and sigstack
        for (ptr, len) in [
            (slot.heap, slot.limits.heap_address_space_size),
            (slot.stack, slot.limits.stack_size),
            (slot.globals, slot.limits.globals_size),
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting none {:p}[{:x}]", *ptr, len);
            unsafe {
                mprotect(*ptr, *len, ProtFlags::PROT_NONE).expect("mprotect succeeds during drop");
                madvise(*ptr, *len, MmapAdvise::MADV_DONTNEED)
                    .expect("madvise succeeds during drop");
            }
        }

        self.freelist.lock().unwrap().push(slot);
    }

    fn expand_heap(&self, slot: &Slot, start: u32, len: u32) -> Result<(), Error> {
        unsafe {
            mprotect(
                (slot.heap as usize + start as usize) as *mut c_void,
                len as usize,
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?;
        }
        Ok(())
    }

    fn reset_heap(&self, alloc: &mut Alloc, module: &dyn Module) -> Result<(), Error> {
        let heap = alloc.slot().heap;

        if alloc.heap_accessible_size > 0 {
            // zero the whole heap, if any of it is currently accessible
            let heap_size = alloc.slot().limits.heap_address_space_size;

            unsafe {
                // `mprotect()` and `madvise()` are sufficient to zero a page on Linux,
                // but not necessarily on all POSIX operating systems, and on macOS in particular.
                #[cfg(not(target_os = "linux"))]
                {
                    mprotect(
                        heap,
                        alloc.heap_accessible_size,
                        ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                    )?;
                    memset(heap, 0, alloc.heap_accessible_size);
                }
                mprotect(heap, heap_size, ProtFlags::PROT_NONE)?;
                madvise(heap, heap_size, MmapAdvise::MADV_DONTNEED)?;
            }
        }

        let initial_size = module
            .heap_spec()
            .map(|h| h.initial_size as usize)
            .unwrap_or(0);

        // reset the heap to the initial size, and mprotect those pages appropriately
        if initial_size > 0 {
            unsafe {
                mprotect(
                    heap,
                    initial_size,
                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                )?
            };
        }
        alloc.heap_accessible_size = initial_size;
        alloc.heap_inaccessible_size = alloc.slot().limits.heap_address_space_size - initial_size;

        // Initialize the heap using the module sparse page data. There cannot be more pages in the
        // sparse page data than will fit in the initial heap size.
        //
        // Pages with a corresponding Some entry in the sparse page data are initialized with
        // the contents of that data.
        //
        // Any pages which don't have an entry in the sparse page data, either because their entry
        // is None, or because the sparse data has fewer pages than the initial heap, are zeroed.
        let heap = unsafe { alloc.heap_mut() };
        let initial_pages =
            initial_size
                .checked_div(host_page_size())
                .ok_or(lucet_incorrect_module!(
                    "initial heap size {} is not divisible by host page size ({})",
                    initial_size,
                    host_page_size()
                ))?;
        for page_num in 0..initial_pages {
            let page_base = page_num * host_page_size();
            if heap.len() < page_base {
                return Err(lucet_incorrect_module!(
                    "sparse page data length exceeded initial heap size"
                ));
            }
            if let Some(contents) = module.get_sparse_page_data(page_num) {
                // otherwise copy in the page data
                heap[page_base..page_base + host_page_size()].copy_from_slice(contents);
            }
        }

        Ok(())
    }

    fn as_dyn_internal(&self) -> &dyn RegionInternal {
        self
    }
}

impl Drop for MmapRegion {
    fn drop(&mut self) {
        for slot in self.freelist.get_mut().unwrap().drain(0..) {
            Self::free_slot(slot);
        }
    }
}

impl RegionCreate for MmapRegion {
    const TYPE_NAME: &'static str = "MmapRegion";

    fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        MmapRegion::create(instance_capacity, limits)
    }
}

impl MmapRegion {
    /// Create a new `MmapRegion` that can support a given number instances, each subject to the
    /// same runtime limits.
    ///
    /// The region is returned in an `Arc`, because any instances created from it carry a reference
    /// back to the region.
    pub fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        assert!(
            SIGSTKSZ % host_page_size() == 0,
            "signal stack size is a multiple of host page size"
        );
        limits.validate()?;

        let region = Arc::new(MmapRegion {
            capacity: instance_capacity,
            freelist: Mutex::new(Vec::with_capacity(instance_capacity)),
            limits: limits.clone(),
        });
        {
            let mut freelist = region.freelist.lock().unwrap();
            for _ in 0..instance_capacity {
                freelist.push(MmapRegion::create_slot(&region)?);
            }
        }

        Ok(region)
    }

    fn create_slot(region: &Arc<MmapRegion>) -> Result<Slot, Error> {
        // get the chunk of virtual memory that the `Slot` will manage
        let mem = unsafe {
            mmap(
                ptr::null_mut(),
                region.limits.total_memory_size(),
                ProtFlags::PROT_NONE,
                MapFlags::MAP_ANON | MapFlags::MAP_PRIVATE,
                0,
                0,
            )?
        };

        // set the first part of the memory to read/write so that the `Instance` can be stored there
        // TODO: post slot refactor, is this necessary/desirable?
        unsafe {
            mprotect(
                mem,
                instance_heap_offset(),
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?
        };

        // lay out the other sections in memory
        let heap = mem as usize + instance_heap_offset();
        let stack = heap + region.limits.heap_address_space_size;
        let globals = stack + region.limits.stack_size + host_page_size();
        let sigstack = globals + host_page_size();

        Ok(Slot {
            start: mem,
            heap: heap as *mut c_void,
            stack: stack as *mut c_void,
            globals: globals as *mut c_void,
            sigstack: sigstack as *mut c_void,
            limits: region.limits.clone(),
            region: Arc::downgrade(region) as Weak<dyn RegionInternal>,
        })
    }

    fn free_slot(slot: Slot) {
        // eprintln!(
        //     "unmapping {:p}[{:x}]",
        //     slot.start,
        //     slot.limits.total_memory_size()
        // );
        let res = unsafe { munmap(slot.start, slot.limits.total_memory_size()) };
        res.expect("munmap succeeded");
    }
}

// TODO: remove this once `nix` PR https://github.com/nix-rust/nix/pull/991 is merged
unsafe fn mprotect(addr: *mut c_void, length: libc::size_t, prot: ProtFlags) -> nix::Result<()> {
    nix::errno::Errno::result(libc::mprotect(addr, length, prot.bits())).map(drop)
}
<|endoftext|>


<commit_before>#![feature(macro_rules)]

//! Implementation of a Micro Transport Protocol library,
//! as well as a small client and server.
//!
//! http://www.bittorrent.org/beps/bep_0029.html

extern crate utp;
use std::io::net::ip::{Ipv4Addr, SocketAddr};

fn usage() {
    println!("Usage: utp [-s|-c] <address> <port>");
}

fn main() {
    use utp::{UtpStream, UtpSocket};
    use std::from_str::FromStr;

    // Defaults
    static BUF_SIZE: uint = 4096;
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

	if args.len() != 2 && args.len() != 4 {
		usage();
		return;
	}

    if args.len() == 4 {
        addr = SocketAddr {
            ip:   FromStr::from_str(args[2].as_slice()).unwrap(),
            port: FromStr::from_str(args[3].as_slice()).unwrap(),
        };
    }

    match args[1].as_slice() {
        "-s" => {
            let mut buf = [0, ..BUF_SIZE];
            let mut sock = UtpSocket::bind(addr).unwrap();
            println!("Serving on {}", addr);

            loop {
                match sock.recv_from(buf) {
                    Ok((_read, _src)) => {}
                    Err(e) => fail!("{}", e),
                }
            }
        }
        "-c" => {
            let mut stream = UtpStream::connect(addr);

            let buf = [0xF, ..BUF_SIZE];
            match stream.write(buf) {
                Ok(_) => {},
                Err(e) => fail!("{}", e),
            }

            let mut buf = [0u8, ..BUF_SIZE];
            match stream.read(buf) {
                Ok(n) => println!("{}", Vec::from_slice(buf.slice(0, n))),
                Err(e) => println!("{}", e),
            }

            match stream.close() {
                Ok(_) => {},
                Err(e) => fail!("{}", e),
            }
            drop(stream);
        }
        _ => usage(),
    }
}
<commit_msg>Update example code using exclusively UtpStream methods.<commit_after>#![feature(macro_rules)]

//! Implementation of a Micro Transport Protocol library,
//! as well as a small client and server.
//!
//! http://www.bittorrent.org/beps/bep_0029.html

extern crate utp;
use std::io::net::ip::{Ipv4Addr, SocketAddr};

fn usage() {
    println!("Usage: utp [-s|-c] <address> <port>");
}

fn main() {
    use utp::UtpStream;
    use std::from_str::FromStr;

    // Defaults
    static BUF_SIZE: uint = 4096;
    let mut addr = SocketAddr { ip: Ipv4Addr(127,0,0,1), port: 8080 };

    let args = std::os::args();

	if args.len() != 2 && args.len() != 4 {
		usage();
		return;
	}

    if args.len() == 4 {
        addr = SocketAddr {
            ip:   FromStr::from_str(args[2].as_slice()).unwrap(),
            port: FromStr::from_str(args[3].as_slice()).unwrap(),
        };
    }

    match args[1].as_slice() {
        "-s" => {
            let mut buf = [0, ..BUF_SIZE];
            let mut stream = UtpStream::bind(addr);

            println!("Serving on {}", addr);

            loop {
                match stream.read(buf) {
                    Ok(_read) => {}
                    Err(e) => fail!("{}", e),
                }
            }
        }
        "-c" => {
            let mut stream = UtpStream::connect(addr);

            let buf = [0xF, ..BUF_SIZE];
            match stream.write(buf) {
                Ok(_) => {},
                Err(e) => fail!("{}", e),
            }

            let mut buf = [0u8, ..BUF_SIZE];
            match stream.read(buf) {
                Ok(n) => println!("{}", Vec::from_slice(buf.slice(0, n))),
                Err(e) => println!("{}", e),
            }

            match stream.close() {
                Ok(_) => {},
                Err(e) => fail!("{}", e),
            }
            drop(stream);
        }
        _ => usage(),
    }
}
<|endoftext|>


<commit_before>use std::iter::{Iterator, Peekable};
use std::collections::btree_map::{BTreeMap, Range};
use std::cmp::Ordering::*;

use super::Result;

use self::NextIterValue::*;

pub type Patch = BTreeMap<Vec<u8>, Change>;
pub type Iter<'a> = Box<Iterator<Item = (&'a [u8], &'a [u8])> + 'a>;

#[derive(Debug, Clone)]
pub enum Change {
    Put(Vec<u8>),
    Delete,
}

pub struct Fork {
    snapshot: Box<Snapshot>,
    changes: Patch,
    changelog: Vec<(Vec<u8>, Option<Change>)>,
}

pub struct ForkIter<'a> {
    snapshot: Peekable<Iter<'a>>,
    changes: Peekable<Range<'a, Vec<u8>, Change>>,
}

#[derive(Debug)]
enum NextIterValue<'a> {
    Stored(&'a [u8], &'a [u8]),
    Replaced(&'a [u8], &'a [u8]),
    Inserted(&'a [u8], &'a [u8]),
    Deleted,
    MissDeleted,
}

pub trait Database: Send + Sync + 'static {
    fn clone(&self) -> Box<Database>;
    fn snapshot(&self) -> Box<Snapshot>;
    fn fork(&self) -> Fork {
        Fork {
            snapshot: self.snapshot(),
            changes: Patch::new(),
            changelog: Vec::new(),
        }
    }
    fn merge(&mut self, patch: Patch) -> Result<()>;
}

pub trait Snapshot {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>>;
    fn contains(&self, key: &[u8]) -> bool {
        self.get(key).is_some()
    }
    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a>;
}

impl Snapshot for Fork {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(ref v) => Some(v.clone()),
                    Change::Delete => None,
                }
            }
            None => self.snapshot.get(key),
        }
    }

    fn contains(&self, key: &[u8]) -> bool {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(..) => true,
                    Change::Delete => false,
                }
            }
            None => self.snapshot.get(key).is_some(),
        }
    }

    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a> {
        use std::collections::Bound::*;
        let range = (Included(from), Unbounded);
        Box::new(ForkIter {
                     snapshot: self.snapshot.iter(from).peekable(),
                     changes: self.changes.range::<[u8], _>(range).peekable(),
                 })
    }
}

impl Fork {
    pub fn checkpoint(&mut self) {
        self.changelog.clear()
    }

    pub fn rollback(&mut self) {
        for (k, c) in self.changelog.drain(..).rev() {
            match c {
                Some(change) => self.changes.insert(k, change),
                None => self.changes.remove(&k),
            };
        }
    }

    pub fn put(&mut self, key: Vec<u8>, value: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Put(value))))
    }

    pub fn remove(&mut self, key: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Delete)));
    }

    pub fn remove_by_prefix(&mut self, prefix: &[u8]) {
        for (k, _) in self.snapshot.iter(prefix) {
            if !k.starts_with(prefix) {
                return;
            }
            self.changes.insert(k.to_vec(), Change::Delete);
        }
    }

    pub fn into_patch(self) -> Patch {
        self.changes
    }

    pub fn merge(&mut self, patch: Patch) {
        self.changes.extend(patch)
    }
}

// TODO: Does we needed this AsRef / AsMut impls?

impl AsRef<Snapshot> for Snapshot {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        self
    }
}

impl AsRef<Snapshot + 'static> for Fork {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        &*self
    }
}

impl ::std::fmt::Debug for Fork {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "Fork(..)")
    }
}

impl<'a> NextIterValue<'a> {
    fn skip_changes(&self) -> bool {
        match *self {
            Replaced(..) | Inserted(..) | Deleted | MissDeleted => true,
            Stored(..) => false,
        }
    }

    fn skip_snapshot(&self) -> bool {
        match *self {
            Stored(..) | Replaced(..) | Deleted => true,
            Inserted(..) | MissDeleted => false,
        }
    }

    fn value(&self) -> Option<(&'a [u8], &'a [u8])> {
        match *self {
            Stored(k, v) => Some((k, v)),
            Replaced(k, v) => Some((k, v)),
            Inserted(k, v) => Some((k, v)),
            Deleted | MissDeleted => None,
        }
    }
}

impl<'a> Iterator for ForkIter<'a> {
    type Item = (&'a [u8], &'a [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let next = match self.changes.peek() {
                Some(&(k, ref change)) => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => {
                            match **change {
                                Change::Put(ref v) => {
                                    match k[..].cmp(key) {
                                        Equal => Replaced(k, v),
                                        Less => Inserted(k, v),
                                        Greater => Stored(key, value),
                                    }
                                }
                                Change::Delete => {
                                    match k[..].cmp(key) {
                                        Equal => Deleted,
                                        Less => MissDeleted,
                                        Greater => Stored(key, value),
                                    }
                                }
                            }
                        }
                        None => {
                            match **change {
                                Change::Put(ref v) => Inserted(k, v),
                                Change::Delete => MissDeleted,
                            }
                        }
                    }
                }
                None => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => Stored(key, value),
                        None => return None,
                    }
                }
            };
            if next.skip_changes() {
                self.changes.next();
            }
            if next.skip_snapshot() {
                self.snapshot.next();
            }
            if let Some(value) = next.value() {
                return Some(value);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::{MemoryDB, Database, Snapshot, Fork};

    #[test]
    fn fork_iter() {
        let mut db = MemoryDB::new();
        let mut fork = db.fork();

        fork.put(vec![10], vec![10]);
        fork.put(vec![20], vec![20]);
        fork.put(vec![30], vec![30]);

        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();

        fn assert_iter(fork: &Fork, from: u8, assumed: &[(u8, u8)]) {
            assert_eq!(fork.iter(&[from])
                           .map(|(k, v)| (k[0], v[0]))
                           .collect::<Vec<_>>(),
                       assumed);
        }

        // Stored
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 5, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 10, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 11, &[(20, 20), (30, 30)]);
        assert_iter(&fork, 31, &[]);

        // Inserted
        fork.put(vec![5], vec![5]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (30, 30)]);
        fork.put(vec![25], vec![25]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30)]);
        fork.put(vec![35], vec![35]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30), (35, 35)]);

        // Double inserted
        fork.put(vec![25], vec![23]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (30, 30), (35, 35)]);
        fork.put(vec![26], vec![26]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (26, 26), (30, 30), (35, 35)]);

        // Replaced
        let mut fork = db.fork();
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 30)]);
        fork.put(vec![30], vec![31]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 31)]);

        // Deleted
        let mut fork = db.fork();
        fork.remove(vec![20]);
        assert_iter(&fork, 0, &[(10, 10), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);

        // MissDeleted
        let mut fork = db.fork();
        fork.remove(vec![5]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![15]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![35]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
    }
}
<commit_msg>Panic when trying to merge into unfinalized fork<commit_after>use std::iter::{Iterator, Peekable};
use std::collections::btree_map::{BTreeMap, Range};
use std::cmp::Ordering::*;

use super::Result;

use self::NextIterValue::*;

pub type Patch = BTreeMap<Vec<u8>, Change>;
pub type Iter<'a> = Box<Iterator<Item = (&'a [u8], &'a [u8])> + 'a>;

#[derive(Debug, Clone)]
pub enum Change {
    Put(Vec<u8>),
    Delete,
}

pub struct Fork {
    snapshot: Box<Snapshot>,
    changes: Patch,
    changelog: Vec<(Vec<u8>, Option<Change>)>,
}

pub struct ForkIter<'a> {
    snapshot: Peekable<Iter<'a>>,
    changes: Peekable<Range<'a, Vec<u8>, Change>>,
}

#[derive(Debug)]
enum NextIterValue<'a> {
    Stored(&'a [u8], &'a [u8]),
    Replaced(&'a [u8], &'a [u8]),
    Inserted(&'a [u8], &'a [u8]),
    Deleted,
    MissDeleted,
}

pub trait Database: Send + Sync + 'static {
    fn clone(&self) -> Box<Database>;
    fn snapshot(&self) -> Box<Snapshot>;
    fn fork(&self) -> Fork {
        Fork {
            snapshot: self.snapshot(),
            changes: Patch::new(),
            changelog: Vec::new(),
        }
    }
    fn merge(&mut self, patch: Patch) -> Result<()>;
}

pub trait Snapshot {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>>;
    fn contains(&self, key: &[u8]) -> bool {
        self.get(key).is_some()
    }
    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a>;
}

impl Snapshot for Fork {
    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(ref v) => Some(v.clone()),
                    Change::Delete => None,
                }
            }
            None => self.snapshot.get(key),
        }
    }

    fn contains(&self, key: &[u8]) -> bool {
        match self.changes.get(key) {
            Some(change) => {
                match *change {
                    Change::Put(..) => true,
                    Change::Delete => false,
                }
            }
            None => self.snapshot.get(key).is_some(),
        }
    }

    fn iter<'a>(&'a self, from: &[u8]) -> Iter<'a> {
        use std::collections::Bound::*;
        let range = (Included(from), Unbounded);
        Box::new(ForkIter {
                     snapshot: self.snapshot.iter(from).peekable(),
                     changes: self.changes.range::<[u8], _>(range).peekable(),
                 })
    }
}

impl Fork {
    pub fn checkpoint(&mut self) {
        self.changelog.clear()
    }

    pub fn rollback(&mut self) {
        for (k, c) in self.changelog.drain(..).rev() {
            match c {
                Some(change) => self.changes.insert(k, change),
                None => self.changes.remove(&k),
            };
        }
    }

    pub fn put(&mut self, key: Vec<u8>, value: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Put(value))))
    }

    pub fn remove(&mut self, key: Vec<u8>) {
        self.changelog
            .push((key.clone(), self.changes.insert(key, Change::Delete)));
    }

    pub fn remove_by_prefix(&mut self, prefix: &[u8]) {
        for (k, _) in self.snapshot.iter(prefix) {
            if !k.starts_with(prefix) {
                return;
            }
            self.changes.insert(k.to_vec(), Change::Delete);
        }
    }

    pub fn into_patch(self) -> Patch {
        self.changes
    }

    pub fn merge(&mut self, patch: Patch) {
        if !self.changelog.is_empty() {
            panic!("merge into a fork is impossible because it has unfinalized changes");
        }
        self.changes.extend(patch)
    }
}

// TODO: Does we needed this AsRef / AsMut impls?

impl AsRef<Snapshot> for Snapshot {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        self
    }
}

impl AsRef<Snapshot + 'static> for Fork {
    fn as_ref<'a>(&'a self) -> &'a (Snapshot + 'static) {
        &*self
    }
}

impl ::std::fmt::Debug for Fork {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "Fork(..)")
    }
}

impl<'a> NextIterValue<'a> {
    fn skip_changes(&self) -> bool {
        match *self {
            Replaced(..) | Inserted(..) | Deleted | MissDeleted => true,
            Stored(..) => false,
        }
    }

    fn skip_snapshot(&self) -> bool {
        match *self {
            Stored(..) | Replaced(..) | Deleted => true,
            Inserted(..) | MissDeleted => false,
        }
    }

    fn value(&self) -> Option<(&'a [u8], &'a [u8])> {
        match *self {
            Stored(k, v) => Some((k, v)),
            Replaced(k, v) => Some((k, v)),
            Inserted(k, v) => Some((k, v)),
            Deleted | MissDeleted => None,
        }
    }
}

impl<'a> Iterator for ForkIter<'a> {
    type Item = (&'a [u8], &'a [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let next = match self.changes.peek() {
                Some(&(k, ref change)) => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => {
                            match **change {
                                Change::Put(ref v) => {
                                    match k[..].cmp(key) {
                                        Equal => Replaced(k, v),
                                        Less => Inserted(k, v),
                                        Greater => Stored(key, value),
                                    }
                                }
                                Change::Delete => {
                                    match k[..].cmp(key) {
                                        Equal => Deleted,
                                        Less => MissDeleted,
                                        Greater => Stored(key, value),
                                    }
                                }
                            }
                        }
                        None => {
                            match **change {
                                Change::Put(ref v) => Inserted(k, v),
                                Change::Delete => MissDeleted,
                            }
                        }
                    }
                }
                None => {
                    match self.snapshot.peek() {
                        Some(&(key, ref value)) => Stored(key, value),
                        None => return None,
                    }
                }
            };
            if next.skip_changes() {
                self.changes.next();
            }
            if next.skip_snapshot() {
                self.snapshot.next();
            }
            if let Some(value) = next.value() {
                return Some(value);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::{MemoryDB, Database, Snapshot, Fork};

    #[test]
    fn fork_iter() {
        let mut db = MemoryDB::new();
        let mut fork = db.fork();

        fork.put(vec![10], vec![10]);
        fork.put(vec![20], vec![20]);
        fork.put(vec![30], vec![30]);

        db.merge(fork.into_patch()).unwrap();

        let mut fork = db.fork();

        fn assert_iter(fork: &Fork, from: u8, assumed: &[(u8, u8)]) {
            assert_eq!(fork.iter(&[from])
                           .map(|(k, v)| (k[0], v[0]))
                           .collect::<Vec<_>>(),
                       assumed);
        }

        // Stored
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 5, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 10, &[(10, 10), (20, 20), (30, 30)]);
        assert_iter(&fork, 11, &[(20, 20), (30, 30)]);
        assert_iter(&fork, 31, &[]);

        // Inserted
        fork.put(vec![5], vec![5]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (30, 30)]);
        fork.put(vec![25], vec![25]);
        assert_iter(&fork, 0, &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30)]);
        fork.put(vec![35], vec![35]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 25), (30, 30), (35, 35)]);

        // Double inserted
        fork.put(vec![25], vec![23]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (30, 30), (35, 35)]);
        fork.put(vec![26], vec![26]);
        assert_iter(&fork,
                    0,
                    &[(5, 5), (10, 10), (20, 20), (25, 23), (26, 26), (30, 30), (35, 35)]);

        // Replaced
        let mut fork = db.fork();
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 30)]);
        fork.put(vec![30], vec![31]);
        assert_iter(&fork, 0, &[(10, 11), (20, 20), (30, 31)]);

        // Deleted
        let mut fork = db.fork();
        fork.remove(vec![20]);
        assert_iter(&fork, 0, &[(10, 10), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);
        fork.put(vec![10], vec![11]);
        assert_iter(&fork, 0, &[(10, 11), (30, 30)]);
        fork.remove(vec![10]);
        assert_iter(&fork, 0, &[(30, 30)]);

        // MissDeleted
        let mut fork = db.fork();
        fork.remove(vec![5]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![15]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
        fork.remove(vec![35]);
        assert_iter(&fork, 0, &[(10, 10), (20, 20), (30, 30)]);
    }
}
<|endoftext|>


<commit_before>use color::{Color, PuyoColor, RealColor};
use field;
use field_checker::FieldChecker;
use position::Position;

pub struct PlainField<C: Color<C>> {
    field: [[C; field::MAP_HEIGHT]; field::MAP_WIDTH],
}

impl<C: Color<C>> PlainField<C> {
    pub fn new() -> PlainField<C> {
        let w = C::wall_color();
        let e = C::empty_color();
        PlainField {
            field: [
                //        4           8          12          16
                [w, w, w, w, w, w, w, w, w, w, w, w, w, w, w, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, w, w, w, w, w, w, w, w, w, w, w, w, w, w, w],
            ]
        }
    }

    pub fn from_str(s: &str) -> PlainField<C> {
        let mut field = Self::new();

        assert!(s.len() % 6 == 0);

        let mut cnt = 0;
        for b in s.bytes().rev() {
            let x = 6 - (cnt % 6);
            let y = (cnt / 6) + 1;
            field.set_color(x, y, C::from_byte(b));
            cnt += 1;
        }

        field
    }

    pub fn color(&self, x: usize, y: usize) -> C {
        // self.field[x as usize][y as usize]

        debug_assert!(x < 8);
        debug_assert!(y < 16);
        unsafe {
            *self.field.get_unchecked(x).get_unchecked(y)
        }
    }

    pub fn set_color(&mut self, x: usize, y: usize, color: C) {
        // self.field[x as usize][y as usize] = color

        debug_assert!(x < 8);
        debug_assert!(y < 16);
        unsafe {
            *self.field.get_unchecked_mut(x).get_unchecked_mut(y) = color
        }
    }

    pub fn is_empty(&self, x: usize, y: usize) -> bool {
        self.color(x, y) == C::empty_color()
    }

    pub fn is_color(&self, x: usize, y: usize, color: C) -> bool {
        self.color(x, y) == color
    }

    pub fn drop(&mut self) {
        for x in 1 .. 7 {
            let mut h = 1;
            for y in 1 .. 14 {
                if self.is_empty(x, y) {
                    continue;
                }
                let c = self.color(x, y);
                self.set_color(x, h, c);
                h += 1;
            }

            while h <= 13 {
                self.set_color(x, h, C::empty_color());
                h += 1;
            }
        }
    }

    pub fn fill_same_color_position(&self, x: usize, y: usize, c: C,
                                    head: &mut usize, queue: &mut [Position; 72],
                                    checker: &mut FieldChecker) {
        if y > field::HEIGHT {
            return;
        }

        let mut write_head = *head;
        let mut read_head = *head;

        debug_assert!(!checker.get(x, y));
        queue[write_head] = Position::new(x, y);
        write_head += 1;
        checker.set(x, y);

        while write_head != read_head {
            let p = queue[read_head];
            read_head += 1;

            if self.is_color(p.x + 1, p.y, c) && !checker.get(p.x + 1, p.y) {
                queue[write_head] = Position::new(p.x + 1, p.y);
                write_head += 1;
                checker.set(p.x + 1, p.y);
            }
            if self.is_color(p.x - 1, p.y, c) && !checker.get(p.x - 1, p.y) {
                queue[write_head] = Position::new(p.x - 1, p.y);
                write_head += 1;
                checker.set(p.x - 1, p.y);
            }
            if self.is_color(p.x, p.y + 1, c) && !checker.get(p.x, p.y + 1) && p.y + 1 <= field::HEIGHT {
                queue[write_head] = Position::new(p.x, p.y + 1);
                write_head += 1;
                checker.set(p.x, p.y + 1);
            }
            if self.is_color(p.x, p.y - 1, c) && !checker.get(p.x, p.y - 1) {
                queue[write_head] = Position::new(p.x, p.y - 1);
                write_head += 1;
                checker.set(p.x, p.y - 1);
            }
        }

        *head = write_head;
    }

    pub fn count_connected_puyos(&self, x: usize, y: usize) -> usize {
        let mut checker = FieldChecker::new();
        self.count_connected_puyos_with_checker(x, y, &mut checker)
    }

    pub fn count_connected_puyos_with_checker(&self, x: usize, y: usize, checker: &mut FieldChecker) -> usize {
        let mut positions = [Position::new(0, 0); 72];
        let mut head = 0;
        self.fill_same_color_position(x, y, self.color(x, y), &mut head, &mut positions, checker);
        head
    }
}

impl<C> PartialEq<PlainField<C>> for PlainField<C>
where C: Color<C> {
    fn eq(&self, other: &PlainField<C>) -> bool {
        // TODO(mayah): Would be good memory comparison.
        for x in 0..8 {
            for y in 0..16 {
                if self.color(x, y) != other.color(x, y) {
                    return false
                }
            }
        }

        true
    }
}

pub type PuyoPlainField = PlainField<PuyoColor>;
pub type RealPlainField = PlainField<RealColor>;

#[cfg(test)]
mod tests {
    use std::mem;
    use color::PuyoColor;
    use plain_field::PuyoPlainField;

    #[test]
    fn test_memory_size() {
        assert_eq!(mem::size_of::<PuyoPlainField>(), 128);
    }

    #[test]
    fn test_color() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert_eq!(pf.color(0, 1), PuyoColor::WALL);
        assert_eq!(pf.color(1, 1), PuyoColor::RED);
        assert_eq!(pf.color(2, 1), PuyoColor::GREEN);
        assert_eq!(pf.color(3, 1), PuyoColor::YELLOW);
        assert_eq!(pf.color(4, 1), PuyoColor::BLUE);
        assert_eq!(pf.color(5, 1), PuyoColor::OJAMA);
        assert_eq!(pf.color(6, 1), PuyoColor::IRON);
        assert_eq!(pf.color(7, 1), PuyoColor::WALL);

        assert_eq!(pf.color(1, 2), PuyoColor::EMPTY);
    }

    #[test]
    fn test_is_color() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert!(pf.is_color(0, 1, PuyoColor::WALL));
        assert!(pf.is_color(1, 1, PuyoColor::RED));
        assert!(pf.is_color(2, 1, PuyoColor::GREEN));
        assert!(pf.is_color(3, 1, PuyoColor::YELLOW));
        assert!(pf.is_color(4, 1, PuyoColor::BLUE));
        assert!(pf.is_color(5, 1, PuyoColor::OJAMA));
        assert!(pf.is_color(6, 1, PuyoColor::IRON));
        assert!(pf.is_color(7, 1, PuyoColor::WALL));

        assert!(pf.is_color(1, 2, PuyoColor::EMPTY));
    }

    #[test]
    fn test_is_empty() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert!(pf.is_empty(1, 2));
        assert!(!pf.is_empty(1, 1));
    }

    #[test]
    fn test_drop() {
        let mut pf = PuyoPlainField::from_str(concat!(
            "RRRBBB",
            "......",
            "RRRBBB",
            "......",
            "RRRBBB"));
        let expected = PuyoPlainField::from_str(concat!(
            "RRRBBB",
            "RRRBBB",
            "RRRBBB"));

        pf.drop();
        assert!(pf == expected);
    }

    #[test]
    fn test_drop_extreme() {
        let mut pf = PuyoPlainField::from_str(concat!(
            "OOOOOO", // 14
            "OOOOOO", // 13
            "......", // 12
            "......",
            "......",
            "......",
            "......", // 8
            "......",
            "......",
            "......",
            "......", // 4
            "......",
            "......",
            "......"));

        let expected = PuyoPlainField::from_str(concat!(
            "OOOOOO", // 14
            "......", // 13
            "......", // 12
            "......",
            "......",
            "......",
            "......", // 8
            "......",
            "......",
            "......",
            "......", // 4
            "......",
            "......",
            "OOOOOO"));

        pf.drop();
        assert!(pf == expected);
    }

    #[test]
    fn test_count_connected_puyos() {
        // I O S Z L J T
        let fi = PuyoPlainField::from_str(concat!(
            "R.....",
            "R.....",
            "R.....",
            "R.....",
            "......",
            "RRRR.."));
        let fo = PuyoPlainField::from_str(concat!(
            "RR....",
            "RR...."));
        let fs = PuyoPlainField::from_str(concat!(
            "....R.",
            ".RR.RR",
            "RR...R"));
        let fz = PuyoPlainField::from_str(concat!(
            ".....R",
            "RR..RR",
            ".RR.R."));
        let fl = PuyoPlainField::from_str(concat!(
            "RR....",
            ".R...R",
            ".R.RRR",
            "R.....",
            "R..RRR",
            "RR.R.."));
        let fj = PuyoPlainField::from_str(concat!(
            "RR....",
            "R..R..",
            "R..RRR",
            ".R....",
            ".R.RRR",
            "RR...R"));
        let ft = PuyoPlainField::from_str(concat!(
            ".R....",
            "RR..R.",
            ".R.RRR",
            "R.....",
            "RR.RRR",
            "R...R."));

        let fields = [fi, fo, fs, fz, fl, fj, ft];
        for pf in fields.iter() {
            for x in 1..7 {
                for y in 1..13 {
                    if !pf.is_color(x, y, PuyoColor::RED) {
                        continue;
                    }

                    assert_eq!(4, pf.count_connected_puyos(x, y));
                }
            }
        }
    }
}
<commit_msg>Add PlainField::count_connected_puyos_max4<commit_after>use color::{Color, PuyoColor, RealColor};
use field;
use field_checker::FieldChecker;
use position::Position;

pub struct PlainField<C: Color<C>> {
    field: [[C; field::MAP_HEIGHT]; field::MAP_WIDTH],
}

impl<C: Color<C>> PlainField<C> {
    pub fn new() -> PlainField<C> {
        let w = C::wall_color();
        let e = C::empty_color();
        PlainField {
            field: [
                //        4           8          12          16
                [w, w, w, w, w, w, w, w, w, w, w, w, w, w, w, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, e, e, e, e, e, e, e, e, e, e, e, e, e, e, w],
                [w, w, w, w, w, w, w, w, w, w, w, w, w, w, w, w],
            ]
        }
    }

    pub fn from_str(s: &str) -> PlainField<C> {
        let mut field = Self::new();

        assert!(s.len() % 6 == 0);

        let mut cnt = 0;
        for b in s.bytes().rev() {
            let x = 6 - (cnt % 6);
            let y = (cnt / 6) + 1;
            field.set_color(x, y, C::from_byte(b));
            cnt += 1;
        }

        field
    }

    pub fn color(&self, x: usize, y: usize) -> C {
        // self.field[x as usize][y as usize]

        debug_assert!(x < 8);
        debug_assert!(y < 16);
        unsafe {
            *self.field.get_unchecked(x).get_unchecked(y)
        }
    }

    pub fn set_color(&mut self, x: usize, y: usize, color: C) {
        // self.field[x as usize][y as usize] = color

        debug_assert!(x < 8);
        debug_assert!(y < 16);
        unsafe {
            *self.field.get_unchecked_mut(x).get_unchecked_mut(y) = color
        }
    }

    pub fn is_empty(&self, x: usize, y: usize) -> bool {
        self.color(x, y) == C::empty_color()
    }

    pub fn is_color(&self, x: usize, y: usize, color: C) -> bool {
        self.color(x, y) == color
    }

    pub fn drop(&mut self) {
        for x in 1 .. 7 {
            let mut h = 1;
            for y in 1 .. 14 {
                if self.is_empty(x, y) {
                    continue;
                }
                let c = self.color(x, y);
                self.set_color(x, h, c);
                h += 1;
            }

            while h <= 13 {
                self.set_color(x, h, C::empty_color());
                h += 1;
            }
        }
    }

    pub fn fill_same_color_position(&self, x: usize, y: usize, c: C,
                                    head: &mut usize, queue: &mut [Position; 72],
                                    checker: &mut FieldChecker) {
        if y > field::HEIGHT {
            return;
        }

        let mut write_head = *head;
        let mut read_head = *head;

        debug_assert!(!checker.get(x, y));
        queue[write_head] = Position::new(x, y);
        write_head += 1;
        checker.set(x, y);

        while write_head != read_head {
            let p = queue[read_head];
            read_head += 1;

            if self.is_color(p.x + 1, p.y, c) && !checker.get(p.x + 1, p.y) {
                queue[write_head] = Position::new(p.x + 1, p.y);
                write_head += 1;
                checker.set(p.x + 1, p.y);
            }
            if self.is_color(p.x - 1, p.y, c) && !checker.get(p.x - 1, p.y) {
                queue[write_head] = Position::new(p.x - 1, p.y);
                write_head += 1;
                checker.set(p.x - 1, p.y);
            }
            if self.is_color(p.x, p.y + 1, c) && !checker.get(p.x, p.y + 1) && p.y + 1 <= field::HEIGHT {
                queue[write_head] = Position::new(p.x, p.y + 1);
                write_head += 1;
                checker.set(p.x, p.y + 1);
            }
            if self.is_color(p.x, p.y - 1, c) && !checker.get(p.x, p.y - 1) {
                queue[write_head] = Position::new(p.x, p.y - 1);
                write_head += 1;
                checker.set(p.x, p.y - 1);
            }
        }

        *head = write_head;
    }

    pub fn count_connected_puyos(&self, x: usize, y: usize) -> usize {
        let mut checker = FieldChecker::new();
        self.count_connected_puyos_with_checker(x, y, &mut checker)
    }

    pub fn count_connected_puyos_with_checker(&self, x: usize, y: usize, checker: &mut FieldChecker) -> usize {
        let mut positions = [Position::new(0, 0); 72];
        let mut head = 0;
        self.fill_same_color_position(x, y, self.color(x, y), &mut head, &mut positions, checker);
        head
    }

    pub fn count_connected_puyos_max4(&self, x: usize, y: usize) -> usize {
        if y > field::HEIGHT {
            return 0;
        }

        let mut left_up = false;
        let mut left_down = false;
        let mut right_up = false;
        let mut right_down = false;

        let mut cnt = 1;
        let c = self.color(x, y);

        if self.is_color(x - 1, y, c) {
            if self.is_color(x - 2, y, c) {
                if self.is_color(x - 3, y, c) {
                    return 4;
                }
                if self.is_color(x - 2, y + 1, c) && y + 1 <= field::HEIGHT {
                    return 4;
                }
                if self.is_color(x - 2, y - 1, c) {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x - 1, y + 1, c) && y + 1 <= field::HEIGHT {
                if self.is_color(x - 2, y + 1, c) {
                    return 4;
                }
                if self.is_color(x - 1, y + 2, c) && y + 2 <= field::HEIGHT {
                    return 4;
                }
                cnt += 1;
                left_up = true;
            }
            if self.is_color(x - 1, y - 1, c) {
                if self.is_color(x - 2, y - 1, c) || self.is_color(x - 1, y - 2, c) {
                    return 4;
                }
                cnt += 1;
                left_down = true;
            }
            cnt += 1;
        }

        if self.is_color(x + 1, y, c) {
            if self.is_color(x + 2, y, c) {
                if self.is_color(x + 3, y, c) {
                    return 4;
                }
                if self.is_color(x + 2, y + 1, c) && y + 1 <= field::HEIGHT {
                    return 4;
                }
                if self.is_color(x + 2, y - 1, c) {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x + 1, y + 1, c) && y + 1 <= field::HEIGHT {
                if self.is_color(x + 2, y + 1, c) {
                    return 4;
                }
                if self.is_color(x + 1, y + 2, c) && y + 2 <= field::HEIGHT {
                    return 4;
                }
                cnt += 1;
                right_up = true;
            }
            if self.is_color(x + 1, y - 1, c) {
                if self.is_color(x + 2, y - 1, c) {
                    return 4;
                }
                if self.is_color(x + 1, y - 2, c) {
                    return 4;
                }
                cnt += 1;
                right_down = true;
            }
            cnt += 1;
        }

        if self.is_color(x, y - 1, c) {
            if self.is_color(x, y - 2, c) {
                if self.is_color(x, y - 3, c) {
                    return 4;
                }
                if self.is_color(x - 1, y - 2, c) {
                    return 4;
                }
                if self.is_color(x + 1, y - 2, c) {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x - 1, y - 1, c) && !left_down {
                if self.is_color(x - 2, y - 1, c) {
                    return 4;
                }
                if self.is_color(x - 1, y - 2, c) {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x + 1, y - 1, c) && !right_down {
                if self.is_color(x + 2, y - 1, c) {
                    return 4;
                }
                if self.is_color(x + 1, y - 2, c) {
                    return 4;
                }
                cnt += 1;
            }
            cnt += 1;
        }

        if self.is_color(x, y + 1, c) && y + 1 <= field::HEIGHT {
            if self.is_color(x, y + 2, c) && y + 2 <= field::HEIGHT {
                if self.is_color(x, y + 3, c) && y + 3 <= field::HEIGHT {
                    return 4;
                }
                if self.is_color(x - 1, y + 2, c) {
                    return 4;
                }
                if self.is_color(x + 1, y + 2, c) {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x - 1, y + 1, c) && !left_up {
                if self.is_color(x - 2, y + 1, c) {
                    return 4;
                }
                if self.is_color(x - 1, y + 2, c) && y + 2 <= field::HEIGHT {
                    return 4;
                }
                cnt += 1;
            }
            if self.is_color(x + 1, y + 1, c) && !right_up {
                if self.is_color(x + 2, y + 1, c) {
                    return 4;
                }
                if self.is_color(x + 1, y + 2, c) && y + 2 <= field::HEIGHT {
                    return 4;
                }
                cnt += 1;
            }
            cnt += 1;
        }

        cnt
    }
}

impl<C> PartialEq<PlainField<C>> for PlainField<C>
where C: Color<C> {
    fn eq(&self, other: &PlainField<C>) -> bool {
        // TODO(mayah): Would be good memory comparison.
        for x in 0..8 {
            for y in 0..16 {
                if self.color(x, y) != other.color(x, y) {
                    return false
                }
            }
        }

        true
    }
}

pub type PuyoPlainField = PlainField<PuyoColor>;
pub type RealPlainField = PlainField<RealColor>;

#[cfg(test)]
mod tests {
    use std::mem;
    use color::PuyoColor;
    use plain_field::PuyoPlainField;

    #[test]
    fn test_memory_size() {
        assert_eq!(mem::size_of::<PuyoPlainField>(), 128);
    }

    #[test]
    fn test_color() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert_eq!(pf.color(0, 1), PuyoColor::WALL);
        assert_eq!(pf.color(1, 1), PuyoColor::RED);
        assert_eq!(pf.color(2, 1), PuyoColor::GREEN);
        assert_eq!(pf.color(3, 1), PuyoColor::YELLOW);
        assert_eq!(pf.color(4, 1), PuyoColor::BLUE);
        assert_eq!(pf.color(5, 1), PuyoColor::OJAMA);
        assert_eq!(pf.color(6, 1), PuyoColor::IRON);
        assert_eq!(pf.color(7, 1), PuyoColor::WALL);

        assert_eq!(pf.color(1, 2), PuyoColor::EMPTY);
    }

    #[test]
    fn test_is_color() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert!(pf.is_color(0, 1, PuyoColor::WALL));
        assert!(pf.is_color(1, 1, PuyoColor::RED));
        assert!(pf.is_color(2, 1, PuyoColor::GREEN));
        assert!(pf.is_color(3, 1, PuyoColor::YELLOW));
        assert!(pf.is_color(4, 1, PuyoColor::BLUE));
        assert!(pf.is_color(5, 1, PuyoColor::OJAMA));
        assert!(pf.is_color(6, 1, PuyoColor::IRON));
        assert!(pf.is_color(7, 1, PuyoColor::WALL));

        assert!(pf.is_color(1, 2, PuyoColor::EMPTY));
    }

    #[test]
    fn test_is_empty() {
        let pf = PuyoPlainField::from_str("RGYB@&");
        assert!(pf.is_empty(1, 2));
        assert!(!pf.is_empty(1, 1));
    }

    #[test]
    fn test_drop() {
        let mut pf = PuyoPlainField::from_str(concat!(
            "RRRBBB",
            "......",
            "RRRBBB",
            "......",
            "RRRBBB"));
        let expected = PuyoPlainField::from_str(concat!(
            "RRRBBB",
            "RRRBBB",
            "RRRBBB"));

        pf.drop();
        assert!(pf == expected);
    }

    #[test]
    fn test_drop_extreme() {
        let mut pf = PuyoPlainField::from_str(concat!(
            "OOOOOO", // 14
            "OOOOOO", // 13
            "......", // 12
            "......",
            "......",
            "......",
            "......", // 8
            "......",
            "......",
            "......",
            "......", // 4
            "......",
            "......",
            "......"));

        let expected = PuyoPlainField::from_str(concat!(
            "OOOOOO", // 14
            "......", // 13
            "......", // 12
            "......",
            "......",
            "......",
            "......", // 8
            "......",
            "......",
            "......",
            "......", // 4
            "......",
            "......",
            "OOOOOO"));

        pf.drop();
        assert!(pf == expected);
    }

    #[test]
    fn test_count_connected_puyos() {
        // I O S Z L J T
        let fi = PuyoPlainField::from_str(concat!(
            "R.....",
            "R.....",
            "R.....",
            "R.....",
            "......",
            "RRRR.."));
        let fo = PuyoPlainField::from_str(concat!(
            "RR....",
            "RR...."));
        let fs = PuyoPlainField::from_str(concat!(
            "....R.",
            ".RR.RR",
            "RR...R"));
        let fz = PuyoPlainField::from_str(concat!(
            ".....R",
            "RR..RR",
            ".RR.R."));
        let fl = PuyoPlainField::from_str(concat!(
            "RR....",
            ".R...R",
            ".R.RRR",
            "R.....",
            "R..RRR",
            "RR.R.."));
        let fj = PuyoPlainField::from_str(concat!(
            "RR....",
            "R..R..",
            "R..RRR",
            ".R....",
            ".R.RRR",
            "RR...R"));
        let ft = PuyoPlainField::from_str(concat!(
            ".R....",
            "RR..R.",
            ".R.RRR",
            "R.....",
            "RR.RRR",
            "R...R."));

        let fields = [fi, fo, fs, fz, fl, fj, ft];
        for pf in fields.iter() {
            for x in 1..7 {
                for y in 1..13 {
                    if !pf.is_color(x, y, PuyoColor::RED) {
                        continue;
                    }

                    assert_eq!(4, pf.count_connected_puyos(x, y));
                    assert_eq!(4, pf.count_connected_puyos_max4(x, y));
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use crate::shell::Key;

/// Contains the state information for the keyboard.
///
/// This currently tracks which keys are currently pressed.
///
/// The key state is stored in a lazy-loaded HashMap.
///
/// There are several convenience methods to check common modifiers (ctrl, shift, alt, etc).
/// This is useful if you don't care which shift key is down.
#[derive(Default, Clone, Debug, PartialEq)]
pub struct KeyboardState {
    key_list: HashMap<Key, bool>,
}

impl KeyboardState {
    /// Sets whether or not the given key is currently pressed
    pub fn set_key_state(&mut self, key: Key, pressed: bool) {
        self.key_list.insert(key, pressed);
    }

    /// Returns whether or not the requested key is pressed
    pub fn is_key_down(&self, key: Key) -> bool {
        match self.key_list.get(&key) {
            // If we have the key on this list, return its state
            Some(item) => *item,
            // Otherwise, it hasn't been set as down
            None => false,
        }
    }
    /// Returns whether or not any shift key is down.
    pub fn is_shift_down(&self) -> bool {
        self.is_key_down(Key::ShiftL) || self.is_key_down(Key::ShiftR)
    }

    /// Returns whether or not any alt key is down.
    pub fn is_alt_down(&self) -> bool {
        self.is_key_down(Key::Alt)
    }

    /// Returns whether or not any control key is down.
    pub fn is_ctrl_down(&self) -> bool {
        self.is_key_down(Key::Control)
    }

    /// Returns whether or not any home key is down.
    pub fn is_home_down(&self) -> bool {
        self.is_key_down(Key::Home)
    }
}

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

    #[test]
    /// A quick test to ensure that the items are properly set.
    fn basic_test() {
        let mut state = KeyboardState::default();
        // Do a quick check beforehand
        assert_eq!(state.is_key_down(Key::ShiftL), false);
        // Set the state down and immediately check
        state.set_key_state(Key::ShiftL, true);
        assert_eq!(state.is_key_down(Key::ShiftL), true);
        state.set_key_state(Key::ShiftL, false);
        assert_eq!(state.is_key_down(Key::ShiftL), false);
        // Set quit a few in a row
        state.set_key_state(Key::ShiftL, true);
        state.set_key_state(Key::ShiftR, true);
        state.set_key_state(Key::Space, true);
        state.set_key_state(Key::Control, true);
        state.set_key_state(Key::Alt, true);
        // Ensure each of these are still down
        assert_eq!(state.is_key_down(Key::ShiftL), true);
        assert_eq!(state.is_key_down(Key::ShiftR), true);
        assert_eq!(state.is_key_down(Key::Space), true);
        assert_eq!(state.is_key_down(Key::Control), true);
        assert_eq!(state.is_key_down(Key::Alt), true);
    }

    #[test]
    /// Test for the convenience methods
    fn test_convenience() {
        let mut state = KeyboardState::default();
        // Check to ensure they are all false
        assert_eq!(state.is_alt_down(), false);
        assert_eq!(state.is_ctrl_down(), false);
        assert_eq!(state.is_shift_down(), false);
        // Set ctrl and alt to true and check
        state.set_key_state(Key::Control, true);
        assert_eq!(state.is_ctrl_down(), true);

        state.set_key_state(Key::Alt, true);
        assert_eq!(state.is_alt_down(), true);

        assert_eq!(state.is_shift_down(), false);
        // Set shift (via L)
        state.set_key_state(Key::ShiftL, true);
        assert_eq!(state.is_shift_down(), true);
        // Set shift (via R and L both set)
        state.set_key_state(Key::ShiftR, true);
        assert_eq!(state.is_shift_down(), true);
        // Disable L Shift to ensure correct result with just R held
        state.set_key_state(Key::ShiftL, false);
        assert_eq!(state.is_shift_down(), true);
        // Disable both shift keys
        state.set_key_state(Key::ShiftR, false);
        assert_eq!(state.is_shift_down(), false);
        // Disable alt and ctrl and check again
        state.set_key_state(Key::Control, false);
        assert_eq!(state.is_ctrl_down(), false);

        state.set_key_state(Key::Alt, false);
        assert_eq!(state.is_alt_down(), false);
    }
}
<commit_msg>orbtk_core: update asset test for keyboard properties<commit_after>use std::collections::HashMap;

use crate::shell::Key;

/// Contains the state information for the keyboard.
///
/// This currently tracks which keys are currently pressed.
///
/// The key state is stored in a lazy-loaded HashMap.
///
/// There are several convenience methods to check common modifiers (ctrl, shift, alt, etc).
/// This is useful if you don't care which shift key is down.
#[derive(Default, Clone, Debug, PartialEq)]
pub struct KeyboardState {
    key_list: HashMap<Key, bool>,
}

impl KeyboardState {
    /// Sets whether or not the given key is currently pressed
    pub fn set_key_state(&mut self, key: Key, pressed: bool) {
        self.key_list.insert(key, pressed);
    }

    /// Returns whether or not the requested key is pressed
    pub fn is_key_down(&self, key: Key) -> bool {
        match self.key_list.get(&key) {
            // If we have the key on this list, return its state
            Some(item) => *item,
            // Otherwise, it hasn't been set as down
            None => false,
        }
    }
    /// Returns whether or not any shift key is down.
    pub fn is_shift_down(&self) -> bool {
        self.is_key_down(Key::ShiftL) || self.is_key_down(Key::ShiftR)
    }

    /// Returns whether or not any alt key is down.
    pub fn is_alt_down(&self) -> bool {
        self.is_key_down(Key::Alt)
    }

    /// Returns whether or not any control key is down.
    pub fn is_ctrl_down(&self) -> bool {
        self.is_key_down(Key::Control)
    }

    /// Returns whether or not any home key is down.
    pub fn is_home_down(&self) -> bool {
        self.is_key_down(Key::Home)
    }
}

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

    #[test]
    /// A quick test to ensure that the items are properly set.
    fn basic_test() {
        let mut state = KeyboardState::default();
        // Do a quick check beforehand
        assert!(!state.is_key_down(Key::ShiftL));
        // Set the state down and immediately check
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_key_down(Key::ShiftL));
        state.set_key_state(Key::ShiftL, false);
        assert!(!state.is_key_down(Key::ShiftL));
        // Set quit a few in a row
        state.set_key_state(Key::ShiftL, true);
        state.set_key_state(Key::ShiftR, true);
        state.set_key_state(Key::Space, true);
        state.set_key_state(Key::Control, true);
        state.set_key_state(Key::Alt, true);
        // Ensure each of these are still down
        assert!(state.is_key_down(Key::ShiftL));
        assert!(state.is_key_down(Key::ShiftR));
        assert!(state.is_key_down(Key::Space));
        assert!(state.is_key_down(Key::Control));
        assert!(state.is_key_down(Key::Alt));
    }

    #[test]
    /// Test for the convenience methods
    fn test_convenience() {
        let mut state = KeyboardState::default();
        // Check to ensure they are all false
        assert!(!state.is_alt_down());
        assert!(!state.is_ctrl_down());
        assert!(!state.is_shift_down());
        // Set ctrl and alt to true and check
        state.set_key_state(Key::Control, true);
        assert!(state.is_ctrl_down());

        state.set_key_state(Key::Alt, true);
        assert!(state.is_alt_down());

        assert!(!state.is_shift_down());
        // Set shift (via L)
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_shift_down());
        // Set shift (via R and L both set)
        state.set_key_state(Key::ShiftR, true);
        assert!(state.is_shift_down(), "{}", true);
        // Disable L Shift to ensure correct result with just R held
        state.set_key_state(Key::ShiftL, false);
        assert!(state.is_shift_down());
        // Disable both shift keys
        state.set_key_state(Key::ShiftR, false);
        assert!(!state.is_shift_down());
        // Disable alt and ctrl and check again
        state.set_key_state(Key::Control, false);
        assert!(!state.is_ctrl_down());

        state.set_key_state(Key::Alt, false);
        assert!(!state.is_alt_down());
    }
}
<|endoftext|>


<commit_before>use std::error::Error as StdError;
use std::fmt;

use futures::{Async, Future, IntoFuture, Poll};

use body::Payload;
use super::Service;

/// An asynchronous constructor of `Service`s.
pub trait MakeService<Ctx> {
    /// The `Payload` body of the `http::Request`.
    type ReqBody: Payload;

    /// The `Payload` body of the `http::Response`.
    type ResBody: Payload;

    /// The error type that can be returned by `Service`s.
    type Error: Into<Box<StdError + Send + Sync>>;

    /// The resolved `Service` from `new_service()`.
    type Service: Service<
        ReqBody=Self::ReqBody,
        ResBody=Self::ResBody,
        Error=Self::Error,
    >;

    /// The future returned from `new_service` of a `Service`.
    type Future: Future<Item=Self::Service, Error=Self::MakeError>;

    /// The error type that can be returned when creating a new `Service`.
    type MakeError: Into<Box<StdError + Send + Sync>>;

    /// Returns `Ready` when the constructor is ready to create a new `Service`.
    ///
    /// The implementation of this method is allowed to return a `Ready` even if
    /// the factory is not ready to create a new service. In this case, the future
    /// returned from `make_service` will resolve to an error.
    fn poll_ready(&mut self) -> Poll<(), Self::MakeError> {
        Ok(Async::Ready(()))
    }

    /// Create a new `Service`.
    fn make_service(&mut self, ctx: Ctx) -> Self::Future;
}

// Just a sort-of "trait alias" of `MakeService`, not to be implemented
// by anyone, only used as bounds.
#[doc(hidden)]
pub trait MakeServiceRef<Ctx>: self::sealed::Sealed<Ctx> {
    type ReqBody: Payload;
    type ResBody: Payload;
    type Error: Into<Box<StdError + Send + Sync>>;
    type Service: Service<
        ReqBody=Self::ReqBody,
        ResBody=Self::ResBody,
        Error=Self::Error,
    >;
    type MakeError: Into<Box<StdError + Send + Sync>>;
    type Future: Future<Item=Self::Service, Error=Self::MakeError>;

    // Acting like a #[non_exhaustive] for associated types of this trait.
    //
    // Basically, no one outside of hyper should be able to set this type
    // or declare bounds on it, so it should prevent people from creating
    // trait objects or otherwise writing code that requires using *all*
    // of the associated types.
    //
    // Why? So we can add new associated types to this alias in the future,
    // if necessary.
    type __DontNameMe: self::sealed::CantImpl;

    fn poll_ready_ref(&mut self) -> Poll<(), Self::MakeError>;

    fn make_service_ref(&mut self, ctx: &Ctx) -> Self::Future;
}

impl<T, Ctx, E, ME, S, F, IB, OB> MakeServiceRef<Ctx> for T
where
    T: for<'a> MakeService<&'a Ctx, Error=E, MakeError=ME, Service=S, Future=F, ReqBody=IB, ResBody=OB>,
    E: Into<Box<StdError + Send + Sync>>,
    ME: Into<Box<StdError + Send + Sync>>,
    S: Service<ReqBody=IB, ResBody=OB, Error=E>,
    F: Future<Item=S, Error=ME>,
    IB: Payload,
    OB: Payload,
{
    type Error = E;
    type Service = S;
    type ReqBody = IB;
    type ResBody = OB;
    type MakeError = ME;
    type Future = F;

    type __DontNameMe = self::sealed::CantName;

    fn poll_ready_ref(&mut self) -> Poll<(), Self::MakeError> {
        self.poll_ready()
    }

    fn make_service_ref(&mut self, ctx: &Ctx) -> Self::Future {
        self.make_service(ctx)
    }
}

impl<T, Ctx, E, ME, S, F, IB, OB> self::sealed::Sealed<Ctx> for T
where
    T: for<'a> MakeService<&'a Ctx, Error=E, MakeError=ME, Service=S, Future=F, ReqBody=IB, ResBody=OB>,
    E: Into<Box<StdError + Send + Sync>>,
    ME: Into<Box<StdError + Send + Sync>>,
    S: Service<ReqBody=IB, ResBody=OB, Error=E>,
    F: Future<Item=S, Error=ME>,
    IB: Payload,
    OB: Payload,
{}

/// Create a `MakeService` from a function.
///
/// # Example
///
/// ```rust
/// use std::net::TcpStream;
/// use hyper::{Body, Request, Response};
/// use hyper::service::{make_service_fn, service_fn_ok};
///
/// let make_svc = make_service_fn(|socket: &TcpStream| {
///     let remote_addr = socket.peer_addr().unwrap();
///     service_fn_ok(move |_: Request<Body>| {
///         Response::new(Body::from(format!("Hello, {}", remote_addr)))
///     })
/// });
/// ```
pub fn make_service_fn<F, Ctx, Ret>(f: F) -> MakeServiceFn<F>
where
    F: FnMut(&Ctx) -> Ret,
    Ret: IntoFuture,
{
    MakeServiceFn {
        f,
    }
}

// Not exported from crate as this will likely be replaced with `impl Service`.
pub struct MakeServiceFn<F> {
    f: F,
}

impl<'c, F, Ctx, Ret, ReqBody, ResBody> MakeService<&'c Ctx> for MakeServiceFn<F>
where
    F: FnMut(&Ctx) -> Ret,
    Ret: IntoFuture,
    Ret::Item: Service<ReqBody=ReqBody, ResBody=ResBody>,
    Ret::Error: Into<Box<StdError + Send + Sync>>,
    ReqBody: Payload,
    ResBody: Payload,
{
    type ReqBody = ReqBody;
    type ResBody = ResBody;
    type Error = <Ret::Item as Service>::Error;
    type Service = Ret::Item;
    type Future = Ret::Future;
    type MakeError = Ret::Error;

    fn make_service(&mut self, ctx: &'c Ctx) -> Self::Future {
        (self.f)(ctx).into_future()
    }
}

impl<F> fmt::Debug for MakeServiceFn<F> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("MakeServiceFn")
            .finish()
    }
}

mod sealed {
    pub trait Sealed<T> {}

    pub trait CantImpl {}

    #[allow(missing_debug_implementations)]
    pub enum CantName {}

    impl CantImpl for CantName {}
}
<commit_msg>docs(service): change the make_service_fn example to a full snippet<commit_after>use std::error::Error as StdError;
use std::fmt;

use futures::{Async, Future, IntoFuture, Poll};

use body::Payload;
use super::Service;

/// An asynchronous constructor of `Service`s.
pub trait MakeService<Ctx> {
    /// The `Payload` body of the `http::Request`.
    type ReqBody: Payload;

    /// The `Payload` body of the `http::Response`.
    type ResBody: Payload;

    /// The error type that can be returned by `Service`s.
    type Error: Into<Box<StdError + Send + Sync>>;

    /// The resolved `Service` from `new_service()`.
    type Service: Service<
        ReqBody=Self::ReqBody,
        ResBody=Self::ResBody,
        Error=Self::Error,
    >;

    /// The future returned from `new_service` of a `Service`.
    type Future: Future<Item=Self::Service, Error=Self::MakeError>;

    /// The error type that can be returned when creating a new `Service`.
    type MakeError: Into<Box<StdError + Send + Sync>>;

    /// Returns `Ready` when the constructor is ready to create a new `Service`.
    ///
    /// The implementation of this method is allowed to return a `Ready` even if
    /// the factory is not ready to create a new service. In this case, the future
    /// returned from `make_service` will resolve to an error.
    fn poll_ready(&mut self) -> Poll<(), Self::MakeError> {
        Ok(Async::Ready(()))
    }

    /// Create a new `Service`.
    fn make_service(&mut self, ctx: Ctx) -> Self::Future;
}

// Just a sort-of "trait alias" of `MakeService`, not to be implemented
// by anyone, only used as bounds.
#[doc(hidden)]
pub trait MakeServiceRef<Ctx>: self::sealed::Sealed<Ctx> {
    type ReqBody: Payload;
    type ResBody: Payload;
    type Error: Into<Box<StdError + Send + Sync>>;
    type Service: Service<
        ReqBody=Self::ReqBody,
        ResBody=Self::ResBody,
        Error=Self::Error,
    >;
    type MakeError: Into<Box<StdError + Send + Sync>>;
    type Future: Future<Item=Self::Service, Error=Self::MakeError>;

    // Acting like a #[non_exhaustive] for associated types of this trait.
    //
    // Basically, no one outside of hyper should be able to set this type
    // or declare bounds on it, so it should prevent people from creating
    // trait objects or otherwise writing code that requires using *all*
    // of the associated types.
    //
    // Why? So we can add new associated types to this alias in the future,
    // if necessary.
    type __DontNameMe: self::sealed::CantImpl;

    fn poll_ready_ref(&mut self) -> Poll<(), Self::MakeError>;

    fn make_service_ref(&mut self, ctx: &Ctx) -> Self::Future;
}

impl<T, Ctx, E, ME, S, F, IB, OB> MakeServiceRef<Ctx> for T
where
    T: for<'a> MakeService<&'a Ctx, Error=E, MakeError=ME, Service=S, Future=F, ReqBody=IB, ResBody=OB>,
    E: Into<Box<StdError + Send + Sync>>,
    ME: Into<Box<StdError + Send + Sync>>,
    S: Service<ReqBody=IB, ResBody=OB, Error=E>,
    F: Future<Item=S, Error=ME>,
    IB: Payload,
    OB: Payload,
{
    type Error = E;
    type Service = S;
    type ReqBody = IB;
    type ResBody = OB;
    type MakeError = ME;
    type Future = F;

    type __DontNameMe = self::sealed::CantName;

    fn poll_ready_ref(&mut self) -> Poll<(), Self::MakeError> {
        self.poll_ready()
    }

    fn make_service_ref(&mut self, ctx: &Ctx) -> Self::Future {
        self.make_service(ctx)
    }
}

impl<T, Ctx, E, ME, S, F, IB, OB> self::sealed::Sealed<Ctx> for T
where
    T: for<'a> MakeService<&'a Ctx, Error=E, MakeError=ME, Service=S, Future=F, ReqBody=IB, ResBody=OB>,
    E: Into<Box<StdError + Send + Sync>>,
    ME: Into<Box<StdError + Send + Sync>>,
    S: Service<ReqBody=IB, ResBody=OB, Error=E>,
    F: Future<Item=S, Error=ME>,
    IB: Payload,
    OB: Payload,
{}

/// Create a `MakeService` from a function.
///
/// # Example
///
/// ```rust,no_run
/// # #[cfg(feature = "runtime")] fn main() {
/// use std::net::TcpStream;
/// use hyper::{Body, Request, Response, Server};
/// use hyper::rt::{self, Future};
/// use hyper::server::conn::AddrStream;
/// use hyper::service::{make_service_fn, service_fn_ok};
///
/// let addr = ([127, 0, 0, 1], 3000).into();
///
/// let make_svc = make_service_fn(|socket: &AddrStream| {
///     let remote_addr = socket.remote_addr();
///     service_fn_ok(move |_: Request<Body>| {
///         Response::new(Body::from(format!("Hello, {}", remote_addr)))
///     })
/// });
///
/// // Then bind and serve...
/// let server = Server::bind(&addr)
///     .serve(make_svc);
///
/// // Finally, spawn `server` onto an Executor...
/// rt::run(server.map_err(|e| {
///     eprintln!("server error: {}", e);
/// }));
/// # }
/// # #[cfg(not(feature = "runtime"))] fn main() {}
/// ```
pub fn make_service_fn<F, Ctx, Ret>(f: F) -> MakeServiceFn<F>
where
    F: FnMut(&Ctx) -> Ret,
    Ret: IntoFuture,
{
    MakeServiceFn {
        f,
    }
}

// Not exported from crate as this will likely be replaced with `impl Service`.
pub struct MakeServiceFn<F> {
    f: F,
}

impl<'c, F, Ctx, Ret, ReqBody, ResBody> MakeService<&'c Ctx> for MakeServiceFn<F>
where
    F: FnMut(&Ctx) -> Ret,
    Ret: IntoFuture,
    Ret::Item: Service<ReqBody=ReqBody, ResBody=ResBody>,
    Ret::Error: Into<Box<StdError + Send + Sync>>,
    ReqBody: Payload,
    ResBody: Payload,
{
    type ReqBody = ReqBody;
    type ResBody = ResBody;
    type Error = <Ret::Item as Service>::Error;
    type Service = Ret::Item;
    type Future = Ret::Future;
    type MakeError = Ret::Error;

    fn make_service(&mut self, ctx: &'c Ctx) -> Self::Future {
        (self.f)(ctx).into_future()
    }
}

impl<F> fmt::Debug for MakeServiceFn<F> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("MakeServiceFn")
            .finish()
    }
}

mod sealed {
    pub trait Sealed<T> {}

    pub trait CantImpl {}

    #[allow(missing_debug_implementations)]
    pub enum CantName {}

    impl CantImpl for CantName {}
}
<|endoftext|>