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

// Coroutines represent nothing more than a context and a stack
// segment.

use std::default::Default;
use std::rt::util::min_stack;
use std::thunk::Thunk;
use std::mem::transmute;
use std::rt::unwind::try;
use std::boxed::BoxAny;
use std::cell::UnsafeCell;
use std::rc::Rc;

use context::Context;
use stack::{StackPool, Stack};

#[derive(Debug, Copy)]
pub enum State {
    Suspended,
    Running,
    Finished,
}

/// Coroutine spawn options
#[derive(Debug)]
pub struct Options {
    stack_size: usize,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            stack_size: min_stack(),
        }
    }
}

/// Handle of a Coroutine
#[derive(Clone)]
pub struct Handle(Rc<UnsafeCell<Coroutine>>);

impl Handle {
    pub fn resume(&self) {
        Coroutine::resume(&self)
    }

    pub fn state(&self) -> State {
        Coroutine::state(&self)
    }
}

/// A coroutine is nothing more than a (register context, stack) pair.
#[allow(raw_pointer_derive)]
pub struct Coroutine {
    /// The segment of stack on which the task is currently running or
    /// if the task is blocked, on which the task will resume
    /// execution.
    current_stack_segment: Option<Stack>,

    /// Always valid if the task is alive and not running.
    saved_context: Context,

    /// Parent coroutine, may always be valid.
    parent: Option<Handle>,

    /// State
    state: State
}

/// Destroy coroutine and try to reuse std::stack segment.
impl Drop for Coroutine {
    fn drop(&mut self) {
        match self.current_stack_segment.take() {
            Some(stack) => {
                COROUTINE_ENVIRONMENT.with(|env| {
                    let env: &mut Environment = unsafe { transmute(env.get()) };
                    env.stack_pool.give_stack(stack);
                });
            },
            None => {}
        }
    }
}

/// Initialization function for make context
extern "C" fn coroutine_initialize(_: usize, f: *mut ()) -> ! {
    let func: Box<Thunk> = unsafe { transmute(f) };

    if let Err(cause) = unsafe { try(move|| func.invoke(())) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    COROUTINE_ENVIRONMENT.with(|env| {
        let env: &mut Environment = unsafe { transmute(env.get()) };
        let cur: &mut Coroutine = unsafe { transmute(env.current_running.0.get()) };
        cur.state = State::Finished;
    });

    loop {
        Coroutine::sched();
    }
}

impl Coroutine {
    pub fn empty() -> Handle {
        Handle(Rc::new(UnsafeCell::new(Coroutine {
            current_stack_segment: None,
            saved_context: Context::empty(),
            parent: None,
            state: State::Running,
        })))
    }

    pub fn spawn_opts<F>(f: F, opts: Options) -> Handle
            where F: FnOnce() + Send + 'static {

        let coro = UnsafeCell::new(Coroutine::empty());
        COROUTINE_ENVIRONMENT.with(|env| {
            unsafe {
                let env: &mut Environment = transmute(env.get());

                let mut stack = env.stack_pool.take_stack(opts.stack_size);

                let ctx = Context::new(coroutine_initialize,
                                   0,
                                   f,
                                   &mut stack);

                let coro_cell: &mut Rc<UnsafeCell<Coroutine>> = transmute(coro.get());
                let coro: &mut Coroutine = transmute(coro_cell.get());
                coro.saved_context = ctx;
                coro.current_stack_segment = Some(stack);
                coro.state = State::Suspended;
            }
        });

        let coro = unsafe { coro.into_inner() };
        Coroutine::resume(&coro);
        coro
    }

    /// Spawn a coroutine with default options
    pub fn spawn<F>(f: F) -> Handle
            where F: FnOnce() + Send + 'static {
        Coroutine::spawn_opts(f, Default::default())
    }

    pub fn resume(coro: &Handle) {
        let to_coro: &mut Coroutine = unsafe { transmute(coro.0.get()) };
        match to_coro.state {
            State::Finished | State::Running => return,
            _ => {}
        }

        COROUTINE_ENVIRONMENT.with(|env| {
            let env: &mut Environment = unsafe { transmute(env.get()) };

            let current_running = env.current_running.clone();

            let from_coro: &mut Coroutine = unsafe { transmute(current_running.0.get()) };

            to_coro.parent = Some(current_running.clone());
            env.current_running = coro.clone();

            from_coro.state = State::Suspended;
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        })
    }

    pub fn sched() {
        COROUTINE_ENVIRONMENT.with(|env| {
            let env: &mut Environment = unsafe { transmute(env.get()) };

            let from_coro: &mut Coroutine = unsafe { transmute(env.current_running.0.get()) };
            let to_coro: &mut Coroutine = match from_coro.parent {
                Some(ref parent) => unsafe {
                    env.current_running = parent.clone();
                    transmute(parent.0.get())
                },
                None => return,
            };

            match from_coro.state {
                State::Finished => {},
                _ => from_coro.state = State::Suspended,
            }
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        });
    }

    pub fn current() -> Handle {
        unsafe {
            let opt = UnsafeCell::new(None);
            COROUTINE_ENVIRONMENT.with(|env| {
                let env: &mut Environment = transmute(env.get());
                let opt: &mut Option<Handle> = transmute(opt.get());

                *opt = Some(env.current_running.clone());
            });
            opt.into_inner().unwrap()
        }
    }

    pub fn state(coro: &Handle) -> State {
        let coro: &mut Coroutine = unsafe { transmute(coro.0.get()) };
        coro.state
    }
}

thread_local!(static COROUTINE_ENVIRONMENT: UnsafeCell<Environment> = UnsafeCell::new(Environment::new()));

/// Coroutine managing environment
#[allow(raw_pointer_derive)]
struct Environment {
    stack_pool: StackPool,
    current_running: Handle,
}

impl Environment {
    /// Initialize a new environment
    fn new() -> Environment {
        Environment {
            stack_pool: StackPool::new(),
            current_running: Coroutine::empty(),
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;

    use test::Bencher;

    use coroutine::Coroutine;

    #[test]
    fn test_coroutine_basic() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
    }

    #[test]
    fn test_coroutine_yield() {
        let (tx, rx) = channel();
        let coro = Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::sched();

            tx.send(2).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert!(rx.try_recv().is_err());

        Coroutine::resume(&coro);

        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    fn test_coroutine_spawn_inside() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::spawn(move|| {
                tx.send(2).unwrap();
            });
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    #[should_fail]
    fn test_coroutine_panic() {
        Coroutine::spawn(move|| {
            panic!("Panic inside a coroutine!!");
        });

        unreachable!();
    }

    #[test]
    #[should_fail]
    fn test_coroutine_child_panic() {
        Coroutine::spawn(move|| {

            Coroutine::spawn(move|| {
                panic!("Panic inside a coroutine's child!!");
            });

        });

        unreachable!();
    }

    #[test]
    fn test_coroutine_resume_after_finished() {
        let coro = Coroutine::spawn(move|| {});

        // It is already finished, but we try to resume it
        // Idealy it would come back immediately
        Coroutine::resume(&coro);

        // Again?
        Coroutine::resume(&coro);
    }

    #[bench]
    fn bench_coroutine_spawning_with_recycle(b: &mut Bencher) {
        b.iter(|| {
            Coroutine::spawn(move|| {});
        });
    }

    #[bench]
    fn bench_normal_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;

            let (tx, rx) = channel();

            let mut result = 0;
            for _ in 0..MAX_NUMBER {
                tx.send(1).unwrap();
                result += rx.recv().unwrap();
            }
            assert_eq!(result, MAX_NUMBER);
        });
    }

    #[bench]
    fn bench_coroutine_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;
            let (tx, rx) = channel();

            let coro = Coroutine::spawn(move|| {
                for _ in 0..MAX_NUMBER {
                    tx.send(1).unwrap();
                    Coroutine::sched();
                }
            });

            let result = rx.iter().fold(0, |a, b| {
                Coroutine::resume(&coro);
                a + b
            });
            assert_eq!(result, MAX_NUMBER);
        });
    }
}
<commit_msg>add two tests<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Coroutines represent nothing more than a context and a stack
// segment.

use std::default::Default;
use std::rt::util::min_stack;
use std::thunk::Thunk;
use std::mem::transmute;
use std::rt::unwind::try;
use std::boxed::BoxAny;
use std::cell::UnsafeCell;
use std::rc::Rc;

use context::Context;
use stack::{StackPool, Stack};

#[derive(Debug, Copy)]
pub enum State {
    Suspended,
    Running,
    Finished,
}

/// Coroutine spawn options
#[derive(Debug)]
pub struct Options {
    stack_size: usize,
}

impl Default for Options {
    fn default() -> Options {
        Options {
            stack_size: min_stack(),
        }
    }
}

/// Handle of a Coroutine
#[derive(Clone)]
pub struct Handle(Rc<UnsafeCell<Coroutine>>);

impl Handle {
    pub fn resume(&self) {
        Coroutine::resume(&self)
    }

    pub fn state(&self) -> State {
        Coroutine::state(&self)
    }
}

/// A coroutine is nothing more than a (register context, stack) pair.
#[allow(raw_pointer_derive)]
pub struct Coroutine {
    /// The segment of stack on which the task is currently running or
    /// if the task is blocked, on which the task will resume
    /// execution.
    current_stack_segment: Option<Stack>,

    /// Always valid if the task is alive and not running.
    saved_context: Context,

    /// Parent coroutine, may always be valid.
    parent: Option<Handle>,

    /// State
    state: State
}

/// Destroy coroutine and try to reuse std::stack segment.
impl Drop for Coroutine {
    fn drop(&mut self) {
        match self.current_stack_segment.take() {
            Some(stack) => {
                COROUTINE_ENVIRONMENT.with(|env| {
                    let env: &mut Environment = unsafe { transmute(env.get()) };
                    env.stack_pool.give_stack(stack);
                });
            },
            None => {}
        }
    }
}

/// Initialization function for make context
extern "C" fn coroutine_initialize(_: usize, f: *mut ()) -> ! {
    let func: Box<Thunk> = unsafe { transmute(f) };

    if let Err(cause) = unsafe { try(move|| func.invoke(())) } {
        error!("Panicked inside: {:?}", cause.downcast::<&str>());
    }

    COROUTINE_ENVIRONMENT.with(|env| {
        let env: &mut Environment = unsafe { transmute(env.get()) };
        let cur: &mut Coroutine = unsafe { transmute(env.current_running.0.get()) };
        cur.state = State::Finished;
    });

    loop {
        Coroutine::sched();
    }
}

impl Coroutine {
    pub fn empty() -> Handle {
        Handle(Rc::new(UnsafeCell::new(Coroutine {
            current_stack_segment: None,
            saved_context: Context::empty(),
            parent: None,
            state: State::Running,
        })))
    }

    pub fn spawn_opts<F>(f: F, opts: Options) -> Handle
            where F: FnOnce() + Send + 'static {

        let coro = UnsafeCell::new(Coroutine::empty());
        COROUTINE_ENVIRONMENT.with(|env| {
            unsafe {
                let env: &mut Environment = transmute(env.get());

                let mut stack = env.stack_pool.take_stack(opts.stack_size);

                let ctx = Context::new(coroutine_initialize,
                                   0,
                                   f,
                                   &mut stack);

                let coro_cell: &mut Rc<UnsafeCell<Coroutine>> = transmute(coro.get());
                let coro: &mut Coroutine = transmute(coro_cell.get());
                coro.saved_context = ctx;
                coro.current_stack_segment = Some(stack);
                coro.state = State::Suspended;
            }
        });

        let coro = unsafe { coro.into_inner() };
        Coroutine::resume(&coro);
        coro
    }

    /// Spawn a coroutine with default options
    pub fn spawn<F>(f: F) -> Handle
            where F: FnOnce() + Send + 'static {
        Coroutine::spawn_opts(f, Default::default())
    }

    pub fn resume(coro: &Handle) {
        let to_coro: &mut Coroutine = unsafe { transmute(coro.0.get()) };
        match to_coro.state {
            State::Finished | State::Running => return,
            _ => {}
        }

        COROUTINE_ENVIRONMENT.with(|env| {
            let env: &mut Environment = unsafe { transmute(env.get()) };

            let current_running = env.current_running.clone();

            let from_coro: &mut Coroutine = unsafe { transmute(current_running.0.get()) };

            to_coro.parent = Some(current_running.clone());
            env.current_running = coro.clone();

            from_coro.state = State::Suspended;
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        })
    }

    pub fn sched() {
        COROUTINE_ENVIRONMENT.with(|env| {
            let env: &mut Environment = unsafe { transmute(env.get()) };

            let from_coro: &mut Coroutine = unsafe { transmute(env.current_running.0.get()) };
            let to_coro: &mut Coroutine = match from_coro.parent {
                Some(ref parent) => unsafe {
                    env.current_running = parent.clone();
                    transmute(parent.0.get())
                },
                None => return,
            };

            match from_coro.state {
                State::Finished => {},
                _ => from_coro.state = State::Suspended,
            }
            to_coro.state = State::Running;

            Context::swap(&mut from_coro.saved_context, &to_coro.saved_context);
        });
    }

    pub fn current() -> Handle {
        unsafe {
            let opt = UnsafeCell::new(None);
            COROUTINE_ENVIRONMENT.with(|env| {
                let env: &mut Environment = transmute(env.get());
                let opt: &mut Option<Handle> = transmute(opt.get());

                *opt = Some(env.current_running.clone());
            });
            opt.into_inner().unwrap()
        }
    }

    pub fn state(coro: &Handle) -> State {
        let coro: &mut Coroutine = unsafe { transmute(coro.0.get()) };
        coro.state
    }
}

thread_local!(static COROUTINE_ENVIRONMENT: UnsafeCell<Environment> = UnsafeCell::new(Environment::new()));

/// Coroutine managing environment
#[allow(raw_pointer_derive)]
struct Environment {
    stack_pool: StackPool,
    current_running: Handle,
}

impl Environment {
    /// Initialize a new environment
    fn new() -> Environment {
        Environment {
            stack_pool: StackPool::new(),
            current_running: Coroutine::empty(),
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;

    use test::Bencher;

    use coroutine::Coroutine;

    #[test]
    fn test_coroutine_basic() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
    }

    #[test]
    fn test_coroutine_yield() {
        let (tx, rx) = channel();
        let coro = Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::sched();

            tx.send(2).unwrap();
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert!(rx.try_recv().is_err());

        Coroutine::resume(&coro);

        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    fn test_coroutine_spawn_inside() {
        let (tx, rx) = channel();
        Coroutine::spawn(move|| {
            tx.send(1).unwrap();

            Coroutine::spawn(move|| {
                tx.send(2).unwrap();
            });
        });

        assert_eq!(rx.recv().unwrap(), 1);
        assert_eq!(rx.recv().unwrap(), 2);
    }

    #[test]
    #[should_fail]
    fn test_coroutine_panic() {
        Coroutine::spawn(move|| {
            panic!("Panic inside a coroutine!!");
        });

        unreachable!();
    }

    #[test]
    #[should_fail]
    fn test_coroutine_child_panic() {
        Coroutine::spawn(move|| {

            Coroutine::spawn(move|| {
                panic!("Panic inside a coroutine's child!!");
            });

        });

        unreachable!();
    }

    #[test]
    fn test_coroutine_resume_after_finished() {
        let coro = Coroutine::spawn(move|| {});

        // It is already finished, but we try to resume it
        // Idealy it would come back immediately
        Coroutine::resume(&coro);

        // Again?
        Coroutine::resume(&coro);
    }

    #[test]
    fn test_coroutine_resume_itself() {
        Coroutine::spawn(move|| {
            // Resume itself
            Coroutine::current().resume();
        });
    }

    #[test]
    fn test_coroutine_yield_in_main() {
        Coroutine::sched();
    }

    #[bench]
    fn bench_coroutine_spawning_with_recycle(b: &mut Bencher) {
        b.iter(|| {
            Coroutine::spawn(move|| {});
        });
    }

    #[bench]
    fn bench_normal_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;

            let (tx, rx) = channel();

            let mut result = 0;
            for _ in 0..MAX_NUMBER {
                tx.send(1).unwrap();
                result += rx.recv().unwrap();
            }
            assert_eq!(result, MAX_NUMBER);
        });
    }

    #[bench]
    fn bench_coroutine_counting(b: &mut Bencher) {
        b.iter(|| {
            const MAX_NUMBER: usize = 100;
            let (tx, rx) = channel();

            let coro = Coroutine::spawn(move|| {
                for _ in 0..MAX_NUMBER {
                    tx.send(1).unwrap();
                    Coroutine::sched();
                }
            });

            let result = rx.iter().fold(0, |a, b| {
                Coroutine::resume(&coro);
                a + b
            });
            assert_eq!(result, MAX_NUMBER);
        });
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
use rand::Rand;

use rand::Rng;

/// Modifies ind1 and ind2 in-place.
pub fn simulated_binary_crossover<R: Rng>(rng: &mut R,
                                          ind1: &mut [f32],
                                          ind2: &mut [f32],
                                          eta: f32) {
    assert!(ind1.len() == ind2.len());

    let mut iter = ind1.iter_mut().zip(ind2.iter_mut());
    for (x1, x2) in iter {
        let (old_x1, old_x2) = (*x1, *x2);

        let u = rng.gen::<f32>();
        debug_assert!(u >= 0.0 && u < 1.0);

        let beta = if u <= 0.5 {
                       2.0 * u
                   } else {
                       1.0 / (2.0 * (1.0 - u))
                   }
                   .powf(1.0 / (eta + 1.0));

        *x1 = 0.5 * (((1.0 + beta) * old_x1) + ((1.0 - beta) * old_x2));
        *x2 = 0.5 * (((1.0 - beta) * old_x1) + ((1.0 + beta) * old_x2));
    }
}

#[test]
fn test_sbx() {
    let mut rng = ::rand::isaac::Isaac64Rng::new_unseeded();

    let parent1 = vec![1.0, 2.0];
    let parent2 = vec![100.0, 1.0];

    let mut child1 = parent1.clone();
    let mut child2 = parent2.clone();

    simulated_binary_crossover(&mut rng, &mut child1[..], &mut child2[..], 2.0);

    let beta: Vec<_> = (0..2)
                           .map(|i| (child2[i] - child1[i]) / (parent2[i] - parent1[i]))
                           .collect();

    println!("({:?}, {:?}) => ({:?}, {:?})",
             parent1,
             parent2,
             child1,
             child2);
    println!("beta: {:?}", beta);
}
<commit_msg>Factor out some code<commit_after>#[cfg(test)]
use rand::Rand;

use rand::Rng;

#[inline]
fn sbx_beta(u: f32, eta: f32) -> f32 {
    debug_assert!(u >= 0.0 && u < 1.0);

    if u <= 0.5 {
        2.0 * u
    } else {
        1.0 / (2.0 * (1.0 - u))
    }
    .powf(1.0 / (eta + 1.0))
}

/// Modifies ind1 and ind2 in-place.
/// Reference: http://www.iitk.ac.in/kangal/papers/k2011017.pdf
pub fn simulated_binary_crossover<R: Rng>(rng: &mut R,
                                          ind1: &mut [f32],
                                          ind2: &mut [f32],
                                          eta: f32) {
    assert!(ind1.len() == ind2.len());

    let mut iter = ind1.iter_mut().zip(ind2.iter_mut());
    for (x1, x2) in iter {
        let (p1, p2) = (*x1, *x2);

        let u = rng.gen::<f32>();
        let beta = sbx_beta(u, eta);

        let c1 = 0.5 * (((1.0 + beta) * p1) + ((1.0 - beta) * p2));
        let c2 = 0.5 * (((1.0 - beta) * p1) + ((1.0 + beta) * p2));

        *x1 = c1;
        *x2 = c2;
    }
}

#[test]
fn test_sbx() {
    let mut rng = ::rand::isaac::Isaac64Rng::new_unseeded();

    let parent1 = vec![1.0, 2.0];
    let parent2 = vec![100.0, 1.0];

    let mut child1 = parent1.clone();
    let mut child2 = parent2.clone();

    simulated_binary_crossover(&mut rng, &mut child1[..], &mut child2[..], 2.0);

    let beta: Vec<_> = (0..2)
                           .map(|i| (child2[i] - child1[i]) / (parent2[i] - parent1[i]))
                           .collect();

    println!("({:?}, {:?}) => ({:?}, {:?})",
             parent1,
             parent2,
             child1,
             child2);
    println!("beta: {:?}", beta);
}
<|endoftext|>


<commit_before>use std::num;
use std::num::{One, Zero};

use super::super::matrix::*;
use super::super::util::{alloc_dirty_vec};

pub struct LUDecomposition<T> {
  lu : Matrix<T>,
  pospivsign : bool,
  piv : ~[uint]
}

// Ported from JAMA.
// LU Decomposition.
//
// For an m-by-n matrix A with m >= n, the LU decomposition is an m-by-n
// unit lower triangular matrix L, an n-by-n upper triangular matrix U,
// and a permutation vector piv of length m so that A(piv,:) = L*U.
// If m < n, then L is m-by-m and U is m-by-n.
//
// The LU decompostion with pivoting always exists, even if the matrix is
// singular. The primary use of the LU decomposition is in the solution of
// square systems of simultaneous linear equations. This will fail if the
// matrix is singular.
impl<T : Add<T, T> + Sub<T, T> + Mul<T, T> + Div<T, T> + Neg<T> + Eq + Ord + ApproxEq<T> + One + Zero + Clone + Signed> LUDecomposition<T> {
  pub fn new(a : &Matrix<T>) -> LUDecomposition<T> {
    // Use a "left-looking", dot-product, Crout/Doolittle algorithm.
    let mut ludata = a.data.clone();
    let m = a.noRows as int;
    let n = a.noCols as int;
    let mut pivdata = alloc_dirty_vec(m as uint);
    for i in range(0, m) {
      pivdata[i] = i as uint;
    }

    let mut pospivsign = true;
    let mut lucolj = alloc_dirty_vec(m as uint);

    // Outer loop.
    for j in range(0, n) {
      // Make a copy of the j-th column to localize references.
      for i in range(0, m) {
        lucolj[i] = ludata[i * n + j].clone();
      }

      // Apply previous transformations.
      for i in range(0, m) {
        let lurowiIdx = i * n;

        // Most of the time is spent in the following dot product.
        let kmax = num::min(i, j);
        let mut s : T = num::zero();
        for k in range(0, kmax) {
          s = s + ludata[lurowiIdx + k] * lucolj[k];
        }

        lucolj[i] = lucolj[i] - s;
        ludata[lurowiIdx + j] = lucolj[i].clone();
      }

      // Find pivot and exchange if necessary.
      let mut p = j;
      for i in range(j + 1, m) {
        if(num::abs(lucolj[i].clone()) > num::abs(lucolj[p].clone())) {
          p = i;
        }
      }

      if(p != j) {
        for k in range(0, n) {
          let t = ludata[p * n + k].clone();
          ludata[p * n + k] = ludata[j * n + k].clone();
          ludata[j * n + k] = t;
        }

        let k = pivdata[p];
        pivdata[p] = pivdata[j];
        pivdata[j] = k;
        pospivsign = !pospivsign;
      }

      // Compute multipliers.
      if((j < m) && (ludata[j * n + j] != num::zero())) {
        for i in range(j + 1, m) {
          ludata[i * n + j] = ludata[i * n + j] / ludata[j * n + j];
        }
      }
    }

    LUDecomposition {
      lu : Matrix { noRows : m as uint, noCols : n as uint, data : ludata },
      pospivsign : pospivsign,
      piv : pivdata
    }
  }

  pub fn is_singular(&self) -> bool {
    !self.is_non_singular()
  }

  pub fn is_non_singular(&self) -> bool {
    let n = self.lu.noCols;
    for j in range(0, n) {
      if(self.lu.data[j * n + j] == num::zero()) {
        return false;
      }
    }
    true
  }

  pub fn get_l(&self) -> Matrix<T> {
    let m = self.lu.noRows;
    let n = if self.lu.noRows >= self.lu.noCols { self.lu.noCols } else { self.lu.noRows };
    let mut ldata = alloc_dirty_vec(m * n);
    for i in range(0, m) {
      for j in range(0, n) {
        ldata[i * n + j] = if(i > j) {
                             self.lu.data[i * self.lu.noCols + j].clone()
                           } else if(i == j) {
                             num::one()
                           } else {
                             num::zero()
                           }
      }
    }
    Matrix { noRows : m, noCols : n, data : ldata }
  }

  pub fn get_u(&self) -> Matrix<T> {
    let m = if self.lu.noRows >= self.lu.noCols { self.lu.noCols as int } else { self.lu.noRows as int };
    let n = self.lu.noCols as int;
    let mut udata = alloc_dirty_vec((m * n) as uint);
    for i in range(0, m) {
      for j in range(0, n) {
        udata[i * n + j] = if(i <= j) { self.lu.data[i * n + j].clone() } else { num::zero() };
      }
    }
    Matrix { noRows : m as uint, noCols : n as uint, data : udata }
  }

  pub fn get_p(&self) -> Matrix<T> {
    let len = self.piv.len();
    id(len, len).permute_rows(self.piv)
  }

  pub fn get_piv<'lt>(&'lt self) -> &'lt ~[uint] { &self.piv }

  pub fn det(&self) -> T {
    assert!(self.lu.noRows == self.lu.noCols);
    let n = self.lu.noCols as int;
    let mut d = if self.pospivsign { num::one::<T>() } else { - num::one::<T>() };
    for j in range(0, n) {
      d = d * self.lu.data[j * n + j];
    }
    d
  }

  // Solve A*X = B
  // B   A Matrix with as many rows as A and any number of columns.
  // Returns X so that L*U*X = B(piv,:)
  pub fn solve(&self, b : &Matrix<T>) -> Option<Matrix<T>> {
    let m = self.lu.noRows as int;
    let n = self.lu.noCols as int;
    assert!(b.noRows == m as uint);
    if(!self.is_non_singular()) {
      return None
    }

    // Copy right hand side with pivoting
    let nx = b.noCols as int;
    let mut xdata = alloc_dirty_vec((m * nx) as uint);
    let mut destIdx = 0;
    for i in range(0, self.piv.len()) {
      for j in range(0, nx) {
        xdata[destIdx] = b.data[(self.piv[i] as int) * (b.noCols as int) + j].clone();
        destIdx += 1;
      }
    }

    // Solve L*Y = B(piv,:)
    for k in range(0, n) {
      for i in range(k + 1, n) {
        for j in range(0, nx) {
          xdata[i * nx + j] = xdata[i * nx + j] - xdata[k * nx + j] * self.lu.data[i * (self.lu.noCols as int) + k];
        }
      }
    }

    // Solve U*X = Y;
    for k in range(0, n).invert() {
      for j in range(0, nx) {
        xdata[k * nx + j] = xdata[k * nx + j] / self.lu.data[k * (self.lu.noCols as int) + k];
      }
      for i in range(0, k) {
        for j in range(0, nx) {
          xdata[i * nx + j] = xdata[i * nx + j] - xdata[k * nx + j] * self.lu.data[i * (self.lu.noCols as int) + k];
        }
      }
    }

    Some(Matrix {
      noRows : self.lu.noRows,
      noCols : b.noCols,
      data : xdata
    })
  }
}

#[test]
fn test_lu__square() {
  let a = matrix(3, 3, ~[1.0, 2.0, 0.0, 3.0, 6.0, -1.0, 1.0, 2.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn test_lu2__m_over_n() {
  let a = matrix(3, 2, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn test_lu2__m_under_n() {
  let a = matrix(2, 3, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn lu_solve_test() {
  let a = matrix(3, 3, ~[2.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0, 3.0]);
  assert!(lu.solve(&b).unwrap().approx_eq(&vector(~[-1.0, 3.0, 3.0])));
}

#[test]
#[should_fail]
fn lu_solve_test__incompatible() {
  let a = matrix(3, 3, ~[2.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0, 3.0, 4.0]);
  let _ = lu.solve(&b);
}

#[test]
fn lu_solve_test__singular() {
  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0]);
  assert!(lu.solve(&b).is_none());
}

#[test]
fn lu_is_singular_test() {
  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.is_singular());

  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(!lu.is_singular());
}

#[test]
fn lu_is_non_singular_test() {
  let a = matrix(2, 2, ~[4.0, 8.0, 3.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.is_non_singular());

  let a = matrix(2, 2, ~[4.0, 6.0, 2.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  assert!(!lu.is_non_singular());
}

#[test]
fn lu_det_test() {
  let a = matrix(2, 2, ~[4.0, 8.0, 3.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.det() == -8.0);

  let a = matrix(2, 2, ~[4.0, 8.0, 2.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.det() == 0.0);
}

#[test]
#[should_fail]
fn lu_det_test__not_square() {
  let a = matrix(2, 3, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let _ = lu.det();
}
<commit_msg>More documentation for LU decomposition.<commit_after>use std::num;
use std::num::{One, Zero};

use super::super::matrix::*;
use super::super::util::{alloc_dirty_vec};

// Ported from JAMA.
// LU Decomposition.
//
// For an m-by-n matrix A with m >= n, the LU decomposition is an m-by-n
// unit lower triangular matrix L, an n-by-n upper triangular matrix U,
// and a permutation vector piv of length m so that A(piv,:) = L*U.
// If m < n, then L is m-by-m and U is m-by-n.
//
// The LU decompostion with pivoting always exists, even if the matrix is
// singular. The primary use of the LU decomposition is in the solution of
// square systems of simultaneous linear equations. This will fail if the
// matrix is singular.
//
// LU Decomposition:
// A = LU,
//   where L is a lower triangular matrix
//   and U is an upper triangular matrix.
//
// Step 1: Set diagonals of L to 1:
// Step 2: Solve columns of L and U in order from left to right.
//         For a specific column, solve in order from top to bottom, first solving the column for U and then for L.
//         For solving element e[i][j], write out the dot between i:th row of L and j:th column of U. By following
//         this specific order, you'll notice that by the time you are solving for that specific element, you have
//         already solved all the other variables in the equation, thus you just need to plug in the numbers to
//         solve the element.
//
//  Example: 3x3 matrix LU decomposition:
//    l11   0   0   u11 u12 u13   a11 a12 a13
//    l21 l22   0 *   0 u22 u23 = a21 a22 a23
//    l31 l32 l33     0   0 u33   a31 a32 a33
//
//  Step 1:
//    l11 = 1
//    l22 = 1
//    l33 = 1
//
//  Step 2:
//    u11 = a11 / l11
//    l21 = a21 / u11
//    l31 = a31 / u11
//
//    u12 = a12 / l11
//    u22 = (a22 - u12 * l21) / l22
//    l32 = (a32 - l31 * u12) / u22
//
//    u13 = a13 / l11
//    u23 = (a23 - u13 * l21) / l22
//    u33 = (a33 - u13 * l31 - u23 * l32) / l33
//
//  In general, from above we can derive the following equations for solving the specific elements:
//  l[i][i] = 1
//
//  u[i][j] = 1 / l[i][i] * (a[i][j] - <l[0 .. (i - 1)][:], u[:][0 .. (i - 1)]>)
//          = a[i][j] - <l[0 .. (i - 1)][:], u[:][0 .. (i - 1)]>                  , (as l[i][i] = 1).
//
//  l[i][j] = 1 / u[j][j] * (a[i][j] - <l[0 .. (j - 1)][:], u[:][0 .. (j - 1)]>)
//
//  The above algorithm is generally known as Crout's algorithm.
//
//  Examining the above equations closely, we can determine that the algorithm works as long as u[j][j]
//  is not zero (which would cause a division by zero). We can generalize this by adding a permutation
//  matrix P, giving us LUP decomposition, which always exists:
//
//  PA = LU,
//    where P is a permutation matrix,
//    L is a lower triangular matrix
//    and U is an upper triangular matrix.
//
//  Applying pivoting also makes the algorithm numerically stable. We can change the algorithm to do
//  (partial) pivoting in the following way:
//    1. When calculating l[i][j], do not perform the division with u[j][j].
//    2. Once you have calculated all the values for the column j, find the maximum value of lu[i..(n-1)][j].
//    3. Swap the pivot row with the row with the maximum value, so that the maximum value is at the pivot row.
//    4. Perform the division steps of the original algorithm with the new value of u[j][j] (the maximum).
//
//  LUP decomposition can be used to solve a set of linear equations 'Ax = b' in the following manner
//  (assuming A is invertible):
//    Ax = b
//    PAx = Pb
//    LUx = Pb
//    (LU)^-1 LUx = (LU)^-1 Pb
//    x = (LU)^-1 Pb
//      = U^-1 L^-1 Pb
//

pub struct LUDecomposition<T> {
  // U is stored in diagonals and above. Elements below diagonals are zero for U.
  // L is stored below diagonals. Diagonal elements are one for U and elements above diagonals are zero.
  lu : Matrix<T>,
  pospivsign : bool,
  piv : ~[uint]
}

impl<T : Add<T, T> + Sub<T, T> + Mul<T, T> + Div<T, T> + Neg<T> + Eq + Ord + ApproxEq<T> + One + Zero + Clone + Signed> LUDecomposition<T> {
  pub fn new(a : &Matrix<T>) -> LUDecomposition<T> {
    // Use a "left-looking", dot-product, Crout/Doolittle algorithm.
    let mut ludata = a.data.clone();
    let m = a.noRows as int;
    let n = a.noCols as int;
    let mut pivdata = alloc_dirty_vec(m as uint);
    for i in range(0, m) {
      pivdata[i] = i as uint;
    }

    let mut pospivsign = true;
    let mut lucolj = alloc_dirty_vec(m as uint);

    // Solve columns 0..(n-1) of L and U in order.
    for j in range(0, n) {
      // Solve column j of L and U.

      // Make a copy of the j-th column to localize references.
      for i in range(0, m) {
        lucolj[i] = ludata[i * n + j].clone();
      }

      // Apply previous transformations.
      for i in range(0, m) {
        let lurowiIdx = i * n;

        // Most of the time is spent in the following dot product.
        let kmax = num::min(i, j);
        let mut s : T = num::zero();
        for k in range(0, kmax) {
          s = s + ludata[lurowiIdx + k] * lucolj[k];
        }

        lucolj[i] = lucolj[i] - s;
        ludata[lurowiIdx + j] = lucolj[i].clone();
      }

      // Find pivot and exchange if necessary.
      let mut p = j;
      for i in range(j + 1, m) {
        if(num::abs(lucolj[i].clone()) > num::abs(lucolj[p].clone())) {
          p = i;
        }
      }

      if(p != j) {
        for k in range(0, n) {
          let t = ludata[p * n + k].clone();
          ludata[p * n + k] = ludata[j * n + k].clone();
          ludata[j * n + k] = t;
        }

        let k = pivdata[p];
        pivdata[p] = pivdata[j];
        pivdata[j] = k;
        pospivsign = !pospivsign;
      }

      // Compute multipliers.
      if((j < m) && (ludata[j * n + j] != num::zero())) {
        for i in range(j + 1, m) {
          ludata[i * n + j] = ludata[i * n + j] / ludata[j * n + j];
        }
      }
    }

    LUDecomposition {
      lu : Matrix { noRows : m as uint, noCols : n as uint, data : ludata },
      pospivsign : pospivsign,
      piv : pivdata
    }
  }

  pub fn is_singular(&self) -> bool {
    !self.is_non_singular()
  }

  pub fn is_non_singular(&self) -> bool {
    let n = self.lu.noCols;
    for j in range(0, n) {
      if(self.lu.data[j * n + j] == num::zero()) {
        return false;
      }
    }
    true
  }

  pub fn get_l(&self) -> Matrix<T> {
    // L is stored below diagonals. Diagonal elements are one for U and elements above diagonals are zero.
    let m = self.lu.noRows;
    let n = if self.lu.noRows >= self.lu.noCols { self.lu.noCols } else { self.lu.noRows };
    let mut ldata = alloc_dirty_vec(m * n);
    for i in range(0, m) {
      for j in range(0, n) {
        ldata[i * n + j] = if(i > j) {
                             self.lu.data[i * self.lu.noCols + j].clone()
                           } else if(i == j) {
                             num::one()
                           } else {
                             num::zero()
                           }
      }
    }
    Matrix { noRows : m, noCols : n, data : ldata }
  }

  pub fn get_u(&self) -> Matrix<T> {
    // U is stored in diagonals and above. Elements below diagonals are zero for U.
    let m = if self.lu.noRows >= self.lu.noCols { self.lu.noCols as int } else { self.lu.noRows as int };
    let n = self.lu.noCols as int;
    let mut udata = alloc_dirty_vec((m * n) as uint);
    for i in range(0, m) {
      for j in range(0, n) {
        udata[i * n + j] = if(i <= j) { self.lu.data[i * n + j].clone() } else { num::zero() };
      }
    }
    Matrix { noRows : m as uint, noCols : n as uint, data : udata }
  }

  pub fn get_p(&self) -> Matrix<T> {
    let len = self.piv.len();
    id(len, len).permute_rows(self.piv)
  }

  pub fn get_piv<'lt>(&'lt self) -> &'lt ~[uint] { &self.piv }

  pub fn det(&self) -> T {
    assert!(self.lu.noRows == self.lu.noCols);
    let n = self.lu.noCols as int;
    let mut d = if self.pospivsign { num::one::<T>() } else { - num::one::<T>() };
    for j in range(0, n) {
      d = d * self.lu.data[j * n + j];
    }
    d
  }

  // Solve A*X = B
  // B   A Matrix with as many rows as A and any number of columns.
  // Returns X so that L*U*X = B(piv,:)
  pub fn solve(&self, b : &Matrix<T>) -> Option<Matrix<T>> {
    let m = self.lu.noRows as int;
    let n = self.lu.noCols as int;
    assert!(b.noRows == m as uint);
    if(!self.is_non_singular()) {
      return None
    }

    // Copy right hand side with pivoting
    let nx = b.noCols as int;
    let mut xdata = alloc_dirty_vec((m * nx) as uint);
    let mut destIdx = 0;
    for i in range(0, self.piv.len()) {
      for j in range(0, nx) {
        xdata[destIdx] = b.data[(self.piv[i] as int) * (b.noCols as int) + j].clone();
        destIdx += 1;
      }
    }

    // Solve L*Y = B(piv,:)
    for k in range(0, n) {
      for i in range(k + 1, n) {
        for j in range(0, nx) {
          xdata[i * nx + j] = xdata[i * nx + j] - xdata[k * nx + j] * self.lu.data[i * (self.lu.noCols as int) + k];
        }
      }
    }

    // Solve U*X = Y;
    for k in range(0, n).invert() {
      for j in range(0, nx) {
        xdata[k * nx + j] = xdata[k * nx + j] / self.lu.data[k * (self.lu.noCols as int) + k];
      }
      for i in range(0, k) {
        for j in range(0, nx) {
          xdata[i * nx + j] = xdata[i * nx + j] - xdata[k * nx + j] * self.lu.data[i * (self.lu.noCols as int) + k];
        }
      }
    }

    Some(Matrix {
      noRows : self.lu.noRows,
      noCols : b.noCols,
      data : xdata
    })
  }
}

#[test]
fn test_lu__square() {
  let a = matrix(3, 3, ~[1.0, 2.0, 0.0, 3.0, 6.0, -1.0, 1.0, 2.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn test_lu2__m_over_n() {
  let a = matrix(3, 2, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn test_lu2__m_under_n() {
  let a = matrix(2, 3, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let l = lu.get_l();
  let u = lu.get_u();
  let p = lu.get_p();
  assert!(l * u == p * a);
}

#[test]
fn lu_solve_test() {
  let a = matrix(3, 3, ~[2.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0, 3.0]);
  assert!(lu.solve(&b).unwrap().approx_eq(&vector(~[-1.0, 3.0, 3.0])));
}

#[test]
#[should_fail]
fn lu_solve_test__incompatible() {
  let a = matrix(3, 3, ~[2.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0, 3.0, 4.0]);
  let _ = lu.solve(&b);
}

#[test]
fn lu_solve_test__singular() {
  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  let b = vector(~[1.0, 2.0]);
  assert!(lu.solve(&b).is_none());
}

#[test]
fn lu_is_singular_test() {
  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.is_singular());

  let a = matrix(2, 2, ~[2.0, 6.0, 1.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(!lu.is_singular());
}

#[test]
fn lu_is_non_singular_test() {
  let a = matrix(2, 2, ~[4.0, 8.0, 3.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.is_non_singular());

  let a = matrix(2, 2, ~[4.0, 6.0, 2.0, 3.0]);
  let lu = LUDecomposition::new(&a);
  assert!(!lu.is_non_singular());
}

#[test]
fn lu_det_test() {
  let a = matrix(2, 2, ~[4.0, 8.0, 3.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.det() == -8.0);

  let a = matrix(2, 2, ~[4.0, 8.0, 2.0, 4.0]);
  let lu = LUDecomposition::new(&a);
  assert!(lu.det() == 0.0);
}

#[test]
#[should_fail]
fn lu_det_test__not_square() {
  let a = matrix(2, 3, ~[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
  let lu = LUDecomposition::new(&a);
  let _ = lu.det();
}
<|endoftext|>


<commit_before>use base64;
use crypto::digest::Digest;
use crypto::mac::Mac;
use crypto;
use futures::sync::mpsc;
use futures::{Future, Stream, Poll};
use hyper::server::{Service, Request, Response, Http};
use hyper::{self, Get, Post, StatusCode};

#[cfg(feature = "with-dns-sd")]
use dns_sd::DNSService;

#[cfg(not(feature = "with-dns-sd"))]
use mdns;

use num_bigint::BigUint;
use rand;
use std::collections::BTreeMap;
use std::io;
use std::sync::Arc;
use tokio_core::reactor::Handle;
use url;

use core::diffie_hellman::{DH_GENERATOR, DH_PRIME};
use core::authentication::Credentials;
use core::util;
use core::config::ConnectConfig;

#[derive(Clone)]
struct Discovery(Arc<DiscoveryInner>);
struct DiscoveryInner {
    config: ConnectConfig,
    device_id: String,
    private_key: BigUint,
    public_key: BigUint,
    tx: mpsc::UnboundedSender<Credentials>,
}

impl Discovery {
    fn new(config: ConnectConfig, device_id: String)
        -> (Discovery, mpsc::UnboundedReceiver<Credentials>)
    {
        let (tx, rx) = mpsc::unbounded();

        let key_data = util::rand_vec(&mut rand::thread_rng(), 95);
        let private_key = BigUint::from_bytes_be(&key_data);
        let public_key = util::powm(&DH_GENERATOR, &private_key, &DH_PRIME);

        let discovery = Discovery(Arc::new(DiscoveryInner {
            config: config,
            device_id: device_id,
            private_key: private_key,
            public_key: public_key,
            tx: tx,
        }));

        (discovery, rx)
    }
}

impl Discovery {
    fn handle_get_info(&self, _params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let public_key = self.0.public_key.to_bytes_be();
        let public_key = base64::encode(&public_key);

        let result = json!({
            "status": 101,
            "statusString": "ERROR-OK",
            "spotifyError": 0,
            "version": "2.1.0",
            "deviceID": (self.0.device_id),
            "remoteName": (self.0.config.name),
            "activeUser": "",
            "publicKey": (public_key),
            "deviceType": (self.0.config.device_type.to_string().to_uppercase()),
            "libraryVersion": "0.1.0",
            "accountReq": "PREMIUM",
            "brandDisplayName": "librespot",
            "modelDisplayName": "librespot",
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn handle_add_user(&self, params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let username = params.get("userName").unwrap();
        let encrypted_blob = params.get("blob").unwrap();
        let client_key = params.get("clientKey").unwrap();

        let encrypted_blob = base64::decode(encrypted_blob).unwrap();

        let client_key = base64::decode(client_key).unwrap();
        let client_key = BigUint::from_bytes_be(&client_key);

        let shared_key = util::powm(&client_key, &self.0.private_key, &DH_PRIME);

        let iv = &encrypted_blob[0..16];
        let encrypted = &encrypted_blob[16..encrypted_blob.len() - 20];
        let cksum = &encrypted_blob[encrypted_blob.len() - 20..encrypted_blob.len()];

        let base_key = {
            let mut data = [0u8; 20];
            let mut h = crypto::sha1::Sha1::new();
            h.input(&shared_key.to_bytes_be());
            h.result(&mut data);
            data[..16].to_owned()
        };

        let checksum_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"checksum");
            h.result().code().to_owned()
        };

        let encryption_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"encryption");
            h.result().code().to_owned()
        };

        let mac = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &checksum_key);
            h.input(encrypted);
            h.result().code().to_owned()
        };

        assert_eq!(&mac[..], cksum);

        let decrypted = {
            let mut data = vec![0u8; encrypted.len()];
            let mut cipher = crypto::aes::ctr(crypto::aes::KeySize::KeySize128,
                                              &encryption_key[0..16], iv);
            cipher.process(encrypted, &mut data);
            String::from_utf8(data).unwrap()
        };

        let credentials = Credentials::with_blob(username.to_owned(), &decrypted, &self.0.device_id);

        self.0.tx.unbounded_send(credentials).unwrap();

        let result = json!({
            "status": 101,
            "spotifyError": 0,
            "statusString": "ERROR-OK"
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn not_found(&self)
        -> ::futures::Finished<Response, hyper::Error>
    {
        ::futures::finished(Response::new().with_status(StatusCode::NotFound))
    }
}

impl Service for Discovery {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item = Response, Error = hyper::Error>>;

    fn call(&self, request: Request) -> Self::Future {
        let mut params = BTreeMap::new();

        let (method, uri, _, _, body) = request.deconstruct();
        if let Some(query) = uri.query() {
            params.extend(url::form_urlencoded::parse(query.as_bytes()).into_owned());
        }

        if method != Get {
            debug!("{:?} {:?} {:?}", method, uri.path(), params);
        }

        let this = self.clone();
        Box::new(body.fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(chunk.as_ref());
            Ok::<_, hyper::Error>(acc)
        }).map(move |body| {
            params.extend(url::form_urlencoded::parse(&body).into_owned());
            params
        }).and_then(move |params| {
            match (method, params.get("action").map(AsRef::as_ref)) {
                (Get, Some("getInfo")) => this.handle_get_info(&params),
                (Post, Some("addUser")) => this.handle_add_user(&params),
                _ => this.not_found(),
            }
        }))
    }
}

#[cfg(feature = "with-dns-sd")]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: DNSService,
}

#[cfg(not(feature = "with-dns-sd"))]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: mdns::Service,
}

pub fn discovery(handle: &Handle, config: ConnectConfig, device_id: String)
    -> io::Result<DiscoveryStream>
{
    let (discovery, creds_rx) = Discovery::new(config.clone(), device_id);

    let serve = {
        let http = Http::new();
        http.serve_addr_handle(&"0.0.0.0:0".parse().unwrap(), &handle, move || Ok(discovery.clone())).unwrap()
    };

    #[cfg(feature = "with-dns-sd")]
    let port = serve.incoming_ref().local_addr().port();

    let server_future = {
        let handle = handle.clone();
        serve.for_each(move |connection| {
                handle.spawn(connection.then(|_| Ok(())));
                Ok(())
            })
            .then(|_| Ok(()))
    };
    handle.spawn(server_future);

    #[cfg(feature = "with-dns-sd")]
    let svc = DNSService::register(Some(&*config.name),
       "_spotify-connect._tcp",
       None,
       None,
       port,
       &["VERSION=1.0", "CPath=/"]).unwrap();

    #[cfg(not(feature = "with-dns-sd"))]
    let addr = serve.incoming_ref().local_addr();
    #[cfg(not(feature = "with-dns-sd"))]
    let responder = mdns::Responder::spawn(&handle)?;
    #[cfg(not(feature = "with-dns-sd"))]
    let svc = responder.register(
        "_spotify-connect._tcp".to_owned(),
        config.name,
        addr.port(),
        &["VERSION=1.0", "CPath=/"]);

    Ok(DiscoveryStream {
        credentials: creds_rx,
        _svc: svc,
    })
}

impl Stream for DiscoveryStream {
    type Item = Credentials;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.credentials.poll()
    }
}
<commit_msg>fix addr in wrong place<commit_after>use base64;
use crypto::digest::Digest;
use crypto::mac::Mac;
use crypto;
use futures::sync::mpsc;
use futures::{Future, Stream, Poll};
use hyper::server::{Service, Request, Response, Http};
use hyper::{self, Get, Post, StatusCode};

#[cfg(feature = "with-dns-sd")]
use dns_sd::DNSService;

#[cfg(not(feature = "with-dns-sd"))]
use mdns;

use num_bigint::BigUint;
use rand;
use std::collections::BTreeMap;
use std::io;
use std::sync::Arc;
use tokio_core::reactor::Handle;
use url;

use core::diffie_hellman::{DH_GENERATOR, DH_PRIME};
use core::authentication::Credentials;
use core::util;
use core::config::ConnectConfig;

#[derive(Clone)]
struct Discovery(Arc<DiscoveryInner>);
struct DiscoveryInner {
    config: ConnectConfig,
    device_id: String,
    private_key: BigUint,
    public_key: BigUint,
    tx: mpsc::UnboundedSender<Credentials>,
}

impl Discovery {
    fn new(config: ConnectConfig, device_id: String)
        -> (Discovery, mpsc::UnboundedReceiver<Credentials>)
    {
        let (tx, rx) = mpsc::unbounded();

        let key_data = util::rand_vec(&mut rand::thread_rng(), 95);
        let private_key = BigUint::from_bytes_be(&key_data);
        let public_key = util::powm(&DH_GENERATOR, &private_key, &DH_PRIME);

        let discovery = Discovery(Arc::new(DiscoveryInner {
            config: config,
            device_id: device_id,
            private_key: private_key,
            public_key: public_key,
            tx: tx,
        }));

        (discovery, rx)
    }
}

impl Discovery {
    fn handle_get_info(&self, _params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let public_key = self.0.public_key.to_bytes_be();
        let public_key = base64::encode(&public_key);

        let result = json!({
            "status": 101,
            "statusString": "ERROR-OK",
            "spotifyError": 0,
            "version": "2.1.0",
            "deviceID": (self.0.device_id),
            "remoteName": (self.0.config.name),
            "activeUser": "",
            "publicKey": (public_key),
            "deviceType": (self.0.config.device_type.to_string().to_uppercase()),
            "libraryVersion": "0.1.0",
            "accountReq": "PREMIUM",
            "brandDisplayName": "librespot",
            "modelDisplayName": "librespot",
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn handle_add_user(&self, params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let username = params.get("userName").unwrap();
        let encrypted_blob = params.get("blob").unwrap();
        let client_key = params.get("clientKey").unwrap();

        let encrypted_blob = base64::decode(encrypted_blob).unwrap();

        let client_key = base64::decode(client_key).unwrap();
        let client_key = BigUint::from_bytes_be(&client_key);

        let shared_key = util::powm(&client_key, &self.0.private_key, &DH_PRIME);

        let iv = &encrypted_blob[0..16];
        let encrypted = &encrypted_blob[16..encrypted_blob.len() - 20];
        let cksum = &encrypted_blob[encrypted_blob.len() - 20..encrypted_blob.len()];

        let base_key = {
            let mut data = [0u8; 20];
            let mut h = crypto::sha1::Sha1::new();
            h.input(&shared_key.to_bytes_be());
            h.result(&mut data);
            data[..16].to_owned()
        };

        let checksum_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"checksum");
            h.result().code().to_owned()
        };

        let encryption_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"encryption");
            h.result().code().to_owned()
        };

        let mac = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &checksum_key);
            h.input(encrypted);
            h.result().code().to_owned()
        };

        assert_eq!(&mac[..], cksum);

        let decrypted = {
            let mut data = vec![0u8; encrypted.len()];
            let mut cipher = crypto::aes::ctr(crypto::aes::KeySize::KeySize128,
                                              &encryption_key[0..16], iv);
            cipher.process(encrypted, &mut data);
            String::from_utf8(data).unwrap()
        };

        let credentials = Credentials::with_blob(username.to_owned(), &decrypted, &self.0.device_id);

        self.0.tx.unbounded_send(credentials).unwrap();

        let result = json!({
            "status": 101,
            "spotifyError": 0,
            "statusString": "ERROR-OK"
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn not_found(&self)
        -> ::futures::Finished<Response, hyper::Error>
    {
        ::futures::finished(Response::new().with_status(StatusCode::NotFound))
    }
}

impl Service for Discovery {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item = Response, Error = hyper::Error>>;

    fn call(&self, request: Request) -> Self::Future {
        let mut params = BTreeMap::new();

        let (method, uri, _, _, body) = request.deconstruct();
        if let Some(query) = uri.query() {
            params.extend(url::form_urlencoded::parse(query.as_bytes()).into_owned());
        }

        if method != Get {
            debug!("{:?} {:?} {:?}", method, uri.path(), params);
        }

        let this = self.clone();
        Box::new(body.fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(chunk.as_ref());
            Ok::<_, hyper::Error>(acc)
        }).map(move |body| {
            params.extend(url::form_urlencoded::parse(&body).into_owned());
            params
        }).and_then(move |params| {
            match (method, params.get("action").map(AsRef::as_ref)) {
                (Get, Some("getInfo")) => this.handle_get_info(&params),
                (Post, Some("addUser")) => this.handle_add_user(&params),
                _ => this.not_found(),
            }
        }))
    }
}

#[cfg(feature = "with-dns-sd")]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: DNSService,
}

#[cfg(not(feature = "with-dns-sd"))]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: mdns::Service,
}

pub fn discovery(handle: &Handle, config: ConnectConfig, device_id: String)
    -> io::Result<DiscoveryStream>
{
    let (discovery, creds_rx) = Discovery::new(config.clone(), device_id);

    let serve = {
        let http = Http::new();
        http.serve_addr_handle(&"0.0.0.0:0".parse().unwrap(), &handle, move || Ok(discovery.clone())).unwrap()
    };

    #[cfg(feature = "with-dns-sd")]
    let port = serve.incoming_ref().local_addr().port();

    #[cfg(not(feature = "with-dns-sd"))]
    let addr = serve.incoming_ref().local_addr();

    let server_future = {
        let handle = handle.clone();
        serve.for_each(move |connection| {
                handle.spawn(connection.then(|_| Ok(())));
                Ok(())
            })
            .then(|_| Ok(()))
    };
    handle.spawn(server_future);

    #[cfg(feature = "with-dns-sd")]
    let svc = DNSService::register(Some(&*config.name),
       "_spotify-connect._tcp",
       None,
       None,
       port,
       &["VERSION=1.0", "CPath=/"]).unwrap();

    #[cfg(not(feature = "with-dns-sd"))]
    let responder = mdns::Responder::spawn(&handle)?;

    #[cfg(not(feature = "with-dns-sd"))]
    let svc = responder.register(
        "_spotify-connect._tcp".to_owned(),
        config.name,
        addr.port(),
        &["VERSION=1.0", "CPath=/"]);

    Ok(DiscoveryStream {
        credentials: creds_rx,
        _svc: svc,
    })
}

impl Stream for DiscoveryStream {
    type Item = Credentials;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.credentials.poll()
    }
}
<|endoftext|>


<commit_before>use super::{Function, Parameterised, VFunction, QFunction};
use geometry::Space;
use std::marker::PhantomData;

pub struct VFunctionGroup<S: Space, V: VFunction<S>>(Vec<V>, PhantomData<S>);

impl<S: Space, V: VFunction<S>> VFunctionGroup<S, V>
{
    pub fn new(functions: Vec<V>) -> Self {
        VFunctionGroup(functions, PhantomData)
    }
}

impl<S: Space, V: VFunction<S>> Function<S::Repr, Vec<f64>> for VFunctionGroup<S, V>
{
    fn evaluate(&self, state: &S::Repr) -> Vec<f64> {
        self.0.iter().map(|f| f.evaluate(state)).collect()
    }
}

impl<S: Space, V: VFunction<S>> Parameterised<S::Repr, Vec<f64>> for VFunctionGroup<S, V>
{
    fn update(&mut self, state: &S::Repr, mut errors: Vec<f64>) {
        let mut index = 0;

        for e in errors.drain(..) {
            self.0[index].update(state, e);
            index += 1;
        }
    }

    fn equivalent(&self, other: &Self) -> bool {
        self.0.iter().zip(other.0.iter())
            .any(|(ref f1, ref f2)| !f1.equivalent(f2))
    }
}

impl<S: Space, V: VFunction<S>> QFunction<S> for VFunctionGroup<S, V>
{
    fn evaluate_action(&self, state: &S::Repr, action: usize) -> f64 {
        self.0[action].evaluate(state)
    }

    fn update_action(&mut self, state: &S::Repr, action: usize, error: f64) {
        self.0[action].update(state, error);
    }
}

#[cfg(test)]
mod tests {
    use super::{VFunctionGroup, Function, Parameterised, VFunction, QFunction};
    use geometry::RegularSpace;
    use geometry::dimensions::Continuous;

    pub struct Mock(f64);

    impl Function<Vec<f64>, f64> for Mock {
        fn evaluate(&self, _: &Vec<f64>) -> f64 { self.0 }
    }

    impl Parameterised<Vec<f64>, f64> for Mock {
        fn update(&mut self, _: &Vec<f64>, e: f64) { self.0 = e; }
        fn equivalent(&self, other: &Self) -> bool { self.0 == other.0 }
    }

    impl VFunction<RegularSpace<Continuous>> for Mock {}

    #[test]
    fn test_function() {
        let fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        assert_eq!(fg.evaluate(&vec![]), vec![0.0, 1.0, 2.0]);
    }

    #[test]
    fn test_parameterised() {
        let mut fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        fg.update(&vec![], vec![4.0, -1.0, 8.0]);

        assert_eq!(fg.evaluate(&vec![]), vec![4.0, -1.0, 8.0]);
    }

    #[test]
    fn test_qfunction() {
        let fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        assert_eq!(fg.evaluate_action(&vec![], 0), 0.0);
        assert_eq!(fg.evaluate_action(&vec![], 1), 1.0);
        assert_eq!(fg.evaluate_action(&vec![], 2), 2.0);
    }

    #[test]
    fn test_boxed_vfuncs() {
        let mut fg = VFunctionGroup::new(vec![
            Box::new(Mock(0.0)),
            Box::new(Mock(1.0)),
            Box::new(Mock(2.0))
        ]);

        assert_eq!(fg.evaluate(&vec![]), vec![0.0, 1.0, 2.0]);

        assert_eq!(fg.evaluate_action(&vec![], 0), 0.0);
        assert_eq!(fg.evaluate_action(&vec![], 1), 1.0);
        assert_eq!(fg.evaluate_action(&vec![], 2), 2.0);

        fg.update(&vec![], vec![4.0, -1.0, 8.0]);

        assert_eq!(fg.evaluate(&vec![]), vec![4.0, -1.0, 8.0]);
    }
}
<commit_msg>Add test for equivalency for VFunctionGroup<commit_after>use super::{Function, Parameterised, VFunction, QFunction};
use geometry::Space;
use std::marker::PhantomData;

pub struct VFunctionGroup<S: Space, V: VFunction<S>>(Vec<V>, PhantomData<S>);

impl<S: Space, V: VFunction<S>> VFunctionGroup<S, V>
{
    pub fn new(functions: Vec<V>) -> Self {
        VFunctionGroup(functions, PhantomData)
    }
}

impl<S: Space, V: VFunction<S>> Function<S::Repr, Vec<f64>> for VFunctionGroup<S, V>
{
    fn evaluate(&self, state: &S::Repr) -> Vec<f64> {
        self.0.iter().map(|f| f.evaluate(state)).collect()
    }
}

impl<S: Space, V: VFunction<S>> Parameterised<S::Repr, Vec<f64>> for VFunctionGroup<S, V>
{
    fn update(&mut self, state: &S::Repr, mut errors: Vec<f64>) {
        let mut index = 0;

        for e in errors.drain(..) {
            self.0[index].update(state, e);
            index += 1;
        }
    }

    fn equivalent(&self, other: &Self) -> bool {
        !self.0.iter().zip(other.0.iter())
            .any(|(ref f1, ref f2)| !f1.equivalent(f2))
    }
}

impl<S: Space, V: VFunction<S>> QFunction<S> for VFunctionGroup<S, V>
{
    fn evaluate_action(&self, state: &S::Repr, action: usize) -> f64 {
        self.0[action].evaluate(state)
    }

    fn update_action(&mut self, state: &S::Repr, action: usize, error: f64) {
        self.0[action].update(state, error);
    }
}

#[cfg(test)]
mod tests {
    use super::{VFunctionGroup, Function, Parameterised, VFunction, QFunction};
    use geometry::RegularSpace;
    use geometry::dimensions::Continuous;

    pub struct Mock(f64);

    impl Function<Vec<f64>, f64> for Mock {
        fn evaluate(&self, _: &Vec<f64>) -> f64 { self.0 }
    }

    impl Parameterised<Vec<f64>, f64> for Mock {
        fn update(&mut self, _: &Vec<f64>, e: f64) { self.0 = e; }
        fn equivalent(&self, other: &Self) -> bool {
            (self.0 - other.0).abs() < 1e-5
        }
    }

    impl VFunction<RegularSpace<Continuous>> for Mock {}

    #[test]
    fn test_function() {
        let fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        assert_eq!(fg.evaluate(&vec![]), vec![0.0, 1.0, 2.0]);
    }

    #[test]
    fn test_parameterised() {
        let mut fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        fg.update(&vec![], vec![4.0, -1.0, 8.0]);

        assert_eq!(fg.evaluate(&vec![]), vec![4.0, -1.0, 8.0]);
    }

    #[test]
    fn test_equivalency() {
        let fg1 = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);
        let fg2 = VFunctionGroup::new(vec![Mock(2.0), Mock(1.0), Mock(2.0)]);
        let fg3 = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        assert!(fg1.equivalent(&fg1));
        assert!(fg1.equivalent(&fg3));
        assert!(fg2.equivalent(&fg2));
        assert!(fg3.equivalent(&fg1));
        assert!(fg3.equivalent(&fg3));

        assert!(!fg1.equivalent(&fg2));
        assert!(!fg2.equivalent(&fg1));
        assert!(!fg2.equivalent(&fg3));
        assert!(!fg3.equivalent(&fg2));
    }

    #[test]
    fn test_qfunction() {
        let fg = VFunctionGroup::new(vec![Mock(0.0), Mock(1.0), Mock(2.0)]);

        assert_eq!(fg.evaluate_action(&vec![], 0), 0.0);
        assert_eq!(fg.evaluate_action(&vec![], 1), 1.0);
        assert_eq!(fg.evaluate_action(&vec![], 2), 2.0);
    }

    #[test]
    fn test_boxed_vfuncs() {
        let mut fg = VFunctionGroup::new(vec![
            Box::new(Mock(0.0)),
            Box::new(Mock(1.0)),
            Box::new(Mock(2.0))
        ]);

        assert_eq!(fg.evaluate(&vec![]), vec![0.0, 1.0, 2.0]);

        assert_eq!(fg.evaluate_action(&vec![], 0), 0.0);
        assert_eq!(fg.evaluate_action(&vec![], 1), 1.0);
        assert_eq!(fg.evaluate_action(&vec![], 2), 2.0);

        fg.update(&vec![], vec![4.0, -1.0, 8.0]);

        assert_eq!(fg.evaluate(&vec![]), vec![4.0, -1.0, 8.0]);
    }
}
<|endoftext|>


<commit_before>use nickel::{Request, Response, MiddlewareResult};
use url::form_urlencoded;
use std::io::Read;
use nickel::status::StatusCode;
use hyper::header::{SetCookie, Location};
use cookie::Cookie as CookiePair;
use mustache::MapBuilder;
use common;

fn get_from_post(parse: &mut form_urlencoded::Parse, key: &str) -> Result<String, String> {
    match parse.find(|k| k.0 == key) {
        Some((_, value)) => Ok(value.into_owned()),
        None => Err("No URL in request".to_string()),
    }
}

pub fn add_favourite<'a, D>(request: &mut Request<D>,
                            mut response: Response<'a, D>)
                            -> MiddlewareResult<'a, D> {
    let ref mut original_req = request.origin;
    let mut existing = common::favourites(original_req);
    let mut buffer = String::new();
    original_req.read_to_string(&mut buffer).unwrap();
    let parse = &mut form_urlencoded::parse(buffer.as_bytes());
    let stopid = get_from_post(parse, "stopid");
    let pretty_name = get_from_post(parse, "prettyName");
    existing[stopid.unwrap()] = pretty_name.unwrap().into();
    response.set(SetCookie(vec![CookiePair::new(common::KEY.to_owned(), existing.dump())]));
    response.set(Location("/favourites".to_string()))
        .set(StatusCode::Found);
    return response.send("");
}

pub fn remove_favourite<'a, D>(request: &mut Request<D>,
                               mut response: Response<'a, D>)
                               -> MiddlewareResult<'a, D> {
    let ref mut original_req = request.origin;
    let mut existing = common::favourites(original_req);
    let mut buffer = String::new();
    original_req.read_to_string(&mut buffer).unwrap();
    let parse = &mut form_urlencoded::parse(buffer.as_bytes());
    let stopid = get_from_post(parse, "stopid");
    existing.remove(&stopid.unwrap());
    response.set(SetCookie(vec![CookiePair::new(common::KEY.to_owned(), existing.dump())]));
    response.set(Location("/favourites".to_string()))
        .set(StatusCode::Found);
    return response.send("");
}

pub fn list_favourites<'a, D>(request: &mut Request<D>,
                              response: Response<'a, D>)
                              -> MiddlewareResult<'a, D> {
    let favourites = common::favourites(&request.origin);
    let data = MapBuilder::new()
        .insert_vec("favourites", |vecbuilder| {
            let mut vecb = vecbuilder;
            for fav in favourites.entries() {
                vecb = vecb.push_map(|mapbuilder| {
                    mapbuilder.insert_str("key", fav.0)
                        .insert_str("value", fav.1)
                });
            }
            vecb
        })
        .build();
    common::render_to_response(response, "resources/templates/favourites.mustache", &data)
}
<commit_msg>Make 1-year cookies not session ones (fixes #10)<commit_after>use nickel::{Request, Response, MiddlewareResult};
use url::form_urlencoded;
use std::io::Read;
use nickel::status::StatusCode;
use hyper::header::{SetCookie, Location};
use cookie::Cookie as CookiePair;
use mustache::MapBuilder;
use common;
use time;
use std::collections::BTreeMap;
use json;

fn get_from_post(parse: &mut form_urlencoded::Parse, key: &str) -> Result<String, String> {
    match parse.find(|k| k.0 == key) {
        Some((_, value)) => Ok(value.into_owned()),
        None => Err("No URL in request".to_string()),
    }
}

fn set_cookie<'a, D>(mut response: &mut Response<'a, D>, existing: json::JsonValue) {
    response.set(SetCookie(vec![CookiePair {
                                    name: common::KEY.to_owned(),
                                    value: existing.dump(),
                                    path: None,
                                    expires: Some(time::now() + time::Duration::days(365)),
                                    max_age: None,
                                    secure: false,
                                    httponly: false,
                                    domain: None,
                                    custom: BTreeMap::new(),
                                }]));
    response.set(Location("/favourites".to_string()))
        .set(StatusCode::Found);
}

pub fn add_favourite<'a, D>(request: &mut Request<D>,
                            mut response: Response<'a, D>)
                            -> MiddlewareResult<'a, D> {
    let ref mut original_req = request.origin;
    let mut existing = common::favourites(original_req);
    let mut buffer = String::new();
    original_req.read_to_string(&mut buffer).unwrap();
    let parse = &mut form_urlencoded::parse(buffer.as_bytes());
    let stopid = get_from_post(parse, "stopid");
    let pretty_name = get_from_post(parse, "prettyName");
    existing[stopid.unwrap()] = pretty_name.unwrap().into();
    set_cookie(&mut response, existing);
    return response.send("");
}

pub fn remove_favourite<'a, D>(request: &mut Request<D>,
                               mut response: Response<'a, D>)
                               -> MiddlewareResult<'a, D> {
    let ref mut original_req = request.origin;
    let mut existing = common::favourites(original_req);
    let mut buffer = String::new();
    original_req.read_to_string(&mut buffer).unwrap();
    let parse = &mut form_urlencoded::parse(buffer.as_bytes());
    let stopid = get_from_post(parse, "stopid");
    existing.remove(&stopid.unwrap());
    set_cookie(&mut response, existing);
    return response.send("");
}

pub fn list_favourites<'a, D>(request: &mut Request<D>,
                              response: Response<'a, D>)
                              -> MiddlewareResult<'a, D> {
    let favourites = common::favourites(&request.origin);
    let data = MapBuilder::new()
        .insert_vec("favourites", |vecbuilder| {
            let mut vecb = vecbuilder;
            for fav in favourites.entries() {
                vecb = vecb.push_map(|mapbuilder| {
                    mapbuilder.insert_str("key", fav.0)
                        .insert_str("value", fav.1)
                });
            }
            vecb
        })
        .build();
    common::render_to_response(response, "resources/templates/favourites.mustache", &data)
}
<|endoftext|>


<commit_before>//! Stand-alone WebAssembly to Cretonne IL translator.
//!
//! This module defines the `FuncTranslator` type which can translate a single WebAssembly
//! function to Cretonne IL guided by a `FuncEnvironment` which provides information about the
//! WebAssembly module and the runtime environment.

use code_translator::translate_operator;
use cretonne::entity::EntityRef;
use cretonne::ir::{self, InstBuilder, Ebb};
use cretonne::result::{CtonResult, CtonError};
use cton_frontend::{ILBuilder, FunctionBuilder};
use environ::FuncEnvironment;
use state::TranslationState;
use translation_utils::Local;
use wasmparser::{self, BinaryReader};

/// Maximum number of local variables permitted in a function. The translation fails with a
/// `CtonError::ImplLimitExceeded` error if the limit is exceeded.
const MAX_LOCALS: usize = 50_000;

/// WebAssembly to Cretonne IL function translator.
///
/// A `FuncTranslator` is used to translate a binary WebAssembly function into Cretonne IL guided
/// by a `FuncEnvironment` object. A single translator instance can be reused to translate multiple
/// functions which will reduce heap allocation traffic.
pub struct FuncTranslator {
    il_builder: ILBuilder<Local>,
    state: TranslationState,
}

impl FuncTranslator {
    /// Create a new translator.
    pub fn new() -> Self {
        Self {
            il_builder: ILBuilder::new(),
            state: TranslationState::new(),
        }
    }

    /// Translate a binary WebAssembly function.
    ///
    /// The `code` slice contains the binary WebAssembly *function code* as it appears in the code
    /// section of a WebAssembly module, not including the initial size of the function code. The
    /// slice is expected to contain two parts:
    ///
    /// - The declaration of *locals*, and
    /// - The function *body* as an expression.
    ///
    /// See [the WebAssembly specification]
    /// (http://webassembly.github.io/spec/binary/modules.html#code-section).
    ///
    /// The Cretonne IR function `func` should be completely empty except for the `func.signature`
    /// and `func.name` fields. The signature may contain special-purpose arguments which are not
    /// regarded as WebAssembly local variables. Any signature arguments marked as
    /// `ArgumentPurpose::Normal` are made accessible as WebAssembly local variables.
    ///
    pub fn translate<FE: FuncEnvironment + ?Sized>(
        &mut self,
        code: &[u8],
        func: &mut ir::Function,
        environ: &mut FE,
    ) -> CtonResult {
        self.translate_from_reader(BinaryReader::new(code), func, environ)
    }

    /// Translate a binary WebAssembly function from a `BinaryReader`.
    pub fn translate_from_reader<FE: FuncEnvironment + ?Sized>(
        &mut self,
        mut reader: BinaryReader,
        func: &mut ir::Function,
        environ: &mut FE,
    ) -> CtonResult {
        dbg!(
            "translate({} bytes, {}{})",
            reader.bytes_remaining(),
            func.name,
            func.signature
        );
        assert_eq!(func.dfg.num_ebbs(), 0, "Function must be empty");
        assert_eq!(func.dfg.num_insts(), 0, "Function must be empty");

        // This clears the `ILBuilder`.
        let mut builder = FunctionBuilder::new(func, &mut self.il_builder);
        let entry_block = builder.create_ebb();
        builder.append_ebb_params_for_function_params(entry_block);
        builder.switch_to_block(entry_block); // This also creates values for the arguments.
        builder.seal_block(entry_block);
        // Make sure the entry block is inserted in the layout before we make any callbacks to
        // `environ`. The callback functions may need to insert things in the entry block.
        builder.ensure_inserted_ebb();

        let num_params = declare_wasm_parameters(&mut builder, entry_block);

        // Set up the translation state with a single pushed control block representing the whole
        // function and its return values.
        let exit_block = builder.create_ebb();
        builder.append_ebb_params_for_function_returns(exit_block);
        self.state.initialize(&builder.func.signature, exit_block);

        parse_local_decls(&mut reader, &mut builder, num_params)?;
        parse_function_body(reader, &mut builder, &mut self.state, environ)
    }
}

/// Declare local variables for the signature parameters that correspond to WebAssembly locals.
///
/// Return the number of local variables declared.
fn declare_wasm_parameters(builder: &mut FunctionBuilder<Local>, entry_block: Ebb) -> usize {
    let sig_len = builder.func.signature.params.len();
    let mut next_local = 0;
    for i in 0..sig_len {
        let param_type = builder.func.signature.params[i];
        // There may be additional special-purpose parameters following the normal WebAssembly
        // signature parameters. For example, a `vmctx` pointer.
        if param_type.purpose == ir::ArgumentPurpose::Normal {
            // This is a normal WebAssembly signature parameter, so create a local for it.
            let local = Local::new(next_local);
            builder.declare_var(local, param_type.value_type);
            next_local += 1;

            let param_value = builder.ebb_params(entry_block)[i];
            builder.def_var(local, param_value);
        }
    }

    next_local
}

/// Parse the local variable declarations that precede the function body.
///
/// Declare local variables, starting from `num_params`.
fn parse_local_decls(
    reader: &mut BinaryReader,
    builder: &mut FunctionBuilder<Local>,
    num_params: usize,
) -> CtonResult {
    let mut next_local = num_params;
    let local_count = reader.read_local_count().map_err(
        |_| CtonError::InvalidInput,
    )?;

    let mut locals_total = 0;
    for _ in 0..local_count {
        builder.set_srcloc(cur_srcloc(reader));
        let (count, ty) = reader.read_local_decl(&mut locals_total).map_err(|_| {
            CtonError::InvalidInput
        })?;
        declare_locals(builder, count, ty, &mut next_local)?;
    }

    Ok(())
}

/// Declare `count` local variables of the same type, starting from `next_local`.
///
/// Fail of too many locals are declared in the function, or if the type is not valid for a local.
fn declare_locals(
    builder: &mut FunctionBuilder<Local>,
    count: u32,
    wasm_type: wasmparser::Type,
    next_local: &mut usize,
) -> CtonResult {
    // All locals are initialized to 0.
    use wasmparser::Type::*;
    let zeroval = match wasm_type {
        I32 => builder.ins().iconst(ir::types::I32, 0),
        I64 => builder.ins().iconst(ir::types::I64, 0),
        F32 => builder.ins().f32const(ir::immediates::Ieee32::with_bits(0)),
        F64 => builder.ins().f64const(ir::immediates::Ieee64::with_bits(0)),
        _ => return Err(CtonError::InvalidInput),
    };

    let ty = builder.func.dfg.value_type(zeroval);
    for _ in 0..count {
        // This implementation limit is arbitrary, but it ensures that a small function can't blow
        // up the compiler by declaring millions of locals.
        if *next_local >= MAX_LOCALS {
            return Err(CtonError::ImplLimitExceeded);
        }

        let local = Local::new(*next_local);
        builder.declare_var(local, ty);
        builder.def_var(local, zeroval);
        *next_local += 1;
    }

    Ok(())
}

/// Parse the function body in `reader`.
///
/// This assumes that the local variable declarations have already been parsed and function
/// arguments and locals are declared in the builder.
fn parse_function_body<FE: FuncEnvironment + ?Sized>(
    mut reader: BinaryReader,
    builder: &mut FunctionBuilder<Local>,
    state: &mut TranslationState,
    environ: &mut FE,
) -> CtonResult {
    // The control stack is initialized with a single block representing the whole function.
    assert_eq!(state.control_stack.len(), 1, "State not initialized");

    // Keep going until the final `End` operator which pops the outermost block.
    while !state.control_stack.is_empty() {
        builder.set_srcloc(cur_srcloc(&reader));
        let op = reader.read_operator().map_err(|_| CtonError::InvalidInput)?;
        translate_operator(&op, builder, state, environ);
    }

    // The final `End` operator left us in the exit block where we need to manually add a return
    // instruction.
    //
    // If the exit block is unreachable, it may not have the correct arguments, so we would
    // generate a return instruction that doesn't match the signature.
    debug_assert!(builder.is_pristine());
    if !builder.is_unreachable() {
        builder.ins().return_(&state.stack);
    }

    // Discard any remaining values on the stack. Either we just returned them,
    // or the end of the function is unreachable.
    state.stack.clear();

    debug_assert!(reader.eof());

    Ok(())
}

/// Get the current source location from a reader.
fn cur_srcloc(reader: &BinaryReader) -> ir::SourceLoc {
    // We record source locations as byte code offsets relative to the beginning of the function.
    // This will wrap around of a single function's byte code is larger than 4 GB, but a) the
    // WebAssembly format doesn't allow for that, and b) that would hit other Cretonne
    // implementation limits anyway.
    ir::SourceLoc::new(reader.current_position() as u32)
}

#[cfg(test)]
mod tests {
    use cretonne::{ir, Context};
    use cretonne::ir::types::I32;
    use environ::{DummyEnvironment, FuncEnvironment};
    use super::FuncTranslator;

    #[test]
    fn small1() {
        // Implicit return.
        //
        // (func $small1 (param i32) (result i32)
        //     (i32.add (get_local 0) (i32.const 1))
        // )
        const BODY: [u8; 7] = [
            0x00,       // local decl count
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 1
            0x6a,       // i32.add
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("small1");
        ctx.func.signature.params.push(ir::AbiParam::new(I32));
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }

    #[test]
    fn small2() {
        // Same as above, but with an explicit return instruction.
        //
        // (func $small2 (param i32) (result i32)
        //     (return (i32.add (get_local 0) (i32.const 1)))
        // )
        const BODY: [u8; 8] = [
            0x00,       // local decl count
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 1
            0x6a,       // i32.add
            0x0f,       // return
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("small2");
        ctx.func.signature.params.push(ir::AbiParam::new(I32));
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }

    #[test]
    fn infloop() {
        // An infinite loop, no return instructions.
        //
        // (func $infloop (result i32)
        //     (local i32)
        //     (loop (result i32)
        //         (i32.add (get_local 0) (i32.const 1))
        //         (set_local 0)
        //         (br 0)
        //     )
        // )
        const BODY: [u8; 16] = [
            0x01,       // 1 local decl.
            0x01, 0x7f, // 1 i32 local.
            0x03, 0x7f, // loop i32
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 0
            0x6a,       // i32.add
            0x21, 0x00, // set_local 0
            0x0c, 0x00, // br 0
            0x0b,       // end
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("infloop");
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }
}
<commit_msg>Remove redundant validation for function locals.<commit_after>//! Stand-alone WebAssembly to Cretonne IL translator.
//!
//! This module defines the `FuncTranslator` type which can translate a single WebAssembly
//! function to Cretonne IL guided by a `FuncEnvironment` which provides information about the
//! WebAssembly module and the runtime environment.

use code_translator::translate_operator;
use cretonne::entity::EntityRef;
use cretonne::ir::{self, InstBuilder, Ebb};
use cretonne::result::{CtonResult, CtonError};
use cton_frontend::{ILBuilder, FunctionBuilder};
use environ::FuncEnvironment;
use state::TranslationState;
use translation_utils::Local;
use wasmparser::{self, BinaryReader};

/// WebAssembly to Cretonne IL function translator.
///
/// A `FuncTranslator` is used to translate a binary WebAssembly function into Cretonne IL guided
/// by a `FuncEnvironment` object. A single translator instance can be reused to translate multiple
/// functions which will reduce heap allocation traffic.
pub struct FuncTranslator {
    il_builder: ILBuilder<Local>,
    state: TranslationState,
}

impl FuncTranslator {
    /// Create a new translator.
    pub fn new() -> Self {
        Self {
            il_builder: ILBuilder::new(),
            state: TranslationState::new(),
        }
    }

    /// Translate a binary WebAssembly function.
    ///
    /// The `code` slice contains the binary WebAssembly *function code* as it appears in the code
    /// section of a WebAssembly module, not including the initial size of the function code. The
    /// slice is expected to contain two parts:
    ///
    /// - The declaration of *locals*, and
    /// - The function *body* as an expression.
    ///
    /// See [the WebAssembly specification]
    /// (http://webassembly.github.io/spec/binary/modules.html#code-section).
    ///
    /// The Cretonne IR function `func` should be completely empty except for the `func.signature`
    /// and `func.name` fields. The signature may contain special-purpose arguments which are not
    /// regarded as WebAssembly local variables. Any signature arguments marked as
    /// `ArgumentPurpose::Normal` are made accessible as WebAssembly local variables.
    ///
    pub fn translate<FE: FuncEnvironment + ?Sized>(
        &mut self,
        code: &[u8],
        func: &mut ir::Function,
        environ: &mut FE,
    ) -> CtonResult {
        self.translate_from_reader(BinaryReader::new(code), func, environ)
    }

    /// Translate a binary WebAssembly function from a `BinaryReader`.
    pub fn translate_from_reader<FE: FuncEnvironment + ?Sized>(
        &mut self,
        mut reader: BinaryReader,
        func: &mut ir::Function,
        environ: &mut FE,
    ) -> CtonResult {
        dbg!(
            "translate({} bytes, {}{})",
            reader.bytes_remaining(),
            func.name,
            func.signature
        );
        assert_eq!(func.dfg.num_ebbs(), 0, "Function must be empty");
        assert_eq!(func.dfg.num_insts(), 0, "Function must be empty");

        // This clears the `ILBuilder`.
        let mut builder = FunctionBuilder::new(func, &mut self.il_builder);
        let entry_block = builder.create_ebb();
        builder.append_ebb_params_for_function_params(entry_block);
        builder.switch_to_block(entry_block); // This also creates values for the arguments.
        builder.seal_block(entry_block);
        // Make sure the entry block is inserted in the layout before we make any callbacks to
        // `environ`. The callback functions may need to insert things in the entry block.
        builder.ensure_inserted_ebb();

        let num_params = declare_wasm_parameters(&mut builder, entry_block);

        // Set up the translation state with a single pushed control block representing the whole
        // function and its return values.
        let exit_block = builder.create_ebb();
        builder.append_ebb_params_for_function_returns(exit_block);
        self.state.initialize(&builder.func.signature, exit_block);

        parse_local_decls(&mut reader, &mut builder, num_params)?;
        parse_function_body(reader, &mut builder, &mut self.state, environ)
    }
}

/// Declare local variables for the signature parameters that correspond to WebAssembly locals.
///
/// Return the number of local variables declared.
fn declare_wasm_parameters(builder: &mut FunctionBuilder<Local>, entry_block: Ebb) -> usize {
    let sig_len = builder.func.signature.params.len();
    let mut next_local = 0;
    for i in 0..sig_len {
        let param_type = builder.func.signature.params[i];
        // There may be additional special-purpose parameters following the normal WebAssembly
        // signature parameters. For example, a `vmctx` pointer.
        if param_type.purpose == ir::ArgumentPurpose::Normal {
            // This is a normal WebAssembly signature parameter, so create a local for it.
            let local = Local::new(next_local);
            builder.declare_var(local, param_type.value_type);
            next_local += 1;

            let param_value = builder.ebb_params(entry_block)[i];
            builder.def_var(local, param_value);
        }
    }

    next_local
}

/// Parse the local variable declarations that precede the function body.
///
/// Declare local variables, starting from `num_params`.
fn parse_local_decls(
    reader: &mut BinaryReader,
    builder: &mut FunctionBuilder<Local>,
    num_params: usize,
) -> CtonResult {
    let mut next_local = num_params;
    let local_count = reader.read_local_count().map_err(
        |_| CtonError::InvalidInput,
    )?;

    let mut locals_total = 0;
    for _ in 0..local_count {
        builder.set_srcloc(cur_srcloc(reader));
        let (count, ty) = reader.read_local_decl(&mut locals_total).map_err(|_| {
            CtonError::InvalidInput
        })?;
        declare_locals(builder, count, ty, &mut next_local);
    }

    Ok(())
}

/// Declare `count` local variables of the same type, starting from `next_local`.
///
/// Fail of too many locals are declared in the function, or if the type is not valid for a local.
fn declare_locals(
    builder: &mut FunctionBuilder<Local>,
    count: u32,
    wasm_type: wasmparser::Type,
    next_local: &mut usize,
) {
    // All locals are initialized to 0.
    use wasmparser::Type::*;
    let zeroval = match wasm_type {
        I32 => builder.ins().iconst(ir::types::I32, 0),
        I64 => builder.ins().iconst(ir::types::I64, 0),
        F32 => builder.ins().f32const(ir::immediates::Ieee32::with_bits(0)),
        F64 => builder.ins().f64const(ir::immediates::Ieee64::with_bits(0)),
        _ => panic!("invalid local type"),
    };

    let ty = builder.func.dfg.value_type(zeroval);
    for _ in 0..count {
        let local = Local::new(*next_local);
        builder.declare_var(local, ty);
        builder.def_var(local, zeroval);
        *next_local += 1;
    }
}

/// Parse the function body in `reader`.
///
/// This assumes that the local variable declarations have already been parsed and function
/// arguments and locals are declared in the builder.
fn parse_function_body<FE: FuncEnvironment + ?Sized>(
    mut reader: BinaryReader,
    builder: &mut FunctionBuilder<Local>,
    state: &mut TranslationState,
    environ: &mut FE,
) -> CtonResult {
    // The control stack is initialized with a single block representing the whole function.
    assert_eq!(state.control_stack.len(), 1, "State not initialized");

    // Keep going until the final `End` operator which pops the outermost block.
    while !state.control_stack.is_empty() {
        builder.set_srcloc(cur_srcloc(&reader));
        let op = reader.read_operator().map_err(|_| CtonError::InvalidInput)?;
        translate_operator(&op, builder, state, environ);
    }

    // The final `End` operator left us in the exit block where we need to manually add a return
    // instruction.
    //
    // If the exit block is unreachable, it may not have the correct arguments, so we would
    // generate a return instruction that doesn't match the signature.
    debug_assert!(builder.is_pristine());
    if !builder.is_unreachable() {
        builder.ins().return_(&state.stack);
    }

    // Discard any remaining values on the stack. Either we just returned them,
    // or the end of the function is unreachable.
    state.stack.clear();

    debug_assert!(reader.eof());

    Ok(())
}

/// Get the current source location from a reader.
fn cur_srcloc(reader: &BinaryReader) -> ir::SourceLoc {
    // We record source locations as byte code offsets relative to the beginning of the function.
    // This will wrap around of a single function's byte code is larger than 4 GB, but a) the
    // WebAssembly format doesn't allow for that, and b) that would hit other Cretonne
    // implementation limits anyway.
    ir::SourceLoc::new(reader.current_position() as u32)
}

#[cfg(test)]
mod tests {
    use cretonne::{ir, Context};
    use cretonne::ir::types::I32;
    use environ::{DummyEnvironment, FuncEnvironment};
    use super::FuncTranslator;

    #[test]
    fn small1() {
        // Implicit return.
        //
        // (func $small1 (param i32) (result i32)
        //     (i32.add (get_local 0) (i32.const 1))
        // )
        const BODY: [u8; 7] = [
            0x00,       // local decl count
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 1
            0x6a,       // i32.add
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("small1");
        ctx.func.signature.params.push(ir::AbiParam::new(I32));
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }

    #[test]
    fn small2() {
        // Same as above, but with an explicit return instruction.
        //
        // (func $small2 (param i32) (result i32)
        //     (return (i32.add (get_local 0) (i32.const 1)))
        // )
        const BODY: [u8; 8] = [
            0x00,       // local decl count
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 1
            0x6a,       // i32.add
            0x0f,       // return
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("small2");
        ctx.func.signature.params.push(ir::AbiParam::new(I32));
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }

    #[test]
    fn infloop() {
        // An infinite loop, no return instructions.
        //
        // (func $infloop (result i32)
        //     (local i32)
        //     (loop (result i32)
        //         (i32.add (get_local 0) (i32.const 1))
        //         (set_local 0)
        //         (br 0)
        //     )
        // )
        const BODY: [u8; 16] = [
            0x01,       // 1 local decl.
            0x01, 0x7f, // 1 i32 local.
            0x03, 0x7f, // loop i32
            0x20, 0x00, // get_local 0
            0x41, 0x01, // i32.const 0
            0x6a,       // i32.add
            0x21, 0x00, // set_local 0
            0x0c, 0x00, // br 0
            0x0b,       // end
            0x0b,       // end
        ];

        let mut trans = FuncTranslator::new();
        let runtime = DummyEnvironment::default();
        let mut ctx = Context::new();

        ctx.func.name = ir::ExternalName::new("infloop");
        ctx.func.signature.returns.push(ir::AbiParam::new(I32));

        trans
            .translate(&BODY, &mut ctx.func, &mut runtime.func_env())
            .unwrap();
        dbg!("{}", ctx.func.display(None));
        ctx.verify(runtime.func_env().flags()).unwrap();
    }
}
<|endoftext|>