<commit_before>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};
use core::cell::Cell;

use stack;
use debug;
use arch::{self, StackPointer};

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:     State,
  stack:     Stack,
  stack_id:  debug::StackId,
  stack_ptr: arch::StackPointer,
  phantom:   PhantomData<(*const Input, *const Output)>
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  ///
  /// See also the [contract](../trait.GuardedStack.html) that needs to be fulfilled by `stack`.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  ///
  /// See also the [contract](../trait.Stack.html) that needs to be fulfilled by `stack`.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize, stack_ptr: StackPointer) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let f = ptr::read(env as *const F);
      let (data, stack_ptr) = arch::swap(0, stack_ptr, None);
      // See the second half of Yielder::suspend_bare.
      let input = ptr::read(data as *const Input);
      // Run the body of the generator.
      let mut yielder = Yielder::new(stack_ptr);
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let stack_id  = debug::StackId::register(&stack);
    let stack_ptr = arch::init(&stack, generator_wrapper::<Input, Output, Stack, F>);

    // Transfer environment to the callee.
    let stack_ptr = arch::swap(&f as *const F as usize, stack_ptr, Some(&stack)).1;
    mem::forget(f);

    Generator {
      state:     State::Runnable,
      stack:     stack,
      stack_id:  stack_id,
      stack_ptr: stack_ptr,
      phantom:   PhantomData
    }
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let (data_out, stack_ptr) = arch::swap(&input as *const Input as usize, self.stack_ptr, Some(&self.stack));
          self.stack_ptr = stack_ptr;
          mem::forget(input);
          ptr::read(data_out as *const Option<Output>)
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => self.stack
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send> {
  stack_ptr: Cell<StackPointer>,
  phantom: PhantomData<(*const Input, *const Output)>
}

impl<Input, Output> Yielder<Input, Output>
    where Input: Send, Output: Send {
  fn new(stack_ptr: StackPointer) -> Yielder<Input, Output> {
    Yielder {
      stack_ptr: Cell::new(stack_ptr),
      phantom: PhantomData
    }
  }

  #[inline(always)]
  fn suspend_bare(&self, val: Option<Output>) -> Input {
    unsafe {
      let (data, stack_ptr) = arch::swap(&val as *const Option<Output> as usize, self.stack_ptr.get(), None);
      self.stack_ptr.set(stack_ptr);
      mem::forget(val);
      ptr::read(data as *const Input)
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<commit_msg>Fix the variance of Input/Output<commit_after>// This file is part of libfringe, a low-level green threading library.
// Copyright (c) whitequark <whitequark@whitequark.org>
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Generators.
//!
//! Generators allow repeatedly suspending the execution of a function,
//! returning a value to the caller, and resuming the suspended function
//! afterwards.

use core::marker::PhantomData;
use core::{ptr, mem};
use core::cell::Cell;

use stack;
use debug;
use arch::{self, StackPointer};

#[derive(Debug, Clone, Copy)]
pub enum State {
  /// Generator can be resumed. This is the initial state.
  Runnable,
  /// Generator cannot be resumed. This is the state of the generator after
  /// the generator function has returned or panicked.
  Unavailable
}

/// Generator wraps a function and allows suspending its execution more than once, returning
/// a value each time.
///
/// The first time `resume(input0)` is called, the function is called as `f(yielder, input0)`.
/// It runs until it suspends its execution through `yielder.suspend(output0)`, after which
/// `resume(input0)` returns `output0`. The function can be resumed again using `resume(input1)`,
/// after which `yielder.suspend(output0)` returns `input1`, and so on. Once the function returns,
/// the `resume()` call will return `None`, and it will return `None` every time it is called
/// after that.
///
/// If the generator function panics, the panic is propagated through the `resume()` call as usual.
///
/// After the generator function returns or panics, it is safe to reclaim the generator stack
/// using `unwrap()`.
///
/// `state()` can be used to determine whether the generator function has returned;
/// the state is `State::Runnable` after creation and suspension, and `State::Unavailable`
/// once the generator function returns or panics.
///
/// When the input type is `()`, a generator implements the Iterator trait.
///
/// # Example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut add_one = Generator::new(stack, move |yielder, mut input| {
///   loop {
///     if input == 0 { break }
///     input = yielder.suspend(input + 1)
///   }
/// });
/// println!("{:?}", add_one.resume(2)); // prints Some(3)
/// println!("{:?}", add_one.resume(3)); // prints Some(4)
/// println!("{:?}", add_one.resume(0)); // prints None
/// ```
///
/// # Iterator example
///
/// ```
/// use fringe::{OsStack, Generator};
///
/// let stack = OsStack::new(0).unwrap();
/// let mut nat = Generator::new(stack, move |yielder, ()| {
///   for i in 1.. { yielder.suspend(i) }
/// });
/// println!("{:?}", nat.next()); // prints Some(0)
/// println!("{:?}", nat.next()); // prints Some(1)
/// println!("{:?}", nat.next()); // prints Some(2)
/// ```
#[derive(Debug)]
pub struct Generator<Input: Send, Output: Send, Stack: stack::Stack> {
  state:     State,
  stack:     Stack,
  stack_id:  debug::StackId,
  stack_ptr: arch::StackPointer,
  phantom:   PhantomData<(*mut Input, *const Output)>
}

impl<Input, Output, Stack> Generator<Input, Output, Stack>
    where Input: Send, Output: Send, Stack: stack::Stack {
  /// Creates a new generator.
  ///
  /// See also the [contract](../trait.GuardedStack.html) that needs to be fulfilled by `stack`.
  pub fn new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where Stack: stack::GuardedStack,
            F: FnOnce(&mut Yielder<Input, Output>, Input) + Send {
    unsafe { Generator::unsafe_new(stack, f) }
  }

  /// Same as `new`, but does not require `stack` to have a guard page.
  ///
  /// This function is unsafe because the generator function can easily violate
  /// memory safety by overflowing the stack. It is useful in environments where
  /// guarded stacks do not exist, e.g. in absence of an MMU.
  ///
  /// See also the [contract](../trait.Stack.html) that needs to be fulfilled by `stack`.
  pub unsafe fn unsafe_new<F>(stack: Stack, f: F) -> Generator<Input, Output, Stack>
      where F: FnOnce(&mut Yielder<Input, Output>, Input) + Send {
    unsafe extern "C" fn generator_wrapper<Input, Output, Stack, F>(env: usize, stack_ptr: StackPointer) -> !
        where Input: Send, Output: Send, Stack: stack::Stack,
              F: FnOnce(&mut Yielder<Input, Output>, Input) {
      // Retrieve our environment from the callee and return control to it.
      let f = ptr::read(env as *const F);
      let (data, stack_ptr) = arch::swap(0, stack_ptr, None);
      // See the second half of Yielder::suspend_bare.
      let input = ptr::read(data as *const Input);
      // Run the body of the generator.
      let mut yielder = Yielder::new(stack_ptr);
      f(&mut yielder, input);
      // Past this point, the generator has dropped everything it has held.
      loop { yielder.suspend_bare(None); }
    }

    let stack_id  = debug::StackId::register(&stack);
    let stack_ptr = arch::init(&stack, generator_wrapper::<Input, Output, Stack, F>);

    // Transfer environment to the callee.
    let stack_ptr = arch::swap(&f as *const F as usize, stack_ptr, Some(&stack)).1;
    mem::forget(f);

    Generator {
      state:     State::Runnable,
      stack:     stack,
      stack_id:  stack_id,
      stack_ptr: stack_ptr,
      phantom:   PhantomData
    }
  }

  /// Resumes the generator and return the next value it yields.
  /// If the generator function has returned, returns `None`.
  #[inline]
  pub fn resume(&mut self, input: Input) -> Option<Output> {
    match self.state {
      State::Runnable => {
        // Set the state to Unavailable. Since we have exclusive access to the generator,
        // the only case where this matters is the generator function panics, after which
        // it must not be invocable again.
        self.state = State::Unavailable;

        // Switch to the generator function, and retrieve the yielded value.
        let val = unsafe {
          let (data_out, stack_ptr) = arch::swap(&input as *const Input as usize, self.stack_ptr, Some(&self.stack));
          self.stack_ptr = stack_ptr;
          mem::forget(input);
          ptr::read(data_out as *const Option<Output>)
        };

        // Unless the generator function has returned, it can be switched to again, so
        // set the state to Runnable.
        if val.is_some() { self.state = State::Runnable }

        val
      }
      State::Unavailable => None
    }
  }

  /// Returns the state of the generator.
  #[inline]
  pub fn state(&self) -> State { self.state }

  /// Extracts the stack from a generator when the generator function has returned.
  /// If the generator function has not returned
  /// (i.e. `self.state() == State::Runnable`), panics.
  pub fn unwrap(self) -> Stack {
    match self.state {
      State::Runnable    => panic!("Argh! Bastard! Don't touch that!"),
      State::Unavailable => self.stack
    }
  }
}

/// Yielder is an interface provided to every generator through which it
/// returns a value.
#[derive(Debug)]
pub struct Yielder<Input: Send, Output: Send> {
  stack_ptr: Cell<StackPointer>,
  phantom: PhantomData<(*const Input, *mut Output)>
}

impl<Input, Output> Yielder<Input, Output>
    where Input: Send, Output: Send {
  fn new(stack_ptr: StackPointer) -> Yielder<Input, Output> {
    Yielder {
      stack_ptr: Cell::new(stack_ptr),
      phantom: PhantomData
    }
  }

  #[inline(always)]
  fn suspend_bare(&self, val: Option<Output>) -> Input {
    unsafe {
      let (data, stack_ptr) = arch::swap(&val as *const Option<Output> as usize, self.stack_ptr.get(), None);
      self.stack_ptr.set(stack_ptr);
      mem::forget(val);
      ptr::read(data as *const Input)
    }
  }

  /// Suspends the generator and returns `Some(item)` from the `resume()`
  /// invocation that resumed the generator.
  #[inline(always)]
  pub fn suspend(&self, item: Output) -> Input {
    self.suspend_bare(Some(item))
  }
}

impl<Output, Stack> Iterator for Generator<(), Output, Stack>
    where Output: Send, Stack: stack::Stack {
  type Item = Output;

  fn next(&mut self) -> Option<Self::Item> { self.resume(()) }
}
<|endoftext|>


<commit_before>/// Wrapper around the afl-fuzz binary

use libc;
use std::env;
use std::ffi::CString;
use std::path::PathBuf;
use std::ptr;

extern "C" {
    // `main` function with afl-fuzz.c
    fn afl_fuzz_main(argc: libc::c_int,
                     argv: *const *const libc::c_char) -> libc::c_int;
}

fn call_afl_fuzz_main(args: Vec<CString>) -> Result<(), libc::c_int> {
    assert!(!args.is_empty());
    // `main` functions in C expect an array of pointers to the arguments
    let mut args_ptrs = args.iter()
                        .map(|arg| arg.as_ptr())
                        .collect::<Vec<_>>();
    let argc = args_ptrs.len() as libc::c_int;
    // argv is a null-terminated array
    // http://stackoverflow.com/a/11020198
    args_ptrs.push(ptr::null());
    let ret = unsafe {
        afl_fuzz_main(argc, args_ptrs.as_ptr())
    };

    match ret {
        0 => Ok(()),
        n => Err(n),
    }
}

// TODO: use builder pattern
pub struct AflFuzzConfig {
    pub in_dir: PathBuf,
    pub out_dir: PathBuf,
    pub target: PathBuf,
}

impl AflFuzzConfig {
    fn into_args(self) -> Vec<CString> {
        let mut args = vec![];

        // -i
        args.push("-i");
        args.push(self.in_dir.to_str().unwrap());

        // -o
        args.push("-o");
        args.push(self.out_dir.to_str().unwrap());

        // target
        args.push(self.target.to_str().unwrap());

        // Convert args from `str`s to `CString`s
        args.into_iter()
            .map(|s| CString::new(s).unwrap())
            .collect()
    }
}

pub fn afl_fuzz_env() -> Result<(), libc::c_int> {
    let args = env::args();

    // don't include anything before 'afl-fuzz' or 'cargo-afl-fuzz'
    let mut args = args.skip_while(|a| !a.ends_with("afl-fuzz"))
                       .map(|arg| CString::new(arg).unwrap())
                       .collect::<Vec<_>>();
    assert!(!args.is_empty(), "Error generating afl-fuzz arguments");

    args[0] = CString::new("afl-fuzz").unwrap();

    call_afl_fuzz_main(args)
}

pub fn afl_fuzz(config: AflFuzzConfig) -> Result<(), libc::c_int> {
    let mut args = vec![];

    // Fake the first argument
    let arg0 = CString::new("afl-fuzz").unwrap();
    args.push(arg0);

    args.extend(config.into_args());

    call_afl_fuzz_main(args)
}

/// FIXME: figure out a way to make it not run indefinitely
#[test]
fn testsomething() {
    let config = AflFuzzConfig {
        in_dir: ".".into(),
        out_dir: "testout".into(),
        target: "../target/debug/examples/hello".into(),
    };
    afl_fuzz(config);
}
<commit_msg>Ignore indefinite test case.<commit_after>/// Wrapper around the afl-fuzz binary

use libc;
use std::env;
use std::ffi::CString;
use std::path::PathBuf;
use std::ptr;

extern "C" {
    // `main` function with afl-fuzz.c
    fn afl_fuzz_main(argc: libc::c_int,
                     argv: *const *const libc::c_char) -> libc::c_int;
}

fn call_afl_fuzz_main(args: Vec<CString>) -> Result<(), libc::c_int> {
    assert!(!args.is_empty());
    // `main` functions in C expect an array of pointers to the arguments
    let mut args_ptrs = args.iter()
                        .map(|arg| arg.as_ptr())
                        .collect::<Vec<_>>();
    let argc = args_ptrs.len() as libc::c_int;
    // argv is a null-terminated array
    // http://stackoverflow.com/a/11020198
    args_ptrs.push(ptr::null());
    let ret = unsafe {
        afl_fuzz_main(argc, args_ptrs.as_ptr())
    };

    match ret {
        0 => Ok(()),
        n => Err(n),
    }
}

// TODO: use builder pattern
pub struct AflFuzzConfig {
    pub in_dir: PathBuf,
    pub out_dir: PathBuf,
    pub target: PathBuf,
}

impl AflFuzzConfig {
    fn into_args(self) -> Vec<CString> {
        let mut args = vec![];

        // -i
        args.push("-i");
        args.push(self.in_dir.to_str().unwrap());

        // -o
        args.push("-o");
        args.push(self.out_dir.to_str().unwrap());

        // target
        args.push(self.target.to_str().unwrap());

        // Convert args from `str`s to `CString`s
        args.into_iter()
            .map(|s| CString::new(s).unwrap())
            .collect()
    }
}

pub fn afl_fuzz_env() -> Result<(), libc::c_int> {
    let args = env::args();

    // don't include anything before 'afl-fuzz' or 'cargo-afl-fuzz'
    let mut args = args.skip_while(|a| !a.ends_with("afl-fuzz"))
                       .map(|arg| CString::new(arg).unwrap())
                       .collect::<Vec<_>>();
    assert!(!args.is_empty(), "Error generating afl-fuzz arguments");

    args[0] = CString::new("afl-fuzz").unwrap();

    call_afl_fuzz_main(args)
}

pub fn afl_fuzz(config: AflFuzzConfig) -> Result<(), libc::c_int> {
    let mut args = vec![];

    // Fake the first argument
    let arg0 = CString::new("afl-fuzz").unwrap();
    args.push(arg0);

    args.extend(config.into_args());

    call_afl_fuzz_main(args)
}

/// FIXME: figure out a way to make it not run indefinitely
#[test]
#[ignore]
fn testsomething() {
    let config = AflFuzzConfig {
        in_dir: ".".into(),
        out_dir: "testout".into(),
        target: "../target/debug/examples/hello".into(),
    };
    afl_fuzz(config);
}
<|endoftext|>


<commit_before><commit_msg>Update test<commit_after><|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// 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.

//! Implementation of cross-process surfaces for Android. This uses EGL surface.

use platform::surface::NativeSurfaceMethods;
use texturegl::Texture;

use geom::size::Size2D;
use opengles::gl2::{egl_image_target_texture2d_oes, TEXTURE_2D, glTexImage2D, BGRA, UNSIGNED_BYTE};
use egl::egl::EGLDisplay;
use egl::eglext::{EGLImageKHR, DestroyImageKHR};
use libc::c_void;
use std::mem;
use std::ptr;
use std::vec::Vec;

/// FIXME(Aydin Kim) :Currently, native surface is consist of 2 types of hybrid image buffer. EGLImageKHR is used to GPU rendering and vector is used to CPU rendering. EGL extension seems not provide simple way to accessing its bitmap directly. In the future, we need to find out the way to integrate them.

pub struct NativeGraphicsMetadata {
    pub display: EGLDisplay,
}

pub struct NativePaintingGraphicsContext{
    display : EGLDisplay,
}

impl NativePaintingGraphicsContext {
    pub fn from_metadata(metadata: &NativeGraphicsMetadata) -> NativePaintingGraphicsContext {
        NativePaintingGraphicsContext {
            display : metadata.display,
        }
    }
}

impl Drop for NativePaintingGraphicsContext {
    fn drop(&mut self) {}
}

pub struct NativeCompositingGraphicsContext;

impl NativeCompositingGraphicsContext {
    pub fn new() -> NativeCompositingGraphicsContext {
        NativeCompositingGraphicsContext
    }
}

pub struct NativeSurface {
    image: Option<EGLImageKHR>, // For GPU rendering
    bitmap: Option<Vec<u8>>, // For CPU rendering
    will_leak: bool,
}

impl NativeSurface {
    pub fn from_image_khr(image_khr: EGLImageKHR) -> NativeSurface {
        let mut _image: Option<EGLImageKHR> = None;
        if image_khr != ptr::null() {
            _image = Some(image_khr);
        }
        NativeSurface {
            image : _image,
            bitmap: None,
            will_leak: true,
        }
    }
}

impl NativeSurfaceMethods for NativeSurface {
    /// This may only be called on the case of CPU rendering.
    fn new(_native_context: &NativePaintingGraphicsContext, size: Size2D<i32>, _stride: i32) -> NativeSurface {
        let len = size.width * size.height * 4;
        let bitmap: Vec<u8> = Vec::from_elem(len as uint, 0 as u8);

        NativeSurface {
            image: None,
            bitmap: Some(bitmap),
            will_leak : true,
        }
    }

    /// This may only be called on the compositor side.
    fn bind_to_texture(&self,
                       _native_context: &NativeCompositingGraphicsContext,
                       texture: &Texture,
                       _size: Size2D<int>) {
        let _bound = texture.bind();

        unsafe {
            match self.image {
                None => match self.bitmap {
                    Some(ref bitmap) => {
                        let data: *c_void = mem::transmute(bitmap.as_ptr());
                        glTexImage2D(TEXTURE_2D, 0, BGRA as i32, _size.width as i32, _size.height as i32, 0, BGRA as u32, UNSIGNED_BYTE, data);
                    }
                    None => {
                        debug!("Cannot bind the buffer(CPU rendering), there is no bitmap");
                    }
                },
                Some(image_khr) => {
                    egl_image_target_texture2d_oes(TEXTURE_2D, image_khr);
                }
            }
        }
    }

    /// This may only be called on the painting side.
    fn upload(&mut self, _graphics_context: &NativePaintingGraphicsContext, data: &[u8]) {
        match self.bitmap {
            Some(ref mut bitmap) => {
                unsafe {
                    let dest: &mut [u8] = mem::transmute((bitmap.as_ptr(), bitmap.len()));
                    dest.copy_memory(data);
                }
            }
            None => {
                debug!("Cannot upload the buffer(CPU rendering), there is no bitmap");
            }
        }
    }

    fn get_id(&self) -> int {
        match self.image {
            None => 0,
            Some(image_khr) => image_khr as int,
        }
    }

    fn destroy(&mut self, graphics_context: &NativePaintingGraphicsContext) {
        match self.image {
            None => {},
            Some(image_khr) => {
                DestroyImageKHR(graphics_context.display, image_khr);
                mem::replace(&mut self.image, None);
            }
        }
        self.mark_wont_leak()
    }

    fn mark_will_leak(&mut self) {
        self.will_leak = true
    }
    fn mark_wont_leak(&mut self) {
        self.will_leak = false
    }
}
<commit_msg>Replace transmute with as and as_mut_slice<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// 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.

//! Implementation of cross-process surfaces for Android. This uses EGL surface.

use platform::surface::NativeSurfaceMethods;
use texturegl::Texture;

use geom::size::Size2D;
use opengles::gl2::{egl_image_target_texture2d_oes, TEXTURE_2D, glTexImage2D, BGRA, UNSIGNED_BYTE};
use egl::egl::EGLDisplay;
use egl::eglext::{EGLImageKHR, DestroyImageKHR};
use libc::c_void;
use std::mem;
use std::ptr;
use std::vec::Vec;

/// FIXME(Aydin Kim) :Currently, native surface is consist of 2 types of hybrid image buffer. EGLImageKHR is used to GPU rendering and vector is used to CPU rendering. EGL extension seems not provide simple way to accessing its bitmap directly. In the future, we need to find out the way to integrate them.

pub struct NativeGraphicsMetadata {
    pub display: EGLDisplay,
}

pub struct NativePaintingGraphicsContext{
    display : EGLDisplay,
}

impl NativePaintingGraphicsContext {
    pub fn from_metadata(metadata: &NativeGraphicsMetadata) -> NativePaintingGraphicsContext {
        NativePaintingGraphicsContext {
            display : metadata.display,
        }
    }
}

impl Drop for NativePaintingGraphicsContext {
    fn drop(&mut self) {}
}

pub struct NativeCompositingGraphicsContext;

impl NativeCompositingGraphicsContext {
    pub fn new() -> NativeCompositingGraphicsContext {
        NativeCompositingGraphicsContext
    }
}

pub struct NativeSurface {
    image: Option<EGLImageKHR>, // For GPU rendering
    bitmap: Option<Vec<u8>>, // For CPU rendering
    will_leak: bool,
}

impl NativeSurface {
    pub fn from_image_khr(image_khr: EGLImageKHR) -> NativeSurface {
        let mut _image: Option<EGLImageKHR> = None;
        if image_khr != ptr::null() {
            _image = Some(image_khr);
        }
        NativeSurface {
            image : _image,
            bitmap: None,
            will_leak: true,
        }
    }
}

impl NativeSurfaceMethods for NativeSurface {
    /// This may only be called on the case of CPU rendering.
    fn new(_native_context: &NativePaintingGraphicsContext, size: Size2D<i32>, _stride: i32) -> NativeSurface {
        let len = size.width * size.height * 4;
        let bitmap: Vec<u8> = Vec::from_elem(len as uint, 0 as u8);

        NativeSurface {
            image: None,
            bitmap: Some(bitmap),
            will_leak : true,
        }
    }

    /// This may only be called on the compositor side.
    fn bind_to_texture(&self,
                       _native_context: &NativeCompositingGraphicsContext,
                       texture: &Texture,
                       _size: Size2D<int>) {
        let _bound = texture.bind();
        match self.image {
            None => match self.bitmap {
                Some(ref bitmap) => {
                    let data = bitmap.as_ptr() as *c_void;
                    unsafe {
                        glTexImage2D(TEXTURE_2D, 0, BGRA as i32, _size.width as i32, _size.height as i32,
                                     0, BGRA as u32, UNSIGNED_BYTE, data);
                    }
                }
                None => {
                    debug!("Cannot bind the buffer(CPU rendering), there is no bitmap");
                }
            },
            Some(image_khr) => {
                egl_image_target_texture2d_oes(TEXTURE_2D, image_khr);
            }
        }
    }

    /// This may only be called on the painting side.
    fn upload(&mut self, _graphics_context: &NativePaintingGraphicsContext, data: &[u8]) {
        match self.bitmap {
            Some(ref mut bitmap) => {
                unsafe {
                    bitmap.as_mut_slice().copy_memory(data);
                }
            }
            None => {
                debug!("Cannot upload the buffer(CPU rendering), there is no bitmap");
            }
        }
    }

    fn get_id(&self) -> int {
        match self.image {
            None => 0,
            Some(image_khr) => image_khr as int,
        }
    }

    fn destroy(&mut self, graphics_context: &NativePaintingGraphicsContext) {
        match self.image {
            None => {},
            Some(image_khr) => {
                DestroyImageKHR(graphics_context.display, image_khr);
                mem::replace(&mut self.image, None);
            }
        }
        self.mark_wont_leak()
    }

    fn mark_will_leak(&mut self) {
        self.will_leak = true
    }
    fn mark_wont_leak(&mut self) {
        self.will_leak = false
    }
}
<|endoftext|>


<commit_before>extern crate sodiumoxide;

use self::sodiumoxide::crypto::aead::chacha20poly1305_ietf;
use self::sodiumoxide::utils;

use std::cmp;
use std::io;
use std::io::{Read, Write};

use utils::crypto::pwhash_argon2i13;
use errors::common::CommonError;
use domain::wallet::KeyDerivationMethod;

pub const KEYBYTES: usize = chacha20poly1305_ietf::KEYBYTES;
pub const NONCEBYTES: usize = chacha20poly1305_ietf::NONCEBYTES;
pub const TAGBYTES: usize = chacha20poly1305_ietf::TAGBYTES;

sodium_type!(Key, chacha20poly1305_ietf::Key, KEYBYTES);
sodium_type!(Nonce, chacha20poly1305_ietf::Nonce, NONCEBYTES);

impl Nonce {
    pub fn increment(&mut self) {
        utils::increment_le(&mut (self.0).0);
    }
}

pub fn gen_key() -> Key {
    Key(chacha20poly1305_ietf::gen_key())
}

pub fn derive_key(passphrase: &str, salt: &pwhash_argon2i13::Salt, key_derivation_method: &KeyDerivationMethod) -> Result<Key, CommonError> {
    let mut key_bytes = [0u8; chacha20poly1305_ietf::KEYBYTES];
    pwhash_argon2i13::pwhash(&mut key_bytes, passphrase.as_bytes(), salt, key_derivation_method)
        .map_err(|err| CommonError::InvalidStructure(format!("Can't derive key: {}", err)))?;
    Ok(Key::new(key_bytes))
}

pub fn gen_nonce() -> Nonce {
    Nonce(chacha20poly1305_ietf::gen_nonce())
}

pub fn gen_nonce_and_encrypt(data: &[u8], key: &Key) -> (Vec<u8>, Nonce) {
    let nonce = gen_nonce();

    let encrypted_data = chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0
    );

    (encrypted_data, nonce)
}

pub fn encrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Vec<u8> {
    chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0,
    )
}

pub fn decrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Result<Vec<u8>, CommonError> {
    chacha20poly1305_ietf::open(
        &data,
        None,
        &nonce.0,
        &key.0,
    )
        .map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))
}

pub struct Writer<W: Write> {
    buffer: Vec<u8>,
    chunk_size: usize,
    key: Key,
    nonce: Nonce,
    inner: W,
}

impl<W: Write> Writer<W> {
    pub fn new(inner: W, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Writer {
            buffer: Vec::new(),
            chunk_size,
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for Writer<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buffer.write_all(buf)?; // TODO: Small optimizations are possible

        let mut chunk_start = 0;

        while self.buffer.len() >= chunk_start + self.chunk_size {
            let chunk = &self.buffer[chunk_start..chunk_start + self.chunk_size];
            self.inner.write_all(&encrypt(chunk, &self.key, &self.nonce))?;
            self.nonce.increment();
            chunk_start += self.chunk_size;
        }

        if chunk_start > 0 {
            self.buffer.drain(..chunk_start);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.buffer.len() > 0 {
            self.inner.write_all(&encrypt(&self.buffer, &self.key, &self.nonce))?;
            self.nonce.increment();
        }

        self.buffer.flush()
    }
}

pub struct Reader<R: Read> {
    rest_buffer: Vec<u8>,
    chunk_buffer: Vec<u8>,
    key: Key,
    nonce: Nonce,
    inner: R,
}

impl<R: Read> Reader<R> {
    pub fn new(inner: R, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Reader {
            rest_buffer: Vec::new(),
            chunk_buffer: vec![0; chunk_size + TAGBYTES],
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> R {
        self.inner
    }

    fn _read_chunk(&mut self) -> io::Result<usize> {
        let mut read = 0;

        while read < self.chunk_buffer.len() {
            match self.inner.read(&mut self.chunk_buffer[read..]) {
                Ok(0) => break,
                Ok(n) => read += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => Err(e)?
            }
        }

        if read == 0 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "No more crypto chucks to consume"))
        } else {
            Ok(read)
        }
    }
}

impl<R: Read> Read for Reader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut pos = 0;

        // Consume from rest buffer
        if self.rest_buffer.len() > 0 {
            let to_copy = cmp::min(self.rest_buffer.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&self.rest_buffer[..to_copy]);
            pos = pos + to_copy;
            self.rest_buffer.drain(..to_copy);
        }

        // Consume from chunks
        while pos < buf.len() {
            let chunk_size = self._read_chunk()?;

            let chunk = decrypt(&self.chunk_buffer[..chunk_size], &self.key, &self.nonce)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid data in crypto chunk"))?;

            self.nonce.increment();

            let to_copy = cmp::min(chunk.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&chunk[..to_copy]);
            pos = pos + to_copy;

            // Save rest in rest buffer
            if pos == buf.len() && to_copy < chunk.len() {
                self.rest_buffer.extend(&chunk[to_copy..]);
            }
        }

        Ok(buf.len())
    }
}

#[cfg(test)]
mod tests {
    extern crate rmp_serde;

    use super::*;
    use utils::crypto::randombytes::randombytes;

    #[test]
    fn derivation_argon2i_mod_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            148, 89, 76, 239, 127, 103, 13, 86, 84, 217, 216, 13, 223, 141, 225, 41,
            223, 126, 145, 138, 174, 31, 142, 199, 81, 12, 40, 201, 67, 8, 6, 251
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_MOD
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn derivation_argon2i_int_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            247, 55, 177, 252, 244, 130, 218, 129, 113, 206, 72, 44, 29, 68, 134, 215,
            249, 233, 131, 199, 38, 87, 69, 217, 156, 217, 10, 160, 30, 148, 80, 160
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_INT
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn encrypt_decrypt_works() {
        let data = randombytes(100);
        let key = gen_key();

        let (c, nonce) = gen_nonce_and_encrypt(&data, &key);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u);
    }

    #[test]
    fn encrypt_decrypt_works_for_nonce() {
        let data = randombytes(16);

        let key = gen_key();
        let nonce = gen_nonce();
        let c = encrypt(&data, &key, &nonce);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u)
    }

    #[test]
    fn nonce_serialize_deserialize_works() {
        let nonce = gen_nonce();
        let serialized = rmp_serde::to_vec(&nonce).unwrap();
        let deserialized: Nonce = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), NONCEBYTES + 2);
        assert_eq!(nonce, deserialized)
    }

    #[test]
    fn key_serialize_deserialize_works() {
        let key = gen_key();
        let serialized = rmp_serde::to_vec(&key).unwrap();
        let deserialized: Key = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), KEYBYTES + 2);
        assert_eq!(key, deserialized)
    }

    #[test]
    fn writer_reader_works_for_less_than_one_chunk() {
        let plain = randombytes(7);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 7 + TAGBYTES);

        let mut decrypted = vec![0u8; 7];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_one_chunk() {
        let plain = randombytes(10);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 10 + TAGBYTES);

        let mut decrypted = vec![0u8; 10];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_one_to_two_chunks() {
        let plain = randombytes(13);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 13 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 13];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_two_chunks() {
        let plain = randombytes(20);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 20 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 20];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }
}
<commit_msg>added detached with aad functons<commit_after>extern crate sodiumoxide;

use self::sodiumoxide::crypto::aead::chacha20poly1305_ietf;
use self::sodiumoxide::utils;

use std::cmp;
use std::io;
use std::io::{Read, Write};

use utils::crypto::pwhash_argon2i13;
use errors::common::CommonError;
use domain::wallet::KeyDerivationMethod;

pub const KEYBYTES: usize = chacha20poly1305_ietf::KEYBYTES;
pub const NONCEBYTES: usize = chacha20poly1305_ietf::NONCEBYTES;
pub const TAGBYTES: usize = chacha20poly1305_ietf::TAGBYTES;

sodium_type!(Key, chacha20poly1305_ietf::Key, KEYBYTES);
sodium_type!(Nonce, chacha20poly1305_ietf::Nonce, NONCEBYTES);
sodium_type!(Tag, chacha20poly1305_ietf::Tag, TAGBYTES);

impl Nonce {
    pub fn increment(&mut self) {
        utils::increment_le(&mut (self.0).0);
    }
}

pub fn gen_key() -> Key {
    Key(chacha20poly1305_ietf::gen_key())
}

pub fn derive_key(passphrase: &str, salt: &pwhash_argon2i13::Salt, key_derivation_method: &KeyDerivationMethod) -> Result<Key, CommonError> {
    let mut key_bytes = [0u8; chacha20poly1305_ietf::KEYBYTES];
    pwhash_argon2i13::pwhash(&mut key_bytes, passphrase.as_bytes(), salt, key_derivation_method)
        .map_err(|err| CommonError::InvalidStructure(format!("Can't derive key: {}", err)))?;
    Ok(Key::new(key_bytes))
}

pub fn gen_nonce() -> Nonce {
    Nonce(chacha20poly1305_ietf::gen_nonce())
}

pub fn gen_nonce_and_encrypt(data: &[u8], key: &Key) -> (Vec<u8>, Nonce) {
    let nonce = gen_nonce();

    let encrypted_data = chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0
    );

    (encrypted_data, nonce)
}

pub fn gen_nonce_and_encrypt_detached(data: &[u8], aad: &[u8], key: &Key) -> (Vec<u8>, Nonce, Tag) {
    let nonce = gen_nonce();

    let mut plain = data.to_vec();
    let tag = chacha20poly1305_ietf::seal_detached(
        plain.as_mut_slice(),
        Some(aad),
        &nonce.0,
        &key.0
    );

    (plain.to_vec(), nonce, Tag(tag))
}

pub fn decrypt_detached(data: &[u8], key: &Key, nonce: &Nonce, tag: &Tag, ad: Option<&[u8]>) -> Result<Vec<u8>, CommonError> {
    let mut plain = data.to_vec();
    chacha20poly1305_ietf::open_detached(plain.as_mut_slice(),
        ad,
        &tag.0,
        &nonce.0,
        &key.0,
    ).map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))?;

    Ok(plain.to_vec())
}

pub fn encrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Vec<u8> {
    chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0,
    )
}

pub fn decrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Result<Vec<u8>, CommonError> {
    chacha20poly1305_ietf::open(
        &data,
        None,
        &nonce.0,
        &key.0,
    )
        .map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))
}

pub struct Writer<W: Write> {
    buffer: Vec<u8>,
    chunk_size: usize,
    key: Key,
    nonce: Nonce,
    inner: W,
}

impl<W: Write> Writer<W> {
    pub fn new(inner: W, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Writer {
            buffer: Vec::new(),
            chunk_size,
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for Writer<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buffer.write_all(buf)?; // TODO: Small optimizations are possible

        let mut chunk_start = 0;

        while self.buffer.len() >= chunk_start + self.chunk_size {
            let chunk = &self.buffer[chunk_start..chunk_start + self.chunk_size];
            self.inner.write_all(&encrypt(chunk, &self.key, &self.nonce))?;
            self.nonce.increment();
            chunk_start += self.chunk_size;
        }

        if chunk_start > 0 {
            self.buffer.drain(..chunk_start);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.buffer.len() > 0 {
            self.inner.write_all(&encrypt(&self.buffer, &self.key, &self.nonce))?;
            self.nonce.increment();
        }

        self.buffer.flush()
    }
}

pub struct Reader<R: Read> {
    rest_buffer: Vec<u8>,
    chunk_buffer: Vec<u8>,
    key: Key,
    nonce: Nonce,
    inner: R,
}

impl<R: Read> Reader<R> {
    pub fn new(inner: R, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Reader {
            rest_buffer: Vec::new(),
            chunk_buffer: vec![0; chunk_size + TAGBYTES],
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> R {
        self.inner
    }

    fn _read_chunk(&mut self) -> io::Result<usize> {
        let mut read = 0;

        while read < self.chunk_buffer.len() {
            match self.inner.read(&mut self.chunk_buffer[read..]) {
                Ok(0) => break,
                Ok(n) => read += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => Err(e)?
            }
        }

        if read == 0 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "No more crypto chucks to consume"))
        } else {
            Ok(read)
        }
    }
}

impl<R: Read> Read for Reader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut pos = 0;

        // Consume from rest buffer
        if self.rest_buffer.len() > 0 {
            let to_copy = cmp::min(self.rest_buffer.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&self.rest_buffer[..to_copy]);
            pos = pos + to_copy;
            self.rest_buffer.drain(..to_copy);
        }

        // Consume from chunks
        while pos < buf.len() {
            let chunk_size = self._read_chunk()?;

            let chunk = decrypt(&self.chunk_buffer[..chunk_size], &self.key, &self.nonce)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid data in crypto chunk"))?;

            self.nonce.increment();

            let to_copy = cmp::min(chunk.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&chunk[..to_copy]);
            pos = pos + to_copy;

            // Save rest in rest buffer
            if pos == buf.len() && to_copy < chunk.len() {
                self.rest_buffer.extend(&chunk[to_copy..]);
            }
        }

        Ok(buf.len())
    }
}

#[cfg(test)]
mod tests {
    extern crate rmp_serde;

    use super::*;
    use utils::crypto::randombytes::randombytes;

    #[test]
    fn derivation_argon2i_mod_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            148, 89, 76, 239, 127, 103, 13, 86, 84, 217, 216, 13, 223, 141, 225, 41,
            223, 126, 145, 138, 174, 31, 142, 199, 81, 12, 40, 201, 67, 8, 6, 251
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_MOD
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn derivation_argon2i_int_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            247, 55, 177, 252, 244, 130, 218, 129, 113, 206, 72, 44, 29, 68, 134, 215,
            249, 233, 131, 199, 38, 87, 69, 217, 156, 217, 10, 160, 30, 148, 80, 160
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_INT
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn gen_nonce_and_encrypt_decrypt_works() {
        let data = randombytes(100);
        let key = gen_key();

        let (c, nonce) = gen_nonce_and_encrypt(&data, &key);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u);
    }

    #[test]
    pub fn gen_nonce_and_encrypt_detached_decrypt_detached_works() {
        let data = randombytes(100);
        let key = gen_key();
        let aad= randombytes(100);

        let (c, nonce, tag) = gen_nonce_and_encrypt_detached(&data, aad.as_slice(), &key);
        let u = decrypt_detached(&c, &key, &nonce, &tag, Some(aad.as_slice())).unwrap();
        assert_eq!(data, u);
}

    #[test]
    fn encrypt_decrypt_works_for_nonce() {
        let data = randombytes(16);

        let key = gen_key();
        let nonce = gen_nonce();
        let c = encrypt(&data, &key, &nonce);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u)
    }

    #[test]
    fn nonce_serialize_deserialize_works() {
        let nonce = gen_nonce();
        let serialized = rmp_serde::to_vec(&nonce).unwrap();
        let deserialized: Nonce = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), NONCEBYTES + 2);
        assert_eq!(nonce, deserialized)
    }

    #[test]
    fn key_serialize_deserialize_works() {
        let key = gen_key();
        let serialized = rmp_serde::to_vec(&key).unwrap();
        let deserialized: Key = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), KEYBYTES + 2);
        assert_eq!(key, deserialized)
    }

    #[test]
    fn writer_reader_works_for_less_than_one_chunk() {
        let plain = randombytes(7);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 7 + TAGBYTES);

        let mut decrypted = vec![0u8; 7];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_one_chunk() {
        let plain = randombytes(10);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 10 + TAGBYTES);

        let mut decrypted = vec![0u8; 10];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_one_to_two_chunks() {
        let plain = randombytes(13);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 13 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 13];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_two_chunks() {
        let plain = randombytes(20);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 20 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 20];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }
}
<|endoftext|>


<commit_before>use super::*;

use crate::ok_or_shutdown;

impl TaskHandler {
    /// See if we can start a new queued task.
    pub fn spawn_new(&mut self) {
        // Get the next task id that can be started
        if let Some(id) = self.get_next_task_id() {
            let cloned_state_mutex = self.state.clone();
            let mut state = cloned_state_mutex.lock().unwrap();
            self.start_process(id, &mut state);
        }
    }

    /// Search and return the next task that can be started.
    /// Precondition for a task to be started:
    /// - is in Queued state
    /// - There are free slots in the task's group
    /// - The group is running
    /// - has all its dependencies in `Done` state
    pub fn get_next_task_id(&mut self) -> Option<usize> {
        let state = self.state.lock().unwrap();
        // Check how many tasks are running in each group
        let mut running_tasks_per_group: HashMap<String, usize> = HashMap::new();

        // Create a default group for tasks without an explicit group
        running_tasks_per_group.insert("default".into(), 0);

        // Walk through all tasks and save the number of running tasks by group
        for (_, task) in state.tasks.iter() {
            // We are only interested in currently running tasks.
            if ![TaskStatus::Running, TaskStatus::Paused].contains(&task.status) {
                continue;
            }

            match running_tasks_per_group.get_mut(&task.group) {
                Some(count) => {
                    *count += 1;
                }
                None => {
                    running_tasks_per_group.insert(task.group.clone(), 1);
                }
            }
        }

        state
            .tasks
            .iter()
            .filter(|(_, task)| task.status == TaskStatus::Queued)
            .filter(|(_, task)| {
                // The task is assigned to a group.
                // First let's check if the group is running. If it isn't, simply return false.
                if !matches!(state.groups.get(&task.group), Some(GroupStatus::Running)) {
                    return false;
                }

                // If there's no running task for the group yet, we can safely return true
                //
                // If there are running tasks for this group, we have to ensure that there are
                // fewer running tasks than allowed for this group.
                match running_tasks_per_group.get(&task.group) {
                    None => true,
                    Some(count) => match state.settings.daemon.groups.get(&task.group) {
                        Some(allowed) => count < allowed,
                        None => {
                            error!(
                                "Got task with unknown group {}. Please report this!",
                                &task.group
                            );
                            false
                        }
                    },
                }
            })
            .find(|(_, task)| {
                // Check whether all dependencies for this task are fulfilled.
                task.dependencies
                    .iter()
                    .flat_map(|id| state.tasks.get(id))
                    .all(|task| task.status == TaskStatus::Done && !task.failed())
            })
            .map(|(id, _)| *id)
    }

    /// Actually spawn a new sub process
    /// The output of subprocesses is piped into a seperate file for easier access
    pub fn start_process(&mut self, task_id: usize, state: &mut LockedState) {
        // Check if the task exists and can actually be spawned. Otherwise do an early return.
        match state.tasks.get(&task_id) {
            Some(task) => {
                if !vec![TaskStatus::Stashed, TaskStatus::Queued, TaskStatus::Paused]
                    .contains(&task.status)
                {
                    info!("Tried to start task with status: {}", task.status);
                    return;
                }
            }
            None => {
                info!("Tried to start non-existing task: {}", task_id);
                return;
            }
        };

        // Try to get the log files to which the output of the process
        // will be written to. Error if this doesn't work!
        let (stdout_log, stderr_log) = match create_log_file_handles(task_id, &self.pueue_directory)
        {
            Ok((out, err)) => (out, err),
            Err(err) => {
                error!("Failed to create child log files: {:?}", err);
                return;
            }
        };

        // Get all necessary info for starting the task
        let (command, path, envs) = {
            let task = state.tasks.get(&task_id).unwrap();
            (task.command.clone(), task.path.clone(), task.envs.clone())
        };

        // Spawn the actual subprocess
        let mut command = compile_shell_command(&command);

        let spawned_command = command
            .current_dir(path)
            .stdin(Stdio::piped())
            .envs(envs)
            .stdout(Stdio::from(stdout_log))
            .stderr(Stdio::from(stderr_log))
            .spawn();

        // Check if the task managed to spawn
        let child = match spawned_command {
            Ok(child) => child,
            Err(err) => {
                let error = format!("Failed to spawn child {} with err: {:?}", task_id, err);
                error!("{}", error);
                clean_log_handles(task_id, &self.pueue_directory);

                // Update all necessary fields on the task.
                let group = {
                    let task = state.tasks.get_mut(&task_id).unwrap();
                    task.status = TaskStatus::Done;
                    task.result = Some(TaskResult::FailedToSpawn(error));
                    task.start = Some(Local::now());
                    task.end = Some(Local::now());
                    task.enqueue_at = None;
                    self.spawn_callback(&task);

                    task.group.clone()
                };

                state.handle_task_failure(group);
                ok_or_shutdown!(self, state.save());
                return;
            }
        };
        self.children.insert(task_id, child);

        let task = state.tasks.get_mut(&task_id).unwrap();

        task.start = Some(Local::now());
        task.status = TaskStatus::Running;
        task.enqueue_at = None;

        info!("Started task: {}", task.command);
        ok_or_shutdown!(self, state.save());
    }
}
<commit_msg>Early mutex lock in spawn_new<commit_after>use super::*;

use crate::ok_or_shutdown;

impl TaskHandler {
    /// See if we can start a new queued task.
    pub fn spawn_new(&mut self) {
        let cloned_state_mutex = self.state.clone();
        let mut state = cloned_state_mutex.lock().unwrap();
        // Get the next task id that can be started
        if let Some(id) = self.get_next_task_id(&state) {
            self.start_process(id, &mut state);
        }
    }

    /// Search and return the next task that can be started.
    /// Precondition for a task to be started:
    /// - is in Queued state
    /// - There are free slots in the task's group
    /// - The group is running
    /// - has all its dependencies in `Done` state
    pub fn get_next_task_id(&mut self, state: &LockedState) -> Option<usize> {
        // Check how many tasks are running in each group
        let mut running_tasks_per_group: HashMap<String, usize> = HashMap::new();

        // Create a default group for tasks without an explicit group
        running_tasks_per_group.insert("default".into(), 0);

        // Walk through all tasks and save the number of running tasks by group
        for (_, task) in state.tasks.iter() {
            // We are only interested in currently running tasks.
            if ![TaskStatus::Running, TaskStatus::Paused].contains(&task.status) {
                continue;
            }

            match running_tasks_per_group.get_mut(&task.group) {
                Some(count) => {
                    *count += 1;
                }
                None => {
                    running_tasks_per_group.insert(task.group.clone(), 1);
                }
            }
        }

        state
            .tasks
            .iter()
            .filter(|(_, task)| task.status == TaskStatus::Queued)
            .filter(|(_, task)| {
                // The task is assigned to a group.
                // First let's check if the group is running. If it isn't, simply return false.
                if !matches!(state.groups.get(&task.group), Some(GroupStatus::Running)) {
                    return false;
                }

                // If there's no running task for the group yet, we can safely return true
                //
                // If there are running tasks for this group, we have to ensure that there are
                // fewer running tasks than allowed for this group.
                match running_tasks_per_group.get(&task.group) {
                    None => true,
                    Some(count) => match state.settings.daemon.groups.get(&task.group) {
                        Some(allowed) => count < allowed,
                        None => {
                            error!(
                                "Got task with unknown group {}. Please report this!",
                                &task.group
                            );
                            false
                        }
                    },
                }
            })
            .find(|(_, task)| {
                // Check whether all dependencies for this task are fulfilled.
                task.dependencies
                    .iter()
                    .flat_map(|id| state.tasks.get(id))
                    .all(|task| task.status == TaskStatus::Done && !task.failed())
            })
            .map(|(id, _)| *id)
    }

    /// Actually spawn a new sub process
    /// The output of subprocesses is piped into a seperate file for easier access
    pub fn start_process(&mut self, task_id: usize, state: &mut LockedState) {
        // Check if the task exists and can actually be spawned. Otherwise do an early return.
        match state.tasks.get(&task_id) {
            Some(task) => {
                if !vec![TaskStatus::Stashed, TaskStatus::Queued, TaskStatus::Paused]
                    .contains(&task.status)
                {
                    info!("Tried to start task with status: {}", task.status);
                    return;
                }
            }
            None => {
                info!("Tried to start non-existing task: {}", task_id);
                return;
            }
        };

        // Try to get the log files to which the output of the process
        // will be written to. Error if this doesn't work!
        let (stdout_log, stderr_log) = match create_log_file_handles(task_id, &self.pueue_directory)
        {
            Ok((out, err)) => (out, err),
            Err(err) => {
                error!("Failed to create child log files: {:?}", err);
                return;
            }
        };

        // Get all necessary info for starting the task
        let (command, path, envs) = {
            let task = state.tasks.get(&task_id).unwrap();
            (task.command.clone(), task.path.clone(), task.envs.clone())
        };

        // Spawn the actual subprocess
        let mut command = compile_shell_command(&command);

        let spawned_command = command
            .current_dir(path)
            .stdin(Stdio::piped())
            .envs(envs)
            .stdout(Stdio::from(stdout_log))
            .stderr(Stdio::from(stderr_log))
            .spawn();

        // Check if the task managed to spawn
        let child = match spawned_command {
            Ok(child) => child,
            Err(err) => {
                let error = format!("Failed to spawn child {} with err: {:?}", task_id, err);
                error!("{}", error);
                clean_log_handles(task_id, &self.pueue_directory);

                // Update all necessary fields on the task.
                let group = {
                    let task = state.tasks.get_mut(&task_id).unwrap();
                    task.status = TaskStatus::Done;
                    task.result = Some(TaskResult::FailedToSpawn(error));
                    task.start = Some(Local::now());
                    task.end = Some(Local::now());
                    task.enqueue_at = None;
                    self.spawn_callback(&task);

                    task.group.clone()
                };

                state.handle_task_failure(group);
                ok_or_shutdown!(self, state.save());
                return;
            }
        };
        self.children.insert(task_id, child);

        let task = state.tasks.get_mut(&task_id).unwrap();

        task.start = Some(Local::now());
        task.status = TaskStatus::Running;
        task.enqueue_at = None;

        info!("Started task: {}", task.command);
        ok_or_shutdown!(self, state.save());
    }
}
<|endoftext|>


<commit_before>use crate::syntax::{NamedType, Type};
use proc_macro2::Ident;
use std::hash::{Hash, Hasher};
use syn::Token;

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub enum ImplKey<'a> {
    RustBox(NamedImplKey<'a>),
    RustVec(NamedImplKey<'a>),
    UniquePtr(NamedImplKey<'a>),
    SharedPtr(NamedImplKey<'a>),
    WeakPtr(NamedImplKey<'a>),
    CxxVector(NamedImplKey<'a>),
}

#[derive(Copy, Clone)]
pub struct NamedImplKey<'a> {
    pub rust: &'a Ident,
    pub lt_token: Option<Token![<]>,
    pub gt_token: Option<Token![>]>,
}

impl Type {
    pub(crate) fn impl_key(&self) -> Option<ImplKey> {
        if let Type::RustBox(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::RustBox(NamedImplKey::from(ident)));
            }
        } else if let Type::RustVec(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::RustVec(NamedImplKey::from(ident)));
            }
        } else if let Type::UniquePtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::UniquePtr(NamedImplKey::from(ident)));
            }
        } else if let Type::SharedPtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::SharedPtr(NamedImplKey::from(ident)));
            }
        } else if let Type::WeakPtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::WeakPtr(NamedImplKey::from(ident)));
            }
        } else if let Type::CxxVector(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::CxxVector(NamedImplKey::from(ident)));
            }
        }
        None
    }
}

impl<'a> PartialEq for NamedImplKey<'a> {
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(self.rust, other.rust)
    }
}

impl<'a> Eq for NamedImplKey<'a> {}

impl<'a> Hash for NamedImplKey<'a> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.rust.hash(hasher);
    }
}

impl<'a> From<&'a NamedType> for NamedImplKey<'a> {
    fn from(ty: &'a NamedType) -> Self {
        NamedImplKey {
            rust: &ty.rust,
            lt_token: ty.generics.lt_token,
            gt_token: ty.generics.gt_token,
        }
    }
}
<commit_msg>Track the outer type's spans too for implicit impls<commit_after>use crate::syntax::{NamedType, Ty1, Type};
use proc_macro2::{Ident, Span};
use std::hash::{Hash, Hasher};
use syn::Token;

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub enum ImplKey<'a> {
    RustBox(NamedImplKey<'a>),
    RustVec(NamedImplKey<'a>),
    UniquePtr(NamedImplKey<'a>),
    SharedPtr(NamedImplKey<'a>),
    WeakPtr(NamedImplKey<'a>),
    CxxVector(NamedImplKey<'a>),
}

#[derive(Copy, Clone)]
pub struct NamedImplKey<'a> {
    pub begin_span: Span,
    pub rust: &'a Ident,
    pub lt_token: Option<Token![<]>,
    pub gt_token: Option<Token![>]>,
    pub end_span: Span,
}

impl Type {
    pub(crate) fn impl_key(&self) -> Option<ImplKey> {
        if let Type::RustBox(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::RustBox(NamedImplKey::new(ty, ident)));
            }
        } else if let Type::RustVec(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::RustVec(NamedImplKey::new(ty, ident)));
            }
        } else if let Type::UniquePtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::UniquePtr(NamedImplKey::new(ty, ident)));
            }
        } else if let Type::SharedPtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::SharedPtr(NamedImplKey::new(ty, ident)));
            }
        } else if let Type::WeakPtr(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::WeakPtr(NamedImplKey::new(ty, ident)));
            }
        } else if let Type::CxxVector(ty) = self {
            if let Type::Ident(ident) = &ty.inner {
                return Some(ImplKey::CxxVector(NamedImplKey::new(ty, ident)));
            }
        }
        None
    }
}

impl<'a> PartialEq for NamedImplKey<'a> {
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(self.rust, other.rust)
    }
}

impl<'a> Eq for NamedImplKey<'a> {}

impl<'a> Hash for NamedImplKey<'a> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.rust.hash(hasher);
    }
}

impl<'a> NamedImplKey<'a> {
    fn new(outer: &Ty1, inner: &'a NamedType) -> Self {
        NamedImplKey {
            begin_span: outer.name.span(),
            rust: &inner.rust,
            lt_token: inner.generics.lt_token,
            gt_token: inner.generics.gt_token,
            end_span: outer.rangle.span,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

#[inline]
pub fn interpolate(a: f64, b: f64, fraction: f64) -> f64 {
    a * (1.0 - fraction) + b * fraction
}

/// Fast lookup table for arbitrary floating point functions.
/// This can be used to e.g., provide fast lookups of distribution values.
/// Input values are sampled with a given precision and results are stored in a vector.
/// During lookup, infimum and supremum of a given value are calculated and the result is
/// interpolated.
pub struct InterpolationTable<F: Fn(f64) -> f64> {
    inner: Vec<f64>,
    func: F,
    offset: usize,
    min_x: f64,
    max_x: f64,
    shift: f64,
}

impl<F: Fn(f64) -> f64> InterpolationTable<F> {
    /// Create a new `InterpolationTable`.
    ///
    /// # Arguments
    ///
    /// * `min_x` - minimum sample value
    /// * `max_x` - maximum sample value
    /// * `frac_digits` - number of fraction digits to store in sample
    /// * `func` - Function to emulate.
    ///
    /// If given value is outside of min_x and max_x, the lookup falls back to applying the
    /// function itself.
    /// The table size grows with the number of fraction digits.
    pub fn new(min_x: f64, max_x: f64, frac_digits: i32, func: F) -> Self {
        let shift = 10.0_f64.powi(frac_digits);
        let offset = (min_x * shift) as usize;
        let mut table = InterpolationTable {
            inner: Vec::new(),
            func,
            min_x,
            max_x,
            shift,
            offset,
        };

        let mut i = table.index(min_x);
        while i < table.index(max_x) {
            let x = i as f64 / shift;
            table.inner.push((table.func)(x));
            i += 1;
        }

        table
    }

    /// Return lower bound index for given f64.
    #[inline]
    fn index(&self, x: f64) -> usize {
        (x * self.shift) as usize - self.offset
    }

    /// Lookup given value in table, and interpolate the result between the sampled values if
    /// necessary. This provides an approximation that is better the more fraction digits are
    /// used to generate this table.
    pub fn get(&self, x: f64) -> f64 {
        if x < self.min_x || x >= self.max_x {
            (self.func)(x)
        } else {
            let i = self.index(x);
            // interpolate
            let fraction = (x * self.shift - i as f64) / self.shift;
            interpolate(self.inner[i], self.inner[i + 1], fraction)
        }
    }
}

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

    #[test]
    fn test_interpolation_table() {
        let table = InterpolationTable::new(0.0, 10.0, 5, |x| x.ln_1p());

        for &x in &[0.02, 0.04, 0.45678686, 0.23875, 1.45345e-6] {
            assert_relative_eq!(table.get(x), x.ln_1p(), epsilon = 0.00001);
        }
    }
}
<commit_msg>Doctests and Complexity for data_structures::interpolation_table (#318)<commit_after>// Copyright 2018 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Fast lookup table for arbitrary floating point functions.
//! # Examples
//! ## Easy:
//! ```
//! use bio::data_structures::interpolation_table::*;
//!
//! let table = InterpolationTable::new(0.0, 10.0, 5, |x| x.powf(2.0));
//! assert_eq!(table.get(3.0), 9.0);
//! assert_eq!(table.get(5.0), 25.0);
//! ```
//! ## More complicated:
//! ```
//! extern crate approx;
//! fn main() {
//!     use bio::data_structures::interpolation_table::*;
//!     use approx::assert_relative_eq;
//!
//!     let table = InterpolationTable::new(0.0, 10.0, 5, |x| x.ln_1p());
//!     for &x in &[0.02, 0.04, 0.45678686, 0.23875, 1.45345e-6] {
//!         assert_relative_eq!(table.get(x), x.ln_1p(), epsilon = 0.00001);
//!     }
//! }

#[inline]
pub fn interpolate(a: f64, b: f64, fraction: f64) -> f64 {
    a * (1.0 - fraction) + b * fraction
}

/// Fast lookup table for arbitrary floating point functions.
/// This can be used to e.g., provide fast lookups of distribution values.
/// Input values are sampled with a given precision and results are stored in a vector.
/// During lookup, infimum and supremum of a given value are calculated and the result is
/// interpolated.
pub struct InterpolationTable<F: Fn(f64) -> f64> {
    inner: Vec<f64>,
    func: F,
    offset: usize,
    min_x: f64,
    max_x: f64,
    shift: f64,
}

impl<F: Fn(f64) -> f64> InterpolationTable<F> {
    /// Create a new `InterpolationTable`.
    ///
    /// # Arguments
    ///
    /// * `min_x` - minimum sample value
    /// * `max_x` - maximum sample value
    /// * `frac_digits` - number of fraction digits to store in sample
    /// * `func` - Function to emulate.
    ///
    /// If given value is outside of min_x and max_x, the lookup falls back to applying the
    /// function itself.
    /// The table size grows with the number of fraction digits.
    /// Space Complexity: O(m * 10^n), where `m = max_x - min_x` and `n = frac_digits`
    pub fn new(min_x: f64, max_x: f64, frac_digits: i32, func: F) -> Self {
        let shift = 10.0_f64.powi(frac_digits);
        let offset = (min_x * shift) as usize;
        let mut table = InterpolationTable {
            inner: Vec::new(),
            func,
            min_x,
            max_x,
            shift,
            offset,
        };

        let mut i = table.index(min_x);
        while i < table.index(max_x) {
            let x = i as f64 / shift;
            table.inner.push((table.func)(x));
            i += 1;
        }

        table
    }

    /// Return lower bound index for given f64.
    #[inline]
    fn index(&self, x: f64) -> usize {
        (x * self.shift) as usize - self.offset
    }

    /// Lookup given value in table, and interpolate the result between the sampled values if
    /// necessary. This provides an approximation that is better the more fraction digits are
    /// used to generate this table.
    /// Time Complexity for lookup: O(1) if `min_x <= x < max_x` and O(func(x)) otherwise.
    pub fn get(&self, x: f64) -> f64 {
        if x < self.min_x || x >= self.max_x {
            (self.func)(x)
        } else {
            let i = self.index(x);
            // interpolate
            let fraction = (x * self.shift - i as f64) / self.shift;
            interpolate(self.inner[i], self.inner[i + 1], fraction)
        }
    }
}

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

    #[test]
    fn test_interpolation_table() {
        let table = InterpolationTable::new(0.0, 10.0, 5, |x| x.ln_1p());

        for &x in &[0.02, 0.04, 0.45678686, 0.23875, 1.45345e-6] {
            assert_relative_eq!(table.get(x), x.ln_1p(), epsilon = 0.00001);
        }
    }
}
<|endoftext|>


<commit_before>use crate::iter::Iter;
use crate::{CopyInWriter, CopyOutReader, Portal, Statement, ToStatement};
use fallible_iterator::FallibleIterator;
use futures::executor;
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::{Error, Row, SimpleQueryMessage};

/// A representation of a PostgreSQL database transaction.
///
/// Transactions will implicitly roll back by default when dropped. Use the `commit` method to commit the changes made
/// in the transaction. Transactions can be nested, with inner transactions implemented via safepoints.
pub struct Transaction<'a>(tokio_postgres::Transaction<'a>);

impl<'a> Transaction<'a> {
    pub(crate) fn new(transaction: tokio_postgres::Transaction<'a>) -> Transaction<'a> {
        Transaction(transaction)
    }

    /// Consumes the transaction, committing all changes made within it.
    pub fn commit(self) -> Result<(), Error> {
        executor::block_on(self.0.commit())
    }

    /// Rolls the transaction back, discarding all changes made within it.
    ///
    /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
    pub fn rollback(self) -> Result<(), Error> {
        executor::block_on(self.0.rollback())
    }

    /// Like `Client::prepare`.
    pub fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
        executor::block_on(self.0.prepare(query))
    }

    /// Like `Client::prepare_typed`.
    pub fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
        executor::block_on(self.0.prepare_typed(query, types))
    }

    /// Like `Client::execute`.
    pub fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.execute(query, params))
    }

    /// Like `Client::query`.
    pub fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Row>, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.query(query, params))
    }

    /// Like `Client::query_one`.
    pub fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.query_one(query, params))
    }

    /// Like `Client::query_raw`.
    pub fn query_raw<'b, T, I>(
        &mut self,
        query: &T,
        params: I,
    ) -> Result<impl FallibleIterator<Item = Row, Error = Error>, Error>
    where
        T: ?Sized + ToStatement,
        I: IntoIterator<Item = &'b dyn ToSql>,
        I::IntoIter: ExactSizeIterator,
    {
        let stream = executor::block_on(self.0.query_raw(query, params))?;
        Ok(Iter::new(stream))
    }

    /// Binds parameters to a statement, creating a "portal".
    ///
    /// Portals can be used with the `query_portal` method to page through the results of a query without being forced
    /// to consume them all immediately.
    ///
    /// Portals are automatically closed when the transaction they were created in is closed.
    ///
    /// # Panics
    ///
    /// Panics if the number of parameters provided does not match the number expected.
    pub fn bind<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Portal, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.bind(query, params))
    }

    /// Continues execution of a portal, returning the next set of rows.
    ///
    /// Unlike `query`, portals can be incrementally evaluated by limiting the number of rows returned in each call to
    /// `query_portal`. If the requested number is negative or 0, all remaining rows will be returned.
    pub fn query_portal(&mut self, portal: &Portal, max_rows: i32) -> Result<Vec<Row>, Error> {
        executor::block_on(self.0.query_portal(portal, max_rows))
    }

    /// The maximally flexible version of `query_portal`.
    pub fn query_portal_raw(
        &mut self,
        portal: &Portal,
        max_rows: i32,
    ) -> Result<impl FallibleIterator<Item = Row, Error = Error>, Error> {
        let stream = executor::block_on(self.0.query_portal_raw(portal, max_rows))?;
        Ok(Iter::new(stream))
    }

    /// Like `Client::copy_in`.
    pub fn copy_in<T>(
        &mut self,
        query: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<CopyInWriter<'_>, Error>
    where
        T: ?Sized + ToStatement,
    {
        let sink = executor::block_on(self.0.copy_in(query, params))?;
        Ok(CopyInWriter::new(sink))
    }

    /// Like `Client::copy_out`.
    pub fn copy_out<T>(
        &mut self,
        query: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<CopyOutReader<'_>, Error>
    where
        T: ?Sized + ToStatement,
    {
        let stream = executor::block_on(self.0.copy_out(query, params))?;
        CopyOutReader::new(stream)
    }

    /// Like `Client::simple_query`.
    pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
        executor::block_on(self.0.simple_query(query))
    }

    /// Like `Client::batch_execute`.
    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
        executor::block_on(self.0.batch_execute(query))
    }

    /// Like `Client::transaction`.
    pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
        let transaction = executor::block_on(self.0.transaction())?;
        Ok(Transaction(transaction))
    }
}
<commit_msg>doc fix<commit_after>use crate::iter::Iter;
use crate::{CopyInWriter, CopyOutReader, Portal, Statement, ToStatement};
use fallible_iterator::FallibleIterator;
use futures::executor;
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::{Error, Row, SimpleQueryMessage};

/// A representation of a PostgreSQL database transaction.
///
/// Transactions will implicitly roll back by default when dropped. Use the `commit` method to commit the changes made
/// in the transaction. Transactions can be nested, with inner transactions implemented via savepoints.
pub struct Transaction<'a>(tokio_postgres::Transaction<'a>);

impl<'a> Transaction<'a> {
    pub(crate) fn new(transaction: tokio_postgres::Transaction<'a>) -> Transaction<'a> {
        Transaction(transaction)
    }

    /// Consumes the transaction, committing all changes made within it.
    pub fn commit(self) -> Result<(), Error> {
        executor::block_on(self.0.commit())
    }

    /// Rolls the transaction back, discarding all changes made within it.
    ///
    /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
    pub fn rollback(self) -> Result<(), Error> {
        executor::block_on(self.0.rollback())
    }

    /// Like `Client::prepare`.
    pub fn prepare(&mut self, query: &str) -> Result<Statement, Error> {
        executor::block_on(self.0.prepare(query))
    }

    /// Like `Client::prepare_typed`.
    pub fn prepare_typed(&mut self, query: &str, types: &[Type]) -> Result<Statement, Error> {
        executor::block_on(self.0.prepare_typed(query, types))
    }

    /// Like `Client::execute`.
    pub fn execute<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<u64, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.execute(query, params))
    }

    /// Like `Client::query`.
    pub fn query<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Vec<Row>, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.query(query, params))
    }

    /// Like `Client::query_one`.
    pub fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.query_one(query, params))
    }

    /// Like `Client::query_raw`.
    pub fn query_raw<'b, T, I>(
        &mut self,
        query: &T,
        params: I,
    ) -> Result<impl FallibleIterator<Item = Row, Error = Error>, Error>
    where
        T: ?Sized + ToStatement,
        I: IntoIterator<Item = &'b dyn ToSql>,
        I::IntoIter: ExactSizeIterator,
    {
        let stream = executor::block_on(self.0.query_raw(query, params))?;
        Ok(Iter::new(stream))
    }

    /// Binds parameters to a statement, creating a "portal".
    ///
    /// Portals can be used with the `query_portal` method to page through the results of a query without being forced
    /// to consume them all immediately.
    ///
    /// Portals are automatically closed when the transaction they were created in is closed.
    ///
    /// # Panics
    ///
    /// Panics if the number of parameters provided does not match the number expected.
    pub fn bind<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Portal, Error>
    where
        T: ?Sized + ToStatement,
    {
        executor::block_on(self.0.bind(query, params))
    }

    /// Continues execution of a portal, returning the next set of rows.
    ///
    /// Unlike `query`, portals can be incrementally evaluated by limiting the number of rows returned in each call to
    /// `query_portal`. If the requested number is negative or 0, all remaining rows will be returned.
    pub fn query_portal(&mut self, portal: &Portal, max_rows: i32) -> Result<Vec<Row>, Error> {
        executor::block_on(self.0.query_portal(portal, max_rows))
    }

    /// The maximally flexible version of `query_portal`.
    pub fn query_portal_raw(
        &mut self,
        portal: &Portal,
        max_rows: i32,
    ) -> Result<impl FallibleIterator<Item = Row, Error = Error>, Error> {
        let stream = executor::block_on(self.0.query_portal_raw(portal, max_rows))?;
        Ok(Iter::new(stream))
    }

    /// Like `Client::copy_in`.
    pub fn copy_in<T>(
        &mut self,
        query: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<CopyInWriter<'_>, Error>
    where
        T: ?Sized + ToStatement,
    {
        let sink = executor::block_on(self.0.copy_in(query, params))?;
        Ok(CopyInWriter::new(sink))
    }

    /// Like `Client::copy_out`.
    pub fn copy_out<T>(
        &mut self,
        query: &T,
        params: &[&(dyn ToSql + Sync)],
    ) -> Result<CopyOutReader<'_>, Error>
    where
        T: ?Sized + ToStatement,
    {
        let stream = executor::block_on(self.0.copy_out(query, params))?;
        CopyOutReader::new(stream)
    }

    /// Like `Client::simple_query`.
    pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
        executor::block_on(self.0.simple_query(query))
    }

    /// Like `Client::batch_execute`.
    pub fn batch_execute(&mut self, query: &str) -> Result<(), Error> {
        executor::block_on(self.0.batch_execute(query))
    }

    /// Like `Client::transaction`.
    pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
        let transaction = executor::block_on(self.0.transaction())?;
        Ok(Transaction(transaction))
    }
}
<|endoftext|>