<commit_before>//! Functions for suppressing non-maximal values.

use definitions::{
    Position,
    Score
};

use image::{
    GenericImage,
    ImageBuffer,
    Luma,
    Primitive};

/// Zeroes all pixels which do not have the greatest intensity in the
/// (2 * radius + 1) square block centred on them. Ties are resolved lexicographically.
// TODO: Implement a more efficient version.
// TODO: e.g. https://www.vision.ee.ethz.ch/en/publications/papers/proceedings/eth_biwi_00446.pdf
pub fn suppress_non_maximum_mut<I, C>(image: &mut I, radius: u32)
    where I: GenericImage<Pixel = Luma<C>>,
          C: Primitive + Ord + 'static
{
    let (width, height) = image.dimensions();
    let irad = radius as i32;

    for y in 0..height {
        for x in 0..width {
            let intensity = image.get_pixel(x, y)[0];
            let mut is_max = true;

            for dy in -irad..irad {
                let py = (y as i32) + dy;
                if py < 0 || py >= height as i32 {
                    continue;
                }
                for dx in -irad..irad {
                    let px = (x as i32) + dx;
                    if px < 0 || px >= width as i32 {
                        continue;
                    }

                    let v = image.get_pixel(px as u32, py as u32)[0];
                    if v > intensity ||
                       (v == intensity && (px as u32, py as u32) < (x, y)){
                        is_max = false;
                        break;
                    }
                }
            }
            if !is_max {
                image.put_pixel(x, y, Luma([C::zero()]));
            }
        }
    }
}

/// Returned image has zeroes for all inputs pixels which do not have the greatest
/// intensity in the (2 * radius + 1) square block centred on them.
/// Ties are resolved lexicographically.
pub fn suppress_non_maximum<I, C>(image: &I, radius: u32) -> ImageBuffer<Luma<C>, Vec<C>>
    where I: GenericImage<Pixel = Luma<C>>,
          C: Primitive + Ord + 'static
{
    let mut out: ImageBuffer<Luma<C>, Vec<C>> = ImageBuffer::new(image.width(), image.height());
    out.copy_from(image, 0, 0);
    suppress_non_maximum_mut(&mut out, radius);
    out
}

/// Returns all items which have the highest score in the
/// (2 * radius + 1) square block centred on them. Ties are resolved lexicographically.
pub fn local_maxima<T>(ts: &[T], radius: u32) -> Vec<T>
    where T: Position + Score + Copy
 {
    let mut ordered_ts = ts.to_vec();
    ordered_ts.sort_by(|c, d| {(c.y(), c.x()).cmp(&(d.y(), d.x()))});
    let height = match ordered_ts.last() {
        Some(t) => t.y(),
        None => 0
    };

    let mut ts_by_row = vec![vec![]; (height + 1) as usize];
    for t in ordered_ts.iter() {
        ts_by_row[t.y() as usize].push(t);
    }

    let mut max_ts = vec![];
    for t in ordered_ts.iter() {
        let cx = t.x();
        let cy = t.y();
        let cs = t.score();

        let mut is_max = true;
        let min_row = if radius > cy {0} else {cy - radius};
        let max_row = if cy + radius > height {height} else {cy + radius};
        for y in min_row..max_row {
            for c in ts_by_row[y as usize].iter() {
                if c.x() + radius < cx {
                    continue;
                }
                if c.x() > cx + radius {
                    break;
                }
                if c.score() > cs {
                    is_max = false;
                    break;
                }
                if c.score() < cs {
                    continue;
                }
                // Break tiebreaks lexicographically
                if (c.y(), c.x()) < (cy, cx) {
                    is_max = false;
                    break;
                }
            }
            if !is_max {
                break;
            }
        }

        if is_max {
            max_ts.push(*t);
        }
    }

    max_ts
}

#[cfg(test)]
mod test {

    use super::{
        local_maxima,
        suppress_non_maximum
    };
    use definitions::{
        Position,
        Score
    };
    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };
    use noise::{
        gaussian_noise_mut
    };
    use test::Bencher;

    #[derive(PartialEq, Debug, Copy, Clone)]
    struct T {
        x: u32,
        y: u32,
        score: f32
    }

    impl T {
        fn new(x: u32, y: u32, score: f32) -> T {
            T { x: x, y: y, score: score}
        }
    }

    impl Position for T {
        fn x(&self) -> u32 { self.x }
        fn y(&self) -> u32 { self.y }
    }

    impl Score for T {
        fn score(&self) -> f32 { self.score }
    }

    #[test]
    fn test_local_maxima() {
        let ts = vec![
            // Suppress vertically
            T::new(0, 0, 10f32),
            T::new(0, 2, 8f32),
            // Suppress horizontally
            T::new(5, 5, 10f32),
            T::new(7, 5, 15f32),
            // Tiebreak
            T::new(12, 20, 10f32),
            T::new(13, 20, 10f32),
            T::new(13, 21, 10f32)
        ];

        let expected = vec![
            T::new(0, 0, 10f32),
            T::new(7, 5, 15f32),
            T::new(12, 20, 10f32)
        ];

        let max = local_maxima(&ts, 3);
        assert_eq!(max, expected);
    }

    #[bench]
    fn bench_local_maxima_dense(b: &mut Bencher) {
        let mut ts = vec![];
        for x in 0..20 {
            for y in 0..20 {
                let score = (x * y) % 15;
                ts.push(T::new(x, y, score as f32));
            }
        }
        b.iter(|| local_maxima(&ts, 15));
    }

    #[bench]
    fn bench_local_maxima_sparse(b: &mut Bencher) {
        let mut ts = vec![];
        for x in 0..20 {
            for y in 0..20 {
                ts.push(T::new(50 * x, 50 * y, 50f32));
            }
        }
        b.iter(|| local_maxima(&ts, 15));
    }

    #[test]
    fn test_suppress_non_maximum() {
        let mut image = GrayImage::new(25, 25);
        // Suppress vertically
        image.put_pixel(0, 0, Luma([10u8]));
        image.put_pixel(0, 2, Luma([8u8]));
        // Suppress horizontally
        image.put_pixel(5, 5, Luma([10u8]));
        image.put_pixel(7, 5, Luma([15u8]));
        // Tiebreak
        image.put_pixel(12, 20, Luma([10u8]));
        image.put_pixel(13, 20, Luma([10u8]));
        image.put_pixel(13, 21, Luma([10u8]));

        let mut expected = GrayImage::new(25, 25);
        expected.put_pixel(0, 0, Luma([10u8]));
        expected.put_pixel(7, 5, Luma([15u8]));
        expected.put_pixel(12, 20, Luma([10u8]));

        let actual = suppress_non_maximum(&image, 3);
        assert_pixels_eq!(actual, expected);
    }

    #[bench]
    fn bench_suppress_non_maximum_increasing_gradient(b: &mut Bencher) {
        // Increasing gradient in both directions. This can be a worst-case for
        // early-abort strategies.
        let img = ImageBuffer::from_fn(40, 20, |x, y| Luma([(x + y) as u8]));
        b.iter(|| suppress_non_maximum(&img, 7));
    }

    #[bench]
    fn bench_suppress_non_maximum_decreasing_gradient(b: &mut Bencher) {
        let width = 40u32;
        let height = 20u32;
        let img = ImageBuffer::from_fn(width, height, |x, y| {
            Luma([((width - x) + (height - y)) as u8])
        });
        b.iter(|| suppress_non_maximum(&img, 7));
    }

    #[bench]
    fn bench_suppress_non_maximum_noise(b: &mut Bencher) {
        let mut img: GrayImage = ImageBuffer::new(40, 20);
        gaussian_noise_mut(&mut img, 128f64, 30f64, 1);
        b.iter(|| suppress_non_maximum(&img, 7));
    }
}
<commit_msg>simpler suppress_non_maximum_mut<commit_after>//! Functions for suppressing non-maximal values.

use std::cmp;

use definitions::{
    Position,
    Score
};

use image::{
    GenericImage,
    ImageBuffer,
    Luma,
    Primitive};

/// Zeroes all pixels which do not have the greatest intensity in the
/// (2 * radius + 1) square block centred on them. Ties are resolved lexicographically.
// TODO: Implement a more efficient version.
// TODO: e.g. https://www.vision.ee.ethz.ch/en/publications/papers/proceedings/eth_biwi_00446.pdf
pub fn suppress_non_maximum_mut<I, C>(image: &mut I, radius: u32)
    where I: GenericImage<Pixel = Luma<C>>,
          C: Primitive + Ord + 'static
{
    let (width, height) = image.dimensions();
    let irad = radius as i32;

    for y in 0..height {
        'test: for x in 0..width {
            let intensity = unsafe { image.unsafe_get_pixel(x, y)[0] };

            // if any pixel in the box has a higer intensity, don't modify the pixel
            for py in cmp::max(0, y as i32 - irad)..cmp::min(y as i32 + irad, height as i32) {
                for px in cmp::max(0, x as i32 - irad)..cmp::min(x as i32 + irad, width as i32) {
                    let v = unsafe { image.unsafe_get_pixel(px as u32, py as u32)[0] };
                    if v > intensity ||
                       (v == intensity && (px as u32, py as u32) < (x, y)){
                        unsafe { image.unsafe_put_pixel(x, y, Luma([C::zero()])) };
                        continue 'test;
                    }
                }
            }
        }
    }
}

/// Returned image has zeroes for all inputs pixels which do not have the greatest
/// intensity in the (2 * radius + 1) square block centred on them.
/// Ties are resolved lexicographically.
pub fn suppress_non_maximum<I, C>(image: &I, radius: u32) -> ImageBuffer<Luma<C>, Vec<C>>
    where I: GenericImage<Pixel = Luma<C>>,
          C: Primitive + Ord + 'static
{
    let mut out: ImageBuffer<Luma<C>, Vec<C>> = ImageBuffer::new(image.width(), image.height());
    out.copy_from(image, 0, 0);
    suppress_non_maximum_mut(&mut out, radius);
    out
}

/// Returns all items which have the highest score in the
/// (2 * radius + 1) square block centred on them. Ties are resolved lexicographically.
pub fn local_maxima<T>(ts: &[T], radius: u32) -> Vec<T>
    where T: Position + Score + Copy
 {
    let mut ordered_ts = ts.to_vec();
    ordered_ts.sort_by(|c, d| {(c.y(), c.x()).cmp(&(d.y(), d.x()))});
    let height = match ordered_ts.last() {
        Some(t) => t.y(),
        None => 0
    };

    let mut ts_by_row = vec![vec![]; (height + 1) as usize];
    for t in ordered_ts.iter() {
        ts_by_row[t.y() as usize].push(t);
    }

    let mut max_ts = vec![];
    for t in ordered_ts.iter() {
        let cx = t.x();
        let cy = t.y();
        let cs = t.score();

        let mut is_max = true;
        let min_row = if radius > cy {0} else {cy - radius};
        let max_row = if cy + radius > height {height} else {cy + radius};
        for y in min_row..max_row {
            for c in ts_by_row[y as usize].iter() {
                if c.x() + radius < cx {
                    continue;
                }
                if c.x() > cx + radius {
                    break;
                }
                if c.score() > cs {
                    is_max = false;
                    break;
                }
                if c.score() < cs {
                    continue;
                }
                // Break tiebreaks lexicographically
                if (c.y(), c.x()) < (cy, cx) {
                    is_max = false;
                    break;
                }
            }
            if !is_max {
                break;
            }
        }

        if is_max {
            max_ts.push(*t);
        }
    }

    max_ts
}

#[cfg(test)]
mod test {

    use super::{
        local_maxima,
        suppress_non_maximum
    };
    use definitions::{
        Position,
        Score
    };
    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };
    use noise::{
        gaussian_noise_mut
    };
    use test::Bencher;

    #[derive(PartialEq, Debug, Copy, Clone)]
    struct T {
        x: u32,
        y: u32,
        score: f32
    }

    impl T {
        fn new(x: u32, y: u32, score: f32) -> T {
            T { x: x, y: y, score: score}
        }
    }

    impl Position for T {
        fn x(&self) -> u32 { self.x }
        fn y(&self) -> u32 { self.y }
    }

    impl Score for T {
        fn score(&self) -> f32 { self.score }
    }

    #[test]
    fn test_local_maxima() {
        let ts = vec![
            // Suppress vertically
            T::new(0, 0, 10f32),
            T::new(0, 2, 8f32),
            // Suppress horizontally
            T::new(5, 5, 10f32),
            T::new(7, 5, 15f32),
            // Tiebreak
            T::new(12, 20, 10f32),
            T::new(13, 20, 10f32),
            T::new(13, 21, 10f32)
        ];

        let expected = vec![
            T::new(0, 0, 10f32),
            T::new(7, 5, 15f32),
            T::new(12, 20, 10f32)
        ];

        let max = local_maxima(&ts, 3);
        assert_eq!(max, expected);
    }

    #[bench]
    fn bench_local_maxima_dense(b: &mut Bencher) {
        let mut ts = vec![];
        for x in 0..20 {
            for y in 0..20 {
                let score = (x * y) % 15;
                ts.push(T::new(x, y, score as f32));
            }
        }
        b.iter(|| local_maxima(&ts, 15));
    }

    #[bench]
    fn bench_local_maxima_sparse(b: &mut Bencher) {
        let mut ts = vec![];
        for x in 0..20 {
            for y in 0..20 {
                ts.push(T::new(50 * x, 50 * y, 50f32));
            }
        }
        b.iter(|| local_maxima(&ts, 15));
    }

    #[test]
    fn test_suppress_non_maximum() {
        let mut image = GrayImage::new(25, 25);
        // Suppress vertically
        image.put_pixel(0, 0, Luma([10u8]));
        image.put_pixel(0, 2, Luma([8u8]));
        // Suppress horizontally
        image.put_pixel(5, 5, Luma([10u8]));
        image.put_pixel(7, 5, Luma([15u8]));
        // Tiebreak
        image.put_pixel(12, 20, Luma([10u8]));
        image.put_pixel(13, 20, Luma([10u8]));
        image.put_pixel(13, 21, Luma([10u8]));

        let mut expected = GrayImage::new(25, 25);
        expected.put_pixel(0, 0, Luma([10u8]));
        expected.put_pixel(7, 5, Luma([15u8]));
        expected.put_pixel(12, 20, Luma([10u8]));

        let actual = suppress_non_maximum(&image, 3);
        assert_pixels_eq!(actual, expected);
    }

    #[bench]
    fn bench_suppress_non_maximum_increasing_gradient(b: &mut Bencher) {
        // Increasing gradient in both directions. This can be a worst-case for
        // early-abort strategies.
        let img = ImageBuffer::from_fn(40, 20, |x, y| Luma([(x + y) as u8]));
        b.iter(|| suppress_non_maximum(&img, 7));
    }

    #[bench]
    fn bench_suppress_non_maximum_decreasing_gradient(b: &mut Bencher) {
        let width = 40u32;
        let height = 20u32;
        let img = ImageBuffer::from_fn(width, height, |x, y| {
            Luma([((width - x) + (height - y)) as u8])
        });
        b.iter(|| suppress_non_maximum(&img, 7));
    }

    #[bench]
    fn bench_suppress_non_maximum_noise(b: &mut Bencher) {
        let mut img: GrayImage = ImageBuffer::new(40, 20);
        gaussian_noise_mut(&mut img, 128f64, 30f64, 1);
        b.iter(|| suppress_non_maximum(&img, 7));
    }
}
<|endoftext|>


<commit_before>use core::fail::{abort, out_of_memory};
use core::ptr::offset;
use core::ptr::{set_memory, copy_memory};
use core::i32::ctlz32;

use kernel::ptr::mut_offset;

enum Node {
    UNUSED = 0,
    USED = 1,
    SPLIT = 2,
    FULL = 3
}

pub trait Allocator {
    unsafe fn alloc(&mut self, size: uint) -> (*mut u8, uint);
    unsafe fn zero_alloc(&mut self, size: uint) -> (*mut u8, uint);
    unsafe fn realloc(&mut self, ptr: *mut u8, size: uint) -> (*mut u8, uint);
    unsafe fn free(&mut self, ptr: *mut u8);
}

trait BitvTrait {
    fn get(&self, i: uint) -> Node;
    fn set(&self, i: uint, x: Node);
    fn to_bytes(&self) -> *mut u8;
    fn size(&self) -> uint;
}

static BITV_SIZE: uint = 0x10_000;
pub type BitvStorage = *mut [u32, ..BITV_SIZE / 4];

// vector of 2-bit
pub struct Bitv {
    storage: BitvStorage
}

impl BitvTrait for Bitv {
    #[inline]
    fn get(&self, i: uint) -> Node {
        let w = i / 16;
        let b = (i % 16) * 2;
        unsafe { [UNUSED, USED, SPLIT, FULL][((*self.storage)[w] as uint >> b) & 3] }
    }

    #[inline]
    fn set(&self, i: uint, x: Node) {
        let w = i / 16;
        let b = (i % 16) * 2;
        unsafe { (*self.storage)[w] = (((*self.storage)[w] & !(3 << b)) | (x as u32 << b)); }
    }

    #[inline]
    fn to_bytes(&self) -> *mut u8 {
        self.storage as *mut u8
    }

    #[inline]
    fn size(&self) -> uint {
        BITV_SIZE
    }
}

pub struct BuddyAlloc {
    base: *mut u8,
    order: uint,
    tree: Bitv
}

impl BuddyAlloc {
    pub fn new(base: *mut u8, order: uint, storage: Bitv) -> BuddyAlloc {
        unsafe { set_memory(storage.to_bytes(), 0, storage.size()); }

        BuddyAlloc { base: base, order: order, tree: storage }
    }

    #[inline]
    fn offset(&self, index: uint, level: uint) -> *mut u8 {
        unsafe {
            mut_offset(self.base, (index + 1 - (1 << (self.order - level))) as int << level)
        }
    }
}

impl Allocator for BuddyAlloc {
    fn alloc(&mut self, mut size: uint) -> (*mut u8, uint) {
        if size == 0 {
            size = 1;
        }
        // smallest power of 2 >= size
        let lg2_size = 32 - unsafe { ctlz32(size as i32 - 1) } as uint;

        let mut index = 0; // points to current tree node
        let mut level = self.order; // current height

        loop {
            match (self.tree.get(index), level == lg2_size) {
                (UNUSED, true) => {
                    // Found appropriate unused node
                    self.tree.set(index, USED); // use

                    let mut parent = index;
                    loop {
                        let buddy = parent - 1 + (parent & 1) * 2;
                        match self.tree.get(buddy) {
                            USED | FULL if parent > 0 => {
                                parent = (parent + 1) / 2 - 1;
                                self.tree.set(parent, FULL);
                            }
                            _ => break
                        }
                    }
                    return (
                        self.offset(index, level),
                        1 << lg2_size
                    );
                }
                (UNUSED, false) => {
                    // This large node is unused, split it!
                    self.tree.set(index, SPLIT);
                    self.tree.set(index*2 + 1, UNUSED);
                    self.tree.set(index*2 + 2, UNUSED);
                    index = index * 2 + 1; // left child
                    level -= 1;
                }
                (SPLIT, false) => {
                    // Traverse children
                    index = index * 2 + 1; // left child
                    level -= 1;
                }
                _ => loop {
                    // Go either right or back up
                    if index & 1 == 1 {
                        // right sibling
                        index += 1;
                        break;
                    }

                    // go up by one level
                    level += 1;

                    if index == 0 {
                        // out of memory -- back at tree's root after traversal
                        return (self.base, 0);
                    }

                    index = (index + 1) / 2 - 1; // parent
                }
            }
        }
    }

    fn zero_alloc(&mut self, s: uint) -> (*mut u8, uint) {
        let (ptr, size) = self.alloc(s);
        unsafe { set_memory(ptr, 0, size); }
        (ptr, size)
    }

    fn realloc(&mut self, src: *mut u8, size: uint) -> (*mut u8, uint) {
        self.free(src);
        let (ptr, sz) = self.alloc(size);
        unsafe { copy_memory(ptr, src as *u8, sz); }
        (ptr, sz)
    }

    fn free(&mut self, ptr: *mut u8) {
        let mut length = 1 << self.order;
        let mut left = 0;
        let mut index = 0;

        if ((ptr as uint) < self.base as uint) || (ptr as uint >= self.base as uint + length) {
            return;
        }
        let offset = ptr as uint - self.base as uint;

        loop {
            match self.tree.get(index) {
                UNUSED => return,
                USED => loop {
                    if index == 0 {
                        self.tree.set(0, UNUSED);
                        return;
                    }

                    let buddy = index - 1 + (index & 1) * 2;
                    match self.tree.get(buddy) {
                        UNUSED => {}
                        _ => {
                            self.tree.set(index, UNUSED);
                            loop {
                                let parent = (index + 1) / 2 - 1; // parent
                                match self.tree.get(parent) {
                                    FULL if index > 0 => {
                                        self.tree.set(parent, SPLIT);
                                    }
                                    _ => return
                                }
                                index = parent;
                            }
                        }
                    }
                    index = (index + 1) / 2 - 1; // parent
                },
                _ => {
                    length /= 2;
                    if offset < left + length {
                        index = index * 2 + 1; // left child
                    }
                    else {
                        left += length;
                        index = index * 2 + 2; // right child
                    }
                }
            }
        }
    }
}
<commit_msg>Buddy: optimize<commit_after>use core::mem::transmute;
use core::ptr::{set_memory, copy_memory, offset};
use core::i32::ctlz32;

use kernel::ptr::mut_offset;

#[repr(u8)]
enum Node {
    UNUSED = 0,
    USED = 1,
    SPLIT = 2,
    FULL = 3
}

pub trait Allocator {
    unsafe fn alloc(&mut self, size: uint) -> (*mut u8, uint);
    unsafe fn zero_alloc(&mut self, size: uint) -> (*mut u8, uint);
    unsafe fn realloc(&mut self, ptr: *mut u8, size: uint) -> (*mut u8, uint);
    unsafe fn free(&mut self, ptr: *mut u8);
}

trait BitvTrait {
    fn get(&self, i: uint) -> Node;
    fn set(&self, i: uint, x: Node);
    fn to_bytes(&self) -> *mut u8;
    fn size(&self) -> uint;
}

static BITV_SIZE: uint = 0x10_000;
pub type BitvStorage = *mut [u32, ..BITV_SIZE / 4];

// vector of 2-bit
pub struct Bitv {
    storage: BitvStorage
}

impl BitvTrait for Bitv {
    #[inline]
    fn get(&self, i: uint) -> Node {
        let w = i / 16;
        let b = (i % 16) * 2;
        unsafe { transmute(((*self.storage)[w] as uint >> b) as u8 & 3) }
    }

    #[inline]
    fn set(&self, i: uint, x: Node) {
        let w = i / 16;
        let b = (i % 16) * 2;
        unsafe { (*self.storage)[w] = (((*self.storage)[w] & !(3 << b)) | (x as u32 << b)); }
    }

    #[inline]
    fn to_bytes(&self) -> *mut u8 {
        self.storage as *mut u8
    }

    #[inline]
    fn size(&self) -> uint {
        BITV_SIZE
    }
}

pub struct BuddyAlloc {
    base: *mut u8,
    order: uint,
    tree: Bitv
}

impl BuddyAlloc {
    pub fn new(base: *mut u8, order: uint, storage: Bitv) -> BuddyAlloc {
        unsafe { set_memory(storage.to_bytes(), 0, storage.size()); }

        BuddyAlloc { base: base, order: order, tree: storage }
    }

    #[inline]
    fn offset(&self, index: uint, level: uint) -> *mut u8 {
        unsafe {
            mut_offset(self.base, (index + 1 - (1 << (self.order - level))) as int << level)
        }
    }
}

impl Allocator for BuddyAlloc {
    fn alloc(&mut self, mut size: uint) -> (*mut u8, uint) {
        if size == 0 {
            size = 1;
        }
        // smallest power of 2 >= size
        let lg2_size = 32 - unsafe { ctlz32(size as i32 - 1) } as uint;

        let mut index = 0; // points to current tree node
        let mut level = self.order; // current height

        loop {
            match (self.tree.get(index), level == lg2_size) {
                (UNUSED, true) => {
                    // Found appropriate unused node
                    self.tree.set(index, USED); // use

                    let mut parent = index;
                    loop {
                        let buddy = parent - 1 + (parent & 1) * 2;
                        match self.tree.get(buddy) {
                            USED | FULL if parent > 0 => {
                                parent = (parent + 1) / 2 - 1;
                                self.tree.set(parent, FULL);
                            }
                            _ => break
                        }
                    }
                    return (
                        self.offset(index, level),
                        1 << lg2_size
                    );
                }
                (UNUSED, false) => {
                    // This large node is unused, split it!
                    self.tree.set(index, SPLIT);
                    self.tree.set(index*2 + 1, UNUSED);
                    self.tree.set(index*2 + 2, UNUSED);
                    index = index * 2 + 1; // left child
                    level -= 1;
                }
                (SPLIT, false) => {
                    // Traverse children
                    index = index * 2 + 1; // left child
                    level -= 1;
                }
                _ => loop {
                    // Go either right or back up
                    if index & 1 == 1 {
                        // right sibling
                        index += 1;
                        break;
                    }

                    // go up by one level
                    level += 1;

                    if index == 0 {
                        // out of memory -- back at tree's root after traversal
                        return (self.base, 0);
                    }

                    index = (index + 1) / 2 - 1; // parent
                }
            }
        }
    }

    fn zero_alloc(&mut self, s: uint) -> (*mut u8, uint) {
        let (ptr, size) = self.alloc(s);
        unsafe { set_memory(ptr, 0, size); }
        (ptr, size)
    }

    fn realloc(&mut self, src: *mut u8, size: uint) -> (*mut u8, uint) {
        self.free(src);
        let (ptr, sz) = self.alloc(size);
        unsafe { copy_memory(ptr, src as *u8, sz); }
        (ptr, sz)
    }

    fn free(&mut self, ptr: *mut u8) {
        let mut length = 1 << self.order;
        let mut left = 0;
        let mut index = 0;

        if ((ptr as uint) < self.base as uint) || (ptr as uint >= self.base as uint + length) {
            return;
        }
        let offset = ptr as uint - self.base as uint;

        loop {
            match self.tree.get(index) {
                UNUSED => return,
                USED => loop {
                    if index == 0 {
                        self.tree.set(0, UNUSED);
                        return;
                    }

                    let buddy = index - 1 + (index & 1) * 2;
                    match self.tree.get(buddy) {
                        UNUSED => {}
                        _ => {
                            self.tree.set(index, UNUSED);
                            loop {
                                let parent = (index + 1) / 2 - 1; // parent
                                match self.tree.get(parent) {
                                    FULL if index > 0 => {
                                        self.tree.set(parent, SPLIT);
                                    }
                                    _ => return
                                }
                                index = parent;
                            }
                        }
                    }
                    index = (index + 1) / 2 - 1; // parent
                },
                _ => {
                    length /= 2;
                    if offset < left + length {
                        index = index * 2 + 1; // left child
                    }
                    else {
                        left += length;
                        index = index * 2 + 2; // right child
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::rc::Rc;
use std::cell::RefCell;

type WordCell = Rc<RefCell<Word>>;
pub type AuthorCell = Rc<RefCell<Author>>;
pub type SourceCell = Rc<RefCell<Source>>;
type MessageCell = Rc<RefCell<Message>>;
type CategoryCell = Rc<RefCell<Category>>;
type ConversationCell = Rc<RefCell<Conversation>>;
type InstanceCell = Rc<RefCell<WordInstance>>;

fn wrap<T>(t: T) -> Rc<RefCell<T>> {
    Rc::new(RefCell::new(t))
}

// const RATIO_CONTAINED_BEFORE_COMBINATION: f64 = 0.8;

#[derive(Default)]
pub struct Lexicon {
    words: BTreeMap<String, WordCell>,
    sources: BTreeMap<String, SourceCell>,
    conversations: Vec<ConversationCell>,

    active_conversations: BTreeMap<*const Source, ConversationCell>,
}

impl Lexicon {
    pub fn source(&mut self, name: String) -> SourceCell {
        match self.sources.entry(name.clone()) {
            Entry::Vacant(v) => v.insert(wrap(Source{
                name: name,
                messages: 0,
                authors: BTreeMap::default(),
            })).clone(),
            Entry::Occupied(o) => o.get().clone(),
        }
    }
    pub fn author(&mut self, source: SourceCell, name: String) -> AuthorCell {
        match source.borrow_mut().authors.entry(name.clone()) {
            Entry::Vacant(v) => {
                let a = wrap(Author{
                    source: source.clone(),
                    name: name.clone(),
                });
                v.insert(a.clone());
                a
            },
            Entry::Occupied(o) => o.get().clone(),
        }
    }
    pub fn tell(&mut self, source: SourceCell, author: AuthorCell, content: String) -> Option<String> {
        let conversation = match self.active_conversations.entry(&*source.borrow() as *const Source) {
            Entry::Vacant(v) => {
                let c = wrap(Conversation{
                    source: source.clone(),
                    messages: Vec::new(),
                });
                v.insert(c.clone());
                self.conversations.push(c.clone());
                c
            },
            Entry::Occupied(o) => {
                o.get().clone()
            },
        };

        let message = wrap(Message{
            author: author.clone(),
            conversation: conversation.clone(),
            index: conversation.borrow().messages.len(),
            words: Vec::new(),
        });

        for word in content.split(' ').map(|s| {
            let string = s.to_string();
            match self.words.entry(string.clone()) {
                Entry::Vacant(v) => {
                    let w = wrap(Word{
                        name: string,
                        instances: Vec::new(),
                    });
                    v.insert(w.clone());
                    w
                },
                Entry::Occupied(o) => {
                    o.get().clone()
                },
            }
        }) {
            // Create empty category for the word
            let category = wrap(Category::default());
            // Create instance of the word
            let instance = wrap(WordInstance{
                word: word.clone(),
                category: category.clone(),
                message: message.clone(),
                index: message.borrow().words.len(),
            });
            // Insert word instance into the word, category, and message for future reference
            message.borrow_mut().words.push(instance.clone());
            category.borrow_mut().instances.push(instance.clone());
            word.borrow_mut().instances.push(instance);
        }

        // Increment the messages by 1 for the source
        source.borrow_mut().messages += 1;

        // TODO: Determine what to say back
        None
    }
}

struct Conversation {
    source: SourceCell,
    messages: Vec<MessageCell>,
}

pub struct Author {
    source: SourceCell,
    name: String,
}

pub struct Source {
    name: String,
    messages: u64,
    authors: BTreeMap<String, AuthorCell>,
}

struct WordInstance {
    word: WordCell,
    category: CategoryCell,
    message: MessageCell,
    index: usize,
}

struct Message {
    author: AuthorCell,
    conversation: ConversationCell,
    index: usize,
    words: Vec<InstanceCell>,
}

#[derive(Default)]
struct Category {
    supercategory: Option<CategoryCell>,
    instances: Vec<InstanceCell>,
    subcategories: Vec<CategoryCell>,
}

struct Word {
    name: String,
    instances: Vec<InstanceCell>,
}
<commit_msg>add message list to lexicon<commit_after>use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::rc::Rc;
use std::cell::RefCell;

type WordCell = Rc<RefCell<Word>>;
pub type AuthorCell = Rc<RefCell<Author>>;
pub type SourceCell = Rc<RefCell<Source>>;
type MessageCell = Rc<RefCell<Message>>;
type CategoryCell = Rc<RefCell<Category>>;
type ConversationCell = Rc<RefCell<Conversation>>;
type InstanceCell = Rc<RefCell<WordInstance>>;

fn wrap<T>(t: T) -> Rc<RefCell<T>> {
    Rc::new(RefCell::new(t))
}

// const RATIO_CONTAINED_BEFORE_COMBINATION: f64 = 0.8;

#[derive(Default)]
pub struct Lexicon {
    words: BTreeMap<String, WordCell>,
    sources: BTreeMap<String, SourceCell>,
    conversations: Vec<ConversationCell>,
    messages: Vec<MessageCell>,

    active_conversations: BTreeMap<*const Source, ConversationCell>,
}

impl Lexicon {
    pub fn source(&mut self, name: String) -> SourceCell {
        match self.sources.entry(name.clone()) {
            Entry::Vacant(v) => v.insert(wrap(Source{
                name: name,
                messages: 0,
                authors: BTreeMap::default(),
            })).clone(),
            Entry::Occupied(o) => o.get().clone(),
        }
    }
    pub fn author(&mut self, source: SourceCell, name: String) -> AuthorCell {
        match source.borrow_mut().authors.entry(name.clone()) {
            Entry::Vacant(v) => {
                let a = wrap(Author{
                    source: source.clone(),
                    name: name.clone(),
                });
                v.insert(a.clone());
                a
            },
            Entry::Occupied(o) => o.get().clone(),
        }
    }
    pub fn tell(&mut self, source: SourceCell, author: AuthorCell, content: String) -> Option<String> {
        let conversation = match self.active_conversations.entry(&*source.borrow() as *const Source) {
            Entry::Vacant(v) => {
                let c = wrap(Conversation{
                    source: source.clone(),
                    messages: Vec::new(),
                });
                v.insert(c.clone());
                self.conversations.push(c.clone());
                c
            },
            Entry::Occupied(o) => {
                o.get().clone()
            },
        };

        let message = wrap(Message{
            author: author.clone(),
            conversation: conversation.clone(),
            index: conversation.borrow().messages.len(),
            words: Vec::new(),
        });

        self.messages.push(message.clone());

        for word in content.split(' ').map(|s| {
            let string = s.to_string();
            match self.words.entry(string.clone()) {
                Entry::Vacant(v) => {
                    let w = wrap(Word{
                        name: string,
                        instances: Vec::new(),
                    });
                    v.insert(w.clone());
                    w
                },
                Entry::Occupied(o) => {
                    o.get().clone()
                },
            }
        }) {
            // Create empty category for the word
            let category = wrap(Category::default());
            // Create instance of the word
            let instance = wrap(WordInstance{
                word: word.clone(),
                category: category.clone(),
                message: message.clone(),
                index: message.borrow().words.len(),
            });
            // Insert word instance into the word, category, and message for future reference
            message.borrow_mut().words.push(instance.clone());
            category.borrow_mut().instances.push(instance.clone());
            word.borrow_mut().instances.push(instance);
        }

        // Increment the messages by 1 for the source
        source.borrow_mut().messages += 1;

        // TODO: Determine what to say back
        None
    }
}

struct Conversation {
    source: SourceCell,
    messages: Vec<MessageCell>,
}

pub struct Author {
    source: SourceCell,
    name: String,
}

pub struct Source {
    name: String,
    messages: u64,
    authors: BTreeMap<String, AuthorCell>,
}

struct WordInstance {
    word: WordCell,
    category: CategoryCell,
    message: MessageCell,
    index: usize,
}

struct Message {
    author: AuthorCell,
    conversation: ConversationCell,
    index: usize,
    words: Vec<InstanceCell>,
}

#[derive(Default)]
struct Category {
    supercategory: Option<CategoryCell>,
    instances: Vec<InstanceCell>,
    subcategories: Vec<CategoryCell>,
}

struct Word {
    name: String,
    instances: Vec<InstanceCell>,
}
<|endoftext|>


<commit_before>struct Graph<T> {
    nodes: Vec<Node<T>>,
}

fn main() {

}<commit_msg>added bash graph class<commit_after>use std::iter::repeat;

fn max(a: u32, b: u32) -> u32 { if a > b { a } else { b } }

#[derive(Clone)]
struct Edge {
    vertex: u32,
    weight: f32,
}

impl Edge {
    pub fn new(vertex_: u32, weight_: f32) -> Edge {
        Edge{vertex: vertex_, weight: weight_,}
    }
}

struct Graph {
    nodes: Vec<Vec<Edge>>,
}

impl Graph {
    pub fn new() -> Graph { Graph{nodes: Vec::new(),} }

    pub fn add_edge(&mut self, src: u32, dst: u32, weight: f32) {
        // only works with rust nightly:
        // if (max(src,dst)) > self.nodes.len() as u32 {
        //  self.nodes.resize((max(src, dst) + 1) as usize, Vec::new());
        //}

        let len = self.nodes.len();
        if (max(src, dst)) > len as u32 {
            use std::iter::repeat;
            let new_len = (max(src, dst) + 1) as usize;
            self.nodes.extend(repeat(Vec::new()).take(new_len - len))
        }
        self.nodes[src as usize].push(Edge::new(dst, weight));
    }
    pub fn bfs(&self, src: u32) {
        let mut current: u32 = src;
        let mut stack: Vec<u32> = Vec::new();
        let mut visited: Vec<bool> = Vec::new();

        // only works with rust nightly:
        // visited.resize(self.nodes.len(), false);

        let len = self.nodes.len() as usize;
        visited.extend(repeat(false).take(len));
        stack.push(current);

        while stack.len() > 0 {
            current = stack[0 as usize];
            stack.remove(0);
            if !visited[current as usize] {
                println!("current: {}" , current);
                visited[current as usize] = true;
            }

            for n in &self.nodes[current as usize] {
                let neighbor: u32 = n.vertex;
                if !visited[neighbor as usize] { stack.push(neighbor); }
            }
        }
    }
}

fn main() {
    let mut g1 = Graph::new();
    g1.add_edge(0, 1, 1.0);
    g1.add_edge(0, 2, 4.0);
    g1.add_edge(1, 0, 1.0);
    g1.add_edge(1, 2, 2.0);
    g1.add_edge(1, 3, 6.0);
    g1.add_edge(2, 0, 4.0);
    g1.add_edge(2, 1, 2.0);
    g1.add_edge(2, 3, 3.0);
    g1.add_edge(3, 1, 6.0);
    g1.add_edge(3, 2, 3.0);
    g1.bfs(0);
}<|endoftext|>


<commit_before>extern crate rsoundio;

use std::os::raw::c_int;
use std::f32::consts;
use std::thread;
use std::time::Duration;

#[test]
fn test_outstream() {
    println!("test: OutStream");
    let sio = rsoundio::SoundIo::new();
    assert!(sio.connect_backend(rsoundio::ffi::SioBackend::Dummy).is_none());
    println!("current backend: {}", sio.current_backend().unwrap());
    sio.flush_events();
    let fmt;
    let dev_idx = sio.default_output_device_index().unwrap();
    let mut frames: Vec<Vec<f32>> = vec![vec![], vec![]];
    let dev = sio.get_output_device(dev_idx).unwrap();
    println!("device: {}, ref_count: {}", dev, dev.ref_count());
    let mut stream = dev.create_outstream().unwrap();
    assert!(stream.open().is_none());
    fmt = stream.current_format().unwrap();
    println!("current format: {}", fmt);
    let f = 4400f32;
    let sr = stream.get_sample_rate();
    let layout = stream.get_layout();
    let channels = layout.channel_count();
    let phi = 2.0 * f * consts::PI / (sr as f32);
    let l: Vec<f32> = (0..4096).map(|i| f32::sin(i as f32 * phi)).collect();
    let r = l.clone();
    frames = vec![l, r];
    let cb = |out: rsoundio::OutStream, min_frame_count: i32, max_frame_count: i32| {
        let frames_written = out.write_stream_f32(min_frame_count, &frames).unwrap();
        assert!(frames_written > 0);
        out.get_latency().map(|latency| assert!(latency >= 0.0));
    };
    stream.register_write_callback(Box::new(cb));
    let ucb = |out: rsoundio::OutStream| println!("Underflow!");
    stream.register_underflow_callback(Box::new(ucb));
    let ecb = |out: rsoundio::OutStream, err: rsoundio::ffi::SioError| println!("Error: {}", err);
    stream.register_error_callback(Box::new(ecb));
    assert!(stream.start().is_none());
    sio.wait_events();
    println!("Received event!");
    assert!(stream.clear_buffer().is_none());
    assert!(stream.pause(true).is_none());
    assert!(stream.pause(false).is_none());
    println!("latency: {}", stream.get_latency().unwrap());
    stream.destroy();
}
<commit_msg>Pause and unpause playback in test run<commit_after>extern crate rsoundio;

use std::os::raw::c_int;
use std::f32::consts;
use std::thread;
use std::time::Duration;

#[test]
fn test_outstream() {
    println!("test: OutStream");
    let sio = rsoundio::SoundIo::new();
    assert!(sio.connect_backend(rsoundio::ffi::SioBackend::Dummy).is_none());
    println!("current backend: {}", sio.current_backend().unwrap());
    sio.flush_events();
    let fmt;
    let dev_idx = sio.default_output_device_index().unwrap();
    let mut frames: Vec<Vec<f32>> = vec![vec![], vec![]];
    let dev = sio.get_output_device(dev_idx).unwrap();
    println!("device: {}, ref_count: {}", dev, dev.ref_count());
    let mut stream = dev.create_outstream().unwrap();
    assert!(stream.open().is_none());
    fmt = stream.current_format().unwrap();
    println!("current format: {}", fmt);
    let f = 4400f32;
    let sr = stream.get_sample_rate();
    let layout = stream.get_layout();
    let channels = layout.channel_count();
    let phi = 2.0 * f * consts::PI / (sr as f32);
    let l: Vec<f32> = (0..4096).map(|i| f32::sin(i as f32 * phi)).collect();
    let r = l.clone();
    frames = vec![l, r];
    let cb = |out: rsoundio::OutStream, min_frame_count: i32, max_frame_count: i32| {
        let frames_written = out.write_stream_f32(min_frame_count, &frames).unwrap();
        assert!(frames_written > 0);
        out.get_latency().map(|latency| assert!(latency >= 0.0));
    };
    stream.register_write_callback(Box::new(cb));
    let ucb = |out: rsoundio::OutStream| println!("Underflow!");
    stream.register_underflow_callback(Box::new(ucb));
    let ecb = |out: rsoundio::OutStream, err: rsoundio::ffi::SioError| println!("Error: {}", err);
    stream.register_error_callback(Box::new(ecb));
    assert!(stream.start().is_none());
    //sio.wait_events();
    thread::sleep(Duration::new(1,0));
    assert!(stream.pause(true).is_none());
    thread::sleep(Duration::new(1,0));
    assert!(stream.pause(false).is_none());
    assert!(stream.clear_buffer().is_none());
    thread::sleep(Duration::new(1,0));
    stream.destroy();
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate cfg_if;

cfg_if! {
    if #[cfg(foo)] {
        fn works() -> bool { false }
    } else if #[cfg(test)] {
        fn works() -> bool { true }
    } else {
        fn works() -> bool { false }
    }
}

#[test]
fn smoke() {
    assert!(works());
}
<commit_msg>Update xcrate test to 2018<commit_after>cfg_if::cfg_if! {
    if #[cfg(foo)] {
        fn works() -> bool { false }
    } else if #[cfg(test)] {
        fn works() -> bool { true }
    } else {
        fn works() -> bool { false }
    }
}

#[test]
fn smoke() {
    assert!(works());
}
<|endoftext|>


<commit_before>use flowgger::config::Config;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::path::{Path, PathBuf};
use std::sync::Arc;

pub mod tls_input;
#[cfg(feature = "coroutines")]
pub mod tlsco_input;

pub use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_FRAMING: &'static str = "line";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
#[cfg(feature = "coroutines")]
const DEFAULT_THREADS: usize = 1;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
pub struct TlsConfig {
    framing: String,
    threads: usize,
    arc_ctx: Arc<SslContext>
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

#[cfg(feature = "coroutines")]
fn get_default_threads(config: &Config) -> usize {
    config.lookup("input.tls_threads").map_or(DEFAULT_THREADS, |x| x.as_integer().
        expect("input.tls_threads must be an unsigned integer") as usize)
}

#[cfg(not(feature = "coroutines"))]
fn get_default_threads(_config: &Config) -> usize {
    1
}

pub fn config_parse(config: &Config) -> (TlsConfig, String, u64) {
    let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
        expect("input.listen must be an ip:port string")).to_owned();
    let threads = get_default_threads(&config);
    let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("input.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("input.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("input.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("input.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("input.tls_compression must be a boolean"));
    let timeout = config.lookup("input.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("input.timeout must be an integer") as u64);
    let framing = if config.lookup("input.framed").map_or(false, |x| x.as_bool().
        expect("input.framed must be a boolean")) {
        "syslen"
    } else {
        DEFAULT_FRAMING
    };
    let framing = config.lookup("input.framing").map_or(framing, |x| x.as_str().
        expect(r#"input.framing must be a string set to "line", "nul" or "syslen""#)).to_owned();
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            ctx.set_CA_file(&ca_file).expect("Unable to read the trusted CA file");
        }
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).
        expect("Unable to read the TLS certificate");
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).
        expect("Unable to read the TLS key");
    ctx.set_cipher_list(&ciphers).expect("Unsupported cipher suite");
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        framing: framing,
        threads: threads,
        arc_ctx: arc_ctx
    };
    (tls_config, listen, timeout)
}
<commit_msg>adding support for cert chain file in tls inputs<commit_after>use flowgger::config::Config;
use openssl::bn::BigNum;
use openssl::dh::DH;
use openssl::ssl::*;
use openssl::x509::X509FileType;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use openssl::x509::X509;

pub mod tls_input;
#[cfg(feature = "coroutines")]
pub mod tlsco_input;

pub use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES256-GCM-SHA384:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:ECDHE-ECDSA-DES-CBC3-SHA:ECDHE-RSA-DES-CBC3-SHA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SHA";
const DEFAULT_COMPRESSION: bool = false;
const DEFAULT_FRAMING: &'static str = "line";
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
#[cfg(feature = "coroutines")]
const DEFAULT_THREADS: usize = 1;
const DEFAULT_TIMEOUT: u64 = 3600;
const DEFAULT_TLS_METHOD: &'static str = "any";
const DEFAULT_VERIFY_PEER: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
pub struct TlsConfig {
    framing: String,
    threads: usize,
    arc_ctx: Arc<SslContext>
}

#[cfg(feature = "ecdh")]
fn set_ecdh(ctx: &mut SslContext) {
    ctx.set_ecdh_auto(true).unwrap();
}

#[cfg(not(feature = "ecdh"))]
fn set_ecdh(ctx: &mut SslContext) {
    let _ = ctx;
}

fn set_fs(ctx: &mut SslContext) {
    let p = BigNum::from_hex_str("87A8E61DB4B6663CFFBBD19C651959998CEEF608660DD0F25D2CEED4435E3B00E00DF8F1D61957D4FAF7DF4561B2AA3016C3D91134096FAA3BF4296D830E9A7C209E0C6497517ABD5A8A9D306BCF67ED91F9E6725B4758C022E0B1EF4275BF7B6C5BFC11D45F9088B941F54EB1E59BB8BC39A0BF12307F5C4FDB70C581B23F76B63ACAE1CAA6B7902D52526735488A0EF13C6D9A51BFA4AB3AD8347796524D8EF6A167B5A41825D967E144E5140564251CCACB83E6B486F6B3CA3F7971506026C0B857F689962856DED4010ABD0BE621C3A3960A54E710C375F26375D7014103A4B54330C198AF126116D2276E11715F693877FAD7EF09CADB094AE91E1A1597").unwrap();
    let g = BigNum::from_hex_str("3FB32C9B73134D0B2E77506660EDBD484CA7B18F21EF205407F4793A1A0BA12510DBC15077BE463FFF4FED4AAC0BB555BE3A6C1B0C6B47B1BC3773BF7E8C6F62901228F8C28CBB18A55AE31341000A650196F931C77A57F2DDF463E5E9EC144B777DE62AAAB8A8628AC376D282D6ED3864E67982428EBC831D14348F6F2F9193B5045AF2767164E1DFC967C1FB3F2E55A4BD1BFFE83B9C80D052B985D182EA0ADB2A3B7313D3FE14C8484B1E052588B9B7D2BBD2DF016199ECD06E1557CD0915B3353BBB64E0EC377FD028370DF92B52C7891428CDC67EB6184B523D1DB246C32F63078490F00EF8D647D148D47954515E2327CFEF98C582664B4C0F6CC41659").unwrap();
    let q = BigNum::from_hex_str("8CF83642A709A097B447997640129DA299B1A47D1EB3750BA308B0FE64F5FBD3").unwrap();
    let dh = DH::from_params(p, g, q).unwrap();
    ctx.set_tmp_dh(dh).unwrap();
    set_ecdh(ctx);
}

#[cfg(feature = "coroutines")]
fn get_default_threads(config: &Config) -> usize {
    config.lookup("input.tls_threads").map_or(DEFAULT_THREADS, |x| x.as_integer().
        expect("input.tls_threads must be an unsigned integer") as usize)
}

#[cfg(not(feature = "coroutines"))]
fn get_default_threads(_config: &Config) -> usize {
    1
}

pub fn config_parse(config: &Config) -> (TlsConfig, String, u64) {
    let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
        expect("input.listen must be an ip:port string")).to_owned();
    let threads = get_default_threads(&config);
    let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
        expect("input.tls_cert must be a path to a .pem file")).to_owned();
    let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
        expect("input.tls_key must be a path to a .pem file")).to_owned();
    let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
        expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
    let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
        expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
            "any" | "sslv23" => SslMethod::Sslv23,
            "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
            "tlsv1.1" => SslMethod::Tlsv1_1,
            "tlsv1.2" => SslMethod::Tlsv1_2,
            _ => panic!(r#"TLS method must be "any", "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
    };
    let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
        expect("input.tls_verify_peer must be a boolean"));
    let ca_file: Option<PathBuf> = config.lookup("input.tls_ca_file").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_ca_file must be a path to a file"))));
    let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
        expect("input.tls_compression must be a boolean"));
    let extra_chain_path: Option<PathBuf> = config.lookup("input.tls_cert_chain").map_or(None, |x|
        Some(PathBuf::from(x.as_str().expect("input.tls_extra_chain_path must be a path to a file"))));
    let timeout = config.lookup("input.timeout").map_or(DEFAULT_TIMEOUT, |x| x.as_integer().
        expect("input.timeout must be an integer") as u64);
    let framing = if config.lookup("input.framed").map_or(false, |x| x.as_bool().
        expect("input.framed must be a boolean")) {
        "syslen"
    } else {
        DEFAULT_FRAMING
    };
    let framing = config.lookup("input.framing").map_or(framing, |x| x.as_str().
        expect(r#"input.framing must be a string set to "line", "nul" or "syslen""#)).to_owned();
    let mut ctx = SslContext::new(tls_method).unwrap();
    if verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        if let Some(ca_file) = ca_file {
            ctx.set_CA_file(&ca_file).expect("Unable to read the trusted CA file");
        }
    }
    if let Some(extra_chain_path) = extra_chain_path {
         let mut extra_chain_file = File::open(extra_chain_path.as_path())
             .ok()
             .expect("Failed to open the extra chain cert file");
         let extra_chain_cert = X509::from_pem(&mut extra_chain_file).ok().expect("unable to load the provided extra chain cert");
         ctx.add_extra_chain_cert(&extra_chain_cert).expect("Unable to read the certificate chain");
    }

    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    set_fs(&mut ctx);
    ctx.set_certificate_file(&Path::new(&cert), X509FileType::PEM).
        expect("Unable to read the TLS certificate");
    ctx.set_private_key_file(&Path::new(&key), X509FileType::PEM).
        expect("Unable to read the TLS key");
    ctx.set_cipher_list(&ciphers).expect("Unsupported cipher suite");
    let arc_ctx = Arc::new(ctx);
    let tls_config = TlsConfig {
        framing: framing,
        threads: threads,
        arc_ctx: arc_ctx
    };
    (tls_config, listen, timeout)
}
<|endoftext|>


<commit_before>use std::iter::successors;

use ra_syntax::{
    algo::{neighbor, SyntaxRewriter},
    ast::{self, edit::AstNodeEdit, make},
    AstNode, Direction, InsertPosition, SyntaxElement, T,
};

use crate::{Assist, AssistCtx, AssistId};

// Assist: merge_imports
//
// Merges two imports with a common prefix.
//
// ```
// use std::<|>fmt::Formatter;
// use std::io;
// ```
// ->
// ```
// use std::{fmt::Formatter, io};
// ```
pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
    let tree: ast::UseTree = ctx.find_node_at_offset()?;
    let mut rewriter = SyntaxRewriter::default();
    let mut offset = ctx.frange.range.start();

    if let Some(use_item) = tree.syntax().parent().and_then(ast::UseItem::cast) {
        let (merged, to_delete) = next_prev()
            .filter_map(|dir| neighbor(&use_item, dir))
            .filter_map(|it| Some((it.clone(), it.use_tree()?)))
            .find_map(|(use_item, use_tree)| {
                Some((try_merge_trees(&tree, &use_tree)?, use_item.clone()))
            })?;

        rewriter.replace_ast(&tree, &merged);
        rewriter += to_delete.remove();

        if to_delete.syntax().text_range().end() < offset {
            offset -= to_delete.syntax().text_range().len();
        }
    } else {
        let (merged, to_delete) = next_prev()
            .filter_map(|dir| neighbor(&tree, dir))
            .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;

        rewriter.replace_ast(&tree, &merged);
        rewriter += to_delete.remove();

        if to_delete.syntax().text_range().end() < offset {
            offset -= to_delete.syntax().text_range().len();
        }
    };

    ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
        edit.rewrite(rewriter);
        // FIXME: we only need because our diff is imprecise
        edit.set_cursor(offset);
    })
}

fn next_prev() -> impl Iterator<Item = Direction> {
    [Direction::Next, Direction::Prev].iter().copied()
}

fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
    let lhs_path = old.path()?;
    let rhs_path = new.path()?;

    let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;

    let lhs = old.split_prefix(&lhs_prefix);
    let rhs = new.split_prefix(&rhs_prefix);

    let mut to_insert: Vec<SyntaxElement> = Vec::new();
    to_insert.push(make::token(T![,]).into());
    to_insert.push(make::tokens::single_space().into());
    to_insert.extend(
        rhs.use_tree_list()?
            .syntax()
            .children_with_tokens()
            .filter(|it| it.kind() != T!['{'] && it.kind() != T!['}']),
    );
    let use_tree_list = lhs.use_tree_list()?;
    let pos = InsertPosition::Before(use_tree_list.r_curly_token()?.into());
    let use_tree_list = use_tree_list.insert_children(pos, to_insert);
    Some(lhs.with_use_tree_list(use_tree_list))
}

fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
    let mut res = None;
    let mut lhs_curr = first_path(&lhs);
    let mut rhs_curr = first_path(&rhs);
    loop {
        match (lhs_curr.segment(), rhs_curr.segment()) {
            (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
            _ => break,
        }
        res = Some((lhs_curr.clone(), rhs_curr.clone()));

        match (lhs_curr.parent_path(), rhs_curr.parent_path()) {
            (Some(lhs), Some(rhs)) => {
                lhs_curr = lhs;
                rhs_curr = rhs;
            }
            _ => break,
        }
    }

    res
}

fn first_path(path: &ast::Path) -> ast::Path {
    successors(Some(path.clone()), |it| it.qualifier()).last().unwrap()
}

#[cfg(test)]
mod tests {
    use crate::helpers::check_assist;

    use super::*;

    #[test]
    fn test_merge_first() {
        check_assist(
            merge_imports,
            r"
use std::fmt<|>::Debug;
use std::fmt::Display;
",
            r"
use std::fmt<|>::{Debug, Display};
",
        )
    }

    #[test]
    fn test_merge_second() {
        check_assist(
            merge_imports,
            r"
use std::fmt::Debug;
use std::fmt<|>::Display;
",
            r"
use std::fmt:<|>:{Display, Debug};
",
        );
    }

    #[test]
    fn test_merge_nested() {
        check_assist(
            merge_imports,
            r"
use std::{fmt<|>::Debug, fmt::Display};
",
            r"
use std::{fmt<|>::{Debug, Display}};
",
        );
        check_assist(
            merge_imports,
            r"
use std::{fmt::Debug, fmt<|>::Display};
",
            r"
use std::{fmt::<|>{Display, Debug}};
",
        );
    }

    #[test]
    fn test_merge_single_wildcard_diff_prefixes() {
        check_assist(
            merge_imports,
            r"
use std<|>::cell::*;
use std::str;
",
            r"
use std<|>::{cell::*, str};
",
        )
    }

    #[test]
    fn test_merge_both_wildcard_diff_prefixes() {
        check_assist(
            merge_imports,
            r"
use std<|>::cell::*;
use std::str::*;
",
            r"
use std<|>::{cell::*, str::*};
",
        )
    }

    #[test]
    fn removes_just_enough_whitespace() {
        check_assist(
            merge_imports,
            r"
use foo<|>::bar;
use foo::baz;

/// Doc comment
",
            r"
use foo<|>::{bar, baz};

/// Doc comment
",
        );
    }

    #[test]
    fn works_with_trailing_comma() {
        check_assist(
            merge_imports,
            r"
use {
    foo<|>::bar,
    foo::baz,
};
",
            r"
use {
    foo<|>::{bar, baz},
};
",
        );
        check_assist(
            merge_imports,
            r"
use {
    foo::baz,
    foo<|>::bar,
};
",
            r"
use {
    foo::{bar<|>, baz},
};
",
        );
    }
}
<commit_msg>Fix double comma when merge imports on second line<commit_after>use std::iter::successors;

use ra_syntax::{
    algo::{neighbor, SyntaxRewriter},
    ast::{self, edit::AstNodeEdit, make},
    AstNode, Direction, InsertPosition, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxToken, Token,
    T,
};

use crate::{Assist, AssistCtx, AssistId};

// Assist: merge_imports
//
// Merges two imports with a common prefix.
//
// ```
// use std::<|>fmt::Formatter;
// use std::io;
// ```
// ->
// ```
// use std::{fmt::Formatter, io};
// ```
pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
    let tree: ast::UseTree = ctx.find_node_at_offset()?;
    let mut rewriter = SyntaxRewriter::default();
    let mut offset = ctx.frange.range.start();

    if let Some(use_item) = tree.syntax().parent().and_then(ast::UseItem::cast) {
        let (merged, to_delete) = next_prev()
            .filter_map(|dir| neighbor(&use_item, dir))
            .filter_map(|it| Some((it.clone(), it.use_tree()?)))
            .find_map(|(use_item, use_tree)| {
                Some((try_merge_trees(&tree, &use_tree)?, use_item.clone()))
            })?;

        rewriter.replace_ast(&tree, &merged);
        rewriter += to_delete.remove();

        if to_delete.syntax().text_range().end() < offset {
            offset -= to_delete.syntax().text_range().len();
        }
    } else {
        let (merged, to_delete) = next_prev()
            .filter_map(|dir| neighbor(&tree, dir))
            .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;

        rewriter.replace_ast(&tree, &merged);
        rewriter += to_delete.remove();

        if to_delete.syntax().text_range().end() < offset {
            offset -= to_delete.syntax().text_range().len();
        }
    };

    ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
        edit.rewrite(rewriter);
        // FIXME: we only need because our diff is imprecise
        edit.set_cursor(offset);
    })
}

fn next_prev() -> impl Iterator<Item = Direction> {
    [Direction::Next, Direction::Prev].iter().copied()
}

fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
    let lhs_path = old.path()?;
    let rhs_path = new.path()?;

    let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;

    let lhs = old.split_prefix(&lhs_prefix);
    let rhs = new.split_prefix(&rhs_prefix);

    let mut should_insert_comma = true;

    lhs.syntax()
        .last_child()
        .and_then(|child| child.children().last())
        .and_then(|last| last.next_sibling_or_token())
        .map(|next_sibling| {
            // FIXME: there's probably a better way to check if it's a COMMA
            if let NodeOrToken::Token(maybe_comma) = next_sibling {
                should_insert_comma = maybe_comma.to_string() != ",";
            }
        });

    let mut to_insert: Vec<SyntaxElement> = Vec::new();
    if should_insert_comma {
        to_insert.push(make::token(T![,]).into());
        to_insert.push(make::tokens::single_space().into());
    }
    to_insert.extend(
        rhs.use_tree_list()?
            .syntax()
            .children_with_tokens()
            .filter(|it| it.kind() != T!['{'] && it.kind() != T!['}']),
    );
    let use_tree_list = lhs.use_tree_list()?;
    let pos = InsertPosition::Before(use_tree_list.r_curly_token()?.into());
    let use_tree_list = use_tree_list.insert_children(pos, to_insert);
    Some(lhs.with_use_tree_list(use_tree_list))
}

fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
    let mut res = None;
    let mut lhs_curr = first_path(&lhs);
    let mut rhs_curr = first_path(&rhs);
    loop {
        match (lhs_curr.segment(), rhs_curr.segment()) {
            (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
            _ => break,
        }
        res = Some((lhs_curr.clone(), rhs_curr.clone()));

        match (lhs_curr.parent_path(), rhs_curr.parent_path()) {
            (Some(lhs), Some(rhs)) => {
                lhs_curr = lhs;
                rhs_curr = rhs;
            }
            _ => break,
        }
    }

    res
}

fn first_path(path: &ast::Path) -> ast::Path {
    successors(Some(path.clone()), |it| it.qualifier()).last().unwrap()
}

#[cfg(test)]
mod tests {
    use crate::helpers::check_assist;

    use super::*;

    #[test]
    fn test_merge_first() {
        check_assist(
            merge_imports,
            r"
use std::fmt<|>::Debug;
use std::fmt::Display;
",
            r"
use std::fmt<|>::{Debug, Display};
",
        )
    }

    #[test]
    fn test_merge_second() {
        check_assist(
            merge_imports,
            r"
use std::fmt::Debug;
use std::fmt<|>::Display;
",
            r"
use std::fmt:<|>:{Display, Debug};
",
        );
    }

    #[test]
    fn test_merge_nested() {
        check_assist(
            merge_imports,
            r"
use std::{fmt<|>::Debug, fmt::Display};
",
            r"
use std::{fmt<|>::{Debug, Display}};
",
        );
        check_assist(
            merge_imports,
            r"
use std::{fmt::Debug, fmt<|>::Display};
",
            r"
use std::{fmt::<|>{Display, Debug}};
",
        );
    }

    #[test]
    fn test_merge_single_wildcard_diff_prefixes() {
        check_assist(
            merge_imports,
            r"
use std<|>::cell::*;
use std::str;
",
            r"
use std<|>::{cell::*, str};
",
        )
    }

    #[test]
    fn test_merge_both_wildcard_diff_prefixes() {
        check_assist(
            merge_imports,
            r"
use std<|>::cell::*;
use std::str::*;
",
            r"
use std<|>::{cell::*, str::*};
",
        )
    }

    #[test]
    fn removes_just_enough_whitespace() {
        check_assist(
            merge_imports,
            r"
use foo<|>::bar;
use foo::baz;

/// Doc comment
",
            r"
use foo<|>::{bar, baz};

/// Doc comment
",
        );
    }

    #[test]
    fn works_with_trailing_comma() {
        check_assist(
            merge_imports,
            r"
use {
    foo<|>::bar,
    foo::baz,
};
",
            r"
use {
    foo<|>::{bar, baz},
};
",
        );
        check_assist(
            merge_imports,
            r"
use {
    foo::baz,
    foo<|>::bar,
};
",
            r"
use {
    foo::{bar<|>, baz},
};
",
        );
    }

    #[test]
    fn test_double_comma() {
        check_assist(
            merge_imports,
            r"
use hyper::service::make_service_fn;
use hyper::<|>{
    StatusCode,
};
",
            r"
use hyper::{<|>
    StatusCode,
service::make_service_fn};
",
        )
    }
}
<|endoftext|>


<commit_before>use af;
use af::{Array, Dim4, MatProp};
use std::sync::{Arc, Mutex};

use utils;
use activations;
use params::Params;
use layer::Layer;

pub struct RNN {
  pub input_size: usize,
  pub output_size: usize,
}

impl Layer for RNN
{
  fn forward(&self, params: Arc<Mutex<Params>>, inputs: &Array, state: Option<Array>) -> (Array, Option<Array>)
  {
    // get a handle to the underlying params
    let mut ltex = params.lock().unwrap();
    let current_unroll = ltex.current_unroll;

    // z_t = xW [bias added in h_t]
    let wx = af::matmul(&inputs            //activated_input
                        , &ltex.weights[0]
                        , MatProp::NONE
                        , MatProp::NONE);

    // uh_tm1 = htm1 U [bias added in h_t]
    // handle case where time = 0
    let uhtm1 = match ltex.recurrences.len(){
      0 => {
        let output_size = ltex.weights[1].dims()[0]; // is [M x M]
        let init_h_dims = Dim4::new(&[inputs.dims()[0], output_size, 1, 1]);
        if let Some(init_state) = state {
          ltex.recurrences.push(init_state);
        }else {
          ltex.recurrences.push(utils::constant(init_h_dims, inputs.get_type(), 0f32));
        }
        af::matmul(&ltex.recurrences.last().unwrap()
                   , &ltex.weights[1]
                   , MatProp::NONE
                   , MatProp::NONE)
      },
      _ => {
        af::matmul(&ltex.recurrences[ltex.current_unroll]
                   , &ltex.weights[1]
                   , MatProp::NONE
                   , MatProp::NONE)
      }
    };

    // h_t = uh_tm1 + wx + b
    let h_t = af::transpose(&af::add(&af::transpose(&uhtm1, false),
                                     &af::add(&af::transpose(&wx, false)
                                              , &ltex.biases[0], true), false)
                            , false);

    // a_t = sigma(h_t)
    let a_t = activations::get_activation(&ltex.activations[0], &h_t).unwrap();

    // parameter manager keeps the output & inputs
    if ltex.inputs.len() > current_unroll { // store in existing
      ltex.inputs[current_unroll] = inputs.clone();
      ltex.outputs[current_unroll] = a_t.clone();
      ltex.recurrences[current_unroll] = h_t.clone();
    }else{                                  // add new
      ltex.inputs.push(inputs.clone());
      ltex.outputs.push(a_t.clone());
      ltex.recurrences.push(h_t.clone());
    }

    // update location in vector
    ltex.current_unroll += 1;

    (a_t.clone(), Some(h_t.clone())) // clone just increases the ref count
  }

  // x, prev_h, Wx, Wh, next_h = cache
  // dout_dnext_h = dnext_h * (1 - np.power(next_h, 2))
  // dx = dout_dnext_h.dot(Wx.T)
  // dprev_h = dout_dnext_h.(Wh.T)
  // dWx = x.T.dot(dout_dnext_h)
  // dWh = prev_h.T.dot(dout_dnext_h)
  // db = dout_dnext_h.sum(axis=0)

  fn backward(&self, params: Arc<Mutex<Params>>, delta: &Array) -> Array
  {
    // get a handle to the underlying params
    let mut ltex = params.lock().unwrap();
    let current_unroll = ltex.current_unroll;
    assert!(current_unroll > 0
            , "Cannot call backward pass without at least 1 forward pass");

    // dz          = grad(a_t)
    // delta_t     = (delta_{t+1} + dh{t+1}) .* dz
    // dh          = delta_{t} * U
    // dW          = x_t^T * delta_t
    // dU          = h_t^T * delta_t
    // db          = sum_{batch} (delta_t)
    // delta_{t-1} = (transpose(W_{t}) * d_{l})
    let dz = activations::get_derivative(&ltex.activations[0]
                                         , &ltex.outputs[current_unroll - 1]).unwrap();

    // check to see if we already have a state derivative, else add one
    if ltex.state_derivatives.len() == 0 {
      let h_size = ltex.recurrences[current_unroll - 1].dims();
      let h_type = ltex.recurrences[current_unroll - 1].get_type();
      ltex.state_derivatives.push(utils::constant(h_size, h_type, 0.0f32));
    }

    // update the delta for the current time-step by combining with state derivative
    let delta_t = af::mul(&af::add(&ltex.state_derivatives[0]
                                   , delta, false), &dz, false);

    //println!("dh dims = {:?} | delta dims = {:?} | delta prop dims = {:?}", dh.dims(), delta_t.dims(), delta.dims());
    let dw = af::matmul(&ltex.inputs[current_unroll - 1], &delta_t       // delta_w = delta_t * a_{t}
                        , af::MatProp::TRANS
                        , af::MatProp::NONE);
    let du = af::matmul(&ltex.recurrences[current_unroll - 1], &delta_t  // delta_u = delta_t * h_{t-1}
                        , af::MatProp::TRANS
                        , af::MatProp::NONE);
    let db = af::transpose(&af::sum(&delta_t, 0), false);                // delta_b = \sum_{batch} delta
    ltex.deltas[0] = af::add(&ltex.deltas[0], &dw, false);
    ltex.deltas[1] = af::add(&ltex.deltas[1], &du, false);
    ltex.deltas[2] = af::add(&ltex.deltas[2], &db, false);

    // add the current state derivative in
    ltex.state_derivatives[0] = af::matmul(&delta_t, &ltex.weights[1], af::MatProp::NONE, af::MatProp::TRANS);

    // update location in vector
    ltex.current_unroll -= 1;

    //println!("dims = {:?}", af::matmul(&delta_t, &ltex.weights[0], af::MatProp::NONE, af::MatProp::TRANS).dims());
    af::matmul(&delta_t, &ltex.weights[0], af::MatProp::NONE, af::MatProp::TRANS)    // delta_{t-1}
  }
}
<commit_msg>cleanup comments<commit_after>use af;
use af::{Array, Dim4, MatProp};
use std::sync::{Arc, Mutex};

use utils;
use activations;
use params::Params;
use layer::Layer;

pub struct RNN {
  pub input_size: usize,
  pub output_size: usize,
}

impl Layer for RNN
{
  fn forward(&self, params: Arc<Mutex<Params>>, inputs: &Array, state: Option<Array>) -> (Array, Option<Array>)
  {
    // get a handle to the underlying params
    let mut ltex = params.lock().unwrap();
    let current_unroll = ltex.current_unroll;

    // z_t = xW [bias added in h_t]
    let wx = af::matmul(&inputs            //activated_input
                        , &ltex.weights[0]
                        , MatProp::NONE
                        , MatProp::NONE);

    // uh_tm1 = htm1 U [bias added in h_t]
    // handle case where time = 0
    let uhtm1 = match ltex.recurrences.len(){
      0 => {
        let output_size = ltex.weights[1].dims()[0]; // is [M x M]
        let init_h_dims = Dim4::new(&[inputs.dims()[0], output_size, 1, 1]);
        if let Some(init_state) = state {
          ltex.recurrences.push(init_state);
        }else {
          ltex.recurrences.push(utils::constant(init_h_dims, inputs.get_type(), 0f32));
        }
        af::matmul(&ltex.recurrences.last().unwrap()
                   , &ltex.weights[1]
                   , MatProp::NONE
                   , MatProp::NONE)
      },
      _ => {
        af::matmul(&ltex.recurrences[ltex.current_unroll]
                   , &ltex.weights[1]
                   , MatProp::NONE
                   , MatProp::NONE)
      }
    };

    // h_t = uh_tm1 + wx + b
    let h_t = af::transpose(&af::add(&af::transpose(&uhtm1, false),
                                     &af::add(&af::transpose(&wx, false)
                                              , &ltex.biases[0], true), false)
                            , false);

    // a_t = sigma(h_t)
    let a_t = activations::get_activation(&ltex.activations[0], &h_t).unwrap();

    // parameter manager keeps the output & inputs
    if ltex.inputs.len() > current_unroll { // store in existing
      ltex.inputs[current_unroll] = inputs.clone();
      ltex.outputs[current_unroll] = a_t.clone();
      ltex.recurrences[current_unroll] = h_t.clone();
    }else{                                  // add new
      ltex.inputs.push(inputs.clone());
      ltex.outputs.push(a_t.clone());
      ltex.recurrences.push(h_t.clone());
    }

    // update location in vector
    ltex.current_unroll += 1;

    (a_t.clone(), Some(h_t.clone())) // clone just increases the ref count
  }

  fn backward(&self, params: Arc<Mutex<Params>>, delta: &Array) -> Array
  {
    // get a handle to the underlying params
    let mut ltex = params.lock().unwrap();
    let current_unroll = ltex.current_unroll;
    assert!(current_unroll > 0
            , "Cannot call backward pass without at least 1 forward pass");

    // dz          = grad(a_t)
    // delta_t     = (delta_{t+1} + dh{t+1}) .* dz
    // dh          = delta_{t} * U
    // dW          = x_t^T * delta_t
    // dU          = h_t^T * delta_t
    // db          = sum_{batch} (delta_t)
    // delta_{t-1} = (transpose(W_{t}) * d_{l})
    let dz = activations::get_derivative(&ltex.activations[0]
                                         , &ltex.outputs[current_unroll - 1]).unwrap();

    // check to see if we already have a state derivative, else add one
    if ltex.state_derivatives.len() == 0 {
      let h_size = ltex.recurrences[current_unroll - 1].dims();
      let h_type = ltex.recurrences[current_unroll - 1].get_type();
      ltex.state_derivatives.push(utils::constant(h_size, h_type, 0.0f32));
    }

    // update the delta for the current time-step by combining with state derivative
    let delta_t = af::mul(&af::add(&ltex.state_derivatives[0]
                                   , delta, false), &dz, false);

    //println!("dh dims = {:?} | delta dims = {:?} | delta prop dims = {:?}", dh.dims(), delta_t.dims(), delta.dims());
    let dw = af::matmul(&ltex.inputs[current_unroll - 1], &delta_t       // delta_w = delta_t * a_{t}
                        , af::MatProp::TRANS
                        , af::MatProp::NONE);
    let du = af::matmul(&ltex.recurrences[current_unroll - 1], &delta_t  // delta_u = delta_t * h_{t-1}
                        , af::MatProp::TRANS
                        , af::MatProp::NONE);
    let db = af::transpose(&af::sum(&delta_t, 0), false);                // delta_b = \sum_{batch} delta
    ltex.deltas[0] = af::add(&ltex.deltas[0], &dw, false);
    ltex.deltas[1] = af::add(&ltex.deltas[1], &du, false);
    ltex.deltas[2] = af::add(&ltex.deltas[2], &db, false);

    // add the current state derivative in
    ltex.state_derivatives[0] = af::matmul(&delta_t, &ltex.weights[1], af::MatProp::NONE, af::MatProp::TRANS);

    // update location in vector
    ltex.current_unroll -= 1;

    //println!("dims = {:?}", af::matmul(&delta_t, &ltex.weights[0], af::MatProp::NONE, af::MatProp::TRANS).dims());
    af::matmul(&delta_t, &ltex.weights[0], af::MatProp::NONE, af::MatProp::TRANS)    // delta_{t-1}
  }
}
<|endoftext|>


<commit_before>#![allow(bad_style)]

extern crate libc;
use libc::*;

pub type rs_magic_number = c_int;
pub const RS_DELTA_MAGIC: c_int = 0x72730236;
pub const RS_MD4_SIG_MAGIC: c_int = 0x72730136;
pub const RS_BLAKE2_SIG_MAGIC: c_int = 0x72730137;

pub type rs_result = c_int;
pub const RS_DONE: c_int = 0;
pub const RS_BLOCKED: c_int = 1;
pub const RS_RUNNING: c_int = 2;
pub const RS_TEST_SKIPPED: c_int = 77;
pub const RS_IO_ERROR: c_int = 100;
pub const RS_SYNTAX_ERROR: c_int = 101;
pub const RS_MEM_ERROR: c_int = 102;
pub const RS_INPUT_ENDED: c_int = 103;
pub const RS_BAD_MAGIC: c_int = 104;
pub const RS_UNIMPLEMENTED: c_int = 105;
pub const RS_CORRUPT: c_int = 106;
pub const RS_INTERNAL_ERROR: c_int = 107;
pub const RS_PARAM_ERROR: c_int = 108;

pub const RS_DEFAULT_BLOCK_LEN: size_t = 2048;

pub type rs_long_t = c_longlong;

pub enum rs_job_t {}
pub enum rs_signature_t {}

#[repr(C)]
pub struct rs_buffers_t {
    pub next_in: *const c_char,
    pub avail_in: size_t,
    pub eof_in: c_int,
    pub next_out: *mut c_char,
    pub avail_out: size_t,
}

pub type rs_driven_cb = extern "C" fn(job: *mut rs_job_t,
                                      buf: *mut rs_buffers_t,
                                      opaque: *mut c_void)
                                      -> rs_result;
pub type rs_copy_cb = extern "C" fn(opaque: *mut c_void,
                                    pos: rs_long_t,
                                    len: *mut size_t,
                                    buf: *mut *mut c_void)
                                    -> rs_result;

extern "C" {
    pub fn rs_job_iter(job: *mut rs_job_t, buffers: *mut rs_buffers_t) -> rs_result;
    pub fn fs_job_drive(job: *mut rs_job_t,
                        buf: *mut rs_buffers_t,
                        in_cb: rs_driven_cb,
                        in_opaque: *mut c_void,
                        out_cb: rs_driven_cb,
                        out_opaque: *mut c_void)
                        -> rs_result;
    pub fn rs_job_free(job: *mut rs_job_t) -> rs_result;

    pub fn rs_sig_begin(new_block_len: size_t,
                        strong_sum_len: size_t,
                        sig_magic: rs_magic_number)
                        -> *mut rs_job_t;
    pub fn rs_delta_begin(sig: *mut rs_signature_t) -> *mut rs_job_t;
    pub fn rs_loadsig_begin(sig: *mut *mut rs_signature_t) -> *mut rs_job_t;
    pub fn rs_build_hash_table(sums: *mut rs_signature_t) -> rs_result;
    pub fn rs_free_sumset(sums: *mut rs_signature_t);
    pub fn rs_patch_begin(copy_cb: rs_copy_cb, copy_arg: *mut c_void) -> *mut rs_job_t;
}
<commit_msg>Added logging apis to librsync-sys<commit_after>#![allow(bad_style)]

extern crate libc;
use libc::*;

pub type rs_magic_number = c_int;
pub const RS_DELTA_MAGIC: c_int = 0x72730236;
pub const RS_MD4_SIG_MAGIC: c_int = 0x72730136;
pub const RS_BLAKE2_SIG_MAGIC: c_int = 0x72730137;

pub type rs_result = c_int;
pub const RS_DONE: c_int = 0;
pub const RS_BLOCKED: c_int = 1;
pub const RS_RUNNING: c_int = 2;
pub const RS_TEST_SKIPPED: c_int = 77;
pub const RS_IO_ERROR: c_int = 100;
pub const RS_SYNTAX_ERROR: c_int = 101;
pub const RS_MEM_ERROR: c_int = 102;
pub const RS_INPUT_ENDED: c_int = 103;
pub const RS_BAD_MAGIC: c_int = 104;
pub const RS_UNIMPLEMENTED: c_int = 105;
pub const RS_CORRUPT: c_int = 106;
pub const RS_INTERNAL_ERROR: c_int = 107;
pub const RS_PARAM_ERROR: c_int = 108;

pub type rs_loglevel = c_int;
pub const RS_LOG_EMERG: c_int = 0;
pub const RS_LOG_ALERT: c_int = 1;
pub const RS_LOG_CRIT: c_int = 2;
pub const RS_LOG_ERR: c_int = 3;
pub const RS_LOG_WARNING: c_int = 4;
pub const RS_LOG_NOTICE: c_int = 5;
pub const RS_LOG_INFO: c_int = 6;
pub const RS_LOG_DEBUG: c_int = 7;

pub const RS_DEFAULT_BLOCK_LEN: size_t = 2048;

pub type rs_long_t = c_longlong;

pub enum rs_job_t {}
pub enum rs_signature_t {}

#[repr(C)]
pub struct rs_buffers_t {
    pub next_in: *const c_char,
    pub avail_in: size_t,
    pub eof_in: c_int,
    pub next_out: *mut c_char,
    pub avail_out: size_t,
}

pub type rs_driven_cb = extern "C" fn(job: *mut rs_job_t,
                                      buf: *mut rs_buffers_t,
                                      opaque: *mut c_void)
                                      -> rs_result;
pub type rs_copy_cb = extern "C" fn(opaque: *mut c_void,
                                    pos: rs_long_t,
                                    len: *mut size_t,
                                    buf: *mut *mut c_void)
                                    -> rs_result;
pub type rs_trace_fn_t = extern "C" fn(level: rs_loglevel, *const c_char);

extern "C" {
    pub fn rs_job_iter(job: *mut rs_job_t, buffers: *mut rs_buffers_t) -> rs_result;
    pub fn fs_job_drive(job: *mut rs_job_t,
                        buf: *mut rs_buffers_t,
                        in_cb: rs_driven_cb,
                        in_opaque: *mut c_void,
                        out_cb: rs_driven_cb,
                        out_opaque: *mut c_void)
                        -> rs_result;
    pub fn rs_job_free(job: *mut rs_job_t) -> rs_result;

    pub fn rs_sig_begin(new_block_len: size_t,
                        strong_sum_len: size_t,
                        sig_magic: rs_magic_number)
                        -> *mut rs_job_t;
    pub fn rs_delta_begin(sig: *mut rs_signature_t) -> *mut rs_job_t;
    pub fn rs_loadsig_begin(sig: *mut *mut rs_signature_t) -> *mut rs_job_t;
    pub fn rs_build_hash_table(sums: *mut rs_signature_t) -> rs_result;
    pub fn rs_free_sumset(sums: *mut rs_signature_t);
    pub fn rs_patch_begin(copy_cb: rs_copy_cb, copy_arg: *mut c_void) -> *mut rs_job_t;

    pub fn rs_trace_set_level(level: rs_loglevel);
    pub fn rs_trace_to(f: rs_trace_fn_t);
}
<|endoftext|>