<commit_before>use alloc::boxed::*;

use core::mem;
use core::ptr;

use common::memory;
use common::paging::*;
use common::resource::*;
use common::scheduler;
use common::string::*;
use common::vec::*;

pub const CONTEXT_STACK_SIZE: usize = 1024 * 1024;

pub static mut contexts_ptr: *mut Vec<Box<Context>> = 0 as *mut Vec<Box<Context>>;
pub static mut context_i: usize = 0;
pub static mut context_enabled: bool = false;

/// Switch context
pub unsafe fn context_switch(interrupted: bool) {
    let reenable = scheduler::start_no_ints();

    let contexts = &mut *contexts_ptr;
    if context_enabled {
        let current_i = context_i;
        context_i += 1;
        // The only garbage collection in Redox
        loop {
            if context_i >= contexts.len() {
                context_i -= contexts.len();
            }

            let mut remove = false;
            if let Some(next) = contexts.get(context_i) {
                if next.exited {
                    remove = true;
                }
            }

            if remove {
                drop(contexts.remove(context_i));
            } else {
                break;
            }
        }

        if context_i >= contexts.len() {
            context_i -= contexts.len();
        }

        if context_i != current_i {
            match contexts.get(current_i) {
                Some(current) => match contexts.get(context_i) {
                    Some(next) => {
                        current.interrupted = interrupted;
                        next.interrupted = false;
                        current.remap(next);
                        current.switch(next);
                    }
                    None => (),
                },
                None => (),
            }
        }
    }

    scheduler::end_no_ints(reenable);
}

pub unsafe extern "cdecl" fn context_fork(parent_i: usize){
    let reenable = scheduler::start_no_ints();

    let contexts = &mut *contexts_ptr;
    let mut context_option: Option<Box<Context>> = None;
    if let Some(parent) = contexts.get(parent_i) {
        let stack = memory::alloc(CONTEXT_STACK_SIZE + 512);
        if stack > 0 {
            ::memcpy(stack as *mut u8, parent.stack as *const u8, CONTEXT_STACK_SIZE + 512);

            let mut mem: Vec<ContextMemory> = Vec::new();
            for entry in parent.memory.iter() {
                let physical_address = memory::alloc(entry.virtual_size);
                if physical_address > 0 {
                    ::memcpy(physical_address as *mut u8, entry.physical_address as *const u8, entry.virtual_size);
                    mem.push(ContextMemory {
                        physical_address: physical_address,
                        virtual_address: entry.virtual_address,
                        virtual_size: entry.virtual_size,
                    });
                }
            }

            let mut files: Vec<ContextFile> = Vec::new();
            for file in parent.files.iter() {
                if let Some(resource) = file.resource.dup() {
                    files.push(ContextFile {
                        fd: file.fd,
                        resource: resource
                    });
                }
            }

            context_option = Some(box Context {
                stack: stack,
                stack_ptr: (parent.stack_ptr - parent.stack) + stack,
                fx: stack + CONTEXT_STACK_SIZE,
                fx_enabled: parent.fx_enabled,
                memory: mem,
                cwd: parent.cwd.clone(),
                files: files,
                interrupted: parent.interrupted,
                exited: parent.exited,
            });
        }
    }

    if let Some(context) = context_option {
        contexts.push(context);
    }

    scheduler::end_no_ints(reenable);
}

//TODO: To clean up memory leak, current must be destroyed!
pub unsafe extern "cdecl" fn context_exit() {
    let reenable = scheduler::start_no_ints();

    let contexts = &*contexts_ptr;
    if context_enabled && context_i > 1 {
        match contexts.get(context_i) {
            Some(mut current) => current.exited = true,
            None => (),
        }
    }

    scheduler::end_no_ints(reenable);

    context_switch(false);
}

pub unsafe extern "cdecl" fn context_box(box_fn_ptr: usize) {
    let box_fn = ptr::read(box_fn_ptr as *mut Box<FnBox()>);
    memory::unalloc(box_fn_ptr);
    box_fn();
}

pub struct ContextMemory {
    pub physical_address: usize,
    pub virtual_address: usize,
    pub virtual_size: usize,
}

pub struct ContextFile {
    pub fd: usize,
    pub resource: Box<Resource>,
}

pub struct Context {
    pub stack: usize,
    pub stack_ptr: usize,
    pub fx: usize,
    pub fx_enabled: bool,
    pub memory: Vec<ContextMemory>,
    pub cwd: String,
    pub files: Vec<ContextFile>,
    pub interrupted: bool,
    pub exited: bool,
}

impl Context {
    pub unsafe fn root() -> Box<Self> {
        box Context {
            stack: 0,
            stack_ptr: 0,
            fx: memory::alloc(512),
            fx_enabled: false,
            memory: Vec::new(),
            cwd: String::new(),
            files: Vec::new(),
            interrupted: false,
            exited: false,
        }
    }

    pub unsafe fn new(call: usize, args: &Vec<usize>) -> Box<Self> {
        let stack = memory::alloc(CONTEXT_STACK_SIZE + 512);

        let mut ret = box Context {
            stack: stack,
            stack_ptr: stack + CONTEXT_STACK_SIZE,
            fx: stack + CONTEXT_STACK_SIZE,
            fx_enabled: false,
            memory: Vec::new(),
            cwd: String::new(),
            files: Vec::new(),
            interrupted: false,
            exited: false,
        };

        let ebp = ret.stack_ptr;

        for arg in args.iter() {
            ret.push(*arg);
        }

        ret.push(call); //We will ret into this function call

        ret.push(0); //ESI is a param used in the switch function

        ret.push(1 << 9); //Flags

        let esp = ret.stack_ptr;

        ret.push(0); //EAX
        ret.push(0); //ECX
        ret.push(0); //EDX
        ret.push(0); //EBX
        ret.push(esp); //ESP (ignored)
        ret.push(ebp); //EBP
        ret.push(0); //ESI
        ret.push(0); //EDI

        ret
    }

    pub fn spawn(box_fn: Box<FnBox()>) {
        unsafe {
            let box_fn_ptr: *mut Box<FnBox()> = memory::alloc_type();
            ptr::write(box_fn_ptr, box_fn);

            let mut context_box_args: Vec<usize> = Vec::new();
            context_box_args.push(box_fn_ptr as usize);
            context_box_args.push(context_exit as usize);

            let reenable = scheduler::start_no_ints();
            if contexts_ptr as usize > 0 {
                (*contexts_ptr).push(Context::new(context_box as usize, &context_box_args));
            }
            scheduler::end_no_ints(reenable);
        }
    }

    pub unsafe fn push(&mut self, data: usize) {
        self.stack_ptr -= mem::size_of::<usize>();
        ptr::write(self.stack_ptr as *mut usize, data);
    }

    pub unsafe fn map(&mut self) {
        for entry in self.memory.iter() {
            for i in 0..(entry.virtual_size + 4095) / 4096 {
                Page::new(entry.virtual_address + i * 4096).map(entry.physical_address + i * 4096);
            }
        }
    }

    pub unsafe fn unmap(&mut self) {
        for entry in self.memory.iter() {
            for i in 0..(entry.virtual_size + 4095) / 4096 {
                Page::new(entry.virtual_address + i * 4096).map_identity();
            }
        }
    }

    pub unsafe fn remap(&mut self, other: &mut Self) {
        self.unmap();
        other.map();
    }

    //Warning: This function MUST be inspected in disassembly for correct push/pop
    //It should have exactly one extra push/pop of ESI
    #[cold]
    #[inline(never)]
    #[cfg(target_arch = "x86")]
    pub unsafe fn switch(&mut self, other: &mut Self) {
        asm!("pushfd
            pushad
            mov [esi], esp"
            :
            : "{esi}"(&mut self.stack_ptr)
            : "memory"
            : "intel", "volatile");

        asm!("fxsave [esi]"
            :
            : "{esi}"(self.fx)
            : "memory"
            : "intel", "volatile");
        self.fx_enabled = true;

        //TODO: Clear registers
        if other.fx_enabled {
            asm!("fxrstor [esi]"
                :
                : "{esi}"(other.fx)
                : "memory"
                : "intel", "volatile");
        }

        asm!("mov esp, [esi]
            popad
            popfd"
            :
            : "{esi}"(&mut other.stack_ptr)
            : "memory"
            : "intel", "volatile");
    }

    //Warning: This function MUST be inspected in disassembly for correct push/pop
    //It should have exactly one extra push/pop of ESI
    #[cold]
    #[inline(never)]
    #[cfg(target_arch = "x86_64")]
    pub unsafe fn switch(&mut self, other: &mut Self) {
        asm!("pushfq
            push rax
            push rcx
            push rdx
            push rbx
            push rbp
            push rsi
            push rdi
            push r8
            push r9
            push r10
            push r11
            push r12
            push r13
            push r14
            push r15
            mov [rsi], rsp"
            :
            : "{rsi}"(&mut self.stack_ptr)
            : "memory"
            : "intel", "volatile");

        asm!("fxsave [rsi]"
            :
            : "{rsi}"(self.fx)
            : "memory"
            : "intel", "volatile");
        self.fx_enabled = true;

        //TODO: Clear registers
        if other.fx_enabled {
            asm!("fxrstor [rsi]"
                :
                : "{rsi}"(other.fx)
                : "memory"
                : "intel", "volatile");
        }

        asm!("mov rsp, [rsi]
            pop r15
            pop r14
            pop r13
            pop r12
            pop r11
            pop r10
            pop r9
            pop r8
            pop rdi
            pop rsi
            pop rbp
            pop rbx
            pop rdx
            pop rcx
            pop rax
            popfq"
            :
            : "{rsi}"(&mut other.stack_ptr)
            : "memory"
            : "intel", "volatile");
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        while let Some(file) = self.files.remove(0) {
            drop(file);
        }

        while let Some(entry) = self.memory.remove(0) {
            unsafe {
                memory::unalloc(entry.physical_address);
            }
        }

        if self.stack > 0 {
            unsafe {
                memory::unalloc(self.stack);
            }
            self.stack = 0;
        }
    }
}
<commit_msg>Context Improvements<commit_after>use alloc::boxed::*;

use core::mem;
use core::ptr;

use common::memory;
use common::paging::*;
use common::resource::*;
use common::scheduler;
use common::string::*;
use common::vec::*;

pub const CONTEXT_STACK_SIZE: usize = 1024 * 1024;

pub static mut contexts_ptr: *mut Vec<Box<Context>> = 0 as *mut Vec<Box<Context>>;
pub static mut context_i: usize = 0;
pub static mut context_enabled: bool = false;

/// Switch context
///
/// Unsafe due to interrupt disabling, raw pointers, and unsafe Context functions
pub unsafe fn context_switch(interrupted: bool) {
    let reenable = scheduler::start_no_ints();

    let contexts = &mut *contexts_ptr;
    if context_enabled {
        let current_i = context_i;
        context_i += 1;
        // The only garbage collection in Redox
        loop {
            if context_i >= contexts.len() {
                context_i -= contexts.len();
            }

            let mut remove = false;
            if let Some(next) = contexts.get(context_i) {
                if next.exited {
                    remove = true;
                }
            }

            if remove {
                drop(contexts.remove(context_i));
            } else {
                break;
            }
        }

        if context_i >= contexts.len() {
            context_i -= contexts.len();
        }

        if context_i != current_i {
            match contexts.get(current_i) {
                Some(current) => match contexts.get(context_i) {
                    Some(next) => {
                        current.interrupted = interrupted;
                        next.interrupted = false;
                        current.remap(next);
                        current.switch(next);
                    }
                    None => (),
                },
                None => (),
            }
        }
    }

    scheduler::end_no_ints(reenable);
}

/// Fork context
///
/// Unsafe due to interrupt disabling, C memory handling, and raw pointers
pub unsafe extern "cdecl" fn context_fork(parent_i: usize){
    let reenable = scheduler::start_no_ints();

    let contexts = &mut *contexts_ptr;
    let mut context_option: Option<Box<Context>> = None;
    if let Some(parent) = contexts.get(parent_i) {
        let stack = memory::alloc(CONTEXT_STACK_SIZE + 512);
        if stack > 0 {
            ::memcpy(stack as *mut u8, parent.stack as *const u8, CONTEXT_STACK_SIZE + 512);

            let mut mem: Vec<ContextMemory> = Vec::new();
            for entry in parent.memory.iter() {
                let physical_address = memory::alloc(entry.virtual_size);
                if physical_address > 0 {
                    ::memcpy(physical_address as *mut u8, entry.physical_address as *const u8, entry.virtual_size);
                    mem.push(ContextMemory {
                        physical_address: physical_address,
                        virtual_address: entry.virtual_address,
                        virtual_size: entry.virtual_size,
                    });
                }
            }

            let mut files: Vec<ContextFile> = Vec::new();
            for file in parent.files.iter() {
                if let Some(resource) = file.resource.dup() {
                    files.push(ContextFile {
                        fd: file.fd,
                        resource: resource
                    });
                }
            }

            context_option = Some(box Context {
                stack: stack,
                stack_ptr: (parent.stack_ptr - parent.stack) + stack,
                fx: stack + CONTEXT_STACK_SIZE,
                fx_enabled: parent.fx_enabled,
                memory: mem,
                cwd: parent.cwd.clone(),
                files: files,
                interrupted: parent.interrupted,
                exited: parent.exited,
            });
        }
    }

    if let Some(context) = context_option {
        contexts.push(context);
    }

    scheduler::end_no_ints(reenable);
}

//TODO: To clean up memory leak, current must be destroyed!
/// Exit context
///
/// Unsafe due to interrupt disabling and raw pointers
pub unsafe fn context_exit() {
    let reenable = scheduler::start_no_ints();

    let contexts = &*contexts_ptr;
    if context_enabled && context_i > 1 {
        match contexts.get(context_i) {
            Some(mut current) => current.exited = true,
            None => (),
        }
    }

    scheduler::end_no_ints(reenable);

    context_switch(false);
}

// Currently unused?
/// Reads a Boxed function and executes it
///
/// Unsafe due to raw memory handling and FnBox
pub unsafe extern "cdecl" fn context_box(box_fn_ptr: usize) {
    let box_fn = ptr::read(box_fn_ptr as *mut Box<FnBox()>);
    memory::unalloc(box_fn_ptr);
    box_fn();
}

pub struct ContextMemory {
    pub physical_address: usize,
    pub virtual_address: usize,
    pub virtual_size: usize,
}

pub struct ContextFile {
    pub fd: usize,
    pub resource: Box<Resource>,
}

pub struct Context {
    pub stack: usize,
    pub stack_ptr: usize,
    pub fx: usize,
    pub fx_enabled: bool,
    pub memory: Vec<ContextMemory>,
    pub cwd: String,
    pub files: Vec<ContextFile>,
    pub interrupted: bool,
    pub exited: bool,
}

impl Context {
    pub unsafe fn root() -> Box<Self> {
        box Context {
            stack: 0,
            stack_ptr: 0,
            fx: memory::alloc(512),
            fx_enabled: false,
            memory: Vec::new(),
            cwd: String::new(),
            files: Vec::new(),
            interrupted: false,
            exited: false,
        }
    }

    pub unsafe fn new(call: usize, args: &Vec<usize>) -> Box<Self> {
        let stack = memory::alloc(CONTEXT_STACK_SIZE + 512);

        let mut ret = box Context {
            stack: stack,
            stack_ptr: stack + CONTEXT_STACK_SIZE,
            fx: stack + CONTEXT_STACK_SIZE,
            fx_enabled: false,
            memory: Vec::new(),
            cwd: String::new(),
            files: Vec::new(),
            interrupted: false,
            exited: false,
        };

        let ebp = ret.stack_ptr;

        for arg in args.iter() {
            ret.push(*arg);
        }

        ret.push(call); //We will ret into this function call

        ret.push(0); //ESI is a param used in the switch function

        ret.push(1 << 9); //Flags

        let esp = ret.stack_ptr;

        ret.push(0); //EAX
        ret.push(0); //ECX
        ret.push(0); //EDX
        ret.push(0); //EBX
        ret.push(esp); //ESP (ignored)
        ret.push(ebp); //EBP
        ret.push(0); //ESI
        ret.push(0); //EDI

        ret
    }

    pub fn spawn(box_fn: Box<FnBox()>) {
        unsafe {
            let box_fn_ptr: *mut Box<FnBox()> = memory::alloc_type();
            ptr::write(box_fn_ptr, box_fn);

            let mut context_box_args: Vec<usize> = Vec::new();
            context_box_args.push(box_fn_ptr as usize);
            context_box_args.push(context_exit as usize);

            let reenable = scheduler::start_no_ints();
            if contexts_ptr as usize > 0 {
                (*contexts_ptr).push(Context::new(context_box as usize, &context_box_args));
            }
            scheduler::end_no_ints(reenable);
        }
    }

    pub unsafe fn push(&mut self, data: usize) {
        self.stack_ptr -= mem::size_of::<usize>();
        ptr::write(self.stack_ptr as *mut usize, data);
    }

    pub unsafe fn map(&mut self) {
        for entry in self.memory.iter() {
            for i in 0..(entry.virtual_size + 4095) / 4096 {
                Page::new(entry.virtual_address + i * 4096).map(entry.physical_address + i * 4096);
            }
        }
    }

    pub unsafe fn unmap(&mut self) {
        for entry in self.memory.iter() {
            for i in 0..(entry.virtual_size + 4095) / 4096 {
                Page::new(entry.virtual_address + i * 4096).map_identity();
            }
        }
    }

    pub unsafe fn remap(&mut self, other: &mut Self) {
        self.unmap();
        other.map();
    }

    //Warning: This function MUST be inspected in disassembly for correct push/pop
    //It should have exactly one extra push/pop of ESI
    #[cold]
    #[inline(never)]
    #[cfg(target_arch = "x86")]
    pub unsafe fn switch(&mut self, other: &mut Self) {
        asm!("pushfd
            pushad
            mov [esi], esp"
            :
            : "{esi}"(&mut self.stack_ptr)
            : "memory"
            : "intel", "volatile");

        asm!("fxsave [esi]"
            :
            : "{esi}"(self.fx)
            : "memory"
            : "intel", "volatile");
        self.fx_enabled = true;

        //TODO: Clear registers
        if other.fx_enabled {
            asm!("fxrstor [esi]"
                :
                : "{esi}"(other.fx)
                : "memory"
                : "intel", "volatile");
        }

        asm!("mov esp, [esi]
            popad
            popfd"
            :
            : "{esi}"(&mut other.stack_ptr)
            : "memory"
            : "intel", "volatile");
    }

    //Warning: This function MUST be inspected in disassembly for correct push/pop
    //It should have exactly one extra push/pop of ESI
    #[cold]
    #[inline(never)]
    #[cfg(target_arch = "x86_64")]
    pub unsafe fn switch(&mut self, other: &mut Self) {
        asm!("pushfq
            push rax
            push rcx
            push rdx
            push rbx
            push rbp
            push rsi
            push rdi
            push r8
            push r9
            push r10
            push r11
            push r12
            push r13
            push r14
            push r15
            mov [rsi], rsp"
            :
            : "{rsi}"(&mut self.stack_ptr)
            : "memory"
            : "intel", "volatile");

        asm!("fxsave [rsi]"
            :
            : "{rsi}"(self.fx)
            : "memory"
            : "intel", "volatile");
        self.fx_enabled = true;

        //TODO: Clear registers
        if other.fx_enabled {
            asm!("fxrstor [rsi]"
                :
                : "{rsi}"(other.fx)
                : "memory"
                : "intel", "volatile");
        }

        asm!("mov rsp, [rsi]
            pop r15
            pop r14
            pop r13
            pop r12
            pop r11
            pop r10
            pop r9
            pop r8
            pop rdi
            pop rsi
            pop rbp
            pop rbx
            pop rdx
            pop rcx
            pop rax
            popfq"
            :
            : "{rsi}"(&mut other.stack_ptr)
            : "memory"
            : "intel", "volatile");
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        while let Some(file) = self.files.remove(0) {
            drop(file);
        }

        while let Some(entry) = self.memory.remove(0) {
            unsafe {
                memory::unalloc(entry.physical_address);
            }
        }

        if self.stack > 0 {
            unsafe {
                memory::unalloc(self.stack);
            }
            self.stack = 0;
        }
    }
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

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

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        println!("poll reading");
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                println!("got length: {}", buf[0]);
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W> where W: ::std::io::Write {
    writer: Option<W>,
    message: Vec<u8>,
    pos: usize,
    wrote_header: bool,
}

impl <W> Writing<W> where W: ::std::io::Write {
    fn new(writer: W, message: Vec<u8>) -> Writing<W> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W> Future for Writing<W> where W: ::std::io::Write {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr) -> Box<Future<Item=(), Error=::std::io::Error> + Send> {
    use all::All;

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
        tokio_core::io::write_all(stream, [0]).map(|(a,_)| a)
    });

    let mut subscribers = Vec::new();
    for _ in 0..1 {
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
            tokio_core::io::write_all(stream, [1]).map(|(a,_)| a)
        }));
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(|(publisher, subscribers)| {
        println!("connected");

        tie_knot((publisher, subscribers, 10i32), |(publisher, subscribers, n)| {
            println!("looping {}", n);
            Writing::new(publisher, vec![n as u8, 1,2,3]).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(|s| {
                    Reading::new(s).and_then(|(s, m)| {
                        println!("got: {:?}", m);
                        Ok(s)
                    })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n - 1), n > 0))
                })
            })
        }).map(|_| ())
    }))
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = "127.0.0.1:8080";
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let _child = ::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::inherit())
        .stderr(::std::process::Stdio::inherit())
        .spawn();

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();

    let f = pool.spawn(new_task(&handle, &addr));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<commit_msg>optional address flag<commit_after>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

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

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        println!("poll reading");
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                println!("got length: {}", buf[0]);
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W> where W: ::std::io::Write {
    writer: Option<W>,
    message: Vec<u8>,
    pos: usize,
    wrote_header: bool,
}

impl <W> Writing<W> where W: ::std::io::Write {
    fn new(writer: W, message: Vec<u8>) -> Writing<W> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W> Future for Writing<W> where W: ::std::io::Write {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr) -> Box<Future<Item=(), Error=::std::io::Error> + Send> {
    use all::All;

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
        tokio_core::io::write_all(stream, [0]).map(|(a,_)| a)
    });

    let mut subscribers = Vec::new();
    for _ in 0..1 {
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(|stream| {
            tokio_core::io::write_all(stream, [1]).map(|(a,_)| a)
        }));
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(|(publisher, subscribers)| {
        println!("connected");

        tie_knot((publisher, subscribers, 10i32), |(publisher, subscribers, n)| {
            println!("looping {}", n);
            Writing::new(publisher, vec![n as u8, 1,2,3]).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(|s| {
                    Reading::new(s).and_then(|(s, m)| {
                        println!("got: {:?}", m);
                        Ok(s)
                    })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n - 1), n > 0))
                })
            })
        }).map(|_| ())
    }))
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .arg(Arg::with_name("server")
             .required(false)
             .long("server")
             .short("s")
             .value_name("address")
             .default_value("127.0.0.1:8080")
             .help("TODO"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = matches.value_of("server").unwrap();
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let _child = ::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::inherit())
        .stderr(::std::process::Stdio::inherit())
        .spawn();

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();

    let f = pool.spawn(new_task(&handle, &addr));

    try!(core.run(f));

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<|endoftext|>


<commit_before>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

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

use byteorder::{LittleEndian, ByteOrder};

use std::cell::{Cell};
use std::rc::Rc;
use std::time::Duration;

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    writer: Option<W>,
    message: B,
    pos: usize,
    wrote_header: bool,
}

impl <W, B> Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    fn new(writer: W, message: B) -> Writing<W, B> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W, B> Future for Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.as_ref().len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message.as_ref()[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.as_ref().len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

#[derive(Clone)]
struct ConnectionIdSource {
    next_id: Rc<Cell<u64>>,
}

impl ConnectionIdSource {
    fn new() -> ConnectionIdSource {
        ConnectionIdSource {
            next_id: Rc::new(Cell::new(1)),
        }
    }

    fn next(&self) -> u64 {
        let result = self.next_id.get();
        self.next_id.set(result + 1);
        result
    }
}

static CLOCK_PREFIX: &'static [u8] = &[0,0,0,0,0,0,0,0];

fn read_until_message_with_prefix<R, B>(
    reader: R,
    prefix: B,
    timeout_seconds: u64)
    -> Box<Future<Item=(R, Option<Vec<u8>>), Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
          B: AsRef<[u8]> + 'static + Send,
{
    Box::new(tie_knot((reader, prefix, 0, None), move |(reader, prefix, ticks, _)| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(message) => {
                    let len = prefix.as_ref().len();
                    let mut new_ticks = ticks;
                    if  &message[0..len] != prefix.as_ref() {
                        Ok(((reader, prefix, new_ticks, Some(message)), false))
                    } else {
                        if &message[..CLOCK_PREFIX.len()] == CLOCK_PREFIX {
                            new_ticks = ticks + 1;
                        }

                        if new_ticks > timeout_seconds {
                            Ok(((reader, prefix, new_ticks, None), false))
                        }  else {
                            Ok(((reader, prefix, new_ticks, Some(message)), true))
                        }
                    }
                }
                None => {
                    Ok(((reader, prefix, ticks, None), false))
                }
            }
        })
    }).map(|(reader, _, _, message)| (reader, message)))
}

fn read_until_eof<R>(reader: R)
                     -> Box<Future<Item=R, Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
{
    Box::new(tie_knot(reader, move |reader| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(_) => {
                    Ok((reader, true))
                }
                None => {
                    Ok((reader, false))
                }
            }
        })
    }))
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr,
            connection_id_source: ConnectionIdSource,
            number_of_messages: u64)
            -> Box<Future<Item=usize, Error=::std::io::Error> + Send>
{
    use all::All;
    use std::sync::atomic::{AtomicUsize, Ordering};

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle);
    let publisher_id = connection_id_source.next();

    let mut subscribers = Vec::new();
    for _ in 0..2 {
        let subscriber_id = connection_id_source.next();
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(move |socket| {
            let mut buf = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..], subscriber_id);
            Writing::new(socket, buf).and_then(move |socket| {
                read_until_message_with_prefix(socket, buf, 2)
                    .map(|(socket, _message)| { socket })
            })
        }))
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(move |(publisher, subscribers)| {
        let successful_message_count = ::futures::task::TaskRc::new(AtomicUsize::new(0));
        let smc = successful_message_count.clone();
        tie_knot((publisher, subscribers, 0u64), move |(publisher, subscribers, n)| {
            println!("looping {}", n);
            let mut buf = vec![255; 16];
            let mut prefix = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..8], publisher_id);
            <LittleEndian as ByteOrder>::write_u64(&mut prefix[..], publisher_id);

            let smc = smc.clone();
            Writing::new(publisher, buf).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(move |s| {
                    let smc = smc.clone();
                    read_until_message_with_prefix(s, prefix, 2)
                        .map(move |(stream, message)| {
                            match message {
                                Some(_) => {
                                    // TODO check that the message is what the publisher sent.
                                    smc.with(|x| {
                                        x.fetch_add(1, Ordering::SeqCst);
                                    });
                                }
                                None => (),
                            }
                            stream
                        })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n + 1), n < number_of_messages))
                })
            })
        }).and_then(|(publisher, subscribers, _)| {
            try!(publisher.shutdown(::std::net::Shutdown::Write));
            for sub in &subscribers {
                try!(sub.shutdown(::std::net::Shutdown::Write));
            }
            Ok((publisher, subscribers))
        }).and_then(|(publisher, subscribers)| {
            read_until_eof(publisher).join(
                All::new(subscribers.into_iter().map(|sub| {
                    read_until_eof(sub)
                })))
        }).map(move |_| successful_message_count.with(|x| x.load(Ordering::SeqCst)))
    }))
}

struct ChildProcess {
    child: ::std::process::Child
}

impl ChildProcess {
    fn new(child: ::std::process::Child) -> ChildProcess {
        ChildProcess { child: child }
    }
}

impl Drop for ChildProcess {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .arg(Arg::with_name("server")
             .required(false)
             .long("server")
             .short("s")
             .value_name("address")
             .default_value("127.0.0.1:8080")
             .help("address to use to connect to server"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = matches.value_of("server").unwrap();
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let mut child = try!(::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::piped())
        .stderr(::std::process::Stdio::inherit())
        .spawn());

    let mut child_stdout = ::std::io::BufReader::new(child.stdout.take().unwrap());
    let _wrapped_child = ChildProcess::new(child);

    let mut first_line = String::new();
    try!(::std::io::BufRead::read_line(&mut child_stdout, &mut first_line));

    if !first_line.starts_with("listening on") {
        return Err(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            format!(
                "expected first line from server to start with 'listening on ', but got {}",
                first_line)))
    }

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();
    let connection_id_source = ConnectionIdSource::new();

    // Start a connection whose sole job is to send periodic "tick" messages.
    let clock_connection = ::tokio_core::net::TcpStream::connect(&addr, &handle);
    let handle1 = handle.clone();
    handle.spawn(clock_connection.and_then(move |stream| {
        tie_knot((stream, handle1), move |(stream, handle)| {
            use tokio_core::reactor::Timeout;
            Timeout::new(Duration::from_secs(1), &handle).expect("creating timeout").and_then(move |()| {
                Writing::new(stream, CLOCK_PREFIX).map(move |stream| {
                    ((stream, handle), true)
                })
            })
        })
    }).map(|_| ()).map_err(|e| { println!("error from clock task: {}", e); () }));

    let f = pool.spawn(new_task(&handle, &addr, connection_id_source.clone(), 5)
                       .join(new_task(&core.handle(), &addr, connection_id_source, 5)));

    let x = try!(core.run(f));
    println!("x = {:?}", x);

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<commit_msg>number_of_subscribers<commit_after>extern crate byteorder;
extern crate clap;

#[macro_use]
extern crate futures;
extern crate futures_cpupool;

#[macro_use]
extern crate tokio_core;

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

use byteorder::{LittleEndian, ByteOrder};

use std::cell::{Cell};
use std::rc::Rc;
use std::time::Duration;

mod all {
    use futures::{Async, Poll, Future};
    enum ElemState<T> where T: Future {
        Pending(T),
        Done(T::Item),
    }

    pub struct All<T> where T: Future {
        elems: Vec<ElemState<T>>,
    }

    impl <T> All<T> where T: Future {
        pub fn new<I>(futures: I) -> All<T> where I: Iterator<Item=T> {
            let mut result = All { elems: Vec::new() };
            for f in futures {
                result.elems.push(ElemState::Pending(f))
            }
            result
        }
    }

    impl <T> Future for All<T> where T: Future {
        type Item = Vec<T::Item>;
        type Error = T::Error;

        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            let mut all_done = true;

            for idx in 0 .. self.elems.len() {
                let done_val = match &mut self.elems[idx] {
                    &mut ElemState::Pending(ref mut t) => {
                        match t.poll() {
                            Ok(Async::Ready(t)) => t,
                            Ok(Async::NotReady) => {
                                all_done = false;
                                continue
                            }
                            Err(e) => return Err(e),
                        }
                    }
                    &mut ElemState::Done(ref mut _v) => continue,
                };

                self.elems[idx] = ElemState::Done(done_val);
            }

            if all_done {
                let mut result = Vec::new();
                let elems = ::std::mem::replace(&mut self.elems, Vec::new());
                for e in elems.into_iter() {
                    match e {
                        ElemState::Done(t) => result.push(t),
                        _ => unreachable!(),
                    }
                }
                Ok(Async::Ready(result))
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}

struct Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    f: F,
    in_progress: T,
}

fn tie_knot<F, S, T, E>(initial_state: S, f: F) -> Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>,
{
    let in_progress = f(initial_state);
    Knot {
        f: f,
        in_progress: in_progress,
    }
}

impl <F, S, T, E> Future for Knot<F, S, T, E>
    where F: Fn(S) -> T,
          T: Future<Item=(S, bool), Error=E>
{
    type Item = S;
    type Error = E;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let (s, more) = try_ready!(self.in_progress.poll());
            if more {
                self.in_progress = (self.f)(s);
            } else {
                return Ok(Async::Ready(s))
            }
        }
    }
}

struct Reading<R> where R: ::std::io::Read {
    reader: Option<R>,
    buffer: Vec<u8>,
    pos: usize,
    frame_end: Option<u8>,
}

impl <R> Reading<R> where R: ::std::io::Read {
    fn new(reader: R) -> Reading<R> {
        Reading {
            reader: Some(reader),
            buffer: Vec::new(),
            pos: 0,
            frame_end: None,
        }
    }
}

impl <R> Future for Reading<R> where R: ::std::io::Read {
    type Item = (R, Option<Vec<u8>>);
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let Some(frame_end) = self.frame_end {
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut self.buffer[self.pos..frame_end as usize]));
                self.pos += n;
                if self.pos == frame_end as usize {
                    self.pos = 0;
                    let result = ::std::mem::replace(&mut self.buffer, Vec::new());
                    self.frame_end = None;
                    return Ok(Async::Ready((self.reader.take().unwrap(), Some(result))))
                }
            } else {
                let mut buf = [0u8];
                let n = try_nb!(self.reader.as_mut().unwrap().read(&mut buf));
                if n == 0 { // EOF
                    return Ok(Async::Ready((self.reader.take().unwrap(), None)))
                }
                self.frame_end = Some(buf[0]);
                self.buffer = vec![0; buf[0] as usize];
            }
        }
    }
}

pub struct Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    writer: Option<W>,
    message: B,
    pos: usize,
    wrote_header: bool,
}

impl <W, B> Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    fn new(writer: W, message: B) -> Writing<W, B> {
        Writing {
            writer: Some(writer),
            message: message,
            pos: 0,
            wrote_header: false,
        }
    }
}

impl <W, B> Future for Writing<W, B> where W: ::std::io::Write, B: AsRef<[u8]> {
    type Item = W;
    type Error = ::std::io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if !self.wrote_header {
                let buf = [self.message.as_ref().len() as u8];
                try_nb!(self.writer.as_mut().unwrap().write(&buf));
                self.wrote_header = true;
            } else {
                let n = try_nb!(self.writer.as_mut().unwrap().write(&self.message.as_ref()[self.pos..]));
                self.pos += n;
                if self.pos >= self.message.as_ref().len() {
                    return Ok(Async::Ready(self.writer.take().unwrap()))
                }
            }
        }
    }
}

#[derive(Clone)]
struct ConnectionIdSource {
    next_id: Rc<Cell<u64>>,
}

impl ConnectionIdSource {
    fn new() -> ConnectionIdSource {
        ConnectionIdSource {
            next_id: Rc::new(Cell::new(1)),
        }
    }

    fn next(&self) -> u64 {
        let result = self.next_id.get();
        self.next_id.set(result + 1);
        result
    }
}

static CLOCK_PREFIX: &'static [u8] = &[0,0,0,0,0,0,0,0];

fn read_until_message_with_prefix<R, B>(
    reader: R,
    prefix: B,
    timeout_seconds: u64)
    -> Box<Future<Item=(R, Option<Vec<u8>>), Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
          B: AsRef<[u8]> + 'static + Send,
{
    Box::new(tie_knot((reader, prefix, 0, None), move |(reader, prefix, ticks, _)| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(message) => {
                    let len = prefix.as_ref().len();
                    let mut new_ticks = ticks;
                    if  &message[0..len] != prefix.as_ref() {
                        Ok(((reader, prefix, new_ticks, Some(message)), false))
                    } else {
                        if &message[..CLOCK_PREFIX.len()] == CLOCK_PREFIX {
                            new_ticks = ticks + 1;
                        }

                        if new_ticks > timeout_seconds {
                            Ok(((reader, prefix, new_ticks, None), false))
                        }  else {
                            Ok(((reader, prefix, new_ticks, Some(message)), true))
                        }
                    }
                }
                None => {
                    Ok(((reader, prefix, ticks, None), false))
                }
            }
        })
    }).map(|(reader, _, _, message)| (reader, message)))
}

fn read_until_eof<R>(reader: R)
                     -> Box<Future<Item=R, Error=::std::io::Error> + 'static + Send>
    where R: ::std::io::Read + 'static + Send,
{
    Box::new(tie_knot(reader, move |reader| {
        Reading::new(reader).and_then(move |(reader, message)| {
            match message {
                Some(_) => {
                    Ok((reader, true))
                }
                None => {
                    Ok((reader, false))
                }
            }
        })
    }))
}

fn new_task(handle: &::tokio_core::reactor::Handle,
            addr: &::std::net::SocketAddr,
            connection_id_source: ConnectionIdSource,
            number_of_messages: u64,
            number_of_subscribers: u64)
            -> Box<Future<Item=usize, Error=::std::io::Error> + Send>
{
    use all::All;
    use std::sync::atomic::{AtomicUsize, Ordering};

    let publisher = ::tokio_core::net::TcpStream::connect(addr, handle);
    let publisher_id = connection_id_source.next();

    let mut subscribers = Vec::new();
    for _ in 0..number_of_subscribers {
        let subscriber_id = connection_id_source.next();
        subscribers.push(::tokio_core::net::TcpStream::connect(addr, handle).and_then(move |socket| {
            let mut buf = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..], subscriber_id);
            Writing::new(socket, buf).and_then(move |socket| {
                read_until_message_with_prefix(socket, buf, 2)
                    .map(|(socket, _message)| { socket })
            })
        }))
    }

    Box::new(publisher.join(All::new(subscribers.into_iter())).and_then(move |(publisher, subscribers)| {
        let successful_message_count = ::futures::task::TaskRc::new(AtomicUsize::new(0));
        let smc = successful_message_count.clone();
        tie_knot((publisher, subscribers, 0u64), move |(publisher, subscribers, n)| {
            println!("looping {}", n);
            let mut buf = vec![255; 16];
            let mut prefix = [0; 8];
            <LittleEndian as ByteOrder>::write_u64(&mut buf[..8], publisher_id);
            <LittleEndian as ByteOrder>::write_u64(&mut prefix[..], publisher_id);

            let smc = smc.clone();
            Writing::new(publisher, buf).and_then(move |publisher| {
                All::new(subscribers.into_iter().map(move |s| {
                    let smc = smc.clone();
                    read_until_message_with_prefix(s, prefix, 2)
                        .map(move |(stream, message)| {
                            match message {
                                Some(_) => {
                                    // TODO check that the message is what the publisher sent.
                                    smc.with(|x| {
                                        x.fetch_add(1, Ordering::SeqCst);
                                    });
                                }
                                None => (),
                            }
                            stream
                        })
                })).and_then(move |subscribers| {
                    futures::finished(((publisher, subscribers, n + 1), n < number_of_messages))
                })
            })
        }).and_then(|(publisher, subscribers, _)| {
            try!(publisher.shutdown(::std::net::Shutdown::Write));
            for sub in &subscribers {
                try!(sub.shutdown(::std::net::Shutdown::Write));
            }
            Ok((publisher, subscribers))
        }).and_then(|(publisher, subscribers)| {
            read_until_eof(publisher).join(
                All::new(subscribers.into_iter().map(|sub| {
                    read_until_eof(sub)
                })))
        }).map(move |_| successful_message_count.with(|x| x.load(Ordering::SeqCst)))
    }))
}

struct ChildProcess {
    child: ::std::process::Child
}

impl ChildProcess {
    fn new(child: ::std::process::Child) -> ChildProcess {
        ChildProcess { child: child }
    }
}

impl Drop for ChildProcess {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

pub fn run() -> Result<(), ::std::io::Error> {
    use clap::{App, Arg};
    let matches = App::new("Zillions benchmarker")
        .version("0.0.0")
        .about("Does awesome things")
        .arg(Arg::with_name("EXECUTABLE")
             .required(true)
             .index(1)
             .help("The executable to benchmark"))
        .arg(Arg::with_name("server")
             .required(false)
             .long("server")
             .short("s")
             .value_name("address")
             .default_value("127.0.0.1:8080")
             .help("address to use to connect to server"))
        .get_matches();

    let executable = matches.value_of("EXECUTABLE").unwrap();

    println!("exectuable: {}", executable);

    let addr_str = matches.value_of("server").unwrap();
    let addr = match addr_str.parse::<::std::net::SocketAddr>() {
        Ok(a) => a,
        Err(e) => {
            panic!("failed to parse socket address {}", e);
        }
    };

    let mut child = try!(::std::process::Command::new(executable)
        .arg(addr_str)
        .stdout(::std::process::Stdio::piped())
        .stderr(::std::process::Stdio::inherit())
        .spawn());

    let mut child_stdout = ::std::io::BufReader::new(child.stdout.take().unwrap());
    let _wrapped_child = ChildProcess::new(child);

    let mut first_line = String::new();
    try!(::std::io::BufRead::read_line(&mut child_stdout, &mut first_line));

    if !first_line.starts_with("listening on") {
        return Err(::std::io::Error::new(
            ::std::io::ErrorKind::Other,
            format!(
                "expected first line from server to start with 'listening on ', but got {}",
                first_line)))
    }

    // start tokio reactor
    let mut core = try!(::tokio_core::reactor::Core::new());

    let handle = core.handle();

    let pool = ::futures_cpupool::CpuPool::new_num_cpus();
    let connection_id_source = ConnectionIdSource::new();

    // Start a connection whose sole job is to send periodic "tick" messages.
    let clock_connection = ::tokio_core::net::TcpStream::connect(&addr, &handle);
    let handle1 = handle.clone();
    handle.spawn(clock_connection.and_then(move |stream| {
        tie_knot((stream, handle1), move |(stream, handle)| {
            use tokio_core::reactor::Timeout;
            Timeout::new(Duration::from_secs(1), &handle).expect("creating timeout").and_then(move |()| {
                Writing::new(stream, CLOCK_PREFIX).map(move |stream| {
                    ((stream, handle), true)
                })
            })
        })
    }).map(|_| ()).map_err(|e| { println!("error from clock task: {}", e); () }));

    let f = pool.spawn(new_task(&handle, &addr, connection_id_source.clone(), 5, 3)
                       .join(new_task(&core.handle(), &addr, connection_id_source, 5, 10)));

    let x = try!(core.run(f));
    println!("x = {:?}", x);

    Ok(())
}

pub fn main() {
    run().expect("top level error");
}
<|endoftext|>


<commit_before>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![feature(old_io, os, old_path)]

extern crate claxon;

fn main() {
    use std::env;
    use std::fs;
    use std::io;
    use std::path;
    use claxon::FlacStream;
    // NOTE: args() is deprecated and should be replaced with std::env::args().nth(1).
    // However, then you get an OsString, which does not work well with old_path.
    // So we should upgrade to the new path. But the new path does not work well
    // with old File, so just wait until std::io is stable ...
    let input = fs::File::open(path::Path::new(&env::args().nth(1).unwrap())).unwrap();
    let mut reader = io::BufReader::new(input);
    let mut stream = FlacStream::new(&mut reader).unwrap();
    let n_samples = stream.streaminfo().n_samples.unwrap();
    let mut blocks = stream.blocks::<i16>();
    let mut sample = 0u64;
    let mut i = 0u64;
    while sample < n_samples {
        let block = blocks.read_next().unwrap();
        let left = block.channel(0);
        let right = block.channel(1);
        println!("block {} decoded\nleft: {:?} .. {:?}\nright: {:?} .. {:?}",
                 i, left.slice_to(12), left.slice_from(block.len() as usize - 12),
                 right.slice_to(12), right.slice_from(block.len() as usize - 12));
        sample = sample + block.len() as u64;
        i = i + 1;
    }
}
<commit_msg>update test program to new slice syntax and features<commit_after>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

extern crate claxon;

fn main() {
    use std::env;
    use std::fs;
    use std::io;
    use std::path;
    use claxon::FlacStream;
    let input = fs::File::open(path::Path::new(&env::args().nth(1).unwrap())).unwrap();
    let mut reader = io::BufReader::new(input);
    let mut stream = FlacStream::new(&mut reader).unwrap();
    let n_samples = stream.streaminfo().n_samples.unwrap();
    let mut blocks = stream.blocks::<i16>();
    let mut sample = 0u64;
    let mut i = 0u64;
    while sample < n_samples {
        let block = blocks.read_next().unwrap();
        let left = block.channel(0);
        let right = block.channel(1);
        println!("block {} decoded\nleft: {:?} .. {:?}\nright: {:?} .. {:?}",
                 i, &left[..12], &left[block.len() as usize - 12 ..],
                 &right[..12], &right[block.len() as usize - 12 ..]);
        sample = sample + block.len() as u64;
        i = i + 1;
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]
use std::fmt;
use std::rc::{Rc, Weak};
use entity::{Vec3f64, Entity};

struct Address {
    x: u64,
    y: u64,
    z: u64,
}

struct Element {
    name: String,
}

pub enum ChunkOption {
    Local(Rc<LocalChunk>),
    Remote(Rc<RemoteChunk>)
}

pub enum ChildOption {
    Local(Rc<LocalChunk>),
    Remote(Rc<RemoteChunk>),
    None
}

pub enum ParentOption {
    Local(Weak<LocalChunk>),
    Remote(Rc<RemoteChunk>),
    None
}

pub struct LocalChunk {
    address: Address,
    center: Vec3f64,
    //children: [ChildOption; 8],
    //FIXME: remove these and uncomment above
    child0: ChildOption,
    child1: ChildOption,
    child2: ChildOption,
    child3: ChildOption,
    child4: ChildOption,
    child5: ChildOption,
    child6: ChildOption,
    child7: ChildOption,
    //END FIXME
    child_number: Option<u8>,
    entities: Vec<Rc<Entity>>,
    mass: f64,
    parent: ParentOption,
    scale: u8,
    structure: u64,
    summary: Rc<Element>,
}

pub struct RemoteChunk {
    connection: i32,
}

impl Address {
    fn from_string(address: String) -> Option<Address> {
        let mut xyz: (u64, u64, u64) = (1, 1, 1);
        let mut xyz_add: (u64, u64, u64);

        for (i, c) in address.chars().enumerate() {
            if i > 63 {
                break;
            }
            xyz_add =  match c.to_digit(8) {
                Some(0) => (0, 0, 0),
                Some(1) => (0, 0, 1),
                Some(2) => (0, 1, 0),
                Some(3) => (0, 1, 1),
                Some(4) => (1, 0, 0),
                Some(5) => (1, 0, 1),
                Some(6) => (1, 1, 0),
                Some(7) => (1, 1, 1),
                _ => return None,
            };
            xyz = (xyz.0 << 1, xyz.1 << 1, xyz.2 << 1);
            xyz = (xyz.0 + xyz_add.0, xyz.1 + xyz_add.1, xyz.2 + xyz_add.2);
        }
        Some(Address {x: xyz.0, y: xyz.1, z: xyz.2})
    }
}

impl ToString for Address {
    fn to_string(&self) -> String {
        String::from_str("")
    }
}

impl PartialEq for Address {
    fn eq(&self, other: &Self) -> bool {
        (((self.x == other.x) && (self.y == other.y)) && (self.z == other.z))
    }
}

impl fmt::Show for Address {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Address({}, {}, {})", self.x, self.y, self.z)
    }
}

pub trait Chunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption;
    fn tick(&self, time_delta: f64);
}

impl Chunk for LocalChunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption {

        /* FIXME: remove
        let index = match (x, y, z) {
            (0, 0, 0) => Some(0),
            (0, 0, 1) => Some(1),
            (0, 1, 0) => Some(2),
            (0, 1, 1) => Some(3),
            (1, 0, 0) => Some(4),
            (1, 0, 1) => Some(5),
            (1, 1, 0) => Some(6),
            (1, 1, 1) => Some(7),
            _ => None,
        };

        if index == None {
            return ChildOption::None
        }

        match self.children[index] {
            ChildOption::Local(ref chunk) => ChildOption::Local(*chunk),
            ChildOption::Remote(ref chunk) => ChildOption::Remote(*chunk),
            ChildOption::None => ChildOption::None
        }*/
        ChildOption::None
    }

    fn tick(&self, time_delta: f64) {
    }
}

impl Chunk for RemoteChunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption {
        ChildOption::None
    }
    fn tick (&self, time_delta: f64) {
    }
}

#[test]
fn test_address_from_string() {
    let known = Option::Some(Address {x: 0b11, y: 0b11, z: 0b11});
    let unknown = Address::from_string(String::from_str("7"));
    assert_eq!(known, unknown)
}
<commit_msg>Use new core::fmt, moved from std::fmt<commit_after>#![allow(dead_code)]
use core::fmt;
use std::rc::{Rc, Weak};
use entity::{Vec3f64, Entity};

struct Address {
    x: u64,
    y: u64,
    z: u64,
}

struct Element {
    name: String,
}

pub enum ChunkOption {
    Local(Rc<LocalChunk>),
    Remote(Rc<RemoteChunk>)
}

pub enum ChildOption {
    Local(Rc<LocalChunk>),
    Remote(Rc<RemoteChunk>),
    None
}

pub enum ParentOption {
    Local(Weak<LocalChunk>),
    Remote(Rc<RemoteChunk>),
    None
}

pub struct LocalChunk {
    address: Address,
    center: Vec3f64,
    //children: [ChildOption; 8],
    //FIXME: remove these and uncomment above
    child0: ChildOption,
    child1: ChildOption,
    child2: ChildOption,
    child3: ChildOption,
    child4: ChildOption,
    child5: ChildOption,
    child6: ChildOption,
    child7: ChildOption,
    //END FIXME
    child_number: Option<u8>,
    entities: Vec<Rc<Entity>>,
    mass: f64,
    parent: ParentOption,
    scale: u8,
    structure: u64,
    summary: Rc<Element>,
}

pub struct RemoteChunk {
    connection: i32,
}

impl Address {
    fn from_string(address: String) -> Option<Address> {
        let mut xyz: (u64, u64, u64) = (1, 1, 1);
        let mut xyz_add: (u64, u64, u64);

        for (i, c) in address.chars().enumerate() {
            if i > 63 {
                break;
            }
            xyz_add =  match c.to_digit(8) {
                Some(0) => (0, 0, 0),
                Some(1) => (0, 0, 1),
                Some(2) => (0, 1, 0),
                Some(3) => (0, 1, 1),
                Some(4) => (1, 0, 0),
                Some(5) => (1, 0, 1),
                Some(6) => (1, 1, 0),
                Some(7) => (1, 1, 1),
                _ => return None,
            };
            xyz = (xyz.0 << 1, xyz.1 << 1, xyz.2 << 1);
            xyz = (xyz.0 + xyz_add.0, xyz.1 + xyz_add.1, xyz.2 + xyz_add.2);
        }
        Some(Address {x: xyz.0, y: xyz.1, z: xyz.2})
    }
}

impl ToString for Address {
    fn to_string(&self) -> String {
        String::from_str("")
    }
}

impl PartialEq for Address {
    fn eq(&self, other: &Self) -> bool {
        (((self.x == other.x) && (self.y == other.y)) && (self.z == other.z))
    }
}

impl fmt::Show for Address {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Address({}, {}, {})", self.x, self.y, self.z)
    }
}

pub trait Chunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption;
    fn tick(&self, time_delta: f64);
}

impl Chunk for LocalChunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption {

        /* FIXME: remove
        let index = match (x, y, z) {
            (0, 0, 0) => Some(0),
            (0, 0, 1) => Some(1),
            (0, 1, 0) => Some(2),
            (0, 1, 1) => Some(3),
            (1, 0, 0) => Some(4),
            (1, 0, 1) => Some(5),
            (1, 1, 0) => Some(6),
            (1, 1, 1) => Some(7),
            _ => None,
        };

        if index == None {
            return ChildOption::None
        }

        match self.children[index] {
            ChildOption::Local(ref chunk) => ChildOption::Local(*chunk),
            ChildOption::Remote(ref chunk) => ChildOption::Remote(*chunk),
            ChildOption::None => ChildOption::None
        }*/
        ChildOption::None
    }

    fn tick(&self, time_delta: f64) {
    }
}

impl Chunk for RemoteChunk {
    fn get_child(&self, x: u8, y: u8, z: u8) -> ChildOption {
        ChildOption::None
    }
    fn tick (&self, time_delta: f64) {
    }
}

#[test]
fn test_address_from_string() {
    let known = Option::Some(Address {x: 0b11, y: 0b11, z: 0b11});
    let unknown = Address::from_string(String::from_str("7"));
    assert_eq!(known, unknown)
}
<|endoftext|>


<commit_before>use std::{collections::BTreeMap, fmt};

use ruma_identifiers::EventId;
use serde::{
    de::{Deserializer, MapAccess, Visitor},
    ser::{SerializeMap, Serializer},
    Deserialize, Serialize,
};

pub fn serialize<S>(
    response: &BTreeMap<EventId, Result<(), String>>,
    serializer: S,
) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let mut map = serializer.serialize_map(Some(response.len()))?;
    for (key, value) in response {
        let wrapped_error = WrappedError {
            error: match value {
                Ok(_) => None,
                Err(error) => Some(error.clone()),
            },
        };
        map.serialize_entry(&key, &wrapped_error)?;
    }
    map.end()
}

pub fn deserialize<'de, D>(
    deserializer: D,
) -> Result<BTreeMap<EventId, Result<(), String>>, D::Error>
where
    D: Deserializer<'de>,
{
    deserializer.deserialize_map(PduProcessResponseVisitor)
}

#[derive(Deserialize, Serialize)]
struct WrappedError {
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

struct PduProcessResponseVisitor;

impl<'de> Visitor<'de> for PduProcessResponseVisitor {
    type Value = BTreeMap<EventId, Result<(), String>>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("A map of EventIds to a map of optional errors")
    }

    fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
    where
        M: MapAccess<'de>,
    {
        let mut map = BTreeMap::new();

        while let Some((key, value)) = access.next_entry::<EventId, WrappedError>()? {
            let v = match value.error {
                None => Ok(()),
                Some(error) => Err(error),
            };
            map.insert(key, v);
        }
        Ok(map)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use ruma_identifiers::{event_id, EventId};
    use serde_json::{json, value::Serializer as JsonSerializer};

    use super::{deserialize, serialize};

    #[test]
    fn serialize_error() {
        let mut response: BTreeMap<EventId, Result<(), String>> = BTreeMap::new();
        response.insert(event_id!("$someevent:matrix.org"), Err("Some processing error.".into()));

        let serialized = serialize(&response, JsonSerializer).unwrap();
        let json = json!({
            "$someevent:matrix.org": { "error": "Some processing error." }
        });
        assert_eq!(serialized, json);
    }

    #[test]
    fn serialize_ok() {
        let mut response: BTreeMap<EventId, Result<(), String>> = BTreeMap::new();
        response.insert(event_id!("$someevent:matrix.org"), Ok(()));

        let serialized = serialize(&response, serde_json::value::Serializer).unwrap();
        let json = json!({
            "$someevent:matrix.org": {}
        });
        assert_eq!(serialized, json);
    }

    #[test]
    fn deserialize_error() {
        let json = json!({
            "$someevent:matrix.org": { "error": "Some processing error." }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        let event_response = response.get(&event_id).unwrap().clone().unwrap_err();
        assert_eq!(event_response, "Some processing error.");
    }

    #[test]
    fn deserialize_null_error_is_ok() {
        let json = json!({
            "$someevent:matrix.org": { "error": null }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        assert!(response.get(&event_id).unwrap().is_ok());
    }

    #[test]
    fn desieralize_empty_error_is_err() {
        let json = json!({
            "$someevent:matrix.org": { "error": "" }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        let event_response = response.get(&event_id).unwrap().clone().unwrap_err();
        assert_eq!(event_response, "");
    }

    #[test]
    fn deserialize_ok() {
        let json = json!({
            "$someevent:matrix.org": {}
        });
        let response = deserialize(json).unwrap();
        assert!(response.get(&event_id!("$someevent:matrix.org")).unwrap().is_ok());
    }
}
<commit_msg>federation-api: Move Visitor declaration into deserialize function<commit_after>use std::{collections::BTreeMap, fmt};

use ruma_identifiers::EventId;
use serde::{
    de::{Deserializer, MapAccess, Visitor},
    ser::{SerializeMap, Serializer},
    Deserialize, Serialize,
};

#[derive(Deserialize, Serialize)]
struct WrappedError {
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

pub fn serialize<S>(
    response: &BTreeMap<EventId, Result<(), String>>,
    serializer: S,
) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    let mut map = serializer.serialize_map(Some(response.len()))?;
    for (key, value) in response {
        let wrapped_error = WrappedError {
            error: match value {
                Ok(_) => None,
                Err(error) => Some(error.clone()),
            },
        };
        map.serialize_entry(&key, &wrapped_error)?;
    }
    map.end()
}

pub fn deserialize<'de, D>(
    deserializer: D,
) -> Result<BTreeMap<EventId, Result<(), String>>, D::Error>
where
    D: Deserializer<'de>,
{
    struct PduProcessResponseVisitor;

    impl<'de> Visitor<'de> for PduProcessResponseVisitor {
        type Value = BTreeMap<EventId, Result<(), String>>;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("A map of EventIds to a map of optional errors")
        }

        fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
        where
            M: MapAccess<'de>,
        {
            let mut map = BTreeMap::new();

            while let Some((key, value)) = access.next_entry::<EventId, WrappedError>()? {
                let v = match value.error {
                    None => Ok(()),
                    Some(error) => Err(error),
                };
                map.insert(key, v);
            }
            Ok(map)
        }
    }

    deserializer.deserialize_map(PduProcessResponseVisitor)
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use ruma_identifiers::{event_id, EventId};
    use serde_json::{json, value::Serializer as JsonSerializer};

    use super::{deserialize, serialize};

    #[test]
    fn serialize_error() {
        let mut response: BTreeMap<EventId, Result<(), String>> = BTreeMap::new();
        response.insert(event_id!("$someevent:matrix.org"), Err("Some processing error.".into()));

        let serialized = serialize(&response, JsonSerializer).unwrap();
        let json = json!({
            "$someevent:matrix.org": { "error": "Some processing error." }
        });
        assert_eq!(serialized, json);
    }

    #[test]
    fn serialize_ok() {
        let mut response: BTreeMap<EventId, Result<(), String>> = BTreeMap::new();
        response.insert(event_id!("$someevent:matrix.org"), Ok(()));

        let serialized = serialize(&response, serde_json::value::Serializer).unwrap();
        let json = json!({
            "$someevent:matrix.org": {}
        });
        assert_eq!(serialized, json);
    }

    #[test]
    fn deserialize_error() {
        let json = json!({
            "$someevent:matrix.org": { "error": "Some processing error." }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        let event_response = response.get(&event_id).unwrap().clone().unwrap_err();
        assert_eq!(event_response, "Some processing error.");
    }

    #[test]
    fn deserialize_null_error_is_ok() {
        let json = json!({
            "$someevent:matrix.org": { "error": null }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        assert!(response.get(&event_id).unwrap().is_ok());
    }

    #[test]
    fn desieralize_empty_error_is_err() {
        let json = json!({
            "$someevent:matrix.org": { "error": "" }
        });

        let response = deserialize(json).unwrap();
        let event_id = event_id!("$someevent:matrix.org");

        let event_response = response.get(&event_id).unwrap().clone().unwrap_err();
        assert_eq!(event_response, "");
    }

    #[test]
    fn deserialize_ok() {
        let json = json!({
            "$someevent:matrix.org": {}
        });
        let response = deserialize(json).unwrap();
        assert!(response.get(&event_id!("$someevent:matrix.org")).unwrap().is_ok());
    }
}
<|endoftext|>