<commit_before>//! Defines the trait that must be implemented by ISPC task execution systems
//! and provides a default threaded one for use.

use libc;
use num_cpus;

use std::mem;
use std::time::Duration;
use std::cell::RefCell;
use std::sync::{Arc, RwLock, Mutex};
use std::sync::atomic::{self, AtomicUsize};
use std::thread::{self, JoinHandle};

use task::{ISPCTaskFn, Context};

/// Trait to be implemented to provide ISPC task execution functionality.
///
/// The runtime [required functions](http://ispc.github.io/ispc.html#task-parallelism-runtime-requirements)
/// will be forwarded directly to your struct, making this interface unsafe.
pub trait TaskSystem {
    /// Alloc is called when memory must be allocated to store parameters to pass to a task
    /// and must return a pointer to an allocation of `size` bytes aligned to `align`.
    ///
    /// The `handle_ptr` will be `NULL` if this is the first time launch has been called in
    /// the function or is the first launch call after an explicit `sync` statement. Both
    /// situations should be treated equivalently as creating a new exeuction context for tasks.
    /// The `handle_ptr` should be set to some context tracking facility so that you can later
    /// track task groups launched in the context and perform finer grained synchronization in
    /// `sync`.
    unsafe fn alloc(&self, handle_ptr: *mut *mut libc::c_void, size: i64, align: i32) -> *mut libc::c_void;
    /// Launch is called when a new group of tasks is being launched and should schedule them to
    /// be executed in some way.
    ///
    /// The `handle_ptr` will point to the same handle you set up in `alloc` and can be used to
    /// associate groups of tasks with a context of execution as mentioned before. The function `f`
    /// should be executed `count0 * count1 * count2` times and indices passed to the function
    /// should be as if running in a nested for loop:
    ///
    /// ```ignore
    /// let total_tasks = count0 * count1 * count2;
    /// for z in 0..count2 {
    ///     for y in 0..count1 {
    ///         for x in 0..count0 {
    ///             let task_id = x + y * count0 + z * count0 * count1;
    ///             f(data, thread_id, total_threads, task_id, total_tasks,
    ///               x, y, z, count0, count1, count2);
    ///         }
    ///     }
    /// }
    /// ```
    ///
    /// The `data` pointer points to the ISPC task specific parameter pointer and should be passed
    /// through to the function.
    unsafe fn launch(&self, handle_ptr: *mut *mut libc::c_void, f: ISPCTaskFn, data: *mut libc::c_void,
                     count0: i32, count1: i32, count2: i32);
    /// Synchronize an execution context with the tasks it's launched. Use `handle` to determine
    /// the task context being synchronized.
    ///
    /// This function should not return until all tasks launched within the context being
    /// synchronized with have been completed. You can use the `handle` to determine which context
    /// is being synchronized with and thus which tasks must be completed before returning.
    unsafe fn sync(&self, handle: *mut libc::c_void);
    /// Return a context that has remaining tasks left to be executed by a thread, returns None
    /// if no contexts have remaining tasks.
    ///
    /// Note that due to threading issues you shouldn't assume the context returned actually has
    /// outstanding tasks by the time it's returned to the caller and a chunk is requested.
    fn get_context(&self) -> Option<Arc<Context>>;
}

// Thread local storage to store the thread's id, otherwise we don't know
// who we are in sync. The thread id starts at an invalid value but will be set
// upon thread launch
thread_local!(static THREAD_ID: RefCell<usize> = RefCell::new(0));

/// A multithreaded execution environment for the tasks.
pub struct Parallel {
    context_list: RwLock<Vec<Arc<Context>>>,
    next_context_id: AtomicUsize,
    threads: Mutex<Vec<JoinHandle<()>>>,
    chunk_size: usize,
}

impl Parallel {
    /// Create a parallel task execution environment that will use `num_cpus` threads
    /// to run tasks
    pub fn new() -> Parallel {
        let par = Parallel { context_list: RwLock::new(Vec::new()),
                             next_context_id: AtomicUsize::new(0),
                             threads: Mutex::new(Vec::new()),
                             chunk_size: 16 };
        {
            let mut threads = par.threads.lock().unwrap();
            let num_threads = num_cpus::get();
            let chunk_size = par.chunk_size;
            //println!("Launching {} threads", num_threads);
            for i in 0..num_threads {
                // Note that the spawned thread ids start at 1 since the main thread is 0
                threads.push(thread::spawn(move || worker_thread(i + 1, num_threads + 1, chunk_size)));
            }
        }
        par
    }
    /// Create an oversubscribued parallel task execution environment that will use
    /// `oversubscribe * num_cpus` threads to run tasks.
    pub fn oversubscribed(oversubscribe: f32) -> Parallel {
        assert!(oversubscribe >= 1.0);
        let par = Parallel { context_list: RwLock::new(Vec::new()),
                             next_context_id: AtomicUsize::new(0),
                             threads: Mutex::new(Vec::new()),
                             chunk_size: 16 };
        {
            let mut threads = par.threads.lock().unwrap();
            let num_threads = (oversubscribe * num_cpus::get() as f32) as usize;
            let chunk_size = par.chunk_size;
            //println!("Launching {} threads", num_threads);
            for i in 0..num_threads {
                // Note that the spawned thread ids start at 1 since the main thread is 0
                threads.push(thread::spawn(move || worker_thread(i + 1, num_threads + 1, chunk_size)));
            }
            println!("have {} threads", num_threads);
        }
        par
    }
}

impl TaskSystem for Parallel {
    unsafe fn alloc(&self, handle_ptr: *mut *mut libc::c_void, size: i64, align: i32) -> *mut libc::c_void {
        //println!("ISPCAlloc, size: {}, align: {}", size, align);
        // If the handle is null this is the first time this function has spawned tasks
        // and we should create a new Context structure in the TASK_LIST for it, otherwise
        // it's the pointer to where we should append the new Group
        if (*handle_ptr).is_null() {
            let mut context_list = self.context_list.write().unwrap();
            //println!("handle ptr is null");
            // This is a bit hairy. We allocate the new task context in a box, then
            // unbox it into a raw ptr to get a ptr we can pass back to ISPC through
            // the handle_ptr and then re-box it into our TASK_LIST so it will
            // be free'd properly when we erase it from the vector in ISPCSync
            let c = Arc::new(Context::new(self.next_context_id.fetch_add(1, atomic::Ordering::SeqCst)));
            {
                let h = &*c;
                *handle_ptr = mem::transmute(h);
            }
            context_list.push(c);
            let ctx = context_list.last().unwrap();
            //println!("context.id = {}", ctx.id);
            ctx.alloc(size as usize, align as usize)
        } else {
            let context_list = self.context_list.read().unwrap();
            //println!("handle ptr is not null");
            let handle_ctx: *mut Context = mem::transmute(*handle_ptr);
            let ctx = context_list.iter().find(|c| (*handle_ctx).id == c.id).unwrap();
            //println!("context.id = {}", ctx.id);
            ctx.alloc(size as usize, align as usize)
        }
    }
    unsafe fn launch(&self, handle_ptr: *mut *mut libc::c_void, f: ISPCTaskFn, data: *mut libc::c_void,
                     count0: i32, count1: i32, count2: i32) {
        // Push the tasks being launched on to the list of task groups for this function
        let context: &mut Context = mem::transmute(*handle_ptr);
        //println!("ISPCLaunch, context.id = {}, counts: [{}, {}, {}]", context.id, count0, count1, count2);
        context.launch((count0, count1, count2), data, f);
        // Unpark any sleeping threads since we have jobs for them
        let threads = self.threads.lock().unwrap();
        for t in threads.iter() {
            t.thread().unpark();
        }
    }
    unsafe fn sync(&self, handle: *mut libc::c_void) {
        let context: &mut Context = mem::transmute(handle);
        let thread = THREAD_ID.with(|f| *f.borrow());
        let total_threads = num_cpus::get();
        //println!("Thread {} is syncing", thread);
        // Make sure all tasks are done, and execute them if not for this simple
        // serial version. TODO: In the future we'd wait on each Group's semaphore or atomic bool
        // Maybe the waiting thread could help execute tasks as well, otherwise it might be
        // possible to deadlock, where all threads are waiting for some enqueue'd tasks but no
        // threads are available to run them. Just running tasks in our context is not sufficient
        // to prevent deadlock actually, because those tasks could in turn launch & sync and get stuck
        // so if our tasks aren't done and there's none left to run in our context we should start
        // running tasks from other contexts to help out
        //println!("ISPCSync, context.id = {}", context.id);
        for tg in context.iter() {
            for chunk in tg.chunks(self.chunk_size) {
                //println!("syncing thread {} running local chunk {:?}", thread, chunk);
                // TODO: We need to figure out which thread we are
                chunk.execute(thread as i32, total_threads as i32);
            }
        }
        // If all the tasks for this context have been finished we're done sync'ing and can
        // clean up memory and remove the context from the TASK_LIST. Otherwise there are some
        // unfinished groups further down the the tree that were spawned by our direct tasks that
        // those are now sync'ing on and we need to help out. However since we don't know the tree
        // our best option is to just start grabbing chunks from unfinished groups in the TASK_LIST
        // and running them to at least ensure global forward progress, which will eventually get
        // the stuff we're waiting on to finish. After each chunk execution we should check if
        // our sync'ing context is done and break
        while !context.current_tasks_done() {
            //println!("Not all tasks for context {} are done, thread {} helping out!", context.id, thread);
            // Get a task group to run
            while let Some(c) = self.get_context() {
                //println!("syncing thread {} got a context {:?}", thread, c);
                let mut ran_some = false;
                for tg in c.iter() {
                    //println!("syncing thread {} looking at task group {:?}", thread, tg);
                    for chunk in tg.chunks(self.chunk_size) {
                        ran_some = true;
                        //println!("syncing thread {} running nonlocal chunk {:?}", thread, chunk);
                        chunk.execute(thread as i32, total_threads as i32);
                    }
                }
                if !ran_some {
                    //println!("Syncing thread ran nothing, waiting a bit");
                    thread::sleep(Duration::from_millis(50));
                }
            }
        }
        // Now erase this context from our vector
        let mut context_list = self.context_list.write().unwrap();
        let pos = context_list.iter().position(|c| context.id == c.id).unwrap();
        context_list.remove(pos);
    }
    fn get_context(&self) -> Option<Arc<Context>> {
        self.context_list.read().unwrap().iter()
            .find(|c| !c.current_tasks_done()).map(|c| c.clone())
    }
}

fn worker_thread(thread: usize, total_threads: usize, chunk_size: usize) {
    THREAD_ID.with(|f| *f.borrow_mut() = thread);
    let task_sys = ::get_task_system();
    //println!("thread {} is spawned", thread);
    loop {
        // Get a task group to run
        while let Some(c) = task_sys.get_context() {
            for tg in c.iter() {
                //println!("thread checking task group {:?}", tg);
                for chunk in tg.chunks(chunk_size) {
                    //println!("thread {} running chunk {:?}", thread, chunk);
                    chunk.execute(thread as i32, total_threads as i32);
                }
            }
        }
        // We ran out of contexts to get, so wait a bit for a new group to get
        // launched
        thread::park();
        //println!("thread {} was woken up to work!", thread);
    }
}

<commit_msg>Using a slightly finer chunk size<commit_after>//! Defines the trait that must be implemented by ISPC task execution systems
//! and provides a default threaded one for use.

use libc;
use num_cpus;

use std::mem;
use std::time::Duration;
use std::cell::RefCell;
use std::sync::{Arc, RwLock, Mutex};
use std::sync::atomic::{self, AtomicUsize};
use std::thread::{self, JoinHandle};

use task::{ISPCTaskFn, Context};

/// Trait to be implemented to provide ISPC task execution functionality.
///
/// The runtime [required functions](http://ispc.github.io/ispc.html#task-parallelism-runtime-requirements)
/// will be forwarded directly to your struct, making this interface unsafe.
pub trait TaskSystem {
    /// Alloc is called when memory must be allocated to store parameters to pass to a task
    /// and must return a pointer to an allocation of `size` bytes aligned to `align`.
    ///
    /// The `handle_ptr` will be `NULL` if this is the first time launch has been called in
    /// the function or is the first launch call after an explicit `sync` statement. Both
    /// situations should be treated equivalently as creating a new exeuction context for tasks.
    /// The `handle_ptr` should be set to some context tracking facility so that you can later
    /// track task groups launched in the context and perform finer grained synchronization in
    /// `sync`.
    unsafe fn alloc(&self, handle_ptr: *mut *mut libc::c_void, size: i64, align: i32) -> *mut libc::c_void;
    /// Launch is called when a new group of tasks is being launched and should schedule them to
    /// be executed in some way.
    ///
    /// The `handle_ptr` will point to the same handle you set up in `alloc` and can be used to
    /// associate groups of tasks with a context of execution as mentioned before. The function `f`
    /// should be executed `count0 * count1 * count2` times and indices passed to the function
    /// should be as if running in a nested for loop:
    ///
    /// ```ignore
    /// let total_tasks = count0 * count1 * count2;
    /// for z in 0..count2 {
    ///     for y in 0..count1 {
    ///         for x in 0..count0 {
    ///             let task_id = x + y * count0 + z * count0 * count1;
    ///             f(data, thread_id, total_threads, task_id, total_tasks,
    ///               x, y, z, count0, count1, count2);
    ///         }
    ///     }
    /// }
    /// ```
    ///
    /// The `data` pointer points to the ISPC task specific parameter pointer and should be passed
    /// through to the function.
    unsafe fn launch(&self, handle_ptr: *mut *mut libc::c_void, f: ISPCTaskFn, data: *mut libc::c_void,
                     count0: i32, count1: i32, count2: i32);
    /// Synchronize an execution context with the tasks it's launched. Use `handle` to determine
    /// the task context being synchronized.
    ///
    /// This function should not return until all tasks launched within the context being
    /// synchronized with have been completed. You can use the `handle` to determine which context
    /// is being synchronized with and thus which tasks must be completed before returning.
    unsafe fn sync(&self, handle: *mut libc::c_void);
    /// Return a context that has remaining tasks left to be executed by a thread, returns None
    /// if no contexts have remaining tasks.
    ///
    /// Note that due to threading issues you shouldn't assume the context returned actually has
    /// outstanding tasks by the time it's returned to the caller and a chunk is requested.
    fn get_context(&self) -> Option<Arc<Context>>;
}

// Thread local storage to store the thread's id, otherwise we don't know
// who we are in sync. The thread id starts at an invalid value but will be set
// upon thread launch
thread_local!(static THREAD_ID: RefCell<usize> = RefCell::new(0));

/// A multithreaded execution environment for the tasks.
pub struct Parallel {
    context_list: RwLock<Vec<Arc<Context>>>,
    next_context_id: AtomicUsize,
    threads: Mutex<Vec<JoinHandle<()>>>,
    chunk_size: usize,
}

impl Parallel {
    /// Create a parallel task execution environment that will use `num_cpus` threads
    /// to run tasks
    pub fn new() -> Parallel {
        let par = Parallel { context_list: RwLock::new(Vec::new()),
                             next_context_id: AtomicUsize::new(0),
                             threads: Mutex::new(Vec::new()),
                             chunk_size: 8 };
        {
            let mut threads = par.threads.lock().unwrap();
            let num_threads = num_cpus::get();
            let chunk_size = par.chunk_size;
            //println!("Launching {} threads", num_threads);
            for i in 0..num_threads {
                // Note that the spawned thread ids start at 1 since the main thread is 0
                threads.push(thread::spawn(move || worker_thread(i + 1, num_threads + 1, chunk_size)));
            }
        }
        par
    }
    /// Create an oversubscribued parallel task execution environment that will use
    /// `oversubscribe * num_cpus` threads to run tasks.
    pub fn oversubscribed(oversubscribe: f32) -> Parallel {
        assert!(oversubscribe >= 1.0);
        let par = Parallel { context_list: RwLock::new(Vec::new()),
                             next_context_id: AtomicUsize::new(0),
                             threads: Mutex::new(Vec::new()),
                             chunk_size: 8 };
        {
            let mut threads = par.threads.lock().unwrap();
            let num_threads = (oversubscribe * num_cpus::get() as f32) as usize;
            let chunk_size = par.chunk_size;
            //println!("Launching {} threads", num_threads);
            for i in 0..num_threads {
                // Note that the spawned thread ids start at 1 since the main thread is 0
                threads.push(thread::spawn(move || worker_thread(i + 1, num_threads + 1, chunk_size)));
            }
            println!("have {} threads", num_threads);
        }
        par
    }
}

impl TaskSystem for Parallel {
    unsafe fn alloc(&self, handle_ptr: *mut *mut libc::c_void, size: i64, align: i32) -> *mut libc::c_void {
        //println!("ISPCAlloc, size: {}, align: {}", size, align);
        // If the handle is null this is the first time this function has spawned tasks
        // and we should create a new Context structure in the TASK_LIST for it, otherwise
        // it's the pointer to where we should append the new Group
        if (*handle_ptr).is_null() {
            let mut context_list = self.context_list.write().unwrap();
            //println!("handle ptr is null");
            // This is a bit hairy. We allocate the new task context in a box, then
            // unbox it into a raw ptr to get a ptr we can pass back to ISPC through
            // the handle_ptr and then re-box it into our TASK_LIST so it will
            // be free'd properly when we erase it from the vector in ISPCSync
            let c = Arc::new(Context::new(self.next_context_id.fetch_add(1, atomic::Ordering::SeqCst)));
            {
                let h = &*c;
                *handle_ptr = mem::transmute(h);
            }
            context_list.push(c);
            let ctx = context_list.last().unwrap();
            //println!("context.id = {}", ctx.id);
            ctx.alloc(size as usize, align as usize)
        } else {
            let context_list = self.context_list.read().unwrap();
            //println!("handle ptr is not null");
            let handle_ctx: *mut Context = mem::transmute(*handle_ptr);
            let ctx = context_list.iter().find(|c| (*handle_ctx).id == c.id).unwrap();
            //println!("context.id = {}", ctx.id);
            ctx.alloc(size as usize, align as usize)
        }
    }
    unsafe fn launch(&self, handle_ptr: *mut *mut libc::c_void, f: ISPCTaskFn, data: *mut libc::c_void,
                     count0: i32, count1: i32, count2: i32) {
        // Push the tasks being launched on to the list of task groups for this function
        let context: &mut Context = mem::transmute(*handle_ptr);
        //println!("ISPCLaunch, context.id = {}, counts: [{}, {}, {}]", context.id, count0, count1, count2);
        context.launch((count0, count1, count2), data, f);
        // Unpark any sleeping threads since we have jobs for them
        let threads = self.threads.lock().unwrap();
        for t in threads.iter() {
            t.thread().unpark();
        }
    }
    unsafe fn sync(&self, handle: *mut libc::c_void) {
        let context: &mut Context = mem::transmute(handle);
        let thread = THREAD_ID.with(|f| *f.borrow());
        let total_threads = num_cpus::get();
        //println!("Thread {} is syncing", thread);
        // Make sure all tasks are done, and execute them if not for this simple
        // serial version. TODO: In the future we'd wait on each Group's semaphore or atomic bool
        // Maybe the waiting thread could help execute tasks as well, otherwise it might be
        // possible to deadlock, where all threads are waiting for some enqueue'd tasks but no
        // threads are available to run them. Just running tasks in our context is not sufficient
        // to prevent deadlock actually, because those tasks could in turn launch & sync and get stuck
        // so if our tasks aren't done and there's none left to run in our context we should start
        // running tasks from other contexts to help out
        //println!("ISPCSync, context.id = {}", context.id);
        for tg in context.iter() {
            for chunk in tg.chunks(self.chunk_size) {
                //println!("syncing thread {} running local chunk {:?}", thread, chunk);
                // TODO: We need to figure out which thread we are
                chunk.execute(thread as i32, total_threads as i32);
            }
        }
        // If all the tasks for this context have been finished we're done sync'ing and can
        // clean up memory and remove the context from the TASK_LIST. Otherwise there are some
        // unfinished groups further down the the tree that were spawned by our direct tasks that
        // those are now sync'ing on and we need to help out. However since we don't know the tree
        // our best option is to just start grabbing chunks from unfinished groups in the TASK_LIST
        // and running them to at least ensure global forward progress, which will eventually get
        // the stuff we're waiting on to finish. After each chunk execution we should check if
        // our sync'ing context is done and break
        while !context.current_tasks_done() {
            //println!("Not all tasks for context {} are done, thread {} helping out!", context.id, thread);
            // Get a task group to run
            while let Some(c) = self.get_context() {
                //println!("syncing thread {} got a context {:?}", thread, c);
                let mut ran_some = false;
                for tg in c.iter() {
                    //println!("syncing thread {} looking at task group {:?}", thread, tg);
                    for chunk in tg.chunks(self.chunk_size) {
                        ran_some = true;
                        //println!("syncing thread {} running nonlocal chunk {:?}", thread, chunk);
                        chunk.execute(thread as i32, total_threads as i32);
                    }
                }
                if !ran_some {
                    //println!("Syncing thread ran nothing, waiting a bit");
                    thread::sleep(Duration::from_millis(50));
                }
            }
        }
        // Now erase this context from our vector
        let mut context_list = self.context_list.write().unwrap();
        let pos = context_list.iter().position(|c| context.id == c.id).unwrap();
        context_list.remove(pos);
    }
    fn get_context(&self) -> Option<Arc<Context>> {
        self.context_list.read().unwrap().iter()
            .find(|c| !c.current_tasks_done()).map(|c| c.clone())
    }
}

fn worker_thread(thread: usize, total_threads: usize, chunk_size: usize) {
    THREAD_ID.with(|f| *f.borrow_mut() = thread);
    let task_sys = ::get_task_system();
    //println!("thread {} is spawned", thread);
    loop {
        // Get a task group to run
        while let Some(c) = task_sys.get_context() {
            for tg in c.iter() {
                //println!("thread checking task group {:?}", tg);
                for chunk in tg.chunks(chunk_size) {
                    //println!("thread {} running chunk {:?}", thread, chunk);
                    chunk.execute(thread as i32, total_threads as i32);
                }
            }
        }
        // We ran out of contexts to get, so wait a bit for a new group to get
        // launched
        thread::park();
        //println!("thread {} was woken up to work!", thread);
    }
}

<|endoftext|>


<commit_before>use std::sync::mpsc::{Receiver, Sender, channel};
use std::iter::repeat;
use eval::DataMode;

use ast::Value;
use eval::{ self, Expr };

pub trait PrimitiveStep {
    fn display(&self) -> String;
    fn body<'a>(&'a self) -> Option<&'a Step> { None }
    fn exec(&self);
}

#[derive(Debug)]
pub struct Message {
    pub components: Vec<Expr>
}

impl Message {
    pub fn eval_down(&self, state: &eval::State) -> Vec<Value> {
        self.components.iter().map(|e| {
            if e.exists_down() {
                e.eval_down(state)
            } else {
                Value::Number(0.0)
            }
        }).collect()
    }

    pub fn eval_up(&self, state: &mut eval::State, values: Vec<Value>) -> bool {
        debug!("eval_up: {:?}", self);
        let mut matched = true;
        let mut vi = values.into_iter().fuse();
        for e in &self.components {
            let v = vi.next().unwrap_or(Value::Number(0.0));
            matched &= e.eval_up(state, v);
        }
        matched
    }
}

#[derive(Debug)]
pub enum Step {
    Nop,
    Token(Message),
    TokenTop(Message, Box<Step>),
    Seq(Vec<Step>),
    Repeat(Expr, Box<Step>, bool),
    //PrimitiveStep(Box<PrimitiveStep>),
}

impl Step {
    /// Returns true if any data is sent in the up direction in this step or its lexical children
    pub fn any_up(&self) -> bool {
        match *self {
            Step::Nop => false,
            Step::Token(ref m) => m.components.iter().any(Expr::exists_up),
            Step::Seq(ref steps) => steps.iter().any(Step::any_up),
            Step::TokenTop(_, box ref _c) => true, //TODO: works on simple cases, but is this the right heuristic?
            Step::Repeat(_, box ref c, _) => c.any_up(),
        }
    }
}

fn first(s: &Step) -> Option<&Step> {
    match *s {
        Step::Nop => None,
        Step::Token(..) => Some(s),
        Step::TokenTop(..) => Some(s),
        Step::Seq(ref steps) => steps.get(0).and_then(first),
        Step::Repeat(_, box ref inner, _) => first(inner),
    }
}

pub fn print_step_tree(s: &Step, indent: u32) {
    let i: String = repeat(" ").take(indent as usize).collect();
    match *s {
        Step::Nop => println!("{}NOP", i),
        Step::Token(ref message) => {
            println!("{}Token: {:?}", i, message);
        }
        Step::TokenTop(ref message, box ref body) => {
            println!("{}Up: {:?}", i, message);
            print_step_tree(body, indent+1);
        }
        Step::Seq(ref steps) => {
            println!("{}Seq", i);
            for c in steps.iter() {
                print_step_tree(c, indent+1);
            }
        }
        Step::Repeat(ref count, box ref inner, up) => {
            println!("{}Repeat: {:?} {}", i, count, up);
            print_step_tree(inner, indent + 1);
        }
        /*PrimitiveStep(ref h) => {
            println!("{}{}", i, h.display());
            h.body().map(|body| {
                print_step_tree(body, indent+1)
            });
        }*/
    }
}

pub struct Connection {
    rx: Option<Receiver<Vec<Value>>>,
    tx: Option<Sender<Vec<Value>>>,

    lookahead_tx: Option<Vec<Value>>,
    lookahead_rx: Option<Vec<Value>>,

    pub alive: bool,
}

impl Connection {
    pub fn new(direction: DataMode) -> (Connection, Connection) {
        debug!("New connection: {:?}", direction);
        let (s1, r1) = if direction.down {
            let (a, b) = channel();
            (Some(a), Some(b))
        } else { (None, None) };

        let (s2, r2) = if direction.up {
            let (a, b) = channel();
            (Some(a), Some(b))
        } else { (None, None) };

        let alive = direction.down || direction.up;

        (Connection{ tx: s1, rx: r2, lookahead_tx: None, lookahead_rx: None, alive: alive },
         Connection{ tx: s2, rx: r1, lookahead_tx: None, lookahead_rx: None, alive: alive })
    }

    pub fn send(&mut self, v: Vec<Value>) -> Result<(), ()> {
        if let Some(ref tx) = self.tx {
            let r = tx.send(v).map_err(|_| ());
            if r.is_err() {
                self.alive = false;
            }
            r
        } else {
            //assert!(v.len() == 0);
            if self.alive { Ok(()) } else { Err(()) }
        }
    }

    pub fn recv(&mut self) -> Result<Vec<Value>, ()> {
        if let Some(ref rx) = self.rx {
            let r = rx.recv().map_err(|_| ());
            if r.is_err() {
                self.alive = false;
            }
            r
        } else {
            if self.alive { Ok(vec![]) } else { Err(()) }
        }
    }

    pub fn lookahead_send(&mut self, v: Vec<Value>) {
        match self.lookahead_tx {
            Some(ref lv) => {
                if v != *lv { panic!("Committed {:?}, but sending {:?}", lv, v); }
            }
            None => {
                match self.send(v.clone()) {
                    Ok(..) => {
                        self.lookahead_tx = Some(v);
                    },
                    Err(..) => (),
                }
            }
        }
    }

    pub fn lookahead_receive(&mut self) -> Option<Vec<Value>> {
        match self.lookahead_rx {
            Some(ref lv) => Some(lv.clone()),
            None => {
                match self.recv() {
                    Ok(r) => {
                        debug!("recv: {:?}", r);
                        self.lookahead_rx = Some(r.clone());
                        Some(r)
                    },
                    Err(..) => None,
                }
            }
        }
    }

    pub fn accept(&mut self) {
        self.lookahead_rx.take();
        self.lookahead_tx.take();
    }

    pub fn can_tx(&self) -> bool { self.tx.is_some() }
    pub fn can_rx(&self) -> bool { self.rx.is_some() }
}

pub fn try_token(state: &mut eval::State, parent: &mut Connection, msg: &Message) -> bool {
    debug!("tokenstep {:?}", msg);
    let m = msg.eval_down(state);

    debug!("  down: {:?}", m);
    parent.lookahead_send(m);

    match parent.lookahead_receive() {
        Some(m) => {
            debug!("  up: {:?}", m);
            msg.eval_up(state, m)
        }
        None => false
    }
}

pub fn exec(state: &mut eval::State, s: &Step, parent: &mut Connection, child: &mut Connection) -> bool {
    match *s {
        Step::Nop => true,
        Step::Token(ref msg) => {
            let r = try_token(state, parent, msg);
            if r {
                debug!("  matched");
                parent.accept();
            }
            r
        }
        Step::TokenTop(ref msg, box ref body) => {
            debug!("tokentop: {:?}", msg);
            match child.lookahead_receive() {
                Some(m) => {
                    debug!("down: {:?}", m);
                    child.accept();

                    msg.eval_up(state, m);
                    let r = exec(state, body, parent, child);
                    let m = msg.eval_down(state);

                    debug!("up: {:?}", m);
                    if child.send(m).is_err() { return false; }
                    r
                }
                None(..) => false
            }
        }
        Step::Seq(ref steps) => {
            for c in steps.iter() {
                match exec(state, c, parent, child) {
                    true => (),
                    false => return false,
                }
            }
            true
        }
        Step::Repeat(ref count_expr, box ref inner, up) => {
            if up {
                let f = first(inner).expect("Loop has no body");
                let mut count = 0;
                loop {
                    debug!("Loop up check {:?}", f);
                    match *f {
                        Step::Token(ref msg) => {
                            if !try_token(state, parent, msg) {
                                debug!("  loop exit");
                                break;
                            }
                        }
                        Step::TokenTop(..) => {
                            child.lookahead_receive();
                            if !parent.alive || !child.alive {
                                debug!("  loop done");
                                break;
                            }
                        },
                        _ => panic!(),
                    }
                    if !exec(state, inner, parent, child) {
                        debug!("loop up early exit");
                        return false;
                    }
                    count += 1;
                }
                debug!("loop up exit {}", count);
                count_expr.eval_up(state, Value::Integer(count))
            } else {
                let count = match count_expr.eval_down(state) {
                    Value::Integer(count) => count,
                    other => panic!("Count must be an integer, found {}", other),
                };

                for _ in 0..count {
                    if !exec(state, inner, parent, child) {
                        debug!("loop down early exit");
                        return false;
                    }
                }

                debug!("loop down exit {}", count);
                true
            }
        }
        //PrimitiveStep(..) => unimplemented!()
    }
}
<commit_msg>exec: Improve the `first` heuristic to look within TokenTop if possible<commit_after>use std::sync::mpsc::{Receiver, Sender, channel};
use std::iter::repeat;
use eval::DataMode;

use ast::Value;
use eval::{ self, Expr };

pub trait PrimitiveStep {
    fn display(&self) -> String;
    fn body<'a>(&'a self) -> Option<&'a Step> { None }
    fn exec(&self);
}

#[derive(Debug)]
pub struct Message {
    pub components: Vec<Expr>
}

impl Message {
    pub fn eval_down(&self, state: &eval::State) -> Vec<Value> {
        self.components.iter().map(|e| {
            if e.exists_down() {
                e.eval_down(state)
            } else {
                Value::Number(0.0)
            }
        }).collect()
    }

    pub fn eval_up(&self, state: &mut eval::State, values: Vec<Value>) -> bool {
        debug!("eval_up: {:?}", self);
        let mut matched = true;
        let mut vi = values.into_iter().fuse();
        for e in &self.components {
            let v = vi.next().unwrap_or(Value::Number(0.0));
            matched &= e.eval_up(state, v);
        }
        matched
    }
}

#[derive(Debug)]
pub enum Step {
    Nop,
    Token(Message),
    TokenTop(Message, Box<Step>),
    Seq(Vec<Step>),
    Repeat(Expr, Box<Step>, bool),
    //PrimitiveStep(Box<PrimitiveStep>),
}

impl Step {
    /// Returns true if any data is sent in the up direction in this step or its lexical children
    pub fn any_up(&self) -> bool {
        match *self {
            Step::Nop => false,
            Step::Token(ref m) => m.components.iter().any(Expr::exists_up),
            Step::Seq(ref steps) => steps.iter().any(Step::any_up),
            Step::TokenTop(_, box ref _c) => true, //TODO: works on simple cases, but is this the right heuristic?
            Step::Repeat(_, box ref c, _) => c.any_up(),
        }
    }
}

pub fn print_step_tree(s: &Step, indent: u32) {
    let i: String = repeat(" ").take(indent as usize).collect();
    match *s {
        Step::Nop => println!("{}NOP", i),
        Step::Token(ref message) => {
            println!("{}Token: {:?}", i, message);
        }
        Step::TokenTop(ref message, box ref body) => {
            println!("{}Up: {:?}", i, message);
            print_step_tree(body, indent+1);
        }
        Step::Seq(ref steps) => {
            println!("{}Seq", i);
            for c in steps.iter() {
                print_step_tree(c, indent+1);
            }
        }
        Step::Repeat(ref count, box ref inner, up) => {
            println!("{}Repeat: {:?} {}", i, count, up);
            print_step_tree(inner, indent + 1);
        }
        /*PrimitiveStep(ref h) => {
            println!("{}{}", i, h.display());
            h.body().map(|body| {
                print_step_tree(body, indent+1)
            });
        }*/
    }
}

pub struct Connection {
    rx: Option<Receiver<Vec<Value>>>,
    tx: Option<Sender<Vec<Value>>>,

    lookahead_tx: Option<Vec<Value>>,
    lookahead_rx: Option<Vec<Value>>,

    pub alive: bool,
}

impl Connection {
    pub fn new(direction: DataMode) -> (Connection, Connection) {
        debug!("New connection: {:?}", direction);
        let (s1, r1) = if direction.down {
            let (a, b) = channel();
            (Some(a), Some(b))
        } else { (None, None) };

        let (s2, r2) = if direction.up {
            let (a, b) = channel();
            (Some(a), Some(b))
        } else { (None, None) };

        let alive = direction.down || direction.up;

        (Connection{ tx: s1, rx: r2, lookahead_tx: None, lookahead_rx: None, alive: alive },
         Connection{ tx: s2, rx: r1, lookahead_tx: None, lookahead_rx: None, alive: alive })
    }

    pub fn send(&mut self, v: Vec<Value>) -> Result<(), ()> {
        if let Some(ref tx) = self.tx {
            let r = tx.send(v).map_err(|_| ());
            if r.is_err() {
                self.alive = false;
            }
            r
        } else {
            //assert!(v.len() == 0);
            if self.alive { Ok(()) } else { Err(()) }
        }
    }

    pub fn recv(&mut self) -> Result<Vec<Value>, ()> {
        if let Some(ref rx) = self.rx {
            let r = rx.recv().map_err(|_| ());
            if r.is_err() {
                self.alive = false;
            }
            r
        } else {
            if self.alive { Ok(vec![]) } else { Err(()) }
        }
    }

    pub fn lookahead_send(&mut self, v: Vec<Value>) {
        match self.lookahead_tx {
            Some(ref lv) => {
                if v != *lv { panic!("Committed {:?}, but sending {:?}", lv, v); }
            }
            None => {
                match self.send(v.clone()) {
                    Ok(..) => {
                        self.lookahead_tx = Some(v);
                    },
                    Err(..) => (),
                }
            }
        }
    }

    pub fn lookahead_receive(&mut self) -> Option<Vec<Value>> {
        match self.lookahead_rx {
            Some(ref lv) => Some(lv.clone()),
            None => {
                match self.recv() {
                    Ok(r) => {
                        debug!("recv: {:?}", r);
                        self.lookahead_rx = Some(r.clone());
                        Some(r)
                    },
                    Err(..) => None,
                }
            }
        }
    }

    pub fn accept(&mut self) {
        self.lookahead_rx.take();
        self.lookahead_tx.take();
    }

    pub fn can_tx(&self) -> bool { self.tx.is_some() }
    pub fn can_rx(&self) -> bool { self.rx.is_some() }
}

pub fn try_token(state: &mut eval::State, parent: &mut Connection, msg: &Message) -> bool {
    debug!("tokenstep {:?}", msg);
    let m = msg.eval_down(state);

    debug!("  down: {:?}", m);
    parent.lookahead_send(m);

    match parent.lookahead_receive() {
        Some(m) => {
            debug!("  up: {:?}", m);
            msg.eval_up(state, m)
        }
        None => false
    }
}

pub fn exec(state: &mut eval::State, s: &Step, parent: &mut Connection, child: &mut Connection) -> bool {
    match *s {
        Step::Nop => true,
        Step::Token(ref msg) => {
            let r = try_token(state, parent, msg);
            if r {
                debug!("  matched");
                parent.accept();
            }
            r
        }
        Step::TokenTop(ref msg, box ref body) => {
            debug!("tokentop: {:?}", msg);
            match child.lookahead_receive() {
                Some(m) => {
                    debug!("down: {:?}", m);
                    child.accept();

                    msg.eval_up(state, m);
                    let r = exec(state, body, parent, child);
                    let m = msg.eval_down(state);

                    debug!("up: {:?}", m);
                    if child.send(m).is_err() { return false; }
                    r
                }
                None(..) => false
            }
        }
        Step::Seq(ref steps) => {
            for c in steps.iter() {
                match exec(state, c, parent, child) {
                    true => (),
                    false => return false,
                }
            }
            true
        }
        Step::Repeat(ref count_expr, box ref inner, up) => {
            if up {
                let mut count = 0;
                loop {
                    fn try_first(state: &mut eval::State, f: &Step, parent: &mut Connection, child: &mut Connection) -> bool {
                        match *f {
                            Step::Nop => true,
                            Step::Token(ref msg) => {
                                try_token(state, parent, msg)
                            }
                            Step::TokenTop(_, box ref body) => {
                                child.lookahead_receive();
                                // TODO: variables in the block will not be assigned
                                // so we can't evaluate the child expression if they might be used.
                                // This works for unidirectional applications only.
                                let down = parent.tx.is_some();
                                parent.alive && child.alive &&
                                    (down || try_first(state, body, parent, child))
                            },
                            Step::Seq(ref steps) => steps.get(0)
                                .map_or(true, |x| try_first(state, x, parent, child)),
                            Step::Repeat(_, box ref inner, _) =>
                                try_first(state, inner, parent, child),
                            }
                        }
                    }

                    if !try_first(state, inner, parent, child) {
                        debug!("  loop done");
                        break;
                    }

                    if !exec(state, inner, parent, child) {
                        debug!("loop up early exit");
                        return false;
                    }
                    count += 1;
                }
                debug!("loop up exit {}", count);
                count_expr.eval_up(state, Value::Integer(count))
            } else {
                let count = match count_expr.eval_down(state) {
                    Value::Integer(count) => count,
                    other => panic!("Count must be an integer, found {}", other),
                };

                for _ in 0..count {
                    if !exec(state, inner, parent, child) {
                        debug!("loop down early exit");
                        return false;
                    }
                }

                debug!("loop down exit {}", count);
                true
            }
        }
        //PrimitiveStep(..) => unimplemented!()
    }
}
<|endoftext|>


<commit_before>//! A type that can be treated as a difference.
//!
//! Differential dataflow most commonly tracks the counts associated with records in a multiset, but it
//! generalizes to tracking any map from the records to an Abelian group. The most common generalization
//! is when we maintain both a count and another accumulation, for example height. The differential
//! dataflow collections would then track for each record the total of counts and heights, which allows
//! us to track something like the average.

use std::ops::{Add, Sub, Neg, Mul};

use ::Data;

pub use self::Abelian as Diff;

/// A type that can be treated as a difference.
///
/// The mathematical requirements are, I believe, an Abelian group, in that we require addition, inverses,
/// and almost certainly use commutativity somewhere (it isn't clear if it is a requirement, as it isn't
/// clear that there are semantics other than "we accumulate your differences"; I suspect we don't always
/// accumulate them in the right order, so commutativity is important until we conclude otherwise).
pub trait Monoid : Add<Self, Output=Self> + ::std::marker::Sized + Data + Copy {
	/// Returns true if the element is the additive identity.
	///
	/// This is primarily used by differential dataflow to know when it is safe to delete and update.
	/// When a difference accumulates to zero, the difference has no effect on any accumulation and can
	/// be removed.
	fn is_zero(&self) -> bool;
	/// The additive identity.
	///
	/// This method is primarily used by differential dataflow internals as part of consolidation, when
	/// one value is accumulated elsewhere and must be replaced by valid but harmless value.
	fn zero() -> Self;
}

/// A commutative monoid with negation.
///
/// This trait represents a commutative group, here a commutative monoid with
/// the additional support for subtraction and negation. An identity subtracted
/// from itself or added to its negation should be the zero element from the
/// underlying monoid.
pub trait Abelian : Monoid + Sub<Self, Output=Self> + Neg<Output=Self> { }
impl<T: Monoid + Sub<Self, Output=Self> + Neg<Output=Self>> Abelian for T { }

impl Monoid for isize {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

impl Monoid for i64 {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

impl Monoid for i32 {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

/// The difference defined by a pair of difference elements.
///
/// This type is essentially a "pair", though in Rust the tuple types do not derive the numeric
/// traits we require, and so we need to emulate the types ourselves. In the interest of ergonomics,
/// we may eventually replace the numeric traits with our own, so that we can implement them for
/// tuples and allow users to ignore details like these.
#[derive(Abomonation, Copy, Ord, PartialOrd, Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct DiffPair<R1, R2> {
	/// The first element in the pair.
	pub element1: R1,
	/// The second element in the pair.
	pub element2: R2,
}

impl<R1, R2> DiffPair<R1, R2> {
	/// Creates a new Diff pair from two elements.
	#[inline(always)] pub fn new(elt1: R1, elt2: R2) -> Self {
		DiffPair {
			element1: elt1,
			element2: elt2,
		}
	}
}

impl<R1: Monoid, R2: Monoid> Monoid for DiffPair<R1, R2> {
	#[inline(always)] fn is_zero(&self) -> bool { self.element1.is_zero() && self.element2.is_zero() }
	#[inline(always)] fn zero() -> Self { DiffPair { element1: R1::zero(), element2: R2::zero() } }
}

impl<R1: Add, R2: Add> Add<DiffPair<R1, R2>> for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Add>::Output, <R2 as Add>::Output>;
	#[inline(always)] fn add(self, rhs: Self) -> Self::Output {
		DiffPair {
			element1: self.element1 + rhs.element1,
			element2: self.element2 + rhs.element2,
		}
	}
}

impl<R1: Sub, R2: Sub> Sub<DiffPair<R1, R2>> for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Sub>::Output, <R2 as Sub>::Output>;
	#[inline(always)] fn sub(self, rhs: Self) -> Self::Output {
		DiffPair {
			element1: self.element1 - rhs.element1,
			element2: self.element2 - rhs.element2,
		}
	}
}

impl<R1: Neg, R2: Neg> Neg for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Neg>::Output, <R2 as Neg>::Output>;
	#[inline(always)] fn neg(self) -> Self::Output {
		DiffPair {
			element1: -self.element1,
			element2: -self.element2,
		}
	}
}

impl<T: Copy, R1: Mul<T>, R2: Mul<T>> Mul<T> for DiffPair<R1,R2> {
	type Output = DiffPair<<R1 as Mul<T>>::Output, <R2 as Mul<T>>::Output>;
	fn mul(self, other: T) -> Self::Output {
		DiffPair::new(
			self.element1 * other,
			self.element2 * other,
		)
	}
}

// // TODO: This currently causes rustc to trip a recursion limit, because who knows why.
// impl<R1: Diff, R2: Diff> Mul<DiffPair<R1,R2>> for isize
// where isize: Mul<R1>, isize: Mul<R2>, <isize as Mul<R1>>::Output: Diff, <isize as Mul<R2>>::Output: Diff {
// 	type Output = DiffPair<<isize as Mul<R1>>::Output, <isize as Mul<R2>>::Output>;
// 	fn mul(self, other: DiffPair<R1,R2>) -> Self::Output {
// 		DiffPair::new(
// 			self * other.element1,
// 			self * other.element2,
// 		)
// 	}
// }
<commit_msg>default is_zero implementation<commit_after>//! A type that can be treated as a difference.
//!
//! Differential dataflow most commonly tracks the counts associated with records in a multiset, but it
//! generalizes to tracking any map from the records to an Abelian group. The most common generalization
//! is when we maintain both a count and another accumulation, for example height. The differential
//! dataflow collections would then track for each record the total of counts and heights, which allows
//! us to track something like the average.

use std::ops::{Add, Sub, Neg, Mul};

use ::Data;

pub use self::Abelian as Diff;

/// A type that can be treated as a difference.
///
/// The mathematical requirements are, I believe, an Abelian group, in that we require addition, inverses,
/// and almost certainly use commutativity somewhere (it isn't clear if it is a requirement, as it isn't
/// clear that there are semantics other than "we accumulate your differences"; I suspect we don't always
/// accumulate them in the right order, so commutativity is important until we conclude otherwise).
pub trait Monoid : Add<Self, Output=Self> + ::std::marker::Sized + Data + Copy {
	/// Returns true if the element is the additive identity.
	///
	/// This is primarily used by differential dataflow to know when it is safe to delete and update.
	/// When a difference accumulates to zero, the difference has no effect on any accumulation and can
	/// be removed.
	fn is_zero(&self) -> bool { self.eq(&Self::zero()) }
	/// The additive identity.
	///
	/// This method is primarily used by differential dataflow internals as part of consolidation, when
	/// one value is accumulated elsewhere and must be replaced by valid but harmless value.
	fn zero() -> Self;
}

/// A commutative monoid with negation.
///
/// This trait represents a commutative group, here a commutative monoid with
/// the additional support for subtraction and negation. An identity subtracted
/// from itself or added to its negation should be the zero element from the
/// underlying monoid.
pub trait Abelian : Monoid + Sub<Self, Output=Self> + Neg<Output=Self> { }
impl<T: Monoid + Sub<Self, Output=Self> + Neg<Output=Self>> Abelian for T { }

impl Monoid for isize {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

impl Monoid for i64 {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

impl Monoid for i32 {
	#[inline(always)] fn is_zero(&self) -> bool { *self == 0 }
	#[inline(always)] fn zero() -> Self { 0 }
}

/// The difference defined by a pair of difference elements.
///
/// This type is essentially a "pair", though in Rust the tuple types do not derive the numeric
/// traits we require, and so we need to emulate the types ourselves. In the interest of ergonomics,
/// we may eventually replace the numeric traits with our own, so that we can implement them for
/// tuples and allow users to ignore details like these.
#[derive(Abomonation, Copy, Ord, PartialOrd, Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
pub struct DiffPair<R1, R2> {
	/// The first element in the pair.
	pub element1: R1,
	/// The second element in the pair.
	pub element2: R2,
}

impl<R1, R2> DiffPair<R1, R2> {
	/// Creates a new Diff pair from two elements.
	#[inline(always)] pub fn new(elt1: R1, elt2: R2) -> Self {
		DiffPair {
			element1: elt1,
			element2: elt2,
		}
	}
}

impl<R1: Monoid, R2: Monoid> Monoid for DiffPair<R1, R2> {
	#[inline(always)] fn is_zero(&self) -> bool { self.element1.is_zero() && self.element2.is_zero() }
	#[inline(always)] fn zero() -> Self { DiffPair { element1: R1::zero(), element2: R2::zero() } }
}

impl<R1: Add, R2: Add> Add<DiffPair<R1, R2>> for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Add>::Output, <R2 as Add>::Output>;
	#[inline(always)] fn add(self, rhs: Self) -> Self::Output {
		DiffPair {
			element1: self.element1 + rhs.element1,
			element2: self.element2 + rhs.element2,
		}
	}
}

impl<R1: Sub, R2: Sub> Sub<DiffPair<R1, R2>> for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Sub>::Output, <R2 as Sub>::Output>;
	#[inline(always)] fn sub(self, rhs: Self) -> Self::Output {
		DiffPair {
			element1: self.element1 - rhs.element1,
			element2: self.element2 - rhs.element2,
		}
	}
}

impl<R1: Neg, R2: Neg> Neg for DiffPair<R1, R2> {
	type Output = DiffPair<<R1 as Neg>::Output, <R2 as Neg>::Output>;
	#[inline(always)] fn neg(self) -> Self::Output {
		DiffPair {
			element1: -self.element1,
			element2: -self.element2,
		}
	}
}

impl<T: Copy, R1: Mul<T>, R2: Mul<T>> Mul<T> for DiffPair<R1,R2> {
	type Output = DiffPair<<R1 as Mul<T>>::Output, <R2 as Mul<T>>::Output>;
	fn mul(self, other: T) -> Self::Output {
		DiffPair::new(
			self.element1 * other,
			self.element2 * other,
		)
	}
}

// // TODO: This currently causes rustc to trip a recursion limit, because who knows why.
// impl<R1: Diff, R2: Diff> Mul<DiffPair<R1,R2>> for isize
// where isize: Mul<R1>, isize: Mul<R2>, <isize as Mul<R1>>::Output: Diff, <isize as Mul<R2>>::Output: Diff {
// 	type Output = DiffPair<<isize as Mul<R1>>::Output, <isize as Mul<R2>>::Output>;
// 	fn mul(self, other: DiffPair<R1,R2>) -> Self::Output {
// 		DiffPair::new(
// 			self * other.element1,
// 			self * other.element2,
// 		)
// 	}
// }
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Parsing and loading Executable and Linkable Format (ELF) 32- and 64-bit
//! binaries.
//!
//! For more information on the ELF format, refer to:
//!
//!  + [Wikipedia](https://en.wikipedia.org/wiki/Executable_and_Linkable_Format)
//!  + The [OS Dev Wiki](http://wiki.osdev.org/ELF)
//!  + The [ELF Format Specification](elfspec)
//!
//! [elfspec]: http://www.skyfree.org/linux/references/ELF_Format.pdf
#![feature(core_intrinsics)]
#![feature(try_from)]
#![no_std]

#[macro_use] extern crate bitflags;
#[macro_use] extern crate macro_attr;

extern crate memory;

use core::{ intrinsics, ops, mem, slice };

pub mod section;
pub mod file;
pub mod program;

/// An ELF section header.
pub type Section<'a> = section::Header<'a>;
/// An ELF header file.
pub type FileHeader<W> = file::HeaderRepr<W>;

/// TODO: should ELF have its own error type?
pub type ElfResult<T> = Result<T, &'static str>;

pub trait ElfWord: Sized + Copy + Clone
                         + ops::Add<Self> + ops::Sub<Self>
                         + ops::Mul<Self> + ops::Div<Self>
                         + ops::Shl<Self> + ops::Shr<Self> { }
impl ElfWord for u64 { }
impl ElfWord for u32 { }

/// Hack to make the type-system let me do what I want
trait ValidatesWord<Word: ElfWord> {
    fn check(&self) -> ElfResult<()>;
}

/// A handle on a parsed ELF binary
///  TODO: do we want this to own a HashMap of section names to section headers,
///        to speed up section lookup?
//          - eliza, 03/08/2017
#[derive(Debug)]
pub struct Image<'a, Word, Header = file::HeaderRepr<Word>>
where Word: ElfWord + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// the binary's [file header](file/trait.Header.html)
    pub header: &'a Header
  , /// references to each [section header](section/struct.Header.html)
    pub sections: &'a [section::Header<'a>]
  , /// the raw binary contents of the ELF binary.
    /// note that this includes the _entire_ binary contents of the file,
    /// so the file header and each section header is included in this slice.
    binary: &'a [u8]
}

impl<'a, Word, Header> Image<'a, Word, Header>
where Word: ElfWord + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// Returns the section header [string table].
    ///
    /// [string table]: section/struct.StrTable.html
    pub fn sh_str_table(&'a self) -> section::StrTable<'a> {
        // TODO: do we want to validate that the string table index is
        //       reasonable (e.g. it's not longer than the binary)?
        //          - eliza, 03/08/2017
        // TODO: do we want to cache a ref to the string table?
        //          - eliza, 03/08/2017
        section::StrTable::from(&self.binary[self.header.sh_str_idx()..])
    }

}

/// Extract `n` instances of type `T` from a byte slice.
///
/// This is essentially just a _slightly_ safer wrapper around
/// [`slice::from_raw_parts`]. Unlike `from_raw_parts`, this function takes
/// a valid byte slice, rather than a pointer. Therefore, some of the safety
/// issues with `from_raw_parts` are avoided:
///
/// + the lifetime (`'slice`) of the returned slice should be the same as the
///   lifetime of the input slice (`data`), rather than inferred arbitrarily.
/// + this function will panic rather than reading past the end of the slice.
///
/// # Arguments
///
/// + `data`: the byte slice to extract a slice of `&[T]`s from
/// + `offset`: a start offset into `data`
/// + `n`: the number of instances of `T` which should be contained
///        in `data[offset..]`
///
/// # Safety
///
/// While this function is safer than [`slice::from_raw_parts`],
/// it is still unsafe for the following reasons:
///
/// + The contents of `data` may not be able to be interpreted as instances of
///   type `T`.
/// + `offset` may not be aligned on a `T`-sized boundary.
///
/// # Caveats
///
/// + If `n` == 0, this will give you an `&[]`. Just a warning.
//    thanks to Max for making  me figure this out.
/// + `offset` must be aligned on a `T`-sized boundary.
///
/// # Panics
///
/// + If the slice `data` is not long enough to contain `n` instances of `T`.
/// + If the index `offset` is longer than `T`
///
/// TODO: rewrite this as a `TryFrom` implementation (see issue #85)
//          - eliza, 03/09/2017
///       wait, possibly we should NOT do that. actually we should
///       almost certainly not do that. since this function is unsafe,
///       but `TryFrom` is not, and because this would be WAY generic.
//          - eliza, 03/09/2017
/// TODO: is this general enough to move into util?
//          - eliza, 03/09/2017
/// TODO: should this be refactored to return a `Result`?
//          - eliza, 03/13/2017
/// TODO: assert that `offset` is aligned on a `T`-sized boundary
//          - eliza, 03/13/2017
/// TODO: do we want to assert that `offset` is less than the length of `data`
///       separately from asserting that the slice is long enough, so that
///       we can panic with different messages?
//          - eliza, 03/13/2017
/// TODO: refactor this to take a `RangeArgument`?
//          - eliza, 03/13/2017
///
/// [`slice::from_raw_parts`]: https://doc.rust-lang.org/stable/std/slice/fn.from_raw_parts.html
unsafe fn extract_from_slice<'slice, T: Sized>( data: &'slice [u8]
                                              , offset: usize
                                              , n: usize)
                                              -> &'slice [T] {
    assert!( data.len() - offset >= mem::size_of::<T>() * n
           , "Slice too short to contain {} objects of type {}"
           , n
           , intrinsics::type_name::<T>()
           );
    slice::from_raw_parts(data[offset..].as_ptr() as *const T, n)
}
<commit_msg>fix(elf): add assertion that extract_from_slice offset is T-aligned<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Parsing and loading Executable and Linkable Format (ELF) 32- and 64-bit
//! binaries.
//!
//! For more information on the ELF format, refer to:
//!
//!  + [Wikipedia](https://en.wikipedia.org/wiki/Executable_and_Linkable_Format)
//!  + The [OS Dev Wiki](http://wiki.osdev.org/ELF)
//!  + The [ELF Format Specification](elfspec)
//!
//! [elfspec]: http://www.skyfree.org/linux/references/ELF_Format.pdf
#![feature(core_intrinsics)]
#![feature(try_from)]
#![no_std]

#[macro_use] extern crate bitflags;
#[macro_use] extern crate macro_attr;

extern crate memory;

use core::{ ops, mem, slice };

pub mod section;
pub mod file;
pub mod program;

/// An ELF section header.
pub type Section<'a> = section::Header<'a>;
/// An ELF header file.
pub type FileHeader<W> = file::HeaderRepr<W>;

/// TODO: should ELF have its own error type?
pub type ElfResult<T> = Result<T, &'static str>;

pub trait ElfWord: Sized + Copy + Clone
                         + ops::Add<Self> + ops::Sub<Self>
                         + ops::Mul<Self> + ops::Div<Self>
                         + ops::Shl<Self> + ops::Shr<Self> { }
impl ElfWord for u64 { }
impl ElfWord for u32 { }

/// Hack to make the type-system let me do what I want
trait ValidatesWord<Word: ElfWord> {
    fn check(&self) -> ElfResult<()>;
}

/// A handle on a parsed ELF binary
///  TODO: do we want this to own a HashMap of section names to section headers,
///        to speed up section lookup?
//          - eliza, 03/08/2017
#[derive(Debug)]
pub struct Image<'a, Word, Header = file::HeaderRepr<Word>>
where Word: ElfWord + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// the binary's [file header](file/trait.Header.html)
    pub header: &'a Header
  , /// references to each [section header](section/struct.Header.html)
    pub sections: &'a [section::Header<'a>]
  , /// the raw binary contents of the ELF binary.
    /// note that this includes the _entire_ binary contents of the file,
    /// so the file header and each section header is included in this slice.
    binary: &'a [u8]
}

impl<'a, Word, Header> Image<'a, Word, Header>
where Word: ElfWord + 'a
    , Header: file::Header<Word = Word> + 'a
    {
    /// Returns the section header [string table].
    ///
    /// [string table]: section/struct.StrTable.html
    pub fn sh_str_table(&'a self) -> section::StrTable<'a> {
        // TODO: do we want to validate that the string table index is
        //       reasonable (e.g. it's not longer than the binary)?
        //          - eliza, 03/08/2017
        // TODO: do we want to cache a ref to the string table?
        //          - eliza, 03/08/2017
        section::StrTable::from(&self.binary[self.header.sh_str_idx()..])
    }

}

/// Extract `n` instances of type `T` from a byte slice.
///
/// This is essentially just a _slightly_ safer wrapper around
/// [`slice::from_raw_parts`]. Unlike `from_raw_parts`, this function takes
/// a valid byte slice, rather than a pointer. Therefore, some of the safety
/// issues with `from_raw_parts` are avoided:
///
/// + the lifetime (`'slice`) of the returned slice should be the same as the
///   lifetime of the input slice (`data`), rather than inferred arbitrarily.
/// + this function will panic rather than reading past the end of the slice.
///
/// # Arguments
///
/// + `data`: the byte slice to extract a slice of `&[T]`s from
/// + `offset`: a start offset into `data`
/// + `n`: the number of instances of `T` which should be contained
///        in `data[offset..]`
///
/// # Safety
///
/// While this function is safer than [`slice::from_raw_parts`],
/// it is still unsafe for the following reasons:
///
/// + The contents of `data` may not be able to be interpreted as instances of
///   type `T`.
///
/// # Caveats
///
/// + If `n` == 0, this will give you an `&[]`. Just a warning.
//    thanks to Max for making  me figure this out.
/// + `offset` must be aligned on a `T`-sized boundary.
///
/// # Panics
///
/// + If the slice `data` is not long enough to contain `n` instances of `T`.
/// + If the index `offset` is longer than `T`
/// + If `offset` is not aligned on a `T`-sized boundary.
///
/// TODO: rewrite this as a `TryFrom` implementation (see issue #85)
//          - eliza, 03/09/2017
///       wait, possibly we should NOT do that. actually we should
///       almost certainly not do that. since this function is unsafe,
///       but `TryFrom` is not, and because this would be WAY generic.
//          - eliza, 03/09/2017
/// TODO: is this general enough to move into util?
//          - eliza, 03/09/2017
/// TODO: should this be refactored to return a `Result`?
//          - eliza, 03/13/2017
/// TODO: assert that `offset` is aligned on a `T`-sized boundary
//          - eliza, 03/13/2017
/// TODO: refactor this to take a `RangeArgument`?
//          - eliza, 03/13/2017
///
/// [`slice::from_raw_parts`]: https://doc.rust-lang.org/stable/std/slice/fn.from_raw_parts.html
unsafe fn extract_from_slice<'slice, T: Sized>( data: &'slice [u8]
                                              , offset: usize
                                              , n: usize)
                                              -> &'slice [T] {
    use core::intrinsics::type_name;
    assert!( offset % mem::align_of::<T>() == 0
           , "Offset {} not aligned on a {}-sized boundary (must be \
              divisible by {})."
           , offset, type_name::<T>(), mem::align_of::<T>()
           );
    assert!( data.len() - offset >= mem::size_of::<T>() * n
           , "Slice too short to contain {} objects of type {}"
           , n, type_name::<T>()
           );
    slice::from_raw_parts(data[offset..].as_ptr() as *const T, n)
}
<|endoftext|>


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

//! The dirichlet distribution.

use Rng;
use distributions::Distribution;
use distributions::gamma::Gamma;

/// The dirichelet distribution `Dirichlet(alpha)`.
///
/// The Dirichlet distribution is a family of continuous multivariate probability distributions parameterized by
/// a vector alpha of positive reals. https://en.wikipedia.org/wiki/Dirichlet_distribution
/// It is a multivariate generalization of the beta distribution.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::Dirichlet;
///
/// let dirichlet = Dirichlet::new(vec![1.0, 2.0, 3.0]);
/// let samples = dirichlet.sample(&mut rand::thread_rng());
/// println!("{:?} is from a Dirichlet([1.0, 2.0, 3.0]) distribution", samples);
/// ```

#[derive(Clone, Debug)]
pub struct Dirichlet {
    /// Concentration parameters (alpha)
    alpha: Vec<f64>,
}

impl Dirichlet {
    /// Construct a new `Dirichlet` with the given alpha parameter
    /// `alpha`. Panics if `alpha.len() < 2`.
    #[inline]
    pub fn new<V: Into<Vec<f64>>>(alpha: V) -> Dirichlet {
        let a = alpha.into();
        assert!(a.len() > 1);
        for i in 0..a.len() {
            assert!(a[i] > 0.0);
        }

        Dirichlet { alpha: a }
    }

    /// Construct a new `Dirichlet` with the given shape parameter and size
    /// `alpha`. Panics if `alpha <= 0.0`.
    /// `size` . Panic if `size < 2`
    #[inline]
    pub fn new_with_param(alpha: f64, size: usize) -> Dirichlet {
        assert!(alpha > 0.0);
        assert!(size > 1);
        Dirichlet {
            alpha: vec![alpha; size],
        }
    }
}

impl Distribution<Vec<f64>> for Dirichlet {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<f64> {
        let n = self.alpha.len();
        let mut samples = vec![0.0f64; n];
        let mut sum = 0.0f64;

        for i in 0..n {
            let g = Gamma::new(self.alpha[i], 1.0);
            samples[i] = g.sample(rng);
            sum += samples[i];
        }
        let invacc = 1.0 / sum;
        for i in 0..n {
            samples[i] *= invacc;
        }
        samples
    }
}

#[cfg(test)]
mod test {
    use super::Dirichlet;
    use distributions::Distribution;

    #[test]
    fn test_dirichlet() {
        let d = Dirichlet::new(vec![1.0, 2.0, 3.0]);
        let mut rng = ::test::rng(221);
        let samples = d.sample(&mut rng);
        let _: Vec<f64> = samples
            .into_iter()
            .map(|x| {
                assert!(x > 0.0);
                x
            })
            .collect();
    }

    #[test]
    fn test_dirichlet_with_param() {
        let alpha = 0.5f64;
        let size = 2;
        let d = Dirichlet::new_with_param(alpha, size);
        let mut rng = ::test::rng(221);
        let samples = d.sample(&mut rng);
        let _: Vec<f64> = samples
            .into_iter()
            .map(|x| {
                assert!(x > 0.0);
                x
            })
            .collect();
    }

    #[test]
    #[should_panic]
    fn test_dirichlet_invalid_length() {
        Dirichlet::new_with_param(0.5f64, 1);
    }

    #[test]
    #[should_panic]
    fn test_dirichlet_invalid_alpha() {
        Dirichlet::new_with_param(0.0f64, 2);
    }
}
<commit_msg>doc markup update<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The dirichlet distribution.

use Rng;
use distributions::Distribution;
use distributions::gamma::Gamma;

/// The dirichelet distribution `Dirichlet(alpha)`.
///
/// The Dirichlet distribution is a family of continuous multivariate probability distributions parameterized by
/// a vector alpha of positive reals. https://en.wikipedia.org/wiki/Dirichlet_distribution
/// It is a multivariate generalization of the beta distribution.
///
/// # Example
///
/// ```
/// use rand::prelude::*;
/// use rand::distributions::Dirichlet;
///
/// let dirichlet = Dirichlet::new(vec![1.0, 2.0, 3.0]);
/// let samples = dirichlet.sample(&mut rand::thread_rng());
/// println!("{:?} is from a Dirichlet([1.0, 2.0, 3.0]) distribution", samples);
/// ```

#[derive(Clone, Debug)]
pub struct Dirichlet {
    /// Concentration parameters (alpha)
    alpha: Vec<f64>,
}

impl Dirichlet {
    /// Construct a new `Dirichlet` with the given alpha parameter `alpha`.
    ///
    /// # Panics
    /// - if `alpha.len() < 2`
    ///
    #[inline]
    pub fn new<V: Into<Vec<f64>>>(alpha: V) -> Dirichlet {
        let a = alpha.into();
        assert!(a.len() > 1);
        for i in 0..a.len() {
            assert!(a[i] > 0.0);
        }

        Dirichlet { alpha: a }
    }

    /// Construct a new `Dirichlet` with the given shape parameter `alpha` and `size`.
    ///
    /// # Panics
    /// - if `alpha <= 0.0`
    /// - if `size < 2`
    ///
    #[inline]
    pub fn new_with_param(alpha: f64, size: usize) -> Dirichlet {
        assert!(alpha > 0.0);
        assert!(size > 1);
        Dirichlet {
            alpha: vec![alpha; size],
        }
    }
}

impl Distribution<Vec<f64>> for Dirichlet {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Vec<f64> {
        let n = self.alpha.len();
        let mut samples = vec![0.0f64; n];
        let mut sum = 0.0f64;

        for i in 0..n {
            let g = Gamma::new(self.alpha[i], 1.0);
            samples[i] = g.sample(rng);
            sum += samples[i];
        }
        let invacc = 1.0 / sum;
        for i in 0..n {
            samples[i] *= invacc;
        }
        samples
    }
}

#[cfg(test)]
mod test {
    use super::Dirichlet;
    use distributions::Distribution;

    #[test]
    fn test_dirichlet() {
        let d = Dirichlet::new(vec![1.0, 2.0, 3.0]);
        let mut rng = ::test::rng(221);
        let samples = d.sample(&mut rng);
        let _: Vec<f64> = samples
            .into_iter()
            .map(|x| {
                assert!(x > 0.0);
                x
            })
            .collect();
    }

    #[test]
    fn test_dirichlet_with_param() {
        let alpha = 0.5f64;
        let size = 2;
        let d = Dirichlet::new_with_param(alpha, size);
        let mut rng = ::test::rng(221);
        let samples = d.sample(&mut rng);
        let _: Vec<f64> = samples
            .into_iter()
            .map(|x| {
                assert!(x > 0.0);
                x
            })
            .collect();
    }

    #[test]
    #[should_panic]
    fn test_dirichlet_invalid_length() {
        Dirichlet::new_with_param(0.5f64, 1);
    }

    #[test]
    #[should_panic]
    fn test_dirichlet_invalid_alpha() {
        Dirichlet::new_with_param(0.0f64, 2);
    }
}
<|endoftext|>


<commit_before>//! Module for working with `Function`s
//!
//! This module contains a type alias for `Rc<Fn(f64) -> f64>`,
//! which is used in many other modules, and functions for
//! working with the alias.

pub use std::rc::Rc;

/// Type alias used to represent functions.
///
/// A `Function` is a `Fn` that takes a single `f64`,
/// does something with it, and returns another `f64`.
///
/// Functions are stored in an `Rc` so they can be `cloned()`
/// and subsequently consumed in other functions.
pub type Function = Rc<Fn(f64) -> f64>;

/// Macro for creating a `Function`.
///
/// More idiomatic than calling `Rc::new()`.
///
/// # Examples
///
/// ```
/// # #[macro_use]
/// # extern crate reikna;
/// # fn main() {
/// use reikna::func::*;
/// let f: Function = func!(|x| x * x);
/// assert_eq!(f(5.0), 25.0);
/// # }
/// ```
#[macro_export]
macro_rules! func {
    ($e:expr) => (Rc::new($e) as Function);
}
<commit_msg>Add `dyn` to trait object<commit_after>//! Module for working with `Function`s
//!
//! This module contains a type alias for `Rc<Fn(f64) -> f64>`,
//! which is used in many other modules, and functions for
//! working with the alias.

pub use std::rc::Rc;

/// Type alias used to represent functions.
///
/// A `Function` is a `Fn` that takes a single `f64`,
/// does something with it, and returns another `f64`.
///
/// Functions are stored in an `Rc` so they can be `cloned()`
/// and subsequently consumed in other functions.
pub type Function = Rc<dyn Fn(f64) -> f64>;

/// Macro for creating a `Function`.
///
/// More idiomatic than calling `Rc::new()`.
///
/// # Examples
///
/// ```
/// # #[macro_use]
/// # extern crate reikna;
/// # fn main() {
/// use reikna::func::*;
/// let f: Function = func!(|x| x * x);
/// assert_eq!(f(5.0), 25.0);
/// # }
/// ```
#[macro_export]
macro_rules! func {
    ($e:expr) => (Rc::new($e) as Function);
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => () // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<commit_msg>Fix style for consistency<commit_after>use std::cell::RefCell;
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = TCPConnection::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                };
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        TCPConnection::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = TCPConnection::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            TCPConnection::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        TCPConnection::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = TCPConnection::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            TCPConnection::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<|endoftext|>