<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>workaround bug in epoll/mio<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{self, Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&self, error: Error) -> Result<()> {
        for (_, resolver) in self.serialized_frames.iter() {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                self.can_read = true;
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
                self.send_continue()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use libc::{c_int, c_char};

pub static MAX_CPUID_LEVEL: uint = 32u;
pub static MAX_EXT_CPUID_LEVEL: uint = 32u;
pub static MAX_INTELFN4_LEVEL: uint = 4u;
pub static MAX_INTELFN11_LEVEL: uint = 4u;

#[link(name = "cpuid")]
extern {
    pub fn cpuid_present() -> c_int;
    pub fn cpuid_lib_version() -> *const c_char;
    pub fn cpuid_error() -> *const c_char;
}
<commit_msg>Prototype for cpuid_get_raw_data.<commit_after>use libc::{c_int, c_char, uint32_t};

pub static MAX_CPUID_LEVEL: uint = 32u;
pub static MAX_EXT_CPUID_LEVEL: uint = 32u;
pub static MAX_INTELFN4_LEVEL: uint = 4u;
pub static MAX_INTELFN11_LEVEL: uint = 4u;

#[repr(C)]
pub struct cpu_raw_data_t {
    pub basic_cpuid: [uint32_t, ..MAX_CPUID_LEVEL * 4u],
    pub ext_cpuid: [uint32_t, ..MAX_EXT_CPUID_LEVEL * 4u],
    pub intel_fn4: [uint32_t, ..MAX_INTELFN4_LEVEL * 4u],
    pub intel_fn11: [uint32_t, ..MAX_INTELFN11_LEVEL * 4u],
}

#[link(name = "cpuid")]
extern {
    pub fn cpuid_present() -> c_int;
    pub fn cpuid_lib_version() -> *const c_char;
    pub fn cpuid_error() -> *const c_char;
    pub fn cpuid_get_raw_data(raw: *mut cpu_raw_data_t) -> c_int;
}
<|endoftext|>


<commit_before>use crate::client::AsyncClient;
pub use crate::index::ranged::trees::*;
use crate::ram::types::RandValue;
pub use cell_ref::NodeCellRef;
pub use cursor::*;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::{key_hash, PrimitiveArray, Value};
pub use external::page_schema;
use external::*;
use futures::future::BoxFuture;
use futures::FutureExt;
use insert::*;
use internal::*;
use itertools::Itertools;
use level::LEVEL_TREE_DEPTH;
use lightning::map::HashSet;
use merge::merge_into_tree_node;
pub use node::*;
use parking_lot::RwLock;
use search::*;
use std::any::Any;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed, Ordering::SeqCst};
use std::sync::Arc;

pub mod cell_ref;
mod cursor;
mod dump;
mod external;
mod insert;
mod internal;
pub mod level;
mod merge;
mod node;
// mod prune;
mod reconstruct;
mod search;
mod split;
pub mod storage;
pub mod verification;
#[macro_use]
pub mod marco;

const DEL_SET_CAP: usize = 16;

pub type DeletionSetInneer = HashSet<EntryKey>;
pub type DeletionSet = Arc<DeletionSetInneer>;

// Items can be added in real-time
// It is not supposed to hold a lot of items when it is actually feasible
// There will be a limit for maximum items in ths data structure, when the limit exceeds, higher ordering
// items with number of one page will be merged to next level
pub struct BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    root: RwLock<NodeCellRef>,
    root_versioning: NodeCellRef,
    head_page_id: Id,
    len: AtomicUsize,
    deleted: DeletionSet,
    marker: PhantomData<(KS, PS)>,
}

unsafe impl<KS, PS> Sync for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
}

unsafe impl<KS, PS> Send for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
}

impl Default for Ordering {
    fn default() -> Self {
        Ordering::Forward
    }
}

impl<KS, PS> BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new() -> BPlusTree<KS, PS> {
        trace!("Creating B+ Tree, with capacity {}", KS::slice_len());
        let mut tree = BPlusTree {
            root: RwLock::new(NodeCellRef::new(Node::<KS, PS>::new(NodeData::None))),
            root_versioning: NodeCellRef::new(Node::<KS, PS>::new(NodeData::None)),
            head_page_id: Id::unit_id(),
            len: AtomicUsize::new(0),
            deleted: Arc::new(DeletionSetInneer::with_capacity(DEL_SET_CAP)),
            marker: PhantomData,
        };
        let root_id = Self::new_page_id();
        let max_key = max_entry_key();
        trace!("New External L1");
        let root_inner = Node::<KS, PS>::new_external(root_id, max_key);
        trace!("B+ Tree created");
        *tree.root.write() = NodeCellRef::new(root_inner);
        tree.head_page_id = root_id;
        return tree;
    }

    pub async fn persist_root(&self, neb: &Arc<crate::client::AsyncClient>) {
        let root = self.get_root();
        root.persist(&self.deleted, &neb).await
    }

    pub async fn from_head_id(head_id: &Id, neb: &AsyncClient) -> Self {
        reconstruct::reconstruct_from_head_id(*head_id, neb).await
    }

    pub fn from_root(root: NodeCellRef, head_id: Id, len: usize) -> Self {
        BPlusTree {
            root: RwLock::new(root),
            root_versioning: NodeCellRef::default(),
            head_page_id: head_id,
            len: AtomicUsize::new(len),
            deleted: Arc::new(DeletionSetInneer::with_capacity(DEL_SET_CAP)),
            marker: PhantomData,
        }
    }

    pub fn get_root(&self) -> NodeCellRef {
        self.root.read().clone()
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> RTCursor<KS, PS> {
        search_node(&self.get_root(), key, ordering, &self.deleted)
    }

    pub fn insert(&self, key: &EntryKey) -> bool {
        match insert_to_tree_node(&self, &self.get_root(), &self.root_versioning, &key, 0) {
            Some(Some(split)) => {
                trace!("split root with pivot key {:?}", split.pivot);
                let new_node = split.new_right_node;
                let pivot = split.pivot;
                let mut new_in_root: Box<InNode<KS, PS>> = InNode::new(1, max_entry_key());
                let old_root = self.get_root().clone();
                new_in_root.keys.as_slice()[0] = pivot;
                new_in_root.ptrs.as_slice()[0] = old_root;
                new_in_root.ptrs.as_slice()[1] = new_node;
                *self.root.write() = NodeCellRef::new(Node::new(NodeData::Internal(new_in_root)));
            }
            Some(None) => {}
            None => return false,
        }
        self.len.fetch_add(1, Relaxed);
        self.deleted.remove(key);
        return true;
    }

    pub fn merge_with_keys_(&self, keys: Box<Vec<EntryKey>>) {
        let keys_len = keys.len();
        let root = self.get_root();
        let root_new_pages = merge_into_tree_node(self, &root, &self.root_versioning, keys, 0);
        if root_new_pages.len() > 0 {
            trace!("Merge have a root node split");
            debug_assert!(
                verification::is_node_serial(&write_node::<KS, PS>(&self.get_root())),
                "verification failed before merge root split"
            );
            debug_assert!(
                verification::are_keys_serial(
                    root_new_pages
                        .iter()
                        .map(|t| t.0.clone())
                        .collect_vec()
                        .as_slice()
                ),
                "verification failed before merge root split"
            );
            debug_assert!(
                root.ptr_eq(&self.get_root()),
                "Merge target tree should always have a persistent root unless merge split"
            );
            let _root_guard = write_node::<KS, PS>(&root);
            let new_root_len = root_new_pages.len();
            debug_assert!(
                new_root_len + 1 < KS::slice_len(),
                "Radical merge split, cannot handle this for now (need to split more than once)"
            );
            let mut new_in_root: Box<InNode<KS, PS>> = InNode::new(new_root_len, max_entry_key());
            new_in_root.ptrs.as_slice()[0] = root.clone();
            for (i, (key, node)) in root_new_pages.into_iter().enumerate() {
                new_in_root.keys.as_slice()[i] = key;
                new_in_root.ptrs.as_slice()[i + 1] = node;
            }
            // new_in_root.debug_check_integrity();
            *self.root.write() = NodeCellRef::new(Node::new(NodeData::Internal(new_in_root)));
            debug_assert!(
                verification::is_node_serial(&write_node::<KS, PS>(&self.get_root())),
                "verification failed after merge root split"
            );
        }
        self.len.fetch_add(keys_len, Relaxed);
    }

    pub fn flush_all(&self) {
        // unimplemented!()
    }

    pub fn len(&self) -> usize {
        self.len.load(Relaxed)
    }

    fn new_page_id() -> Id {
        // TODO: achieve locality
        Id::rand()
    }
}

pub trait LevelTree: Sync + Send {
    fn size(&self) -> usize;
    fn count(&self) -> usize;
    fn merge_to<'a>(&'a self, level: usize, upper_level: &'a dyn LevelTree)
        -> BoxFuture<'a, usize>;
    fn merge_with_keys(&self, keys: Box<Vec<EntryKey>>);
    fn insert_into(&self, key: &EntryKey) -> bool;
    fn seek_for(&self, key: &EntryKey, ordering: Ordering) -> Box<dyn Cursor>;
    fn mark_key_deleted(&self, key: &EntryKey) -> bool;
    fn dump(&self, f: &str);
    fn mid_key(&self) -> Option<EntryKey>;
    fn head_id(&self) -> Id;
    fn verify(&self, level: usize) -> bool;
    fn ideal_capacity(&self) -> usize {
        ideal_capacity_from_node_size(self.size())
    }
    fn oversized(&self) -> bool {
        self.count() > self.ideal_capacity()
    }
}

pub fn ideal_capacity_from_node_size(size: usize) -> usize {
    size.pow(LEVEL_TREE_DEPTH)
}

impl<KS, PS> LevelTree for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    fn size(&self) -> usize {
        KS::slice_len()
    }

    fn count(&self) -> usize {
        self.len()
    }

    fn merge_to<'a>(
        &'a self,
        level: usize,
        upper_level: &'a dyn LevelTree,
    ) -> BoxFuture<'a, usize> {
        async move { level::level_merge(level, self, upper_level).await }.boxed()
    }

    fn merge_with_keys(&self, keys: Box<Vec<EntryKey>>) {
        self.merge_with_keys_(keys)
    }

    fn insert_into(&self, key: &EntryKey) -> bool {
        self.insert(key)
    }

    fn seek_for(&self, key: &EntryKey, ordering: Ordering) -> Box<dyn Cursor> {
        box self.seek(key, ordering)
    }

    fn mark_key_deleted(&self, key: &EntryKey) -> bool {
        if let Some(seek_key) = self.seek(key, Ordering::Forward).current() {
            if seek_key == key {
                return self.deleted.insert(key);
            }
        }
        false
    }

    fn dump(&self, f: &str) {
        dump::dump_tree(self, f);
    }

    fn mid_key(&self) -> Option<EntryKey> {
        split::mid_key::<KS, PS>(&self.get_root())
    }

    fn head_id(&self) -> Id {
        self.head_page_id
    }

    fn verify(&self, level: usize) -> bool {
        verification::is_tree_in_order(self, level)
    }
}

pub struct DummyLevelTree;

impl LevelTree for DummyLevelTree {
    fn size(&self) -> usize {
        unreachable!()
    }

    fn count(&self) -> usize {
        unreachable!()
    }

    fn merge_to<'a>(
        &'a self,
        _level: usize,
        _upper_level: &'a dyn LevelTree,
    ) -> BoxFuture<'a, usize> {
        unreachable!()
    }

    fn merge_with_keys(&self, _keys: Box<Vec<EntryKey>>) {
        unreachable!()
    }

    fn insert_into(&self, _key: &EntryKey) -> bool {
        unreachable!()
    }

    fn seek_for(&self, _key: &EntryKey, _ordering: Ordering) -> Box<dyn Cursor> {
        unreachable!()
    }

    fn mark_key_deleted(&self, _key: &EntryKey) -> bool {
        unreachable!()
    }

    fn dump(&self, _f: &str) {
        unreachable!()
    }

    fn mid_key(&self) -> Option<EntryKey> {
        unreachable!()
    }

    fn head_id(&self) -> Id {
        unreachable!()
    }

    fn verify(&self, _level: usize) -> bool {
        unreachable!()
    }
}

impl_slice_ops!([EntryKey; 0], EntryKey, 0);
impl_slice_ops!([NodeCellRef; 0], NodeCellRef, 0);

#[cfg(test)]
pub mod test;
<commit_msg>need better merge<commit_after>use crate::client::AsyncClient;
pub use crate::index::ranged::trees::*;
use crate::ram::types::RandValue;
pub use cell_ref::NodeCellRef;
pub use cursor::*;
use dovahkiin::types::custom_types::id::Id;
use dovahkiin::types::{key_hash, PrimitiveArray, Value};
pub use external::page_schema;
use external::*;
use futures::future::BoxFuture;
use futures::FutureExt;
use insert::*;
use internal::*;
use itertools::Itertools;
use level::LEVEL_TREE_DEPTH;
use lightning::map::HashSet;
use merge::merge_into_tree_node;
pub use node::*;
use parking_lot::RwLock;
use search::*;
use std::any::Any;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed, Ordering::SeqCst};
use std::sync::Arc;

pub mod cell_ref;
mod cursor;
mod dump;
mod external;
mod insert;
mod internal;
pub mod level;
mod merge;
mod node;
// mod prune;
mod reconstruct;
mod search;
mod split;
pub mod storage;
pub mod verification;
#[macro_use]
pub mod marco;

const DEL_SET_CAP: usize = 16;

pub type DeletionSetInneer = HashSet<EntryKey>;
pub type DeletionSet = Arc<DeletionSetInneer>;

// Items can be added in real-time
// It is not supposed to hold a lot of items when it is actually feasible
// There will be a limit for maximum items in ths data structure, when the limit exceeds, higher ordering
// items with number of one page will be merged to next level
pub struct BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    root: RwLock<NodeCellRef>,
    root_versioning: NodeCellRef,
    head_page_id: Id,
    len: AtomicUsize,
    deleted: DeletionSet,
    marker: PhantomData<(KS, PS)>,
}

unsafe impl<KS, PS> Sync for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
}

unsafe impl<KS, PS> Send for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
}

impl Default for Ordering {
    fn default() -> Self {
        Ordering::Forward
    }
}

impl<KS, PS> BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    pub fn new() -> BPlusTree<KS, PS> {
        trace!("Creating B+ Tree, with capacity {}", KS::slice_len());
        let mut tree = BPlusTree {
            root: RwLock::new(NodeCellRef::new(Node::<KS, PS>::new(NodeData::None))),
            root_versioning: NodeCellRef::new(Node::<KS, PS>::new(NodeData::None)),
            head_page_id: Id::unit_id(),
            len: AtomicUsize::new(0),
            deleted: Arc::new(DeletionSetInneer::with_capacity(DEL_SET_CAP)),
            marker: PhantomData,
        };
        let root_id = Self::new_page_id();
        let max_key = max_entry_key();
        trace!("New External L1");
        let root_inner = Node::<KS, PS>::new_external(root_id, max_key);
        trace!("B+ Tree created");
        *tree.root.write() = NodeCellRef::new(root_inner);
        tree.head_page_id = root_id;
        return tree;
    }

    pub async fn persist_root(&self, neb: &Arc<crate::client::AsyncClient>) {
        let root = self.get_root();
        root.persist(&self.deleted, &neb).await
    }

    pub async fn from_head_id(head_id: &Id, neb: &AsyncClient) -> Self {
        reconstruct::reconstruct_from_head_id(*head_id, neb).await
    }

    pub fn from_root(root: NodeCellRef, head_id: Id, len: usize) -> Self {
        BPlusTree {
            root: RwLock::new(root),
            root_versioning: NodeCellRef::default(),
            head_page_id: head_id,
            len: AtomicUsize::new(len),
            deleted: Arc::new(DeletionSetInneer::with_capacity(DEL_SET_CAP)),
            marker: PhantomData,
        }
    }

    pub fn get_root(&self) -> NodeCellRef {
        self.root.read().clone()
    }

    pub fn seek(&self, key: &EntryKey, ordering: Ordering) -> RTCursor<KS, PS> {
        search_node(&self.get_root(), key, ordering, &self.deleted)
    }

    pub fn insert(&self, key: &EntryKey) -> bool {
        match insert_to_tree_node(&self, &self.get_root(), &self.root_versioning, &key, 0) {
            Some(Some(split)) => {
                trace!("split root with pivot key {:?}", split.pivot);
                let new_node = split.new_right_node;
                let pivot = split.pivot;
                let mut new_in_root: Box<InNode<KS, PS>> = InNode::new(1, max_entry_key());
                let old_root = self.get_root().clone();
                new_in_root.keys.as_slice()[0] = pivot;
                new_in_root.ptrs.as_slice()[0] = old_root;
                new_in_root.ptrs.as_slice()[1] = new_node;
                *self.root.write() = NodeCellRef::new(Node::new(NodeData::Internal(new_in_root)));
            }
            Some(None) => {}
            None => return false,
        }
        self.len.fetch_add(1, Relaxed);
        self.deleted.remove(key);
        return true;
    }

    pub fn merge_with_keys_(&self, keys: Box<Vec<EntryKey>>) {
        let keys_len = keys.len();
        let root = self.get_root();
        let root_new_pages = merge_into_tree_node(self, &root, &self.root_versioning, keys, 0);
        if root_new_pages.len() > 0 {
            trace!("Merge have a root node split");
            debug_assert!(
                verification::is_node_serial(&write_node::<KS, PS>(&self.get_root())),
                "verification failed before merge root split"
            );
            debug_assert!(
                verification::are_keys_serial(
                    root_new_pages
                        .iter()
                        .map(|t| t.0.clone())
                        .collect_vec()
                        .as_slice()
                ),
                "verification failed before merge root split"
            );
            debug_assert!(
                root.ptr_eq(&self.get_root()),
                "Merge target tree should always have a persistent root unless merge split"
            );
            let _root_guard = write_node::<KS, PS>(&root);
            let new_root_len = root_new_pages.len();
            debug_assert!(
                new_root_len + 1 < KS::slice_len(),
                "Radical merge split, cannot handle this for now (need to split more than once), {} to {}, num keys {}",
                new_root_len + 1,
                KS::slice_len(),
                keys_len
            );
            let mut new_in_root: Box<InNode<KS, PS>> = InNode::new(new_root_len, max_entry_key());
            new_in_root.ptrs.as_slice()[0] = root.clone();
            for (i, (key, node)) in root_new_pages.into_iter().enumerate() {
                new_in_root.keys.as_slice()[i] = key;
                new_in_root.ptrs.as_slice()[i + 1] = node;
            }
            // new_in_root.debug_check_integrity();
            *self.root.write() = NodeCellRef::new(Node::new(NodeData::Internal(new_in_root)));
            debug_assert!(
                verification::is_node_serial(&write_node::<KS, PS>(&self.get_root())),
                "verification failed after merge root split"
            );
        }
        self.len.fetch_add(keys_len, Relaxed);
    }

    pub fn flush_all(&self) {
        // unimplemented!()
    }

    pub fn len(&self) -> usize {
        self.len.load(Relaxed)
    }

    fn new_page_id() -> Id {
        // TODO: achieve locality
        Id::rand()
    }
}

pub trait LevelTree: Sync + Send {
    fn size(&self) -> usize;
    fn count(&self) -> usize;
    fn merge_to<'a>(&'a self, level: usize, upper_level: &'a dyn LevelTree)
        -> BoxFuture<'a, usize>;
    fn merge_with_keys(&self, keys: Box<Vec<EntryKey>>);
    fn insert_into(&self, key: &EntryKey) -> bool;
    fn seek_for(&self, key: &EntryKey, ordering: Ordering) -> Box<dyn Cursor>;
    fn mark_key_deleted(&self, key: &EntryKey) -> bool;
    fn dump(&self, f: &str);
    fn mid_key(&self) -> Option<EntryKey>;
    fn head_id(&self) -> Id;
    fn verify(&self, level: usize) -> bool;
    fn ideal_capacity(&self) -> usize {
        ideal_capacity_from_node_size(self.size())
    }
    fn oversized(&self) -> bool {
        self.count() > self.ideal_capacity()
    }
}

pub fn ideal_capacity_from_node_size(size: usize) -> usize {
    size.pow(LEVEL_TREE_DEPTH)
}

impl<KS, PS> LevelTree for BPlusTree<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    fn size(&self) -> usize {
        KS::slice_len()
    }

    fn count(&self) -> usize {
        self.len()
    }

    fn merge_to<'a>(
        &'a self,
        level: usize,
        upper_level: &'a dyn LevelTree,
    ) -> BoxFuture<'a, usize> {
        async move { level::level_merge(level, self, upper_level).await }.boxed()
    }

    fn merge_with_keys(&self, keys: Box<Vec<EntryKey>>) {
        self.merge_with_keys_(keys)
    }

    fn insert_into(&self, key: &EntryKey) -> bool {
        self.insert(key)
    }

    fn seek_for(&self, key: &EntryKey, ordering: Ordering) -> Box<dyn Cursor> {
        box self.seek(key, ordering)
    }

    fn mark_key_deleted(&self, key: &EntryKey) -> bool {
        if let Some(seek_key) = self.seek(key, Ordering::Forward).current() {
            if seek_key == key {
                return self.deleted.insert(key);
            }
        }
        false
    }

    fn dump(&self, f: &str) {
        dump::dump_tree(self, f);
    }

    fn mid_key(&self) -> Option<EntryKey> {
        split::mid_key::<KS, PS>(&self.get_root())
    }

    fn head_id(&self) -> Id {
        self.head_page_id
    }

    fn verify(&self, level: usize) -> bool {
        verification::is_tree_in_order(self, level)
    }
}

pub struct DummyLevelTree;

impl LevelTree for DummyLevelTree {
    fn size(&self) -> usize {
        unreachable!()
    }

    fn count(&self) -> usize {
        unreachable!()
    }

    fn merge_to<'a>(
        &'a self,
        _level: usize,
        _upper_level: &'a dyn LevelTree,
    ) -> BoxFuture<'a, usize> {
        unreachable!()
    }

    fn merge_with_keys(&self, _keys: Box<Vec<EntryKey>>) {
        unreachable!()
    }

    fn insert_into(&self, _key: &EntryKey) -> bool {
        unreachable!()
    }

    fn seek_for(&self, _key: &EntryKey, _ordering: Ordering) -> Box<dyn Cursor> {
        unreachable!()
    }

    fn mark_key_deleted(&self, _key: &EntryKey) -> bool {
        unreachable!()
    }

    fn dump(&self, _f: &str) {
        unreachable!()
    }

    fn mid_key(&self) -> Option<EntryKey> {
        unreachable!()
    }

    fn head_id(&self) -> Id {
        unreachable!()
    }

    fn verify(&self, _level: usize) -> bool {
        unreachable!()
    }
}

impl_slice_ops!([EntryKey; 0], EntryKey, 0);
impl_slice_ops!([NodeCellRef; 0], NodeCellRef, 0);

#[cfg(test)]
pub mod test;
<|endoftext|>


<commit_before>//! Module for `Rabbits and Recurrence Relations`

/// This function calculates recurrence relation, or value of element at `n` position
/// in Fibonaccie's sequence
///
/// # Examples
/// ```
/// use rosalind::fib::*;
///
/// assert_eq!(recurrence_relation(5, 3), 19);
/// ```
#[allow(unused_variables)]
pub fn recurrence_relation(n: u8, k: u8) -> u16 {
  let mut pair: (u16, u16) = (1, 1);
  for i in 0..n - 1 {
    pair = (pair.1, pair.0 * k as u16 + pair.1)
  }
  return pair.0;
}

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

  #[test]
  fn it_should_return_recurrence_relation() {
    assert_eq!(recurrence_relation(5, 3), 19);
  }
}
<commit_msg>Overflow fix u16 -> u64<commit_after>//! Module for `Rabbits and Recurrence Relations`

/// This function calculates recurrence relation, or value of element at `n` position
/// in Fibonaccie's sequence
///
/// # Examples
/// ```
/// use rosalind::fib::*;
///
/// assert_eq!(recurrence_relation(5, 3), 19);
/// ```
#[allow(unused_variables)]
pub fn recurrence_relation(n: u8, k: u8) -> u64 {
  let mut pair: (u64, u64) = (1, 1);
  for i in 0..n - 1 {
    pair = (pair.1, pair.0 * k as u64 + pair.1)
  }
  return pair.0;
}

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

  #[test]
  fn it_should_return_recurrence_relation() {
    assert_eq!(recurrence_relation(5, 3), 19);
  }
}
<|endoftext|>


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

use rand::{thread_rng, Rng};
use test::Bencher;

fn random(n: usize) -> BTreeSet<usize> {
    let mut rng = thread_rng();
    let mut set = BTreeSet::new();
    while set.len() < n {
        set.insert(rng.gen());
    }
    assert_eq!(set.len(), n);
    set
}

fn neg(n: usize) -> BTreeSet<i32> {
    let set: BTreeSet<i32> = (-(n as i32)..=-1).collect();
    assert_eq!(set.len(), n);
    set
}

fn pos(n: usize) -> BTreeSet<i32> {
    let set: BTreeSet<i32> = (1..=(n as i32)).collect();
    assert_eq!(set.len(), n);
    set
}

fn stagger(n1: usize, factor: usize) -> [BTreeSet<u32>; 2] {
    let n2 = n1 * factor;
    let mut sets = [BTreeSet::new(), BTreeSet::new()];
    for i in 0..(n1 + n2) {
        let b = i % (factor + 1) != 0;
        sets[b as usize].insert(i as u32);
    }
    assert_eq!(sets[0].len(), n1);
    assert_eq!(sets[1].len(), n2);
    sets
}

macro_rules! set_bench {
    ($name: ident, $set_func: ident, $result_func: ident, $sets: expr) => {
        #[bench]
        pub fn $name(b: &mut Bencher) {
            // setup
            let sets = $sets;

            // measure
            b.iter(|| sets[0].$set_func(&sets[1]).$result_func())
        }
    };
}

const BUILD_SET_SIZE: usize = 100;

#[bench]
pub fn build_and_clear(b: &mut Bencher) {
    b.iter(|| pos(BUILD_SET_SIZE).clear())
}

#[bench]
pub fn build_and_drop(b: &mut Bencher) {
    b.iter(|| pos(BUILD_SET_SIZE))
}

#[bench]
pub fn build_and_into_iter(b: &mut Bencher) {
    b.iter(|| pos(BUILD_SET_SIZE).into_iter().count())
}

#[bench]
pub fn build_and_pop_all(b: &mut Bencher) {
    b.iter(|| {
        let mut s = pos(BUILD_SET_SIZE);
        while s.pop_first().is_some() {}
        s
    });
}

#[bench]
pub fn build_and_remove_all(b: &mut Bencher) {
    b.iter(|| {
        let mut s = pos(BUILD_SET_SIZE);
        while let Some(elt) = s.iter().copied().next() {
            s.remove(&elt);
        }
        s
    });
}

set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}
set_bench! {intersection_100_pos_vs_10k_neg, intersection, count, [pos(100), neg(10_000)]}
set_bench! {intersection_10k_neg_vs_100_pos, intersection, count, [neg(10_000), pos(100)]}
set_bench! {intersection_10k_neg_vs_10k_pos, intersection, count, [neg(10_000), pos(10_000)]}
set_bench! {intersection_10k_pos_vs_100_neg, intersection, count, [pos(10_000), neg(100)]}
set_bench! {intersection_10k_pos_vs_10k_neg, intersection, count, [pos(10_000), neg(10_000)]}
set_bench! {intersection_random_100_vs_100, intersection, count, [random(100), random(100)]}
set_bench! {intersection_random_100_vs_10k, intersection, count, [random(100), random(10_000)]}
set_bench! {intersection_random_10k_vs_100, intersection, count, [random(10_000), random(100)]}
set_bench! {intersection_random_10k_vs_10k, intersection, count, [random(10_000), random(10_000)]}
set_bench! {intersection_staggered_100_vs_100, intersection, count, stagger(100, 1)}
set_bench! {intersection_staggered_10k_vs_10k, intersection, count, stagger(10_000, 1)}
set_bench! {intersection_staggered_100_vs_10k, intersection, count, stagger(100, 100)}
set_bench! {difference_random_100_vs_100, difference, count, [random(100), random(100)]}
set_bench! {difference_random_100_vs_10k, difference, count, [random(100), random(10_000)]}
set_bench! {difference_random_10k_vs_100, difference, count, [random(10_000), random(100)]}
set_bench! {difference_random_10k_vs_10k, difference, count, [random(10_000), random(10_000)]}
set_bench! {difference_staggered_100_vs_100, difference, count, stagger(100, 1)}
set_bench! {difference_staggered_10k_vs_10k, difference, count, stagger(10_000, 1)}
set_bench! {difference_staggered_100_vs_10k, difference, count, stagger(100, 100)}
set_bench! {is_subset_100_vs_100, is_subset, clone, [pos(100), pos(100)]}
set_bench! {is_subset_100_vs_10k, is_subset, clone, [pos(100), pos(10_000)]}
set_bench! {is_subset_10k_vs_100, is_subset, clone, [pos(10_000), pos(100)]}
set_bench! {is_subset_10k_vs_10k, is_subset, clone, [pos(10_000), pos(10_000)]}
<commit_msg>Refine and extend benchmarks of mutable BTreeSet methods<commit_after>use std::collections::BTreeSet;

use rand::{thread_rng, Rng};
use test::Bencher;

fn random(n: usize) -> BTreeSet<usize> {
    let mut rng = thread_rng();
    let mut set = BTreeSet::new();
    while set.len() < n {
        set.insert(rng.gen());
    }
    assert_eq!(set.len(), n);
    set
}

fn neg(n: usize) -> BTreeSet<i32> {
    let set: BTreeSet<i32> = (-(n as i32)..=-1).collect();
    assert_eq!(set.len(), n);
    set
}

fn pos(n: usize) -> BTreeSet<i32> {
    let set: BTreeSet<i32> = (1..=(n as i32)).collect();
    assert_eq!(set.len(), n);
    set
}

fn stagger(n1: usize, factor: usize) -> [BTreeSet<u32>; 2] {
    let n2 = n1 * factor;
    let mut sets = [BTreeSet::new(), BTreeSet::new()];
    for i in 0..(n1 + n2) {
        let b = i % (factor + 1) != 0;
        sets[b as usize].insert(i as u32);
    }
    assert_eq!(sets[0].len(), n1);
    assert_eq!(sets[1].len(), n2);
    sets
}

macro_rules! set_bench {
    ($name: ident, $set_func: ident, $result_func: ident, $sets: expr) => {
        #[bench]
        pub fn $name(b: &mut Bencher) {
            // setup
            let sets = $sets;

            // measure
            b.iter(|| sets[0].$set_func(&sets[1]).$result_func())
        }
    };
}

#[bench]
pub fn clone_100(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| src.clone())
}

#[bench]
pub fn clone_100_and_clear(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| src.clone().clear())
}

#[bench]
pub fn clone_100_and_into_iter(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| src.clone().into_iter().count())
}

#[bench]
pub fn clone_100_and_pop_all(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| {
        let mut set = src.clone();
        while set.pop_first().is_some() {}
        set
    });
}

#[bench]
pub fn clone_100_and_remove_all(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| {
        let mut set = src.clone();
        while let Some(elt) = set.iter().copied().next() {
            set.remove(&elt);
        }
        set
    });
}

#[bench]
pub fn clone_100_and_remove_half(b: &mut Bencher) {
    let src = pos(100);
    b.iter(|| {
        let mut set = src.clone();
        for i in (2..=100 as i32).step_by(2) {
            set.remove(&i);
        }
        assert_eq!(set.len(), 100 / 2);
        set
    })
}

#[bench]
pub fn clone_10k(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| src.clone())
}

#[bench]
pub fn clone_10k_and_clear(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| src.clone().clear())
}

#[bench]
pub fn clone_10k_and_into_iter(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| src.clone().into_iter().count())
}

#[bench]
pub fn clone_10k_and_pop_all(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| {
        let mut set = src.clone();
        while set.pop_first().is_some() {}
        set
    });
}

#[bench]
pub fn clone_10k_and_remove_all(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| {
        let mut set = src.clone();
        while let Some(elt) = set.iter().copied().next() {
            set.remove(&elt);
        }
        set
    });
}

#[bench]
pub fn clone_10k_and_remove_half(b: &mut Bencher) {
    let src = pos(10_000);
    b.iter(|| {
        let mut set = src.clone();
        for i in (2..=10_000 as i32).step_by(2) {
            set.remove(&i);
        }
        assert_eq!(set.len(), 10_000 / 2);
        set
    })
}

set_bench! {intersection_100_neg_vs_100_pos, intersection, count, [neg(100), pos(100)]}
set_bench! {intersection_100_neg_vs_10k_pos, intersection, count, [neg(100), pos(10_000)]}
set_bench! {intersection_100_pos_vs_100_neg, intersection, count, [pos(100), neg(100)]}
set_bench! {intersection_100_pos_vs_10k_neg, intersection, count, [pos(100), neg(10_000)]}
set_bench! {intersection_10k_neg_vs_100_pos, intersection, count, [neg(10_000), pos(100)]}
set_bench! {intersection_10k_neg_vs_10k_pos, intersection, count, [neg(10_000), pos(10_000)]}
set_bench! {intersection_10k_pos_vs_100_neg, intersection, count, [pos(10_000), neg(100)]}
set_bench! {intersection_10k_pos_vs_10k_neg, intersection, count, [pos(10_000), neg(10_000)]}
set_bench! {intersection_random_100_vs_100, intersection, count, [random(100), random(100)]}
set_bench! {intersection_random_100_vs_10k, intersection, count, [random(100), random(10_000)]}
set_bench! {intersection_random_10k_vs_100, intersection, count, [random(10_000), random(100)]}
set_bench! {intersection_random_10k_vs_10k, intersection, count, [random(10_000), random(10_000)]}
set_bench! {intersection_staggered_100_vs_100, intersection, count, stagger(100, 1)}
set_bench! {intersection_staggered_10k_vs_10k, intersection, count, stagger(10_000, 1)}
set_bench! {intersection_staggered_100_vs_10k, intersection, count, stagger(100, 100)}
set_bench! {difference_random_100_vs_100, difference, count, [random(100), random(100)]}
set_bench! {difference_random_100_vs_10k, difference, count, [random(100), random(10_000)]}
set_bench! {difference_random_10k_vs_100, difference, count, [random(10_000), random(100)]}
set_bench! {difference_random_10k_vs_10k, difference, count, [random(10_000), random(10_000)]}
set_bench! {difference_staggered_100_vs_100, difference, count, stagger(100, 1)}
set_bench! {difference_staggered_10k_vs_10k, difference, count, stagger(10_000, 1)}
set_bench! {difference_staggered_100_vs_10k, difference, count, stagger(100, 100)}
set_bench! {is_subset_100_vs_100, is_subset, clone, [pos(100), pos(100)]}
set_bench! {is_subset_100_vs_10k, is_subset, clone, [pos(100), pos(10_000)]}
set_bench! {is_subset_10k_vs_100, is_subset, clone, [pos(10_000), pos(100)]}
set_bench! {is_subset_10k_vs_10k, is_subset, clone, [pos(10_000), pos(10_000)]}
<|endoftext|>


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

use std::mem;

/// `BitSlice` provides helper methods for treating a `[usize]`
/// as a bitvector.
pub trait BitSlice {
    fn clear_bit(&mut self, idx: usize) -> bool;
    fn set_bit(&mut self, idx: usize) -> bool;
    fn get_bit(&self, idx: usize) -> bool;
}

impl BitSlice for [usize] {
    /// Clears bit at `idx` to 0; returns true iff this changed `self.`
    fn clear_bit(&mut self, idx: usize) -> bool {
        let words = self;
        debug!("clear_bit: words={} idx={}",
               bits_to_string(words, words.len() * mem::size_of::<usize>()), bit_str(idx));
        let BitLookup { word, bit_in_word, bit_mask } = bit_lookup(idx);
        debug!("word={} bit_in_word={} bit_mask={}", word, bit_in_word, bit_mask);
        let oldv = words[word];
        let newv = oldv & !bit_mask;
        words[word] = newv;
        oldv != newv
    }

    /// Sets bit at `idx` to 1; returns true iff this changed `self.`
    fn set_bit(&mut self, idx: usize) -> bool {
        let words = self;
        debug!("set_bit: words={} idx={}",
               bits_to_string(words, words.len() * mem::size_of::<usize>()), bit_str(idx));
        let BitLookup { word, bit_in_word, bit_mask } = bit_lookup(idx);
        debug!("word={} bit_in_word={} bit_mask={}", word, bit_in_word, bit_mask);
        let oldv = words[word];
        let newv = oldv | bit_mask;
        words[word] = newv;
        oldv != newv
    }

    /// Extracts value of bit at `idx` in `self`.
    fn get_bit(&self, idx: usize) -> bool {
        let words = self;
        let BitLookup { word, bit_mask, .. } = bit_lookup(idx);
        (words[word] & bit_mask) != 0
    }
}

struct BitLookup {
    /// An index of the word holding the bit in original `[usize]` of query.
    word: usize,
    /// Index of the particular bit within the word holding the bit.
    bit_in_word: usize,
    /// Word with single 1-bit set corresponding to where the bit is located.
    bit_mask: usize,
}

#[inline]
fn bit_lookup(bit: usize) -> BitLookup {
    let usize_bits = mem::size_of::<usize>() * 8;
    let word = bit / usize_bits;
    let bit_in_word = bit % usize_bits;
    let bit_mask = 1 << bit_in_word;
    BitLookup { word: word, bit_in_word: bit_in_word, bit_mask: bit_mask }
}

fn bit_str(bit: usize) -> String {
    let byte = bit >> 3;
    let lobits = 1 << (bit & 0b111);
    format!("[{}:{}-{:02x}]", bit, byte, lobits)
}

pub fn bits_to_string(words: &[usize], bits: usize) -> String {
    let mut result = String::new();
    let mut sep = '[';

    // Note: this is a little endian printout of bytes.

    // i tracks how many bits we have printed so far.
    let mut i = 0;
    for &word in words.iter() {
        let mut v = word;
        loop { // for each byte in `v`:
            let remain = bits - i;
            // If less than a byte remains, then mask just that many bits.
            let mask = if remain <= 8 { (1 << remain) - 1 } else { 0xFF };
            assert!(mask <= 0xFF);
            let byte = v & mask;

            result.push(sep);
            result.push_str(&format!("{:02x}", byte));

            if remain <= 8 { break; }
            v >>= 8;
            i += 8;
            sep = '-';
        }
    }
    result.push(']');
    return result
}
<commit_msg>Refactor `bitslice`: distinguish `usize` for indexing vs word type being indexed.<commit_after>// Copyright 2012-2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;

pub type Word = usize;

/// `BitSlice` provides helper methods for treating a `[Word]`
/// as a bitvector.
pub trait BitSlice {
    fn clear_bit(&mut self, idx: usize) -> bool;
    fn set_bit(&mut self, idx: usize) -> bool;
    fn get_bit(&self, idx: usize) -> bool;
}

impl BitSlice for [Word] {
    /// Clears bit at `idx` to 0; returns true iff this changed `self.`
    fn clear_bit(&mut self, idx: usize) -> bool {
        let words = self;
        debug!("clear_bit: words={} idx={}",
               bits_to_string(words, words.len() * mem::size_of::<Word>()), bit_str(idx));
        let BitLookup { word, bit_in_word, bit_mask } = bit_lookup(idx);
        debug!("word={} bit_in_word={} bit_mask={}", word, bit_in_word, bit_mask);
        let oldv = words[word];
        let newv = oldv & !bit_mask;
        words[word] = newv;
        oldv != newv
    }

    /// Sets bit at `idx` to 1; returns true iff this changed `self.`
    fn set_bit(&mut self, idx: usize) -> bool {
        let words = self;
        debug!("set_bit: words={} idx={}",
               bits_to_string(words, words.len() * mem::size_of::<Word>()), bit_str(idx));
        let BitLookup { word, bit_in_word, bit_mask } = bit_lookup(idx);
        debug!("word={} bit_in_word={} bit_mask={}", word, bit_in_word, bit_mask);
        let oldv = words[word];
        let newv = oldv | bit_mask;
        words[word] = newv;
        oldv != newv
    }

    /// Extracts value of bit at `idx` in `self`.
    fn get_bit(&self, idx: usize) -> bool {
        let words = self;
        let BitLookup { word, bit_mask, .. } = bit_lookup(idx);
        (words[word] & bit_mask) != 0
    }
}

struct BitLookup {
    /// An index of the word holding the bit in original `[Word]` of query.
    word: usize,
    /// Index of the particular bit within the word holding the bit.
    bit_in_word: usize,
    /// Word with single 1-bit set corresponding to where the bit is located.
    bit_mask: Word,
}

#[inline]
fn bit_lookup(bit: usize) -> BitLookup {
    let word_bits = mem::size_of::<Word>() * 8;
    let word = bit / word_bits;
    let bit_in_word = bit % word_bits;
    let bit_mask = 1 << bit_in_word;
    BitLookup { word: word, bit_in_word: bit_in_word, bit_mask: bit_mask }
}

fn bit_str(bit: Word) -> String {
    let byte = bit >> 3;
    let lobits = 1 << (bit & 0b111);
    format!("[{}:{}-{:02x}]", bit, byte, lobits)
}

pub fn bits_to_string(words: &[Word], bits: usize) -> String {
    let mut result = String::new();
    let mut sep = '[';

    // Note: this is a little endian printout of bytes.

    // i tracks how many bits we have printed so far.
    let mut i = 0;
    for &word in words.iter() {
        let mut v = word;
        loop { // for each byte in `v`:
            let remain = bits - i;
            // If less than a byte remains, then mask just that many bits.
            let mask = if remain <= 8 { (1 << remain) - 1 } else { 0xFF };
            assert!(mask <= 0xFF);
            let byte = v & mask;

            result.push(sep);
            result.push_str(&format!("{:02x}", byte));

            if remain <= 8 { break; }
            v >>= 8;
            i += 8;
            sep = '-';
        }
    }
    result.push(']');
    return result
}
<|endoftext|>


<commit_before>// Copyright 2014 Birunthan Mohanathas
//
// Licensed under the MIT license <http://opensource.org/licenses/MIT>. This
// file may not be copied, modified, or distributed except according to those
// terms.

pub struct Measure {
    name: String,
    value: f32,
}

impl Measure {
    pub fn new(name: &str) -> Measure {
        Measure {
            name: name.to_string(),
            value: 0.0f32,
        }
    }

    pub fn name(&self) -> &str {
        self.name.as_slice()
    }

    pub fn value(&self) -> f32 {
        self.value
    }

    pub fn update_value(&mut self) {
        self.value += 1.0f32
    }
}
<commit_msg>Store string value as well in Measure<commit_after>// Copyright 2014 Birunthan Mohanathas
//
// Licensed under the MIT license <http://opensource.org/licenses/MIT>. This
// file may not be copied, modified, or distributed except according to those
// terms.

pub struct Measure {
    name: String,
    number_value: f32,
    string_value: String,
}

impl Measure {
    pub fn new(name: &str) -> Measure {
        Measure {
            name: name.to_string(),
            number_value: 0.0f32,
            string_value: "".to_string(),
        }
    }

    pub fn name(&self) -> &str {
        self.name.as_slice()
    }

    pub fn number_value(&self) -> f32 {
        self.number_value
    }

    pub fn string_value(&self) -> &str {
        self.string_value.as_slice()
    }
}
<|endoftext|>


<commit_before>extern crate collections;
use std::io::{File, Open, Write, TempDir};
use std::os::args;
use std::from_str::from_str;
use collections::priority_queue::PriorityQueue;
#[cfg(test)]
use std::io::{ReadWrite, SeekSet};

/*
 * an entry is the number that we read and the file it came from, so after we took
 * out one Entry from the priority queue, we need to know where to take the next number
 * from.
 */
struct Entry {
	value: u64,
	source: File,
}

/* The priority queue depends on the ordering relation and always returns the biggest */
impl Ord for Entry {
	// lesser than implementation
	fn lt(&self, other: &Entry) -> bool {
		// we invert: for us the smallest will have the highest priority
		!(self.value < other.value)
	}
}

/* ordering needs equality for some reason, so there */
impl Eq for Entry {
	fn eq(&self, other: &Entry) -> bool {
		self.value == other.value
	}
}

/* A struct to hold our queue */
struct PriorityFile {
	queue: PriorityQueue<Entry>,
}

/*
 * 'Class' that takes a list of files and does a k-way merge of the smallest values
 * from the files
 */
impl PriorityFile {
	pub fn new(files: Vec<File>) -> PriorityFile {
		let mut queue: PriorityQueue<Entry> = PriorityQueue::new();
		// now that we got the queue, populate it with initial values

		// move_iter is required to get the values themselves, not borrowed pointers to them
		for mut f in files.move_iter() {
			// read from file, append to queue
			let number = match f.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			queue.push(Entry {value: number, source: f});
		};
		PriorityFile {queue: queue}
	}

	/* Returns the next smallest number from the list of files or None if exhausted */
	pub fn next(&mut self) -> Option<u64> {
		match self.queue.maybe_pop() {
			Some(mut res) => {
				let v = res.value;
				match res.source.read_le_u64() {
					Ok(num) => self.queue.push(Entry {value: num, source: res.source}),
					Err(_) => (),
				};
				Some(v)
			},
			None => None,
		}
	}
}

fn externalSort(mut fdInput: File, size: u64, mut fdOutput: File, memSize: u64) {
	let runs = size / memSize;
	let items_per_run = (memSize / 8) as uint;

	println!("There will be {} runs with {} elements each.", runs, items_per_run);

	/* initialize a Vector. Rust Vectors grow but this one will stay fixed */
	let mut run: Vec<u64> = Vec::with_capacity(items_per_run);
	// preallocate it with zeroes
	for _ in range(0, items_per_run) {
		run.push(0_u64);
	}

	// let's create a temporary directory to store the sorted chunks
	// neat: TempDir deletes the directory and its contents when it goes out of
	// scope
	let overflow_dir = match TempDir::new("externalsort") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let overflow_path = overflow_dir.path();
	println!("temp dir path {}", overflow_path.display());

	// start the sorting runs
	for n in range(0, runs) {
		// iterate over our 'run' buffer and read from the input file
		// I suppose the values are little endian, which considering we are most
		// likely on x86, is a reasonable assumption
		for element in run.mut_iter() {
			let number = match fdInput.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			*element = number;
		}
		// O(n log n) sort from the stdlib, hopefully more or less equivalent to
		// C++ std::sort
		run.sort();

		// now write out the results into numbered files
		let file_path = overflow_path.join(n.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};

		for &element in run.iter() {
			match file_file.write_le_u64(element) {
				Ok (_) => (),
				Err(e) => fail!("writing overflow failed: {}", e),
			};
		};
	};

	/* sorting done, now merging */

	// collect all file handles. Hopefully less than 4096
	let mut files = Vec::with_capacity(runs as uint);
	for n in range(0, runs) {
		let file_path = overflow_path.join(n.to_str());
		let fd = match File::open(&file_path) {
			Ok(f) => f,
			Err(e) => fail!("overflow file opening error {}", e),
		};
		files.push(fd);
	}

	// we got a nice class here that spits out sorted numbers until exhausted
	let mut pf = PriorityFile::new(files);
	loop {
		match pf.next() {
			// if we got a number, write it to the output
			Some (number) => match fdOutput.write_le_u64(number) {
				Ok(_) => (),
				Err(e) => fail!("writing output file failed: {}", e),
			},
			None => break,
		};
	};
}

fn main() {
	if args().len() < 4 {
		fail!("Argument error: <inputFile> <outputFile> <memoryBufferInMB>");
	}
	let input_file_path = &Path::new(args()[1]);
	let size = match input_file_path.stat() {
		Ok(stat) => stat.size,
		Err(e) => fail!("Couldn't read {}", e),
	};
	let fin = match File::open(input_file_path) {
		Ok(f) => f,
		Err(e) => fail!("input file error: {}", e),
	};
	let fout = match File::open_mode(&Path::new(args()[2]), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	let buffer_size:u64 = match from_str(args()[3]) {
		Some(num) => num,
		None => fail!("Not numeric input"),
	};
	println!("Input file size {}", size);

	externalSort(fin, size, fout, buffer_size * 1024 * 1024);
}

#[test]
fn generate_5gb_and_sort() {
	let mut random = match File::open(&Path::new("/dev/urandom")) {
		Ok(f) => f,
		Err(e) => fail!("Error opening urandom: {}", e),
	};

	let test_dir = match TempDir::new("externalsort-test") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let test_path = test_dir.path();

	let mut unordered_nums = match File::open_mode(&test_path.join("input"), Open, ReadWrite) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut buffer = [0, .. 1024*1024];
	let mut left = 1 * 1024;
	while left > 0 {
		random.read(buffer);
		unordered_nums.write(buffer);
		left -= 1;
	}

	unordered_nums.seek(0, SeekSet);
	let ordered_nums = match File::open_mode(&test_path.join("output"), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	externalSort(unordered_nums, 1024 * 1024 * 1024 * 1, ordered_nums, 1024 * 1024 * 200);
	let mut ordered_nums = match File::open(&test_path.join("output")) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut last = 0_u64;
	loop {
		let number = match ordered_nums.read_le_u64() {
			Ok(num) => num,
			Err(e) => fail!("failed to read u64 from file: {}", e),
		};
		assert!(number >= last);
		last = number;
	}
}
<commit_msg>Collecting left-over elements too<commit_after>extern crate collections;
use std::io::{File, Open, Write, TempDir};
use std::os::args;
use std::from_str::from_str;
use collections::priority_queue::PriorityQueue;
#[cfg(test)]
use std::io::{ReadWrite, SeekSet};

/*
 * an entry is the number that we read and the file it came from, so after we took
 * out one Entry from the priority queue, we need to know where to take the next number
 * from.
 */
struct Entry {
	value: u64,
	source: File,
}

/* The priority queue depends on the ordering relation and always returns the biggest */
impl Ord for Entry {
	// lesser than implementation
	fn lt(&self, other: &Entry) -> bool {
		// we invert: for us the smallest will have the highest priority
		!(self.value < other.value)
	}
}

/* ordering needs equality for some reason, so there */
impl Eq for Entry {
	fn eq(&self, other: &Entry) -> bool {
		self.value == other.value
	}
}

/* A struct to hold our queue */
struct PriorityFile {
	queue: PriorityQueue<Entry>,
}

/*
 * 'Class' that takes a list of files and does a k-way merge of the smallest values
 * from the files
 */
impl PriorityFile {
	pub fn new(files: Vec<File>) -> PriorityFile {
		let mut queue: PriorityQueue<Entry> = PriorityQueue::new();
		// now that we got the queue, populate it with initial values

		// move_iter is required to get the values themselves, not borrowed pointers to them
		for mut f in files.move_iter() {
			// read from file, append to queue
			let number = match f.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			queue.push(Entry {value: number, source: f});
		};
		PriorityFile {queue: queue}
	}

	/* Returns the next smallest number from the list of files or None if exhausted */
	pub fn next(&mut self) -> Option<u64> {
		match self.queue.maybe_pop() {
			Some(mut res) => {
				let v = res.value;
				match res.source.read_le_u64() {
					Ok(num) => self.queue.push(Entry {value: num, source: res.source}),
					Err(_) => (),
				};
				Some(v)
			},
			None => None,
		}
	}
}

fn externalSort(mut fdInput: File, size: u64, mut fdOutput: File, memSize: u64) {
	let mut runs = size / memSize;
	let items_per_run = (memSize / 8) as uint;
	let over = ((size / 8) - (items_per_run as u64 * runs)) as uint;

	println!("There will be {} runs with {} elements each.", runs, items_per_run);
	if over > 0 {
		println!("And one additional run with {} items.", over);
	}

	/* initialize a Vector. Rust Vectors grow but this one will stay fixed */
	let mut run: Vec<u64> = Vec::with_capacity(items_per_run);
	// preallocate it with zeroes
	for _ in range(0, items_per_run) {
		run.push(0_u64);
	}

	// let's create a temporary directory to store the sorted chunks
	// neat: TempDir deletes the directory and its contents when it goes out of
	// scope
	let overflow_dir = match TempDir::new("externalsort") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let overflow_path = overflow_dir.path();
	println!("temp dir path {}", overflow_path.display());

	// start the sorting runs
	for n in range(0, runs) {
		// iterate over our 'run' buffer and read from the input file
		// I suppose the values are little endian, which considering we are most
		// likely on x86, is a reasonable assumption
		for element in run.mut_iter() {
			let number = match fdInput.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			*element = number;
		};
		// O(n log n) sort from the stdlib, hopefully more or less equivalent to
		// C++ std::sort
		run.sort();

		// now write out the results into numbered files
		let file_path = overflow_path.join(n.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};

		for &element in run.iter() {
			match file_file.write_le_u64(element) {
				Ok (_) => (),
				Err(e) => fail!("writing overflow failed: {}", e),
			};
		};
	};

	/* additional run to catch remaining objects */
	if over > 0 {
		let mut run: Vec<u64> = Vec::with_capacity(over);
		for _ in range(0, over) {
			let number = match fdInput.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			run.push(number);
		}
		run.sort();
		let file_path = overflow_path.join(runs.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};

		for &element in run.iter() {
			match file_file.write_le_u64(element) {
				Ok (_) => (),
				Err(e) => fail!("writing overflow failed: {}", e),
			};
		};
		runs += 1;
	};

	/* sorting done, now merging */

	// collect all file handles. Hopefully less than 4096
	let mut files = Vec::with_capacity(runs as uint);
	for n in range(0, runs) {
		let file_path = overflow_path.join(n.to_str());
		let fd = match File::open(&file_path) {
			Ok(f) => f,
			Err(e) => fail!("overflow file opening error {}", e),
		};
		files.push(fd);
	}

	// we got a nice class here that spits out sorted numbers until exhausted
	let mut pf = PriorityFile::new(files);
	loop {
		match pf.next() {
			// if we got a number, write it to the output
			Some (number) => match fdOutput.write_le_u64(number) {
				Ok(_) => (),
				Err(e) => fail!("writing output file failed: {}", e),
			},
			None => break,
		};
	};
}

fn main() {
	if args().len() < 4 {
		fail!("Argument error: <inputFile> <outputFile> <memoryBufferInMB>");
	}
	let input_file_path = &Path::new(args()[1]);
	let size = match input_file_path.stat() {
		Ok(stat) => stat.size,
		Err(e) => fail!("Couldn't read {}", e),
	};
	let fin = match File::open(input_file_path) {
		Ok(f) => f,
		Err(e) => fail!("input file error: {}", e),
	};
	let fout = match File::open_mode(&Path::new(args()[2]), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	let buffer_size:u64 = match from_str(args()[3]) {
		Some(num) => num,
		None => fail!("Not numeric input"),
	};
	println!("Input file size {}", size);

	externalSort(fin, size, fout, buffer_size * 1024 * 1024);
}

#[test]
fn generate_5gb_and_sort() {
	let mut random = match File::open(&Path::new("/dev/urandom")) {
		Ok(f) => f,
		Err(e) => fail!("Error opening urandom: {}", e),
	};

	let test_dir = match TempDir::new("externalsort-test") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let test_path = test_dir.path();

	let mut unordered_nums = match File::open_mode(&test_path.join("input"), Open, ReadWrite) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut buffer = [0, .. 1024*1024];
	let mut left = 1 * 1024;
	while left > 0 {
		random.read(buffer);
		unordered_nums.write(buffer);
		left -= 1;
	}

	unordered_nums.seek(0, SeekSet);
	let ordered_nums = match File::open_mode(&test_path.join("output"), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	externalSort(unordered_nums, 1024 * 1024 * 1024 * 1, ordered_nums, 1024 * 1024 * 200);
	let mut ordered_nums = match File::open(&test_path.join("output")) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut last = 0_u64;
	loop {
		let number = match ordered_nums.read_le_u64() {
			Ok(num) => num,
			Err(e) => fail!("failed to read u64 from file: {}", e),
		};
		assert!(number >= last);
		last = number;
	}
}
<|endoftext|>


<commit_before>use hyper::status::StatusCode;
use hyper::header::{ContentLength, ContentType};
use time::Duration;
use hyper::mime::{Mime, TopLevel, SubLevel};

use super::Response;

pub trait Context {
    fn emit_error_page(&self, code: StatusCode, response: &mut Response) {
        response.status(code);
        let data = format!("
{}
\n\

Served for you by rotor-http
\n",
            code, code.canonical_reason().unwrap_or("Unknown"));
        let bytes = data.as_bytes();
        response.add_header(ContentLength(bytes.len() as u64)).unwrap();
        response.add_header(ContentType(
            Mime(TopLevel::Text, SubLevel::Html, vec![]))).unwrap();
        response.done_headers().unwrap();
        response.write_body(bytes);
        response.done();
    }
    fn byte_timeout(&self) -> Duration {
        Duration::seconds(10)
    }
}
<commit_msg>Fixup of error page bug<commit_after>use hyper::status::StatusCode;
use hyper::header::{ContentLength, ContentType};
use time::Duration;
use hyper::mime::{Mime, TopLevel, SubLevel};

use super::Response;

pub trait Context {
    fn emit_error_page(&self, code: StatusCode, response: &mut Response) {
        response.status(code);
        let data = format!("
{}
\n\

Served for you by rotor-http
\n", code);
        let bytes = data.as_bytes();
        response.add_header(ContentLength(bytes.len() as u64)).unwrap();
        response.add_header(ContentType(
            Mime(TopLevel::Text, SubLevel::Html, vec![]))).unwrap();
        response.done_headers().unwrap();
        response.write_body(bytes);
        response.done();
    }
    fn byte_timeout(&self) -> Duration {
        Duration::seconds(10)
    }
}
<|endoftext|>