<commit_before>use libc::c_void;

use parse_code;
use std::io::{self, Write};
use zstd_sys;

struct EncoderContext {
    s: *mut zstd_sys::ZSTD_CStream,
}

impl Default for EncoderContext {
    fn default() -> Self {
        EncoderContext { s: unsafe { zstd_sys::ZSTD_createCStream() } }
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        let code = unsafe { zstd_sys::ZSTD_freeCStream(self.s) };
        parse_code(code).unwrap();
    }
}

#[derive(PartialEq)]
enum EncoderState {
    Accepting,
    // Finished just makes sure writes are no longer accepted.
    Finished,
    StreamEnd,
    // The "stream ended and flushed" state corresponds to the object being
    // destroyed.
}

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call `finish()` before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: W,
    // output buffer
    buffer: Vec<u8>,
    // offset in the output buffer
    offset: usize,

    // compression context
    context: EncoderContext,
    state: EncoderState,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,
    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F: 'static + FnMut(io::Result<W>)>(encoder: Encoder<W>,
                                              on_finish: F)
                                              -> Self {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder
            .as_ref()
            .unwrap()
            .get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder
            .as_mut()
            .unwrap()
            .get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder
            .take()
            .unwrap()
            .finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder
            .as_mut()
            .unwrap()
            .write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder
            .as_mut()
            .unwrap()
            .flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21)
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, level: i32, dictionary: &[u8])
                           -> io::Result<Self> {
        let context = EncoderContext::default();

        // Initialize the stream with an existing dictionary
        parse_code(unsafe {
            zstd_sys::ZSTD_initCStream_usingDict(context.s,
                                           dictionary.as_ptr() as *const c_void,
                                           dictionary.len(),
                                           level)
        })?;

        Encoder::with_context(writer, context)
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| { result.unwrap(); })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>
        (self, f: F)
         -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    fn with_context(writer: W, context: EncoderContext) -> io::Result<Self> {
        // This is the output buffer size,
        // for compressed data we get from zstd.
        let buffer_size = unsafe { zstd_sys::ZSTD_CStreamOutSize() };

        Ok(Encoder {
               writer: writer,
               buffer: Vec::with_capacity(buffer_size),
               offset: 0,
               context: context,
               state: EncoderState::Accepting,
           })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.writer
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Finishes the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// To get back `self` in case an error happened, use `try_finish`.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after writing your data,
    /// consider using `auto_finish()` to get an `AutoFinishEncoder`.
    pub fn finish(mut self) -> io::Result<W> {
        self.do_finish()?;
        // Return the writer, because why not
        Ok(self.writer)
    }

    /// Attempts to finish the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer if the finish was successful, or the
    /// object plus an error if it wasn't.
    ///
    /// `write` on this object will panic after `try_finish` has been called,
    /// even if it fails.
    pub fn try_finish(mut self) -> Result<W, (Self, io::Error)> {
        match self.do_finish() {
            Ok(()) => Ok(self.writer),
            Err(e) => Err((self, e))
        }
    }

    fn do_finish(&mut self) -> io::Result<()> {
        if self.state != EncoderState::StreamEnd {
            self.state = EncoderState::Finished;
            // First, closes the stream.
            let mut buffer = zstd_sys::ZSTD_outBuffer {
                dst: self.buffer.as_mut_ptr() as *mut c_void,
                size: self.buffer.capacity(),
                pos: 0,
            };
            let remaining = parse_code(unsafe {
                zstd_sys::ZSTD_endStream(self.context.s,
                                   &mut buffer as *mut zstd_sys::ZSTD_outBuffer)
            })?;
            unsafe {
                self.buffer.set_len(buffer.pos);
            }
            if remaining != 0 {
                // Need to flush?
                panic!("Need to flush, but I'm lazy.");
            }
            self.offset = 0;
            self.state = EncoderState::StreamEnd;
        }

        self.write_from_offset()?;

        Ok(())
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        unsafe { zstd_sys::ZSTD_CStreamInSize() }
    }

    /// write_all, except keep track of partial writes for non-blocking IO.
    fn write_from_offset(&mut self) -> io::Result<()> {
        while self.offset < self.buffer.len() {
            match self.writer.write(&self.buffer[self.offset..]) {
                Ok(n) => self.offset += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.state != EncoderState::Accepting {
            panic!("write called after try_finish attempted");
        }

        // Write any data pending in `self.buffer`.
        self.write_from_offset()?;

        // If we get to here, `self.buffer` can safely be discarded.

        let mut in_buffer = zstd_sys::ZSTD_inBuffer {
            src: buf.as_ptr() as *const c_void,
            size: buf.len(),
            pos: 0,
        };

        let mut out_buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };

        unsafe {
            // Time to fill our input buffer
            let code =
                zstd_sys::ZSTD_compressStream(self.context.s,
                                              &mut out_buffer as
                                              *mut zstd_sys::ZSTD_outBuffer,
                                              &mut in_buffer as
                                              *mut zstd_sys::ZSTD_inBuffer);
            // Note: this may very well be empty,
            // if it doesn't exceed zstd's own buffer
            self.buffer.set_len(out_buffer.pos);

            // Do we care about the hint?
            let _ = parse_code(code)?;
        }

        // This is the first time he sees this buffer.
        // Remember his delicate touch.
        self.offset = self.writer.write(&self.buffer)?;

        Ok(in_buffer.pos)
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.state == EncoderState::Accepting {
            let mut buffer = zstd_sys::ZSTD_outBuffer {
                dst: self.buffer.as_mut_ptr() as *mut c_void,
                size: self.buffer.capacity(),
                pos: 0,
            };
            unsafe {
                let code =
                    zstd_sys::ZSTD_flushStream(self.context.s,
                                               &mut buffer as
                                               *mut zstd_sys::ZSTD_outBuffer);
                self.buffer.set_len(buffer.pos);
                let _ = parse_code(code)?;
            }
            self.offset = 0;
        }

        self.write_from_offset()?;
        Ok(())
    }
}
<commit_msg>encoder: tries to write internal buffer immediately.<commit_after>use libc::c_void;

use parse_code;
use std::io::{self, Write};
use zstd_sys;

struct EncoderContext {
    s: *mut zstd_sys::ZSTD_CStream,
}

impl Default for EncoderContext {
    fn default() -> Self {
        EncoderContext { s: unsafe { zstd_sys::ZSTD_createCStream() } }
    }
}

impl Drop for EncoderContext {
    fn drop(&mut self) {
        let code = unsafe { zstd_sys::ZSTD_freeCStream(self.s) };
        parse_code(code).unwrap();
    }
}

#[derive(PartialEq)]
enum EncoderState {
    Accepting,
    // Finished just makes sure writes are no longer accepted.
    Finished,
    StreamEnd,
    // The "stream ended and flushed" state corresponds to the object being
    // destroyed.
}

/// An encoder that compress and forward data to another writer.
///
/// This allows to compress a stream of data
/// (good for files or heavy network stream).
///
/// Don't forget to call `finish()` before dropping it!
///
/// Note: The zstd library has its own internal input buffer (~128kb).
pub struct Encoder<W: Write> {
    // output writer (compressed data)
    writer: W,
    // output buffer
    buffer: Vec<u8>,
    // offset in the output buffer
    offset: usize,

    // compression context
    context: EncoderContext,
    state: EncoderState,
}

/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
pub struct AutoFinishEncoder<W: Write> {
    // We wrap this in an option to take it during drop.
    encoder: Option<Encoder<W>>,
    // TODO: make this a FnOnce once it works in a Box
    on_finish: Option<Box<FnMut(io::Result<W>)>>,
}

impl<W: Write> AutoFinishEncoder<W> {
    fn new<F: 'static + FnMut(io::Result<W>)>(encoder: Encoder<W>,
                                              on_finish: F)
                                              -> Self {
        AutoFinishEncoder {
            encoder: Some(encoder),
            on_finish: Some(Box::new(on_finish)),
        }
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        self.encoder
            .as_ref()
            .unwrap()
            .get_ref()
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        self.encoder
            .as_mut()
            .unwrap()
            .get_mut()
    }
}

impl<W: Write> Drop for AutoFinishEncoder<W> {
    fn drop(&mut self) {
        let result = self.encoder
            .take()
            .unwrap()
            .finish();
        if let Some(mut on_finish) = self.on_finish.take() {
            on_finish(result);
        }
    }
}

impl<W: Write> Write for AutoFinishEncoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.encoder
            .as_mut()
            .unwrap()
            .write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.encoder
            .as_mut()
            .unwrap()
            .flush()
    }
}

impl<W: Write> Encoder<W> {
    /// Creates a new encoder.
    ///
    /// `level`: compression level (1-21)
    pub fn new(writer: W, level: i32) -> io::Result<Self> {
        Self::with_dictionary(writer, level, &[])
    }

    /// Creates a new encoder, using an existing dictionary.
    ///
    /// (Provides better compression ratio for small files,
    /// but requires the dictionary to be present during decompression.)
    pub fn with_dictionary(writer: W, level: i32, dictionary: &[u8])
                           -> io::Result<Self> {
        let context = EncoderContext::default();

        // Initialize the stream with an existing dictionary
        parse_code(unsafe {
            zstd_sys::ZSTD_initCStream_usingDict(context.s,
                                           dictionary.as_ptr() as *const c_void,
                                           dictionary.len(),
                                           level)
        })?;

        Encoder::with_context(writer, context)
    }

    /// Returns a wrapper around `self` that will finish the stream on drop.
    ///
    /// # Panic
    ///
    /// Panics on drop if an error happens when finishing the stream.
    pub fn auto_finish(self) -> AutoFinishEncoder<W> {
        self.on_finish(|result| { result.unwrap(); })
    }

    /// Returns an encoder that will finish the stream on drop.
    ///
    /// Calls the given callback with the result from `finish()`.
    pub fn on_finish<F: 'static + FnMut(io::Result<W>)>
        (self, f: F)
         -> AutoFinishEncoder<W> {
        AutoFinishEncoder::new(self, f)
    }

    fn with_context(writer: W, context: EncoderContext) -> io::Result<Self> {
        // This is the output buffer size,
        // for compressed data we get from zstd.
        let buffer_size = unsafe { zstd_sys::ZSTD_CStreamOutSize() };

        Ok(Encoder {
               writer: writer,
               buffer: Vec::with_capacity(buffer_size),
               offset: 0,
               context: context,
               state: EncoderState::Accepting,
           })
    }

    /// Acquires a reference to the underlying writer.
    pub fn get_ref(&self) -> &W {
        &self.writer
    }

    /// Acquires a mutable reference to the underlying writer.
    ///
    /// Note that mutation of the writer may result in surprising results if
    /// this encoder is continued to be used.
    pub fn get_mut(&mut self) -> &mut W {
        &mut self.writer
    }

    /// Finishes the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer in case you need it.
    ///
    /// To get back `self` in case an error happened, use `try_finish`.
    ///
    /// **Note**: If you don't want (or can't) call `finish()` manually after writing your data,
    /// consider using `auto_finish()` to get an `AutoFinishEncoder`.
    pub fn finish(mut self) -> io::Result<W> {
        self.do_finish()?;
        // Return the writer, because why not
        Ok(self.writer)
    }

    /// Attempts to finish the stream. You *need* to call this after writing your stuff.
    ///
    /// This returns the inner writer if the finish was successful, or the
    /// object plus an error if it wasn't.
    ///
    /// `write` on this object will panic after `try_finish` has been called,
    /// even if it fails.
    pub fn try_finish(mut self) -> Result<W, (Self, io::Error)> {
        match self.do_finish() {
            Ok(()) => Ok(self.writer),
            Err(e) => Err((self, e))
        }
    }

    fn do_finish(&mut self) -> io::Result<()> {
        if self.state != EncoderState::StreamEnd {
            self.state = EncoderState::Finished;
            // First, closes the stream.
            let mut buffer = zstd_sys::ZSTD_outBuffer {
                dst: self.buffer.as_mut_ptr() as *mut c_void,
                size: self.buffer.capacity(),
                pos: 0,
            };
            let remaining = parse_code(unsafe {
                zstd_sys::ZSTD_endStream(self.context.s,
                                   &mut buffer as *mut zstd_sys::ZSTD_outBuffer)
            })?;
            unsafe {
                self.buffer.set_len(buffer.pos);
            }
            if remaining != 0 {
                // Need to flush?
                panic!("Need to flush, but I'm lazy.");
            }
            self.offset = 0;
            self.state = EncoderState::StreamEnd;
        }

        self.write_from_offset()?;

        Ok(())
    }

    /// Return a recommendation for the size of data to write at once.
    pub fn recommended_input_size() -> usize {
        unsafe { zstd_sys::ZSTD_CStreamInSize() }
    }

    /// write_all, except keep track of partial writes for non-blocking IO.
    fn write_from_offset(&mut self) -> io::Result<()> {
        while self.offset < self.buffer.len() {
            match self.writer.write(&self.buffer[self.offset..]) {
                Ok(n) => self.offset += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {},
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }
}

impl<W: Write> Write for Encoder<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.state != EncoderState::Accepting {
            panic!("write called after try_finish attempted");
        }

        // Write any data pending in `self.buffer`.
        self.write_from_offset()?;

        // If we get to here, `self.buffer` can safely be discarded.

        let mut in_buffer = zstd_sys::ZSTD_inBuffer {
            src: buf.as_ptr() as *const c_void,
            size: buf.len(),
            pos: 0,
        };

        let mut out_buffer = zstd_sys::ZSTD_outBuffer {
            dst: self.buffer.as_mut_ptr() as *mut c_void,
            size: self.buffer.capacity(),
            pos: 0,
        };

        unsafe {
            // Time to fill our input buffer
            let code =
                zstd_sys::ZSTD_compressStream(self.context.s,
                                              &mut out_buffer as
                                              *mut zstd_sys::ZSTD_outBuffer,
                                              &mut in_buffer as
                                              *mut zstd_sys::ZSTD_inBuffer);
            // Note: this may very well be empty,
            // if it doesn't exceed zstd's own buffer
            self.buffer.set_len(out_buffer.pos);

            // Do we care about the hint?
            let _ = parse_code(code)?;
        }

        // This is the first time he sees this buffer.
        // Remember his delicate touch.
        self.offset = 0;
        self.write_from_offset()?;

        Ok(in_buffer.pos)
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.state == EncoderState::Accepting {
            let mut buffer = zstd_sys::ZSTD_outBuffer {
                dst: self.buffer.as_mut_ptr() as *mut c_void,
                size: self.buffer.capacity(),
                pos: 0,
            };
            unsafe {
                let code =
                    zstd_sys::ZSTD_flushStream(self.context.s,
                                               &mut buffer as
                                               *mut zstd_sys::ZSTD_outBuffer);
                self.buffer.set_len(buffer.pos);
                let _ = parse_code(code)?;
            }
            self.offset = 0;
        }

        self.write_from_offset()?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use {Poll, Async, Future, AsyncSink};
use stream::{Stream, Fuse};
use sink::Sink;

/// Future for the `Stream::forward` combinator, which sends a stream of values
/// to a sink and then waits until the sink has fully flushed those values.
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Forward<T: Stream, U> {
    sink: Option<U>,
    stream: Option<Fuse<T>>,
    buffered: Option<T::Item>,
}

pub fn new<T, U>(stream: T, sink: U) -> Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    Forward {
        sink: Some(sink),
        stream: Some(stream.fuse()),
        buffered: None,
    }
}

impl<T, U> Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    fn sink_mut(&mut self) -> &mut U {
        self.sink.as_mut().take()
            .expect("Attempted to poll Forward after completion")
    }

    fn stream_mut(&mut self) -> &mut Fuse<T> {
        self.stream.as_mut().take()
            .expect("Attempted to poll Forward after completion")
    }

    fn take_result(&mut self) -> (T, U) {
        let sink = self.sink.take()
            .expect("Attempted to poll Forward after completion");
        let fuse = self.stream.take()
            .expect("Attempted to poll Forward after completion");
        (fuse.into_inner(), sink)
    }

    fn try_start_send(&mut self, item: T::Item) -> Poll<(), U::SinkError> {
        debug_assert!(self.buffered.is_none());
        if let AsyncSink::NotReady(item) = self.sink_mut().start_send(item)? {
            self.buffered = Some(item);
            return Ok(Async::NotReady)
        }
        Ok(Async::Ready(()))
    }
}

impl<T, U> Future for Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    type Item = (T, U);
    type Error = T::Error;

    fn poll(&mut self) -> Poll<(T, U), T::Error> {
        // If we've got an item buffered already, we need to write it to the
        // sink before we can do anything else
        if let Some(item) = self.buffered.take() {
            try_ready!(self.try_start_send(item))
        }

        loop {
            match self.stream_mut().poll()? {
                Async::Ready(Some(item)) => try_ready!(self.try_start_send(item)),
                Async::Ready(None) => {
                    try_ready!(self.sink_mut().close());
                    return Ok(Async::Ready(self.take_result()))
                }
                Async::NotReady => {
                    try_ready!(self.sink_mut().poll_complete());
                    return Ok(Async::NotReady)
                }
            }
        }
    }
}
<commit_msg>Add standard stream accessors to Forward future<commit_after>use {Poll, Async, Future, AsyncSink};
use stream::{Stream, Fuse};
use sink::Sink;

/// Future for the `Stream::forward` combinator, which sends a stream of values
/// to a sink and then waits until the sink has fully flushed those values.
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Forward<T: Stream, U> {
    sink: Option<U>,
    stream: Option<Fuse<T>>,
    buffered: Option<T::Item>,
}

pub fn new<T, U>(stream: T, sink: U) -> Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    Forward {
        sink: Some(sink),
        stream: Some(stream.fuse()),
        buffered: None,
    }
}

impl<T, U> Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    /// Get a shared reference to the inner sink.
    /// If this combintator has already been polled to completetion, None will be returned.
    pub fn sink_ref(&self) -> Option<&U> {
        self.sink.as_ref()
    }

    /// Get a mutable reference to the inner sink.
    /// If this combintator has already been polled to completetion, None will be returned.
    pub fn sink_mut(&mut self) -> Option<&mut U> {
        self.sink.as_mut()
    }

    /// Get a shared reference to the inner stream.
    /// If this combintator has already been polled to completetion, None will be returned.
    pub fn stream_ref(&self) -> Option<&T> {
        self.stream.as_ref().map(|x| x.get_ref())
    }

    /// Get a mutable reference to the inner stream.
    /// If this combintator has already been polled to completetion, None will be returned.
    pub fn stream_mut(&mut self) -> Option<&mut T> {
        self.stream.as_mut().map(|x| x.get_mut())
    }

    fn take_result(&mut self) -> (T, U) {
        let sink = self.sink.take()
            .expect("Attempted to poll Forward after completion");
        let fuse = self.stream.take()
            .expect("Attempted to poll Forward after completion");
        (fuse.into_inner(), sink)
    }

    fn try_start_send(&mut self, item: T::Item) -> Poll<(), U::SinkError> {
        debug_assert!(self.buffered.is_none());
        if let AsyncSink::NotReady(item) = self.sink_mut()
            .take().expect("Attempted to poll Forward after completion")
            .start_send(item)?
        {
            self.buffered = Some(item);
            return Ok(Async::NotReady)
        }
        Ok(Async::Ready(()))
    }
}

impl<T, U> Future for Forward<T, U>
    where U: Sink<SinkItem=T::Item>,
          T: Stream,
          T::Error: From<U::SinkError>,
{
    type Item = (T, U);
    type Error = T::Error;

    fn poll(&mut self) -> Poll<(T, U), T::Error> {
        // If we've got an item buffered already, we need to write it to the
        // sink before we can do anything else
        if let Some(item) = self.buffered.take() {
            try_ready!(self.try_start_send(item))
        }

        loop {
            match self.stream_mut()
                .take().expect("Attempted to poll Forward after completion")
                .poll()?
            {
                Async::Ready(Some(item)) => try_ready!(self.try_start_send(item)),
                Async::Ready(None) => {
                    try_ready!(self.sink_mut().take().expect("Attempted to poll Forward after completion").close());
                    return Ok(Async::Ready(self.take_result()))
                }
                Async::NotReady => {
                    try_ready!(self.sink_mut().take().expect("Attempted to poll Forward after completion").poll_complete());
                    return Ok(Async::NotReady)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use itertools;

use error::{Error, Result};

use interpreter::Argument;
use interpreter::Context;
use interpreter::Renderable;
use compiler::Token;
use compiler::LiquidOptions;
use compiler::{consume_value_token, value_token};

#[derive(Clone, Debug)]
struct Cycle {
    name: String,
    values: Vec<Argument>,
}

impl Renderable for Cycle {
    fn render(&self, context: &mut Context) -> Result<Option<String>> {
        let value = context.cycle_element(&self.name, &self.values)?;
        Ok(value.map(|v| v.to_string()))
    }
}

/// Internal implementation of cycle, to allow easier testing.
fn parse_cycle(arguments: &[Token], _options: &LiquidOptions) -> Result<Cycle> {
    let mut args = arguments.iter();
    let mut name = String::new();
    let mut values = Vec::new();
    let first = consume_value_token(&mut args)?;

    match args.next() {
        Some(&Token::Colon) => {
            // the first argument is the name of the cycle block
            name = first.to_string();
        }
        Some(&Token::Comma) |
        None => {
            // first argument is the first item in the cycle
            values.push(first.to_arg()?);
        }
        x => return Error::parser(": | Number | String | Identifier", x),
    }

    loop {
        match args.next() {
            Some(a) => {
                let v = value_token(a.clone())?.to_arg()?;
                values.push(v);
            }
            None => break,
        }

        match args.next() {
            Some(&Token::Comma) => {}
            None => break,
            x => return Error::parser("Comma", x),
        }
    }

    if name.is_empty() {
        name = itertools::join(values.iter(), "-");
    }
    println!("name={}", name);

    Ok(Cycle {
           name: name,
           values: values,
       })
}

pub fn cycle_tag(_tag_name: &str,
                 arguments: &[Token],
                 options: &LiquidOptions)
                 -> Result<Box<Renderable>> {
    parse_cycle(arguments, options).map(|opt| Box::new(opt) as Box<Renderable>)
}

#[cfg(test)]
mod test {
    use super::*;
    use value::Value;
    use compiler;
    use interpreter;

    fn options() -> LiquidOptions {
        let mut options = LiquidOptions::default();
        options.tags.insert("cycle".to_owned(),
                            (cycle_tag as compiler::FnParseTag).into());
        options
    }

    #[test]
    fn unnamed_cycle_gets_a_name() {
        let tokens = vec![Token::Identifier("this".to_owned()),
                          Token::Comma,
                          Token::StringLiteral("cycle".to_owned()),
                          Token::Comma,
                          Token::Identifier("has".to_owned()),
                          Token::Comma,
                          Token::Identifier("no".to_owned()),
                          Token::Comma,
                          Token::Identifier("name".to_owned())];
        let options = LiquidOptions::default();
        let cycle = parse_cycle(&tokens[..], &options).unwrap();
        assert!(!cycle.name.is_empty());
    }

    #[test]
    fn named_values_are_independent() {
        let text = concat!("{% cycle 'a': 'one', 'two', 'three' %}\n",
                           "{% cycle 'a': 'one', 'two', 'three' %}\n",
                           "{% cycle 'b': 'one', 'two', 'three' %}\n",
                           "{% cycle 'b': 'one', 'two', 'three' %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("one\ntwo\none\ntwo\n".to_owned()));
    }

    #[test]
    fn values_are_cycled() {
        let text = concat!("{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("one\ntwo\nthree\none\n".to_owned()));
    }

    #[test]
    fn values_can_be_variables() {
        let text = concat!("{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        context.set_global_val("alpha", Value::Num(1f32));
        context.set_global_val("beta", Value::Num(2f32));
        context.set_global_val("gamma", Value::Num(3f32));

        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("1\n2\n3\n1\n".to_owned()));
    }

    #[test]
    fn bad_cycle_indices_dont_crash() {
        // note the pair of cycle tags with the same name but a differing
        // number of elements
        let text = concat!("{% cycle c: 1, 2 %}\n", "{% cycle c: 1 %}\n").to_owned();

        let tokens = compiler::tokenize(&text).unwrap();
        let options = options();
        let template = compiler::parse(&tokens, &options)
            .map(interpreter::Template::new)
            .unwrap();
        let output = template.render(&mut Default::default());
        assert!(output.is_err());
    }
}
<commit_msg>fix(dbg): Remove debug code<commit_after>use itertools;

use error::{Error, Result};

use interpreter::Argument;
use interpreter::Context;
use interpreter::Renderable;
use compiler::Token;
use compiler::LiquidOptions;
use compiler::{consume_value_token, value_token};

#[derive(Clone, Debug)]
struct Cycle {
    name: String,
    values: Vec<Argument>,
}

impl Renderable for Cycle {
    fn render(&self, context: &mut Context) -> Result<Option<String>> {
        let value = context.cycle_element(&self.name, &self.values)?;
        Ok(value.map(|v| v.to_string()))
    }
}

/// Internal implementation of cycle, to allow easier testing.
fn parse_cycle(arguments: &[Token], _options: &LiquidOptions) -> Result<Cycle> {
    let mut args = arguments.iter();
    let mut name = String::new();
    let mut values = Vec::new();
    let first = consume_value_token(&mut args)?;

    match args.next() {
        Some(&Token::Colon) => {
            // the first argument is the name of the cycle block
            name = first.to_string();
        }
        Some(&Token::Comma) |
        None => {
            // first argument is the first item in the cycle
            values.push(first.to_arg()?);
        }
        x => return Error::parser(": | Number | String | Identifier", x),
    }

    loop {
        match args.next() {
            Some(a) => {
                let v = value_token(a.clone())?.to_arg()?;
                values.push(v);
            }
            None => break,
        }

        match args.next() {
            Some(&Token::Comma) => {}
            None => break,
            x => return Error::parser("Comma", x),
        }
    }

    if name.is_empty() {
        name = itertools::join(values.iter(), "-");
    }

    Ok(Cycle {
           name: name,
           values: values,
       })
}

pub fn cycle_tag(_tag_name: &str,
                 arguments: &[Token],
                 options: &LiquidOptions)
                 -> Result<Box<Renderable>> {
    parse_cycle(arguments, options).map(|opt| Box::new(opt) as Box<Renderable>)
}

#[cfg(test)]
mod test {
    use super::*;
    use value::Value;
    use compiler;
    use interpreter;

    fn options() -> LiquidOptions {
        let mut options = LiquidOptions::default();
        options.tags.insert("cycle".to_owned(),
                            (cycle_tag as compiler::FnParseTag).into());
        options
    }

    #[test]
    fn unnamed_cycle_gets_a_name() {
        let tokens = vec![Token::Identifier("this".to_owned()),
                          Token::Comma,
                          Token::StringLiteral("cycle".to_owned()),
                          Token::Comma,
                          Token::Identifier("has".to_owned()),
                          Token::Comma,
                          Token::Identifier("no".to_owned()),
                          Token::Comma,
                          Token::Identifier("name".to_owned())];
        let options = LiquidOptions::default();
        let cycle = parse_cycle(&tokens[..], &options).unwrap();
        assert!(!cycle.name.is_empty());
    }

    #[test]
    fn named_values_are_independent() {
        let text = concat!("{% cycle 'a': 'one', 'two', 'three' %}\n",
                           "{% cycle 'a': 'one', 'two', 'three' %}\n",
                           "{% cycle 'b': 'one', 'two', 'three' %}\n",
                           "{% cycle 'b': 'one', 'two', 'three' %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("one\ntwo\none\ntwo\n".to_owned()));
    }

    #[test]
    fn values_are_cycled() {
        let text = concat!("{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n",
                           "{% cycle 'one', 'two', 'three' %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("one\ntwo\nthree\none\n".to_owned()));
    }

    #[test]
    fn values_can_be_variables() {
        let text = concat!("{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n",
                           "{% cycle alpha, beta, gamma %}\n")
            .to_owned();
        let tokens = compiler::tokenize(&text).unwrap();
        let template = compiler::parse(&tokens, &options())
            .map(interpreter::Template::new)
            .unwrap();

        let mut context = Context::new();
        context.set_global_val("alpha", Value::Num(1f32));
        context.set_global_val("beta", Value::Num(2f32));
        context.set_global_val("gamma", Value::Num(3f32));

        let output = template.render(&mut context);

        assert_eq!(output.unwrap(), Some("1\n2\n3\n1\n".to_owned()));
    }

    #[test]
    fn bad_cycle_indices_dont_crash() {
        // note the pair of cycle tags with the same name but a differing
        // number of elements
        let text = concat!("{% cycle c: 1, 2 %}\n", "{% cycle c: 1 %}\n").to_owned();

        let tokens = compiler::tokenize(&text).unwrap();
        let options = options();
        let template = compiler::parse(&tokens, &options)
            .map(interpreter::Template::new)
            .unwrap();
        let output = template.render(&mut Default::default());
        assert!(output.is_err());
    }
}
<|endoftext|>


<commit_before>//! Provides a Vec-like container for large sizes that is split to blocks.

use std::fmt;
use std::ops;
use std::ops::{Range, RangeFrom, RangeTo, RangeFull};
use std::cmp;

use itertools;
use odds::vec::VecExt;

/// A generic trait over Rust's built types.
pub trait FromRange {
    #[inline(always)]
    #[doc(hidden)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize);
}

impl FromRange for RangeFull {
    #[inline(always)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize) {
        return (0, seg.len());
    }
}

impl FromRange for Range<usize> {
    #[inline(always)]
    fn from_range(&self, _: &SplitVec) -> (usize, usize) {
        return (self.start, self.end);
    }
}

impl FromRange for RangeFrom<usize> {
    #[inline(always)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize) {
        return (self.start, seg.len());
    }
}

impl FromRange for RangeTo<usize> {
    #[inline(always)]
    fn from_range(&self, _: &SplitVec) -> (usize, usize) {
        return (0, self.end);
    }
}

/// A Vec for large sizes split into smaller blocks.
///
/// Splits a large continuos memory into block sizes to allow insertion/deletion to have an
/// upper bound, based on the block size. It is implemented as a Vec of Vec's, where each Vec
/// has a m inimum and maximum block size. As data is inserted and deleted, from any location,
/// SplitVec will split up and merge the blocks to try and stay between those minimum and maximum
/// block sizes.
pub struct SplitVec {
    vecs: Vec<Vec<u8>>,
    length: usize,
}

#[derive(Copy, Clone, Debug)]
struct Index {
    outer: usize,
    inner: usize,
}

/// An iterator over SplitVec contents.
pub struct Items<'a> {
    seg: &'a SplitVec,
    index: Index,
    num_elem: Option<usize>,
}

/// A mutable iterator over SplitVec contents.
pub struct MutItems<'a> {
    seg: &'a mut SplitVec,
    index: Index,
    num_elem: Option<usize>,
}

/// An iterator over the blocks in SplitVec.
pub struct Slices<'a> {
    seg: &'a SplitVec,
    outer: usize,
}

static MIN_BLOCK_SIZE: usize = 1024 * 1024;
static MAX_BLOCK_SIZE: usize = 4 * 1024 * 1024;

impl SplitVec {
    /// Create a new, empty SplitVec
    pub fn new() -> SplitVec {
        SplitVec {
            vecs: Vec::new(),
            length: 0,
        }
    }

    /// Create a SplitVec by consuming a vec as the initial data vector
    pub fn from_vec(values: Vec<u8>) -> SplitVec {
        let len = values.len();
        SplitVec {
            vecs: vec![values],
            length: len,
        }
    }

    /// Create a SplitVec by copying in values from a slice
    pub fn from_slice(values: &[u8]) -> SplitVec {
        SplitVec {
            vecs: vec![values.into()],
            length: values.len(),
        }
    }

    /// Return the length.
    pub fn len(&self) -> usize {
        self.length
    }

    /// Update the saved length value so that the len func will be -O(1)
    fn calc_len(&mut self) {
        self.length = 0;
        for len in self.vecs.iter().map(|v| v.len()) {
            self.length += len
        }
    }

    /// Convert a global pos to a locall index
    fn pos_to_index(&self, pos: usize, for_insert: bool) -> Index {
        if pos == 0 {
            return Index { outer: 0, inner: 0 };
        }

        let mut cur_pos = pos;
        for (i, vec) in self.vecs.iter().enumerate() {
            if cur_pos < vec.len() || (for_insert && cur_pos == vec.len()) {
                return Index {
                    outer: i,
                    inner: cur_pos,
                }
            }
            cur_pos -= vec.len();
        }

        panic!("Position {} is out of bounds", pos);
    }

    /// Give an iterator over a given range
    pub fn iter_range<'a, R: FromRange>(&'a self, range: R) -> Items<'a> {
        let (from, to) = range.from_range(self);
        if to < from {
            panic!("to ({}) is smaller than from ({})!", to, from);
        }

        let idx = self.pos_to_index(from, false);
        Items {
            seg: self,
            index: idx,
            num_elem: Some(to - from),
        }
    }

    /// Give a mutable iterator over a given range
    pub fn mut_iter_range<'a, R: FromRange>(&'a mut self, range: R) -> MutItems<'a> {
        let (from, to) = range.from_range(self);
        if to < from {
            panic!("to ({}) is smaller than from ({})!", to, from);
        }

        let idx = self.pos_to_index(from, false);
        MutItems {
            seg: self,
            index: idx,
            num_elem: Some(to - from),
        }
    }

    /// Provide an iterator over continuous memory slices
    ///
    /// This is useful for doing an efficient save to disk
    pub fn iter_slices<'a>(&'a self) -> Slices<'a> {
        Slices {
            seg: self,
            outer: 0,
        }
    }

    /// Prepare an index for future text insertion, splitting/merging big/small sections respectively
    fn prepare_insert(&mut self, index: Index) -> Index {
        if index.outer >= self.vecs.len() {
            self.vecs.push(Vec::new());
        }

        if self.vecs[index.outer].len() < MAX_BLOCK_SIZE {
            return index;
        }

        let page_start_idx = (index.inner / MIN_BLOCK_SIZE) * MIN_BLOCK_SIZE;
        if page_start_idx == 0 {
            if self.vecs[index.outer].len() > MAX_BLOCK_SIZE {
                let insert_vec: Vec < _ >= self.vecs[index.outer][MIN_BLOCK_SIZE..].into();
                self.vecs.insert(index.outer + 1, insert_vec);
                self.vecs[index.outer].truncate(MIN_BLOCK_SIZE);
            }

            return index;
        } else {
            let insert_vec: Vec<_> = self.vecs[index.outer][page_start_idx..].into();
            self.vecs.insert(index.outer + 1, insert_vec);
            self.vecs[index.outer].truncate(page_start_idx);
            return self.prepare_insert(Index {
                outer: index.outer + 1,
                inner: index.inner - page_start_idx
            })
        }
    }

    /// insert all values from a slice at an offset.
    pub fn insert(&mut self, offset: usize, values: &[u8]) {
        let mut index = self.pos_to_index(offset, true);
        index = self.prepare_insert(index);

        // This is needed for the mut borrow vec
        {
            self.vecs[index.outer].splice(index.inner..index.inner, values.into_iter().cloned());
        }

        self.calc_len();
    }

    /// Moves data out from the supplied range.
    pub fn move_out<R: FromRange>(&mut self, range: R) -> Vec<u8> {
        let (from, to) = range.from_range(self);
        // TODO: Convert to drain when that settles
        assert!(from <= to);
        let mut res = Vec::new();
        let mut index = self.pos_to_index(from, false);
        let num_elem = to - from;

        for _ in 0..num_elem {
            let c = self.vecs[index.outer].remove(index.inner);
            res.push(c);

            if index.inner >= self.vecs[index.outer].len() {
                if self.vecs[index.outer].len() == 0 {
                    self.vecs.remove(index.outer);
                } else {
                    index.inner = 0;
                    index.outer += 1;
                }
            }
        }

        self.calc_len();

        res
    }

    /// Produce of copy of the supplied range
    pub fn copy_out<R: FromRange>(&mut self, range: R) -> Vec<u8> {
        let (from, to) = range.from_range(self);

        // TODO: Make this use direct cocpy rather than iterators
        self.iter_range(from..to).map(|x| *x).collect::<Vec<u8>>()
    }

    /// Copy data from a slice in.
    pub fn copy_in(&mut self, offset: usize, val: &[u8]) {
        for (s, d) in val.iter().zip(self.mut_iter_range(offset..(offset + val.len()))) {
            *d = s.clone();
        }
    }

    /// Replace values in range with the supplied values
    pub fn splice<R: FromRange>(&mut self, range: R, values: &[u8]) -> Vec<u8> {
        // TODO: This is terribly inefficient, will need a reimplementation
        let (from, to) = range.from_range(self);
        let res;

        // Make sure that when we pull data out for the splice, we don't go over the end
        if from < self.len() {
            let move_end = cmp::min(self.len(), to);
            res = self.move_out(from..move_end);
        } else {
            res = vec![];
        }

        self.insert(from, values);

        res
    }

    /// Find a slice.
    pub fn find_slice(&self, needle: &[u8]) -> Option<usize> {
        self.find_slice_from(0, needle)
    }

    /// Find a slice from a certain index and onward
    pub fn find_slice_from(&self, from: usize, needle: &[u8]) -> Option<usize> {
        for i in from..self.len() {
            if itertools::equal(self.iter_range(i..i+needle.len()), needle.iter()) {
                return Some(i);
            }
        }
        None
    }

    #[cfg(test)]
    fn get_lengths(&self) -> Vec<usize> {
        self.vecs.iter().map(|v| v.len()).collect::<Vec<usize>>()
    }
}

impl ops::Index<usize> for SplitVec {
    type Output = u8;
    fn index<'a>(&'a self, _index: usize) -> &'a u8 {
        let idx = self.pos_to_index(_index, false);
        &self.vecs[idx.outer][idx.inner]
    }
}

impl ops::IndexMut<usize> for SplitVec {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut u8 {
        let idx = self.pos_to_index(_index, false);
        &mut self.vecs[idx.outer][idx.inner]
    }
}

impl fmt::Debug for SplitVec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.vecs.fmt(f)
    }
}

impl<'a> Iterator for Items<'a> {
    type Item = &'a u8;
    fn next(&mut self) -> Option<&'a u8> {
        if self.index.outer >= self.seg.vecs.len() {
            return None;
        }
        if let Some(ref mut num_elem) = self.num_elem {
            if *num_elem <= 0 {
                return None;
            }
            *num_elem -= 1;
        }

        let elem = {
            let vv = &self.seg.vecs[self.index.outer];
            &vv[self.index.inner]
        };

        self.index.inner += 1;
        if self.index.inner >= self.seg.vecs[self.index.outer].len() {
            self.index.inner = 0;
            self.index.outer += 1;
        }

        Some(elem)
    }
}

impl<'a> Iterator for MutItems<'a> {
    type Item = &'a mut u8;
    fn next(&mut self) -> Option<&'a mut u8> {
        if self.index.outer >= self.seg.vecs.len() {
            return None;
        }
        if let Some(ref mut num_elem) = self.num_elem {
            if *num_elem <= 0 {
                return None;
            }
            *num_elem -= 1;
        }

        let elem_raw: *mut u8 = {
            let vv = &mut self.seg.vecs[self.index.outer];
            &mut vv[self.index.inner]
        };

        self.index.inner += 1;
        if self.index.inner >= self.seg.vecs[self.index.outer].len() {
            self.index.inner = 0;
            self.index.outer += 1;
        }

        Some(unsafe { &mut *elem_raw })
    }
}

impl<'a> Iterator for Slices<'a> {
    type Item = &'a [u8];
    fn next(&mut self) -> Option<&'a [u8]> {
        if self.outer >= self.seg.vecs.len() {
            None
        } else {
            let i = self.outer;
            self.outer += 1;
            Some(&self.seg.vecs[i])
        }
    }
}

#[test]
fn test_small_splitvec() {
    let size = 1024;
    let mut seg = SplitVec::from_vec(vec![1, 2, 3, 4, 5]);
    assert_eq!(Some(4), seg.find_slice(&[5]));

    let seg_len = seg.len();
    seg.insert(seg_len/2, &vec![1 as u8; size]);

    assert_eq!(Some(size + 4), seg.find_slice(&[5]));
}

#[test]
fn test_large_splitvec() {
    let big_size = 4*1024*1024;
    let small_size = 1024;
    let mut seg = SplitVec::from_vec(vec![0; big_size]);

    seg.insert(big_size/2, &vec![1 as u8; small_size]);

    assert_eq!(Some(big_size/2 -1), seg.find_slice(&[0, 1]));

    // Make sure we actually tested a "split" version
    let seg_lengths = seg.get_lengths();
    assert_eq!(2, seg_lengths.len());
    let index = seg_lengths[0];
    let sentinal = 100;
    seg[index] = sentinal;
    seg[index+1] = sentinal +1;
    assert_eq!(Some(index), seg.find_slice(&[sentinal, sentinal+1]));
}
<commit_msg>Make unused externally methods, private<commit_after>//! Provides a Vec-like container for large sizes that is split to blocks.

use std::fmt;
use std::ops;
use std::ops::{Range, RangeFrom, RangeTo, RangeFull};
use std::cmp;

use itertools;
use odds::vec::VecExt;

/// A generic trait over Rust's built types.
pub trait FromRange {
    #[inline(always)]
    #[doc(hidden)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize);
}

impl FromRange for RangeFull {
    #[inline(always)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize) {
        return (0, seg.len());
    }
}

impl FromRange for Range<usize> {
    #[inline(always)]
    fn from_range(&self, _: &SplitVec) -> (usize, usize) {
        return (self.start, self.end);
    }
}

impl FromRange for RangeFrom<usize> {
    #[inline(always)]
    fn from_range(&self, seg: &SplitVec) -> (usize, usize) {
        return (self.start, seg.len());
    }
}

impl FromRange for RangeTo<usize> {
    #[inline(always)]
    fn from_range(&self, _: &SplitVec) -> (usize, usize) {
        return (0, self.end);
    }
}

/// A Vec for large sizes split into smaller blocks.
///
/// Splits a large continuos memory into block sizes to allow insertion/deletion to have an
/// upper bound, based on the block size. It is implemented as a Vec of Vec's, where each Vec
/// has a m inimum and maximum block size. As data is inserted and deleted, from any location,
/// SplitVec will split up and merge the blocks to try and stay between those minimum and maximum
/// block sizes.
pub struct SplitVec {
    vecs: Vec<Vec<u8>>,
    length: usize,
}

#[derive(Copy, Clone, Debug)]
struct Index {
    outer: usize,
    inner: usize,
}

/// An iterator over SplitVec contents.
pub struct Items<'a> {
    seg: &'a SplitVec,
    index: Index,
    num_elem: Option<usize>,
}

/// A mutable iterator over SplitVec contents.
pub struct MutItems<'a> {
    seg: &'a mut SplitVec,
    index: Index,
    num_elem: Option<usize>,
}

/// An iterator over the blocks in SplitVec.
pub struct Slices<'a> {
    seg: &'a SplitVec,
    outer: usize,
}

static MIN_BLOCK_SIZE: usize = 1024 * 1024;
static MAX_BLOCK_SIZE: usize = 4 * 1024 * 1024;

impl SplitVec {
    /// Create a new, empty SplitVec
    pub fn new() -> SplitVec {
        SplitVec {
            vecs: Vec::new(),
            length: 0,
        }
    }

    /// Create a SplitVec by consuming a vec as the initial data vector
    pub fn from_vec(values: Vec<u8>) -> SplitVec {
        let len = values.len();
        SplitVec {
            vecs: vec![values],
            length: len,
        }
    }

    /// Create a SplitVec by copying in values from a slice
    pub fn from_slice(values: &[u8]) -> SplitVec {
        SplitVec {
            vecs: vec![values.into()],
            length: values.len(),
        }
    }

    /// Return the length.
    pub fn len(&self) -> usize {
        self.length
    }

    /// Update the saved length value so that the len func will be -O(1)
    fn calc_len(&mut self) {
        self.length = 0;
        for len in self.vecs.iter().map(|v| v.len()) {
            self.length += len
        }
    }

    /// Convert a global pos to a locall index
    fn pos_to_index(&self, pos: usize, for_insert: bool) -> Index {
        if pos == 0 {
            return Index { outer: 0, inner: 0 };
        }

        let mut cur_pos = pos;
        for (i, vec) in self.vecs.iter().enumerate() {
            if cur_pos < vec.len() || (for_insert && cur_pos == vec.len()) {
                return Index {
                    outer: i,
                    inner: cur_pos,
                }
            }
            cur_pos -= vec.len();
        }

        panic!("Position {} is out of bounds", pos);
    }

    /// Give an iterator over a given range
    pub fn iter_range<'a, R: FromRange>(&'a self, range: R) -> Items<'a> {
        let (from, to) = range.from_range(self);
        if to < from {
            panic!("to ({}) is smaller than from ({})!", to, from);
        }

        let idx = self.pos_to_index(from, false);
        Items {
            seg: self,
            index: idx,
            num_elem: Some(to - from),
        }
    }

    /// Give a mutable iterator over a given range
    fn mut_iter_range<'a, R: FromRange>(&'a mut self, range: R) -> MutItems<'a> {
        let (from, to) = range.from_range(self);
        if to < from {
            panic!("to ({}) is smaller than from ({})!", to, from);
        }

        let idx = self.pos_to_index(from, false);
        MutItems {
            seg: self,
            index: idx,
            num_elem: Some(to - from),
        }
    }

    /// Provide an iterator over continuous memory slices
    ///
    /// This is useful for doing an efficient save to disk
    pub fn iter_slices<'a>(&'a self) -> Slices<'a> {
        Slices {
            seg: self,
            outer: 0,
        }
    }

    /// Prepare an index for future text insertion, splitting/merging big/small sections respectively
    fn prepare_insert(&mut self, index: Index) -> Index {
        if index.outer >= self.vecs.len() {
            self.vecs.push(Vec::new());
        }

        if self.vecs[index.outer].len() < MAX_BLOCK_SIZE {
            return index;
        }

        let page_start_idx = (index.inner / MIN_BLOCK_SIZE) * MIN_BLOCK_SIZE;
        if page_start_idx == 0 {
            if self.vecs[index.outer].len() > MAX_BLOCK_SIZE {
                let insert_vec: Vec < _ >= self.vecs[index.outer][MIN_BLOCK_SIZE..].into();
                self.vecs.insert(index.outer + 1, insert_vec);
                self.vecs[index.outer].truncate(MIN_BLOCK_SIZE);
            }

            return index;
        } else {
            let insert_vec: Vec<_> = self.vecs[index.outer][page_start_idx..].into();
            self.vecs.insert(index.outer + 1, insert_vec);
            self.vecs[index.outer].truncate(page_start_idx);
            return self.prepare_insert(Index {
                outer: index.outer + 1,
                inner: index.inner - page_start_idx
            })
        }
    }

    /// insert all values from a slice at an offset.
    fn insert(&mut self, offset: usize, values: &[u8]) {
        let mut index = self.pos_to_index(offset, true);
        index = self.prepare_insert(index);

        // This is needed for the mut borrow vec
        {
            self.vecs[index.outer].splice(index.inner..index.inner, values.into_iter().cloned());
        }

        self.calc_len();
    }

    /// Moves data out from the supplied range.
    fn move_out<R: FromRange>(&mut self, range: R) -> Vec<u8> {
        let (from, to) = range.from_range(self);
        // TODO: Convert to drain when that settles
        assert!(from <= to);
        let mut res = Vec::new();
        let mut index = self.pos_to_index(from, false);
        let num_elem = to - from;

        for _ in 0..num_elem {
            let c = self.vecs[index.outer].remove(index.inner);
            res.push(c);

            if index.inner >= self.vecs[index.outer].len() {
                if self.vecs[index.outer].len() == 0 {
                    self.vecs.remove(index.outer);
                } else {
                    index.inner = 0;
                    index.outer += 1;
                }
            }
        }

        self.calc_len();

        res
    }

    /// Produce of copy of the supplied range
    pub fn copy_out<R: FromRange>(&mut self, range: R) -> Vec<u8> {
        let (from, to) = range.from_range(self);

        // TODO: Make this use direct cocpy rather than iterators
        self.iter_range(from..to).map(|x| *x).collect::<Vec<u8>>()
    }

    /// Copy data from a slice in.
    fn copy_in(&mut self, offset: usize, val: &[u8]) {
        for (s, d) in val.iter().zip(self.mut_iter_range(offset..(offset + val.len()))) {
            *d = s.clone();
        }
    }

    /// Replace values in range with the supplied values
    pub fn splice<R: FromRange>(&mut self, range: R, values: &[u8]) -> Vec<u8> {
        // TODO: This is terribly inefficient, will need a reimplementation
        let (from, to) = range.from_range(self);
        let res;

        // Make sure that when we pull data out for the splice, we don't go over the end
        if from < self.len() {
            let move_end = cmp::min(self.len(), to);
            res = self.move_out(from..move_end);
        } else {
            res = vec![];
        }

        self.insert(from, values);

        res
    }

    /// Find a slice.
    pub fn find_slice(&self, needle: &[u8]) -> Option<usize> {
        self.find_slice_from(0, needle)
    }

    /// Find a slice from a certain index and onward
    pub fn find_slice_from(&self, from: usize, needle: &[u8]) -> Option<usize> {
        for i in from..self.len() {
            if itertools::equal(self.iter_range(i..i+needle.len()), needle.iter()) {
                return Some(i);
            }
        }
        None
    }

    #[cfg(test)]
    fn get_lengths(&self) -> Vec<usize> {
        self.vecs.iter().map(|v| v.len()).collect::<Vec<usize>>()
    }
}

impl ops::Index<usize> for SplitVec {
    type Output = u8;
    fn index<'a>(&'a self, _index: usize) -> &'a u8 {
        let idx = self.pos_to_index(_index, false);
        &self.vecs[idx.outer][idx.inner]
    }
}

impl ops::IndexMut<usize> for SplitVec {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut u8 {
        let idx = self.pos_to_index(_index, false);
        &mut self.vecs[idx.outer][idx.inner]
    }
}

impl fmt::Debug for SplitVec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.vecs.fmt(f)
    }
}

impl<'a> Iterator for Items<'a> {
    type Item = &'a u8;
    fn next(&mut self) -> Option<&'a u8> {
        if self.index.outer >= self.seg.vecs.len() {
            return None;
        }
        if let Some(ref mut num_elem) = self.num_elem {
            if *num_elem <= 0 {
                return None;
            }
            *num_elem -= 1;
        }

        let elem = {
            let vv = &self.seg.vecs[self.index.outer];
            &vv[self.index.inner]
        };

        self.index.inner += 1;
        if self.index.inner >= self.seg.vecs[self.index.outer].len() {
            self.index.inner = 0;
            self.index.outer += 1;
        }

        Some(elem)
    }
}

impl<'a> Iterator for MutItems<'a> {
    type Item = &'a mut u8;
    fn next(&mut self) -> Option<&'a mut u8> {
        if self.index.outer >= self.seg.vecs.len() {
            return None;
        }
        if let Some(ref mut num_elem) = self.num_elem {
            if *num_elem <= 0 {
                return None;
            }
            *num_elem -= 1;
        }

        let elem_raw: *mut u8 = {
            let vv = &mut self.seg.vecs[self.index.outer];
            &mut vv[self.index.inner]
        };

        self.index.inner += 1;
        if self.index.inner >= self.seg.vecs[self.index.outer].len() {
            self.index.inner = 0;
            self.index.outer += 1;
        }

        Some(unsafe { &mut *elem_raw })
    }
}

impl<'a> Iterator for Slices<'a> {
    type Item = &'a [u8];
    fn next(&mut self) -> Option<&'a [u8]> {
        if self.outer >= self.seg.vecs.len() {
            None
        } else {
            let i = self.outer;
            self.outer += 1;
            Some(&self.seg.vecs[i])
        }
    }
}

#[test]
fn test_small_splitvec() {
    let size = 1024;
    let mut seg = SplitVec::from_vec(vec![1, 2, 3, 4, 5]);
    assert_eq!(Some(4), seg.find_slice(&[5]));

    let seg_len = seg.len();
    seg.insert(seg_len/2, &vec![1 as u8; size]);

    assert_eq!(Some(size + 4), seg.find_slice(&[5]));
}

#[test]
fn test_large_splitvec() {
    let big_size = 4*1024*1024;
    let small_size = 1024;
    let mut seg = SplitVec::from_vec(vec![0; big_size]);

    seg.insert(big_size/2, &vec![1 as u8; small_size]);

    assert_eq!(Some(big_size/2 -1), seg.find_slice(&[0, 1]));

    // Make sure we actually tested a "split" version
    let seg_lengths = seg.get_lengths();
    assert_eq!(2, seg_lengths.len());
    let index = seg_lengths[0];
    let sentinal = 100;
    seg[index] = sentinal;
    seg[index+1] = sentinal +1;
    assert_eq!(Some(index), seg.find_slice(&[sentinal, sentinal+1]));
}
<|endoftext|>


<commit_before>//! Server latency statistic

use std::collections::VecDeque;

/// Interval between each check
pub const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10;
/// Timeout of each check
pub const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5; // A common connection timeout of 5 seconds.

const MAX_SERVER_RTT: u32 = DEFAULT_CHECK_TIMEOUT_SEC as u32 * 1000;
const MAX_LATENCY_QUEUE_SIZE: usize = 59; // Account for the last 10 minutes.

/// Statistic score
#[derive(Debug, Copy, Clone)]
pub enum Score {
    /// Unified latency
    Latency(u32),
    /// Request error
    Errored,
}

/// Statistic of a remote server
#[derive(Debug)]
pub struct ServerStat {
    /// Median of latency time (in millisec)
    ///
    /// Use median instead of average time,
    /// because probing result may have some really bad cases
    rtt: u32,
    /// Total_Fail / Total_Probe
    fail_rate: f64,
    /// Recently probe data
    latency_queue: VecDeque<Score>,
    /// Score's standard deviation
    latency_stdev: f64,
    /// Score's average
    latency_mean: f64,
    /// User's customized weight
    user_weight: f32,
}

fn max_latency_stdev() -> f64 {
    let mrtt = MAX_SERVER_RTT as f64;
    let avg = (0.0 + mrtt) / 2.0;
    let diff1 = (0.0 - avg) * (0.0 - avg);
    let diff2 = (mrtt - avg) * (mrtt - avg);
    // (1.0 / (2.0 - 1.0)) * (diff1 + diff2).sqrt()
    (diff1 + diff2).sqrt()
}

impl ServerStat {
    pub fn new(user_weight: f32) -> ServerStat {
        assert!((0.0..=1.0).contains(&user_weight));

        ServerStat {
            rtt: MAX_SERVER_RTT,
            fail_rate: 1.0,
            latency_queue: VecDeque::new(),
            latency_stdev: 0.0,
            latency_mean: 0.0,
            user_weight,
        }
    }

    fn score(&self) -> u32 {
        // Normalize rtt
        let nrtt = self.rtt as f64 / MAX_SERVER_RTT as f64;

        // Normalize stdev
        let nstdev = self.latency_stdev / max_latency_stdev();

        const SCORE_RTT_WEIGHT: f64 = 1.0;
        const SCORE_FAIL_WEIGHT: f64 = 3.0;
        const SCORE_STDEV_WEIGHT: f64 = 1.0;
        const SCORE_USER_WEIGHT: f64 = 5.0;

        // Score = (norm_lat * 1.0 + prop_err * 3.0 + stdev * 1.0 + user_weight * 5.0) / 10.0
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        // 3. The lower latency's stdev, the better
        // 4. The higher user's weight, the better
        let score = (nrtt * SCORE_RTT_WEIGHT
            + self.fail_rate * SCORE_FAIL_WEIGHT
            + nstdev * SCORE_STDEV_WEIGHT
            + self.user_weight as f64 * SCORE_USER_WEIGHT)
            / (SCORE_RTT_WEIGHT + SCORE_FAIL_WEIGHT + SCORE_STDEV_WEIGHT + SCORE_USER_WEIGHT);

        // Times 10000 converts to u32, for 0.0001 precision
        (score * 10000.0) as u32
    }

    pub fn push_score(&mut self, score: Score) -> u32 {
        self.latency_queue.push_back(score);

        // Only records recently MAX_LATENCY_QUEUE_SIZE probe data
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.recalculate_score()
    }

    fn recalculate_score(&mut self) -> u32 {
        if self.latency_queue.is_empty() {
            return self.score();
        }

        let mut vlat = Vec::with_capacity(self.latency_queue.len());
        let mut cerr = 0;
        for s in &self.latency_queue {
            match *s {
                Score::Errored => cerr += 1,
                Score::Latency(lat) => vlat.push(lat),
            }
        }

        // Error rate
        self.fail_rate = cerr as f64 / self.latency_queue.len() as f64;

        if !vlat.is_empty() {
            vlat.sort_unstable();

            // Find median of latency
            let mid = vlat.len() / 2;

            self.rtt = if vlat.len() % 2 == 0 {
                (vlat[mid] + vlat[mid - 1]) / 2
            } else {
                vlat[mid]
            };

            if vlat.len() > 1 {
                // STDEV
                let n = vlat.len() as f64;

                let mut total_lat = 0;
                for s in &vlat {
                    total_lat += *s;
                }
                self.latency_mean = total_lat as f64 / n;
                let mut acc_diff = 0.0;
                for s in &vlat {
                    let diff = *s as f64 - self.latency_mean;
                    acc_diff += diff * diff;
                }
                // Corrected Sample Standard Deviation
                self.latency_stdev = ((1.0 / (n - 1.0)) * acc_diff).sqrt();
            }
        }

        self.score()
    }
}
<commit_msg>fixed score calculation, user_weight should be divider<commit_after>//! Server latency statistic

use std::collections::VecDeque;

/// Interval between each check
pub const DEFAULT_CHECK_INTERVAL_SEC: u64 = 10;
/// Timeout of each check
pub const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 5; // A common connection timeout of 5 seconds.

const MAX_SERVER_RTT: u32 = DEFAULT_CHECK_TIMEOUT_SEC as u32 * 1000;
const MAX_LATENCY_QUEUE_SIZE: usize = 59; // Account for the last 10 minutes.

/// Statistic score
#[derive(Debug, Copy, Clone)]
pub enum Score {
    /// Unified latency
    Latency(u32),
    /// Request error
    Errored,
}

/// Statistic of a remote server
#[derive(Debug)]
pub struct ServerStat {
    /// Median of latency time (in millisec)
    ///
    /// Use median instead of average time,
    /// because probing result may have some really bad cases
    rtt: u32,
    /// Total_Fail / Total_Probe
    fail_rate: f64,
    /// Recently probe data
    latency_queue: VecDeque<Score>,
    /// Score's standard deviation
    latency_stdev: f64,
    /// Score's average
    latency_mean: f64,
    /// User's customized weight
    user_weight: f32,
}

fn max_latency_stdev() -> f64 {
    let mrtt = MAX_SERVER_RTT as f64;
    let avg = (0.0 + mrtt) / 2.0;
    let diff1 = (0.0 - avg) * (0.0 - avg);
    let diff2 = (mrtt - avg) * (mrtt - avg);
    // (1.0 / (2.0 - 1.0)) * (diff1 + diff2).sqrt()
    (diff1 + diff2).sqrt()
}

impl ServerStat {
    pub fn new(user_weight: f32) -> ServerStat {
        assert!((0.0..=1.0).contains(&user_weight));

        ServerStat {
            rtt: MAX_SERVER_RTT,
            fail_rate: 1.0,
            latency_queue: VecDeque::new(),
            latency_stdev: 0.0,
            latency_mean: 0.0,
            user_weight,
        }
    }

    fn score(&self) -> u32 {
        // Normalize rtt
        let nrtt = self.rtt as f64 / MAX_SERVER_RTT as f64;

        // Normalize stdev
        let nstdev = self.latency_stdev / max_latency_stdev();

        const SCORE_RTT_WEIGHT: f64 = 1.0;
        const SCORE_FAIL_WEIGHT: f64 = 3.0;
        const SCORE_STDEV_WEIGHT: f64 = 1.0;

        // Score = (norm_lat * 1.0 + prop_err * 3.0 + stdev * 1.0) / 5.0 / user_weight
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        // 3. The lower latency's stdev, the better
        // 4. The higher user's weight, the better
        let score = (nrtt * SCORE_RTT_WEIGHT + self.fail_rate * SCORE_FAIL_WEIGHT + nstdev * SCORE_STDEV_WEIGHT)
            / (SCORE_RTT_WEIGHT + SCORE_FAIL_WEIGHT + SCORE_STDEV_WEIGHT)
            / self.user_weight as f64;

        // Times 10000 converts to u32, for 0.0001 precision
        (score * 10000.0) as u32
    }

    pub fn push_score(&mut self, score: Score) -> u32 {
        self.latency_queue.push_back(score);

        // Only records recently MAX_LATENCY_QUEUE_SIZE probe data
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.recalculate_score()
    }

    fn recalculate_score(&mut self) -> u32 {
        if self.latency_queue.is_empty() {
            return self.score();
        }

        let mut vlat = Vec::with_capacity(self.latency_queue.len());
        let mut cerr = 0;
        for s in &self.latency_queue {
            match *s {
                Score::Errored => cerr += 1,
                Score::Latency(lat) => vlat.push(lat),
            }
        }

        // Error rate
        self.fail_rate = cerr as f64 / self.latency_queue.len() as f64;

        if !vlat.is_empty() {
            vlat.sort_unstable();

            // Find median of latency
            let mid = vlat.len() / 2;

            self.rtt = if vlat.len() % 2 == 0 {
                (vlat[mid] + vlat[mid - 1]) / 2
            } else {
                vlat[mid]
            };

            if vlat.len() > 1 {
                // STDEV
                let n = vlat.len() as f64;

                let mut total_lat = 0;
                for s in &vlat {
                    total_lat += *s;
                }
                self.latency_mean = total_lat as f64 / n;
                let mut acc_diff = 0.0;
                for s in &vlat {
                    let diff = *s as f64 - self.latency_mean;
                    acc_diff += diff * diff;
                }
                // Corrected Sample Standard Deviation
                self.latency_stdev = ((1.0 / (n - 1.0)) * acc_diff).sqrt();
            }
        }

        self.score()
    }
}
<|endoftext|>


<commit_before>mod mdb;
pub use self::mdb::{MDB, MDBAddr, CapFull, IntoFull, CapNearlyFull};

use common::*;
use cap::{Cap, UntypedFull, UntypedNearlyFull, CapReadObject};
use core::mem::{size_of, align_of};
use core::ops::{Index, IndexMut, Deref, DerefMut};
use core::slice::Iter;
use util::{RwLock, SharedReadGuard, SharedWriteGuard, MemoryObject};

pub type CPoolFull = CapFull<CPoolHalf, [MDB; 1]>;
pub type CPoolNearlyFull<'a> = CapNearlyFull<CPoolHalf, [Option<&'a mut MDB>; 1]>;
pub type CPool = [RwLock<Option<Cap>>; 256];

macro_rules! cpool_default {
    () => ({
        [RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None)]
    })
}

impl CPoolFull {
    pub unsafe fn bootstrap(mut untyped: UntypedFull) -> CPoolHalf {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let (start_paddr, _) = untyped.allocate(length, alignment);

        let mut cap = unsafe { CPoolHalf::new(start_paddr, 256) };

        unsafe {
            let obj = MemoryObject::<CPool>::new(cap.start_paddr);
            *(obj.as_mut().unwrap()) =
                cpool_default!();
        }

        cap.insert(untyped);

        let mut untyped_cap = cap.write(0);
        let mut untyped = match untyped_cap.as_mut().unwrap() {
            &mut Cap::Untyped(ref mut untyped) => untyped,
            _ => panic!(),
        };
        let cloned_cap = cap.clone();
        cap.insert(CPoolNearlyFull::new(cloned_cap, [ Some(untyped.mdb_mut(0)) ]));

        cap
    }

    pub fn retype<'a>(untyped: &'a mut UntypedFull) -> CPoolNearlyFull<'a> {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let (start_paddr, mdb) = untyped.allocate(length, alignment);

        let mut cap = unsafe { CPoolHalf::new(start_paddr, 256) };

        unsafe {
            let obj = MemoryObject::<CPool>::new(cap.start_paddr);
            *(obj.as_mut().unwrap()) =
                cpool_default!();
        }

        CPoolNearlyFull::new(cap, [ mdb ])
    }
}

#[derive(Debug, Clone)]
pub struct CPoolHalf {
    start_paddr: PAddr,
    size: usize,
}

impl CPoolHalf {
    pub unsafe fn new(start_paddr: PAddr, size: usize) -> Self {
        CPoolHalf {
            start_paddr: start_paddr,
            size: size
        }
    }

    pub fn start_paddr(&self) -> PAddr {
        self.start_paddr
    }

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

    pub fn insert<Half, M, U>(&mut self, mut cap: U)
        where U: IntoFull<Half, M>, Cap: From<CapFull<Half, M>> {
        let cpool = self.clone();
        for index in 0..self.size {
            let mut item = self.try_write(index);
            if item.is_some() {
                let mut item = item.unwrap();
                if item.is_none() {
                    *item = unsafe { Some(cap.into_full(cpool, index).into()) };
                    return;
                }
            }
        }
        assert!(false);
    }

    fn item_paddr(&self, index: usize) -> PAddr {
        self.start_paddr + size_of::<RwLock<Option<Cap>>>() * index
    }

    pub fn read<'a, 'b>(&'a self, index: usize) -> SharedReadGuard<'b, Option<Cap>> {
        let paddr = self.item_paddr(index);
        unsafe { SharedReadGuard::new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn try_read<'a, 'b>(&'a self, index: usize) -> Option<SharedReadGuard<'b, Option<Cap>>> {
        let paddr = self.item_paddr(index);
        unsafe { SharedReadGuard::try_new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn write<'a, 'b>(&'a mut self, index: usize) -> SharedWriteGuard<'b, Option<Cap>> {
        let paddr = self.item_paddr(index);
        unsafe { SharedWriteGuard::new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn try_write<'a, 'b>(&'a mut self, index: usize) -> Option<SharedWriteGuard<'b, Option<Cap>>> {
        let paddr = self.item_paddr(index);
        unsafe { SharedWriteGuard::try_new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    // pub fn traverse(&self, routes: &[u8]) -> SharedReadGuard<Option<Cap>> {
    //     if routes.len() == 0 {
    //         assert!(false);
    //     }
    //     let (first, rest) = routes.split_first().unwrap();
    //     let mut current_cap = self.read(*first);

    //     for path in rest {
    //         let new_current_cap = if let &Some(Cap::CPool(ref cpool_full)) = current_cap.deref() {
    //             Some(cpool_full.read(*path))
    //         } else {
    //             None
    //         };
    //         if new_current_cap.is_some() {
    //             current_cap = new_current_cap.unwrap();
    //         }
    //     }

    //     return current_cap;
    // }
}
<commit_msg>Make CPool access not overflow<commit_after>mod mdb;
pub use self::mdb::{MDB, MDBAddr, CapFull, IntoFull, CapNearlyFull};

use common::*;
use cap::{Cap, UntypedFull, UntypedNearlyFull, CapReadObject};
use core::mem::{size_of, align_of};
use core::ops::{Index, IndexMut, Deref, DerefMut};
use core::slice::Iter;
use util::{RwLock, SharedReadGuard, SharedWriteGuard, MemoryObject};

pub type CPoolFull = CapFull<CPoolHalf, [MDB; 1]>;
pub type CPoolNearlyFull<'a> = CapNearlyFull<CPoolHalf, [Option<&'a mut MDB>; 1]>;
pub type CPool = [RwLock<Option<Cap>>; 256];

macro_rules! cpool_default {
    () => ({
        [RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None),
         RwLock::new(None), RwLock::new(None), RwLock::new(None), RwLock::new(None)]
    })
}

impl CPoolFull {
    pub unsafe fn bootstrap(mut untyped: UntypedFull) -> CPoolHalf {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let (start_paddr, _) = untyped.allocate(length, alignment);

        let mut cap = unsafe { CPoolHalf::new(start_paddr, 256) };

        unsafe {
            let obj = MemoryObject::<CPool>::new(cap.start_paddr);
            *(obj.as_mut().unwrap()) =
                cpool_default!();
        }

        cap.insert(untyped);

        let mut untyped_cap = cap.write(0);
        let mut untyped = match untyped_cap.as_mut().unwrap() {
            &mut Cap::Untyped(ref mut untyped) => untyped,
            _ => panic!(),
        };
        let cloned_cap = cap.clone();
        cap.insert(CPoolNearlyFull::new(cloned_cap, [ Some(untyped.mdb_mut(0)) ]));

        cap
    }

    pub fn retype<'a>(untyped: &'a mut UntypedFull) -> CPoolNearlyFull<'a> {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let (start_paddr, mdb) = untyped.allocate(length, alignment);

        let mut cap = unsafe { CPoolHalf::new(start_paddr, 256) };

        unsafe {
            let obj = MemoryObject::<CPool>::new(cap.start_paddr);
            *(obj.as_mut().unwrap()) =
                cpool_default!();
        }

        CPoolNearlyFull::new(cap, [ mdb ])
    }
}

#[derive(Debug, Clone)]
pub struct CPoolHalf {
    start_paddr: PAddr,
    size: usize,
}

impl CPoolHalf {
    pub unsafe fn new(start_paddr: PAddr, size: usize) -> Self {
        CPoolHalf {
            start_paddr: start_paddr,
            size: size
        }
    }

    pub fn start_paddr(&self) -> PAddr {
        self.start_paddr
    }

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

    pub fn insert<Half, M, U>(&mut self, mut cap: U)
        where U: IntoFull<Half, M>, Cap: From<CapFull<Half, M>> {
        let cpool = self.clone();
        for index in 0..self.size {
            let mut item = self.try_write(index);
            if item.is_some() {
                let mut item = item.unwrap();
                if item.is_none() {
                    *item = unsafe { Some(cap.into_full(cpool, index).into()) };
                    return;
                }
            }
        }
        assert!(false);
    }

    fn item_paddr(&self, index: usize) -> PAddr {
        self.start_paddr + size_of::<RwLock<Option<Cap>>>() * index
    }

    pub fn read<'a, 'b>(&'a self, index: usize) -> SharedReadGuard<'b, Option<Cap>> {
        assert!(index < self.size);
        let paddr = self.item_paddr(index);
        unsafe { SharedReadGuard::new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn try_read<'a, 'b>(&'a self, index: usize) -> Option<SharedReadGuard<'b, Option<Cap>>> {
        assert!(index < self.size);
        let paddr = self.item_paddr(index);
        unsafe { SharedReadGuard::try_new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn write<'a, 'b>(&'a mut self, index: usize) -> SharedWriteGuard<'b, Option<Cap>> {
        assert!(index < self.size);
        let paddr = self.item_paddr(index);
        unsafe { SharedWriteGuard::new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    pub fn try_write<'a, 'b>(&'a mut self, index: usize) -> Option<SharedWriteGuard<'b, Option<Cap>>> {
        assert!(index < self.size);
        let paddr = self.item_paddr(index);
        unsafe { SharedWriteGuard::try_new(MemoryObject::<RwLock<Option<Cap>>>::new(paddr)) }
    }

    // pub fn traverse(&self, routes: &[u8]) -> SharedReadGuard<Option<Cap>> {
    //     if routes.len() == 0 {
    //         assert!(false);
    //     }
    //     let (first, rest) = routes.split_first().unwrap();
    //     let mut current_cap = self.read(*first);

    //     for path in rest {
    //         let new_current_cap = if let &Some(Cap::CPool(ref cpool_full)) = current_cap.deref() {
    //             Some(cpool_full.read(*path))
    //         } else {
    //             None
    //         };
    //         if new_current_cap.is_some() {
    //             current_cap = new_current_cap.unwrap();
    //         }
    //     }

    //     return current_cap;
    // }
}
<|endoftext|>