<commit_before>// Copyright (c) 2016 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! A Rust library that provides an interface for the Pimoroni Blinkt!, and any
//! similar APA102 strips or boards, on a Raspberry Pi.
//!
//! Blinkt accesses the BCM2708/BCM2709 GPIO peripheral either through
//! `/dev/gpiomem` (preferred) or `/dev/mem`. Both the original APA102, and the
//! alternate version with a smaller, darker die, are supported. The APA102 RGB
//! LED/driver ICs are referred to as pixels throughout the code and documentation.
//!
//! Each pixel has a red, green and blue LED with possible values between 0-255.
//! Additionally, the overall brightness of each pixel can be set to 0.0-1.0, which
//! is converted to a 5-bit value.
//!
//! Blinkt stores all color and brightness changes in a local buffer. Use
//! `show()` to send the buffered values to the pixels.
//!
//! # Examples
//!
//! A complete example that cycles all pixels through red, green and blue.
//!
//! ```rust,no_run
//! extern crate blinkt;
//!
//! use std::{thread, mem};
//! use std::time::Duration;
//!
//! use blinkt::Blinkt;
//!
//! fn main() {
//!     let mut blinkt = Blinkt::new().unwrap();
//!     let (red, green, blue) = (&mut 255, &mut 0, &mut 0);
//!
//!     loop {
//!         blinkt.set_all_pixels(*red, *green, *blue);
//!         blinkt.show();
//!
//!         thread::sleep(Duration::from_millis(250));
//!
//!         mem::swap(red, green);
//!         mem::swap(red, blue);
//!     }
//! }
//! ```
//!
//! By default, all pixels are cleared when Blinkt goes out of
//! scope. Use `set_clear_on_drop(false)` to disable this behavior. Note that
//! drop methods aren't called when a program is abnormally terminated (for
//! instance when a SIGINT isn't caught).
//!
//! ```rust,no_run
//! use blinkt::Blinkt;
//!
//! let mut blinkt = Blinkt::new().unwrap();
//! blinkt.set_clear_on_drop(false);
//!
//! for n in 0..8 {
//!     blinkt.set_pixel(n, 36 * n as u8, 0, 255 - (36 * n as u8));
//! }
//!
//! blinkt.show();
//! ```
//!

#![recursion_limit="128"]   // Needed for the quick_error! macro

extern crate libc;
extern crate num;
#[macro_use]
extern crate enum_primitive;
#[macro_use]
extern crate quick_error;

mod gpio;
mod system;

use std::result;

use gpio::{GPIO, Level, Mode};

pub use gpio::Error as GPIOError;

// Default values for the Pimoroni Blinkt! board
const DAT: u8 = 23;
const CLK: u8 = 24;
const NUM_PIXELS: usize = 8;

const DEFAULT_BRIGHTNESS: u8 = 7;

quick_error! {
    #[derive(Debug)]
/// Errors that can occur when creating a new Blinkt.
    pub enum Error {
/// Accessing the GPIO peripheral returned an error.
///
/// Some of these errors can be fixed by changing file permissions, or upgrading
/// to a newer version of Raspbian.
        GPIO(err: GPIOError) { description(err.description()) from() }
    }
}

/// Result type returned from methods that can have `blinkt::Error`s.
pub type Result<T> = result::Result<T, Error>;

#[derive(Debug, Copy, Clone)]
struct Pixel {
    red: u8,
    green: u8,
    blue: u8,
    brightness: u8,
}

impl Default for Pixel {
    fn default() -> Pixel {
        Pixel {
            red: 0,
            green: 0,
            blue: 0,
            brightness: DEFAULT_BRIGHTNESS,
        }
    }
}

/// Interface for a Blinkt! or any similar APA102 strips/boards.
///
/// By default, Blinkt is set up to communicate with an 8-pixel board through
/// data pin 23 and clock pin 24. These settings can be changed to support
/// alternate configurations.
pub struct Blinkt {
    gpio: GPIO,
    pixels: Vec<Pixel>,
    clear_on_drop: bool,
    pin_data: u8,
    pin_clock: u8,
    endframe_pulses: usize,
}

impl Blinkt {
    /// Creates a new `Blinkt` using the default settings for a Pimoroni
    /// Blinkt! board.
    ///
    /// This sets the data pin to 23, the clock pin to 24, and number of pixels
    /// to 8.
    pub fn new() -> Result<Blinkt> {
        Blinkt::with_settings(DAT, CLK, NUM_PIXELS)
    }

    /// Creates a new `Blinkt` using custom settings for the data pin, clock
    /// pin, and number of pixels.
    pub fn with_settings(pin_data: u8, pin_clock: u8, num_pixels: usize) -> Result<Blinkt> {
        // GPIO init might fail with an error the user could solve
        let mut gpio = try!(GPIO::new());

        gpio.set_mode(pin_data, Mode::Output);
        gpio.write(pin_data, Level::Low);
        gpio.set_mode(pin_clock, Mode::Output);
        gpio.write(pin_clock, Level::Low);

        let blinkt = Blinkt {
            gpio: gpio,
            pixels: vec![Pixel::default(); num_pixels],
            clear_on_drop: true,
            pin_data: pin_data,
            pin_clock: pin_clock,
            endframe_pulses: ((num_pixels as f32 * 0.5) + 0.5) as usize,
        };

        // Send a start frame (32x0) here, because show() doesn't
        // start with one anymore. See show() for details.
        blinkt.write_byte(0);
        blinkt.write_byte(0);
        blinkt.write_byte(0);
        blinkt.write_byte(0);

        Ok(blinkt)
    }

    /// When enabled, clears all pixels when the `Blinkt` goes out of scope.
    ///
    /// Drop methods aren't called when a program is abnormally terminated,
    /// for instance when a user presses Ctrl-C, and the SIGINT signal isn't
    /// caught. You'll either have to catch those using crates such as
    /// `simple_signal`, or manually call `cleanup()`.
    ///
    /// Enabled by default.
    pub fn set_clear_on_drop(&mut self, clear_on_drop: bool) {
        self.clear_on_drop = clear_on_drop;
    }

    /// Changes the GPIO pin mode for the data and clock pins back to their
    /// original state, and optionally clears all pixels.
    ///
    /// Normally, this method is automatically called when Blinkt goes out of
    /// scope, but you can manually call it to handle early/abnormal termination.
    /// After calling this method, any future calls to `show()` won't have any
    /// result.
    pub fn cleanup(&mut self) {
        if self.clear_on_drop {
            self.clear();
            self.show();
        }

        self.gpio.cleanup();
    }

    /// Sets the red, green and blue values for a single pixel in the local
    /// buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid values
    /// for red, green and blue are 0-255.
    pub fn set_pixel(&mut self, pixel: usize, red: u8, green: u8, blue: u8) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
        }
    }

    /// Sets the red, green, blue and brightness values for a single pixel in
    /// the local buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid
    /// values for red, green and blue are 0-255. Valid values for brightness
    /// are 0.0-1.0, which is converted to a 5-bit value.
    pub fn set_pixel_rgbb(&mut self, pixel: usize, red: u8, green: u8, blue: u8, brightness: f32) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
            pixel.brightness = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        }
    }

    /// Sets the brightness value for a single pixel in the local buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid
    /// values for brightness are 0.0-1.0, which is converted to a
    /// 5-bit value.
    pub fn set_pixel_brightness(&mut self, pixel: usize, brightness: f32) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.brightness = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        }
    }

    /// Sets the red, green and blue values for all pixels in the local buffer.
    ///
    /// Valid values for red, green and blue are 0-255.
    pub fn set_all_pixels(&mut self, red: u8, green: u8, blue: u8) {
        for pixel in &mut self.pixels {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
        }
    }

    /// Sets the red, green, blue and brightness values for all pixels in the
    /// local buffer.
    ///
    /// Valid values for red, green and blue are 0-255. Valid values for
    /// brightness are 0.0-1.0, which is converted to a 5-bit value.
    pub fn set_all_pixels_rgbb(&mut self, red: u8, green: u8, blue: u8, brightness: f32) {
        let brightness: u8 = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        for pixel in &mut self.pixels {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
            pixel.brightness = brightness;
        }
    }

    /// Sets the brightness value for all pixels in the local buffer.
    ///
    /// Valid values for brightness are 0.0-1.0, which is converted to a 5-bit
    /// value.
    pub fn set_all_pixels_brightness(&mut self, brightness: f32) {
        let brightness: u8 = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        for pixel in &mut self.pixels {
            pixel.brightness = brightness;
        }
    }

    /// Sets the red, green and blue values to 0 for all pixels in the local
    /// buffer.
    pub fn clear(&mut self) {
        self.set_all_pixels(0, 0, 0);
    }

    /// Sends the contents of the local buffer to the pixels, updating their
    /// LED colors and brightness.
    pub fn show(&self) {
        // We don't start with a start frame, because we already sent one during setup, and
        // our end frame incorporates the next start frame.

        // LED frames
        for pixel in &self.pixels {
            self.write_byte(0b11100000 | pixel.brightness); // 3-bit header + 5-bit brightness
            self.write_byte(pixel.blue);
            self.write_byte(pixel.green);
            self.write_byte(pixel.red);
        }

        // End frame + start frame. One version of the APA102 (smaller, darker
        // die, possibly counterfeit) won't change its LED colors until it
        // receives a start frame after the end frame. This workaround is
        // compatible with both the normal and the 'different' APA102.
        self.gpio.write(self.pin_data, Level::Low);
        for _ in 0..32 + self.endframe_pulses {
            self.gpio.write(self.pin_clock, Level::High);
            self.gpio.write(self.pin_clock, Level::Low);
        }
    }

    fn write_byte(&self, byte: u8) {
        for n in 0..8 {
            if (byte & (1 << (7 - n))) > 0 {
                self.gpio.write(self.pin_data, Level::High);
            } else {
                self.gpio.write(self.pin_data, Level::Low);
            }

            self.gpio.write(self.pin_clock, Level::High);
            self.gpio.write(self.pin_clock, Level::Low);
        }
    }
}

impl Drop for Blinkt {
    fn drop(&mut self) {
        self.cleanup();
    }
}

#[test]
fn test_new() {
    let mut blinkt = match Blinkt::new() {
        // GPIO errors are acceptable, since they're likely caused by outside
        // distro/filesystem issues.
        Err(Error::GPIO(_)) => return,
        Ok(blinkt) => blinkt,
    };

    blinkt.set_clear_on_drop(false);
}
<commit_msg>Add back start frame to show()<commit_after>// Copyright (c) 2016 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! A Rust library that provides an interface for the Pimoroni Blinkt!, and any
//! similar APA102 strips or boards, on a Raspberry Pi.
//!
//! Blinkt accesses the BCM2708/BCM2709 GPIO peripheral either through
//! `/dev/gpiomem` (preferred) or `/dev/mem`. Both the original APA102, and the
//! alternate version with a smaller, darker die, are supported. The APA102 RGB
//! LED/driver ICs are referred to as pixels throughout the code and documentation.
//!
//! Each pixel has a red, green and blue LED with possible values between 0-255.
//! Additionally, the overall brightness of each pixel can be set to 0.0-1.0, which
//! is converted to a 5-bit value.
//!
//! Blinkt stores all color and brightness changes in a local buffer. Use
//! `show()` to send the buffered values to the pixels.
//!
//! # Examples
//!
//! A complete example that cycles all pixels through red, green and blue.
//!
//! ```rust,no_run
//! extern crate blinkt;
//!
//! use std::{thread, mem};
//! use std::time::Duration;
//!
//! use blinkt::Blinkt;
//!
//! fn main() {
//!     let mut blinkt = Blinkt::new().unwrap();
//!     let (red, green, blue) = (&mut 255, &mut 0, &mut 0);
//!
//!     loop {
//!         blinkt.set_all_pixels(*red, *green, *blue);
//!         blinkt.show();
//!
//!         thread::sleep(Duration::from_millis(250));
//!
//!         mem::swap(red, green);
//!         mem::swap(red, blue);
//!     }
//! }
//! ```
//!
//! By default, all pixels are cleared when Blinkt goes out of
//! scope. Use `set_clear_on_drop(false)` to disable this behavior. Note that
//! drop methods aren't called when a program is abnormally terminated (for
//! instance when a SIGINT isn't caught).
//!
//! ```rust,no_run
//! use blinkt::Blinkt;
//!
//! let mut blinkt = Blinkt::new().unwrap();
//! blinkt.set_clear_on_drop(false);
//!
//! for n in 0..8 {
//!     blinkt.set_pixel(n, 36 * n as u8, 0, 255 - (36 * n as u8));
//! }
//!
//! blinkt.show();
//! ```
//!

#![recursion_limit="128"]   // Needed for the quick_error! macro

extern crate libc;
extern crate num;
#[macro_use]
extern crate enum_primitive;
#[macro_use]
extern crate quick_error;

mod gpio;
mod system;

use std::result;

use gpio::{GPIO, Level, Mode};

pub use gpio::Error as GPIOError;

// Default values for the Pimoroni Blinkt! board
const DAT: u8 = 23;
const CLK: u8 = 24;
const NUM_PIXELS: usize = 8;

const DEFAULT_BRIGHTNESS: u8 = 7;

quick_error! {
    #[derive(Debug)]
/// Errors that can occur when creating a new Blinkt.
    pub enum Error {
/// Accessing the GPIO peripheral returned an error.
///
/// Some of these errors can be fixed by changing file permissions, or upgrading
/// to a newer version of Raspbian.
        GPIO(err: GPIOError) { description(err.description()) from() }
    }
}

/// Result type returned from methods that can have `blinkt::Error`s.
pub type Result<T> = result::Result<T, Error>;

#[derive(Debug, Copy, Clone)]
struct Pixel {
    red: u8,
    green: u8,
    blue: u8,
    brightness: u8,
}

impl Default for Pixel {
    fn default() -> Pixel {
        Pixel {
            red: 0,
            green: 0,
            blue: 0,
            brightness: DEFAULT_BRIGHTNESS,
        }
    }
}

/// Interface for a Blinkt! or any similar APA102 strips/boards.
///
/// By default, Blinkt is set up to communicate with an 8-pixel board through
/// data pin 23 and clock pin 24. These settings can be changed to support
/// alternate configurations.
pub struct Blinkt {
    gpio: GPIO,
    pixels: Vec<Pixel>,
    clear_on_drop: bool,
    pin_data: u8,
    pin_clock: u8,
    endframe_pulses: usize,
}

impl Blinkt {
    /// Creates a new `Blinkt` using the default settings for a Pimoroni
    /// Blinkt! board.
    ///
    /// This sets the data pin to 23, the clock pin to 24, and number of pixels
    /// to 8.
    pub fn new() -> Result<Blinkt> {
        Blinkt::with_settings(DAT, CLK, NUM_PIXELS)
    }

    /// Creates a new `Blinkt` using custom settings for the data pin, clock
    /// pin, and number of pixels.
    pub fn with_settings(pin_data: u8, pin_clock: u8, num_pixels: usize) -> Result<Blinkt> {
        // GPIO init might fail with an error the user could solve
        let mut gpio = try!(GPIO::new());

        gpio.set_mode(pin_data, Mode::Output);
        gpio.write(pin_data, Level::Low);
        gpio.set_mode(pin_clock, Mode::Output);
        gpio.write(pin_clock, Level::Low);

        let blinkt = Blinkt {
            gpio: gpio,
            pixels: vec![Pixel::default(); num_pixels],
            clear_on_drop: true,
            pin_data: pin_data,
            pin_clock: pin_clock,
            endframe_pulses: ((num_pixels as f32 * 0.5) + 0.5) as usize,
        };

        Ok(blinkt)
    }

    /// When enabled, clears all pixels when the `Blinkt` goes out of scope.
    ///
    /// Drop methods aren't called when a program is abnormally terminated,
    /// for instance when a user presses Ctrl-C, and the SIGINT signal isn't
    /// caught. You'll either have to catch those using crates such as
    /// `simple_signal`, or manually call `cleanup()`.
    ///
    /// Enabled by default.
    pub fn set_clear_on_drop(&mut self, clear_on_drop: bool) {
        self.clear_on_drop = clear_on_drop;
    }

    /// Changes the GPIO pin mode for the data and clock pins back to their
    /// original state, and optionally clears all pixels.
    ///
    /// Normally, this method is automatically called when Blinkt goes out of
    /// scope, but you can manually call it to handle early/abnormal termination.
    /// After calling this method, any future calls to `show()` won't have any
    /// result.
    pub fn cleanup(&mut self) {
        if self.clear_on_drop {
            self.clear();
            self.show();
        }

        self.gpio.cleanup();
    }

    /// Sets the red, green and blue values for a single pixel in the local
    /// buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid values
    /// for red, green and blue are 0-255.
    pub fn set_pixel(&mut self, pixel: usize, red: u8, green: u8, blue: u8) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
        }
    }

    /// Sets the red, green, blue and brightness values for a single pixel in
    /// the local buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid
    /// values for red, green and blue are 0-255. Valid values for brightness
    /// are 0.0-1.0, which is converted to a 5-bit value.
    pub fn set_pixel_rgbb(&mut self, pixel: usize, red: u8, green: u8, blue: u8, brightness: f32) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
            pixel.brightness = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        }
    }

    /// Sets the brightness value for a single pixel in the local buffer.
    ///
    /// For an 8-pixel board, valid values for pixel are 0-7. Valid
    /// values for brightness are 0.0-1.0, which is converted to a
    /// 5-bit value.
    pub fn set_pixel_brightness(&mut self, pixel: usize, brightness: f32) {
        if let Some(pixel) = self.pixels.get_mut(pixel) {
            pixel.brightness = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        }
    }

    /// Sets the red, green and blue values for all pixels in the local buffer.
    ///
    /// Valid values for red, green and blue are 0-255.
    pub fn set_all_pixels(&mut self, red: u8, green: u8, blue: u8) {
        for pixel in &mut self.pixels {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
        }
    }

    /// Sets the red, green, blue and brightness values for all pixels in the
    /// local buffer.
    ///
    /// Valid values for red, green and blue are 0-255. Valid values for
    /// brightness are 0.0-1.0, which is converted to a 5-bit value.
    pub fn set_all_pixels_rgbb(&mut self, red: u8, green: u8, blue: u8, brightness: f32) {
        let brightness: u8 = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        for pixel in &mut self.pixels {
            pixel.red = red;
            pixel.green = green;
            pixel.blue = blue;
            pixel.brightness = brightness;
        }
    }

    /// Sets the brightness value for all pixels in the local buffer.
    ///
    /// Valid values for brightness are 0.0-1.0, which is converted to a 5-bit
    /// value.
    pub fn set_all_pixels_brightness(&mut self, brightness: f32) {
        let brightness: u8 = (31.0 * if brightness > 1.0 { 1.0 } else if brightness < 0.0 { 0.0 } else { brightness }) as u8;
        for pixel in &mut self.pixels {
            pixel.brightness = brightness;
        }
    }

    /// Sets the red, green and blue values to 0 for all pixels in the local
    /// buffer.
    pub fn clear(&mut self) {
        self.set_all_pixels(0, 0, 0);
    }

    /// Sends the contents of the local buffer to the pixels, updating their
    /// LED colors and brightness.
    pub fn show(&self) {
        // Start frame (32x0)
        self.write_byte(0);
        self.write_byte(0);
        self.write_byte(0);
        self.write_byte(0);

        // LED frames
        for pixel in &self.pixels {
            self.write_byte(0b11100000 | pixel.brightness); // 3-bit header + 5-bit brightness
            self.write_byte(pixel.blue);
            self.write_byte(pixel.green);
            self.write_byte(pixel.red);
        }

        // We send another start frame immediately after our end frame, because
        // the APA102 version with a smaller, darker die, won't update the
        // pixels until it receives the next start frame. We still start show()
        // with a start frame, basically sending it twice, in case the user
        // connects a Blinkt! while the code is already running. This workaround
        // is compatible with both the normal and the 'different' APA102.
        self.gpio.write(self.pin_data, Level::Low);
        for _ in 0..32 + self.endframe_pulses {
            self.gpio.write(self.pin_clock, Level::High);
            self.gpio.write(self.pin_clock, Level::Low);
        }
    }

    fn write_byte(&self, byte: u8) {
        for n in 0..8 {
            if (byte & (1 << (7 - n))) > 0 {
                self.gpio.write(self.pin_data, Level::High);
            } else {
                self.gpio.write(self.pin_data, Level::Low);
            }

            self.gpio.write(self.pin_clock, Level::High);
            self.gpio.write(self.pin_clock, Level::Low);
        }
    }
}

impl Drop for Blinkt {
    fn drop(&mut self) {
        self.cleanup();
    }
}

#[test]
fn test_new() {
    let mut blinkt = match Blinkt::new() {
        // GPIO errors are acceptable, since they're likely caused by outside
        // distro/filesystem issues.
        Err(Error::GPIO(_)) => return,
        Ok(blinkt) => blinkt,
    };

    blinkt.set_clear_on_drop(false);
}
<|endoftext|>


<commit_before>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Debug)]
pub struct NodeSimilarityMatrix {
    // number of nodes in graph A
    na: usize,
    // number of nodes in graph B
    nb: usize,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iters: usize,
}

impl NodeSimilarityMatrix {
    pub fn init<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let (na, nb) = (in_a.len(), in_b.len());
        assert!((na, nb) == (out_a.len(), out_b.len()));

        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(na, nb, |i, j| {
            node_color_scale((i, j)) *
            if in_a[i].len() + out_a[i].len() > 0 && in_b[j].len() + out_b[j].len() > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(na, nb);

        NodeSimilarityMatrix {
            na: na,
            nb: nb,
            current: x,
            previous: new_x,
            num_iters: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    fn next<F>(&mut self,
               in_a: &[Vec<Idx>],
               in_b: &[Vec<Idx>],
               out_a: &[Vec<Idx>],
               out_b: &[Vec<Idx>],
               node_color_scale: &F)
        where F: Fn((usize, usize)) -> f32
    {
        next_x(&self.current,
               &mut self.previous,
               in_a,
               in_b,
               out_a,
               out_b,
               node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iters += 1;
    }

    #[inline]
    /// Calculates the similarity matrix for two graphs A and B.
    ///
    /// `in_a`:  Incoming edge list for each node of graph A
    /// `in_b`:  Incoming edge list for each node of graph B
    /// `out_a`: Outgoing edge list for each node of graph A
    /// `out_b`: Outgoing edge list for each node of graph B
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn calc<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   eps: f32,
                   stop_after_iter: usize,
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let mut x = NodeSimilarityMatrix::init(in_a, in_b, out_a, out_b, node_color_scale);

        for _ in 0..stop_after_iter {
            if x.in_eps(eps) {
                break;
            }
            x.next(in_a, in_b, out_a, out_b, node_color_scale);
        }

        return x;
    }

    fn matrix(self) -> DMat<f32> {
        self.current
    }

    pub fn num_iterations(&self) -> usize {
        self.num_iters
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        let score: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[ab]);
        score
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);
    println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 1, &|_| 1.0);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<commit_msg>Factor out optimal_node_assignment()<commit_after>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Debug)]
pub struct NodeSimilarityMatrix {
    // number of nodes in graph A
    na: usize,
    // number of nodes in graph B
    nb: usize,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iters: usize,
}

impl NodeSimilarityMatrix {
    pub fn init<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let (na, nb) = (in_a.len(), in_b.len());
        assert!((na, nb) == (out_a.len(), out_b.len()));

        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(na, nb, |i, j| {
            node_color_scale((i, j)) *
            if in_a[i].len() + out_a[i].len() > 0 && in_b[j].len() + out_b[j].len() > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(na, nb);

        NodeSimilarityMatrix {
            na: na,
            nb: nb,
            current: x,
            previous: new_x,
            num_iters: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    fn next<F>(&mut self,
               in_a: &[Vec<Idx>],
               in_b: &[Vec<Idx>],
               out_a: &[Vec<Idx>],
               out_b: &[Vec<Idx>],
               node_color_scale: &F)
        where F: Fn((usize, usize)) -> f32
    {
        next_x(&self.current,
               &mut self.previous,
               in_a,
               in_b,
               out_a,
               out_b,
               node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iters += 1;
    }

    #[inline]
    /// Calculates the similarity matrix for two graphs A and B.
    ///
    /// `in_a`:  Incoming edge list for each node of graph A
    /// `in_b`:  Incoming edge list for each node of graph B
    /// `out_a`: Outgoing edge list for each node of graph A
    /// `out_b`: Outgoing edge list for each node of graph B
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn calc<F>(in_a: &[Vec<Idx>],
                   in_b: &[Vec<Idx>],
                   out_a: &[Vec<Idx>],
                   out_b: &[Vec<Idx>],
                   eps: f32,
                   stop_after_iter: usize,
                   node_color_scale: &F)
                   -> NodeSimilarityMatrix
        where F: Fn((usize, usize)) -> f32
    {
        let mut x = NodeSimilarityMatrix::init(in_a, in_b, out_a, out_b, node_color_scale);

        for _ in 0..stop_after_iter {
            if x.in_eps(eps) {
                break;
            }
            x.next(in_a, in_b, out_a, out_b, node_color_scale);
        }

        return x;
    }

    fn matrix(self) -> DMat<f32> {
        self.current
    }

    pub fn num_iterations(&self) -> usize {
        self.num_iters
    }

    fn optimal_node_assignment(&self, n: usize) -> Vec<(usize, usize)> {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        assignment
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        self.optimal_node_assignment(n).iter().fold(0.0, |acc, &ab| acc + self.current[ab])
    }

    /// Calculate a measure how good the edge weights match up.
    ///
    /// We start by calculating the optimal node assignment between nodes of graph A and
    /// graph B, then compare all outgoing edges of similar nodes by again using an
    /// assignment.
    pub fn score_outgoing_edge_weights(&self) -> f32 {
        // XXX
        0.0
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);
    println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 1, &|_| 1.0);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let s = NodeSimilarityMatrix::calc(&in_a, &in_b, &out_a, &out_b, 0.1, 100, &|_| 1.0);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<|endoftext|>


<commit_before>//! # The rulinalg crate.
//!
//! A crate that provides high-dimensional linear algebra
//! implemented entirely in Rust.
//!
//! ---
//!
//! This crate provides two core data structures: `Matrix` and
//! `Vector`. These structs are designed to behave as you would expect
//! with relevant operator overloading.
//!
//! The library currently contains (at least) the following linear algebra
//! methods:
//!
//! - Matrix inversion
//! - LUP decomposition
//! - QR decomposition
//! - SVD decomposition
//! - Cholesky decomposition
//! - Eigenvalue decomposition
//! - Upper Hessenberg decomposition
//! - Linear system solver
//! - Other standard transformations, e.g. Transposing, concatenation, etc.
//!
//! ---
//!
//! ## Usage
//!
//! Specific usage of modules is described within the modules themselves. This section
//! will highlight the basic usage.
//!
//! We can create new matrices.
//!
//! ```
//! use rulinalg::matrix::Matrix;
//!
//! // A new matrix with 3 rows and 2 columns.
//! let a = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
//! ```
//!
//! The matrices are stored in row-major order. This means in the example above the top
//! row will be [1,2,3].
//!
//! We can perform operations on matrices.
//!
//! ```
//! use rulinalg::matrix::Matrix;
//!
//! // A new matrix with 3 rows and 2 columns.
//! let a = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
//! let b = Matrix::new(3, 2, vec![6.0, 5.0, 4.0, 3.0, 2.0, 1.0]);
//!
//! // Produces a 3x2 matrix filled with sevens.
//! let c = a + b;
//! ```
//!
//! Sometimes we want to construct small matrices by hand, usually for writing unit tests
//! or examples. For this purpose, `rulinalg` provides the `matrix!` macro:
//!
//! ```
//! // Remember to enable macro usage in rulinalg!
//! #[macro_use]
//! extern crate rulinalg;
//!
//! # fn main() {
//! // Construct a 3x3 matrix of f64
//! // Commas separate columns and semi-colons separate rows
//! let mat = matrix![1.0, 2.0, 3.0;
//!                   4.0, 5.0, 6.0;
//!                   7.0, 8.0, 9.0];
//! # }
//! ```
//!
//! Of course the library can support more complex operations but you should check the individual
//! modules for more information.
//!
//! # Matrix Slices
//!
//! Often times it is desirable to operate on only a sub-section of a `Matrix` without copying this block.
//! Rulinalg allows this via the `MatrixSlice` and `MatrixSliceMut` structs. These structs can be created
//! from `Matrix` structs and follow all of the borrowing rules of Rust.
//!
//! Note finally that much of the `Matrix`/`MatrixSlice`/`MatrixSliceMut` functionality is contained behind
//! the `BaseMatrix`/`BaseMatrixMut` traits. This allows us to be generic over matrices or slices.

#![deny(missing_docs)]
#![warn(missing_debug_implementations)]

extern crate num as libnum;
extern crate matrixmultiply;

// macros should be at the top in order for macros to be accessible in subsequent modules
#[macro_use]
pub mod macros;
pub mod matrix;
pub mod convert;
pub mod error;
pub mod utils;
pub mod vector;
pub mod ulp;
pub mod norm;

mod internal_utils;

#[cfg(test)]
mod testsupport;

#[cfg(test)]
#[macro_use]
extern crate quickcheck;

#[cfg(test)]
extern crate itertools;

pub use norm::{VectorNorm, MatrixNorm};
pub use norm::{VectorMetric, MatrixMetric};

#[cfg(feature = "io")]
extern crate csv as libcsv;

#[cfg(feature = "io")]
extern crate rustc_serialize;

#[cfg(feature = "io")]
pub mod io;
<commit_msg>Fix matrix example in docs (#185)<commit_after>//! # The rulinalg crate.
//!
//! A crate that provides high-dimensional linear algebra
//! implemented entirely in Rust.
//!
//! ---
//!
//! This crate provides two core data structures: `Matrix` and
//! `Vector`. These structs are designed to behave as you would expect
//! with relevant operator overloading.
//!
//! The library currently contains (at least) the following linear algebra
//! methods:
//!
//! - Matrix inversion
//! - LUP decomposition
//! - QR decomposition
//! - SVD decomposition
//! - Cholesky decomposition
//! - Eigenvalue decomposition
//! - Upper Hessenberg decomposition
//! - Linear system solver
//! - Other standard transformations, e.g. Transposing, concatenation, etc.
//!
//! ---
//!
//! ## Usage
//!
//! Specific usage of modules is described within the modules themselves. This section
//! will highlight the basic usage.
//!
//! We can create new matrices.
//!
//! ```
//! use rulinalg::matrix::Matrix;
//!
//! // A new matrix with 3 rows and 2 columns.
//! let a = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
//! ```
//!
//! The matrices are stored in row-major order. This means in the example above the top
//! row will be [1, 2].
//!
//! We can perform operations on matrices.
//!
//! ```
//! use rulinalg::matrix::Matrix;
//!
//! // A new matrix with 3 rows and 2 columns.
//! let a = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
//! let b = Matrix::new(3, 2, vec![6.0, 5.0, 4.0, 3.0, 2.0, 1.0]);
//!
//! // Produces a 3x2 matrix filled with sevens.
//! let c = a + b;
//! ```
//!
//! Sometimes we want to construct small matrices by hand, usually for writing unit tests
//! or examples. For this purpose, `rulinalg` provides the `matrix!` macro:
//!
//! ```
//! // Remember to enable macro usage in rulinalg!
//! #[macro_use]
//! extern crate rulinalg;
//!
//! # fn main() {
//! // Construct a 3x3 matrix of f64
//! // Commas separate columns and semi-colons separate rows
//! let mat = matrix![1.0, 2.0, 3.0;
//!                   4.0, 5.0, 6.0;
//!                   7.0, 8.0, 9.0];
//! # }
//! ```
//!
//! Of course the library can support more complex operations but you should check the individual
//! modules for more information.
//!
//! # Matrix Slices
//!
//! Often times it is desirable to operate on only a sub-section of a `Matrix` without copying this block.
//! Rulinalg allows this via the `MatrixSlice` and `MatrixSliceMut` structs. These structs can be created
//! from `Matrix` structs and follow all of the borrowing rules of Rust.
//!
//! Note finally that much of the `Matrix`/`MatrixSlice`/`MatrixSliceMut` functionality is contained behind
//! the `BaseMatrix`/`BaseMatrixMut` traits. This allows us to be generic over matrices or slices.

#![deny(missing_docs)]
#![warn(missing_debug_implementations)]

extern crate num as libnum;
extern crate matrixmultiply;

// macros should be at the top in order for macros to be accessible in subsequent modules
#[macro_use]
pub mod macros;
pub mod matrix;
pub mod convert;
pub mod error;
pub mod utils;
pub mod vector;
pub mod ulp;
pub mod norm;

mod internal_utils;

#[cfg(test)]
mod testsupport;

#[cfg(test)]
#[macro_use]
extern crate quickcheck;

#[cfg(test)]
extern crate itertools;

pub use norm::{VectorNorm, MatrixNorm};
pub use norm::{VectorMetric, MatrixMetric};

#[cfg(feature = "io")]
extern crate csv as libcsv;

#[cfg(feature = "io")]
extern crate rustc_serialize;

#[cfg(feature = "io")]
pub mod io;
<|endoftext|>


<commit_before>#[derive(Debug, PartialEq, Eq)]
struct GitStatus {
    branch: Option<BranchInfo>,
    upstream: Option<UpstreamInfo>,
    staged: bool,
    unstaged: bool,
    untracked: bool,
    ignored: bool,
}

#[derive(Debug, PartialEq, Eq)]
struct BranchInfo {
    object_id: String,
    branch: Option<String>,
}

#[derive(Debug, PartialEq, Eq)]
struct UpstreamInfo {
    upstream: String,
    ahead: bool,
    behind: bool,
}

impl GitStatus {
    fn new(status_txt: &str) -> Result<GitStatus, &'static str> {
        let mut s = GitStatus {
            branch: None,
            upstream: None,
            staged: false,
            unstaged: false,
            untracked: false,
            ignored: false,
        };
        for line in status_txt.lines() {
            let mut words = line.split(' ');
            match words.next() {
                Some("#") => {
                },
                Some("1") | Some("2") => {
                    if let Some(changes) = words.next() {
                        let mut changes = changes.chars();
                        match changes.next() {
                            Some('.') => {},
                            Some(_) => s.staged = true,
                            _ => {},
                        }
                        match changes.next() {
                            Some('.') => {},
                            Some(_) => s.unstaged = true,
                            _ => {},
                        }
                    }
                },
                Some("?") => s.untracked = true,
                Some("!") => s.ignored = true,
                _ => {},
            }
        }
        Ok(s)
    }
}

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

    #[test]
    fn parse_clean() {
        let test_status = "\n";
        assert_eq!(
            GitStatus::new(test_status).unwrap(),
            GitStatus {
                branch: None,
                upstream: None,
                staged: false,
                unstaged: false,
                untracked: false,
                ignored: false,
            }
        );
    }

    #[test]
    fn parse_branch() {
        let test_status = "\
# branch.oid 3845e7a3c3aadaaebb2d1b261bf07a9357d35a79
# branch.head master
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().branch,
            Some(BranchInfo {
                object_id: String::from("3845e7a3c3aadaaebb2d1b261bf07a9357d35a79"),
                branch: Some(String::from("master")),
            })
        );
    }

    #[test]
    fn parse_detached_head() {
        let test_status = "\
# branch.oid 2a7bb916bba69a5fc9d428acc80b1bce64e3e0bc
# branch.head (detached)
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().branch,
            Some(BranchInfo {
                object_id: String::from("2a7bb916bba69a5fc9d428acc80b1bce64e3e0bc"),
                branch: None
            })
        );
    }

    #[test]
    fn parse_upstream() {
        let test_status = "\
# branch.upstream origin/master
# branch.ab +1 -1
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().upstream,
            Some(UpstreamInfo {
                upstream: String::from("origin/master"),
                ahead: true,
                behind: true,
            })
        );
    }

    #[test]
    fn parse_staged_only() {
        let test_status = "1 M. N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 311c77295c5b6056f4599c2b8d0a019d4c76746a README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_unstaged_only() {
        let test_status = "1 .M N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(!s.staged);
        assert!(s.unstaged);
    }

    #[test]
    fn parse_both_staged_and_unstaged() {
        let test_status = "1 MM N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 311c77295c5b6056f4599c2b8d0a019d4c76746a README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(s.unstaged);
    }

    #[test]
    fn parse_deleted() {
        let test_status = "1 D. N... 100644 000000 000000 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 0000000000000000000000000000000000000000 README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_renamed() {
        let test_status = "2 R. N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 R100 README	README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_untracked() {
        let test_status = "? foo.txt\n";
        assert!(GitStatus::new(test_status).unwrap().untracked);
    }

    #[test]
    fn parse_ignored() {
        let test_status = "! node_modules/\n";
        assert!(GitStatus::new(test_status).unwrap().ignored);
    }

    #[test]
    fn parse_very_dirty() {
        let test_status = "\
# branch.oid 3845e7a3c3aadaaebb2d1b261bf07a9357d35a79
# branch.head master
# branch.upstream origin/master
# branch.ab +1 -1
1 D. N... 100644 000000 000000 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 0000000000000000000000000000000000000000 README.md
1 .M N... 100644 100644 100644 5e8a8090976077ddf16252a560460a20dbbdd6a5 5e8a8090976077ddf16252a560460a20dbbdd6a5 gh-pages.sh
? foo.txt
! ignored.txt
";
        assert_eq!(
            GitStatus::new(test_status).unwrap(),
            GitStatus {
                branch: Some(BranchInfo {
                    object_id: String::from("3845e7a3c3aadaaebb2d1b261bf07a9357d35a79"),
                    branch: Some(String::from("master")),
                }),
                upstream: Some(UpstreamInfo {
                    upstream: String::from("origin/master"),
                    ahead: true,
                    behind: true,
                }),
                staged: true,
                unstaged: true,
                untracked: true,
                ignored: true,
            }
        );
    }
}
<commit_msg>Refactor upstream to be inside branch info<commit_after>#[derive(Debug, PartialEq, Eq)]
struct GitStatus {
    branch: Option<BranchInfo>,
    staged: bool,
    unstaged: bool,
    untracked: bool,
    ignored: bool,
}

#[derive(Debug, PartialEq, Eq)]
struct BranchInfo {
    object_id: String,
    name: Option<String>,
    upstream: Option<UpstreamInfo>,
}

#[derive(Debug, PartialEq, Eq)]
struct UpstreamInfo {
    name: String,
    ahead: bool,
    behind: bool,
}

impl GitStatus {
    fn new(status_txt: &str) -> Result<GitStatus, &'static str> {
        let mut s = GitStatus {
            branch: None,
            staged: false,
            unstaged: false,
            untracked: false,
            ignored: false,
        };
        for line in status_txt.lines() {
            let mut words = line.split(' ');
            match words.next() {
                Some("#") => {
                },
                Some("1") | Some("2") => {
                    if let Some(changes) = words.next() {
                        let mut changes = changes.chars();
                        match changes.next() {
                            Some('.') => {},
                            Some(_) => s.staged = true,
                            _ => {},
                        }
                        match changes.next() {
                            Some('.') => {},
                            Some(_) => s.unstaged = true,
                            _ => {},
                        }
                    }
                },
                Some("?") => s.untracked = true,
                Some("!") => s.ignored = true,
                _ => {},
            }
        }
        Ok(s)
    }
}

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

    #[test]
    fn parse_clean() {
        let test_status = "\n";
        assert_eq!(
            GitStatus::new(test_status).unwrap(),
            GitStatus {
                branch: None,
                staged: false,
                unstaged: false,
                untracked: false,
                ignored: false,
            }
        );
    }

    #[test]
    fn parse_branch() {
        let test_status = "\
# branch.oid 3845e7a3c3aadaaebb2d1b261bf07a9357d35a79
# branch.head master
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().branch,
            Some(BranchInfo {
                object_id: String::from("3845e7a3c3aadaaebb2d1b261bf07a9357d35a79"),
                name: Some(String::from("master")),
                upstream: None,

            })
        );
    }

    #[test]
    fn parse_detached_head() {
        let test_status = "\
# branch.oid 2a7bb916bba69a5fc9d428acc80b1bce64e3e0bc
# branch.head (detached)
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().branch,
            Some(BranchInfo {
                object_id: String::from("2a7bb916bba69a5fc9d428acc80b1bce64e3e0bc"),
                name: None,
                upstream: None,
            })
        );
    }

    #[test]
    fn parse_upstream() {
        let test_status = "\
# branch.oid 3845e7a3c3aadaaebb2d1b261bf07a9357d35a79
# branch.head master
# branch.upstream origin/master
# branch.ab +1 -1
";
        assert_eq!(
            GitStatus::new(test_status).unwrap().branch,
            Some(BranchInfo {
                object_id: String::from("3845e7a3c3aadaaebb2d1b261bf07a9357d35a79"),
                name: Some(String::from("master")),
                upstream: Some(UpstreamInfo {
                    name: String::from("origin/master"),
                    ahead: true,
                    behind: true,
                }),
            })
        );
    }

    #[test]
    fn parse_staged_only() {
        let test_status = "1 M. N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 311c77295c5b6056f4599c2b8d0a019d4c76746a README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_unstaged_only() {
        let test_status = "1 .M N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(!s.staged);
        assert!(s.unstaged);
    }

    #[test]
    fn parse_both_staged_and_unstaged() {
        let test_status = "1 MM N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 311c77295c5b6056f4599c2b8d0a019d4c76746a README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(s.unstaged);
    }

    #[test]
    fn parse_deleted() {
        let test_status = "1 D. N... 100644 000000 000000 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 0000000000000000000000000000000000000000 README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_renamed() {
        let test_status = "2 R. N... 100644 100644 100644 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 R100 README	README.md\n";
        let s = GitStatus::new(test_status).unwrap();
        assert!(s.staged);
        assert!(!s.unstaged);
    }

    #[test]
    fn parse_untracked() {
        let test_status = "? foo.txt\n";
        assert!(GitStatus::new(test_status).unwrap().untracked);
    }

    #[test]
    fn parse_ignored() {
        let test_status = "! node_modules/\n";
        assert!(GitStatus::new(test_status).unwrap().ignored);
    }

    #[test]
    fn parse_very_dirty() {
        let test_status = "\
# branch.oid 3845e7a3c3aadaaebb2d1b261bf07a9357d35a79
# branch.head master
# branch.upstream origin/master
# branch.ab +1 -1
1 D. N... 100644 000000 000000 1290f45e7ad7575848a436d8febbd6c4ba07f1f3 0000000000000000000000000000000000000000 README.md
1 .M N... 100644 100644 100644 5e8a8090976077ddf16252a560460a20dbbdd6a5 5e8a8090976077ddf16252a560460a20dbbdd6a5 gh-pages.sh
? foo.txt
! ignored.txt
";
        assert_eq!(
            GitStatus::new(test_status).unwrap(),
            GitStatus {
                branch: Some(BranchInfo {
                    object_id: String::from("3845e7a3c3aadaaebb2d1b261bf07a9357d35a79"),
                    name: Some(String::from("master")),
                    upstream: Some(UpstreamInfo {
                        name: String::from("origin/master"),
                        ahead: true,
                        behind: true,
                    }),
                }),
                staged: true,
                unstaged: true,
                untracked: true,
                ignored: true,
            }
        );
    }
}
<|endoftext|>


<commit_before>#[test]
fn it_works() {
}
<commit_msg>Write Rust usage stats backend<commit_after>#![feature(cstr_memory)]

extern crate libc;

use std::collections::HashMap;
use std::ffi::{CStr, CString};
use libc::c_char;
use std::str;

#[derive(Debug)]
pub struct Usage {
    calls: HashMap<String, u32>,
    total: u32
}

impl Usage {
    fn new() -> Usage {
        Usage {
            calls: HashMap::new(),
            total: 0
        }
    }

    fn record(&mut self, method_name: String) {
        *self.calls.entry(method_name).or_insert(0) += 1;
        self.total += 1;
    }
}

#[derive(Debug)]
pub struct CallCount {
    count: u32,
    method_name: String
}

impl CallCount {
    fn new(method: &str, count: &u32) -> CallCount {
        CallCount {
            method_name: method.to_string(),
            count: *count
        }
    }

    fn to_c_struct(&self) -> CCallCount {
        let method = self.method_name.clone();
        let method = CString::new(method).unwrap();
        CCallCount {
            method_name: method.into_ptr(),
            count: self.count
        }
    }
}

#[repr(C)]
pub struct CCallCount {
    count: u32,
    method_name: *const c_char
}

#[repr(C)]
pub struct Report {
    length: u32,
    call_counts: *const CCallCount
}

#[no_mangle]
pub extern "C" fn new_usage() -> Box<Usage> {
    Box::new(Usage::new())
}

#[no_mangle]
pub extern "C" fn record(usage: &mut Usage,
                         event: *const c_char,
                         file: *const c_char,
                         line: u32,
                         id: *const c_char,
                         classname: *const c_char) {
    let event = unsafe { CStr::from_ptr(event) };
    let event = str::from_utf8(event.to_bytes()).unwrap();

    if event == "call" || event == "c-call" {
        let method_name = unsafe { CStr::from_ptr(id) };
        let method_name = str::from_utf8(method_name.to_bytes()).unwrap();
        let class_name = unsafe { CStr::from_ptr(classname) };
        let class_name = str::from_utf8(class_name.to_bytes()).unwrap();
        usage.record(class_name.to_string() + "#" + method_name)
    }
}

#[no_mangle]
pub extern "C" fn report(usage: &mut Usage) -> Report {
    let mut counts: Vec<CallCount> = usage.calls
        .iter()
        .map(|(method, count)| CallCount::new(method, count) )
        .collect();
    counts.sort_by(|a, b| a.count.cmp(&b.count).reverse());

    let c_counts: Vec<CCallCount> = counts
        .iter()
        .map(|cc| cc.to_c_struct())
        .collect();

    Report {
        length: c_counts.len() as u32,
        call_counts: c_counts.as_ptr()
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015-2018 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! A client for ElasticSearch's REST API
//!
//! The `Client` itself is used as the central access point, from which numerous
//! operations are defined implementing each of the specific ElasticSearch APIs.
//!
//! Warning: at the time of writing the majority of such APIs are currently
//! unimplemented.

#[macro_use]
pub mod util;

#[macro_use]
pub mod json;

pub mod error;
pub mod operations;
pub mod query;
pub mod units;

use std::time;

use reqwest::{header::CONTENT_TYPE, RequestBuilder, StatusCode, Url};

use serde::{de::DeserializeOwned, ser::Serialize};

use crate::error::EsError;

pub trait EsResponse {
    fn status_code(&self) -> StatusCode;
    fn read_response<R>(self) -> Result<R, EsError>
    where
        R: DeserializeOwned;
}

impl EsResponse for reqwest::Response {
    fn status_code(&self) -> StatusCode {
        self.status()
    }

    fn read_response<R>(self) -> Result<R, EsError>
    where
        R: DeserializeOwned,
    {
        Ok(serde_json::from_reader(self)?)
    }
}

// The client

/// Process the result of an HTTP request, returning the status code and the
/// `Json` result (if the result had a body) or an `EsError` if there were any
/// errors
///
/// This function is exposed to allow extensions to certain operations, it is
/// not expected to be used by consumers of the library
fn do_req(resp: reqwest::Response) -> Result<reqwest::Response, EsError> {
    let mut resp = resp;
    let status = resp.status();
    match status {
        StatusCode::OK | StatusCode::CREATED | StatusCode::NOT_FOUND => Ok(resp),
        _ => Err(EsError::from(&mut resp)),
    }
}

/// The core of the ElasticSearch client, owns a HTTP connection.
///
/// Each instance of `Client` is reusable, but only one thread can use each one
/// at once.  This will be enforced by the borrow-checker as most methods are
/// defined on `&mut self`.
///
/// To create a `Client`, the URL needs to be specified.
///
/// Each ElasticSearch API operation is defined as a method on `Client`.  Any
/// compulsory parameters must be given as arguments to this method.  It returns
/// an operation builder that can be used to add any optional parameters.
///
/// Finally `send` is called to submit the operation:
///
/// # Examples
///
/// ```
/// use rs_es::Client;
///
/// let mut client = Client::init("http://localhost:9200");
/// ```
///
/// See the specific operations and their builder objects for details.
#[derive(Debug)]
pub struct Client {
    base_url: Url,
    http_client: reqwest::Client,
}

impl Client {
    fn do_es_op(
        &self,
        url: &str,
        action: impl FnOnce(Url) -> RequestBuilder,
    ) -> Result<reqwest::Response, EsError> {
        let url = self.full_url(url);
        let username = self.base_url.username();
        let mut method = action(url);
        if !username.is_empty() {
            method = method.basic_auth(username, self.base_url.password());
        }
        let result = method.header(CONTENT_TYPE, "application/json").send()?;
        do_req(result)
    }
}

/// Create a HTTP function for the given method (GET/PUT/POST/DELETE)
macro_rules! es_op {
    ($n:ident,$cn:ident) => {
        fn $n(&self, url: &str) -> Result<reqwest::Response, EsError> {
            log::info!("Doing {} on {}", stringify!($n), url);
            self.do_es_op(url, |url| self.http_client.$cn(url.clone()))
        }
    }
}

/// Create a HTTP function with a request body for the given method
/// (GET/PUT/POST/DELETE)
///
macro_rules! es_body_op {
    ($n:ident,$cn:ident) => {
        fn $n<E>(&mut self, url: &str, body: &E) -> Result<reqwest::Response, EsError>
            where E: Serialize {

            log::info!("Doing {} on {}", stringify!($n), url);
            let json_string = serde_json::to_string(body)?;

            self.do_es_op(url, |url| {
                self.http_client.$cn(url.clone()).body(json_string)
            })
        }
    }
}

impl Client {
    /// Create a new client
    pub fn init(url_s: &str) -> Result<Client, reqwest::UrlError> {
        let url = Url::parse(url_s)?;

        Ok(Client {
            http_client: reqwest::Client::new(),
            base_url: url,
        })
    }

    // TODO - this should be replaced with a builder object, especially if more options are going
    // to be allowed
    pub fn init_with_timeout(
        url_s: &str,
        timeout: Option<time::Duration>,
    ) -> Result<Client, reqwest::UrlError> {
        let url = Url::parse(url_s)?;

        Ok(Client {
            http_client: reqwest::Client::builder()
                .timeout(timeout)
                .build()
                .expect("Failed to build client"),
            base_url: url,
        })
    }

    /// Take a nearly complete ElasticSearch URL, and stick
    /// the URL on the front.
    pub fn full_url(&self, suffix: &str) -> Url {
        self.base_url.join(suffix).expect("Invalid URL created")
    }

    es_op!(get_op, get);

    es_op!(post_op, post);
    es_body_op!(post_body_op, post);
    es_op!(put_op, put);
    es_body_op!(put_body_op, put);
    es_op!(delete_op, delete);
}

#[cfg(test)]
pub mod tests {
    use std::env;

    use serde::{Deserialize, Serialize};
    use serde_json::Value;

    use super::{
        operations::{bulk::Action, search::ScanResult},
        query::Query,
        units::Duration,
        Client,
    };

    // test setup

    pub fn make_client() -> Client {
        let hostname = match env::var("ES_HOST") {
            Ok(val) => val,
            Err(_) => "http://localhost:9200".to_owned(),
        };
        Client::init(&hostname).unwrap()
    }

    #[derive(Debug, Serialize, Deserialize)]
    pub struct TestDocument {
        pub str_field: String,
        pub int_field: i64,
        pub bool_field: bool,
    }

    #[allow(clippy::new_without_default_derive)]
    impl TestDocument {
        pub fn new() -> TestDocument {
            TestDocument {
                str_field: "I am a test".to_owned(),
                int_field: 1,
                bool_field: true,
            }
        }

        pub fn with_str_field(mut self, s: &str) -> TestDocument {
            self.str_field = s.to_owned();
            self
        }

        pub fn with_int_field(mut self, i: i64) -> TestDocument {
            self.int_field = i;
            self
        }

        pub fn with_bool_field(mut self, b: bool) -> TestDocument {
            self.bool_field = b;
            self
        }
    }

    pub fn setup_test_data(client: &mut Client, index_name: &str) {
        // TODO - this should use the Bulk API
        let documents = vec![
            TestDocument::new()
                .with_str_field("Document A123")
                .with_int_field(1),
            TestDocument::new()
                .with_str_field("Document B456")
                .with_int_field(2),
            TestDocument::new()
                .with_str_field("Document 1ABC")
                .with_int_field(3),
        ];
        for doc in documents.iter() {
            client
                .index(index_name, "test_type")
                .with_doc(doc)
                .send()
                .unwrap();
        }
        client.refresh().with_indexes(&[index_name]).send().unwrap();
    }

    pub fn clean_db(mut client: &mut Client, test_idx: &str) {
        let scroll = Duration::minutes(1);
        let mut scan: ScanResult<Value> = match client
            .search_query()
            .with_indexes(&[test_idx])
            .with_query(&Query::build_match_all().build())
            .scan(&scroll)
        {
            Ok(scan) => scan,
            Err(e) => {
                log::warn!("Scan error: {:?}", e);
                return; // Ignore not-found errors
            }
        };

        loop {
            let page = scan.scroll(&mut client, &scroll).unwrap();
            let hits = page.hits.hits;
            if hits.is_empty() {
                break;
            }
            let actions: Vec<Action<()>> = hits
                .into_iter()
                .map(|hit| {
                    Action::delete(hit.id)
                        .with_index(test_idx)
                        .with_doc_type(hit.doc_type)
                })
                .collect();
            client.bulk(&actions).send().unwrap();
        }

        scan.close(&mut client).unwrap();
    }
}
<commit_msg>Add debug message with body content<commit_after>/*
 * Copyright 2015-2018 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! A client for ElasticSearch's REST API
//!
//! The `Client` itself is used as the central access point, from which numerous
//! operations are defined implementing each of the specific ElasticSearch APIs.
//!
//! Warning: at the time of writing the majority of such APIs are currently
//! unimplemented.

#[macro_use]
pub mod util;

#[macro_use]
pub mod json;

pub mod error;
pub mod operations;
pub mod query;
pub mod units;

use std::time;

use reqwest::{header::CONTENT_TYPE, RequestBuilder, StatusCode, Url};

use serde::{de::DeserializeOwned, ser::Serialize};

use crate::error::EsError;

pub trait EsResponse {
    fn status_code(&self) -> StatusCode;
    fn read_response<R>(self) -> Result<R, EsError>
    where
        R: DeserializeOwned;
}

impl EsResponse for reqwest::Response {
    fn status_code(&self) -> StatusCode {
        self.status()
    }

    fn read_response<R>(self) -> Result<R, EsError>
    where
        R: DeserializeOwned,
    {
        Ok(serde_json::from_reader(self)?)
    }
}

// The client

/// Process the result of an HTTP request, returning the status code and the
/// `Json` result (if the result had a body) or an `EsError` if there were any
/// errors
///
/// This function is exposed to allow extensions to certain operations, it is
/// not expected to be used by consumers of the library
fn do_req(resp: reqwest::Response) -> Result<reqwest::Response, EsError> {
    let mut resp = resp;
    let status = resp.status();
    match status {
        StatusCode::OK | StatusCode::CREATED | StatusCode::NOT_FOUND => Ok(resp),
        _ => Err(EsError::from(&mut resp)),
    }
}

/// The core of the ElasticSearch client, owns a HTTP connection.
///
/// Each instance of `Client` is reusable, but only one thread can use each one
/// at once.  This will be enforced by the borrow-checker as most methods are
/// defined on `&mut self`.
///
/// To create a `Client`, the URL needs to be specified.
///
/// Each ElasticSearch API operation is defined as a method on `Client`.  Any
/// compulsory parameters must be given as arguments to this method.  It returns
/// an operation builder that can be used to add any optional parameters.
///
/// Finally `send` is called to submit the operation:
///
/// # Examples
///
/// ```
/// use rs_es::Client;
///
/// let mut client = Client::init("http://localhost:9200");
/// ```
///
/// See the specific operations and their builder objects for details.
#[derive(Debug)]
pub struct Client {
    base_url: Url,
    http_client: reqwest::Client,
}

impl Client {
    fn do_es_op(
        &self,
        url: &str,
        action: impl FnOnce(Url) -> RequestBuilder,
    ) -> Result<reqwest::Response, EsError> {
        let url = self.full_url(url);
        let username = self.base_url.username();
        let mut method = action(url);
        if !username.is_empty() {
            method = method.basic_auth(username, self.base_url.password());
        }
        let result = method.header(CONTENT_TYPE, "application/json").send()?;
        do_req(result)
    }
}

/// Create a HTTP function for the given method (GET/PUT/POST/DELETE)
macro_rules! es_op {
    ($n:ident,$cn:ident) => {
        fn $n(&self, url: &str) -> Result<reqwest::Response, EsError> {
            log::info!("Doing {} on {}", stringify!($n), url);
            self.do_es_op(url, |url| self.http_client.$cn(url.clone()))
        }
    }
}

/// Create a HTTP function with a request body for the given method
/// (GET/PUT/POST/DELETE)
///
macro_rules! es_body_op {
    ($n:ident,$cn:ident) => {
        fn $n<E>(&mut self, url: &str, body: &E) -> Result<reqwest::Response, EsError>
            where E: Serialize {

            log::info!("Doing {} on {}", stringify!($n), url);
            let json_string = serde_json::to_string(body)?;
            log::debug!("With body: {}", &json_string);

            self.do_es_op(url, |url| {
                self.http_client.$cn(url.clone()).body(json_string)
            })
        }
    }
}

impl Client {
    /// Create a new client
    pub fn init(url_s: &str) -> Result<Client, reqwest::UrlError> {
        let url = Url::parse(url_s)?;

        Ok(Client {
            http_client: reqwest::Client::new(),
            base_url: url,
        })
    }

    // TODO - this should be replaced with a builder object, especially if more options are going
    // to be allowed
    pub fn init_with_timeout(
        url_s: &str,
        timeout: Option<time::Duration>,
    ) -> Result<Client, reqwest::UrlError> {
        let url = Url::parse(url_s)?;

        Ok(Client {
            http_client: reqwest::Client::builder()
                .timeout(timeout)
                .build()
                .expect("Failed to build client"),
            base_url: url,
        })
    }

    /// Take a nearly complete ElasticSearch URL, and stick
    /// the URL on the front.
    pub fn full_url(&self, suffix: &str) -> Url {
        self.base_url.join(suffix).expect("Invalid URL created")
    }

    es_op!(get_op, get);

    es_op!(post_op, post);
    es_body_op!(post_body_op, post);
    es_op!(put_op, put);
    es_body_op!(put_body_op, put);
    es_op!(delete_op, delete);
}

#[cfg(test)]
pub mod tests {
    use std::env;

    use serde::{Deserialize, Serialize};
    use serde_json::Value;

    use super::{
        operations::{bulk::Action, search::ScanResult},
        query::Query,
        units::Duration,
        Client,
    };

    // test setup

    pub fn make_client() -> Client {
        let hostname = match env::var("ES_HOST") {
            Ok(val) => val,
            Err(_) => "http://localhost:9200".to_owned(),
        };
        Client::init(&hostname).unwrap()
    }

    #[derive(Debug, Serialize, Deserialize)]
    pub struct TestDocument {
        pub str_field: String,
        pub int_field: i64,
        pub bool_field: bool,
    }

    #[allow(clippy::new_without_default_derive)]
    impl TestDocument {
        pub fn new() -> TestDocument {
            TestDocument {
                str_field: "I am a test".to_owned(),
                int_field: 1,
                bool_field: true,
            }
        }

        pub fn with_str_field(mut self, s: &str) -> TestDocument {
            self.str_field = s.to_owned();
            self
        }

        pub fn with_int_field(mut self, i: i64) -> TestDocument {
            self.int_field = i;
            self
        }

        pub fn with_bool_field(mut self, b: bool) -> TestDocument {
            self.bool_field = b;
            self
        }
    }

    pub fn setup_test_data(client: &mut Client, index_name: &str) {
        // TODO - this should use the Bulk API
        let documents = vec![
            TestDocument::new()
                .with_str_field("Document A123")
                .with_int_field(1),
            TestDocument::new()
                .with_str_field("Document B456")
                .with_int_field(2),
            TestDocument::new()
                .with_str_field("Document 1ABC")
                .with_int_field(3),
        ];
        for doc in documents.iter() {
            client
                .index(index_name, "test_type")
                .with_doc(doc)
                .send()
                .unwrap();
        }
        client.refresh().with_indexes(&[index_name]).send().unwrap();
    }

    pub fn clean_db(mut client: &mut Client, test_idx: &str) {
        let scroll = Duration::minutes(1);
        let mut scan: ScanResult<Value> = match client
            .search_query()
            .with_indexes(&[test_idx])
            .with_query(&Query::build_match_all().build())
            .scan(&scroll)
        {
            Ok(scan) => scan,
            Err(e) => {
                log::warn!("Scan error: {:?}", e);
                return; // Ignore not-found errors
            }
        };

        loop {
            let page = scan.scroll(&mut client, &scroll).unwrap();
            let hits = page.hits.hits;
            if hits.is_empty() {
                break;
            }
            let actions: Vec<Action<()>> = hits
                .into_iter()
                .map(|hit| {
                    Action::delete(hit.id)
                        .with_index(test_idx)
                        .with_doc_type(hit.doc_type)
                })
                .collect();
            client.bulk(&actions).send().unwrap();
        }

        scan.close(&mut client).unwrap();
    }
}
<|endoftext|>