<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(missing_copy_implementations)]

use fmt;
use io::{self, Read, Initializer, Write, ErrorKind, BufRead};
use mem;

/// Copies the entire contents of a reader into a writer.
///
/// This function will continuously read data from `reader` and then
/// write it into `writer` in a streaming fashion until `reader`
/// returns EOF.
///
/// On success, the total number of bytes that were copied from
/// `reader` to `writer` is returned.
///
/// # Errors
///
/// This function will return an error immediately if any call to `read` or
/// `write` returns an error. All instances of `ErrorKind::Interrupted` are
/// handled by this function and the underlying operation is retried.
///
/// # Examples
///
/// ```
/// use std::io;
///
/// # fn foo() -> io::Result<()> {
/// let mut reader: &[u8] = b"hello";
/// let mut writer: Vec<u8> = vec![];
///
/// io::copy(&mut reader, &mut writer)?;
///
/// assert_eq!(reader, &writer[..]);
/// # Ok(())
/// # }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = unsafe {
        let mut buf: [u8; super::DEFAULT_BUF_SIZE] = mem::uninitialized();
        reader.initializer().initialize(&mut buf);
        buf
    };

    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        written += len as u64;
    }
}

/// A reader which is always at EOF.
///
/// This struct is generally created by calling [`empty`][empty]. Please see
/// the documentation of `empty()` for more details.
///
/// [empty]: fn.empty.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Empty { _priv: () }

/// Constructs a new handle to an empty reader.
///
/// All reads from the returned reader will return `Ok(0)`.
///
/// # Examples
///
/// A slightly sad example of not reading anything into a buffer:
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = String::new();
/// io::empty().read_to_string(&mut buffer).unwrap();
/// assert!(buffer.is_empty());
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn empty() -> Empty { Empty { _priv: () } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Empty {
    #[inline]
    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> { Ok(0) }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}
#[stable(feature = "rust1", since = "1.0.0")]
impl BufRead for Empty {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> { Ok(&[]) }
    #[inline]
    fn consume(&mut self, _n: usize) {}
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Empty {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Empty { .. }")
    }
}

/// A reader which yields one byte over and over and over and over and over and...
///
/// This struct is generally created by calling [`repeat`][repeat]. Please
/// see the documentation of `repeat()` for more details.
///
/// [repeat]: fn.repeat.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Repeat { byte: u8 }

/// Creates an instance of a reader that infinitely repeats one byte.
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = [0; 3];
/// io::repeat(0b101).read_exact(&mut buffer).unwrap();
/// assert_eq!(buffer, [0b101, 0b101, 0b101]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn repeat(byte: u8) -> Repeat { Repeat { byte: byte } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Repeat {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        for slot in &mut *buf {
            *slot = self.byte;
        }
        Ok(buf.len())
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Repeat {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Repeat { .. }")
    }
}

/// A writer which will move data into the void.
///
/// This struct is generally created by calling [`sink`][sink]. Please
/// see the documentation of `sink()` for more details.
///
/// [sink]: fn.sink.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Sink { _priv: () }

/// Creates an instance of a writer which will successfully consume all data.
///
/// All calls to `write` on the returned instance will return `Ok(buf.len())`
/// and the contents of the buffer will not be inspected.
///
/// # Examples
///
/// ```rust
/// use std::io::{self, Write};
///
/// let buffer = vec![1, 2, 3, 5, 8];
/// let num_bytes = io::sink().write(&buffer).unwrap();
/// assert_eq!(num_bytes, 5);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn sink() -> Sink { Sink { _priv: () } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Write for Sink {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { Ok(buf.len()) }
    #[inline]
    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Sink {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Sink { .. }")
    }
}

#[cfg(test)]
mod tests {
    use io::prelude::*;
    use io::{copy, sink, empty, repeat};

    #[test]
    fn copy_copies() {
        let mut r = repeat(0).take(4);
        let mut w = sink();
        assert_eq!(copy(&mut r, &mut w).unwrap(), 4);

        let mut r = repeat(0).take(1 << 17);
        assert_eq!(copy(&mut r as &mut Read, &mut w as &mut Write).unwrap(), 1 << 17);
    }

    #[test]
    fn sink_sinks() {
        let mut s = sink();
        assert_eq!(s.write(&[]).unwrap(), 0);
        assert_eq!(s.write(&[0]).unwrap(), 1);
        assert_eq!(s.write(&[0; 1024]).unwrap(), 1024);
        assert_eq!(s.by_ref().write(&[0; 1024]).unwrap(), 1024);
    }

    #[test]
    fn empty_reads() {
        let mut e = empty();
        assert_eq!(e.read(&mut []).unwrap(), 0);
        assert_eq!(e.read(&mut [0]).unwrap(), 0);
        assert_eq!(e.read(&mut [0; 1024]).unwrap(), 0);
        assert_eq!(e.by_ref().read(&mut [0; 1024]).unwrap(), 0);
    }

    #[test]
    fn repeat_repeats() {
        let mut r = repeat(4);
        let mut b = [0; 1024];
        assert_eq!(r.read(&mut b).unwrap(), 1024);
        assert!(b.iter().all(|b| *b == 4));
    }

    #[test]
    fn take_some_bytes() {
        assert_eq!(repeat(4).take(100).bytes().count(), 100);
        assert_eq!(repeat(4).take(100).bytes().next().unwrap().unwrap(), 4);
        assert_eq!(repeat(1).take(10).chain(repeat(2).take(10)).bytes().count(), 20);
    }
}
<commit_msg>fix an incorrect assertion in the doc example for `std::io::copy`<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(missing_copy_implementations)]

use fmt;
use io::{self, Read, Initializer, Write, ErrorKind, BufRead};
use mem;

/// Copies the entire contents of a reader into a writer.
///
/// This function will continuously read data from `reader` and then
/// write it into `writer` in a streaming fashion until `reader`
/// returns EOF.
///
/// On success, the total number of bytes that were copied from
/// `reader` to `writer` is returned.
///
/// # Errors
///
/// This function will return an error immediately if any call to `read` or
/// `write` returns an error. All instances of `ErrorKind::Interrupted` are
/// handled by this function and the underlying operation is retried.
///
/// # Examples
///
/// ```
/// use std::io;
///
/// # fn foo() -> io::Result<()> {
/// let mut reader: &[u8] = b"hello";
/// let mut writer: Vec<u8> = vec![];
///
/// io::copy(&mut reader, &mut writer)?;
///
/// assert_eq!(&b"hello"[..], &writer[..]);
/// # Ok(())
/// # }
/// # foo().unwrap();
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn copy<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = unsafe {
        let mut buf: [u8; super::DEFAULT_BUF_SIZE] = mem::uninitialized();
        reader.initializer().initialize(&mut buf);
        buf
    };

    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        written += len as u64;
    }
}

/// A reader which is always at EOF.
///
/// This struct is generally created by calling [`empty`][empty]. Please see
/// the documentation of `empty()` for more details.
///
/// [empty]: fn.empty.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Empty { _priv: () }

/// Constructs a new handle to an empty reader.
///
/// All reads from the returned reader will return `Ok(0)`.
///
/// # Examples
///
/// A slightly sad example of not reading anything into a buffer:
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = String::new();
/// io::empty().read_to_string(&mut buffer).unwrap();
/// assert!(buffer.is_empty());
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn empty() -> Empty { Empty { _priv: () } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Empty {
    #[inline]
    fn read(&mut self, _buf: &mut [u8]) -> io::Result<usize> { Ok(0) }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}
#[stable(feature = "rust1", since = "1.0.0")]
impl BufRead for Empty {
    #[inline]
    fn fill_buf(&mut self) -> io::Result<&[u8]> { Ok(&[]) }
    #[inline]
    fn consume(&mut self, _n: usize) {}
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Empty {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Empty { .. }")
    }
}

/// A reader which yields one byte over and over and over and over and over and...
///
/// This struct is generally created by calling [`repeat`][repeat]. Please
/// see the documentation of `repeat()` for more details.
///
/// [repeat]: fn.repeat.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Repeat { byte: u8 }

/// Creates an instance of a reader that infinitely repeats one byte.
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```
/// use std::io::{self, Read};
///
/// let mut buffer = [0; 3];
/// io::repeat(0b101).read_exact(&mut buffer).unwrap();
/// assert_eq!(buffer, [0b101, 0b101, 0b101]);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn repeat(byte: u8) -> Repeat { Repeat { byte: byte } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Read for Repeat {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        for slot in &mut *buf {
            *slot = self.byte;
        }
        Ok(buf.len())
    }

    #[inline]
    unsafe fn initializer(&self) -> Initializer {
        Initializer::nop()
    }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Repeat {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Repeat { .. }")
    }
}

/// A writer which will move data into the void.
///
/// This struct is generally created by calling [`sink`][sink]. Please
/// see the documentation of `sink()` for more details.
///
/// [sink]: fn.sink.html
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Sink { _priv: () }

/// Creates an instance of a writer which will successfully consume all data.
///
/// All calls to `write` on the returned instance will return `Ok(buf.len())`
/// and the contents of the buffer will not be inspected.
///
/// # Examples
///
/// ```rust
/// use std::io::{self, Write};
///
/// let buffer = vec![1, 2, 3, 5, 8];
/// let num_bytes = io::sink().write(&buffer).unwrap();
/// assert_eq!(num_bytes, 5);
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub fn sink() -> Sink { Sink { _priv: () } }

#[stable(feature = "rust1", since = "1.0.0")]
impl Write for Sink {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { Ok(buf.len()) }
    #[inline]
    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}

#[stable(feature = "std_debug", since = "1.16.0")]
impl fmt::Debug for Sink {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("Sink { .. }")
    }
}

#[cfg(test)]
mod tests {
    use io::prelude::*;
    use io::{copy, sink, empty, repeat};

    #[test]
    fn copy_copies() {
        let mut r = repeat(0).take(4);
        let mut w = sink();
        assert_eq!(copy(&mut r, &mut w).unwrap(), 4);

        let mut r = repeat(0).take(1 << 17);
        assert_eq!(copy(&mut r as &mut Read, &mut w as &mut Write).unwrap(), 1 << 17);
    }

    #[test]
    fn sink_sinks() {
        let mut s = sink();
        assert_eq!(s.write(&[]).unwrap(), 0);
        assert_eq!(s.write(&[0]).unwrap(), 1);
        assert_eq!(s.write(&[0; 1024]).unwrap(), 1024);
        assert_eq!(s.by_ref().write(&[0; 1024]).unwrap(), 1024);
    }

    #[test]
    fn empty_reads() {
        let mut e = empty();
        assert_eq!(e.read(&mut []).unwrap(), 0);
        assert_eq!(e.read(&mut [0]).unwrap(), 0);
        assert_eq!(e.read(&mut [0; 1024]).unwrap(), 0);
        assert_eq!(e.by_ref().read(&mut [0; 1024]).unwrap(), 0);
    }

    #[test]
    fn repeat_repeats() {
        let mut r = repeat(4);
        let mut b = [0; 1024];
        assert_eq!(r.read(&mut b).unwrap(), 1024);
        assert!(b.iter().all(|b| *b == 4));
    }

    #[test]
    fn take_some_bytes() {
        assert_eq!(repeat(4).take(100).bytes().count(), 100);
        assert_eq!(repeat(4).take(100).bytes().next().unwrap().unwrap(), 4);
        assert_eq!(repeat(1).take(10).chain(repeat(2).take(10)).bytes().count(), 20);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use fmt;
use io::{self, Error, ErrorKind};
use net::{ToSocketAddrs, SocketAddr};
use sys_common::net as net_imp;
use sys_common::{AsInner, FromInner, IntoInner};
use time::Duration;

/// A User Datagram Protocol socket.
///
/// This is an implementation of a bound UDP socket. This supports both IPv4 and
/// IPv6 addresses, and there is no corresponding notion of a server because UDP
/// is a datagram protocol.
///
/// # Examples
///
/// ```no_run
/// use std::net::UdpSocket;
///
/// # fn foo() -> std::io::Result<()> {
/// let mut socket = try!(UdpSocket::bind("127.0.0.1:34254"));
///
/// let mut buf = [0; 10];
/// let (amt, src) = try!(socket.recv_from(&mut buf));
///
/// // Send a reply to the socket we received data from
/// let buf = &mut buf[..amt];
/// buf.reverse();
/// try!(socket.send_to(buf, &src));
///
/// drop(socket); // close the socket
/// # Ok(())
/// # }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub struct UdpSocket(net_imp::UdpSocket);

impl UdpSocket {
    /// Creates a UDP socket from the given address.
    ///
    /// The address type can be any implementor of `ToSocketAddr` trait. See
    /// its documentation for concrete examples.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<UdpSocket> {
        super::each_addr(addr, net_imp::UdpSocket::bind).map(UdpSocket)
    }

    /// Receives data from the socket. On success, returns the number of bytes
    /// read and the address from whence the data came.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        self.0.recv_from(buf)
    }

    /// Sends data on the socket to the given address. On success, returns the
    /// number of bytes written.
    ///
    /// Address type can be any implementor of `ToSocketAddrs` trait. See its
    /// documentation for concrete examples.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn send_to<A: ToSocketAddrs>(&self, buf: &[u8], addr: A)
                                     -> io::Result<usize> {
        match try!(addr.to_socket_addrs()).next() {
            Some(addr) => self.0.send_to(buf, &addr),
            None => Err(Error::new(ErrorKind::InvalidInput,
                                   "no addresses to send data to")),
        }
    }

    /// Returns the socket address that this socket was created from.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.socket_addr()
    }

    /// Creates a new independently owned handle to the underlying socket.
    ///
    /// The returned `UdpSocket` is a reference to the same socket that this
    /// object references. Both handles will read and write the same port, and
    /// options set on one socket will be propagated to the other.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn try_clone(&self) -> io::Result<UdpSocket> {
        self.0.duplicate().map(UdpSocket)
    }

    /// Sets the read timeout to the timeout specified.
    ///
    /// If the value specified is `None`, then `read` calls will block
    /// indefinitely. It is an error to pass the zero `Duration` to this
    /// method.
    ///
    /// # Note
    ///
    /// Platforms may return a different error code whenever a read times out as
    /// a result of setting this option. For example Unix typically returns an
    /// error of the kind `WouldBlock`, but Windows may return `TimedOut`.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.0.set_read_timeout(dur)
    }

    /// Sets the write timeout to the timeout specified.
    ///
    /// If the value specified is `None`, then `write` calls will block
    /// indefinitely. It is an error to pass the zero `Duration` to this
    /// method.
    ///
    /// # Note
    ///
    /// Platforms may return a different error code whenever a write times out
    /// as a result of setting this option. For example Unix typically returns
    /// an error of the kind `WouldBlock`, but Windows may return `TimedOut`.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.0.set_write_timeout(dur)
    }

    /// Returns the read timeout of this socket.
    ///
    /// If the timeout is `None`, then `read` calls will block indefinitely.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn read_timeout(&self) -> io::Result<Option<Duration>> {
        self.0.read_timeout()
    }

    /// Returns the write timeout of this socket.
    ///
    /// If the timeout is `None`, then `write` calls will block indefinitely.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn write_timeout(&self) -> io::Result<Option<Duration>> {
        self.0.write_timeout()
    }
}

impl AsInner<net_imp::UdpSocket> for UdpSocket {
    fn as_inner(&self) -> &net_imp::UdpSocket { &self.0 }
}

impl FromInner<net_imp::UdpSocket> for UdpSocket {
    fn from_inner(inner: net_imp::UdpSocket) -> UdpSocket { UdpSocket(inner) }
}

impl IntoInner<net_imp::UdpSocket> for UdpSocket {
    fn into_inner(self) -> net_imp::UdpSocket { self.0 }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl fmt::Debug for UdpSocket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

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

    use io::ErrorKind;
    use net::*;
    use net::test::{next_test_ip4, next_test_ip6};
    use sync::mpsc::channel;
    use sys_common::AsInner;
    use time::Duration;
    use thread;

    fn each_ip(f: &mut FnMut(SocketAddr, SocketAddr)) {
        f(next_test_ip4(), next_test_ip4());
        f(next_test_ip6(), next_test_ip6());
    }

    macro_rules! t {
        ($e:expr) => {
            match $e {
                Ok(t) => t,
                Err(e) => panic!("received error for `{}`: {}", stringify!($e), e),
            }
        }
    }

    #[test]
    fn bind_error() {
        match UdpSocket::bind("1.1.1.1:9999") {
            Ok(..) => panic!(),
            Err(e) => {
                assert_eq!(e.kind(), ErrorKind::AddrNotAvailable)
            }
        }
    }

    #[test]
    fn socket_smoke_test_ip4() {
        each_ip(&mut |server_ip, client_ip| {
            let (tx1, rx1) = channel();
            let (tx2, rx2) = channel();

            let _t = thread::spawn(move|| {
                let client = t!(UdpSocket::bind(&client_ip));
                rx1.recv().unwrap();
                t!(client.send_to(&[99], &server_ip));
                tx2.send(()).unwrap();
            });

            let server = t!(UdpSocket::bind(&server_ip));
            tx1.send(()).unwrap();
            let mut buf = [0];
            let (nread, src) = t!(server.recv_from(&mut buf));
            assert_eq!(nread, 1);
            assert_eq!(buf[0], 99);
            assert_eq!(src, client_ip);
            rx2.recv().unwrap();
        })
    }

    #[test]
    fn socket_name_ip4() {
        each_ip(&mut |addr, _| {
            let server = t!(UdpSocket::bind(&addr));
            assert_eq!(addr, t!(server.local_addr()));
        })
    }

    #[test]
    fn udp_clone_smoke() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));

            let _t = thread::spawn(move|| {
                let mut buf = [0, 0];
                assert_eq!(sock2.recv_from(&mut buf).unwrap(), (1, addr1));
                assert_eq!(buf[0], 1);
                t!(sock2.send_to(&[2], &addr1));
            });

            let sock3 = t!(sock1.try_clone());

            let (tx1, rx1) = channel();
            let (tx2, rx2) = channel();
            let _t = thread::spawn(move|| {
                rx1.recv().unwrap();
                t!(sock3.send_to(&[1], &addr2));
                tx2.send(()).unwrap();
            });
            tx1.send(()).unwrap();
            let mut buf = [0, 0];
            assert_eq!(sock1.recv_from(&mut buf).unwrap(), (1, addr2));
            rx2.recv().unwrap();
        })
    }

    #[test]
    fn udp_clone_two_read() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));
            let (tx1, rx) = channel();
            let tx2 = tx1.clone();

            let _t = thread::spawn(move|| {
                t!(sock2.send_to(&[1], &addr1));
                rx.recv().unwrap();
                t!(sock2.send_to(&[2], &addr1));
                rx.recv().unwrap();
            });

            let sock3 = t!(sock1.try_clone());

            let (done, rx) = channel();
            let _t = thread::spawn(move|| {
                let mut buf = [0, 0];
                t!(sock3.recv_from(&mut buf));
                tx2.send(()).unwrap();
                done.send(()).unwrap();
            });
            let mut buf = [0, 0];
            t!(sock1.recv_from(&mut buf));
            tx1.send(()).unwrap();

            rx.recv().unwrap();
        })
    }

    #[test]
    fn udp_clone_two_write() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));

            let (tx, rx) = channel();
            let (serv_tx, serv_rx) = channel();

            let _t = thread::spawn(move|| {
                let mut buf = [0, 1];
                rx.recv().unwrap();
                t!(sock2.recv_from(&mut buf));
                serv_tx.send(()).unwrap();
            });

            let sock3 = t!(sock1.try_clone());

            let (done, rx) = channel();
            let tx2 = tx.clone();
            let _t = thread::spawn(move|| {
                match sock3.send_to(&[1], &addr2) {
                    Ok(..) => { let _ = tx2.send(()); }
                    Err(..) => {}
                }
                done.send(()).unwrap();
            });
            match sock1.send_to(&[2], &addr2) {
                Ok(..) => { let _ = tx.send(()); }
                Err(..) => {}
            }
            drop(tx);

            rx.recv().unwrap();
            serv_rx.recv().unwrap();
        })
    }

    #[test]
    fn debug() {
        let name = if cfg!(windows) {"socket"} else {"fd"};
        let socket_addr = next_test_ip4();

        let udpsock = t!(UdpSocket::bind(&socket_addr));
        let udpsock_inner = udpsock.0.socket().as_inner();
        let compare = format!("UdpSocket {{ addr: {:?}, {}: {:?} }}",
                              socket_addr, name, udpsock_inner);
        assert_eq!(format!("{:?}", udpsock), compare);
    }

    // FIXME: re-enabled bitrig/openbsd/netbsd tests once their socket timeout code
    //        no longer has rounding errors.
    #[cfg_attr(any(target_os = "bitrig", target_os = "netbsd", target_os = "openbsd"), ignore)]
    #[test]
    fn timeouts() {
        let addr = next_test_ip4();

        let stream = t!(UdpSocket::bind(&addr));
        let dur = Duration::new(15410, 0);

        assert_eq!(None, t!(stream.read_timeout()));

        t!(stream.set_read_timeout(Some(dur)));
        assert_eq!(Some(dur), t!(stream.read_timeout()));

        assert_eq!(None, t!(stream.write_timeout()));

        t!(stream.set_write_timeout(Some(dur)));
        assert_eq!(Some(dur), t!(stream.write_timeout()));

        t!(stream.set_read_timeout(None));
        assert_eq!(None, t!(stream.read_timeout()));

        t!(stream.set_write_timeout(None));
        assert_eq!(None, t!(stream.write_timeout()));
    }

    #[test]
    fn test_read_timeout() {
        let addr = next_test_ip4();

        let mut stream = t!(UdpSocket::bind(&addr));
        t!(stream.set_read_timeout(Some(Duration::from_millis(1000))));

        let mut buf = [0; 10];
        let wait = Duration::span(|| {
            let kind = stream.recv_from(&mut buf).err().expect("expected error").kind();
            assert!(kind == ErrorKind::WouldBlock || kind == ErrorKind::TimedOut);
        });
        assert!(wait > Duration::from_millis(400));
    }

    #[test]
    fn test_read_with_timeout() {
        let addr = next_test_ip4();

        let mut stream = t!(UdpSocket::bind(&addr));
        t!(stream.set_read_timeout(Some(Duration::from_millis(1000))));

        t!(stream.send_to(b"hello world", &addr));

        let mut buf = [0; 11];
        t!(stream.recv_from(&mut buf));
        assert_eq!(b"hello world", &buf[..]);

        let wait = Duration::span(|| {
            let kind = stream.recv_from(&mut buf).err().expect("expected error").kind();
            assert!(kind == ErrorKind::WouldBlock || kind == ErrorKind::TimedOut);
        });
        assert!(wait > Duration::from_millis(400));
    }
}
<commit_msg>Removed usage of `drop` from `std::net::UdpSocket` docs<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use fmt;
use io::{self, Error, ErrorKind};
use net::{ToSocketAddrs, SocketAddr};
use sys_common::net as net_imp;
use sys_common::{AsInner, FromInner, IntoInner};
use time::Duration;

/// A User Datagram Protocol socket.
///
/// This is an implementation of a bound UDP socket. This supports both IPv4 and
/// IPv6 addresses, and there is no corresponding notion of a server because UDP
/// is a datagram protocol.
///
/// # Examples
///
/// ```no_run
/// use std::net::UdpSocket;
///
/// # fn foo() -> std::io::Result<()> {
/// {
///     let mut socket = try!(UdpSocket::bind("127.0.0.1:34254"));
///
///     // read from the socket
///     let mut buf = [0; 10];
///     let (amt, src) = try!(socket.recv_from(&mut buf));
///
///     // send a reply to the socket we received data from
///     let buf = &mut buf[..amt];
///     buf.reverse();
///     try!(socket.send_to(buf, &src));
///     # Ok(())
/// } // the socket is closed here
/// # }
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
pub struct UdpSocket(net_imp::UdpSocket);

impl UdpSocket {
    /// Creates a UDP socket from the given address.
    ///
    /// The address type can be any implementor of `ToSocketAddr` trait. See
    /// its documentation for concrete examples.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<UdpSocket> {
        super::each_addr(addr, net_imp::UdpSocket::bind).map(UdpSocket)
    }

    /// Receives data from the socket. On success, returns the number of bytes
    /// read and the address from whence the data came.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
        self.0.recv_from(buf)
    }

    /// Sends data on the socket to the given address. On success, returns the
    /// number of bytes written.
    ///
    /// Address type can be any implementor of `ToSocketAddrs` trait. See its
    /// documentation for concrete examples.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn send_to<A: ToSocketAddrs>(&self, buf: &[u8], addr: A)
                                     -> io::Result<usize> {
        match try!(addr.to_socket_addrs()).next() {
            Some(addr) => self.0.send_to(buf, &addr),
            None => Err(Error::new(ErrorKind::InvalidInput,
                                   "no addresses to send data to")),
        }
    }

    /// Returns the socket address that this socket was created from.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn local_addr(&self) -> io::Result<SocketAddr> {
        self.0.socket_addr()
    }

    /// Creates a new independently owned handle to the underlying socket.
    ///
    /// The returned `UdpSocket` is a reference to the same socket that this
    /// object references. Both handles will read and write the same port, and
    /// options set on one socket will be propagated to the other.
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn try_clone(&self) -> io::Result<UdpSocket> {
        self.0.duplicate().map(UdpSocket)
    }

    /// Sets the read timeout to the timeout specified.
    ///
    /// If the value specified is `None`, then `read` calls will block
    /// indefinitely. It is an error to pass the zero `Duration` to this
    /// method.
    ///
    /// # Note
    ///
    /// Platforms may return a different error code whenever a read times out as
    /// a result of setting this option. For example Unix typically returns an
    /// error of the kind `WouldBlock`, but Windows may return `TimedOut`.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.0.set_read_timeout(dur)
    }

    /// Sets the write timeout to the timeout specified.
    ///
    /// If the value specified is `None`, then `write` calls will block
    /// indefinitely. It is an error to pass the zero `Duration` to this
    /// method.
    ///
    /// # Note
    ///
    /// Platforms may return a different error code whenever a write times out
    /// as a result of setting this option. For example Unix typically returns
    /// an error of the kind `WouldBlock`, but Windows may return `TimedOut`.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.0.set_write_timeout(dur)
    }

    /// Returns the read timeout of this socket.
    ///
    /// If the timeout is `None`, then `read` calls will block indefinitely.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn read_timeout(&self) -> io::Result<Option<Duration>> {
        self.0.read_timeout()
    }

    /// Returns the write timeout of this socket.
    ///
    /// If the timeout is `None`, then `write` calls will block indefinitely.
    #[stable(feature = "socket_timeout", since = "1.4.0")]
    pub fn write_timeout(&self) -> io::Result<Option<Duration>> {
        self.0.write_timeout()
    }
}

impl AsInner<net_imp::UdpSocket> for UdpSocket {
    fn as_inner(&self) -> &net_imp::UdpSocket { &self.0 }
}

impl FromInner<net_imp::UdpSocket> for UdpSocket {
    fn from_inner(inner: net_imp::UdpSocket) -> UdpSocket { UdpSocket(inner) }
}

impl IntoInner<net_imp::UdpSocket> for UdpSocket {
    fn into_inner(self) -> net_imp::UdpSocket { self.0 }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl fmt::Debug for UdpSocket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

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

    use io::ErrorKind;
    use net::*;
    use net::test::{next_test_ip4, next_test_ip6};
    use sync::mpsc::channel;
    use sys_common::AsInner;
    use time::Duration;
    use thread;

    fn each_ip(f: &mut FnMut(SocketAddr, SocketAddr)) {
        f(next_test_ip4(), next_test_ip4());
        f(next_test_ip6(), next_test_ip6());
    }

    macro_rules! t {
        ($e:expr) => {
            match $e {
                Ok(t) => t,
                Err(e) => panic!("received error for `{}`: {}", stringify!($e), e),
            }
        }
    }

    #[test]
    fn bind_error() {
        match UdpSocket::bind("1.1.1.1:9999") {
            Ok(..) => panic!(),
            Err(e) => {
                assert_eq!(e.kind(), ErrorKind::AddrNotAvailable)
            }
        }
    }

    #[test]
    fn socket_smoke_test_ip4() {
        each_ip(&mut |server_ip, client_ip| {
            let (tx1, rx1) = channel();
            let (tx2, rx2) = channel();

            let _t = thread::spawn(move|| {
                let client = t!(UdpSocket::bind(&client_ip));
                rx1.recv().unwrap();
                t!(client.send_to(&[99], &server_ip));
                tx2.send(()).unwrap();
            });

            let server = t!(UdpSocket::bind(&server_ip));
            tx1.send(()).unwrap();
            let mut buf = [0];
            let (nread, src) = t!(server.recv_from(&mut buf));
            assert_eq!(nread, 1);
            assert_eq!(buf[0], 99);
            assert_eq!(src, client_ip);
            rx2.recv().unwrap();
        })
    }

    #[test]
    fn socket_name_ip4() {
        each_ip(&mut |addr, _| {
            let server = t!(UdpSocket::bind(&addr));
            assert_eq!(addr, t!(server.local_addr()));
        })
    }

    #[test]
    fn udp_clone_smoke() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));

            let _t = thread::spawn(move|| {
                let mut buf = [0, 0];
                assert_eq!(sock2.recv_from(&mut buf).unwrap(), (1, addr1));
                assert_eq!(buf[0], 1);
                t!(sock2.send_to(&[2], &addr1));
            });

            let sock3 = t!(sock1.try_clone());

            let (tx1, rx1) = channel();
            let (tx2, rx2) = channel();
            let _t = thread::spawn(move|| {
                rx1.recv().unwrap();
                t!(sock3.send_to(&[1], &addr2));
                tx2.send(()).unwrap();
            });
            tx1.send(()).unwrap();
            let mut buf = [0, 0];
            assert_eq!(sock1.recv_from(&mut buf).unwrap(), (1, addr2));
            rx2.recv().unwrap();
        })
    }

    #[test]
    fn udp_clone_two_read() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));
            let (tx1, rx) = channel();
            let tx2 = tx1.clone();

            let _t = thread::spawn(move|| {
                t!(sock2.send_to(&[1], &addr1));
                rx.recv().unwrap();
                t!(sock2.send_to(&[2], &addr1));
                rx.recv().unwrap();
            });

            let sock3 = t!(sock1.try_clone());

            let (done, rx) = channel();
            let _t = thread::spawn(move|| {
                let mut buf = [0, 0];
                t!(sock3.recv_from(&mut buf));
                tx2.send(()).unwrap();
                done.send(()).unwrap();
            });
            let mut buf = [0, 0];
            t!(sock1.recv_from(&mut buf));
            tx1.send(()).unwrap();

            rx.recv().unwrap();
        })
    }

    #[test]
    fn udp_clone_two_write() {
        each_ip(&mut |addr1, addr2| {
            let sock1 = t!(UdpSocket::bind(&addr1));
            let sock2 = t!(UdpSocket::bind(&addr2));

            let (tx, rx) = channel();
            let (serv_tx, serv_rx) = channel();

            let _t = thread::spawn(move|| {
                let mut buf = [0, 1];
                rx.recv().unwrap();
                t!(sock2.recv_from(&mut buf));
                serv_tx.send(()).unwrap();
            });

            let sock3 = t!(sock1.try_clone());

            let (done, rx) = channel();
            let tx2 = tx.clone();
            let _t = thread::spawn(move|| {
                match sock3.send_to(&[1], &addr2) {
                    Ok(..) => { let _ = tx2.send(()); }
                    Err(..) => {}
                }
                done.send(()).unwrap();
            });
            match sock1.send_to(&[2], &addr2) {
                Ok(..) => { let _ = tx.send(()); }
                Err(..) => {}
            }
            drop(tx);

            rx.recv().unwrap();
            serv_rx.recv().unwrap();
        })
    }

    #[test]
    fn debug() {
        let name = if cfg!(windows) {"socket"} else {"fd"};
        let socket_addr = next_test_ip4();

        let udpsock = t!(UdpSocket::bind(&socket_addr));
        let udpsock_inner = udpsock.0.socket().as_inner();
        let compare = format!("UdpSocket {{ addr: {:?}, {}: {:?} }}",
                              socket_addr, name, udpsock_inner);
        assert_eq!(format!("{:?}", udpsock), compare);
    }

    // FIXME: re-enabled bitrig/openbsd/netbsd tests once their socket timeout code
    //        no longer has rounding errors.
    #[cfg_attr(any(target_os = "bitrig", target_os = "netbsd", target_os = "openbsd"), ignore)]
    #[test]
    fn timeouts() {
        let addr = next_test_ip4();

        let stream = t!(UdpSocket::bind(&addr));
        let dur = Duration::new(15410, 0);

        assert_eq!(None, t!(stream.read_timeout()));

        t!(stream.set_read_timeout(Some(dur)));
        assert_eq!(Some(dur), t!(stream.read_timeout()));

        assert_eq!(None, t!(stream.write_timeout()));

        t!(stream.set_write_timeout(Some(dur)));
        assert_eq!(Some(dur), t!(stream.write_timeout()));

        t!(stream.set_read_timeout(None));
        assert_eq!(None, t!(stream.read_timeout()));

        t!(stream.set_write_timeout(None));
        assert_eq!(None, t!(stream.write_timeout()));
    }

    #[test]
    fn test_read_timeout() {
        let addr = next_test_ip4();

        let mut stream = t!(UdpSocket::bind(&addr));
        t!(stream.set_read_timeout(Some(Duration::from_millis(1000))));

        let mut buf = [0; 10];
        let wait = Duration::span(|| {
            let kind = stream.recv_from(&mut buf).err().expect("expected error").kind();
            assert!(kind == ErrorKind::WouldBlock || kind == ErrorKind::TimedOut);
        });
        assert!(wait > Duration::from_millis(400));
    }

    #[test]
    fn test_read_with_timeout() {
        let addr = next_test_ip4();

        let mut stream = t!(UdpSocket::bind(&addr));
        t!(stream.set_read_timeout(Some(Duration::from_millis(1000))));

        t!(stream.send_to(b"hello world", &addr));

        let mut buf = [0; 11];
        t!(stream.recv_from(&mut buf));
        assert_eq!(b"hello world", &buf[..]);

        let wait = Duration::span(|| {
            let kind = stream.recv_from(&mut buf).err().expect("expected error").kind();
            assert!(kind == ErrorKind::WouldBlock || kind == ErrorKind::TimedOut);
        });
        assert!(wait > Duration::from_millis(400));
    }
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// xfail-test
// This checks that preemption works.

fn starve_main(alive: chan<int>) {
    debug!("signalling main");
    alive <| 1;
    debug!("starving main");
    let i: int = 0;
    loop { i += 1; }
}

pub fn main() {
    let alive: port<int> = port();
    debug!("main started");
    let s: task = spawn starve_main(chan(alive));
    let i: int;
    debug!("main waiting for alive signal");
    alive |> i;
    debug!("main got alive signal");
    while i < 50 { debug!("main iterated"); i += 1; }
    debug!("main completed");
}
<commit_msg>This test case now parses<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// xfail-test
// This checks that preemption works.

fn starve_main(alive: chan<int>) {
    debug!("signalling main");
    alive.recv(1);
    debug!("starving main");
    let i: int = 0;
    loop { i += 1; }
}

pub fn main() {
    let alive: port<int> = port();
    debug!("main started");
    let s: task = do task::spawn {
        starve_main(chan(alive));
    };
    let i: int;
    debug!("main waiting for alive signal");
    alive.send(i);
    debug!("main got alive signal");
    while i < 50 { debug!("main iterated"); i += 1; }
    debug!("main completed");
}
<|endoftext|>


<commit_before>//! A map data type which allows the same key to exist at multiple scope levels
use std::borrow::Borrow;
use std::collections::hash_map::{self, IterMut};
use std::fmt;
use std::hash::Hash;

use crate::fnv::FnvMap;

/// A map struct which allows for the introduction of different scopes
/// Introducing a new scope will make it possible to introduce additional
/// variables with names already defined, shadowing the old name
/// After exiting a scope the shadowed variable will again be re introduced
pub struct ScopedMap<K: Eq + Hash, V> {
    /// A hashmap storing a key -> value mapping
    /// Stores a vector of values in which the value at the top is value returned from 'get'
    map: FnvMap<K, Vec<V>>,
    /// A vector of scopes, when entering a scope, None is added as a marker
    /// when later exiting a scope, values are removed from the map until the marker is found
    scopes: Vec<Option<K>>,
}

impl<K: Eq + Hash, V> Default for ScopedMap<K, V> {
    fn default() -> Self {
        ScopedMap {
            map: FnvMap::default(),
            scopes: Vec::default(),
        }
    }
}

impl<K, V> fmt::Debug for ScopedMap<K, V>
where
    K: Eq + Hash + fmt::Debug + Clone,
    V: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

#[allow(dead_code)]
impl<K: Eq + Hash + Clone, V> ScopedMap<K, V> {
    pub fn new() -> ScopedMap<K, V> {
        ScopedMap::default()
    }

    pub fn num_scopes(&self) -> usize {
        self.scopes.iter().filter(|s| s.is_none()).count() + 1
    }

    /// Introduces a new scope
    pub fn enter_scope(&mut self) {
        self.scopes.push(None);
    }

    pub fn current_scope(&self) -> impl Iterator<Item = (&K, &V)> {
        self.scopes
            .iter()
            .rev()
            .take_while(|scope| scope.is_some())
            .map(move |scope| {
                let key = scope.as_ref().unwrap();
                let value = self.map.get(key).and_then(|x| x.last()).expect("Value");
                (key, value)
            })
    }

    /// Exits the current scope, returning an iterator over the (key, value) pairs that are removed
    /// When `ExitScopeIter` is dropped any remaining pairs of the scope is removed as well.
    pub fn exit_scope(&mut self) -> ExitScopeIter<K, V> {
        ExitScopeIter {
            map: self,
            done: false,
        }
    }

    /// Removes a previously inserted value from the map.
    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Eq + Hash,
        K: ::std::fmt::Debug,
        V: ::std::fmt::Debug,
    {
        let x = self.map.get_mut(k).map(|x| x.pop())?;
        let mut i = self.scopes.len() as isize - 1;
        while i >= 0 {
            if self.scopes[i as usize]
                .as_ref()
                .expect("Tried to remove entry not in the current scope")
                .borrow()
                == k
            {
                self.scopes.remove(i as usize);
                break;
            }
            i -= 1;
        }
        x
    }

    /// Returns true if the key has a value declared in the last declared scope
    pub fn in_current_scope<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        for n in self.scopes.iter().rev() {
            match *n {
                Some(ref name) if name.borrow() == k => return true,
                None => break,
                _ => (),
            }
        }
        false
    }

    /// Returns a reference to the last inserted value corresponding to the key
    pub fn get<'a, Q: ?Sized>(&'a self, k: &Q) -> Option<&'a V>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.map.get(k).and_then(|x| x.last())
    }

    /// Returns a reference to the all inserted value corresponding to the key
    pub fn get_all<'a, Q: ?Sized>(&'a self, k: &Q) -> Option<&'a [V]>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.map.get(k).map(|x| &x[..])
    }

    pub fn entry(&mut self, key: K) -> Entry<K, V> {
        match self.map.entry(key) {
            hash_map::Entry::Occupied(entry) => {
                if entry.get().is_empty() {
                    Entry::Vacant(VacantEntry(InnerVacantEntry::Occupied(OccupiedEntry(
                        entry,
                        &mut self.scopes,
                    ))))
                } else {
                    Entry::Occupied(OccupiedEntry(entry, &mut self.scopes))
                }
            }
            hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry(InnerVacantEntry::Vacant(
                entry,
                &mut self.scopes,
            ))),
        }
    }

    pub fn contains_key<'a, Q: ?Sized>(&'a self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.get(k).is_some()
    }

    /// Returns the number of elements in the container.
    /// Shadowed elements are not counted
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Returns the number of elements in the container.
    /// Shadowed elements are counted
    pub fn full_len(&self) -> usize {
        self.map.values().map(|v| v.len()).sum()
    }

    /// Returns true if this map is empty
    pub fn is_empty(&self) -> bool {
        self.map.len() == 0
    }

    /// Removes all elements
    pub fn clear(&mut self) {
        self.map.clear();
        self.scopes.clear();
    }

    /// Swaps the value stored at key, or inserts it if it is not present
    pub fn swap(&mut self, k: K, v: V) -> Option<V> {
        let vec = self.map.entry(k.clone()).or_default();
        if vec.is_empty() {
            vec.push(v);
            self.scopes.push(Some(k));
            None
        } else {
            let r = vec.pop();
            vec.push(v);
            r
        }
    }

    pub fn pop(&mut self, k: &K) -> Option<V> {
        match self.map.get_mut(k).and_then(|x| x.pop()) {
            Some(v) => {
                let mut i = self.scopes.len() as isize - 1;
                while i >= 0 {
                    if self.scopes[i as usize].as_ref().map_or(false, |x| x == k) {
                        self.scopes.remove(i as usize);
                    }
                    i -= 1;
                }
                Some(v)
            }
            None => None,
        }
    }

    pub fn get_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> {
        self.map.get_mut(key).and_then(|x| x.last_mut())
    }

    pub fn insert(&mut self, k: K, v: V) -> bool {
        match self.entry(k) {
            Entry::Occupied(mut entry) => {
                entry.insert(v);
                true
            }
            Entry::Vacant(entry) => {
                entry.insert(v);
                false
            }
        }
    }

    pub fn into_iter(self) -> impl Iterator<Item = (K, V)> {
        self.map
            .into_iter()
            .filter_map(|(k, mut v)| Some((k, v.pop()?)))
    }
}

impl<K: Eq + Hash, V> ScopedMap<K, V> {
    /// Returns an iterator of the (key, values) pairs inserted in the map
    pub fn iter_mut(&mut self) -> IterMut<K, Vec<V>> {
        self.map.iter_mut()
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter {
            iter: self.map.iter(),
        }
    }
}

impl<K: Eq + Hash + Clone, V> Extend<(K, V)> for ScopedMap<K, V> {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = (K, V)>,
    {
        for (k, v) in iter {
            self.insert(k, v);
        }
    }
}

pub enum Entry<'a, K, V> {
    Vacant(VacantEntry<'a, K, V>),
    Occupied(OccupiedEntry<'a, K, V>),
}

impl<'a, K, V> Entry<'a, K, V> {
    pub fn or_insert(self, default: V) -> &'a mut V
    where
        K: Clone,
    {
        self.or_insert_with(|| default)
    }

    pub fn or_insert_with(self, default: impl FnOnce() -> V) -> &'a mut V
    where
        K: Clone,
    {
        match self {
            Entry::Vacant(entry) => entry.insert(default()),
            Entry::Occupied(entry) => entry.into_mut(),
        }
    }
}

pub enum InnerVacantEntry<'a, K, V> {
    Vacant(hash_map::VacantEntry<'a, K, Vec<V>>, &'a mut Vec<Option<K>>),
    Occupied(OccupiedEntry<'a, K, V>),
}
pub struct VacantEntry<'a, K, V>(InnerVacantEntry<'a, K, V>);
pub struct OccupiedEntry<'a, K, V>(
    hash_map::OccupiedEntry<'a, K, Vec<V>>,
    &'a mut Vec<Option<K>>,
);

impl<'a, K, V> VacantEntry<'a, K, V> {
    pub fn insert(self, value: V) -> &'a mut V
    where
        K: Clone,
    {
        match self.0 {
            InnerVacantEntry::Vacant(entry, scopes) => {
                scopes.push(Some(entry.key().clone()));
                &mut entry.insert(vec![value])[0]
            }
            InnerVacantEntry::Occupied(mut entry) => {
                entry.insert(value);
                entry.into_mut()
            }
        }
    }
}

impl<'a, K, V> OccupiedEntry<'a, K, V> {
    pub fn key(&self) -> &K {
        self.0.key()
    }

    pub fn get(&self) -> &V {
        self.0.get().last().unwrap()
    }

    pub fn get_mut(&mut self) -> &mut V {
        self.0.get_mut().last_mut().unwrap()
    }

    pub fn into_mut(self) -> &'a mut V {
        self.0.into_mut().last_mut().unwrap()
    }

    pub fn insert(&mut self, value: V)
    where
        K: Clone,
    {
        self.1.push(Some(self.key().clone()));
        self.0.get_mut().push(value);
    }
}

pub struct Iter<'a, K, V>
where
    K: 'a,
    V: 'a,
{
    iter: hash_map::Iter<'a, K, Vec<V>>,
}

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

    fn next(&mut self) -> Option<(&'a K, &'a V)> {
        loop {
            match self.iter.next() {
                Some((k, vs)) => {
                    if let Some(v) = vs.last() {
                        return Some((k, v));
                    }
                }
                None => return None,
            }
        }
    }
}

pub struct ExitScopeIter<'a, K, V>
where
    K: 'a + Eq + Hash + Clone,
    V: 'a,
{
    map: &'a mut ScopedMap<K, V>,
    done: bool,
}

impl<'a, K, V> Drop for ExitScopeIter<'a, K, V>
where
    K: Eq + Hash + Clone,
{
    fn drop(&mut self) {
        for _ in self {}
    }
}

impl<'a, K, V> Iterator for ExitScopeIter<'a, K, V>
where
    K: Eq + Hash + Clone,
{
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            None
        } else {
            match self.map.scopes.pop() {
                Some(Some(key)) => {
                    let result = self.map.map.get_mut(&key).and_then(|x| x.pop());
                    self.done = result.is_none();
                    result.map(|value| (key, value))
                }
                _ => {
                    self.done = true;
                    None
                }
            }
        }
    }
}

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

    #[test]
    fn test() {
        let mut map = ScopedMap::new();
        map.insert("a", 0);
        map.insert("b", 1);
        map.enter_scope();
        assert_eq!(map.get(&"a"), Some(&0));
        assert_eq!(map.get(&"b"), Some(&1));
        assert_eq!(map.get(&"c"), None);
        map.insert("a", 1);
        map.insert("c", 2);
        assert_eq!(map.get(&"a"), Some(&1));
        assert_eq!(map.get(&"c"), Some(&2));
        map.exit_scope();
        assert_eq!(map.get(&"a"), Some(&0));
        assert_eq!(map.get(&"c"), None);
    }

    #[test]
    fn remove() {
        let mut map = ScopedMap::new();
        map.insert("a", 0);
        map.enter_scope();
        map.insert("a", 1);
        map.insert("b", 2);
        assert_eq!(map.get("a"), Some(&1));
        map.remove("a");
        assert_eq!(map.get("a"), Some(&0));
        assert_eq!(map.get("b"), Some(&2));
        map.remove("b");
        assert_eq!(map.get("b"), None);
        map.exit_scope();
    }
}
<commit_msg>Correct ScopedMap::len<commit_after>//! A map data type which allows the same key to exist at multiple scope levels
use std::borrow::Borrow;
use std::collections::hash_map::{self, IterMut};
use std::fmt;
use std::hash::Hash;

use crate::fnv::FnvMap;

/// A map struct which allows for the introduction of different scopes
/// Introducing a new scope will make it possible to introduce additional
/// variables with names already defined, shadowing the old name
/// After exiting a scope the shadowed variable will again be re introduced
pub struct ScopedMap<K: Eq + Hash, V> {
    /// A hashmap storing a key -> value mapping
    /// Stores a vector of values in which the value at the top is value returned from 'get'
    map: FnvMap<K, Vec<V>>,
    /// A vector of scopes, when entering a scope, None is added as a marker
    /// when later exiting a scope, values are removed from the map until the marker is found
    scopes: Vec<Option<K>>,
}

impl<K: Eq + Hash, V> Default for ScopedMap<K, V> {
    fn default() -> Self {
        ScopedMap {
            map: FnvMap::default(),
            scopes: Vec::default(),
        }
    }
}

impl<K, V> fmt::Debug for ScopedMap<K, V>
where
    K: Eq + Hash + fmt::Debug + Clone,
    V: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

#[allow(dead_code)]
impl<K: Eq + Hash + Clone, V> ScopedMap<K, V> {
    pub fn new() -> ScopedMap<K, V> {
        ScopedMap::default()
    }

    pub fn num_scopes(&self) -> usize {
        self.scopes.iter().filter(|s| s.is_none()).count() + 1
    }

    /// Introduces a new scope
    pub fn enter_scope(&mut self) {
        self.scopes.push(None);
    }

    pub fn current_scope(&self) -> impl Iterator<Item = (&K, &V)> {
        self.scopes
            .iter()
            .rev()
            .take_while(|scope| scope.is_some())
            .map(move |scope| {
                let key = scope.as_ref().unwrap();
                let value = self.map.get(key).and_then(|x| x.last()).expect("Value");
                (key, value)
            })
    }

    /// Exits the current scope, returning an iterator over the (key, value) pairs that are removed
    /// When `ExitScopeIter` is dropped any remaining pairs of the scope is removed as well.
    pub fn exit_scope(&mut self) -> ExitScopeIter<K, V> {
        ExitScopeIter {
            map: self,
            done: false,
        }
    }

    /// Removes a previously inserted value from the map.
    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Eq + Hash,
        K: ::std::fmt::Debug,
        V: ::std::fmt::Debug,
    {
        let x = self.map.get_mut(k).map(|x| x.pop())?;
        let mut i = self.scopes.len() as isize - 1;
        while i >= 0 {
            if self.scopes[i as usize]
                .as_ref()
                .expect("Tried to remove entry not in the current scope")
                .borrow()
                == k
            {
                self.scopes.remove(i as usize);
                break;
            }
            i -= 1;
        }
        x
    }

    /// Returns true if the key has a value declared in the last declared scope
    pub fn in_current_scope<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        for n in self.scopes.iter().rev() {
            match *n {
                Some(ref name) if name.borrow() == k => return true,
                None => break,
                _ => (),
            }
        }
        false
    }

    /// Returns a reference to the last inserted value corresponding to the key
    pub fn get<'a, Q: ?Sized>(&'a self, k: &Q) -> Option<&'a V>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.map.get(k).and_then(|x| x.last())
    }

    /// Returns a reference to the all inserted value corresponding to the key
    pub fn get_all<'a, Q: ?Sized>(&'a self, k: &Q) -> Option<&'a [V]>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.map.get(k).map(|x| &x[..])
    }

    pub fn entry(&mut self, key: K) -> Entry<K, V> {
        match self.map.entry(key) {
            hash_map::Entry::Occupied(entry) => {
                if entry.get().is_empty() {
                    Entry::Vacant(VacantEntry(InnerVacantEntry::Occupied(OccupiedEntry(
                        entry,
                        &mut self.scopes,
                    ))))
                } else {
                    Entry::Occupied(OccupiedEntry(entry, &mut self.scopes))
                }
            }
            hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry(InnerVacantEntry::Vacant(
                entry,
                &mut self.scopes,
            ))),
        }
    }

    pub fn contains_key<'a, Q: ?Sized>(&'a self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        self.get(k).is_some()
    }

    /// Returns the number of elements in the container.
    /// Shadowed elements are counted
    pub fn len(&self) -> usize {
        self.map.values().map(|v| v.len()).sum()
    }

    /// Returns true if this map is empty
    pub fn is_empty(&self) -> bool {
        self.map.len() == 0
    }

    /// Removes all elements
    pub fn clear(&mut self) {
        self.map.clear();
        self.scopes.clear();
    }

    /// Swaps the value stored at key, or inserts it if it is not present
    pub fn swap(&mut self, k: K, v: V) -> Option<V> {
        let vec = self.map.entry(k.clone()).or_default();
        if vec.is_empty() {
            vec.push(v);
            self.scopes.push(Some(k));
            None
        } else {
            let r = vec.pop();
            vec.push(v);
            r
        }
    }

    pub fn pop(&mut self, k: &K) -> Option<V> {
        match self.map.get_mut(k).and_then(|x| x.pop()) {
            Some(v) => {
                let mut i = self.scopes.len() as isize - 1;
                while i >= 0 {
                    if self.scopes[i as usize].as_ref().map_or(false, |x| x == k) {
                        self.scopes.remove(i as usize);
                    }
                    i -= 1;
                }
                Some(v)
            }
            None => None,
        }
    }

    pub fn get_mut<'a>(&'a mut self, key: &K) -> Option<&'a mut V> {
        self.map.get_mut(key).and_then(|x| x.last_mut())
    }

    pub fn insert(&mut self, k: K, v: V) -> bool {
        match self.entry(k) {
            Entry::Occupied(mut entry) => {
                entry.insert(v);
                true
            }
            Entry::Vacant(entry) => {
                entry.insert(v);
                false
            }
        }
    }

    pub fn into_iter(self) -> impl Iterator<Item = (K, V)> {
        self.map
            .into_iter()
            .filter_map(|(k, mut v)| Some((k, v.pop()?)))
    }
}

impl<K: Eq + Hash, V> ScopedMap<K, V> {
    /// Returns an iterator of the (key, values) pairs inserted in the map
    pub fn iter_mut(&mut self) -> IterMut<K, Vec<V>> {
        self.map.iter_mut()
    }

    pub fn iter(&self) -> Iter<K, V> {
        Iter {
            iter: self.map.iter(),
        }
    }
}

impl<K: Eq + Hash + Clone, V> Extend<(K, V)> for ScopedMap<K, V> {
    fn extend<T>(&mut self, iter: T)
    where
        T: IntoIterator<Item = (K, V)>,
    {
        for (k, v) in iter {
            self.insert(k, v);
        }
    }
}

pub enum Entry<'a, K, V> {
    Vacant(VacantEntry<'a, K, V>),
    Occupied(OccupiedEntry<'a, K, V>),
}

impl<'a, K, V> Entry<'a, K, V> {
    pub fn or_insert(self, default: V) -> &'a mut V
    where
        K: Clone,
    {
        self.or_insert_with(|| default)
    }

    pub fn or_insert_with(self, default: impl FnOnce() -> V) -> &'a mut V
    where
        K: Clone,
    {
        match self {
            Entry::Vacant(entry) => entry.insert(default()),
            Entry::Occupied(entry) => entry.into_mut(),
        }
    }
}

pub enum InnerVacantEntry<'a, K, V> {
    Vacant(hash_map::VacantEntry<'a, K, Vec<V>>, &'a mut Vec<Option<K>>),
    Occupied(OccupiedEntry<'a, K, V>),
}
pub struct VacantEntry<'a, K, V>(InnerVacantEntry<'a, K, V>);
pub struct OccupiedEntry<'a, K, V>(
    hash_map::OccupiedEntry<'a, K, Vec<V>>,
    &'a mut Vec<Option<K>>,
);

impl<'a, K, V> VacantEntry<'a, K, V> {
    pub fn insert(self, value: V) -> &'a mut V
    where
        K: Clone,
    {
        match self.0 {
            InnerVacantEntry::Vacant(entry, scopes) => {
                scopes.push(Some(entry.key().clone()));
                &mut entry.insert(vec![value])[0]
            }
            InnerVacantEntry::Occupied(mut entry) => {
                entry.insert(value);
                entry.into_mut()
            }
        }
    }
}

impl<'a, K, V> OccupiedEntry<'a, K, V> {
    pub fn key(&self) -> &K {
        self.0.key()
    }

    pub fn get(&self) -> &V {
        self.0.get().last().unwrap()
    }

    pub fn get_mut(&mut self) -> &mut V {
        self.0.get_mut().last_mut().unwrap()
    }

    pub fn into_mut(self) -> &'a mut V {
        self.0.into_mut().last_mut().unwrap()
    }

    pub fn insert(&mut self, value: V)
    where
        K: Clone,
    {
        self.1.push(Some(self.key().clone()));
        self.0.get_mut().push(value);
    }
}

pub struct Iter<'a, K, V>
where
    K: 'a,
    V: 'a,
{
    iter: hash_map::Iter<'a, K, Vec<V>>,
}

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

    fn next(&mut self) -> Option<(&'a K, &'a V)> {
        loop {
            match self.iter.next() {
                Some((k, vs)) => {
                    if let Some(v) = vs.last() {
                        return Some((k, v));
                    }
                }
                None => return None,
            }
        }
    }
}

pub struct ExitScopeIter<'a, K, V>
where
    K: 'a + Eq + Hash + Clone,
    V: 'a,
{
    map: &'a mut ScopedMap<K, V>,
    done: bool,
}

impl<'a, K, V> Drop for ExitScopeIter<'a, K, V>
where
    K: Eq + Hash + Clone,
{
    fn drop(&mut self) {
        for _ in self {}
    }
}

impl<'a, K, V> Iterator for ExitScopeIter<'a, K, V>
where
    K: Eq + Hash + Clone,
{
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            None
        } else {
            match self.map.scopes.pop() {
                Some(Some(key)) => {
                    let result = self.map.map.get_mut(&key).and_then(|x| x.pop());
                    self.done = result.is_none();
                    result.map(|value| (key, value))
                }
                _ => {
                    self.done = true;
                    None
                }
            }
        }
    }
}

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

    #[test]
    fn test() {
        let mut map = ScopedMap::new();
        map.insert("a", 0);
        map.insert("b", 1);
        map.enter_scope();
        assert_eq!(map.get(&"a"), Some(&0));
        assert_eq!(map.get(&"b"), Some(&1));
        assert_eq!(map.get(&"c"), None);
        map.insert("a", 1);
        map.insert("c", 2);
        assert_eq!(map.get(&"a"), Some(&1));
        assert_eq!(map.get(&"c"), Some(&2));
        map.exit_scope();
        assert_eq!(map.get(&"a"), Some(&0));
        assert_eq!(map.get(&"c"), None);
    }

    #[test]
    fn remove() {
        let mut map = ScopedMap::new();
        map.insert("a", 0);
        map.enter_scope();
        map.insert("a", 1);
        map.insert("b", 2);
        assert_eq!(map.get("a"), Some(&1));
        map.remove("a");
        assert_eq!(map.get("a"), Some(&0));
        assert_eq!(map.get("b"), Some(&2));
        map.remove("b");
        assert_eq!(map.get("b"), None);
        map.exit_scope();
    }
}
<|endoftext|>


<commit_before>use bifrost::utils::async_locks::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use core::borrow::BorrowMut;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::Future;
use index::btree::*;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use owning_ref::{OwningHandle, OwningRef, RcRef};
use ram::cell::Cell;
use ram::schema::{Field, Schema};
use ram::types::*;
use std::cell::Ref;
use std::cell::RefCell;
use std::cell::RefMut;
use std::cell::UnsafeCell;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

const PAGE_SCHEMA: &'static str = "NEB_BTREE_PAGE";
const KEYS_FIELD: &'static str = "keys";
const NEXT_FIELD: &'static str = "next";
const PREV_FIELD: &'static str = "prev";

lazy_static! {
    static ref KEYS_KEY_HASH: u64 = key_hash(KEYS_FIELD);
    static ref NEXT_PAGE_KEY_HASH: u64 = key_hash(NEXT_FIELD);
    static ref PREV_PAGE_KEY_HASH: u64 = key_hash(PREV_FIELD);
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct ExtNode<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub id: Id,
    pub keys: KS,
    pub next: NodeCellRef,
    pub prev: NodeCellRef,
    pub len: usize,
    pub dirty: bool,
    pub right_bound: EntryKey,
    mark: PhantomData<PS>,
}

pub struct ExtNodeSplit<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub node_2: Box<ExtNode<KS, PS>>,
    pub keys_1_len: usize,
}

impl<KS, PS> ExtNode<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(id: Id, right_bound: EntryKey) -> Box<Self> {
        box ExtNode {
            id,
            keys: KS::init(),
            next: Node::<KS, PS>::none_ref(),
            prev: Node::<KS, PS>::none_ref(),
            len: 0,
            dirty: false,
            right_bound,
            mark: PhantomData,
        }
    }

    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        let prev_id = read_node(&self.prev, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().id
        });
        let next_id = read_node(&self.next, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().id
        });

        value[*NEXT_PAGE_KEY_HASH] = Value::Id(prev_id);
        value[*PREV_PAGE_KEY_HASH] = Value::Id(next_id);
        value[*KEYS_KEY_HASH] = self.keys.as_slice_immute()[..self.len]
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        self.keys.as_slice_immute()[..self.len]
            .binary_search(key)
            .unwrap_or_else(|i| i)
    }
    pub fn remove_at(&mut self, pos: usize) {
        let mut cached_len = &mut self.len;
        debug!(
            "Removing from external pos {}, len {}, key {:?}",
            pos,
            cached_len,
            self.keys.as_slice()[pos]
        );
        self.keys.remove_at(pos, cached_len);
    }

    pub fn split_insert(
        &mut self,
        key: EntryKey,
        pos: usize,
        self_ref: &NodeCellRef,
        self_next: &mut NodeWriteGuard<KS, PS>,
        tree: &BPlusTree<KS, PS>,
    ) -> (NodeCellRef, EntryKey) {
        // cached.dump();
        let pivot = self.len / 2;
        let new_page_id = tree.new_page_id();
        let mut keys_1 = &mut self.keys;
        let mut keys_2 = keys_1.split_at_pivot(pivot, self.len);
        let mut keys_1_len = pivot;
        let mut keys_2_len = self.len - pivot;
        // modify next node point previous to new node
        insert_into_split(
            key,
            keys_1,
            &mut keys_2,
            &mut keys_1_len,
            &mut keys_2_len,
            pos,
        );
        let pivot_key = keys_2.as_slice()[0].clone();
        let extnode_2: Box<ExtNode<KS, PS>> = box ExtNode {
            id: new_page_id,
            keys: keys_2,
            next: self.next.clone(),
            prev: self_ref.clone(),
            len: keys_2_len,
            dirty: true,
            right_bound: mem::replace(&mut self.right_bound, pivot_key.clone()),
            mark: PhantomData,
        };
        debug_assert!(pivot_key > smallvec!(0));
        debug_assert!(
            &pivot_key > &keys_1.as_slice()[keys_1_len - 1],
            "{:?} / {:?} @ {}",
            pivot_key,
            &keys_1.as_slice()[keys_1_len - 1],
            pos
        );
        debug_assert!(read_unchecked::<KS, PS>(&extnode_2.prev).is_ext());
        self.len = keys_1_len;
        debug!(
            "Split to left len {}, right len {}, right prev id: {:?}",
            self.len,
            extnode_2.len,
            read_unchecked::<KS, PS>(&extnode_2.prev).ext_id()
        );
        let node_2 = NodeCellRef::new(Node::with_external(extnode_2));
        {
            if !self_next.is_none() {
                self_next.extnode_mut().prev = node_2.clone();
            }
        }
        self.next = node_2.clone();

        (node_2, pivot_key)
    }

    pub fn insert(
        &mut self,
        key: &EntryKey,
        tree: &BPlusTree<KS, PS>,
        self_ref: &NodeCellRef,
        parent: &NodeCellRef,
    ) -> Option<NodeSplit<KS, PS>> {
        let key = key.clone();
        let pos = self.search(&key);
        debug_assert!(self.len <= KS::slice_len());
        debug_assert!(pos <= self.len);
        if self.len == KS::slice_len() {
            // need to split
            debug!("insert to external with split, key {:?}, pos {}", key, pos);
            let mut self_next: NodeWriteGuard<KS, PS> = write_node(&self.next);
            let mut parent_latch: NodeWriteGuard<KS, PS> = write_node(parent);
            let (node_2, pivot_key) = self.split_insert(key, pos, self_ref, &mut self_next, tree);
            Some(NodeSplit {
                new_right_node: node_2,
                pivot: pivot_key,
                parent_latch,
                left_node_latch: NodeWriteGuard::default(),
            })
        } else {
            debug!("insert to external without split at {}, key {:?}", pos, key);
            self.keys.insert_at(key, pos, &mut self.len);
            None
        }
    }

    pub fn merge_with(&mut self, right: &mut Self) {
        debug!(
            "Merge external node, left len {}:{:?}, right len {}:{:?}",
            self.len, self.keys, right.len, right.keys
        );
        let self_len = self.len;
        let new_len = self.len + right.len;
        debug_assert!(new_len <= KS::slice_len());
        for i in self.len..new_len {
            mem::swap(
                &mut self.keys.as_slice()[i],
                &mut right.keys.as_slice()[i - self_len],
            );
        }
        self.len = new_len;
    }

    pub fn remove_contains(&mut self, set: &mut BTreeSet<EntryKey>) {
        let remaining_keys = {
            let remaining = self.keys.as_slice()[..self.len]
                .iter()
                .filter(|&k| !set.remove(k))
                .collect_vec();
            if remaining.len() == self.len {
                return;
            }
            remaining.into_iter().cloned().collect_vec()
        };
        let mut self_key_slice = self.keys.as_slice();
        self.len = remaining_keys.len();
        for (i, k_ref) in remaining_keys.into_iter().enumerate() {
            self_key_slice[i] = k_ref;
        }
    }

    pub fn merge_sort(&mut self, right: &[&EntryKey]) {
        debug!("Merge sort have right nodes {:?}", right);
        let self_len_before_merge = self.len;
        debug_assert!(self_len_before_merge + right.len() <= KS::slice_len());
        let mut pos = 0;
        let mut left_pos = 0;
        let mut right_pos = 0;
        let mut left_keys = mem::replace(&mut self.keys, KS::init());
        let mut left = left_keys.as_slice();
        while left_pos < self.len && right_pos < right.len() {
            let left_key = &left[left_pos];
            let right_key = right[right_pos];
            if left_key < right_key {
                self.keys.as_slice()[pos] = left_key.clone();
                left_pos += 1;
            } else {
                self.keys.as_slice()[pos] = right_key.clone();
                right_pos += 1;
            }
            pos += 1;
        }
        for key in &left[left_pos..self.len] {
            self.keys.as_slice()[pos] = key.clone();
            pos += 1;
            left_pos += 1;
        }
        for key in &right[right_pos..] {
            self.keys.as_slice()[pos] = (*key).clone();
            pos += 1;
            right_pos += 1;
        }
        debug!(
            "Merge sorted have keys {:?}",
            &self.keys.as_slice_immute()[..pos]
        );
        self.len = pos;
        self.dirty = true;
        debug_assert_eq!(self_len_before_merge, left_pos);
        debug_assert_eq!(right.len(), right_pos);
        debug_assert_eq!(self.len, self_len_before_merge + right.len());
        debug!(
            "Merge sorted page have keys: {:?}",
            &self.keys.as_slice_immute()[..self.len]
        )
    }

    pub fn dump(&self) {
        debug!("Dumping {:?}, keys {}", self.id, self.len);
        for i in 0..KS::slice_len() {
            debug!("{}\t- {:?}", i, self.keys.as_slice_immute()[i]);
        }
    }
    pub fn is_dirty(&self) -> bool {
        unimplemented!()
    }
}

pub fn page_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new(
            "*",
            0,
            false,
            false,
            Some(vec![
                Field::new(NEXT_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(PREV_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(KEYS_FIELD, type_id_of(Type::SmallBytes), false, true, None),
            ]),
        ),
    }
}

pub struct NodeCache<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    nodes: RefCell<HashMap<Id, NodeCellRef>>,
    storage: Arc<AsyncClient>,
    marker: PhantomData<(KS, PS)>,
}

impl<KS, PS> NodeCache<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        NodeCache {
            nodes: RefCell::new(HashMap::new()),
            storage: neb_client.clone(),
            marker: PhantomData,
        }
    }

    pub fn get(&self, id: &Id) -> NodeCellRef {
        if id.is_unit_id() {
            return NodeCellRef::new(Node::<KS, PS>::with_none());
        }
        let mut nodes = self.nodes.borrow_mut();
        nodes
            .entry(*id)
            .or_insert_with(|| {
                let cell = self.storage.read_cell(*id).wait().unwrap().unwrap();
                NodeCellRef::new(Node::with_external(self.extnode_from_cell(cell)))
            })
            .clone()
    }

    fn extnode_from_cell(&self, cell: Cell) -> Box<ExtNode<KS, PS>> {
        let cell_id = cell.id();
        let cell_version = cell.header.version;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        let mut key_slice = KS::init();
        let mut key_count = 0;
        for (i, key_val) in keys_array.iter().enumerate() {
            key_slice.as_slice()[i] = EntryKey::from(key_val.as_slice());
            key_count += 1;
        }
        box ExtNode {
            id: cell_id,
            keys: key_slice,
            next: self.get(next),
            prev: self.get(prev),
            len: key_count,
            dirty: false,
            right_bound: max_entry_key(), // TODO: assign a real one on reconstructing
            mark: PhantomData,
        }
    }
}
<commit_msg>Revert "make node split on heap"<commit_after>use bifrost::utils::async_locks::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
use bifrost::utils::fut_exec::wait;
use client::AsyncClient;
use core::borrow::BorrowMut;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::custom_types::map::Map;
use dovahkiin::types::type_id_of;
use dovahkiin::types::value::ToValue;
use futures::Future;
use index::btree::*;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use owning_ref::{OwningHandle, OwningRef, RcRef};
use ram::cell::Cell;
use ram::schema::{Field, Schema};
use ram::types::*;
use std::cell::Ref;
use std::cell::RefCell;
use std::cell::RefMut;
use std::cell::UnsafeCell;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use utils::lru_cache::LRUCache;

const PAGE_SCHEMA: &'static str = "NEB_BTREE_PAGE";
const KEYS_FIELD: &'static str = "keys";
const NEXT_FIELD: &'static str = "next";
const PREV_FIELD: &'static str = "prev";

lazy_static! {
    static ref KEYS_KEY_HASH: u64 = key_hash(KEYS_FIELD);
    static ref NEXT_PAGE_KEY_HASH: u64 = key_hash(NEXT_FIELD);
    static ref PREV_PAGE_KEY_HASH: u64 = key_hash(PREV_FIELD);
    static ref PAGE_SCHEMA_ID: u32 = key_hash(PAGE_SCHEMA) as u32;
}

pub struct ExtNode<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub id: Id,
    pub keys: KS,
    pub next: NodeCellRef,
    pub prev: NodeCellRef,
    pub len: usize,
    pub dirty: bool,
    pub right_bound: EntryKey,
    mark: PhantomData<PS>,
}

pub struct ExtNodeSplit<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub node_2: ExtNode<KS, PS>,
    pub keys_1_len: usize,
}

impl<KS, PS> ExtNode<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(id: Id, right_bound: EntryKey) -> Box<Self> {
        box ExtNode {
            id,
            keys: KS::init(),
            next: Node::<KS, PS>::none_ref(),
            prev: Node::<KS, PS>::none_ref(),
            len: 0,
            dirty: false,
            right_bound,
            mark: PhantomData,
        }
    }

    pub fn to_cell(&self) -> Cell {
        let mut value = Value::Map(Map::new());
        let prev_id = read_node(&self.prev, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().id
        });
        let next_id = read_node(&self.next, |node: &NodeReadHandler<KS, PS>| {
            node.extnode().id
        });

        value[*NEXT_PAGE_KEY_HASH] = Value::Id(prev_id);
        value[*PREV_PAGE_KEY_HASH] = Value::Id(next_id);
        value[*KEYS_KEY_HASH] = self.keys.as_slice_immute()[..self.len]
            .iter()
            .map(|key| SmallBytes::from_vec(key.as_slice().to_vec()))
            .collect_vec()
            .value();
        Cell::new_with_id(*PAGE_SCHEMA_ID, &self.id, value)
    }
    pub fn search(&self, key: &EntryKey) -> usize {
        self.keys.as_slice_immute()[..self.len]
            .binary_search(key)
            .unwrap_or_else(|i| i)
    }
    pub fn remove_at(&mut self, pos: usize) {
        let mut cached_len = &mut self.len;
        debug!(
            "Removing from external pos {}, len {}, key {:?}",
            pos,
            cached_len,
            self.keys.as_slice()[pos]
        );
        self.keys.remove_at(pos, cached_len);
    }

    pub fn split_insert(
        &mut self,
        key: EntryKey,
        pos: usize,
        self_ref: &NodeCellRef,
        self_next: &mut NodeWriteGuard<KS, PS>,
        tree: &BPlusTree<KS, PS>,
    ) -> (NodeCellRef, EntryKey) {
        // cached.dump();
        let pivot = self.len / 2;
        let new_page_id = tree.new_page_id();
        let mut keys_1 = &mut self.keys;
        let mut keys_2 = keys_1.split_at_pivot(pivot, self.len);
        let mut keys_1_len = pivot;
        let mut keys_2_len = self.len - pivot;
        // modify next node point previous to new node
        insert_into_split(
            key,
            keys_1,
            &mut keys_2,
            &mut keys_1_len,
            &mut keys_2_len,
            pos,
        );
        let pivot_key = keys_2.as_slice()[0].clone();
        let extnode_2: Box<ExtNode<KS, PS>> = box ExtNode {
            id: new_page_id,
            keys: keys_2,
            next: self.next.clone(),
            prev: self_ref.clone(),
            len: keys_2_len,
            dirty: true,
            right_bound: mem::replace(&mut self.right_bound, pivot_key.clone()),
            mark: PhantomData,
        };
        debug_assert!(pivot_key > smallvec!(0));
        debug_assert!(
            &pivot_key > &keys_1.as_slice()[keys_1_len - 1],
            "{:?} / {:?} @ {}",
            pivot_key,
            &keys_1.as_slice()[keys_1_len - 1],
            pos
        );
        debug_assert!(read_unchecked::<KS, PS>(&extnode_2.prev).is_ext());
        self.len = keys_1_len;
        debug!(
            "Split to left len {}, right len {}, right prev id: {:?}",
            self.len,
            extnode_2.len,
            read_unchecked::<KS, PS>(&extnode_2.prev).ext_id()
        );
        let node_2 = NodeCellRef::new(Node::with_external(extnode_2));
        {
            if !self_next.is_none() {
                self_next.extnode_mut().prev = node_2.clone();
            }
        }
        self.next = node_2.clone();

        (node_2, pivot_key)
    }

    pub fn insert(
        &mut self,
        key: &EntryKey,
        tree: &BPlusTree<KS, PS>,
        self_ref: &NodeCellRef,
        parent: &NodeCellRef,
    ) -> Option<NodeSplit<KS, PS>> {
        let key = key.clone();
        let pos = self.search(&key);
        debug_assert!(self.len <= KS::slice_len());
        debug_assert!(pos <= self.len);
        if self.len == KS::slice_len() {
            // need to split
            debug!("insert to external with split, key {:?}, pos {}", key, pos);
            let mut self_next: NodeWriteGuard<KS, PS> = write_node(&self.next);
            let mut parent_latch: NodeWriteGuard<KS, PS> = write_node(parent);
            let (node_2, pivot_key) = self.split_insert(key, pos, self_ref, &mut self_next, tree);
            Some(NodeSplit {
                new_right_node: node_2,
                pivot: pivot_key,
                parent_latch,
                left_node_latch: NodeWriteGuard::default(),
            })
        } else {
            debug!("insert to external without split at {}, key {:?}", pos, key);
            self.keys.insert_at(key, pos, &mut self.len);
            None
        }
    }

    pub fn merge_with(&mut self, right: &mut Self) {
        debug!(
            "Merge external node, left len {}:{:?}, right len {}:{:?}",
            self.len, self.keys, right.len, right.keys
        );
        let self_len = self.len;
        let new_len = self.len + right.len;
        debug_assert!(new_len <= KS::slice_len());
        for i in self.len..new_len {
            mem::swap(
                &mut self.keys.as_slice()[i],
                &mut right.keys.as_slice()[i - self_len],
            );
        }
        self.len = new_len;
    }

    pub fn remove_contains(&mut self, set: &mut BTreeSet<EntryKey>) {
        let remaining_keys = {
            let remaining = self.keys.as_slice()[..self.len]
                .iter()
                .filter(|&k| !set.remove(k))
                .collect_vec();
            if remaining.len() == self.len {
                return;
            }
            remaining.into_iter().cloned().collect_vec()
        };
        let mut self_key_slice = self.keys.as_slice();
        self.len = remaining_keys.len();
        for (i, k_ref) in remaining_keys.into_iter().enumerate() {
            self_key_slice[i] = k_ref;
        }
    }

    pub fn merge_sort(&mut self, right: &[&EntryKey]) {
        debug!("Merge sort have right nodes {:?}", right);
        let self_len_before_merge = self.len;
        debug_assert!(self_len_before_merge + right.len() <= KS::slice_len());
        let mut pos = 0;
        let mut left_pos = 0;
        let mut right_pos = 0;
        let mut left_keys = mem::replace(&mut self.keys, KS::init());
        let mut left = left_keys.as_slice();
        while left_pos < self.len && right_pos < right.len() {
            let left_key = &left[left_pos];
            let right_key = right[right_pos];
            if left_key < right_key {
                self.keys.as_slice()[pos] = left_key.clone();
                left_pos += 1;
            } else {
                self.keys.as_slice()[pos] = right_key.clone();
                right_pos += 1;
            }
            pos += 1;
        }
        for key in &left[left_pos..self.len] {
            self.keys.as_slice()[pos] = key.clone();
            pos += 1;
            left_pos += 1;
        }
        for key in &right[right_pos..] {
            self.keys.as_slice()[pos] = (*key).clone();
            pos += 1;
            right_pos += 1;
        }
        debug!(
            "Merge sorted have keys {:?}",
            &self.keys.as_slice_immute()[..pos]
        );
        self.len = pos;
        self.dirty = true;
        debug_assert_eq!(self_len_before_merge, left_pos);
        debug_assert_eq!(right.len(), right_pos);
        debug_assert_eq!(self.len, self_len_before_merge + right.len());
        debug!(
            "Merge sorted page have keys: {:?}",
            &self.keys.as_slice_immute()[..self.len]
        )
    }

    pub fn dump(&self) {
        debug!("Dumping {:?}, keys {}", self.id, self.len);
        for i in 0..KS::slice_len() {
            debug!("{}\t- {:?}", i, self.keys.as_slice_immute()[i]);
        }
    }
    pub fn is_dirty(&self) -> bool {
        unimplemented!()
    }
}

pub fn page_schema() -> Schema {
    Schema {
        id: *PAGE_SCHEMA_ID,
        name: String::from(PAGE_SCHEMA),
        key_field: None,
        str_key_field: None,
        is_dynamic: false,
        fields: Field::new(
            "*",
            0,
            false,
            false,
            Some(vec![
                Field::new(NEXT_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(PREV_FIELD, type_id_of(Type::Id), false, false, None),
                Field::new(KEYS_FIELD, type_id_of(Type::SmallBytes), false, true, None),
            ]),
        ),
    }
}

pub struct NodeCache<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    nodes: RefCell<HashMap<Id, NodeCellRef>>,
    storage: Arc<AsyncClient>,
    marker: PhantomData<(KS, PS)>,
}

impl<KS, PS> NodeCache<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new(neb_client: &Arc<AsyncClient>) -> Self {
        NodeCache {
            nodes: RefCell::new(HashMap::new()),
            storage: neb_client.clone(),
            marker: PhantomData,
        }
    }

    pub fn get(&self, id: &Id) -> NodeCellRef {
        if id.is_unit_id() {
            return NodeCellRef::new(Node::<KS, PS>::with_none());
        }
        let mut nodes = self.nodes.borrow_mut();
        nodes
            .entry(*id)
            .or_insert_with(|| {
                let cell = self.storage.read_cell(*id).wait().unwrap().unwrap();
                NodeCellRef::new(Node::with_external(self.extnode_from_cell(cell)))
            })
            .clone()
    }

    fn extnode_from_cell(&self, cell: Cell) -> Box<ExtNode<KS, PS>> {
        let cell_id = cell.id();
        let cell_version = cell.header.version;
        let next = cell.data[*NEXT_PAGE_KEY_HASH].Id().unwrap();
        let prev = cell.data[*PREV_PAGE_KEY_HASH].Id().unwrap();
        let keys = &cell.data[*KEYS_KEY_HASH];
        let keys_len = keys.len().unwrap();
        let keys_array = if let Value::PrimArray(PrimitiveArray::SmallBytes(ref array)) = keys {
            array
        } else {
            panic!()
        };
        let mut key_slice = KS::init();
        let mut key_count = 0;
        for (i, key_val) in keys_array.iter().enumerate() {
            key_slice.as_slice()[i] = EntryKey::from(key_val.as_slice());
            key_count += 1;
        }
        box ExtNode {
            id: cell_id,
            keys: key_slice,
            next: self.get(next),
            prev: self.get(prev),
            len: key_count,
            dirty: false,
            right_bound: max_entry_key(), // TODO: assign a real one on reconstructing
            mark: PhantomData,
        }
    }
}
<|endoftext|>